diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index de4fded6ae6e66995aa9f1687a9d598017416f7a..3dad41a88c8212b7445c32f241d887306d3c19ad 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -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). @@ -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/README.md b/README.md index c754c3f0db088be8b638c8bc508e1dd765d960d8..3cdb6e478ddf4f18af7f81bb3e321510903beb9d 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,10 @@ ----------------- -| **`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) [ ![Download](https://api.bintray.com/packages/google/tensorflow/tensorflow/images/download.svg) ](https://bintray.com/google/tensorflow/tensorflow/_latestVersion) | + +| **`Documentation`** | **`Linux CPU`** | **`Linux GPU`** | **`Mac OS CPU`** | **`Windows CPU`** | **`Android`** | +|-----------------|---------------------|------------------|-------------------|---------------|---------------| +| [![Documentation](https://img.shields.io/badge/api-reference-blue.svg)](https://www.tensorflow.org/api_docs/) | [![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 @@ -21,19 +22,9 @@ organization for the purposes of conducting machine learning and deep neural networks research. The system is general enough to be applicable in a wide variety of other domains, as well. -**If you want to contribute to TensorFlow, be sure to review the [contribution -guidelines](CONTRIBUTING.md). This project adheres to TensorFlow's -[code of conduct](CODE_OF_CONDUCT.md). By participating, you are expected to -uphold this code.** - -**We use [GitHub issues](https://github.com/tensorflow/tensorflow/issues) for -tracking requests and bugs. So please see -[TensorFlow Discuss](https://groups.google.com/a/tensorflow.org/forum/#!forum/discuss) for general questions -and discussion, and please direct specific questions to [Stack Overflow](https://stackoverflow.com/questions/tagged/tensorflow).** - -The TensorFlow project strives to abide by generally accepted best practices in open-source software development: - -[![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/1486/badge)](https://bestpractices.coreinfrastructure.org/projects/1486) +Keep up to date with release announcements and security updates by +subscribing to +[announce@tensorflow.org](https://groups.google.com/a/tensorflow.org/forum/#!forum/announce). ## 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.* @@ -75,6 +66,22 @@ $ python >>> sess.close() ``` +## Contribution guidelines + +**If you want to contribute to TensorFlow, be sure to review the [contribution +guidelines](CONTRIBUTING.md). This project adheres to TensorFlow's +[code of conduct](CODE_OF_CONDUCT.md). By participating, you are expected to +uphold this code.** + +**We use [GitHub issues](https://github.com/tensorflow/tensorflow/issues) for +tracking requests and bugs. So please see +[TensorFlow Discuss](https://groups.google.com/a/tensorflow.org/forum/#!forum/discuss) for general questions +and discussion, and please direct specific questions to [Stack Overflow](https://stackoverflow.com/questions/tagged/tensorflow).** + +The TensorFlow project strives to abide by generally accepted best practices in open-source software development: + +[![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/1486/badge)](https://bestpractices.coreinfrastructure.org/projects/1486) + ## For more information * [TensorFlow Website](https://www.tensorflow.org) diff --git a/RELEASE.md b/RELEASE.md index e63971ef93e005773a55f8d5be0bba25794a6c7e..6f54dee58f75c29a16545ba25de12fe059baf1eb 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,7 +1,98 @@ +# 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 and cuDNN 7. +* 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. @@ -10,7 +101,7 @@ 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 and cuDNN 7 support. +* 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. @@ -123,6 +214,27 @@ * 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: @@ -521,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/tensorflow/SECURITY.md b/SECURITY.md similarity index 94% rename from tensorflow/SECURITY.md rename to SECURITY.md index 074eed2951526d53ab62515b7b869569a9708299..378e77696725e338e8289cda84dbc543303ae053 100644 --- a/tensorflow/SECURITY.md +++ b/SECURITY.md @@ -6,7 +6,7 @@ report vulnerabilities in TensorFlow. ## TensorFlow models are programs -TensorFlow's runtime system interprets and executes programs. What machine +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 @@ -28,12 +28,12 @@ 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.** +perspective.** ## Running untrusted models As a general rule: **Always** execute untrusted models inside a sandbox (e.g., -[nsjail](https://github.com/google/nsjail)). +[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. @@ -109,11 +109,11 @@ 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**. +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 +reduced permissions). 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. @@ -129,11 +129,11 @@ 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 +which exhibit unexpected or unwanted behavior. 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. +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 @@ -168,7 +168,7 @@ 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. +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 @@ -233,7 +233,7 @@ v//Fw6ZeY+HmRDFdirjD7wXtIuER4vqCryIqR6Xe9X8oJXz9L/Jhslc= ### Known vulnerabilities -| Type | Versions affected | Reported by | Additional Information | -|------|:-----------------:|---------------------------------------| -| out of bounds read| <=1.4 | @zhangbo5891001 | [issue report](https://github.com/tensorflow/tensorflow/issues/14959) | +| Type | Versions affected | Reported by | Additional Information | +|-------------------|:-----------------:|--------------------|-----------------------------| +| out of bounds read| <=1.4 | TenCent Blade Team | [issue report](https://github.com/tensorflow/tensorflow/issues/14959) | diff --git a/configure b/configure index 9c21d2b03a27714f05094667691e74c16fa89f35..66b66ba54ed68a9aa0ee556f84f68c3a83a495ab 100755 --- a/configure +++ b/configure @@ -8,7 +8,8 @@ if [ -z "$PYTHON_BIN_PATH" ]; then fi # Set all env variables -"$PYTHON_BIN_PATH" configure.py +CONFIGURE_DIR=$(dirname "$0") +"$PYTHON_BIN_PATH" "${CONFIGURE_DIR}/configure.py" "$@" echo "Configuration finished" diff --git a/configure.py b/configure.py index 27519b4aba013e53f688513b798de4238e545022..d14edef1be9e31137c96bed7aebf7ba158b3274f 100644 --- a/configure.py +++ b/configure.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import argparse import errno import os import platform @@ -32,10 +33,6 @@ except ImportError: from distutils.spawn import find_executable as which # pylint: enable=g-import-not-at-top -_TF_BAZELRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), - '.tf_configure.bazelrc') -_TF_WORKSPACE = os.path.join(os.path.dirname(os.path.abspath(__file__)), - 'WORKSPACE') _DEFAULT_CUDA_VERSION = '9.0' _DEFAULT_CUDNN_VERSION = '7' _DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,5.2' @@ -51,6 +48,11 @@ _SUPPORTED_ANDROID_NDK_VERSIONS = [10, 11, 12, 13, 14, 15] _DEFAULT_PROMPT_ASK_ATTEMPTS = 10 +_TF_WORKSPACE_ROOT = os.path.abspath(os.path.dirname(__file__)) +_TF_BAZELRC_FILENAME = '.tf_configure.bazelrc' +_TF_BAZELRC = os.path.join(_TF_WORKSPACE_ROOT, _TF_BAZELRC_FILENAME) +_TF_WORKSPACE = os.path.join(_TF_WORKSPACE_ROOT, 'WORKSPACE') + class UserInputError(Exception): pass @@ -119,22 +121,6 @@ def sed_in_place(filename, old, new): f.write(newdata) -def remove_line_with(filename, token): - """Remove lines that contain token from file. - - Args: - filename: string for filename. - token: string token to check if to remove a line from file or not. - """ - with open(filename, 'r') as f: - filedata = f.read() - - with open(filename, 'w') as f: - for line in filedata.strip().split('\n'): - if token not in line: - f.write(line + '\n') - - def write_to_bazelrc(line): with open(_TF_BAZELRC, 'a') as f: f.write(line + '\n') @@ -245,25 +231,30 @@ def setup_python(environ_cp): environ_cp['PYTHON_BIN_PATH'] = python_bin_path # Write tools/python_bin_path.sh - with open('tools/python_bin_path.sh', 'w') as f: + with open(os.path.join( + _TF_WORKSPACE_ROOT, 'tools', 'python_bin_path.sh'), 'w') as f: f.write('export PYTHON_BIN_PATH="%s"' % python_bin_path) -def reset_tf_configure_bazelrc(): +def reset_tf_configure_bazelrc(workspace_path): """Reset file that contains customized config settings.""" open(_TF_BAZELRC, 'w').close() - - home = os.path.expanduser('~') - if not os.path.exists('.bazelrc'): - if os.path.exists(os.path.join(home, '.bazelrc')): - with open('.bazelrc', 'a') as f: - f.write('import %s/.bazelrc\n' % home.replace('\\', '/')) + bazelrc_path = os.path.join(workspace_path, '.bazelrc') + + data = [] + if os.path.exists(bazelrc_path): + with open(bazelrc_path, 'r') as f: + data = f.read().splitlines() + with open(bazelrc_path, 'w') as f: + for l in data: + if _TF_BAZELRC_FILENAME in l: + continue + f.write('%s\n' % l) + if is_windows(): + tf_bazelrc_path = _TF_BAZELRC.replace("\\", "/") else: - open('.bazelrc', 'w').close() - - remove_line_with('.bazelrc', 'tf_configure') - with open('.bazelrc', 'a') as f: - f.write('import %workspace%/.tf_configure.bazelrc\n') + tf_bazelrc_path = _TF_BAZELRC + f.write('import %s\n' % tf_bazelrc_path) def cleanup_makefile(): @@ -271,7 +262,8 @@ def cleanup_makefile(): These files could interfere with Bazel parsing. """ - makefile_download_dir = 'tensorflow/contrib/makefile/downloads' + makefile_download_dir = os.path.join( + _TF_WORKSPACE_ROOT, 'tensorflow', 'contrib', 'makefile', 'downloads') if os.path.isdir(makefile_download_dir): for root, _, filenames in os.walk(makefile_download_dir): for f in filenames: @@ -445,7 +437,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. @@ -456,7 +448,7 @@ def check_bazel_version(min_version): if which('bazel') is None: print('Cannot find bazel. Please install bazel.') sys.exit(0) - curr_version = run_shell(['bazel', '--batch', 'version']) + curr_version = run_shell(['bazel', '--batch', '--bazelrc=/dev/null', 'version']) for line in curr_version.split('\n'): if 'Build label: ' in line: @@ -502,7 +494,8 @@ def set_cc_opt_flags(environ_cp): for opt in cc_opt_flags.split(): write_to_bazelrc('build:opt --copt=%s' % opt) # It should be safe on the same build host. - write_to_bazelrc('build:opt --host_copt=-march=native') + if not is_ppc64le(): + write_to_bazelrc('build:opt --host_copt=-march=native') write_to_bazelrc('build:opt --define with_default_optimizations=true') # TODO(mikecase): Remove these default defines once we are able to get # TF Lite targets building without them. @@ -827,6 +820,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 = ( @@ -837,6 +852,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 @@ -893,6 +909,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 ' @@ -1031,7 +1048,10 @@ def set_tf_tensorrt_install_path(environ_cp): for lib_file in possible_files: if is_compatible(lib_file, cuda_ver, cudnn_ver): - ver_str = nvinfer_pattern.search(lib_file).group(1) + matches = nvinfer_pattern.search(lib_file) + if len(matches.groups()) == 0: + continue + ver_str = matches.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] @@ -1054,12 +1074,22 @@ def set_tf_tensorrt_install_path(environ_cp): break # Reset and Retry - 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) + if 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 ' @@ -1195,7 +1225,7 @@ def set_host_c_compiler(environ_cp): environ_cp, var_name='HOST_C_COMPILER', var_default=default_c_host_compiler, - ask_for_var=('Please specify which C compiler should be used as the host' + ask_for_var=('Please specify which C compiler should be used as the host ' 'C compiler.'), check_success=os.path.exists, error_msg='Invalid C compiler path. %s cannot be found.', @@ -1339,13 +1369,20 @@ def config_info_line(name, help_text): def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--workspace", + type=str, + default=_TF_WORKSPACE_ROOT, + help="The absolute path to your active Bazel workspace.") + args = parser.parse_args() + # Make a copy of os.environ to be clear when functions and getting and setting # environment variables. environ_cp = dict(os.environ) - check_bazel_version('0.5.4') + check_bazel_version('0.10.0') - reset_tf_configure_bazelrc() + reset_tf_configure_bazelrc(args.workspace) cleanup_makefile() setup_python(environ_cp) @@ -1400,6 +1437,10 @@ def main(): 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': @@ -1453,7 +1494,6 @@ def main(): 'more details.') config_info_line('mkl', 'Build with MKL support.') config_info_line('monolithic', 'Config for mostly static monolithic build.') - config_info_line('tensorrt', 'Build with TensorRT support.') if __name__ == '__main__': main() diff --git a/tensorflow/BUILD b/tensorflow/BUILD index 9e69613c79ebd1d63ff052295cdb7acaaea5ff92..9932e5607685b5b8f5900bdfa42363151e57d3f1 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -415,6 +415,17 @@ py_library( deps = ["//tensorflow/python"], ) +py_library( + name = "experimental_tensorflow_py", + srcs = ["experimental_api.py"], + srcs_version = "PY2AND3", + visibility = ["//tensorflow/tools/api/tests:__subpackages__"], + deps = [ + "//tensorflow/python", + "//tensorflow/tools/api/generator:python_api", + ], +) + filegroup( name = "all_opensource_files", data = [ @@ -441,6 +452,7 @@ filegroup( "//tensorflow/compiler/xla:all_files", "//tensorflow/compiler/xla/client:all_files", "//tensorflow/compiler/xla/client/lib:all_files", + "//tensorflow/compiler/xla/client/xla_client:all_files", "//tensorflow/compiler/xla/legacy_flags:all_files", "//tensorflow/compiler/xla/python:all_files", "//tensorflow/compiler/xla/service:all_files", @@ -483,6 +495,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", @@ -542,7 +555,6 @@ filegroup( "//tensorflow/contrib/model_pruning:all_files", "//tensorflow/contrib/model_pruning/examples/cifar10: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", @@ -550,8 +562,10 @@ filegroup( "//tensorflow/contrib/predictor: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", @@ -785,6 +799,7 @@ tf_cc_shared_object( }), deps = [ "//tensorflow/c:c_api", + "//tensorflow/c:c_api_experimental", "//tensorflow/c:exported_symbols.lds", "//tensorflow/c:version_script.lds", "//tensorflow/c/eager:c_api", 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 c46cb32aa46af474c889095564d46c5f2399c3ad..c178d7f81f1d7f67eaa37c86c1fdc9307d32afe6 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -6,23 +6,21 @@ 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", ) -# For platform specific build config -load( - "//tensorflow/core:platform/default/build_config.bzl", - "tf_kernel_tests_linkstatic", -) - # ----------------------------------------------------------------------------- # Public targets filegroup( name = "headers", - srcs = ["c_api.h"], + srcs = [ + "c_api.h", + "c_api_experimental.h", + ], visibility = ["//tensorflow:__subpackages__"], ) @@ -33,7 +31,11 @@ filegroup( "*.cc", "*.h", ], - exclude = ["*test*"], + exclude = [ + "c_api_experimental.cc", + "c_api_experimental.h", + "*test*", + ], ), visibility = ["//visibility:public"], ) @@ -91,9 +93,36 @@ 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": [], }), ) +tf_cuda_library( + name = "c_api_experimental", + srcs = [ + "c_api_experimental.cc", + ], + hdrs = [ + "c_api_experimental.h", + ], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = [ + ":c_api", + ":c_api_internal", + "//tensorflow/compiler/jit/legacy_flags:mark_for_compilation_pass_flags", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + exports_files( [ "version_script.lds", @@ -135,15 +164,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", + ":c_api_experimental", "//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 3c7f041b39f01d9b8b187079b00e0c5ad99a38cc..778cb667e2c0015c6a768ecf3b12b82601764117 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -30,6 +30,7 @@ limitations under the License. #endif #include "tensorflow/c/c_api_internal.h" #include "tensorflow/core/common_runtime/device_mgr.h" +#include "tensorflow/core/common_runtime/eval_const_tensor.h" #include "tensorflow/core/common_runtime/shape_refiner.h" #include "tensorflow/core/framework/allocation_description.pb.h" #include "tensorflow/core/framework/log_memory.h" @@ -62,8 +63,10 @@ limitations under the License. // brain namespace because we are defining 'extern "C"' functions. using tensorflow::AllocationDescription; using tensorflow::DataType; +using tensorflow::ExtendSessionGraphHelper; using tensorflow::Graph; using tensorflow::GraphDef; +using tensorflow::mutex_lock; using tensorflow::NameRangeMap; using tensorflow::NameRangesForNode; using tensorflow::NewSession; @@ -72,11 +75,13 @@ using tensorflow::NodeBuilder; using tensorflow::NodeDef; using tensorflow::OpDef; using tensorflow::OpRegistry; +using tensorflow::OutputTensor; using tensorflow::PartialTensorShape; 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 +92,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 +112,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 +202,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 +218,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) { @@ -628,11 +641,11 @@ Status MessageToBuffer(const tensorflow::protobuf::Message& in, } void RecordMutation(TF_Graph* graph, const TF_Operation& op, - const char* mutation_type) - EXCLUSIVE_LOCKS_REQUIRED(graph->mu) { + const char* mutation_type) { // If any session has already run this node_id, mark this session as // unrunnable. for (auto it : graph->sessions) { + mutex_lock session_lock(it.first->mu); if (it.first->last_num_graph_nodes > op.node.id()) { it.second = FailedPrecondition( "Operation '", op.node.DebugString(), "' was changed by ", @@ -698,6 +711,60 @@ void TF_GraphSetOutputHandleShapesAndTypes(TF_Graph* graph, TF_Output output, Status LoadLibrary(const char* library_filename, void** result, const void** buf, size_t* len); +// TODO(josh11b,mrry): Change Session to be able to use a Graph* +// directly, instead of requiring us to serialize to a GraphDef and +// call Session::Extend(). +bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) { + if (session->graph != nullptr) { + // Take the graph lock before the session lock to avoid deadlock. This is + // safe since session->graph does not change. + session->graph->mu.lock(); + mutex_lock session_lock(session->mu); + const Graph& graph = session->graph->graph; + + status->status = session->graph->sessions[session]; + if (!status->status.ok()) { + session->graph->mu.unlock(); + return false; + } + + const auto num_nodes = graph.num_node_ids(); + if (session->last_num_graph_nodes < num_nodes) { + status->status = tensorflow::ValidateNoCycles(session->graph->graph); + if (!status->status.ok()) { + session->graph->mu.unlock(); + return false; + } + + GraphDef graph_def; + *graph_def.mutable_versions() = graph.versions(); + // Fill graph_def with nodes with ids in the range + // [session->last_num_graph_nodes, num_nodes), that is the nodes + // added since the last TF_SessionRun() call. + for (auto id = session->last_num_graph_nodes; id < num_nodes; ++id) { + Node* const node = graph.FindNodeId(id); + if (node != nullptr && node->IsOp()) { + NodeDef* const node_def = graph_def.add_node(); + *node_def = node->def(); + } + } + *graph_def.mutable_library() = graph.flib_def().ToProto(); + session->graph->mu.unlock(); + status->status = session->session->Extend(graph_def); + if (!status->status.ok()) { + // Contract is we always delete input_values[i]. + return false; + } + // Note: session->session is not modified if Extend() fails, so + // we only set last_num_graph_nodes if it succeeds. + session->last_num_graph_nodes = num_nodes; + } else { + session->graph->mu.unlock(); + } + } + return true; +} + } // namespace tensorflow static void TF_Run_Setup(int noutputs, TF_Tensor** c_outputs, @@ -2144,7 +2211,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)); @@ -2398,11 +2465,7 @@ void TF_AddGradients(TF_Graph* g, TF_Output* y, int ny, TF_Output* x, int nx, // TF_Session functions ---------------------------------------------- TF_Session::TF_Session(tensorflow::Session* s, TF_Graph* g) - : session(s), graph(g), last_num_graph_nodes(0), device_mgr(nullptr) { - if (s->LocalDeviceManager(&device_mgr).ok()) { - devices = device_mgr->ListDevices(); - } -} + : session(s), graph(g), last_num_graph_nodes(0), extend_before_run(true) {} TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opt, TF_Status* status) { @@ -2502,58 +2565,6 @@ void TF_DeleteSession(TF_Session* s, TF_Status* status) { delete s; } -// TODO(josh11b,mrry): Change Session to be able to use a Graph* -// directly, instead of requiring us to serialize to a GraphDef and -// call Session::Extend(). -static bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) { - if (session->graph != nullptr) { - mutex_lock session_lock(session->mu); - session->graph->mu.lock(); - const Graph& graph = session->graph->graph; - - status->status = session->graph->sessions[session]; - if (!status->status.ok()) { - session->graph->mu.unlock(); - return false; - } - - const auto num_nodes = graph.num_node_ids(); - if (session->last_num_graph_nodes < num_nodes) { - status->status = tensorflow::ValidateNoCycles(session->graph->graph); - if (!status->status.ok()) { - session->graph->mu.unlock(); - return false; - } - - GraphDef graph_def; - *graph_def.mutable_versions() = graph.versions(); - // Fill graph_def with nodes with ids in the range - // [session->last_num_graph_nodes, num_nodes), that is the nodes - // added since the last TF_SessionRun() call. - for (auto id = session->last_num_graph_nodes; id < num_nodes; ++id) { - Node* const node = graph.FindNodeId(id); - if (node != nullptr && node->IsOp()) { - NodeDef* const node_def = graph_def.add_node(); - *node_def = node->def(); - } - } - *graph_def.mutable_library() = graph.flib_def().ToProto(); - session->graph->mu.unlock(); - status->status = session->session->Extend(graph_def); - if (!status->status.ok()) { - // Contract is we always delete input_values[i]. - return false; - } - // Note: session->session is not modified if Extend() fails, so - // we only set last_num_graph_nodes if it succeeds. - session->last_num_graph_nodes = num_nodes; - } else { - session->graph->mu.unlock(); - } - } - return true; -} - void TF_SessionRun(TF_Session* session, const TF_Buffer* run_options, const TF_Output* inputs, TF_Tensor* const* input_values, int ninputs, const TF_Output* outputs, @@ -2563,7 +2574,8 @@ void TF_SessionRun(TF_Session* session, const TF_Buffer* run_options, // TODO(josh11b,mrry): Change Session to be able to use a Graph* // directly, instead of requiring us to serialize to a GraphDef and // call Session::Extend(). - if (!ExtendSessionGraphHelper(session, status)) { + if (session->extend_before_run && + !ExtendSessionGraphHelper(session, status)) { return; } @@ -2600,7 +2612,8 @@ void TF_SessionPRunSetup(TF_Session* session, const TF_Output* inputs, const char** handle, TF_Status* status) { *handle = nullptr; - if (!ExtendSessionGraphHelper(session, status)) { + if (session->extend_before_run && + !ExtendSessionGraphHelper(session, status)) { return; } @@ -2643,7 +2656,8 @@ void TF_SessionPRun(TF_Session* session, const char* handle, // TODO(josh11b,mrry): Change Session to be able to use a Graph* // directly, instead of requiring us to serialize to a GraphDef and // call Session::Extend(). - if (!ExtendSessionGraphHelper(session, status)) { + if (session->extend_before_run && + !ExtendSessionGraphHelper(session, status)) { return; } @@ -2672,6 +2686,24 @@ void TF_SessionPRun(TF_Session* session, const char* handle, output_values, target_names, nullptr, status); } +unsigned char TF_TryEvaluateConstant(TF_Graph* graph, TF_Output output, + TF_Tensor** result, TF_Status* status) { + *result = nullptr; + mutex_lock l(graph->mu); + OutputTensor tensor(&output.oper->node, output.index); + bool evaluated; + Tensor result_tensor; + status->status = EvaluateConstantTensor( + tensor, graph->refiner, *graph->graph.op_registry(), + graph->graph.versions().producer(), &evaluated, &result_tensor); + if (evaluated) { + DCHECK(status->status.ok()); + *result = TF_TensorFromTensor(result_tensor, status); + if (!status->status.ok()) evaluated = false; + } + return evaluated; +} + TF_ApiDefMap* TF_NewApiDefMap(TF_Buffer* op_list_buffer, TF_Status* status) { tensorflow::OpList op_list; if (!op_list.ParseFromArray(op_list_buffer->data, op_list_buffer->length)) { diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index d2e45341bf1b9ee4579f84064550ce26041dd04a..b32f574628c4d1dc5c3bb3f1265a1b12adee28bc 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), @@ -1271,23 +1275,33 @@ TF_CAPI_EXPORT extern void TF_FunctionGetAttrValueProto( // Deleting a function does not remove it from any graphs it was copied to. TF_CAPI_EXPORT extern void TF_DeleteFunction(TF_Function* func); +// Attempts to evaluate `output`. This will only be possible if `output` doesn't +// depend on any graph inputs (this function is safe to call if this isn't the +// case though). +// +// If the evaluation is successful, this function returns true and `output`s +// value is returned in `result`. Otherwise returns false. An error status is +// returned if something is wrong with the graph or input. Note that this may +// return false even if no error status is set. +TF_CAPI_EXPORT extern unsigned char TF_TryEvaluateConstant(TF_Graph* graph, + TF_Output output, + TF_Tensor** result, + TF_Status* status); + // TODO(josh11b): Register OpDef, available to all operations added // to this graph. -// The following two may both benefit from a subgraph-definition API -// that re-uses most of the graph-definition API. -// TODO(andydavis): Add functions to a graph. - // -------------------------------------------------------------------------- // API for driving Graph execution. 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_experimental.cc b/tensorflow/c/c_api_experimental.cc new file mode 100644 index 0000000000000000000000000000000000000000..f6d8949bb0d09fde539515aced17932cfca95fcb --- /dev/null +++ b/tensorflow/c/c_api_experimental.cc @@ -0,0 +1,385 @@ +/* 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/c/c_api_experimental.h" + +#include "tensorflow/c/c_api_internal.h" +#include "tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/protobuf/config.pb.h" + +using tensorflow::Node; +using tensorflow::NodeBuilder; +using tensorflow::NodeDef; +using tensorflow::Status; + +namespace { + +const char* const DEVICE_TPU_REPLICATED_CORE = "TPU_REPLICATED_CORE"; +const char* const DEVICE_TPU_SYSTEM = "TPU_SYSTEM"; + +TF_Operation* ToTF_Operation(Node* node) { + return static_cast(static_cast(node)); +} + +// Graph rewrite algorithm (modeled after the python TPU graph rewrite path): +// +// 1. For each input node I, feed it to a new TPUReplicatedInput node, which in +// turn feeds a new Identity node N, and store the mapping I->N. +// +// 2. Rewrite all existing graph nodes by adding a attribute on TPU cluster. For +// each node reading some input node I, rewire it to read from N instead based +// on the I->N mapping in step #1. +// +// 3. For each output node O, feed it to a new Identity node, which in turn +// feeds a new TPUReplicatedOutput node, which in turn feeds a new Identity node +// M. Return the set of new output nodes (the "M" nodes) for caller to fetch +// from. +// +// Limitations compared to the python TPU rewrite path: +// - # replicas is always 1. +// - Less error checking. +// +// TODO(hongm): Simplify the graph rewrite to generating fewer TPUReplicate +// related nodes. +class GraphRewriter { + public: + GraphRewriter(TF_Graph* graph, int num_input_nodes, + const TF_Output* input_nodes, int num_output_nodes, + const TF_Output* output_nodes) + EXCLUSIVE_LOCKS_REQUIRED(graph->mu) + : graph_(graph), input_nodes_(input_nodes) { + for (int i = 0; i < num_input_nodes; ++i) { + // Will fill in the value part later when we create the associated new + // input node. + input_node_map_[input_nodes[i].oper->node.name()] = nullptr; + } + + // Grab all existing nodes for the upcoming rewrite, before mutating the + // graph. + for (Node* n : graph->graph.nodes()) { + nodes_to_rewrite_.push_back(n); + } + + for (int i = 0; i < num_output_nodes; ++i) { + output_node_map_.emplace(output_nodes[i].oper->node.name(), + PortIndexPair{output_nodes[i].index, i}); + } + } + + // On success, sets `config_op` and `shutdown_op` to the corresponding + // "ConfigureDistributedTPU" and "ShutdownDistributedTPU" nodes added to the + // graph. + tensorflow::Status Rewrite(TF_Output* new_output_nodes, TF_Output* config_op, + TF_Output* shutdown_op) + EXCLUSIVE_LOCKS_REQUIRED(graph_->mu) { + TF_RETURN_IF_ERROR(ProcessInputNodes()); + + return RewriteGraphAndAddOutputNodes(new_output_nodes, config_op, + shutdown_op); + } + + private: + // Synthensizes new nodes for the input nodes, and creates a replicated + // metadata node. + tensorflow::Status ProcessInputNodes() EXCLUSIVE_LOCKS_REQUIRED(graph_->mu) { + Node* metadata_node; + TF_RETURN_IF_ERROR( + NodeBuilder(metadata_node_name_.c_str(), "TPUReplicateMetadata") + .Attr("num_replicas", 1) + .Attr("_tpu_replicate", cluster_name_.c_str()) + .Finalize(&graph_->graph, &metadata_node)); + + for (int i = 0; i < input_node_map_.size(); ++i) { + VLOG(1) << "Handling input node " << input_nodes_[i].oper->node.name(); + Node* replicated_input_node; + { + std::string replicated_input_name("TPUReplicate/input" + + std::to_string(i)); + NodeBuilder::NodeOut input(&input_nodes_[i].oper->node, + input_nodes_[i].index); + std::vector input_list; + input_list.push_back(input); + TF_RETURN_IF_ERROR( + NodeBuilder(replicated_input_name.c_str(), "TPUReplicatedInput") + // This op requires an input list. + .Input(input_list) + .Finalize(&graph_->graph, &replicated_input_node)); + } + + { + Node* new_input_node; + const std::string new_input_name("TPUReplicate/replicated_input_" + + std::to_string(i)); + TF_RETURN_IF_ERROR(NodeBuilder(new_input_name.c_str(), "Identity") + .Input(replicated_input_node, 0) + .ControlInput(metadata_node) + .Attr("_tpu_replicate", cluster_name_.c_str()) + .Finalize(&graph_->graph, &new_input_node)); + DCHECK_GT(input_node_map_.count(input_nodes_[i].oper->node.name()), 0); + input_node_map_[input_nodes_[i].oper->node.name()] = new_input_node; + } + } + return Status::OK(); + } + + // On success, sets `config_op` and `shutdown_op` to the corresponding + // "ConfigureDistributedTPU" and "ShutdownDistributedTPU" nodes added to the + // graph. + tensorflow::Status RewriteGraphAndAddOutputNodes(TF_Output* new_output_nodes, + TF_Output* config_op, + TF_Output* shutdown_op) + EXCLUSIVE_LOCKS_REQUIRED(graph_->mu) { + tensorflow::Status s; + // For each non-input node in the input graph, place the node in a "TPU + // replicate cluster" via an attribute, and with the above metadata node + // as a control dependency. + // + // Although we have handled the input nodes in ProcessInputNodes(), some + // of those nodes may also serve as output nodes, which we will handle + // below. + for (Node* n : nodes_to_rewrite_) { + if (n->IsSource()) continue; + VLOG(1) << "Rewriting node " << n->name(); + + if (n->IsSink()) { + // TODO(hongm): Rewire SINK to be control dependent on the new input + // nodes created above? + continue; + } + + const NodeDef& old_def = n->def(); + Node* new_node; + if (input_node_map_.count(n->name())) { + new_node = input_node_map_[n->name()]; + } else { + // This node is to replace `n` in the graph. + NodeDef new_def = n->def(); + const std::string new_node_name = "TPUReplicate/" + n->name(); + new_def.set_name(new_node_name); + new_def.clear_input(); + for (int i = 0; i < old_def.input_size(); ++i) { + const std::string& old_input_name = old_def.input(i); + const std::string new_input_name = + input_node_map_.count(old_input_name) > 0 + ? std::string(input_node_map_[old_input_name]->name()) + : "TPUReplicate/" + old_input_name; + new_def.add_input(new_input_name); + } + if (old_def.input_size() == 0) { + // It is sufficient to only set control dependency of nodes without + // input. Other nodes with input(s) with inherit such control + // dependency. + // e.g. say the graph computes add(x, y). Once we make nodes x and y + // control-dependent on the metadata node, node add will inherit + // such control dependency indirectly. + new_def.add_input( + tensorflow::strings::StrCat("^", metadata_node_name_.c_str())); + } + tensorflow::AddNodeAttr("_tpu_replicate", cluster_name_.c_str(), + &new_def); + new_node = graph_->graph.AddNode(new_def, &s); + if (!s.ok()) { + return s; + } + VLOG(1) << "The rewritten node node is " << new_node->DebugString(); + } + + if (output_node_map_.count(n->name()) > 0) { + VLOG(1) << "Handling output node " << n->name(); + auto range_it = output_node_map_.equal_range(n->name()); + for (auto it = range_it.first; it != range_it.second; ++it) { + const PortIndexPair& pair = it->second; + Node* out_identity_node; + { + VLOG(1) << "Handling its output port " << pair.port + << " at output index " << pair.index; + std::string output_node_name = "TPUReplicate/Identity"; + if (pair.index > 0) { + output_node_name += "_" + std::to_string(pair.index); + } + TF_RETURN_IF_ERROR( + NodeBuilder(output_node_name.c_str(), "Identity") + .Input(new_node, pair.port) + .Device(!old_def.device().empty() + ? old_def.device() + : tensorflow::strings::StrCat( + "/device:", DEVICE_TPU_REPLICATED_CORE)) + .Attr("_tpu_replicate", cluster_name_.c_str()) + .Finalize(&graph_->graph, &out_identity_node)); + VLOG(1) << "out_identity_node: " + << out_identity_node->DebugString(); + } + + Node* replicated_output_node; + { + const std::string replicated_output_node_name = + "TPUReplicate/output" + std::to_string(pair.index); + TF_RETURN_IF_ERROR( + NodeBuilder(replicated_output_node_name.c_str(), + "TPUReplicatedOutput") + .Input(out_identity_node, 0) + .Attr("num_replicas", 1) + .Finalize(&graph_->graph, &replicated_output_node)); + VLOG(1) << "replicated_output_node: " + << replicated_output_node->DebugString(); + } + + Node* final_output_node; + const std::string final_output_node_name = + "TPUReplicate/output_" + std::to_string(pair.index) + "_shard_" + + std::to_string(0); + TF_RETURN_IF_ERROR( + NodeBuilder(final_output_node_name.c_str(), "Identity") + .Input(replicated_output_node, 0) + .Finalize(&graph_->graph, &final_output_node)); + VLOG(1) << "new_output_node: " << final_output_node->DebugString(); + auto oper = ToTF_Operation(final_output_node); + new_output_nodes[pair.index] = {oper, 0}; + } + } + + if (input_node_map_.count(n->name()) == 0) { + graph_->graph.RemoveNode(n); + } + } + + { + Node* config_node; + TF_RETURN_IF_ERROR( + NodeBuilder("ConfigureDistributedTPU", "ConfigureDistributedTPU") + .Device(DEVICE_TPU_SYSTEM) + .Finalize(&graph_->graph, &config_node)); + *config_op = {ToTF_Operation(config_node), 0}; + } + + { + Node* shutdown_node; + TF_RETURN_IF_ERROR( + NodeBuilder("ShutdownDistributedTPU", "ShutdownDistributedTPU") + .Device(DEVICE_TPU_SYSTEM) + .Finalize(&graph_->graph, &shutdown_node)); + *shutdown_op = {ToTF_Operation(shutdown_node), 0}; + } + + return Status::OK(); + } + + TF_Graph* const graph_; + + const TF_Output* const input_nodes_; + + const std::string cluster_name_ = "TPUReplicate/cluster"; + const std::string metadata_node_name_ = "TPUReplicate/TPUReplicateMetadata"; + + // Keep mappings from the current input nodes to newly created input nodes, + // which we will use to rewrite existing nodes that read these + // inputs. e.g. A node that reads input node PlaceHolder could be rewired to + // read the created TPUReplicate/replicated_input_0 node. + std::unordered_map input_node_map_; + + std::vector nodes_to_rewrite_; + + // Map from name to set{(output port, output tensor idx)}. + // e.g. Say ther are 3 output tensors, respectively produced by (node 0, + // port 0), (node 0, port 1), (node 1, port 0). Then the mapping entries + // are: node 0 -> {(port 0, idx 0), (port 1, idx 1)} node 1 -> {(port 0, idx + // 2)} Based on these mappings, we will generated 3 new output nodes. + struct PortIndexPair { + int port; + int index; + }; + std::multimap output_node_map_; +}; + +} // namespace + +void TF_EnableXLACompilation(TF_SessionOptions* options, unsigned char enable) { + tensorflow::ConfigProto& config = options->options.config; + auto* optimizer_options = + config.mutable_graph_options()->mutable_optimizer_options(); + if (enable) { + optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::ON_1); + + // These XLA flags are needed to trigger XLA properly from C (more generally + // non-Python) clients. If this API is called again with `enable` set to + // false, it is safe to keep these flag values as is. + tensorflow::legacy_flags::MarkForCompilationPassFlags* flags = + tensorflow::legacy_flags::GetMarkForCompilationPassFlags(); + flags->tf_xla_cpu_global_jit = true; + flags->tf_xla_min_cluster_size = 1; + } else { + optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF); + } +} + +TF_Output TF_SetupTPUExecution(TF_Session* session, int num_input_nodes, + const TF_Output* input_nodes, + int num_output_nodes, + const TF_Output* output_nodes, + TF_Output* new_output_nodes, TF_Status* status) { + TF_Output config_op, shutdown_op; + { + auto graph = session->graph; + tensorflow::mutex_lock c(graph->mu); + + VLOG(1) << "Graph before TPU rewrite: " + << graph->graph.ToGraphDefDebug().DebugString(); + GraphRewriter rewriter(graph, num_input_nodes, input_nodes, + num_output_nodes, output_nodes); + status->status = + rewriter.Rewrite(new_output_nodes, &config_op, &shutdown_op); + if (!status->status.ok()) { + return shutdown_op; + } + VLOG(1) << "Graph after TPU rewrite: " + << graph->graph.ToGraphDefDebug().DebugString(); + } + + VLOG(1) << "Initializing TPU"; + TF_Tensor* dummy_output; + TF_SessionRun(session, /*run_options*/ nullptr, + // input related parameters + /*inputs*/ nullptr, /*input_values*/ nullptr, /*ninputs*/ 0, + // output related parameters + /*outputs*/ &config_op, /*output_values*/ &dummy_output, + /*noutputs*/ 1, + /*targets*/ nullptr, /*ntargets*/ 0, + /*run_metadata*/ nullptr, status); + if (status->status.ok()) { + TF_DeleteTensor(dummy_output); + } + return shutdown_op; +} + +void TF_ShutdownTPUExecution(TF_Session* session, TF_Output shutdown_node, + TF_Status* status) { + { + tensorflow::mutex_lock c(session->graph->mu); + VLOG(1) << "Shutting down TPU, with input graph: " + << session->graph->graph.ToGraphDefDebug().DebugString(); + } + + TF_SessionRun(session, /*run_options*/ nullptr, + // input related parameters + /*inputs*/ nullptr, /*input_values*/ nullptr, /*ninputs*/ 0, + // output related parameters + /*outputs*/ nullptr, /*output_values*/ nullptr, + /*noutputs*/ 0, + /*targets*/ &shutdown_node.oper, /*ntargets*/ 1, + /*run_metadata*/ nullptr, status); +} diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h new file mode 100644 index 0000000000000000000000000000000000000000..af65123131e0899c6825d2701970808912e71bf4 --- /dev/null +++ b/tensorflow/c/c_api_experimental.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_C_C_API_EXPERIMENTAL_H_ +#define TENSORFLOW_C_C_API_EXPERIMENTAL_H_ + +#include +#include + +#include "tensorflow/c/c_api.h" + +// -------------------------------------------------------------------------- +// Experimental C API for TensorFlow. +// +// The API here is subject to changes in the future. +// -------------------------------------------------------------------------- + +// Macro to control visibility of exported symbols in the shared library (.so, +// .dylib, .dll). +// This duplicates the TF_EXPORT macro definition in +// tensorflow/core/platform/macros.h in order to keep this .h file independent +// of any other includes.$a +#ifdef SWIG +#define TF_CAPI_EXPORT +#else +#if defined(COMPILER_MSVC) +#ifdef TF_COMPILE_LIBRARY +#define TF_CAPI_EXPORT __declspec(dllexport) +#else +#define TF_CAPI_EXPORT __declspec(dllimport) +#endif // TF_COMPILE_LIBRARY +#else +#define TF_CAPI_EXPORT __attribute__((visibility("default"))) +#endif // COMPILER_MSVC +#endif // SWIG + +#ifdef __cplusplus +extern "C" { +#endif + +// When `enable` is true, set +// tensorflow.ConfigProto.OptimizerOptions.global_jit_level to ON_1, and also +// set XLA flag values to prepare for XLA compilation. Otherwise set +// global_jit_level to OFF. +// +// This API is syntax sugar over TF_SetConfig(), and is used by clients that +// cannot read/write the tensorflow.ConfigProto proto. +TF_CAPI_EXPORT extern void TF_EnableXLACompilation(TF_SessionOptions* options, + unsigned char enable); + +// Sets up TPU execution, by rewriting the graph accordingly, and initializing +// TPU system. +// +// On success, returns a shutdown node to be used in a subsequent +// TF_ShutdownTPUExecution(), and sets the new output nodes in +// `new_output_nodes` for caller to fetch from. Must be called exactly once +// before TF_SessionRun(). +// +// The API and logic is modeled after the python counterparts +// tpu.{initialize_system(), rewrite(), shutdown_system()}. +// +// TODO(b/74774824): Create separate APIs for initializing TPU system and graph +// rewrite. +TF_CAPI_EXPORT extern TF_Output TF_SetupTPUExecution( + TF_Session* session, int num_input_nodes, const TF_Output* input_nodes, + int num_output_nodes, const TF_Output* output_nodes, + TF_Output* new_output_nodes, TF_Status* status); + +// Shuts down TPU system. For any `session` where TF_SetupTPUExecution() has +// been successfully called, this call must be made exactly once before the +// session is closed. +TF_CAPI_EXPORT extern void TF_ShutdownTPUExecution(TF_Session* session, + TF_Output shutdown_node, + TF_Status* status); + +#ifdef __cplusplus +} /* end extern "C" */ +#endif + +#endif // TENSORFLOW_C_C_API_EXPERIMENTAL_H_ diff --git a/tensorflow/c/c_api_function.cc b/tensorflow/c/c_api_function.cc index 46271e0514f473099848a8573cb7cb6fad33f7dc..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; } @@ -467,7 +478,7 @@ Status ComputeBodyNodes( return Status::OK(); } -} // anonymous namespace +} // namespace } // namespace tensorflow using tensorflow::Node; diff --git a/tensorflow/c/c_api_function_test.cc b/tensorflow/c/c_api_function_test.cc index dbce66d2317a8e89288fab932cf69055f8b5a7f0..610274696f5940c063e68f2310cfd9cc1e0bd964 100644 --- a/tensorflow/c/c_api_function_test.cc +++ b/tensorflow/c/c_api_function_test.cc @@ -20,6 +20,7 @@ limitations under the License. #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/proto_serialization.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" @@ -331,6 +332,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) { @@ -980,7 +986,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) { @@ -1023,12 +1029,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) { diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h index 91667056e0eeb224b4b8a034766f11a123cd1a03..e885a699274cfae04d5a17c736da1acfddcc7b3b 100644 --- a/tensorflow/c/c_api_internal.h +++ b/tensorflow/c/c_api_internal.h @@ -124,15 +124,16 @@ struct TF_Session { TF_Session(tensorflow::Session* s, TF_Graph* g); tensorflow::Session* session; - TF_Graph* graph; + TF_Graph* const graph; - tensorflow::mutex mu; + tensorflow::mutex mu ACQUIRED_AFTER(TF_Graph::mu); int last_num_graph_nodes; - // NOTE(ashankar): Experimental fields to help keep the - // buffers of a TF_Tensor pinned in device memory. - const tensorflow::DeviceMgr* device_mgr; // Owned by session. - std::vector devices; // Owned by device_mgr. + // If true, TF_SessionRun and similar methods will call + // ExtendSessionGraphHelper before running the graph (this is the default + // public behavior). Can be set to false if the caller needs to call + // ExtendSessionGraphHelper manually. + std::atomic extend_before_run; }; struct TF_ImportGraphDefOptions { @@ -210,7 +211,11 @@ void TF_GraphSetOutputHandleShapesAndTypes(TF_Graph* graph, TF_Output output, TF_Status* status); void RecordMutation(TF_Graph* graph, const TF_Operation& op, - const char* mutation_type); + const char* mutation_type) + EXCLUSIVE_LOCKS_REQUIRED(graph->mu); + +bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) + LOCKS_EXCLUDED(session->graph->mu, session->mu); } // end namespace tensorflow diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc index 01954eb235f1a93d943c2ec7ea4c5ca44785d402..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 = @@ -912,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(); @@ -1956,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 @@ -2054,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); @@ -2084,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..22f77e7b874a13b3b6e0fbe981b4188c634db439 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/c/c_api_experimental.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; @@ -32,6 +34,10 @@ static void DoubleDeallocator(void* data, size_t, void* arg) { delete[] static_cast(data); } +static void FloatDeallocator(void* data, size_t, void* arg) { + delete[] static_cast(data); +} + TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values) { int64_t num_values = 1; for (int i = 0; i < num_dims; ++i) { @@ -76,21 +82,30 @@ TF_Tensor* DoubleTensor(double v) { &DoubleDeallocator, nullptr); } +TF_Tensor* FloatTensor(float v) { + const int num_bytes = sizeof(float); + float* values = new float[1]; + values[0] = v; + return TF_NewTensor(TF_FLOAT, nullptr, 0, values, num_bytes, + &FloatDeallocator, nullptr); +} + // All the *Helper methods are used as a workaround for the restrictions that // one cannot call ASSERT_* methods in non-void-returning functions (when // exceptions are disabled during compilation) void PlaceholderHelper(TF_Graph* graph, TF_Status* s, const char* name, - TF_Operation** op) { + TF_DataType dtype, TF_Operation** op) { TF_OperationDescription* desc = TF_NewOperation(graph, "Placeholder", name); - TF_SetAttrType(desc, "dtype", TF_INT32); + TF_SetAttrType(desc, "dtype", dtype); *op = TF_FinishOperation(desc, s); ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); ASSERT_NE(*op, nullptr); } -TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, const char* name) { +TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, const char* name, + TF_DataType dtype) { TF_Operation* op; - PlaceholderHelper(graph, s, name, &op); + PlaceholderHelper(graph, s, name, dtype, &op); return op; } @@ -124,8 +139,15 @@ 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) { +TF_Operation* ScalarConst(float v, TF_Graph* graph, TF_Status* s, + const char* name) { + unique_tensor_ptr tensor(FloatTensor(v), TF_DeleteTensor); + return Const(tensor.get(), graph, s, name); +} + +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 +161,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 +182,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 +421,9 @@ 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(); + TF_EnableXLACompilation(opts, use_XLA); 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 6acc2fec0063a8592e8e22a00b530df05a08cdb8..d87c57fd5193129665ca65761872a38131ee532b 100644 --- a/tensorflow/c/c_test_util.h +++ b/tensorflow/c/c_test_util.h @@ -44,8 +44,11 @@ TF_Tensor* Int32Tensor(int32_t v); TF_Tensor* DoubleTensor(double v); +TF_Tensor* FloatTensor(float v); + TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, - const char* name = "feed"); + const char* name = "feed", + TF_DataType dtype = TF_INT32); TF_Operation* Const(TF_Tensor* t, TF_Graph* graph, TF_Status* s, const char* name = "const"); @@ -56,6 +59,9 @@ TF_Operation* ScalarConst(int32_t v, TF_Graph* graph, TF_Status* s, TF_Operation* ScalarConst(double v, TF_Graph* graph, TF_Status* s, const char* name = "scalar"); +TF_Operation* ScalarConst(float v, TF_Graph* graph, TF_Status* s, + const char* name = "scalar"); + TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph, TF_Status* s, const char* name = "add"); @@ -69,6 +75,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 +122,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 +138,8 @@ class CSession { TF_Tensor* output_tensor(int i) { return output_values_[i]; } + TF_Session* mutable_session() { return session_; } + private: void DeleteInputValues(); void ResetOutputValues(); diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index e62310d811462f88af93505393b622d9a87c72d3..3046d9064a6d4b39cd8a7209d7f20e1e779c2847 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( @@ -49,6 +58,7 @@ tf_cuda_library( "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:framework_lite", + "//tensorflow/core:lib", "//tensorflow/core:lib_internal", ], ) @@ -56,12 +66,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", diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index fd6cecd77ba59011e252171e8f079d41ea0bd345..455bc19be8c2cbae9241b899adce95cfffedc8cc 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -25,17 +25,24 @@ 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" +#include "tensorflow/core/common_runtime/device_set.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/rendezvous_mgr.h" +#include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/rendezvous.h" #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/env.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/thread_annotations.h" #include "tensorflow/core/public/version.h" @@ -44,15 +51,38 @@ 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 +TFE_ContextDevicePlacementPolicy PlacementPolicy( + bool soft_placement, TFE_ContextDevicePlacementPolicy original_policy) { + if (!soft_placement) { + return original_policy; + } + if (original_policy == TFE_DEVICE_PLACEMENT_EXPLICIT || + original_policy == TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32) { + return TFE_DEVICE_PLACEMENT_SILENT; + } + return original_policy; +} + extern "C" { TFE_ContextOptions* TFE_NewContextOptions() { return new TFE_ContextOptions; } @@ -62,47 +92,64 @@ void TFE_ContextOptionsSetConfig(TFE_ContextOptions* options, const void* proto, TF_SetConfig(&options->session_options, proto, proto_len, status); } +void TFE_ContextOptionsSetAsync(TFE_ContextOptions* options, + unsigned char async) { + options->async = async; +} void TFE_ContextOptionsSetDevicePlacementPolicy( TFE_ContextOptions* options, TFE_ContextDevicePlacementPolicy policy) { options->policy = policy; } +TF_CAPI_EXPORT extern void TFE_ContextSetAsyncForThread(TFE_Context* ctx, + unsigned char async, + TF_Status* status) { + { + tensorflow::mutex_lock l(ctx->async_map_mu); + ctx->thread_local_async[std::this_thread::get_id()] = async; + } + if (async) { + ctx->executor.EnableAsync(); + } else { + // TODO(agarwal): Currently we add a wait here to handle cases where a sync + // op has a control dependency on an async op, and the latter has not + // executed yet. This wait can be removed by storing all the control inputs + // and waiting for them when executing ops. + status->status = ctx->executor.WaitForAllPendingNodes(); + } +} + void TFE_DeleteContextOptions(TFE_ContextOptions* options) { delete options; } TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) { - TF_Graph* graph = TF_NewGraph(); - TF_Session* session = TF_NewSession(graph, &opts->session_options, status); - if (status->status.ok()) { - if (session->device_mgr == nullptr || session->devices.empty()) { - status->status = tensorflow::errors::InvalidArgument( - "Provided TF_SessionOptions are not compatible with eager execution " - "(perhaps the TF_SessionOptions alluded to session execution in a " - "remote address space?)"); - } - } + std::vector devices; + status->status = tensorflow::DeviceFactory::AddDevices( + opts->session_options.options, "/job:localhost/replica:0/task:0", + &devices); if (!status->status.ok()) { - TF_DeleteGraph(graph); return nullptr; } - - return new TFE_Context(*opts, session); + std::unique_ptr device_mgr( + new tensorflow::DeviceMgr(devices)); + tensorflow::Rendezvous* r = + new tensorflow::IntraProcessRendezvous(device_mgr.get()); + return new TFE_Context(*opts, std::move(device_mgr), r); } void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status) { - status->status = tensorflow::Status::OK(); + status->status = ctx->executor.WaitForAllPendingNodes(); { 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); ctx->rendezvous->Unref(); delete ctx; } TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx, TF_Status* status) { - return TF_SessionListDevices(ctx->session, status); + TF_DeviceList* list = new TF_DeviceList; + ctx->device_manager->ListDeviceAttributes(&list->response); + return list; } void TFE_ContextClearCaches(TFE_Context* ctx) { @@ -116,6 +163,9 @@ void TFE_ContextSetThreadLocalDevicePlacementPolicy( ctx->thread_local_policies[std::this_thread::get_id()] = policy; } +// Note: this function looks up a thread local policy. So it should be called in +// the appropriate client thread. In particular, in async mode, it may not be +// safe to call this function from the async TFE_Executor threads. extern TFE_ContextDevicePlacementPolicy TFE_ContextGetDevicePlacementPolicy( TFE_Context* ctx) { tensorflow::mutex_lock ml(ctx->policy_map_mu); @@ -127,75 +177,107 @@ extern TFE_ContextDevicePlacementPolicy TFE_ContextGetDevicePlacementPolicy( return ctx->policy; } +void TFE_ContextAsyncWait(TFE_Context* ctx, TF_Status* status) { + status->status = ctx->executor.WaitForAllPendingNodes(); +} + +void TFE_ContextGetStatus(TFE_Context* ctx, TF_Status* status) { + status->status = ctx->executor.status(); +} + +void TFE_ContextAsyncClearError(TFE_Context* ctx) { + ctx->executor.ClearError(); +} + TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_Status* status) { tensorflow::Tensor tensor; status->status = tensorflow::TF_TensorToTensor(t, &tensor); if (!status->status.ok()) return nullptr; - return new TFE_TensorHandle(tensor, nullptr); + return new TFE_TensorHandle(tensor, nullptr, nullptr); } -void TFE_DeleteTensorHandle(TFE_TensorHandle* h) { delete h; } +void TFE_DeleteTensorHandle(TFE_TensorHandle* h) { + DCHECK(h); + h->Unref(); +} TF_DataType TFE_TensorHandleDataType(TFE_TensorHandle* h) { - return static_cast(h->t.dtype()); + return static_cast(h->dtype); } -int TFE_TensorHandleNumDims(TFE_TensorHandle* h) { return h->t.dims(); } +int TFE_TensorHandleNumDims(TFE_TensorHandle* h, TF_Status* status) { + const tensorflow::Tensor* t = nullptr; + status->status = h->Tensor(&t); + return t == nullptr ? 0 : t->dims(); +} -int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index) { - return h->t.dim_size(dim_index); +int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index, + TF_Status* status) { + const tensorflow::Tensor* t = nullptr; + status->status = h->Tensor(&t); + return t == nullptr ? 0 : t->dim_size(dim_index); } -const char* TFE_TensorHandleDeviceName(TFE_TensorHandle* h) { - // This might be a bit confusing as a tensor on CPU can sometimes return - // "CPU:0" and sometimes "/job:localhost/replica:0/task:0/cpu:0". - // TODO(ashankar): Figure out which one would be nicer. - return (h->d == nullptr) ? "CPU:0" : h->d->name().c_str(); +const char* TFE_TensorHandleDeviceName(TFE_TensorHandle* h, TF_Status* status) { + tensorflow::Device* d = nullptr; + status->status = h->OpDevice(&d); + return (d == nullptr) ? "/job:localhost/replica:0/task:0/device:CPU:0" + : d->name().c_str(); } TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status) { - if (!IsCPU(h->d)) { + // TODO(agarwal): move this implementation inside TFE_TensorHandle. + tensorflow::Device* d = nullptr; + tensorflow::Device* op_device = nullptr; + const tensorflow::Tensor* t = nullptr; + status->status = h->TensorAndDevice(&t, &d, &op_device); + if (!status->status.ok()) return nullptr; + if (!IsCPU(d)) { TF_SetStatus(status, TF_UNIMPLEMENTED, tensorflow::strings::StrCat( "TFE_TensorHandle can be resolved iff it is on CPU (this " "handle is on ", - h->d->name(), + d->name(), "). Consider using TFE_TensorHandleCopyToDevice to get a " "copy of the tensor on CPU") .c_str()); return nullptr; } - return tensorflow::TF_TensorFromTensor(h->t, status); + return tensorflow::TF_TensorFromTensor(*t, status); } +} // extern "C" -TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, - TFE_Context* ctx, - const char* device_name, - TF_Status* status) { - tensorflow::Device* dstd = ctx->devices()[0]; - if (device_name != nullptr && strlen(device_name) > 0) { - status->status = ctx->session->device_mgr->LookupDevice(device_name, &dstd); - if (!status->status.ok()) return nullptr; - } +namespace { - tensorflow::Device* srcd = h->d == nullptr ? ctx->devices()[0] : h->d; +tensorflow::Status TensorHandleCopyToDevice(TFE_TensorHandle* h, + TFE_Context* ctx, + tensorflow::Device* dstd, + TFE_TensorHandle** output) { + const tensorflow::Tensor* src = nullptr; + tensorflow::Device* srcd = nullptr; + // TODO(agarwal): src_opd is unused. Perhaps allow TensorAndDevice to accept + // nullptr. + tensorflow::Device* src_opd = nullptr; + TF_RETURN_IF_ERROR(h->TensorAndDevice(&src, &srcd, &src_opd)); + if (srcd == nullptr) srcd = ctx->devices[0]; bool is_same_device = (srcd == dstd) || (DeviceName(srcd) == DeviceName(dstd)); const bool dst_cpu = IsCPU(dstd); const bool src_cpu = IsCPU(srcd); - if (is_same_device) { - return new TFE_TensorHandle(h->t, dst_cpu ? nullptr : dstd); + // 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) { + dstd = dst_cpu ? nullptr : dstd; + *output = new TFE_TensorHandle(*src, dstd, dstd); + return tensorflow::Status::OK(); } - tensorflow::Tensor* src = &(h->t); 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 ", - tensorflow::DataTypeString(src->dtype()), - " to device ", DeviceName(dstd), ".") - .c_str()); - return nullptr; + return tensorflow::errors::InvalidArgument( + "Can't copy Tensor with type ", + tensorflow::DataTypeString(src->dtype()), " to device ", + DeviceName(dstd), "."); } tensorflow::AllocatorAttributes attr; if (src->dtype() == tensorflow::DT_VARIANT) { @@ -203,7 +285,9 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, } 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); + dstd = dst_cpu ? nullptr : dstd; + *output = new TFE_TensorHandle(dst, dstd, dstd); + return tensorflow::Status::OK(); } tensorflow::DeviceContext* src_device_context = nullptr; if (!src_cpu) { @@ -220,20 +304,26 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, // With that setup, Sync()ing across all 3 streams should be sufficient // but more than necessary (since it waits for operations that might have // nothing to do with this tensor to complete). - status->status = srcd->Sync(); + TF_RETURN_IF_ERROR(srcd->Sync()); tensorflow::Notification n; + tensorflow::Status status; tensorflow::CopyTensor::ViaDMA("copy", src_device_context, dst_device_context, srcd, dstd, tensorflow::AllocatorAttributes(), tensorflow::AllocatorAttributes(), src, &dst, - [status, &n](const tensorflow::Status& s) { - status->status = s; + [&status, &n](const tensorflow::Status& s) { + status = s; n.Notify(); }); n.WaitForNotification(); - return (TF_GetCode(status) == TF_OK) - ? new TFE_TensorHandle(dst, dst_cpu ? nullptr : dstd) - : nullptr; + if (status.ok()) { + dstd = dst_cpu ? nullptr : dstd; + *output = new TFE_TensorHandle(dst, dstd, dstd); + } + return status; } +} // namespace + +extern "C" { TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, TF_Status* status) { @@ -253,30 +343,43 @@ 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) { - status->status = - op->ctx->session->device_mgr->LookupDevice(device_name, &d); + status->status = op->ctx->device_manager->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); - if (!status->status.ok()) return; - op->inputs.push_back(h->t); - op->input_devices.push_back(h->d); + if (op->device == nullptr) { + // Questionable heuristic ... + // - If a device was explicitly set on the op, always use that. + // - If not, place on the first non-host device seen. + tensorflow::Device* d = nullptr; + // TODO(agarwal): This call may block if h is not ready. Avoid this if + // possible. + status->status = h->Device(&d); + if (!status->status.ok()) return; + if (!IsCPU(d)) op->device = d; + } + h->Ref(); + op->inputs.push_back(h); op->attrs.NumInputs(op->inputs.size()); } @@ -290,7 +393,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; } @@ -426,13 +529,26 @@ 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)); +} +} // extern "C" + namespace { tensorflow::Status ValidateInputTypeAndPlacement( TFE_Context* ctx, tensorflow::Device* host_device, tensorflow::Device* op_device, TFE_Op* op, - const tensorflow::OpKernel* kernel, - std::vector* copied_tensors) { + const tensorflow::OpKernel* kernel) { const tensorflow::MemoryTypeVector& memtypes = kernel->input_memory_types(); if (memtypes.size() != op->inputs.size()) { return tensorflow::errors::InvalidArgument( @@ -441,14 +557,17 @@ tensorflow::Status ValidateInputTypeAndPlacement( for (int i = 0; i < op->inputs.size(); ++i) { const tensorflow::Device* expected_device = memtypes[i] == tensorflow::HOST_MEMORY ? host_device : op_device; + TFE_TensorHandle* handle = op->inputs[i]; + tensorflow::Device* handle_device = nullptr; + TF_RETURN_IF_ERROR(handle->Device(&handle_device)); const tensorflow::Device* actual_device = - op->input_devices[i] == nullptr ? host_device : op->input_devices[i]; + handle_device == nullptr ? host_device : handle_device; if (expected_device != actual_device) { 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) { + if (handle->dtype == tensorflow::DT_INT32) { // Note: enabling silent copies of int32 tensors to match behavior // of graph mode. break; @@ -479,86 +598,85 @@ tensorflow::Status ValidateInputTypeAndPlacement( } // We are only here if the policy is warn or silent copies, so we should // trigger a copy. - TFE_TensorHandle original{op->inputs[i], op->input_devices[i]}; TF_Status* s = TF_NewStatus(); TFE_TensorHandle* copied_tensor = TFE_TensorHandleCopyToDevice( - &original, ctx, expected_device->name().c_str(), s); - if (!s->status.ok()) { - tensorflow::Status status = s->status; - delete s; + handle, ctx, expected_device->name().c_str(), s); + tensorflow::Status status = s->status; + TF_DeleteStatus(s); + if (!status.ok()) { + if (copied_tensor != nullptr) copied_tensor->Unref(); return tensorflow::errors::Internal( "Failed copying input tensor from ", actual_device->name(), " to ", expected_device->name(), " in order to run ", op->name, ": ", status.error_message()); } - op->inputs[i] = copied_tensor->t; - copied_tensors->push_back(copied_tensor); - op->input_devices[i] = copied_tensor->d; - delete s; + handle->Unref(); + handle = copied_tensor; + op->inputs[i] = copied_tensor; } - if (op->inputs[i].dtype() != kernel->input_type(i)) { + if (handle->dtype != kernel->input_type(i)) { return tensorflow::errors::InvalidArgument( "cannot compute ", op->name, " as input #", i, " was expected to be a ", tensorflow::DataTypeString(kernel->input_type(i)), - " tensor but is a ", - tensorflow::DataTypeString(op->inputs[i].dtype()), " tensor"); + " tensor but is a ", tensorflow::DataTypeString(handle->dtype), + " tensor"); } } return tensorflow::Status::OK(); } -} // namespace -void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, - TF_Status* status) { - TFE_Context* ctx = op->ctx; - // TODO(ashankar): ASSUMPTION: ctx->devices()[0] is always CPU - tensorflow::Device* device = - (op->device == nullptr) ? ctx->devices()[0] : op->device; - std::vector outputs(1); - const tensorflow::MemoryTypeVector* output_memory_types = nullptr; - tensorflow::Fprint128 cache_key = op->attrs.CacheKey(device->name()); - tensorflow::KernelAndDevice* kernel; - { - tensorflow::tf_shared_lock l(ctx->cache_mu); - kernel = tensorflow::gtl::FindPtrOrNull(ctx->kernel_cache, cache_key); +tensorflow::Device* SelectDevice(const tensorflow::NodeDef& ndef, + TFE_Context* ctx, TF_Status* status) { + tensorflow::DeviceSet ds; + for (tensorflow::Device* d : ctx->devices) { + ds.AddDevice(d); } - if (kernel == nullptr) { - const tensorflow::NodeDef& ndef = op->attrs.BuildNodeDef(); - kernel = new tensorflow::KernelAndDevice(ctx->rendezvous); - // Knowledge of the implementation of Init (and in-turn - // FunctionLibraryRuntime::CreateKernel) tells us that ctx->func_lib_def - // will be accessed, so grab on to the lock. - // See WARNING comment below - would be nice to rework to avoid this - // subtlety. - tensorflow::tf_shared_lock l(ctx->functions_mu); - status->status = - tensorflow::KernelAndDevice::Init(ndef, ctx->func_lib(device), kernel); - if (!status->status.ok()) { - delete kernel; - return; - } - tensorflow::mutex_lock ml(ctx->cache_mu); - tensorflow::gtl::InsertOrUpdate(&(ctx->kernel_cache), cache_key, kernel); - } - std::vector copied_tensors; - status->status = ValidateInputTypeAndPlacement( - ctx, ctx->devices()[0], device, op, kernel->kernel(), &copied_tensors); - output_memory_types = &kernel->kernel()->output_memory_types(); + tensorflow::DeviceTypeVector final_devices; + status->status = tensorflow::SupportedDeviceTypesForNode( + ds.PrioritizedDeviceTypeList(), ndef, &final_devices); if (!status->status.ok()) { - for (auto* t : copied_tensors) { - TFE_DeleteTensorHandle(t); + return nullptr; + } + if (final_devices.empty()) { + status->status = tensorflow::errors::Internal( + "Could not find valid device for node ", ndef.DebugString()); + return nullptr; + } + for (tensorflow::Device* d : ctx->devices) { + if (d->device_type() == final_devices[0].type_string()) { + return d; } - return; } - std::unique_ptr maybe_stats; - if (ctx->should_store_metadata.load()) { - maybe_stats.reset(new tensorflow::NodeExecStats); - maybe_stats->set_node_name(op->name); - maybe_stats->set_all_start_micros(tensorflow::Env::Default()->NowMicros()); - maybe_stats->set_op_start_rel_micros(0); - maybe_stats->set_scheduled_micros(tensorflow::Env::Default()->NowMicros()); - // TODO(apassos) track referenced tensors + status->status = tensorflow::errors::Unknown( + "Could not find a device for node ", ndef.DebugString()); + return nullptr; +} + +tensorflow::Status Execute( + TFE_Context* ctx, tensorflow::Device* device, + const tensorflow::gtl::InlinedVector& op_inputs, + tensorflow::KernelAndDevice* kernel, tensorflow::NodeExecStats* maybe_stats, + TFE_TensorHandle** retvals, int num_retvals) { + if (!ctx->soft_placement && device == nullptr) { + // TODO(ashankar): ASSUMPTION: ctx->devices[0] is always CPU + device = ctx->devices[0]; + } + + if (device == nullptr) { + // TODO(apassos) debug how the assignment below might return a different + // device from the one requested above. + device = kernel->device(); + } + + std::vector outputs(1); + const tensorflow::MemoryTypeVector* output_memory_types = nullptr; + output_memory_types = &kernel->kernel()->output_memory_types(); + std::vector inputs(op_inputs.size()); + for (int i = 0; i < op_inputs.size(); ++i) { + const tensorflow::Tensor* input_tensor = nullptr; + TF_RETURN_IF_ERROR(op_inputs[i]->Tensor(&input_tensor)); + inputs[i] = *input_tensor; } // WARNING: kernel->Run utilizes the FunctionLibraryRuntime // (ctx->func_lib(device)), which in turn holds a pointer to func_lib_def, @@ -569,11 +687,8 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, // sense for FunctionLibraryRuntime to ensure thread-safe access to // FunctionLibraryDefinition?). TODO(apassos) figure out how to record stats // for ops which are a part of functions. - status->status = kernel->Run(&op->inputs, &outputs, maybe_stats.get()); - for (auto* t : copied_tensors) { - TFE_DeleteTensorHandle(t); - } - if (!status->status.ok()) return; + // TODO(agarwal): change Run to take vector of handles ? + TF_RETURN_IF_ERROR(kernel->Run(&inputs, &outputs, maybe_stats)); if (maybe_stats != nullptr) { maybe_stats->set_op_end_rel_micros(tensorflow::Env::Default()->NowMicros() - maybe_stats->all_start_micros()); @@ -582,13 +697,13 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, auto* step_stats = ctx->run_metadata.mutable_step_stats(); // Lazily initialize the RunMetadata with information about all devices if // this is the first call. - while (step_stats->dev_stats_size() < ctx->devices().size()) { + while (step_stats->dev_stats_size() < ctx->devices.size()) { step_stats->add_dev_stats(); } // Find the current device's index. int device_idx = 0; - for (int i = 0; i < ctx->devices().size(); ++i) { - if (ctx->devices()[i] == device) { + for (int i = 0; i < ctx->devices.size(); ++i) { + if (ctx->devices[i] == device) { device_idx = i; break; } @@ -599,14 +714,514 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, *dev_stats->add_node_stats() = *maybe_stats; } } - *num_retvals = std::min(*num_retvals, outputs.size()); - for (int i = 0; i < *num_retvals; ++i) { - tensorflow::Device* d = IsCPU(device) ? nullptr : device; + DCHECK_EQ(num_retvals, outputs.size()); + tensorflow::Device* op_device = IsCPU(device) ? nullptr : device; + for (int i = 0; i < num_retvals; ++i) { + tensorflow::Device* d = op_device; if (d != nullptr && output_memory_types != nullptr && (*output_memory_types)[i] == tensorflow::HOST_MEMORY) { d = nullptr; } - retvals[i] = new TFE_TensorHandle(outputs[i], d); + if (retvals[i] == nullptr) { + retvals[i] = new TFE_TensorHandle(outputs[i], d, op_device); + } else { + retvals[i]->SetTensorAndDevice(outputs[i], d, op_device); + } + } + return tensorflow::Status::OK(); +} + +// TODO(agarwal): move TFE_Executor and TFE_Node related code to a separate +// file. +class ExecuteNode : public TFE_Node { + public: + ExecuteNode(TFE_Op* op, tensorflow::KernelAndDevice* kernel, + tensorflow::NodeExecStats* maybe_stats, + const tensorflow::DataTypeVector& output_dtypes, + TFE_TensorHandle** retvals, int num_retvals) + : TFE_Node(op->ctx->executor.NextId()), + ctx_(op->ctx), + op_device_(op->device), + inputs_(op->inputs), + kernel_(kernel), + maybe_stats_(maybe_stats), + retvals_(num_retvals) { + for (auto handle : inputs_) { + handle->Ref(); + } + TFE_Context* ctx = op->ctx; + for (int i = 0; i < num_retvals; ++i) { + TFE_TensorHandle* h = new TFE_TensorHandle(id, output_dtypes[i], ctx); + h->Ref(); + retvals[i] = h; + retvals_[i] = h; + } + } + + ~ExecuteNode() override { + for (auto handle : inputs_) { + handle->Unref(); + } + for (auto handle : retvals_) { + handle->Unref(); + } + } + + tensorflow::Status Run() override { + const tensorflow::Status status = + Execute(ctx_, op_device_, inputs_, kernel_, maybe_stats_.get(), + retvals_.begin(), retvals_.size()); + if (status.ok()) { + return status; + } else { + return tensorflow::Status( + status.code(), + tensorflow::strings::StrCat("Got error, \"", status.error_message(), + "\" while executing kernel ", + kernel_->kernel()->def().DebugString())); + } + } + + private: + TFE_Context* ctx_; + tensorflow::Device* op_device_; + tensorflow::gtl::InlinedVector inputs_; + tensorflow::KernelAndDevice* kernel_; + std::unique_ptr maybe_stats_; + tensorflow::gtl::InlinedVector retvals_; +}; + +class CopyToDeviceNode : public TFE_Node { + public: + CopyToDeviceNode(TFE_TensorHandle* src, tensorflow::Device* dstd, + TFE_Context* ctx) + : TFE_Node(ctx->executor.NextId()), + src_(src), + dstd_(dstd), + ctx_(ctx), + dst_(new TFE_TensorHandle(id, src_->dtype, ctx)) { + src_->Ref(); + dst_->Ref(); + } + + ~CopyToDeviceNode() override { + src_->Unref(); + dst_->Unref(); + } + + tensorflow::Status Run() override { + TFE_TensorHandle* temp = nullptr; + TF_RETURN_IF_ERROR(TensorHandleCopyToDevice(src_, ctx_, dstd_, &temp)); + const tensorflow::Tensor* tensor = nullptr; + tensorflow::Device* device = nullptr; + tensorflow::Device* op_device = nullptr; + tensorflow::Status status = + temp->TensorAndDevice(&tensor, &device, &op_device); + // `temp` is a ready handle. So the following call should return OK. + TF_DCHECK_OK(status) << status.error_message(); + DCHECK(tensor); + dst_->SetTensorAndDevice(*tensor, device, op_device); + temp->Unref(); + return tensorflow::Status::OK(); + } + + TFE_TensorHandle* dst() { return dst_; } + + private: + TFE_TensorHandle* src_; + tensorflow::Device* dstd_; + TFE_Context* ctx_; + TFE_TensorHandle* dst_; +}; + +#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; + for (TFE_TensorHandle* h : launch_op->inputs) { + h->Ref(); + } + if (!op_input_to_func_input.empty()) { + DCHECK_EQ(op->inputs.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]; + } + } + 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 + +extern "C" { + +void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, + TF_Status* status) { + TFE_Context* ctx = op->ctx; + status->status = ctx->executor.status(); + if (!status->status.ok()) { + return; + } +#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 + // Ensure all resource-touching ops run in the device the resource is, + // regardless of anything else that has been specified. This is identical to + // the graph mode behavior. + for (int i = 0; i < op->inputs.size(); ++i) { + tensorflow::Device* input_op_device = nullptr; + status->status = op->inputs[i]->OpDevice(&input_op_device); + if (!status->status.ok()) return; + if (op->inputs[i]->dtype == tensorflow::DT_RESOURCE && + input_op_device != op->device) { + tensorflow::Device* d = + input_op_device == nullptr ? ctx->devices[0] : input_op_device; + VLOG(1) << "Changing device of operation " << op->name << " to " + << d->name() << " because input #" << i + << " is a resource in this device."; + op->device = d; + } + } + tensorflow::Device* device = op->device; + if (!ctx->soft_placement && device == nullptr) { + // TODO(ashankar): ASSUMPTION: ctx->devices[0] is always CPU + device = ctx->devices[0]; + } + + tensorflow::Fprint128 cache_key = + op->attrs.CacheKey(device == nullptr ? "unspecified" : device->name()); + 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(); + if (ctx->soft_placement && device == nullptr) { + device = SelectDevice(ndef, ctx, status); + if (!status->status.ok()) { + return; + } + } + CHECK(device != nullptr); + if (ctx->log_device_placement) { + LOG(INFO) << "Executing op " << ndef.op() << " in device " + << device->name(); + } + kernel = new tensorflow::KernelAndDevice(ctx->rendezvous); + // Knowledge of the implementation of Init (and in-turn + // FunctionLibraryRuntime::CreateKernel) tells us that ctx->func_lib_def + // will be accessed, so grab on to the lock. + // See WARNING comment in Execute (before kernel->Run) - would be nice to + // rework to avoid this subtlety. + tensorflow::tf_shared_lock l(ctx->functions_mu); + status->status = + tensorflow::KernelAndDevice::Init(ndef, ctx->func_lib(device), kernel); + if (!status->status.ok()) { + delete kernel; + return; + } + // Update output_dtypes inside `kernel`. + const tensorflow::OpDef* op_def = nullptr; + const tensorflow::FunctionDef* function_def = + ctx->func_lib_def.Find(ndef.op()); + if (function_def != nullptr) { + op_def = &(function_def->signature()); + } + if (op_def == nullptr) { + status->status = OpDefForOp(ndef.op().c_str(), &op_def); + if (!status->status.ok()) { + return; + } + } + tensorflow::DataTypeVector input_dtypes; + status->status = InOutTypesForNode(ndef, *op_def, &input_dtypes, + kernel->mutable_output_dtypes()); + if (!status->status.ok()) { + return; + } + tensorflow::mutex_lock ml(ctx->cache_mu); + tensorflow::gtl::InsertOrUpdate(&(ctx->kernel_cache), cache_key, kernel); + } + const tensorflow::DataTypeVector& output_dtypes = kernel->output_dtypes(); + const int output_dtypes_size = output_dtypes.size(); + if (output_dtypes_size > *num_retvals) { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat("Expecting ", output_dtypes.size(), + " outputs, but *num_retvals is ", + *num_retvals) + .c_str()); + return; + } + *num_retvals = output_dtypes_size; + if (device == nullptr) { + // TODO(apassos) debug how the assignment below might return a different + // device from the one requested above. + device = kernel->device(); + } + status->status = ValidateInputTypeAndPlacement(ctx, ctx->devices[0], device, + op, kernel->kernel()); + if (!status->status.ok()) return; + std::unique_ptr maybe_stats; + if (ctx->should_store_metadata.load()) { + maybe_stats.reset(new tensorflow::NodeExecStats); + maybe_stats->set_node_name(op->name); + maybe_stats->set_all_start_micros(tensorflow::Env::Default()->NowMicros()); + maybe_stats->set_op_start_rel_micros(0); + maybe_stats->set_scheduled_micros(tensorflow::Env::Default()->NowMicros()); + // TODO(apassos) track referenced tensors + } + if (ctx->Async()) { + // Note that for async mode, execution order will make sure that all + // input handles are ready before executing them. + // TODO(agarwal): Consider executing "cheap" kernels inline for performance. + TFE_Node* node = new ExecuteNode(op, kernel, maybe_stats.release(), + output_dtypes, retvals, *num_retvals); + ctx->executor.Add(node); + } else { + // Execute checks if retvals[i] is nullptr or not to figure if it needs to + // allocate it. + for (int i = 0; i < *num_retvals; ++i) { + retvals[i] = nullptr; + } + status->status = Execute(op->ctx, op->device, op->inputs, kernel, + maybe_stats.get(), retvals, *num_retvals); + } +} + +TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, + TFE_Context* ctx, + const char* device_name, + TF_Status* status) { + status->status = ctx->executor.status(); + if (!status->status.ok()) { + return nullptr; + } + tensorflow::Device* dstd = ctx->devices[0]; + if (device_name != nullptr && strlen(device_name) > 0) { + status->status = ctx->device_manager->LookupDevice(device_name, &dstd); + if (!status->status.ok()) return nullptr; + } + if (ctx->Async()) { + // Note that `h` may not be currently ready. However execution order will + // make sure that `h` is ready before the copy is actually done. + CopyToDeviceNode* node = new CopyToDeviceNode(h, dstd, ctx); + TFE_TensorHandle* output = node->dst(); + // Note that calling Add makes `node` accessible by the TFE_Executor thread. + // So further accesses need to be thread-safe. + ctx->executor.Add(node); + return output; + } else { + TFE_TensorHandle* output = nullptr; + status->status = TensorHandleCopyToDevice(h, ctx, dstd, &output); + return output; } } @@ -629,36 +1244,319 @@ void TFE_ContextAddFunction(TFE_Context* ctx, TF_Function* function, status->status = ctx->func_lib_def.AddFunctionDef(function->fdef); } +void TFE_ContextEnableRunMetadata(TFE_Context* ctx) { + ctx->should_store_metadata.store(true); +} + +void TFE_ContextDisableRunMetadata(TFE_Context* ctx) { + tensorflow::mutex_lock ml(ctx->metadata_mu); + ctx->should_store_metadata.store(false); + ctx->run_metadata.Clear(); +} + } // extern "C" TFE_TensorHandle* TFE_NewTensorHandle(const tensorflow::Tensor& t) { - return new TFE_TensorHandle(t, nullptr); + return new TFE_TensorHandle(t, nullptr, nullptr); } const tensorflow::Tensor* TFE_TensorHandleUnderlyingTensorInHostMemory( TFE_TensorHandle* h, TF_Status* status) { - if (h->d != nullptr) { + tensorflow::Device* d = nullptr; + tensorflow::Device* op_device = nullptr; + const tensorflow::Tensor* t = nullptr; + status->status = h->TensorAndDevice(&t, &d, &op_device); + if (!status->status.ok()) return nullptr; + if (d != nullptr) { status->status = tensorflow::errors::FailedPrecondition( "TFE_TensorHandle is placed in device (not host) memory. Cannot return " "a tensorflow::Tensor"); return nullptr; } - return &h->t; -} - -void TFE_ContextEnableRunMetadata(TFE_Context* ctx) { - ctx->should_store_metadata.store(true); -} - -void TFE_ContextDisableRunMetadata(TFE_Context* ctx) { - tensorflow::mutex_lock ml(ctx->metadata_mu); - ctx->should_store_metadata.store(false); - ctx->run_metadata.Clear(); + return t; } void TFE_ContextExportRunMetadata(TFE_Context* ctx, TF_Buffer* buf, TF_Status* status) { + TFE_ContextAsyncWait(ctx, status); + if (!status->status.ok()) return; tensorflow::mutex_lock ml(ctx->metadata_mu); status->status = MessageToBuffer(ctx->run_metadata, buf); ctx->run_metadata.Clear(); } + +namespace { +TFE_Op* GetFunc(TFE_Context* ctx, const tensorflow::NameAttrList& func, + TF_Status* status) { + TFE_Op* func_op = TFE_NewOp(ctx, func.name().data(), status); + for (const auto& attr : func.attr()) { + if (TF_GetCode(status) != TF_OK) return nullptr; + SetOpAttrValueScalar(ctx, func_op, attr.second, attr.first.data(), status); + if (TF_GetCode(status) != TF_OK) return nullptr; + } + return func_op; +} +} // namespace + +namespace tensorflow { +void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, + const tensorflow::AttrValue& default_value, + const char* attr_name, TF_Status* status) { + switch (default_value.value_case()) { + case tensorflow::AttrValue::kS: + TFE_OpSetAttrString(op, attr_name, default_value.s().data()); + break; + case tensorflow::AttrValue::kI: + TFE_OpSetAttrInt(op, attr_name, static_cast(default_value.i())); + break; + case tensorflow::AttrValue::kF: + TFE_OpSetAttrFloat(op, attr_name, default_value.f()); + break; + case tensorflow::AttrValue::kB: + TFE_OpSetAttrBool(op, attr_name, default_value.b()); + break; + case tensorflow::AttrValue::kType: + TFE_OpSetAttrType(op, attr_name, + static_cast(default_value.type())); + break; + case tensorflow::AttrValue::kShape: { + const auto& tensor_shape = default_value.shape(); + if (tensor_shape.unknown_rank()) { + TFE_OpSetAttrShape(op, attr_name, nullptr, -1, status); + } else { + const auto num_dims = tensor_shape.dim_size(); + std::unique_ptr dims(new int64_t[num_dims]); + for (int i = 0; i < num_dims; ++i) { + dims[i] = tensor_shape.dim(i).size(); + } + TFE_OpSetAttrShape(op, attr_name, dims.get(), num_dims, status); + } + } break; + case tensorflow::AttrValue::kFunc: { + const auto func_op = GetFunc(ctx, default_value.func(), status); + if (TF_GetCode(status) != TF_OK) return; + // TODO(nareshmodi): TFE_OpSetAttrFunction and TFE_OpSetAttrFunctionList + // require TFE_Op* and just convert it internally a NameAttrValue, so + // consider adding an overload to the C API to make this case easier. + TFE_OpSetAttrFunction(op, attr_name, func_op); + } break; + case tensorflow::AttrValue::kList: + TF_FALLTHROUGH_INTENDED; + case tensorflow::AttrValue::kTensor: + TF_FALLTHROUGH_INTENDED; + case tensorflow::AttrValue::kPlaceholder: + TF_FALLTHROUGH_INTENDED; + case tensorflow::AttrValue::VALUE_NOT_SET: + TF_SetStatus( + status, TF_UNIMPLEMENTED, + tensorflow::strings::StrCat("Unable to get setfor default value: ", + default_value.DebugString()) + .data()); + } +} +} // namespace tensorflow + +TFE_Node::TFE_Node(tensorflow::uint64 id) : id(id) {} + +TFE_Executor::~TFE_Executor() { + tensorflow::mutex_lock l(node_queue_mutex_); + thread_done_ = true; + nodes_pending_.notify_all(); +} + +tensorflow::uint64 TFE_Executor::NextId() { + tensorflow::mutex_lock l(next_id_mutex_); + return next_id_++; +} + +void TFE_Executor::EnableAsync() { + tensorflow::mutex_lock l(node_queue_mutex_); + if (thread_ == nullptr) { + thread_.reset(tensorflow::Env::Default()->StartThread( + tensorflow::ThreadOptions(), "eager_async_executor", + std::bind(&TFE_Executor::Run, this))); + } +} + +void TFE_Executor::Add(TFE_Node* node) { + tensorflow::mutex_lock l(node_queue_mutex_); + DCHECK(thread_) << "EnableAsync should have been called before Add"; + if (!status_.ok()) { + delete node; + return; + } + int qlen = node_queue_.size(); + if (qlen > 0) { + if (node_queue_.back()->id >= node->id) { + status_ = tensorflow::errors::InvalidArgument( + "Inserting TFE_Node with non-increasing ids:", node_queue_.back()->id, + " vs ", node->id); + delete node; + return; + } + node_queue_.push(node); + } else { + node_queue_.push(node); + nodes_pending_.notify_all(); + } +} + +tensorflow::Status TFE_Executor::WaitFor(tensorflow::uint64 node_id) { + return WaitImpl(false, node_id); +} + +tensorflow::Status TFE_Executor::WaitForAllPendingNodes() { + return WaitImpl(true, 0); +} + +tensorflow::Status TFE_Executor::WaitImpl(bool wait_all, + tensorflow::uint64 node_id) { + tensorflow::condition_variable cond; + tensorflow::mutex_lock l(node_queue_mutex_); + // Don't wait if an error is already set. + if (!status_.ok()) return status_; + if (node_queue_.empty()) return tensorflow::Status::OK(); + if (wait_all) { + node_id = node_queue_.back()->id; + } else if (node_id < node_queue_.front()->id) { + // Note that we are relying on the ops being dispatched sequentially from + // the queue. + return tensorflow::Status::OK(); + } + node_done_notifications_.insert(std::make_pair(node_id, &cond)); + cond.wait(l); + // Note that we could be woken up if an error occurs, even though the node has + // not actually executed. + return status_; +} + +void TFE_Executor::ClearError() { + tensorflow::mutex_lock l(node_queue_mutex_); + if (status_.ok()) return; + // If an error was set, node_done_notifications_ and node_queue_ should have + // been cleared, and no new entries should have been added since. + DCHECK(node_done_notifications_.empty()); + DCHECK(node_queue_.empty()); + status_ = tensorflow::Status::OK(); + nodes_pending_.notify_all(); +} + +tensorflow::Status TFE_Executor::status() { + tensorflow::mutex_lock l(node_queue_mutex_); + return status_; +} + +void TFE_Executor::Run() { + while (true) { + std::unique_ptr curr_node; + { + tensorflow::mutex_lock l(node_queue_mutex_); + while (node_queue_.empty() || !status_.ok()) { + if (thread_done_) return; + nodes_pending_.wait(l); + } + curr_node.reset(node_queue_.front()); + } + tensorflow::Status status = curr_node->Run(); + const bool ok = status.ok(); + tensorflow::mutex_lock l(node_queue_mutex_); + node_queue_.pop(); + if (!ok) { + status_ = status; + // TODO(agarwal): mark all affected handles as corrupted before clearing + // this queue. + // We remove any pending ops so that we don't try to execute them if + // ClearError is called. + for (int i = 0; i < node_queue_.size(); ++i) { + delete node_queue_.front(); + node_queue_.pop(); + } + } + if (!node_done_notifications_.empty()) { + tensorflow::uint64 node_id = curr_node->id; + // Note that we notify all waiting threads in case an error has occurred. + // These calling threads are responsible for checking status_ before + // proceeding. + const auto range = ok ? node_done_notifications_.equal_range(node_id) + : make_pair(node_done_notifications_.begin(), + node_done_notifications_.end()); + for (auto it = range.first; it != range.second; ++it) { + it->second->notify_all(); + } + node_done_notifications_.erase(range.first, range.second); + } + } +} + +bool TFE_Context::Async() const { + tensorflow::mutex_lock l(async_map_mu); + return tensorflow::gtl::FindWithDefault( + thread_local_async, std::this_thread::get_id(), async_default); +} + +bool TFE_TensorHandle::IsReady() { + if (node_id == 0) return true; + tensorflow::mutex_lock l(ctx_mutex_); + return ctx_ == nullptr; +} + +tensorflow::Status TFE_TensorHandle::WaitReady() { + if (node_id == 0) return tensorflow::Status::OK(); + TFE_Executor* executor = nullptr; + { + tensorflow::mutex_lock l(ctx_mutex_); + if (ctx_ == nullptr) return tensorflow::Status::OK(); + executor = &ctx_->executor; + } + return executor->WaitFor(node_id); +} + +tensorflow::Status TFE_TensorHandle::Tensor(const tensorflow::Tensor** t) { + TF_RETURN_IF_ERROR(WaitReady()); + DCHECK(IsReady()); + *t = &tensor_; + return tensorflow::Status::OK(); +} + +tensorflow::Status TFE_TensorHandle::Device(tensorflow::Device** d) { + TF_RETURN_IF_ERROR(WaitReady()); + DCHECK(IsReady()); + *d = device_; + return tensorflow::Status::OK(); +} + +tensorflow::Status TFE_TensorHandle::OpDevice(tensorflow::Device** d) { + TF_RETURN_IF_ERROR(WaitReady()); + DCHECK(IsReady()); + *d = op_device_; + return tensorflow::Status::OK(); +} + +tensorflow::Status TFE_TensorHandle::TensorAndDevice( + const tensorflow::Tensor** tensor, tensorflow::Device** device, + tensorflow::Device** op_device) { + TF_RETURN_IF_ERROR(WaitReady()); + DCHECK(IsReady()); + *tensor = &tensor_; + *device = device_; + *op_device = op_device_; + return tensorflow::Status::OK(); +} + +void TFE_TensorHandle::SetTensorAndDevice(const tensorflow::Tensor& tensor, + tensorflow::Device* device, + tensorflow::Device* op_device) { + tensorflow::mutex_lock l(ctx_mutex_); + DCHECK(node_id > 0 && ctx_) << "SetTensorAndDevice should be only called " + << "on non-ready handles."; + ctx_ = nullptr; + tensor_ = tensor; + device_ = device; + op_device_ = op_device; +} + +TFE_Op::~TFE_Op() { + for (TFE_TensorHandle* h : inputs) { + h->Unref(); + } +} diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h index 387de078948e5076d0b069d6380dfc04ea6254df..a5029bf2115c7dac54d03b8bc6397bc63349c068 100644 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -61,7 +61,8 @@ 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 { - // Running operations with input tensors on the wrong device will fail. + // Running operations with input tensors on the wrong device will fail. When + // soft placement is enabled acts like TFE_DEVICE_PLACEMENT_SILENT. TFE_DEVICE_PLACEMENT_EXPLICIT = 0, // Copy the tensor to the right device but log a warning. TFE_DEVICE_PLACEMENT_WARN = 1, @@ -69,10 +70,16 @@ typedef enum TFE_ContextDevicePlacementPolicy { // 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. + // dtypes. When soft placement is enabled acts like + // TFE_DEVICE_PLACEMENT_SILENT. TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32 = 3, } TFE_ContextDevicePlacementPolicy; +// Sets the default execution mode (sync/async). Note that this can be +// overridden per thread using TFE_ContextSetAsyncForThread. +TF_CAPI_EXPORT extern void TFE_ContextOptionsSetAsync(TFE_ContextOptions*, + unsigned char async); + TF_CAPI_EXPORT extern void TFE_ContextOptionsSetDevicePlacementPolicy( TFE_ContextOptions*, TFE_ContextDevicePlacementPolicy); @@ -87,7 +94,8 @@ 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); @@ -107,6 +115,30 @@ TF_CAPI_EXPORT extern void TFE_ContextSetThreadLocalDevicePlacementPolicy( TF_CAPI_EXPORT extern TFE_ContextDevicePlacementPolicy TFE_ContextGetDevicePlacementPolicy(TFE_Context*); +// Overrides the execution mode (sync/async) for the current thread. +TF_CAPI_EXPORT extern void TFE_ContextSetAsyncForThread(TFE_Context*, + unsigned char async, + TF_Status* status); + +// Causes the calling thread to block till all ops dispatched in async mode +// have been executed. Note that "execution" here refers to kernel execution / +// scheduling of copies, etc. Similar to sync execution, it doesn't guarantee +// that lower level device queues (like GPU streams) have been flushed. +// +// This call may not block for execution of ops enqueued concurrently with this +// call. +TF_CAPI_EXPORT extern void TFE_ContextAsyncWait(TFE_Context*, + TF_Status* status); + +// When an error happens, any pending operations are discarded and newly issued +// ops return an error. This call clears the error state and re-enables +// execution of newly issued ops. +// +// Note that outputs of discarded ops remain in a corrupt state and should not +// be used for future calls. +// TODO(agarwal): mark the affected handles and raise errors if they are used. +TF_CAPI_EXPORT extern void TFE_ContextAsyncClearError(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, @@ -116,11 +148,21 @@ typedef struct TFE_TensorHandle TFE_TensorHandle; TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_Status* status); +// Indicates that the caller will not be using `h` any more. 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); +// This function will block till the operation that produces `h` has completed. +TF_CAPI_EXPORT extern int TFE_TensorHandleNumDims(TFE_TensorHandle* h, + TF_Status* status); +// This function will block till the operation that produces `h` has completed. +TF_CAPI_EXPORT extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, + int dim_index, + TF_Status* status); +// This function will block till the operation that produces `h` has completed. +TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName( + TFE_TensorHandle* h, TF_Status* status); + +// This function will block till the operation that produces `h` has completed. TF_CAPI_EXPORT extern TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status); @@ -130,10 +172,12 @@ 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); +// If async execution is enabled, the copy may be enqueued and the call will +// return "non-ready" handle. Else, this function returns after the copy has +// been done. +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. // @@ -148,17 +192,32 @@ 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); + +// 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 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. @@ -166,10 +225,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, @@ -178,7 +240,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'. @@ -189,36 +252,58 @@ 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'. +// tensors in `retvals`. // -// 'retvals' must point to a pre-allocated array of TFE_TensorHandle* -// and '*num_retvals' should be set to the size of this array. +// 'retvals' must point to a pre-allocated array of TFE_TensorHandle* and +// '*num_retvals' should be set to the size of this array. It is an error if +// the size of 'retvals' is less than the number of outputs. This call sets +// *num_retvals to the number of outputs. // -// On return, 'num_retvals' will be set to the actual number of outputs -// returned by the operation. +// If async execution is enabled, the call may simply enqueue the execution +// and return "non-ready" handles in `retvals`. Note that any handles contained +// in 'op' should not be mutated till the kernel execution actually finishes. +// +// For sync execution, if any of the inputs to `op` are not ready, this call +// will block till they become ready and then return when the kernel execution +// is done. +// TODO(agarwal): change num_retvals to int from int*. TF_CAPI_EXPORT extern void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, TF_Status* status); // 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 @@ -236,6 +321,8 @@ TF_CAPI_EXPORT extern void TFE_ContextDisableRunMetadata(TFE_Context* ctx); // Populates the passed-in buffer with a serialized RunMetadata protocol buffer // containing any run metadata information accumulated so far and clears this // information. +// If async mode is enabled, this call blocks till all currently pending ops are +// done. TF_CAPI_EXPORT extern void TFE_ContextExportRunMetadata(TFE_Context* ctx, TF_Buffer* buf, TF_Status* status); diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index dda68471a82f444cde5e74a38423ae5c97b1d7d0..8dba12f47b580c33041cc134c6f07a1fafff7453 100644 --- a/tensorflow/c/eager/c_api_internal.h +++ b/tensorflow/c/eager/c_api_internal.h @@ -19,7 +19,9 @@ limitations under the License. #include #include +#include #include +#include #include #include #include @@ -31,27 +33,140 @@ limitations under the License. #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/rendezvous_mgr.h" #include "tensorflow/core/framework/rendezvous.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/lib/gtl/map_util.h" #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" +// A unit of execution for the TFE_Executor class below. Example subclasses +// encapsulate execution of a TFE_Op, or copying a TFE_TensorHandle from one +// device to another. +class TFE_Node { + public: + explicit TFE_Node(tensorflow::uint64 id); + + virtual ~TFE_Node() {} + + // Runs the computation corresponding to this node and blocks till the + // execution is done. + virtual tensorflow::Status Run() = 0; + + // An id unique to the TFE_Context under which this node is created. Allocated + // monotonically. + const tensorflow::uint64 id; +}; + +// A class for handling async execution (see TFE_ContextSetAsync). +// Note that this class is thread-safe. +// TODO(agarwal): TFE_OpAddInput may currently block if it tries to access the +// device of the input handle. Fix that. +// TODO(agarwal): On error, mark all affected handles as corrupted. +// TODO(agarwal): Implement support for control dependencies. +// TODO(agarwal): Support out-of-order execution and dispatching multiple +// TFE_Node in parallel. +// TODO(agarwal): Implement optimizations over TFE_Node traces. +class TFE_Executor { + public: + ~TFE_Executor(); + + // This is called whenever async mode is enabled. Note that it may be called + // multiple times as different calling threads may switch async mode on or off + // independently. + void EnableAsync(); + + // Helper function to create monotonically increasing ids unique to this + // object. + tensorflow::uint64 NextId(); + + // Schedules `node` for execution. + // Note that Add must be called in monotonically increasing order of node->id. + void Add(TFE_Node* node); + + // Causes the caller to block till node with id `node_id` has finished + // execution. + tensorflow::Status WaitFor(tensorflow::uint64 node_id); + + // Blocks till all currently pending ops are done. + tensorflow::Status WaitForAllPendingNodes(); + + // Clears all currently set errors which re-enables async execution. + void ClearError(); + + // Returns Status based on any errors that occurred during async execution. + tensorflow::Status status(); + + private: + // Starts execution of pending TFE_Nodes. This function loops till + // thread_done_ is set to true. If any errors are encontered, these are set + // inside `status_`. The loop blocks anytime there are no pending nodes, or if + // `status_` is not ok. + void Run(); + + tensorflow::Status WaitImpl(bool wait_all, tensorflow::uint64 node_id); + + tensorflow::mutex node_queue_mutex_; + + // Used to signal that some TFE_Nodes are pending execution. + tensorflow::condition_variable nodes_pending_ GUARDED_BY(node_queue_mutex_); + + // Queue of pending TFE_Nodes. + std::queue node_queue_ GUARDED_BY(node_queue_mutex_); + + // `status_` is set based on any errors raised during execution of a TFE_Node. + // It remains set until ClearError is called. + tensorflow::Status status_ GUARDED_BY(node_queue_mutex_); + + // Map from id of a TFE_Node to condition_variables (not owned by the map). + // These condition_variables are notified and removed when that TFE_Node is + // done executing, or if an error is found in execution of any TFE_Node. + std::multimap + node_done_notifications_ GUARDED_BY(node_queue_mutex_); + + // Thread object that calls the `Run` method. Currently we use only one thread + // for executing the TFE_Nodes one-by-one. + std::unique_ptr thread_ GUARDED_BY(node_queue_mutex_); + + // Indicates that `thread_` should stop as soon as it is done executing the + // current TFE_Node. + bool thread_done_ GUARDED_BY(node_queue_mutex_) = false; + + tensorflow::mutex next_id_mutex_; + tensorflow::uint64 next_id_ GUARDED_BY(next_id_mutex_) = 1; +}; + struct TFE_ContextOptions { TF_SessionOptions session_options; + // true if async execution is enabled. + bool async = false; TFE_ContextDevicePlacementPolicy policy{ TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32}; }; +TFE_ContextDevicePlacementPolicy PlacementPolicy( + bool soft_placement, TFE_ContextDevicePlacementPolicy original_policy); + struct TFE_Context { - explicit TFE_Context(const TFE_ContextOptions& opts, TF_Session* s) - : policy(opts.policy), - session(s), - rendezvous(new tensorflow::IntraProcessRendezvous(s->device_mgr)), + explicit TFE_Context(const TFE_ContextOptions& opts, + std::unique_ptr device_mgr, + tensorflow::Rendezvous* rendezvous) + : soft_placement( + opts.session_options.options.config.allow_soft_placement()), + policy(PlacementPolicy(soft_placement, opts.policy)), + device_manager(std::move(device_mgr)), + devices(device_manager->ListDevices()), + rendezvous(rendezvous), pflr(new tensorflow::ProcessFunctionLibraryRuntime( - session->device_mgr, opts.session_options.options.env, - TF_GRAPH_DEF_VERSION, &func_lib_def, {})) {} + device_manager.get(), opts.session_options.options.env, + TF_GRAPH_DEF_VERSION, &func_lib_def, {})), + log_device_placement( + opts.session_options.options.config.log_device_placement()), + async_default(opts.async) { + if (async_default) executor.EnableAsync(); + } + const bool soft_placement; const TFE_ContextDevicePlacementPolicy policy; // Note: we cannot use C++11 thread_local here as there is no concept of a @@ -60,8 +175,9 @@ struct TFE_Context { 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* const session; + std::unique_ptr device_manager; + // Devices owned by device_manager + const std::vector devices; tensorflow::Rendezvous* const rendezvous; tensorflow::mutex functions_mu; @@ -82,43 +198,129 @@ struct TFE_Context { return pflr->GetFLR(d->name()); } - const std::vector& devices() { return session->devices; } - // Whether we should compute RunMetadata. std::atomic should_store_metadata{false}; tensorflow::mutex metadata_mu; tensorflow::RunMetadata run_metadata GUARDED_BY(metadata_mu); + const bool log_device_placement; + // TFE_Executor for async execution. + TFE_Executor executor; + + // True if running in asynchronous mode. + bool Async() const; + + // True if the default value for execution mode is async. Note that this value + // can be overridden per thread based on `thread_local_async` overrides. + const bool async_default; + mutable tensorflow::mutex async_map_mu; + std::unordered_map thread_local_async + GUARDED_BY(async_map_mu); }; -struct TFE_TensorHandle { - TFE_TensorHandle(const tensorflow::Tensor& t, tensorflow::Device* d) - : t(t), d(d) {} +struct TFE_TensorHandle : public tensorflow::core::RefCounted { + public: + TFE_TensorHandle(const tensorflow::Tensor& t, tensorflow::Device* d, + tensorflow::Device* op_device) + : dtype(t.dtype()), + node_id(0), + tensor_(t), + device_(d), + op_device_(op_device), + ctx_(nullptr) {} + + TFE_TensorHandle(tensorflow::uint64 node_id, tensorflow::DataType dtype, + TFE_Context* ctx) + : dtype(dtype), + node_id(node_id), + tensor_(dtype), + device_(nullptr), + op_device_(nullptr), + ctx_(ctx) { + DCHECK_GT(node_id, 0); + } + + ~TFE_TensorHandle() override {} + + tensorflow::Status Tensor(const tensorflow::Tensor** t); + + tensorflow::Status Device(tensorflow::Device** d); + + tensorflow::Status OpDevice(tensorflow::Device** d); + + tensorflow::Status TensorAndDevice(const tensorflow::Tensor** tensor, + tensorflow::Device** device, + tensorflow::Device** op_device); - tensorflow::Tensor t; - // TODO(ashankar): d == nullptr iff local CPU - // This was expedient, but perhaps worth revisiting ('d' should always be a - // valid pointer?) + // Note that this can be called at most once, and only on non-ready handles, + // and makes them ready. + void SetTensorAndDevice(const tensorflow::Tensor& tensor, + tensorflow::Device* device, + tensorflow::Device* op_device); + + // dtype for the handle. It must be the same as t.dtype() once the handle is + // ready. + const tensorflow::DataType dtype; + + private: + // If the contents of the Tensor pointed to by this handle is yet to be + // computed by a TFE_Node, this function will block till that compuatation is + // done and the handle is "ready". + tensorflow::Status WaitReady(); + + bool IsReady(); + + // Id for the TFE_Node that will compute the value pointed to by this handle. + // If the value is 0, the handle is already ready, but not vice-versa. + const tensorflow::uint64 node_id; + + tensorflow::Tensor tensor_; + + // TODO(ashankar): device_ == nullptr iff local CPU + // This was expedient, but perhaps worth revisiting ('device_' should always + // be a valid pointer?) // This can be done if TFE_NewOp() and the TFE_TensorHandle constructors are // provided with the appropriate TFE_Context. // - // TODO(ashankar): Reference count TFE_Context to ensure that 'd' of a + // TODO(ashankar): Reference count TFE_Context to ensure that 'device_' of a // TFE_TensorHandle does not outlive the TFE_Context from which it came? - tensorflow::Device* d; + tensorflow::Device* device_; + + // Device in which the op producing this tensor was executed. Equals to + // device_ for constant tensors. + tensorflow::Device* op_device_; + + tensorflow::mutex ctx_mutex_; + + // `ctx` is only guaranteed to be set if the handle is not "ready". This is + // typically true when the handle was produced during async execution. + // `ctx` object is not owned and should outlive this handle. + TFE_Context* ctx_ GUARDED_BY(ctx_mutex_); }; 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) {} + ~TFE_Op(); + bool const is_function() const { return attr_types == nullptr; } TFE_Context* ctx; // Must outlive the TFE_Op. const tensorflow::string name; tensorflow::AttrBuilder attrs; const tensorflow::AttrTypeMap* attr_types; - std::vector inputs; - std::vector input_devices; + tensorflow::gtl::InlinedVector inputs; tensorflow::Device* device; + bool use_xla = false; }; +namespace tensorflow { +// Set an AttrValue on the op. Doesn't handle the list types. +void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, + const tensorflow::AttrValue& default_value, + const char* attr_name, TF_Status* status); +} // namespace tensorflow + #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 18e7a64435e6c7e51998a744abd615edc7ad4318..2268aba90d60b7b2f10e99f64fd7aa3ae719badb 100644 --- a/tensorflow/c/eager/c_api_test.cc +++ b/tensorflow/c/eager/c_api_test.cc @@ -29,6 +29,20 @@ using tensorflow::string; namespace { +TFE_TensorHandle* DoubleTestMatrixTensorHandle() { + int64_t dims[] = {2, 2}; + double data[] = {1.0, 2.0, 3.0, 4.0}; + TF_Tensor* t = TF_AllocateTensor( + TF_DOUBLE, &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_TensorHandle* TestMatrixTensorHandle() { int64_t dims[] = {2, 2}; float data[] = {1.0f, 2.0f, 3.0f, 4.0f}; @@ -43,6 +57,20 @@ TFE_TensorHandle* TestMatrixTensorHandle() { return th; } +TFE_TensorHandle* TestMatrixTensorHandle3X2() { + int64_t dims[] = {3, 2}; + double data[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + TF_Tensor* t = TF_AllocateTensor( + TF_FLOAT, &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* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) { TF_Status* status = TF_NewStatus(); @@ -60,6 +88,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(); @@ -82,10 +167,12 @@ void BM_InitOp(int iters) { } BENCHMARK(BM_InitOp); -void BM_Execute(int iters) { +void BM_Execute(int iters, int async) { tensorflow::testing::StopTiming(); + tensorflow::testing::SetLabel(async ? "ExecuteAsync" : "Execute"); TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TFE_DeleteContextOptions(opts); @@ -99,6 +186,9 @@ void BM_Execute(int iters) { TFE_Execute(matmul, &retvals[0], &num_retvals, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); } + if (async) { + TFE_ContextAsyncWait(ctx, status); + } tensorflow::testing::StopTiming(); TFE_DeleteOp(matmul); TFE_DeleteTensorHandle(m); @@ -106,7 +196,7 @@ void BM_Execute(int iters) { CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TF_DeleteStatus(status); } -BENCHMARK(BM_Execute); +BENCHMARK(BM_Execute)->Arg(0)->Arg(1); TEST(CAPI, Context) { TF_Status* status = TF_NewStatus(); @@ -148,10 +238,11 @@ TEST(CAPI, TensorHandle) { TFE_DeleteTensorHandle(h); } -TEST(CAPI, TensorHandleCopyBetweenDevices) { +void TensorHandleCopyBetweenDevices(bool async) { std::unique_ptr status( TF_NewStatus(), TF_DeleteStatus); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status.get()); TFE_DeleteContextOptions(opts); ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); @@ -217,10 +308,56 @@ TEST(CAPI, TensorHandleCopyBetweenDevices) { EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); } -TEST(CAPI, TensorHandleCopyBetweenTwoGPUDevices) { +TEST(CAPI, TensorHandleCopyBetweenDevices) { + TensorHandleCopyBetweenDevices(false); +} + +TEST(CAPI, TensorHandleCopyBetweenDevicesAsync) { + TensorHandleCopyBetweenDevices(true); +} + +void TensorHandleCopyBetweenDevicesError(bool async) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); + TFE_Context* ctx = TFE_NewContext(opts, status.get()); + TFE_DeleteContextOptions(opts); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + TFE_TensorHandle* hcpu = TestMatrixTensorHandle(); + const char* kErrorDevice = "NoSuchDevice:0"; + TFE_TensorHandle* hdevice = + TFE_TensorHandleCopyToDevice(hcpu, ctx, kErrorDevice, status.get()); + EXPECT_NE(TF_OK, TF_GetCode(status.get())); + const char* msg = "NoSuchDevice:0 unknown device"; + EXPECT_TRUE(strstr(TF_Message(status.get()), msg) != nullptr) + << TF_Message(status.get()); + TF_SetStatus(status.get(), TF_OK, ""); + const char* kCPUDevice = "CPU:0"; + TFE_TensorHandle* hcopy = + TFE_TensorHandleCopyToDevice(hcpu, ctx, kCPUDevice, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + TFE_ContextAsyncWait(ctx, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())); + TFE_DeleteTensorHandle(hcopy); + TFE_DeleteTensorHandle(hcpu); + if (hdevice != nullptr) TFE_DeleteTensorHandle(hdevice); + TFE_DeleteContext(ctx, status.get()); +} + +TEST(CAPI, TensorHandleCopyBetweenDevicesError) { + TensorHandleCopyBetweenDevicesError(false); +} + +TEST(CAPI, TensorHandleCopyBetweenDevicesErrorAsync) { + TensorHandleCopyBetweenDevicesError(true); +} + +void TensorHandleCopyBetweenTwoGPUDevices(bool async) { std::unique_ptr status( TF_NewStatus(), TF_DeleteStatus); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status.get()); TFE_DeleteContextOptions(opts); ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); @@ -275,11 +412,20 @@ TEST(CAPI, TensorHandleCopyBetweenTwoGPUDevices) { EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); } -TEST(CAPI, TensorHandleSilentCopy) { +TEST(CAPI, TensorHandleCopyBetweenTwoGPUDevices) { + TensorHandleCopyBetweenTwoGPUDevices(false); +} + +TEST(CAPI, TensorHandleCopyBetweenTwoGPUDevicesAsync) { + TensorHandleCopyBetweenTwoGPUDevices(true); +} + +void TensorHandleSilentCopy(bool async) { std::unique_ptr status( TF_NewStatus(), TF_DeleteStatus); TFE_ContextOptions* opts = TFE_NewContextOptions(); TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status.get()); TFE_DeleteContextOptions(opts); ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); @@ -288,22 +434,15 @@ 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())); - ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); - - TFE_TensorHandle* hgpu = - TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), 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_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,17 +453,22 @@ TEST(CAPI, TensorHandleSilentCopy) { TFE_DeleteTensorHandle(hgpu); } - TF_DeleteDeviceList(devices); TF_DeleteTensor(t); TFE_DeleteTensorHandle(hcpu); + TFE_ContextAsyncWait(ctx, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); TFE_DeleteContext(ctx, status.get()); EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); } -TEST(CAPI, TensorHandleSilentCopyLocal) { +TEST(CAPI, TensorHandleSilentCopy) { TensorHandleSilentCopy(false); } +TEST(CAPI, TensorHandleSilentCopyAsync) { TensorHandleSilentCopy(true); } + +void TensorHandleSilentCopyLocal(bool async) { std::unique_ptr status( TF_NewStatus(), TF_DeleteStatus); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_EXPLICIT); TFE_Context* ctx = TFE_NewContext(opts, status.get()); @@ -337,22 +481,15 @@ TEST(CAPI, TensorHandleSilentCopyLocal) { 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())); - ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); - - TFE_TensorHandle* hgpu = - TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), 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_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; @@ -363,32 +500,236 @@ TEST(CAPI, TensorHandleSilentCopyLocal) { TFE_DeleteTensorHandle(hgpu); } - TF_DeleteDeviceList(devices); TF_DeleteTensor(t); TFE_DeleteTensorHandle(hcpu); + TFE_ContextAsyncWait(ctx, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); TFE_DeleteContext(ctx, status.get()); EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); } +TEST(CAPI, TensorHandleSilentCopyLocal) { TensorHandleSilentCopyLocal(false); } +TEST(CAPI, TensorHandleSilentCopyLocalAsync) { + TensorHandleSilentCopyLocal(true); +} + +void SetAndGetOpDevices(bool async) { + 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) { +void Execute_MatMul_CPU(bool async) { TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); 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_TensorHandle* retvals[2] = {nullptr}; - int num_retvals = 2; // Should be reduced to 1 by the TFE_Execute call. + TFE_TensorHandle* retvals[2] = {nullptr, nullptr}; + int num_retvals = 2; TFE_Execute(matmul, &retvals[0], &num_retvals, status); + EXPECT_EQ(1, num_retvals); EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TFE_DeleteOp(matmul); TFE_DeleteTensorHandle(m); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteContext(ctx, status); + 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_MatMul_CPU) { Execute_MatMul_CPU(false); } +TEST(CAPI, Execute_MatMul_CPUAsync) { Execute_MatMul_CPU(true); } + +void Execute_MatMul_CPU_Runtime_Error(bool async) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* m1 = TestMatrixTensorHandle(); + TFE_TensorHandle* m2 = TestMatrixTensorHandle3X2(); + TFE_Op* matmul = MatMulOp(ctx, m1, m2); + TFE_Op* matmul2 = MatMulOp(ctx, m1, m1); + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + TFE_Execute(matmul, &retvals[0], &num_retvals, status); + TFE_DeleteOp(matmul); + if (!async) { + EXPECT_NE(TF_OK, TF_GetCode(status)); + } else { + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + EXPECT_NE(TF_OK, TF_GetCode(status)); + EXPECT_EQ(nullptr, t); + const char* msg = "Matrix size-incompatible: In[0]: [2,2], In[1]: [3,2]"; + EXPECT_TRUE(strstr(TF_Message(status), msg) != nullptr) + << TF_Message(status); + // Since error is not cleared, the following copy with correct device will + // still fail. + TF_SetStatus(status, TF_OK, ""); + TFE_DeleteTensorHandle(retvals[0]); + retvals[0] = nullptr; + TFE_Execute(matmul2, &retvals[0], &num_retvals, status); + EXPECT_NE(TF_OK, TF_GetCode(status)); + TFE_ContextAsyncClearError(ctx); + TFE_ContextAsyncWait(ctx, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + } + // Following works in async mode since TFE_ContextAsyncClearError was called. + TF_SetStatus(status, TF_OK, ""); + if (retvals[0] != nullptr) { + TFE_DeleteTensorHandle(retvals[0]); + } + retvals[0] = nullptr; + TFE_Execute(matmul2, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + TF_DeleteTensor(t); + TFE_DeleteOp(matmul2); + TFE_DeleteTensorHandle(m1); + TFE_DeleteTensorHandle(m2); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteContext(ctx, status); + TF_DeleteStatus(status); +} +TEST(CAPI, Execute_MatMul_CPU_Runtime_Error) { + Execute_MatMul_CPU_Runtime_Error(false); +} +TEST(CAPI, Execute_MatMul_CPU_Runtime_ErrorAsync) { + Execute_MatMul_CPU_Runtime_Error(true); +} + +void Execute_MatMul_CPU_Type_Error(bool async) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* m1 = TestMatrixTensorHandle(); + TFE_TensorHandle* m2 = DoubleTestMatrixTensorHandle(); + TFE_Op* matmul = MatMulOp(ctx, m1, m2); + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + TFE_Execute(matmul, &retvals[0], &num_retvals, status); + EXPECT_NE(TF_OK, TF_GetCode(status)); + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(m1); + TFE_DeleteTensorHandle(m2); + if (retvals[0] != nullptr) { + TFE_DeleteTensorHandle(retvals[0]); + } + TFE_DeleteContext(ctx, status); + TF_DeleteStatus(status); +} + +TEST(CAPI, Execute_MatMul_CPU_Type_Error) { + Execute_MatMul_CPU_Type_Error(false); +} +TEST(CAPI, Execute_MatMul_CPU_Type_ErrorAsync) { + Execute_MatMul_CPU_Type_Error(true); +} +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[1] = {nullptr}; + int num_retvals = 1; + 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 +void Execute_MatMul_XLA_CPU(bool async) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); + 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[1] = {nullptr}; + int num_retvals = 1; + 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); + 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); @@ -400,12 +741,56 @@ TEST(CAPI, Execute) { EXPECT_EQ(10, product[1]); EXPECT_EQ(15, product[2]); EXPECT_EQ(22, product[3]); + TFE_DeleteContext(ctx, status); TF_DeleteStatus(status); } +TEST(CAPI, Execute_MatMul_XLA_CPU) { Execute_MatMul_XLA_CPU(false); } +TEST(CAPI, Execute_MatMul_XLA_CPUAsync) { Execute_MatMul_XLA_CPU(true); } -TEST(CAPI, ExecuteWithTracing) { +void Execute_Min_XLA_CPU(bool async) { TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); + 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[1] = {nullptr}; + int num_retvals = 1; + 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); + 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]); + TFE_DeleteContext(ctx, status); + TF_DeleteStatus(status); +} +TEST(CAPI, Execute_Min_XLA_CPU) { Execute_Min_XLA_CPU(false); } +TEST(CAPI, Execute_Min_XLA_CPUAsync) { Execute_Min_XLA_CPU(true); } +#endif // TENSORFLOW_EAGER_USE_XLA + +void ExecuteWithTracing(bool async) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status); TFE_ContextEnableRunMetadata(ctx); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); @@ -413,8 +798,8 @@ TEST(CAPI, ExecuteWithTracing) { TFE_TensorHandle* m = TestMatrixTensorHandle(); TFE_Op* matmul = MatMulOp(ctx, m, m); - TFE_TensorHandle* retvals[2] = {nullptr}; - int num_retvals = 2; // Should be reduced to 1 by the TFE_Execute call. + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; TFE_Execute(matmul, &retvals[0], &num_retvals, status); EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TFE_DeleteOp(matmul); @@ -426,12 +811,12 @@ TEST(CAPI, ExecuteWithTracing) { EXPECT_TRUE( rm.ParseFromString({reinterpret_cast(b->data), b->length})); TF_DeleteBuffer(b); - 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]); + TFE_DeleteContext(ctx, status); ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); float product[4] = {0}; EXPECT_EQ(sizeof(product), TF_TensorByteSize(t)); @@ -443,8 +828,10 @@ TEST(CAPI, ExecuteWithTracing) { EXPECT_EQ(22, product[3]); TF_DeleteStatus(status); } +TEST(CAPI, ExecuteWithTracing) { ExecuteWithTracing(false); } +TEST(CAPI, ExecuteWithTracingAsync) { ExecuteWithTracing(true); } -TEST(CAPI, Function) { +TEST(CAPI, Function_ident_CPU) { // First create a simple identity function. TF_Graph* function_graph = TF_NewGraph(); TF_OperationDescription* arg_descr = @@ -474,36 +861,112 @@ TEST(CAPI, Function) { 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); + for (bool async : {false, true, false}) { + TFE_ContextSetAsyncForThread(ctx, static_cast(async), + status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK); + 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_DeleteTensor(t); + TF_DeleteStatus(status); +} - TFE_Op* op = TFE_NewOp(ctx, "ident", 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 = + 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); - TFE_OpAddInput(op, h, 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); - - std::vector result; - result.push_back(nullptr); - int num_retvals = 1; - TFE_Execute(op, result.data(), &num_retvals, status); - TFE_DeleteOp(op); + 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); - ASSERT_EQ(num_retvals, 1); - - TF_Tensor* r = TFE_TensorHandleResolve(result[0], 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); - EXPECT_EQ(*reinterpret_cast(TF_TensorData(r)), 42); - TFE_DeleteTensorHandle(h); - TF_DeleteTensor(r); - TFE_DeleteTensorHandle(result[0]); + TFE_DeleteContextOptions(opts); + TFE_ContextAddFunction(ctx, fn, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_DeleteFunction(fn); + + for (bool async : {false, true, false}) { + TFE_ContextSetAsyncForThread(ctx, static_cast(async), + status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK); + 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); + + // Now run it via XLA. + TFE_OpSetXLACompilation(op, true); + + 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); } +#endif // TENSORFLOW_EAGER_USE_XLA string MatMulFunction() { tensorflow::FunctionDef def; @@ -539,9 +1002,10 @@ string MatMulFunction() { return def.SerializeAsString(); } -TEST(CAPI, FunctionDefAndExecute) { +void FunctionDefAndExecute(bool async) { TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TFE_DeleteContextOptions(opts); @@ -578,11 +1042,16 @@ TEST(CAPI, FunctionDefAndExecute) { EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TF_DeleteStatus(status); } +TEST(CAPI, FunctionDefAndExecute) { FunctionDefAndExecute(false); } +TEST(CAPI, FunctionDefAndExecuteAsync) { FunctionDefAndExecute(true); } -void BM_ExecuteFunction(int iters) { +void BM_ExecuteFunction(int iters, int async) { tensorflow::testing::StopTiming(); + tensorflow::testing::SetLabel(async ? "ExecuteFunctionAsync" + : "ExecuteFunction"); TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); TFE_Context* ctx = TFE_NewContext(opts, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TFE_DeleteContextOptions(opts); @@ -604,6 +1073,9 @@ void BM_ExecuteFunction(int iters) { TFE_Execute(matmul, &retval[0], &num_retvals, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); } + if (async) { + TFE_ContextAsyncWait(ctx, status); + } tensorflow::testing::StopTiming(); TFE_DeleteTensorHandle(m); TFE_DeleteTensorHandle(retval[0]); @@ -611,7 +1083,7 @@ void BM_ExecuteFunction(int iters) { EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TF_DeleteStatus(status); } -BENCHMARK(BM_ExecuteFunction); +BENCHMARK(BM_ExecuteFunction)->Arg(0)->Arg(1); TFE_TensorHandle* CreateVariable(TFE_Context* ctx, float value, TF_Status* status) { @@ -683,7 +1155,8 @@ TEST(CAPI, Variables) { ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); ASSERT_EQ(1, num_retvals); EXPECT_EQ(TF_FLOAT, TFE_TensorHandleDataType(value_handle)); - EXPECT_EQ(0, TFE_TensorHandleNumDims(value_handle)); + EXPECT_EQ(0, TFE_TensorHandleNumDims(value_handle, status)); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); float value = 0.0f; TF_Tensor* t = TFE_TensorHandleResolve(value_handle, status); ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); @@ -725,7 +1198,8 @@ void BM_ReadVariable(int iters) { CHECK_EQ(1, num_retvals); CHECK(h); CHECK_EQ(TF_FLOAT, TFE_TensorHandleDataType(h)); - CHECK_EQ(0, TFE_TensorHandleNumDims(h)); + CHECK_EQ(0, TFE_TensorHandleNumDims(h, status)); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); h = nullptr; } tensorflow::testing::StopTiming(); diff --git a/tensorflow/c/eager/runtime.cc b/tensorflow/c/eager/runtime.cc index 3a9951e14de3a70e0b9e47fa62e6342e063c4bed..9b46cf8245901934c9c4d41a2b7c10c1c5bf7cbd 100644 --- a/tensorflow/c/eager/runtime.cc +++ b/tensorflow/c/eager/runtime.cc @@ -41,17 +41,26 @@ const uint32 kIsList = 1U << 31; } // namespace +Status OpDefForOp(const char* op_name, const OpDef** op_def) { + const OpRegistrationData* op_reg_data = nullptr; + Status s = OpRegistry::Global()->LookUp(op_name, &op_reg_data); + if (s.ok()) { + *op_def = &op_reg_data->op_def; + } + return s; +} + Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out) { mutex_lock l(g_op_name_to_attr_type_map_lock); *out = gtl::FindPtrOrNull(*OpNameToAttrTypeMap(), op_name); if (*out != nullptr) return Status::OK(); - const OpRegistrationData* op_reg_data = nullptr; - Status s = OpRegistry::Global()->LookUp(op_name, &op_reg_data); + const OpDef* op_def = nullptr; + Status s = OpDefForOp(op_name, &op_def); if (!s.ok()) return s; std::unique_ptr m(new AttrTypeMap); // TODO(agarwal): Avoid having to create this "registry" at runtime, // perhaps can be done at op registration time? - for (const auto& attr : op_reg_data->op_def.attr()) { + for (const auto& attr : op_def->attr()) { string type = attr.type(); const bool is_list = (type.length() > 6 && type.compare(0, 4, "list") == 0); if (is_list) { @@ -86,10 +95,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 +181,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}); } @@ -294,7 +302,18 @@ Status KernelAndDevice::Run(std::vector* input_tensors, params.runner = &runner; OpKernelContext context(¶ms); - device_->Compute(kernel_.get(), &context); + + if (kernel_->def().op() == "_Recv") { + // TODO(apassos) do not special-case _Recv. Currently the GPU device fails + // if trying to run _Recv->Compute(), specifically checking for _Recv. To go + // around this we call _Recv->ComputeAsync, to mimic graph mode behavior. + AsyncOpKernel* async = kernel_->AsAsync(); + Notification done; + device_->ComputeAsync(async, &context, [&done]() { done.Notify(); }); + done.WaitForNotification(); + } else { + device_->Compute(kernel_.get(), &context); + } if (!context.status().ok()) return context.status(); output_tensors->clear(); @@ -316,7 +335,7 @@ Status KernelAndDevice::Run(std::vector* input_tensors, allocator_pair.second->GetRecordsAndUnRef(); } auto* ms = stats->mutable_memory_stats(); - ms->set_temp_memory_size(context.temp_memory_size()); + 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); } diff --git a/tensorflow/c/eager/runtime.h b/tensorflow/c/eager/runtime.h index e28a416e67f8382dbd490648106a7eb6e5fcfd13..ad16f65495f8a8193b685c2b13a099232d03a505 100644 --- a/tensorflow/c/eager/runtime.h +++ b/tensorflow/c/eager/runtime.h @@ -39,11 +39,14 @@ namespace tensorflow { // represent the TF_AttrType type of the values in the list. typedef std::unordered_map AttrTypeMap; +// Look up OpDef for `op_name`. +Status OpDefForOp(const char* op_name, const OpDef** op_def); + // Returns the AttrTypeMap for the TensorFlow operation named op_name. 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. @@ -180,12 +183,18 @@ class KernelAndDevice { const OpKernel* kernel() const { return kernel_.get(); } + Device* device() const { return device_; } + + DataTypeVector* mutable_output_dtypes() { return &output_dtypes_; } + const DataTypeVector& output_dtypes() { return output_dtypes_; } + private: std::unique_ptr kernel_; Device* device_; FunctionLibraryRuntime* flib_; checkpoint::TensorSliceReaderCacheWrapper slice_reader_cache_; Rendezvous* rendez_; + DataTypeVector output_dtypes_; }; } // namespace tensorflow diff --git a/tensorflow/c/eager/runtime_test.cc b/tensorflow/c/eager/runtime_test.cc index 2ccca66f672b96b3c782ddbfc828eeda270cebee..4f75d278878d7c8ff6a5e48e5b4e633aa13aedc5 100644 --- a/tensorflow/c/eager/runtime_test.cc +++ b/tensorflow/c/eager/runtime_test.cc @@ -41,7 +41,7 @@ class TestEnv { device_mgr_.reset(new DeviceMgr({device})); flib_runtime_ = NewFunctionLibraryRuntime(device_mgr_.get(), Env::Default(), device, TF_GRAPH_DEF_VERSION, - &flib_def_, {}, nullptr); + &flib_def_, nullptr, {}, nullptr); } FunctionLibraryRuntime* function_library_runtime() const { @@ -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 bdb0815d6b68444ec1c89b835d563db20ce4d8a1..c7bd3bdafd787e5c72625b190ea8bf8b8264d22d 100644 --- a/tensorflow/c/eager/tape.h +++ b/tensorflow/c/eager/tape.h @@ -152,6 +152,8 @@ class GradientTape { gtl::ArraySlice output_gradients, std::vector* result); + bool IsPersistent() const { return persistent_; } + private: TensorTape tensor_tape_; OpTape op_tape_; diff --git a/tensorflow/c/python_api.cc b/tensorflow/c/python_api.cc index 6e37cdb5f4beea53d4a2ded0705ae482d0bc2d68..cd604538f1fa142c6fe6a76624c048baddaa52fb 100644 --- a/tensorflow/c/python_api.cc +++ b/tensorflow/c/python_api.cc @@ -99,4 +99,14 @@ void RemoveAllControlInputs(TF_Graph* graph, TF_Operation* op) { } } +void SetRequireShapeInferenceFns(TF_Graph* graph, bool require) { + mutex_lock l(graph->mu); + graph->refiner.set_require_shape_inference_fns(require); +} + +void ExtendSession(TF_Session* session, TF_Status* status) { + ExtendSessionGraphHelper(session, status); + session->extend_before_run = false; +} + } // namespace tensorflow diff --git a/tensorflow/c/python_api.h b/tensorflow/c/python_api.h index aa9d9e06b28c54cb8869eb547d36ee3cb0d4e6b8..13b680b3a24afa2d285ea18207578aff4350f6d5 100644 --- a/tensorflow/c/python_api.h +++ b/tensorflow/c/python_api.h @@ -37,6 +37,20 @@ void UpdateEdge(TF_Graph* graph, TF_Output new_src, TF_Input dst, void RemoveAllControlInputs(TF_Graph* graph, TF_Operation* op); +// Sets whether ops missing a shape inference function should trigger an +// error. The default is true. +void SetRequireShapeInferenceFns(TF_Graph* graph, bool require); + +// Extends `session` with any new operations added to its associated graph. +// Usually this happens automatically in TF_SessionRun. After this is called, +// TF_SessionRun will no longer extend the session on every call. +// +// We expose this here to allow fine-grained synchronization in multi-threaded +// workloads, which is required since the Python implementation depends on the +// above mutation methods. This allows us to prevent modifications to nodes in +// the graph after the session has been made aware of them. +void ExtendSession(TF_Session* session, TF_Status* status); + } // namespace tensorflow #endif // TENSORFLOW_C_PYTHON_API_H_ diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc index a40ad1ffc3b262840e6ca0043139b1b61e04510d..d73121c7b701ec06c03836d1a765f4b35d88fe92 100644 --- a/tensorflow/cc/framework/cc_op_gen.cc +++ b/tensorflow/cc/framework/cc_op_gen.cc @@ -28,6 +28,7 @@ limitations under the License. #include "tensorflow/core/framework/types.pb_text.h" #include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/gtl/stl_util.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/env.h" @@ -697,7 +698,8 @@ string OpInfo::GetOpAttrStruct() const { attr_comment = MakeComment(attr_comment, " "); strings::StrAppend(&setters, attr_comment); - strings::StrAppend(&setters, " Attrs ", attr_func_def, " x) {\n"); + strings::StrAppend(&setters, " TF_MUST_USE_RESULT Attrs ", attr_func_def, + " x) {\n"); strings::StrAppend(&setters, " Attrs ret = *this;\n"); strings::StrAppend(&setters, " ret.", api_def_attr.rename_to(), "_ = x;\n"); diff --git a/tensorflow/cc/framework/while_gradients.cc b/tensorflow/cc/framework/while_gradients.cc index 0734075fc6144d7c9f4fdb48c5e097faa58b8355..81870a0efa309ae6dbd5cc05a5dbe8c3e2d437c8 100644 --- a/tensorflow/cc/framework/while_gradients.cc +++ b/tensorflow/cc/framework/while_gradients.cc @@ -72,9 +72,9 @@ Status AddForwardLoopCounter(WhileContext* while_ctx, const Scope& scope, }; // Body function that adds one to input. - BodyGraphBuilderFn body_fn = [while_ctx](const Scope& scope, - const std::vector& inputs, - std::vector* outputs) { + BodyGraphBuilderFn body_fn = [](const Scope& scope, + const std::vector& inputs, + std::vector* outputs) { DCHECK_EQ(inputs.size(), 1); outputs->emplace_back(ops::Add(scope, inputs[0], 1)); return scope.status(); diff --git a/tensorflow/cc/gradients/nn_grad.cc b/tensorflow/cc/gradients/nn_grad.cc index 13a3bba5e6d5ca19ff3f0eca76665ba7d3ab628d..0cb3132e94e381f672d69aefe4a199d2b590830c 100644 --- a/tensorflow/cc/gradients/nn_grad.cc +++ b/tensorflow/cc/gradients/nn_grad.cc @@ -48,8 +48,8 @@ Status SoftmaxGrad(const Scope& scope, const Operation& op, REGISTER_GRADIENT_OP("Softmax", SoftmaxGrad); Status LogSoftmaxGrad(const Scope& scope, const Operation& op, - const std::vector& grad_inputs, - std::vector* grad_outputs) { + const std::vector& grad_inputs, + std::vector* grad_outputs) { auto softmax = Exp(scope, op.output(0)); auto sum = Sum(scope, grad_inputs[0], {1}, Sum::KeepDims(true)); auto mul = Mul(scope, sum, softmax); @@ -107,11 +107,10 @@ Status BiasAddGradHelper(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { string data_format; - BiasAddGrad::Attrs input_attrs; TF_RETURN_IF_ERROR( GetNodeAttr(op.output(0).node()->attrs(), "data_format", &data_format)); - input_attrs.DataFormat(data_format); - auto dx_1 = BiasAddGrad(scope, grad_inputs[0], input_attrs); + auto dx_1 = + BiasAddGrad(scope, grad_inputs[0], BiasAddGrad::DataFormat(data_format)); grad_outputs->push_back(Identity(scope, grad_inputs[0])); grad_outputs->push_back(dx_1); return scope.status(); @@ -130,19 +129,16 @@ Status Conv2DGrad(const Scope& scope, const Operation& op, TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding)); TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides)); TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "use_cudnn_on_gpu", &use_cudnn_on_gpu)); - Conv2DBackpropInput::Attrs input_attrs; - input_attrs.DataFormat(data_format); - input_attrs.UseCudnnOnGpu(use_cudnn_on_gpu); - auto dx_1 = Conv2DBackpropInput(scope, Shape(scope, op.input(0)), - op.input(1), grad_inputs[0], - strides, padding, input_attrs); + auto dx_1 = Conv2DBackpropInput(scope, Shape(scope, op.input(0)), op.input(1), + grad_inputs[0], strides, padding, + Conv2DBackpropInput::DataFormat(data_format) + .UseCudnnOnGpu(use_cudnn_on_gpu)); grad_outputs->push_back(dx_1); - Conv2DBackpropFilter::Attrs filter_attrs; - filter_attrs.DataFormat(data_format); - filter_attrs.UseCudnnOnGpu(use_cudnn_on_gpu); - auto dx_2 = Conv2DBackpropFilter(scope, op.input(0), - Shape(scope, op.input(1)), grad_inputs[0], - strides, padding, filter_attrs); + auto dx_2 = + Conv2DBackpropFilter(scope, op.input(0), Shape(scope, op.input(1)), + grad_inputs[0], strides, padding, + Conv2DBackpropFilter::DataFormat(data_format) + .UseCudnnOnGpu(use_cudnn_on_gpu)); grad_outputs->push_back(dx_2); return scope.status(); } @@ -160,13 +156,9 @@ Status MaxPoolGradHelper(const Scope& scope, const Operation& op, TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize)); TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding)); TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides)); - internal::MaxPoolGrad::Attrs grad_attrs; - grad_attrs.DataFormat(data_format); - auto dx = internal::MaxPoolGrad(scope, op.input(0), - op.output(0), - grad_inputs[0], - ksize, strides, - padding, grad_attrs); + auto dx = internal::MaxPoolGrad( + scope, op.input(0), op.output(0), grad_inputs[0], ksize, strides, padding, + internal::MaxPoolGrad::DataFormat(data_format)); grad_outputs->push_back(dx); return scope.status(); } @@ -180,15 +172,9 @@ Status MaxPoolGradV2Helper(const Scope& scope, const Operation& op, auto attrs = op.output(0).node()->attrs(); TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format)); TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding)); - MaxPoolGradV2::Attrs grad_attrs; - grad_attrs.DataFormat(data_format); - auto dx = MaxPoolGradV2(scope, op.input(0), - op.output(0), - grad_inputs[0], - op.input(1), - op.input(2), - padding, - grad_attrs); + auto dx = MaxPoolGradV2(scope, op.input(0), op.output(0), grad_inputs[0], + op.input(1), op.input(2), padding, + MaxPoolGradV2::DataFormat(data_format)); grad_outputs->push_back(dx); grad_outputs->push_back(NoGradient()); grad_outputs->push_back(NoGradient()); @@ -196,13 +182,74 @@ Status MaxPoolGradV2Helper(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("MaxPoolV2", MaxPoolGradV2Helper); +Status MaxPool3DGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + std::vector ksize; + std::vector strides; + string padding; + string data_format; + auto attrs = op.output(0).node()->attrs(); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format)); + MaxPool3DGrad::Attrs grad_attrs; + auto dx = MaxPool3DGrad(scope, op.input(0), op.output(0), grad_inputs[0], + ksize, strides, padding, + grad_attrs.DataFormat(data_format)); + grad_outputs->push_back(dx); + return scope.status(); +} +REGISTER_GRADIENT_OP("MaxPool3D", MaxPool3DGradHelper); + +Status AvgPoolGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + std::vector ksize; + std::vector strides; + string padding; + string data_format; + auto attrs = op.output(0).node()->attrs(); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format)); + internal::AvgPoolGrad::Attrs grad_attrs; + auto dx = + internal::AvgPoolGrad(scope, Shape(scope, op.input(0)), grad_inputs[0], + ksize, strides, padding, + grad_attrs.DataFormat(data_format)); + grad_outputs->push_back(dx); + return scope.status(); +} +REGISTER_GRADIENT_OP("AvgPool", AvgPoolGradHelper); + +Status AvgPool3DGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + std::vector ksize; + std::vector strides; + string padding; + string data_format; + auto attrs = op.output(0).node()->attrs(); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding)); + TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format)); + AvgPool3DGrad::Attrs grad_attrs; + auto dx = AvgPool3DGrad(scope, Shape(scope, op.input(0)), grad_inputs[0], + ksize, strides, padding, + grad_attrs.DataFormat(data_format)); + grad_outputs->push_back(dx); + return scope.status(); +} +REGISTER_GRADIENT_OP("AvgPool3D", AvgPool3DGradHelper); + Status LRNGradHelper(const Scope& scope, const Operation& op, const std::vector& grad_inputs, - std::vector* grad_outputs){ - internal::LRNGrad::Attrs grad_attrs; - - auto dx = internal::LRNGrad(scope, grad_inputs[0], op.input(0), op.output(0), - grad_attrs); + std::vector* grad_outputs) { + auto dx = internal::LRNGrad(scope, grad_inputs[0], op.input(0), op.output(0)); grad_outputs->push_back(dx); return scope.status(); } diff --git a/tensorflow/cc/gradients/nn_grad_test.cc b/tensorflow/cc/gradients/nn_grad_test.cc index 0cfe5f6e3c49f7c4a3cafbf48ff4e54a0ffd0d47..c4eba7ecb017fe4628140d75a63bc7f0f09deb7f 100644 --- a/tensorflow/cc/gradients/nn_grad_test.cc +++ b/tensorflow/cc/gradients/nn_grad_test.cc @@ -31,8 +31,11 @@ using ops::Elu; using ops::L2Loss; using ops::LogSoftmax; using ops::LRN; +using ops::AvgPool; +using ops::AvgPool3D; using ops::MaxPool; using ops::MaxPoolV2; +using ops::MaxPool3D; using ops::Placeholder; using ops::Relu; using ops::Relu6; @@ -70,9 +73,9 @@ class NNGradTest : public ::testing::Test { // Sets tensor with random values, ensuring that the max value is largest by // a reasonable amount. - // This is an issue for MaxPool and MaxPoolV2, in which perturbations by the - // numeric gradient computation in the gradient checker can change the max - // value if values are too close together. + // This is an issue for MaxPool, MaxPoolV2 and MaxPool3D, in which + // perturbations by the numeric gradient computation in the gradient checker + // can change the max value if values are too close together. template void SetRandomValuesWithBumpedMax(Tensor* tensor) { auto tensor_flat = tensor->flat(); @@ -203,6 +206,41 @@ TEST_F(NNGradTest, MaxPoolGradV2Helper) { RunTest(x, x_init_value, y, y_shape); } +TEST_F(NNGradTest, MaxPool3DGradHelper) { + TensorShape x_shape({1, 3, 3, 3, 1}); + TensorShape y_shape({1, 1, 1, 1, 1}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + // Setup window and strides so that we only do one MaxPool3D. + const std::vector ksize{1, 3, 3, 3, 1}; + const std::vector strides{1, 3, 3, 3, 1}; + auto y = MaxPool3D(scope_, x, ksize, strides, "VALID"); + Tensor x_init_value = Tensor(DT_FLOAT, x_shape); + SetRandomValuesWithBumpedMax(&x_init_value); + RunTest(x, x_init_value, y, y_shape); +} + +TEST_F(NNGradTest, AvgPoolGradHelper) { + TensorShape x_shape({1, 2, 2, 1}); + TensorShape y_shape({1, 1, 1, 1}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + // Setup window and strides so that we only do one AvgPool. + const std::vector ksize{1, 2, 2, 1}; + const std::vector strides{1, 2, 2, 1}; + auto y = AvgPool(scope_, x, ksize, strides, "SAME"); + RunTest(x, x_shape, y, y_shape); +} + +TEST_F(NNGradTest, AvgPool3DGradHelper) { + TensorShape x_shape({1, 3, 3, 3, 1}); + TensorShape y_shape({1, 1, 1, 1, 1}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + // Setup window and strides so that we only do one AvgPool3D. + const std::vector ksize{1, 3, 3, 3, 1}; + const std::vector strides{1, 3, 3, 3, 1}; + auto y = AvgPool3D(scope_, x, ksize, strides, "SAME"); + RunTest(x, x_shape, y, y_shape); +} + TEST_F(NNGradTest, LRN){ TensorShape x_shape({1, 1, 2, 1}); auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); diff --git a/tensorflow/cc/profiler/profiler.h b/tensorflow/cc/profiler/profiler.h index 6077c45c5854fd5812ccb7c91522f93ed4e54883..64edbb5766c3604fbe0f15c2299843718381aa3f 100644 --- a/tensorflow/cc/profiler/profiler.h +++ b/tensorflow/cc/profiler/profiler.h @@ -61,18 +61,18 @@ class Profiler { /// Adds tracing information `run_meta` to profiler. A `run_meta` is /// generated by a TensorFlow session run call. `step` is the key /// to the `run_meta`. When calling ProfileXXX methods, caller can specify - /// `step` in `options` to seletively profile the corresponding `run_meta`. + /// `step` in `options` to selectively profile the corresponding `run_meta`. /// Multiple different `run_meta` can be keyed by the same `step` in order /// to group them together. void AddStep(int64 step, const RunMetadata& run_meta); /// Profiles the model by organizing nodes in graph structure. - /// Each node is an op and the nodes are contected by the op inputs/outputs. + /// Each node is an op and the nodes are connected by the op inputs/outputs. GraphNodeProto ProfileGraph(const Options& options); /// Profiles the model by organizing nodes in name scope structure. /// Each node is an op, and nodes are organized by the ops' name - /// scope, similar to a filesystem tree. + /// scope, similar to a file system tree. /// E.g. /foo is the root of operation /foo/matmul_1 and foo/conv_2. GraphNodeProto ProfileNameScope(const Options& options); 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_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/tools/BUILD b/tensorflow/cc/tools/BUILD index 97f66e79b8ad9f383b22f56e9385fc6d2080e1f8..f413a5cc52e9eb4bc393b8186f5b591681fa2e5e 100644 --- a/tensorflow/cc/tools/BUILD +++ b/tensorflow/cc/tools/BUILD @@ -32,6 +32,7 @@ tf_cc_test( deps = [ ":freeze_saved_model", "//tensorflow/cc:cc_ops", + "//tensorflow/cc:resource_variable_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:framework_internal", "//tensorflow/core:protos_all_cc", diff --git a/tensorflow/cc/tools/freeze_saved_model.cc b/tensorflow/cc/tools/freeze_saved_model.cc index ddf372cdef21e1b3892c9a03714478d5a5785517..4ddddcb5863c9ffb1e5367db750b0d2ffd29cd5e 100644 --- a/tensorflow/cc/tools/freeze_saved_model.cc +++ b/tensorflow/cc/tools/freeze_saved_model.cc @@ -75,16 +75,13 @@ void GetNodeNameToNodeDefMap( // variable nodes to convert. void GetReachableNodesAndVariables( GraphDef* graph_def, const std::unordered_set& outputs, + const std::unordered_map& name_to_node_map, 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); + new std::unordered_set({"Variable", "VariableV2", "VarHandleOp"}); + std::queue nodes_to_visit; for (const string& tensor_name : outputs) { // We need to strip off the tensor part to get the node name. @@ -99,7 +96,7 @@ void GetReachableNodesAndVariables( continue; } reachable_node_names->insert(node_name); - NodeDef* node = name_to_node_map[node_name]; + NodeDef* node = name_to_node_map.at(node_name); if (kVariableTypes->find(node->op()) != kVariableTypes->end()) { variable_node_names->insert(node->name()); } @@ -111,7 +108,9 @@ void GetReachableNodesAndVariables( // Gets a map from variable name to variable value. Status GetVariableNameToTensorMap( - Session* session, std::unordered_set variable_names_set, + Session* session, + const std::unordered_map& name_to_node_map, + std::unordered_set variable_names_set, std::unordered_map* variable_name_to_value_map) { if (variable_names_set.empty()) { return Status::OK(); @@ -120,8 +119,14 @@ Status GetVariableNameToTensorMap( 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"); + NodeDef* node_def = name_to_node_map.at(node_name); + if (node_def->op() == "VarHandleOp") { + // If this is a resource variable, we have to run the corresponding + // ReadVariableOp. + tensor_names.push_back(node_name + "/Read/ReadVariableOp:0"); + } else { + tensor_names.push_back(node_name + ":0"); + } } std::vector outputs; TF_RETURN_IF_ERROR( @@ -143,6 +148,15 @@ void ConvertVariableToConstant(const NodeDef& variable_node, (*const_node->mutable_attr())["value"].mutable_tensor()); } +// Converts a ReadVariableOp NodeDef to an Identity NodeDef. +void ConvertReadVariableOpToIdentity(const NodeDef& node, + NodeDef* identity_node) { + identity_node->set_name(node.name()); + identity_node->set_op("Identity"); + (*identity_node->mutable_attr())["T"] = node.attr().at("dtype"); + identity_node->add_input(node.input(0)); +} + // Freezes the subgraph of all nodes needed by `outputs`. Status FreezeGraphDef(const SavedModelBundle& saved_model_bundle, const std::unordered_set& outputs, @@ -155,14 +169,19 @@ Status FreezeGraphDef(const SavedModelBundle& saved_model_bundle, if (graph_def.node_size() == 0) { return Status::OK(); } + // 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::unordered_set reachable_node_names; std::unordered_set variable_node_names; - GetReachableNodesAndVariables(&graph_def, outputs, &reachable_node_names, - &variable_node_names); + GetReachableNodesAndVariables(&graph_def, outputs, name_to_node_map, + &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)); + TF_RETURN_IF_ERROR(GetVariableNameToTensorMap( + saved_model_bundle.session.get(), name_to_node_map, 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()) { @@ -171,6 +190,12 @@ Status FreezeGraphDef(const SavedModelBundle& saved_model_bundle, 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 (node.op() == "ReadVariableOp" && + variable_node_names.find(node.input(0)) != + variable_node_names.end()) { + // If the node is a ReadVariableOp, its input VarHandleOp will be + // converted to a Constant, so we will need to convert it to an Identity. + ConvertReadVariableOpToIdentity(node, 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; diff --git a/tensorflow/cc/tools/freeze_saved_model_test.cc b/tensorflow/cc/tools/freeze_saved_model_test.cc index 52a81a50284aec36bba4e56a0232c886cb0cb6cf..cd35fd3b95deec669218cfa4f25fea2c3ac9e56e 100644 --- a/tensorflow/cc/tools/freeze_saved_model_test.cc +++ b/tensorflow/cc/tools/freeze_saved_model_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/cc/tools/freeze_saved_model.h" +#include "tensorflow/cc/ops/resource_variable_ops.h" #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/core/framework/function_testlib.h" #include "tensorflow/core/framework/graph.pb.h" @@ -113,6 +114,160 @@ class FreezeTest : public ::testing::Test { test::ExpectTensorEqual(unfrozen_outputs[0], frozen_outputs[0]); } + + void TestFreezeGraphWithoutDependentVariables(bool use_resource) { + // 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); + if (use_resource) { + Output var = + ops::VarHandleOp(scope.WithOpName("var"), DataType::DT_FLOAT, {}); + Output read_var = ops::ReadVariableOp( + scope.WithOpName("var/Read/ReadVariableOp"), var, DataType::DT_FLOAT); + auto assign = ops::AssignVariableOp(scope.WithOpName("assign"), var, a); + } else { + 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", &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"); + } + + void TestFreezeGraphWithDependentVariables(bool use_resource) { + // 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 read_var; + if (use_resource) { + Output var = + ops::VarHandleOp(scope.WithOpName("var"), DataType::DT_FLOAT, {}); + read_var = ops::ReadVariableOp( + scope.WithOpName("var/Read/ReadVariableOp"), var, DataType::DT_FLOAT); + auto assign = ops::AssignVariableOp(scope.WithOpName("assign"), var, a); + } else { + Output read_var = + ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT); + Output assign = ops::Assign(scope.WithOpName("assign"), read_var, a); + } + Output c = ops::Mul(scope.WithOpName("c"), a, read_var); + 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", &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)); + + // If using normal variables there should be 3 nodes in the resulting + // graph_def. If using resource variables there should be 4 nodes in the + // resulting graph_def. + // In both cases, none should be variables. + size_t expected_nodes = use_resource ? 4 : 3; + EXPECT_EQ(frozen_graph_def.node_size(), expected_nodes); + for (const NodeDef& node : frozen_graph_def.node()) { + EXPECT_NE(node.op(), "Variable") << node.name(); + EXPECT_NE(node.op(), "VariableV2") << node.name(); + EXPECT_NE(node.op(), "VarHandleOp") << node.name(); + EXPECT_NE(node.op(), "ReadVariableOp") << node.name(); + } + + RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(), + frozen_graph_def, "c:0"); + } + + void TestFreezeGraphWithAndWithoutDependentVariables(bool use_resource) { + // 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 read_var; + + if (use_resource) { + Output var = + ops::VarHandleOp(scope.WithOpName("var"), DataType::DT_FLOAT, {}); + read_var = ops::ReadVariableOp( + scope.WithOpName("var/Read/ReadVariableOp"), var, DataType::DT_FLOAT); + auto assign = ops::AssignVariableOp(scope.WithOpName("assign"), var, a); + Output var_1 = + ops::VarHandleOp(scope.WithOpName("var_1"), DataType::DT_FLOAT, {}); + Output read_var_1 = + ops::ReadVariableOp(scope.WithOpName("var_1/Read/ReadVariableOp"), + var, DataType::DT_FLOAT); + auto assign_1 = + ops::AssignVariableOp(scope.WithOpName("assign_1"), var_1, a); + } else { + read_var = ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT); + Output assign = ops::Assign(scope.WithOpName("assign"), read_var, a); + Output var_1 = + ops::Variable(scope.WithOpName("var_1"), {}, DataType::DT_FLOAT); + Output assign_1 = ops::Assign(scope.WithOpName("assign_1"), var_1, a); + } + + Output c = ops::Mul(scope.WithOpName("c"), a, read_var); + 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", &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. + size_t expected_nodes = use_resource ? 4 : 3; + EXPECT_EQ(frozen_graph_def.node_size(), expected_nodes); + for (const NodeDef& node : frozen_graph_def.node()) { + EXPECT_NE(node.op(), "Variable") << node.name(); + EXPECT_NE(node.op(), "VariableV2") << node.name(); + EXPECT_NE(node.op(), "VarHandleOp") << node.name(); + EXPECT_NE(node.op(), "ReadVariableOp") << node.name(); + } + + RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(), + frozen_graph_def, "c:0"); + } }; TEST_F(FreezeTest, InputsAndOutputsSingleSignatureDef) { @@ -196,111 +351,28 @@ TEST_F(FreezeTest, GraphDefWithNoVariables) { 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, GraphDefWithoutDependentVariables) { + TestFreezeGraphWithoutDependentVariables(false); } -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, GraphDefWithoutDependentResourceVariables) { + TestFreezeGraphWithoutDependentVariables(true); } -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)); +TEST_F(FreezeTest, GraphDefWithDependentVariables) { + TestFreezeGraphWithDependentVariables(false); +} - GraphDef frozen_graph_def; - std::unordered_set inputs; - std::unordered_set outputs; - TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, - &outputs)); +TEST_F(FreezeTest, GraphDefWithDependentResourceVariables) { + TestFreezeGraphWithDependentVariables(true); +} - // 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(); - } +TEST_F(FreezeTest, GraphDefWithAndWithoutDependentVariables) { + TestFreezeGraphWithAndWithoutDependentVariables(false); +} - RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(), - frozen_graph_def, "c:0"); +TEST_F(FreezeTest, GraphDefWithAndWithoutDependentResourceVariables) { + TestFreezeGraphWithAndWithoutDependentVariables(true); } } // namespace diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index 0540260efd83e18258ec6e93c514d14e328791b1..ffa2d088295375bbbcd2cdd9365982907f2bf480 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -72,6 +72,7 @@ cc_library( "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", ], ) @@ -132,7 +133,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 @@ -143,7 +146,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 @@ -155,7 +160,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 @@ -166,7 +173,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 @@ -189,7 +198,6 @@ cc_library( name = "benchmark_extra_android", tags = [ "manual", - "notap", ], visibility = ["//visibility:public"], ) diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index c87f2b75dfa18ad5c3eda4bd6fcbcb3083ef73fd..7c833878818022c86fd3171ec9cef9fcd3217a24 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/lib/strings/proto_serialization.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" 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 2b9c83ba149adf9e089786b91039e256216579c8..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,9 +121,9 @@ 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 @@ -213,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, @@ -254,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, ) @@ -283,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. @@ -301,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 @@ -315,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, ) @@ -330,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/jit/BUILD b/tensorflow/compiler/jit/BUILD index a711319607f4ff2b83aa0ebe50e215b3d0e2258e..39eb390f38867070a4ab56a56c1b427aaf61f09d 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -29,7 +29,10 @@ load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured") # Target that bundles up the XLA CPU and GPU JIT devices. cc_library( name = "jit", - visibility = [":friends"], + visibility = [ + ":friends", + "//learning/tfx:__subpackages__", + ], deps = [ ":xla_cpu_device", ":xla_cpu_jit", @@ -102,12 +105,32 @@ cc_library( cc_library( name = "xla_interpreter_device", srcs = ["xla_interpreter_device.cc"], + visibility = [":friends"], deps = [ + ":jit_compilation_passes", ":xla_device", "//tensorflow/compiler/jit/kernels:xla_launch_op", "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla/kernels:xla_ops", + "//tensorflow/compiler/xla/service:interpreter_plugin", # buildcleaner: keep + "//tensorflow/core:lib", + ], + alwayslink = 1, +) + +cc_library( + name = "xla_tensor_info", + srcs = ["xla_tensor_info.cc"], + hdrs = ["xla_tensor_info.h"], + deps = [ + ":common", + "//tensorflow/compiler/xla/service:shaped_buffer", + "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", ], - alwayslink = True, ) cc_library( @@ -127,6 +150,8 @@ cc_library( deps = [ ":common", ":jit_compilation_passes", + ":xla_launch_util", + ":xla_tensor_info", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:dump_graph", @@ -166,6 +191,29 @@ cc_library( visibility = [":friends"], ) +cc_library( + name = "xla_launch_util", + srcs = ["xla_launch_util.cc"], + hdrs = ["xla_launch_util.h"], + deps = [ + ":common", + ":xla_compilation_cache", + ":xla_tensor_info", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla/client:client_library", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core/kernels:variable_ops", + ], +) + cc_library( name = "xla_compilation_cache", srcs = ["xla_compilation_cache.cc"], @@ -200,6 +248,7 @@ cc_library( name = "graph_to_functiondef", srcs = ["graph_to_functiondef.cc"], hdrs = ["graph_to_functiondef.h"], + visibility = [":friends"], deps = [ "//tensorflow/core:core_cpu", "//tensorflow/core:framework", diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index 0de163d3a8f082eab4d8d802485da1bbc56e8180..0685036c9dd53e9117e9f6c10d7cc4511620536e 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" @@ -51,6 +53,8 @@ namespace tensorflow { const char* const kXlaCompiledKernelAttr = "_XlaCompiledKernel"; const char* const kXlaNumConstantArgsAttr = "_XlaNumConstantArgs"; const char* const kXlaNumResourceArgsAttr = "_XlaNumResourceArgs"; +const char* const kXlaHostTransferSequencerAttr = + "_xla_host_transfer_sequencer"; namespace { @@ -141,8 +145,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 +174,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 +205,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 +258,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,28 +321,28 @@ 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. Status MakeSequencingNode(const string& subgraph_name, Graph* graph_out); // If there is a sequencer node, adds a control edge from the sequencer to - // all the downstream nodes of call_node_outputs. - void ConnectSequencerToOutputs(Graph* graph_out); + // the call node. + void ConnectSequencerToCallNode(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 +351,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 +368,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; @@ -367,12 +383,24 @@ class Encapsulator { Node* send_from_host = nullptr; }; + // Creates an outside_compilation subgraph for outside_compilation_id if + // none exists yet. Returns the (possible newly created) subgraph for + // outside_compilation_id. + OutsideCompilationSubgraph* LookupOrCreateOutsideCompilationSubgraph( + const string& outside_compilation_id); + // Builds a ParallelCheck op that compares the output of the original // subgraph with the encapsulated subgraph. Status BuildParallelCheckOp( const std::unordered_map& node_images, Graph* graph_out); + // Builds a placeholder node used to provide the key input to a RecvAtHost + // or SendFromHost node. This placeholder node will be removed by a later + // pass. + Status AddHostComputeKeyPlaceholder(OutsideCompilationSubgraph* oc_subgraph, + Graph* graph_out); + // Builds a _RecvAtHost node producing all the inputs of an // outside_compilation subgraph and stores it in oc_subgraph.recv_at_host. Status AddRecvAtHostNode(const string& subgraph_name, @@ -399,6 +427,14 @@ class Encapsulator { // NodeDef for the function call node. NodeDef call_node_def_; + // Name that is used for the call node. This may not be + // call_node_def_.name() if the client supplies a rewrite lambda. + string function_def_name_; + + // Placeholder node simulating the host compute key in the output graph. + // Not owned. + Node* host_compute_key_placeholder_ = nullptr; + // Function call node(s) in the output graph. Not owned. // If parallel_checking is enabled, 'call_node_inputs' is the function call // node to which inputs should be fed, and 'call_node_outputs' is the @@ -516,6 +552,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_; @@ -645,53 +734,62 @@ Status Encapsulator::Subgraph::RecordResult( return Status::OK(); } -void Encapsulator::Subgraph::RecordOutsideCompilationInputOrControl( - const string& outside_compilation_id, const Edge* edge) { +Encapsulator::Subgraph::OutsideCompilationSubgraph* +Encapsulator::Subgraph::LookupOrCreateOutsideCompilationSubgraph( + const string& outside_compilation_id) { auto iter = outside_compilation_subgraphs_ .emplace(outside_compilation_id, OutsideCompilationSubgraph()) .first; - OutsideCompilationSubgraph& outside_subgraph = iter->second; + OutsideCompilationSubgraph* outside_subgraph = &iter->second; + return outside_subgraph; +} + +void Encapsulator::Subgraph::RecordOutsideCompilationInputOrControl( + const string& outside_compilation_id, const Edge* edge) { + OutsideCompilationSubgraph* outside_subgraph = + LookupOrCreateOutsideCompilationSubgraph(outside_compilation_id); if (edge->IsControlEdge()) { - outside_subgraph.control_inputs.insert(edge->src()); + outside_subgraph->control_inputs.insert(edge->src()); } else { - int input_index = outside_subgraph.inputs.size(); - outside_subgraph.inputs.emplace(NodeSlot(edge->src(), edge->src_output()), - input_index); + int input_index = outside_subgraph->inputs.size(); + outside_subgraph->inputs.emplace(NodeSlot(edge->src(), edge->src_output()), + input_index); } } void Encapsulator::Subgraph::RecordOutsideCompilationOutputOrControl( const string& outside_compilation_id, const Edge* edge) { - auto subgraph_iter = - outside_compilation_subgraphs_ - .emplace(outside_compilation_id, OutsideCompilationSubgraph()) - .first; - OutsideCompilationSubgraph& outside_subgraph = subgraph_iter->second; + OutsideCompilationSubgraph* outside_subgraph = + LookupOrCreateOutsideCompilationSubgraph(outside_compilation_id); if (edge->IsControlEdge()) { - outside_subgraph.control_outputs.insert(edge->dst()); + outside_subgraph->control_outputs.insert(edge->dst()); } else { DataType dtype = edge->dst()->input_type(edge->dst_input()); auto output_iter = - outside_subgraph.outputs_by_src + outside_subgraph->outputs_by_src .emplace(NodeSlot(edge->src(), edge->src_output(), dtype), - outside_subgraph.outputs_by_src.size()) + outside_subgraph->outputs_by_src.size()) .first; int output_index = output_iter->second; - outside_subgraph.outputs_by_dst[NodeSlot(edge->dst(), edge->dst_input())] = + outside_subgraph->outputs_by_dst[NodeSlot(edge->dst(), edge->dst_input())] = output_index; } } -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 +798,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; + } + + 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 send_def; - NodeDefBuilder builder( - strings::StrCat("outside_compilation_", oc_subgraph_name, "_send"), - kSendToHostOp); - builder.Attr("dtypes", dtypes); + 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); + graph_->AddControlEdge(src_image, host_compute); } - } - } - - 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; - } - - 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); } } } @@ -801,25 +869,21 @@ Status Encapsulator::Subgraph::MakeSequencingNode(const string& subgraph_name, NodeDef seq_def; NodeDefBuilder builder(strings::StrCat(subgraph_name, "_sequencer"), "NoOp"); + builder.Attr(kXlaHostTransferSequencerAttr, subgraph_name); + builder.Device(device_); Status s = builder.Finalize(&seq_def); if (!s.ok()) return s; sequencer_ = graph_out->AddNode(seq_def, &s); if (!s.ok()) return s; - sequencer_->set_assigned_device_name(device_); } return Status::OK(); } -void Encapsulator::Subgraph::ConnectSequencerToOutputs(Graph* graph_out) { +void Encapsulator::Subgraph::ConnectSequencerToCallNode(Graph* graph_out) { if (sequencer_ != nullptr) { - std::unordered_set output_dependencies; - for (Node* node : call_node_outputs_->out_nodes()) { - output_dependencies.insert(node); - } - for (Node* node : output_dependencies) { - graph_out->AddControlEdge(sequencer_, node); - } + VLOG(2) << "ConnectSequencerToCallNode"; + graph_out->AddControlEdge(sequencer_, call_node_inputs_); } } @@ -865,6 +929,8 @@ Status Encapsulator::Subgraph::BuildFunctionDef( name = call_node_def_.op(); } + function_def_name_ = name; + FunctionDef fdef; TF_RETURN_IF_ERROR(GraphToFunctionDef(*graph_, name, &fdef)); @@ -882,6 +948,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 = function_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) { @@ -962,9 +1085,36 @@ Status Encapsulator::Subgraph::AddFunctionCallNode( return Status::OK(); } +Status Encapsulator::Subgraph::AddHostComputeKeyPlaceholder( + OutsideCompilationSubgraph* oc_subgraph, Graph* graph_out) { + TensorShapeProto shape_proto; + TensorShape shape({2}); + shape.AsProto(&shape_proto); + GraphDefBuilder::Options options(graph_out, /*status=*/nullptr); + NodeDef key_def; + NodeDefBuilder builder( + strings::StrCat(call_node_def_.name(), "_key_placeholder"), + "Placeholder"); + builder.Attr("dtype", DT_STRING); + builder.Attr("shape", shape_proto); + builder.Attr("_host_compute_call_node", call_node_def_.name()); + Status s = builder.Finalize(&key_def); + if (!s.ok()) return s; + + host_compute_key_placeholder_ = graph_out->AddNode(key_def, &s); + if (!s.ok()) return s; + host_compute_key_placeholder_->set_assigned_device_name(device_); + + return Status::OK(); +} + Status Encapsulator::Subgraph::AddRecvAtHostNode( const string& subgraph_name, const string& oc_subgraph_name, OutsideCompilationSubgraph* oc_subgraph, Graph* graph_out) { + if (host_compute_key_placeholder_ == nullptr) { + TF_RETURN_IF_ERROR(AddHostComputeKeyPlaceholder(oc_subgraph, graph_out)); + } + std::vector dtypes(oc_subgraph->inputs.size(), DT_INVALID); for (const auto& input : oc_subgraph->inputs) { @@ -980,13 +1130,22 @@ Status Encapsulator::Subgraph::AddRecvAtHostNode( NodeDefBuilder builder(strings::StrCat("outside_compilation_", subgraph_name, "_", oc_subgraph_name, "_recv"), kRecvAtHostOp); - builder.Attr("dtypes", dtypes); + // TODO(misard) When we add replication the device placement will have to be + // redone. + builder.Device(device_); + builder.Attr("Toutputs", dtypes); + // TODO(misard) For now we only support TPU device 0. + builder.Attr("device_ordinal", 0); + builder.Attr("key", strings::StrCat("host_compute_channel_", subgraph_name, + "_", oc_subgraph_name)); + builder.Input(host_compute_key_placeholder_->name(), 0, DT_STRING); Status s = builder.Finalize(&recv_def); if (!s.ok()) return s; oc_subgraph->recv_at_host = graph_out->AddNode(recv_def, &s); if (!s.ok()) return s; - oc_subgraph->recv_at_host->set_assigned_device_name(device_); + graph_out->AddEdge(host_compute_key_placeholder_, 0, + oc_subgraph->recv_at_host, 0); // Add a control dependency forcing the RecvAtHost to run before the subgraph // completes. This has no effect on execution order but prevents the @@ -1001,6 +1160,10 @@ Status Encapsulator::Subgraph::AddSendFromHostNode( const std::unordered_map& node_images, const string& subgraph_name, const string& oc_subgraph_name, OutsideCompilationSubgraph* oc_subgraph, Graph* graph_out) { + if (host_compute_key_placeholder_ == nullptr) { + TF_RETURN_IF_ERROR(AddHostComputeKeyPlaceholder(oc_subgraph, graph_out)); + } + std::vector dtypes(oc_subgraph->outputs_by_src.size(), DT_INVALID); std::vector inputs( oc_subgraph->outputs_by_src.size()); @@ -1020,14 +1183,23 @@ Status Encapsulator::Subgraph::AddSendFromHostNode( NodeDefBuilder builder(strings::StrCat("outside_compilation_", subgraph_name, "_", oc_subgraph_name, "_send"), kSendFromHostOp); - builder.Attr("dtypes", dtypes); + // TODO(misard) When we add replication the device placement will have to be + // redone. + builder.Device(device_); + builder.Attr("Tinputs", dtypes); + builder.Attr("key", strings::StrCat("host_compute_channel_", subgraph_name, + "_", oc_subgraph_name)); + // TODO(misard) For now we only support TPU device 0. + builder.Attr("device_ordinal", 0); builder.Input(inputs); + builder.Input(host_compute_key_placeholder_->name(), 0, DT_STRING); Status s = builder.Finalize(&send_def); if (!s.ok()) return s; oc_subgraph->send_from_host = graph_out->AddNode(send_def, &s); if (!s.ok()) return s; - oc_subgraph->send_from_host->set_assigned_device_name(device_); + graph_out->AddEdge(host_compute_key_placeholder_, 0, + oc_subgraph->send_from_host, inputs.size()); // Add a control dependency forcing the SendFromHost to run before the // subgraph completes. This has no effect on execution order but prevents the @@ -1062,6 +1234,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 +1399,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); @@ -1503,14 +1681,368 @@ Status Encapsulator::AddEdgesToOutputGraph( for (auto& subgraph_entry : subgraphs_) { Subgraph& subgraph = subgraph_entry.second; - subgraph.ConnectSequencerToOutputs(graph_out); + subgraph.ConnectSequencerToCallNode(graph_out); } 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()); + // The final input to the send node is the dynamic key, which we don't include + // in the static shapes. + static_shape_out->resize(send_node->num_inputs() - 1); + + // 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); + if (dummy_node_images.find(src_node) == dummy_node_images.end()) { + dummy_node_images[src_node] = AddDummyShapedNode( + src_node->output_type(src_port), proto, graph_out.get()); + } + // The final input to the send node is the dynamic key, which we + // don't include in the static shapes. + if (n == send_node && + in_edge->dst_input() < static_shape_out->size()) { + (*static_shape_out)[in_edge->dst_input()] = proto; + } + } else { + has_parent_with_unknown_shape = true; + if (!visited[src_node->id()]) { + if (VLOG_IS_ON(2)) { + TensorShapeProto proto; + context->ShapeHandleToProto(shape, &proto); + } + 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)); + + if (VLOG_IS_ON(1)) { + dump_graph::DumpGraphToFile("pruned_graph_for_shape_inference", + *pruned_graph, 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 +2054,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 +2080,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..711b1424c7c121b9fd59981c271dc9f310a0fbdf 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -29,17 +29,184 @@ limitations under the License. namespace tensorflow { namespace { +const char* const kXlaHostTransferSequencerAttr = + "_xla_host_transfer_sequencer"; + +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; } @@ -82,31 +249,72 @@ bool EqualFunctionDefLibrary(const FunctionDefLibrary& expected, << diff << "\nActual: " << actual.DebugString(); \ } while (false) -// 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"); +// These dummy Op registrations are here because the real Op registrations live +// in contrib and there can't be a dependence from this test to contrib. +REGISTER_OP("XlaHostCompute") + .Input("inputs: Tinputs") + .Output("outputs: Toutputs") + .Attr("Tinputs: list(type) >= 0") + .Attr("Toutputs: list(type) >= 0") + .Attr("key: string") + .Attr("shape_inference_graph: string = ''") + .Attr("shapes: list(shape) >= 0") + .SetShapeFn(::tensorflow::shape_inference::UnknownShape); REGISTER_OP("_XlaSendFromHost") - .Input("input: dtypes") - .Attr("dtypes: list(type) >= 0"); + .Input("inputs: Tinputs") + .Input("dynamic_key: string") + .Attr("Tinputs: list(type) >= 0") + .Attr("key: string") + .Attr("device_ordinal: int") + .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"); + .Input("dynamic_key: string") + .Output("outputs: Toutputs") + .Attr("Toutputs: list(type) >= 0") + .Attr("key: string") + .Attr("device_ordinal: int") + .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") @@ -116,30 +324,89 @@ REGISTER_OP("AddNLikeTest") .SetIsCommutative() .SetIsAggregate(); -Node* NoOp(const GraphDefBuilder::Options& opts) { - return ops::SourceOp("NoOp", opts); +Node* Sequencer(const GraphDefBuilder::Options& opts, + const string& call_node_name) { + if (opts.HaveError()) return nullptr; + NodeBuilder node_builder(opts.GetNameForOp("NoOp"), "NoOp", + opts.op_registry()); + return opts.WithAttr(kXlaHostTransferSequencerAttr, call_node_name) + .FinalizeBuilder(&node_builder); } 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* KnownShapeBase(DataType dtype, 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(dtype); + for (int dim : shape) { + value.mutable_tensor_shape()->add_dim()->set_size(dim); + } + return opts.WithAttr("value", value) + .WithAttr("dtype", dtype) + .FinalizeBuilder(&node_builder); +} + +Node* KnownShape(const gtl::ArraySlice& shape, + const GraphDefBuilder::Options& opts) { + return KnownShapeBase(DT_FLOAT, shape, opts); +} + +Node* KeyPlaceholderShape(const GraphDefBuilder::Options& opts) { + return KnownShapeBase(DT_STRING, {2}, opts); +} + +Node* KeyPlaceholder(const string& call_node, + const GraphDefBuilder::Options& opts) { + if (opts.HaveError()) return nullptr; + NodeBuilder node_builder(opts.GetNameForOp("Placeholder"), "Placeholder", + opts.op_registry()); + TensorShapeProto shape; + shape.add_dim()->set_size(2); + return opts.WithAttr("shape", shape) + .WithAttr("dtype", DT_STRING) + .WithAttr("_host_compute_call_node", call_node) + .FinalizeBuilder(&node_builder); +} + +Node* RecvAtHost(ops::NodeOut key_input, 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); + node_builder.Input(std::move(key_input)); + return opts.WithAttr("Toutputs", dtypes) + .WithAttr("key", key) + .WithAttr("device_ordinal", 0) + .FinalizeBuilder(&node_builder); } -Node* SendFromHost(const std::vector& inputs, - const gtl::ArraySlice& dtypes, +Node* SendFromHost(ops::NodeOut key_input, 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); + node_builder.Input(std::move(key_input)); + std::vector dtypes; + for (const auto& node : inputs) { + dtypes.push_back(node.dt); + } + return opts.WithAttr("Tinputs", dtypes) + .WithAttr("key", key) + .WithAttr("device_ordinal", 0) + .FinalizeBuilder(&node_builder); } Node* Unary(ops::NodeOut a, const GraphDefBuilder::Options& opts) { @@ -151,6 +418,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 +848,24 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + string shape_string_expected; + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); + Node* recv = + RecvAtHost(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), "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 +874,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"}}); @@ -607,24 +896,30 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - NodeBuilder node_builder("F1", "F1", lib_def.get()); - node_builder.Input(a).Input(b); - Node* call = b2.opts().FinalizeBuilder(&node_builder); - + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* recv = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), + "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})); + Node* s = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), + "F1"); - Binary(a, call, b2.opts().WithName("G").WithControlInputs({s, e})); + NodeBuilder node_builder("F1", "F1", lib_def.get()); + node_builder.Input(a).Input(b); + Node* call = + b2.opts().WithControlInputs({s}).FinalizeBuilder(&node_builder); + + Binary(a, call, b2.opts().WithName("G").WithControlInputs({e})); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -674,37 +969,78 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + string shape_string_expected_1; + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape1.opts().WithName("KnownShape/_0")); + Node* recv = + RecvAtHost(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), "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* key_constant = + KeyPlaceholderShape(shape2.opts().WithName("KnownShape/_0")); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), "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"}}); @@ -715,34 +1051,41 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - NodeBuilder node_builder("F1", "F1", lib_def.get()); - node_builder.Input(a).Input(b); - Node* call = b2.opts().FinalizeBuilder(&node_builder); - + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* recv1 = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), + "host_compute_channel_F1_O1", {e}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); Node* recv2 = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost(ops::NodeOut(key_constant, 0), "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")); + ops::NodeOut(key_constant, 0), "host_compute_channel_F1_O2", {h}, + b2.opts().WithName("outside_compilation_F1_O2_send")); - Node* s = NoOp(b2.opts() - .WithName("F1_sequencer") - .WithControlInputs({recv1, send1, recv2, send2})); + Node* s = Sequencer(b2.opts() + .WithName("F1_sequencer") + .WithControlInputs({recv1, send1, recv2, send2}), + "F1"); - Binary(g, call, b2.opts().WithName("J").WithControlInput(s)); + NodeBuilder node_builder("F1", "F1", lib_def.get()); + node_builder.Input(a).Input(b); + Node* call = b2.opts().WithControlInput(s).FinalizeBuilder(&node_builder); + + Binary(g, call, b2.opts().WithName("J")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -758,8 +1101,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 +1134,27 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + string shape_string_expected; + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); + Node* recv = + RecvAtHost(ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), "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 +1163,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 +1185,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,40 +1202,50 @@ 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* key_constant1 = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* recv1 = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost(ops::NodeOut(key_constant1, 0), "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(ops::NodeOut(key_constant1, 0), + "host_compute_channel_F1_O1", {e}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}), + "F1"); + NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); - Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Node* s1 = NoOp( - b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1})); + Node* call1 = + b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); + Node* key_constant2 = + KeyPlaceholder("F2", b2.opts().WithName("F2_key_placeholder")); Node* recv2 = RecvAtHost( + ops::NodeOut(key_constant2, 0), "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* h = Binary(ops::NodeOut(call1, 1), recv2, b2.opts().WithName("H")); Node* send2 = SendFromHost( - {h}, {DT_FLOAT}, b2.opts().WithName("outside_compilation_F2_O1_send")); + ops::NodeOut(key_constant2, 0), "host_compute_channel_F2_O1", {h}, + b2.opts().WithName("outside_compilation_F2_O1_send")); + Node* s2 = Sequencer( + b2.opts().WithName("F2_sequencer").WithControlInputs({recv2, send2}), + "F2"); NodeBuilder node_builder2("F2", "F2", lib_def.get()); node_builder2.Input(e).Input(call1); Node* call2 = b2.opts() - .WithControlInputs({s1, e, call1}) + .WithControlInputs({s2, e, call1}) .FinalizeBuilder(&node_builder2); - Node* s2 = NoOp( - b2.opts().WithName("F2_sequencer").WithControlInputs({recv2, send2})); - Binary(call2, ops::NodeOut(call2, 1), - b2.opts().WithName("J").WithControlInput(s2)); + Binary(call2, ops::NodeOut(call2, 1), b2.opts().WithName("J")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -888,7 +1261,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 +1281,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 +1291,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,18 +1308,23 @@ 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* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* send1 = SendFromHost( - {e}, {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_send")); + ops::NodeOut(key_constant, 0), "host_compute_channel_F1_O1", {e}, + b2.opts().WithName("outside_compilation_F1_O1_send")); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInput(send1), "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); - Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Node* s1 = NoOp(b2.opts().WithName("F1_sequencer").WithControlInput(send1)); + Node* call1 = + b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); - Unary(call1, b2.opts().WithName("G").WithControlInput(s1)); + Unary(call1, b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -954,7 +1340,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 +1361,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 +1371,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,21 +1389,27 @@ 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* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* recv1 = - RecvAtHost({}, b2.opts().WithName("outside_compilation_F1_O1_recv")); + RecvAtHost(ops::NodeOut(key_constant, 0), "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")); + ops::NodeOut(key_constant, 0), "host_compute_channel_F1_O1", {e}, + b2.opts().WithName("outside_compilation_F1_O1_send")); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); - Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Node* s1 = NoOp( - b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1})); + Node* call1 = + b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); - Unary(call1, b2.opts().WithName("G").WithControlInput(s1)); + Unary(call1, b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1055,10 +1450,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,15 +1468,20 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* recv1 = RecvAtHost( - {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); + ops::NodeOut(key_constant, 0), "host_compute_channel_F1_O1", {DT_FLOAT}, + b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Unary(recv1, b2.opts().WithName("E")); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInput(recv1), "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); - Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Node* s1 = NoOp(b2.opts().WithName("F1_sequencer").WithControlInput(recv1)); + Node* call1 = + b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("G").WithControlInput(s1)); + Binary(e, call1, b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1118,16 +1522,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,20 +1545,26 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); Node* recv1 = RecvAtHost( - {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); + ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), + "host_compute_channel_F1_O1", {}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); - Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Node* s1 = NoOp( - b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1})); + Node* call1 = + b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("G").WithControlInput(s1)); + Binary(e, call1, b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1215,5 +1628,116 @@ 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* key_constant = + KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); + Node* known = KnownShape({2}, shape.opts().WithName("KnownShape/_1")); + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), "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")); + + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "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(ops::NodeOut(key_constant, 0), + "host_compute_channel_F1_O1", {e}, + b2.opts() + .WithName("outside_compilation_F1_O1_send") + .WithControlInput(e)); + + Node* s = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), + "F1"); + + NodeBuilder node_builder("F1", "F1", lib_def.get()); + node_builder.Input(b).Input(c); + Node* call = + b2.opts().WithControlInputs({s, c}).FinalizeBuilder(&node_builder); + + Binary(a, call, b2.opts().WithName("G").WithControlInputs({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/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 9bea5663319c8a25249fdc265cee0191556a7c04..616a7f8f1541d3debff97a90bd390c76c665d196 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -14,6 +14,7 @@ cc_library( "//tensorflow/compiler/jit:common", "//tensorflow/compiler/jit:xla_compilation_cache", "//tensorflow/compiler/jit:xla_device", + "//tensorflow/compiler/jit:xla_launch_util", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/xla:statusor", diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_launch_op.cc index 17ae2bb25cac94cee0a1f3df66edf1b3a404e3ec..e24a9a075135f077d05b7e06206d132c530785a3 100644 --- a/tensorflow/compiler/jit/kernels/xla_launch_op.cc +++ b/tensorflow/compiler/jit/kernels/xla_launch_op.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/jit/xla_device.h" +#include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" @@ -40,111 +41,6 @@ namespace gpu = perftools::gputools; namespace tensorflow { -// Adapter class that wraps a Tensorflow allocator as an XLA allocator. -// Assumes that the Tensorflow allocator permits asynchronous deallocation: -// see comment on `AllowsAsynchronousDeallocation()`. -class XlaAllocator : public xla::DeviceMemoryAllocator { - public: - XlaAllocator(const gpu::Platform* platform, OpKernelContext* op_context); - ~XlaAllocator() override; - xla::StatusOr Allocate(int device_ordinal, uint64 size, - bool retry_on_failure) override; - Status Deallocate(int device_ordinal, gpu::DeviceMemoryBase* mem) override; - - // Register an Tensor (input or resource variable) with the allocator. If - // the operation returns an alias to one of its inputs, then the allocator - // needs to be able to handle it. - Status RegisterArgument(const Tensor* t); - - // Makes 'tensor' a wrapper around the data buffer at 'ptr'. The buffer is - // interpreted as having data type 'dtype' and shape 'shape'. - Status MakeTensorFromBuffer(gpu::DeviceMemoryBase buffer, DataType dtype, - const TensorShape& shape, Tensor* tensor) const; - - // The Tensorflow BFC allocator used on GPU allows host-side deallocation - // before GPU execution takes place. Tensorflow uses the ordering of the main - // compute stream to enforce a happens-before relationship between a memory - // allocation and code that reuses the same memory. If Tensorflow adds - // support for multiple GPU streams or allocators with different ordering - // requirements, this code may need to change. - // (This attribute has no effect on CPU.) - bool AllowsAsynchronousDeallocation() const override { return true; } - - private: - OpKernelContext* const op_context_; - - // Map from pointer address to the owning Tensor; used by - // MakeTensorFromBuffer. Also used to automatically release Tensors when the - // allocator is freed. - std::unordered_map tensors_; -}; - -XlaAllocator::XlaAllocator(const gpu::Platform* platform, - OpKernelContext* op_context) - : xla::DeviceMemoryAllocator(platform), op_context_(op_context) {} - -XlaAllocator::~XlaAllocator() = default; - -xla::StatusOr XlaAllocator::Allocate( - int device_ordinal, uint64 size, bool retry_on_failure) { - AllocatorAttributes allocator_attrs; - allocator_attrs.set_on_host(false); - - AllocationAttributes allocation_attrs; - allocation_attrs.no_retry_on_failure = !retry_on_failure; - - Tensor t; - Status status = op_context_->allocate_temp( - DT_UINT8, TensorShape({static_cast(size)}), &t, allocator_attrs, - allocation_attrs); - if (!status.ok()) { - VLOG(2) << "Allocation failed " << size; - return status; - } - void* data = - reinterpret_cast(const_cast(t.tensor_data().data())); - tensors_[data] = t; - return gpu::DeviceMemoryBase(data, size); -} - -Status XlaAllocator::RegisterArgument(const Tensor* t) { - void* data = - reinterpret_cast(const_cast(t->tensor_data().data())); - tensors_[data] = *t; - return Status::OK(); -} - -Status XlaAllocator::Deallocate(int device_ordinal, - gpu::DeviceMemoryBase* mem) { - if (mem->opaque() != nullptr) { - if (tensors_.erase(mem->opaque()) == 0) { - return tensorflow::errors::InvalidArgument("Unknown tensor address"); - } - } - return Status::OK(); -} - -Status XlaAllocator::MakeTensorFromBuffer(gpu::DeviceMemoryBase buffer, - DataType dtype, - const TensorShape& shape, - Tensor* out_tensor) const { - void* ptr = const_cast(buffer.opaque()); - auto it = tensors_.find(ptr); - if (it == tensors_.end()) { - return errors::InvalidArgument("Unknown tensor address"); - } - const Tensor& tensor = it->second; - - int64 output_size = DataTypeSize(dtype) * shape.num_elements(); - if (tensor.TotalBytes() == output_size) { - out_tensor->UnsafeCopyFromInternal(tensor, dtype, shape); - } else { - Tensor slice = tensor.Slice(0, output_size); - out_tensor->UnsafeCopyFromInternal(slice, dtype, shape); - } - return Status::OK(); -} - XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx) : OpKernel(ctx), device_type_(ctx->device_type()) { const NameAttrList* func; @@ -196,23 +92,6 @@ Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx, return Status::OK(); } -std::vector SnapshotResourceVariables(OpKernelContext* ctx, - int num_variables) { - std::vector snapshot(num_variables); - int first_variable = ctx->num_inputs() - num_variables; - for (int i = 0; i < num_variables; ++i) { - Var* variable = nullptr; - ResourceHandle handle = HandleFromInput(ctx, first_variable + i); - if (LookupResource(ctx, handle, &variable).ok()) { - tf_shared_lock lock(*variable->mu()); - snapshot[i].name = handle.name(); - snapshot[i].present = true; - snapshot[i].value = *variable->tensor(); - } - } - return snapshot; -} - void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { VLOG(1) << "XlaLocalLaunchOp::Compute " << Canonicalize(function_.name(), AttrSlice(&function_.attr())); @@ -235,16 +114,22 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { // this is more obviously correct.) core::ScopedUnref cache_ref(cache); + const XlaDevice::Metadata* metadata; + Status s = XlaDevice::GetMetadata(ctx, &metadata); + + XlaTensorInfoManager* tensor_info_manager = nullptr; + if (s.ok()) { + tensor_info_manager = &metadata->tensor_info_manager(); + } + // Get the platform_id_ for XLA_* devices. if (platform_id_ == nullptr) { - const XlaDevice::Metadata* metadata; - Status s = XlaDevice::GetMetadata(ctx, &metadata); if (s.ok()) { platform_id_ = metadata->platform()->id(); } } - std::vector variables = + std::map variables = SnapshotResourceVariables(ctx, num_resource_args_); xla::LocalClient* client = static_cast(cache->client()); @@ -269,43 +154,9 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { VLOG(1) << "Executing XLA Computation..."; - std::unique_ptr output; - // Build xla::ShapedBuffers that point directly to the Tensor buffers. - std::vector> arg_buffers; - arg_buffers.reserve(kernel->xla_input_shapes.size() + 1); - arg_buffers.resize(kernel->xla_input_shapes.size()); - std::vector arg_ptrs(arg_buffers.size()); - - const int first_variable_arg = ctx->num_inputs() - num_resource_args_; - // Pass remaining parameters. - const Tensor* t; - for (int i = 0; i < kernel->xla_input_shapes.size(); ++i) { - int arg_num = kernel->input_mapping[i]; - const xla::Shape& shape = kernel->xla_input_shapes[i]; - if (arg_num >= first_variable_arg) { - t = &(variables[arg_num - first_variable_arg].value); - } else { - t = &(ctx->input(arg_num)); - } - - gpu::DeviceMemoryBase dmem = gpu::DeviceMemoryBase( - const_cast(t->tensor_data().data()), t->tensor_data().size()); - - const xla::Shape on_device_shape = - client->backend().transfer_manager()->HostShapeToDeviceShape(shape); - CHECK(xla::ShapeUtil::Equal(shape, on_device_shape)) - << "On-device shape " - << xla::ShapeUtil::HumanStringWithLayout(on_device_shape) - << " not the same as on-host shape " - << xla::ShapeUtil::HumanStringWithLayout(shape); - arg_buffers[i] = xla::MakeUnique( - /*on_host_shape=*/shape, /*on_device_shape=*/shape, client->platform(), - client->default_device_ordinal()); - arg_buffers[i]->set_buffer(dmem, /*index=*/{}); - arg_ptrs[i] = arg_buffers[i].get(); - - OP_REQUIRES_OK(ctx, xla_allocator.RegisterArgument(t)); - } + XlaComputationLaunchContext launch_context( + num_resource_args_, client, &xla_allocator, tensor_info_manager); + launch_context.PopulateInputs(ctx, kernel, variables); // Execute the computation. VLOG(2) << "Executing computation."; @@ -315,95 +166,13 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); Env* env = Env::Default(); auto start_time = env->NowMicros(); - auto run_result = executable->Run(arg_ptrs, run_options); + auto run_result = executable->Run(launch_context.arguments(), run_options); OP_REQUIRES(ctx, run_result.ok(), run_result.status()); - output = run_result.ConsumeValueOrDie()->release(); auto elapsed = env->NowMicros() - start_time; VLOG(2) << "Elapsed time: " << elapsed << "us"; - // Computation output should always be a tuple. - if (VLOG_IS_ON(2)) { - VLOG(2) << "Result tuple shape: " << output->on_host_shape().DebugString(); - } - CHECK_EQ(ctx->num_outputs(), kernel->outputs.size()); - - // Copy XLA results to the OpOutputList. - int output_num = 0; - for (int i = 0; i < ctx->num_outputs(); ++i) { - if (kernel->outputs[i].is_constant) { - // Output is a constant. - const Tensor& const_tensor = kernel->outputs[i].constant_value; - const size_t total_bytes = const_tensor.TotalBytes(); - if (stream && total_bytes > 0) { - // Copy host -> device. (Empty tensors don't have backing buffers.) - VLOG(1) << "Constant output tensor on device"; - Tensor* output_tensor; - TF_CHECK_OK( - ctx->allocate_output(i, const_tensor.shape(), &output_tensor)); - - const void* src_ptr = DMAHelper::base(&const_tensor); - void* dst_ptr = DMAHelper::base(output_tensor); - gpu::DeviceMemoryBase gpu_dst_ptr(dst_ptr, total_bytes); - stream->ThenMemcpy(&gpu_dst_ptr, src_ptr, total_bytes); - } else { - // No copy required. - ctx->set_output(i, const_tensor); - } - } else { - const TensorShape& shape = kernel->outputs[i].shape; - VLOG(2) << "Retval " << i << " shape " << shape.DebugString(); - - gpu::DeviceMemoryBase buffer = output->buffer({output_num}); - Tensor output_tensor; - // Looks up the owning Tensor by buffer address. - OP_REQUIRES_OK(ctx, xla_allocator.MakeTensorFromBuffer( - buffer, ctx->expected_output_dtype(i), shape, - &output_tensor)); - ctx->set_output(i, output_tensor); - ++output_num; - } - - if (VLOG_IS_ON(3)) { - VLOG(3) << ctx->mutable_output(i)->DebugString(); - } - } - - // Apply variable updates, if any. - VLOG(2) << "Applying variable updates"; - for (int i = 0; i < kernel->resource_updates.size(); ++i) { - const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i]; - 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}); - - Var* variable = nullptr; - // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor, not - // a Tensor. - OP_REQUIRES_OK(ctx, LookupOrCreateResource( - ctx, HandleFromInput(ctx, write.input_index), - &variable, [this, ctx, &write](Var** ptr) { - *ptr = new Var(write.type); - return Status::OK(); - })); - - core::ScopedUnref s(variable); - - mutex_lock ml(*variable->mu()); - OP_REQUIRES(ctx, variable->tensor()->dtype() == write.type, - errors::Internal("Mismatched type in variable write")); - - // Looks up the owning Tensor by buffer address. - OP_REQUIRES_OK( - ctx, xla_allocator.MakeTensorFromBuffer(buffer, write.type, write_shape, - variable->tensor())); - ++output_num; - } - + launch_context.PopulateOutputs(ctx, kernel, run_result.ConsumeValueOrDie()); VLOG(1) << "Done"; } diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.h b/tensorflow/compiler/jit/kernels/xla_launch_op.h index 47fd912b12abbbe876e933ab57f6f586fd299909..c6cc0986af0300c51283d432c671e92a1e4d8145 100644 --- a/tensorflow/compiler/jit/kernels/xla_launch_op.h +++ b/tensorflow/compiler/jit/kernels/xla_launch_op.h @@ -26,14 +26,6 @@ limitations under the License. namespace tensorflow { -// Takes a snapshot of the values of resource variable arguments, which are -// the last `num_variables` arguments. We snapshot tensors that back -// resource variables since concurrent updates may modify the shape, and it is -// important that the shapes used for compilation match the true shapes of the -// buffers. -std::vector SnapshotResourceVariables(OpKernelContext* ctx, - int num_variables); - // XlaLocalLaunchOp is used to replace a region of the TensorFlow graph // which will be compiled and executed using XLA. The XlaLocalLaunchOp is // responsible for handling interactions with the TensorFlow executor. diff --git a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc b/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc index 4bc209b7ecf499d82e7567f7eff12b17cefa9863..51384ac2fe6fa70c8a723097093a0a29e7ad2c6b 100644 --- a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc +++ b/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc @@ -40,6 +40,7 @@ static void AllocateFlags() { flags->tf_xla_max_cluster_size = std::numeric_limits::max(); flags->tf_xla_clustering_debug = false; flags->tf_xla_cpu_global_jit = false; + flags->tf_xla_clustering_fuel = std::numeric_limits::max(); flag_list = new std::vector( {Flag("tf_xla_auto_jit", &flags->tf_xla_auto_jit, "Control compilation of operators into XLA computations on CPU and " @@ -55,7 +56,10 @@ static void AllocateFlags() { Flag("tf_xla_clustering_debug", &flags->tf_xla_clustering_debug, "Dump graphs during XLA compilation."), Flag("tf_xla_cpu_global_jit", &flags->tf_xla_cpu_global_jit, - "Enables global JIT compilation for CPU via SessionOptions.")}); + "Enables global JIT compilation for CPU via SessionOptions."), + Flag("tf_xla_clustering_fuel", &flags->tf_xla_clustering_fuel, + "Places an artificial limit on the number of ops marked as " + "eligible for clustering.")}); xla::legacy_flags::ParseFlagsFromEnv(*flag_list); } diff --git a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h b/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h index e1ccd7ddb8706ca445b6811ca1fec369af7cd5d5..170b89c987f30f985f981d7835b4af455922594e 100644 --- a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h +++ b/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h @@ -48,6 +48,9 @@ typedef struct { bool tf_xla_clustering_debug; // Dump graphs during XLA compilation. bool tf_xla_cpu_global_jit; // Enables global JIT compilation for CPU // via SessionOptions. + int64 tf_xla_clustering_fuel; // "Compiler fuel" for clustering. Only this + // many ops will be marked as eligible for + // clustering. } MarkForCompilationPassFlags; // Return a pointer to the MarkForCompilationPassFlags struct; diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 79b02baba83cb47f4f2f16544ad711a4b6937d90..57fb8d242208318a608b1f356bef7a8d39dbdc83 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -174,7 +174,9 @@ bool HasResourceInputOrOutput(const Node& node) { } struct NodeCompare { - bool operator()(const Node* a, const Node* b) { return a->id() < b->id(); } + bool operator()(const Node* a, const Node* b) const { + return a->id() < b->id(); + } }; using OrderedNodeSet = std::set; @@ -189,7 +191,30 @@ Status FindCompilationCandidates( FunctionLibraryRuntime* lib_runtime = pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice); + int64& fuel = + legacy_flags::GetMarkForCompilationPassFlags()->tf_xla_clustering_fuel; + + // Iterate over nodes in sorted order so that compiler fuel is deterministic. + // We can't simply pass op_nodes().begin() and op_nodes().end to the + // std::vector constructor because they're not proper iterators, with + // iterator_traits defined and so on. + std::vector sorted_nodes; for (Node* node : graph.op_nodes()) { + sorted_nodes.push_back(node); + } + std::sort(sorted_nodes.begin(), sorted_nodes.end(), NodeCompare()); + + for (Node* node : sorted_nodes) { + VLOG(2) << "Fuel: " << fuel; + if (fuel <= 0) { + VLOG(2) + << "Hit fuel limit; not marking any remaining ops as clusterable."; + break; + } + + VLOG(2) << "FindCompilationCandidates(): Processing " + << node->DebugString(); + DeviceType device_type(""); TF_RETURN_IF_ERROR( DeviceTypeOfDevice(node->assigned_device_name(), &device_type)); @@ -216,6 +241,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") { @@ -224,7 +256,9 @@ Status FindCompilationCandidates( continue; } candidates->insert(node); + --fuel; } + VLOG(2) << "candidates->size() = " << candidates->size(); return Status::OK(); } @@ -304,6 +338,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 21d3a54f1b8ea59c3da09d8c5d626a9f5bcebbc9..8cc79a9bd0b7aa2098ce177a9d7749f4e6c6ac27 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -93,7 +93,7 @@ uint64 XlaCompilationCache::Signature::Hash::operator()( Status XlaCompilationCache::BuildSignature( const NameAttrList& function, int num_constant_args, - const std::vector& variable_args, OpKernelContext* ctx, + const std::map& variable_args, OpKernelContext* ctx, Signature* signature) { signature->name = Canonicalize(function.name(), AttrSlice(&function.attr())); signature->arg_values.resize(num_constant_args); @@ -115,7 +115,8 @@ Status XlaCompilationCache::BuildSignature( } // For variable signatures, use the type and shape of the variable's // current value. - for (const OptionalTensor& variable : variable_args) { + for (auto& iterator : variable_args) { + const OptionalTensor& variable = iterator.second; TF_RET_CHECK(input_num < ctx->num_inputs()); if (variable.present) { signature->arg_types.emplace_back(variable.value.dtype(), @@ -133,7 +134,7 @@ namespace { // Builds a XlaCompiler::Argument vector from the arguments to the _XlaLaunch // op. The first `num_constant_args` arguments must be host-memory Tensors. Status BuildArguments(int num_constant_args, - const std::vector& variable_args, + const std::map& variable_args, OpKernelContext* ctx, std::vector* args) { args->resize(ctx->num_inputs()); @@ -148,8 +149,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,27 +170,25 @@ 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; } // Handles resource variables. TF_RET_CHECK(input_num + num_variable_args == ctx->num_inputs()); - for (int variable_id = 0; variable_id < num_variable_args; ++variable_id) { + for (auto& iterator : variable_args) { const Tensor& input = ctx->input(input_num); TF_RET_CHECK(input.dtype() == DT_RESOURCE); XlaCompiler::Argument& arg = (*args)[input_num]; - arg.name = variable_args[variable_id].name; + arg.name = iterator.second.name; arg.kind = XlaCompiler::Argument::kResource; arg.resource_kind = XlaResource::kVariable; - if (variable_args[variable_id].present) { - const Tensor& value = variable_args[variable_id].value; + if (iterator.second.present) { + const Tensor& value = iterator.second.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 +197,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; } @@ -236,7 +234,7 @@ Status XlaCompilationCache::BuildExecutable( Status XlaCompilationCache::Compile( const XlaCompiler::Options& options, const NameAttrList& function, - int num_constant_args, const std::vector& variable_args, + int num_constant_args, const std::map& variable_args, OpKernelContext* ctx, const XlaCompiler::CompilationResult** compilation_result, xla::LocalExecutable** executable, @@ -253,10 +251,12 @@ Status XlaCompilationCache::Compile( << " present=" << ctx->has_input(i) << " shape=" << shape.DebugString(); } - for (const OptionalTensor& variable : variable_args) { + for (auto& iterator : variable_args) { + const OptionalTensor& variable = iterator.second; VLOG(2) << "variable present=" << variable.present << " type=" << DataTypeString(variable.value.dtype()) - << " shape=" << variable.value.shape().DebugString(); + << " shape=" << variable.value.shape().DebugString() + << " TF arg= " << iterator.first; } VLOG(2) << "num_outputs = " << ctx->num_outputs(); for (int i = 0; i < ctx->num_outputs(); i++) { diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h index 0858020716fcf4763e42dc0699ad22cfda756942..d5063783140205db54e673a7c7fd8f94b8aa2c65 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.h +++ b/tensorflow/compiler/jit/xla_compilation_cache.h @@ -63,7 +63,7 @@ class XlaCompilationCache : public ResourceBase { // outputs. Status Compile(const XlaCompiler::Options& options, const NameAttrList& function, int num_constant_args, - const std::vector& variable_args, + const std::map& variable_args, OpKernelContext* ctx, const XlaCompiler::CompilationResult** compilation_result, xla::LocalExecutable** executable, @@ -105,7 +105,7 @@ class XlaCompilationCache : public ResourceBase { // Builds the signature for a compilation. Status BuildSignature(const NameAttrList& function, int num_constant_args, - const std::vector& variable_args, + const std::map& variable_args, OpKernelContext* ctx, Signature* signature); // The value associated with a cache entry. diff --git a/tensorflow/compiler/jit/xla_cpu_device.cc b/tensorflow/compiler/jit/xla_cpu_device.cc index e238252751e677eb947f6df03e3b2f2e948ffe19..db3bf3ea33fb1e74a70dee8432cfdbd5a5b26a3a 100644 --- a/tensorflow/compiler/jit/xla_cpu_device.cc +++ b/tensorflow/compiler/jit/xla_cpu_device.cc @@ -39,9 +39,10 @@ Status XlaCpuDeviceFactory::CreateDevices(const SessionOptions& options, (void)registrations; std::unique_ptr device; - TF_RETURN_IF_ERROR(XlaDevice::Create( - "Host", DEVICE_XLA_CPU, 0, DEVICE_CPU_XLA_JIT, options, name_prefix, - /*register_device_for_compilation=*/true, &device)); + TF_RETURN_IF_ERROR(XlaDevice::Create("Host", DEVICE_XLA_CPU, 0, + DEVICE_CPU_XLA_JIT, options, name_prefix, + /*register_device_for_compilation=*/true, + /*transfer_as_literal=*/false, &device)); devices->push_back(device.release()); return Status::OK(); } diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc index d4d8fe1c1d575b4e35d624621cc709e3a16569d5..e4e11d4ce2d389338a01bb9f61dffddc85c1d76a 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -109,7 +109,7 @@ XlaDeviceAllocator* XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( const string& platform_name, const string& device_name, int device_ordinal, const string& jit_device_name, const SessionOptions& options, const string& name_prefix, bool register_device_for_compilation, - std::unique_ptr* device) { + bool transfer_as_literal, std::unique_ptr* device) { VLOG(1) << "XlaDevice::Create " << platform_name << " " << device_name << ":" << device_ordinal; @@ -137,15 +137,17 @@ XlaDeviceAllocator* XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( device->reset(new XlaDevice(options, attrs, device_ordinal, DeviceType(jit_device_name), - platform.ValueOrDie())); + platform.ValueOrDie(), transfer_as_literal)); return Status::OK(); } -XlaDevice::Metadata::Metadata(int device_ordinal, se::Platform* platform, - const DeviceType& device_type) +XlaDevice::Metadata::Metadata( + int device_ordinal, se::Platform* platform, const DeviceType& device_type, + std::unique_ptr* tensor_info_manager) : device_ordinal_(device_ordinal), device_type_(device_type), - platform_(platform) {} + platform_(platform), + tensor_info_manager_(*tensor_info_manager) {} int XlaDevice::Metadata::device_ordinal() const { return device_ordinal_; } @@ -160,6 +162,10 @@ const DeviceType& XlaDevice::Metadata::jit_device_type() const { return device_type_; } +XlaTensorInfoManager& XlaDevice::Metadata::tensor_info_manager() const { + return *tensor_info_manager_; +} + /* static */ Status XlaDevice::GetMetadata(OpKernelContext* ctx, const Metadata** metadata) { XlaDevice* xla_device = @@ -177,13 +183,19 @@ const DeviceType& XlaDevice::Metadata::jit_device_type() const { XlaDevice::XlaDevice(const SessionOptions& options, const DeviceAttributes& attrs, int device_ordinal, - const DeviceType& jit_device_name, se::Platform* platform) + const DeviceType& jit_device_name, se::Platform* platform, + bool transfer_as_literal) : LocalDevice(options, attrs), - xla_metadata_(device_ordinal, platform, jit_device_name), + xla_metadata_( + device_ordinal, platform, jit_device_name, + // Pass tensor_info_manager_ by reference as it is initialized lazily. + &tensor_info_manager_), device_ordinal_(device_ordinal), jit_device_name_(jit_device_name), xla_allocator_(nullptr), - platform_(platform) {} + platform_(platform), + tensor_info_manager_(nullptr), + transfer_as_literal_(transfer_as_literal) {} XlaDevice::~XlaDevice() {} @@ -208,6 +220,7 @@ Allocator* XlaDevice::GetAllocator(AllocatorAttributes attr) { xla::Backend* backend = client()->mutable_backend(); xla_allocator_ = XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( backend, device_ordinal_); + tensor_info_manager_.reset(new XlaTensorInfoManager(xla_allocator_)); } return xla_allocator_; } @@ -225,7 +238,11 @@ Status XlaDevice::FillContextMap(const Graph* graph, VLOG(1) << "XlaDevice::FillContextMap"; device_context_map->resize(graph->num_node_ids()); TF_ASSIGN_OR_RETURN(se::Stream * stream, GetStream()); - auto ctx = new XlaDeviceContext(stream); + // Call GetAllocator for the side-effect of ensuring the allocator and + // XlaTensorInfoManager is created. + (void)GetAllocator({}); + auto ctx = new XlaDeviceContext(stream, tensor_info_manager_.get(), + transfer_as_literal_); for (Node* n : graph->nodes()) { VLOG(2) << n->id() << " : " << n->type_string() << " : " << n->name(); ctx->Ref(); @@ -273,7 +290,8 @@ Status XlaDevice::MakeTensorFromProto(const TensorProto& tensor_proto, Tensor copy(GetAllocator(alloc_attrs), parsed.dtype(), parsed.shape()); Notification n; TF_ASSIGN_OR_RETURN(se::Stream * stream, GetStream()); - XlaTransferManager manager(stream); + XlaTransferManager manager(stream, tensor_info_manager_.get(), + transfer_as_literal_); manager.CopyCPUTensorToDevice(&parsed, this, ©, [&n, &status](const Status& s) { status = s; diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index d2ec38293c429f04f088bf3726ba97eb4e4b0dba..0f4476296b2560d33bc3916cd763ca29861f4ada 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -26,6 +26,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_DEVICE_H_ #define TENSORFLOW_COMPILER_JIT_XLA_DEVICE_H_ +#include "tensorflow/compiler/jit/xla_tensor_info.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/local_device.h" @@ -48,7 +49,8 @@ class XlaDevice : public LocalDevice { class Metadata { public: Metadata(int device_ordinal, perftools::gputools::Platform* platform, - const DeviceType& device_type); + const DeviceType& device_type, + std::unique_ptr* tensor_info_manager); // The index of the device on this host. int device_ordinal() const; @@ -56,11 +58,13 @@ class XlaDevice : public LocalDevice { perftools::gputools::Platform* platform() const; xla::LocalClient* client() const; const DeviceType& jit_device_type() const; + XlaTensorInfoManager& tensor_info_manager() const; private: const int device_ordinal_; const DeviceType device_type_; perftools::gputools::Platform* platform_; // Not owned. + std::unique_ptr& tensor_info_manager_; TF_DISALLOW_COPY_AND_ASSIGN(Metadata); }; @@ -71,15 +75,20 @@ class XlaDevice : public LocalDevice { // Factory function. 'platform_name' is the name of the XLA platform. // 'device_name' is the name of the Tensorflow device to create. // 'jit_device_name' is the name of the corresponding JIT device. + // 'transfer_as_literal' is true if device<->host transfers must be done using + // XLA's TransferLiteral{To,From}Device interface. If false, we can use + // ThenMemcpy instead. static Status Create(const string& platform_name, const string& device_name, int device_ordinal, const string& jit_device_name, const SessionOptions& options, const string& name_prefix, bool register_device_for_compilation, + bool transfer_as_literal, std::unique_ptr* device); XlaDevice(const SessionOptions& options, const DeviceAttributes& attrs, int device_ordinal, const DeviceType& jit_device_name, - ::perftools::gputools::Platform* platform); + ::perftools::gputools::Platform* platform, + bool transfer_as_literal); ~XlaDevice() override; Allocator* GetAllocator(AllocatorAttributes attr) override; @@ -113,6 +122,16 @@ class XlaDevice : public LocalDevice { // copying back and forth between CPU and the device, and // computations enqueued by XLA. xla::Backend::StreamPtr stream_; + // Manages sideband data about tensors, in particular the on-device shape tree + // if the tensor requires multiple device buffers to represent (for example, + // tuple shapes). + // This is a unique_ptr because XlaTensorInfoManager is non-copy-constructible + // and we need to initialize this lazily (as we also lazily initialize the + // underlying allocator). + std::unique_ptr tensor_info_manager_; + // Must we use XLA's transfer manager for correct host<->device transfers? if + // false, we can use ThenMemcpy() instead. + bool transfer_as_literal_; }; // Builds dummy OpKernel registrations on 'device' for the JIT operators diff --git a/tensorflow/compiler/jit/xla_device_context.cc b/tensorflow/compiler/jit/xla_device_context.cc index c936222f32056e92efced82d5adb3a96c8041a17..b57f82f98e730faf8f40af3618e0afe80ffc1ac7 100644 --- a/tensorflow/compiler/jit/xla_device_context.cc +++ b/tensorflow/compiler/jit/xla_device_context.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_device_context.h" +#include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/xla/util.h" @@ -52,7 +53,12 @@ void XlaDeviceAllocator::DeallocateRaw(void* ptr) { void XlaDeviceAllocator::GetStats(AllocatorStats* stats) { stats->Clear(); } -XlaTransferManager::XlaTransferManager(se::Stream* stream) : stream_(stream) {} +XlaTransferManager::XlaTransferManager( + se::Stream* stream, XlaTensorInfoManager* tensor_info_manager, + bool transfer_as_literal) + : stream_(stream), + tensor_info_manager_(tensor_info_manager), + transfer_as_literal_(transfer_as_literal) {} void XlaTransferManager::CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device, @@ -72,13 +78,19 @@ void XlaTransferManager::CopyCPUTensorToDevice(const Tensor* cpu_tensor, se::DeviceMemoryBase dev_dst_ptr(dst_ptr, total_bytes); Status status; - stream_->ThenMemcpy(&dev_dst_ptr, src_ptr, total_bytes); - // TODO(hpucha): Make this asynchronous. - Status block_status = stream_->BlockHostUntilDone(); - if (!block_status.ok()) { - status = xla::InternalError( - "Failed to complete data transfer on stream %p: %s", stream_, - block_status.error_message().c_str()); + if (transfer_as_literal_) { + status = xla::Unimplemented( + "XlaTransferManager::CopyCPUTensorToDevice not implemented for " + "literals"); + } else { + stream_->ThenMemcpy(&dev_dst_ptr, src_ptr, total_bytes); + // TODO(hpucha): Make this asynchronous. + Status block_status = stream_->BlockHostUntilDone(); + if (!block_status.ok()) { + status = xla::InternalError( + "Failed to complete data transfer on stream %p: %s", stream_, + block_status.error_message().c_str()); + } } done(status); @@ -108,13 +120,19 @@ void XlaTransferManager::CopyDeviceTensorToCPU(const Tensor* device_tensor, void* dst_ptr = DMAHelper::base(cpu_tensor); Status status; - stream_->ThenMemcpy(dst_ptr, dev_src_ptr, total_bytes); - // TODO(hpucha): Make this asynchronous. - Status block_status = stream_->BlockHostUntilDone(); - if (!block_status.ok()) { - status = xla::InternalError( - "Failed to complete data transfer on stream %p: %s", stream_, - block_status.error_message().c_str()); + if (transfer_as_literal_) { + status = xla::Unimplemented( + "XlaTransferManager::CopyDeviceTensorToCPU not implemented for " + "literals"); + } else { + stream_->ThenMemcpy(dst_ptr, dev_src_ptr, total_bytes); + // TODO(hpucha): Make this asynchronous. + Status block_status = stream_->BlockHostUntilDone(); + if (!block_status.ok()) { + status = xla::InternalError( + "Failed to complete data transfer on stream %p: %s", stream_, + block_status.error_message().c_str()); + } } done(status); @@ -125,7 +143,10 @@ void XlaTransferManager::CopyDeviceTensorToCPU(const Tensor* device_tensor, done(Status::OK()); } -XlaDeviceContext::XlaDeviceContext(se::Stream* stream) : manager_(stream) {} +XlaDeviceContext::XlaDeviceContext(se::Stream* stream, + XlaTensorInfoManager* tensor_info_manager, + bool transfer_as_literal) + : manager_(stream, tensor_info_manager, transfer_as_literal) {} void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device, diff --git a/tensorflow/compiler/jit/xla_device_context.h b/tensorflow/compiler/jit/xla_device_context.h index c4edcd474e48f791af9340c3cd6e4d031407bb68..df02f4eac482f385f8864476d11c5430971f00c8 100644 --- a/tensorflow/compiler/jit/xla_device_context.h +++ b/tensorflow/compiler/jit/xla_device_context.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "tensorflow/compiler/jit/xla_tensor_info.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/core/framework/allocator.h" @@ -49,7 +50,9 @@ class XlaDeviceAllocator : public Allocator { // Helper class for managing data transfers between host and XLA devices. class XlaTransferManager { public: - explicit XlaTransferManager(perftools::gputools::Stream* stream); + explicit XlaTransferManager(perftools::gputools::Stream* stream, + XlaTensorInfoManager* tensor_info_manager, + bool transfer_as_literal); void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device, Tensor* device_tensor, StatusCallback done) const; @@ -62,6 +65,10 @@ class XlaTransferManager { // Stream obtained from a Device, used to transfer tensors between // CPU and device. perftools::gputools::Stream* stream_; + // The tensor info manager, for access to sideband information about tensors. + XlaTensorInfoManager* tensor_info_manager_; + // True if we must use XLA's TransferManager for correct device transfers. + bool transfer_as_literal_; }; // DeviceContext for operators assigned to XlaDevice devices. The @@ -69,7 +76,9 @@ class XlaTransferManager { // wraps the methods in XlaTransferManager. class XlaDeviceContext : public DeviceContext { public: - explicit XlaDeviceContext(perftools::gputools::Stream* stream); + explicit XlaDeviceContext(perftools::gputools::Stream* stream, + XlaTensorInfoManager* tensor_info_manager, + bool transfer_as_literal); void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device, Tensor* device_tensor, diff --git a/tensorflow/compiler/jit/xla_gpu_device.cc b/tensorflow/compiler/jit/xla_gpu_device.cc index 2326070358d67c0cf30ef17fab5c93862cd8932c..383ed879ef6177c3de9e27d18f6d57e8f64c69eb 100644 --- a/tensorflow/compiler/jit/xla_gpu_device.cc +++ b/tensorflow/compiler/jit/xla_gpu_device.cc @@ -39,9 +39,10 @@ Status XlaGpuDeviceFactory::CreateDevices(const SessionOptions& options, (void)registrations; std::unique_ptr device; - Status status = XlaDevice::Create( - "CUDA", DEVICE_XLA_GPU, 0, DEVICE_GPU_XLA_JIT, options, name_prefix, - /*register_device_for_compilation=*/true, &device); + Status status = XlaDevice::Create("CUDA", DEVICE_XLA_GPU, 0, + DEVICE_GPU_XLA_JIT, options, name_prefix, + /*register_device_for_compilation=*/true, + /*transfer_as_literal=*/false, &device); if (!status.ok()) { // Treat failures as non-fatal; there might not be a GPU in the machine. VLOG(1) << "Failed to create XLA_GPU device: " << status; diff --git a/tensorflow/compiler/jit/xla_interpreter_device.cc b/tensorflow/compiler/jit/xla_interpreter_device.cc index 2614deefd8823dcb8f38e9e22ae4e78145d0d96a..a329451b14a785b17913e3838a6571b62b422804 100644 --- a/tensorflow/compiler/jit/xla_interpreter_device.cc +++ b/tensorflow/compiler/jit/xla_interpreter_device.cc @@ -25,8 +25,8 @@ namespace tensorflow { const char* const DEVICE_XLA_INTERPRETER = "XLA_INTERPRETER"; const char* const DEVICE_INTERPRETER_XLA_JIT = "XLA_INTERPRETER_JIT"; -constexpr std::array kExecAllTypes = { - {DT_INT32, DT_FLOAT, DT_BOOL, DT_DOUBLE, DT_INT64}}; +constexpr std::array kExecAllTypes = { + {DT_INT32, DT_INT64, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_BOOL}}; class XlaInterpreterDeviceFactory : public DeviceFactory { public: diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..689fa3299ca46f7524c83bc5dc3109bacab83889 --- /dev/null +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -0,0 +1,268 @@ +/* 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/jit/xla_launch_util.h" + +#include "tensorflow/compiler/jit/defs.h" +#include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/compiler/xla/client/client_library.h" +#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/common_runtime/dma_helper.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/node_def_util.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/util/stream_executor_util.h" + +namespace gpu = perftools::gputools; + +namespace tensorflow { + +std::map SnapshotResourceVariables(OpKernelContext* ctx, + int num_variables) { + std::map snapshot; + int first_variable = ctx->num_inputs() - num_variables; + for (int i = 0; i < num_variables; ++i) { + Var* variable = nullptr; + ResourceHandle handle = HandleFromInput(ctx, first_variable + i); + OptionalTensor& tensor = snapshot[first_variable + i]; + if (LookupResource(ctx, handle, &variable).ok()) { + tf_shared_lock lock(*variable->mu()); + tensor.name = handle.name(); + tensor.present = true; + tensor.value = *variable->tensor(); + } + } + return snapshot; +} + +XlaAllocator::XlaAllocator(const gpu::Platform* platform, + OpKernelContext* op_context) + : xla::DeviceMemoryAllocator(platform), op_context_(op_context) {} + +XlaAllocator::~XlaAllocator() { CHECK(allocated_.empty()); } + +xla::StatusOr XlaAllocator::Allocate( + int device_ordinal, uint64 size, bool retry_on_failure) { + void* data = op_context_->device()->GetAllocator({})->AllocateRaw( + Allocator::kAllocatorAlignment, size); + allocated_.insert(data); + return gpu::DeviceMemoryBase(data, size); +} + +void XlaAllocator::Release(void* ptr) { allocated_.erase(ptr); } + +Status XlaAllocator::Deallocate(int device_ordinal, + gpu::DeviceMemoryBase* mem) { + if (allocated_.count(mem->opaque())) { + op_context_->device()->GetAllocator({})->DeallocateRaw(mem->opaque()); + allocated_.erase(mem->opaque()); + } + return Status::OK(); +} + +namespace { +// Return the 'index''th subtree of the given ShapedBuffer as a ShapedBuffer. +xla::ShapedBuffer ExtractSubShapedBuffer(const xla::ShapedBuffer& shaped_buffer, + int index) { + xla::Shape on_host_shape = xla::ShapeUtil::GetTupleElementShape( + shaped_buffer.on_host_shape(), index); + xla::Shape on_device_shape = xla::ShapeUtil::GetTupleElementShape( + shaped_buffer.on_device_shape(), index); + + xla::ShapedBuffer sub_shaped_buffer(on_host_shape, on_device_shape, + shaped_buffer.platform(), + shaped_buffer.device_ordinal()); + + auto& shape_tree = shaped_buffer.buffers(); + auto& sub_shape_tree = sub_shaped_buffer.buffers(); + sub_shape_tree.CopySubtreeFrom(shape_tree, + /*source_base_index=*/{index}, + /*target_base_index=*/{}); + return sub_shaped_buffer; +} +} // namespace + +XlaComputationLaunchContext::XlaComputationLaunchContext( + int64 num_resource_args, xla::LocalClient* client, + XlaAllocator* xla_allocator, XlaTensorInfoManager* tensor_info_manager) + : num_resource_args_(num_resource_args), + client_(client), + xla_allocator_(xla_allocator), + tensor_info_manager_(tensor_info_manager) {} + +void XlaComputationLaunchContext::PopulateInputs( + OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, + const std::map& variables) { + // Build xla::ShapedBuffers that point directly to the Tensor buffers. + arg_buffers_.reserve(kernel->xla_input_shapes.size() + 1); + arg_buffers_.resize(kernel->xla_input_shapes.size()); + arg_ptrs_ = std::vector(arg_buffers_.size()); + + // Pass remaining parameters. + const Tensor* t; + for (int i = 0; i < kernel->xla_input_shapes.size(); ++i) { + int arg_num = kernel->input_mapping[i]; + const xla::Shape& shape = kernel->xla_input_shapes[i]; + if (variables.count(arg_num)) { + t = &(variables.at(arg_num).value); + CHECK(t); + } else { + t = &(ctx->input(arg_num)); + } + + const xla::Shape on_device_shape = + client_->backend().transfer_manager()->HostShapeToDeviceShape(shape); + if (xla::ShapeUtil::IsTuple(on_device_shape)) { + CHECK(tensor_info_manager_); + const XlaTensorInfo* tensor_info = + tensor_info_manager_->GetTensorInfo(*t); + CHECK(tensor_info && tensor_info->has_shaped_buffer()); + arg_ptrs_[i] = + const_cast(&tensor_info->shaped_buffer()); + } else { + CHECK(xla::ShapeUtil::Equal(shape, on_device_shape)) + << "On-device shape " + << xla::ShapeUtil::HumanStringWithLayout(on_device_shape) + << " not the same as on-host shape " + << xla::ShapeUtil::HumanStringWithLayout(shape); + gpu::DeviceMemoryBase dmem = gpu::DeviceMemoryBase( + const_cast(t->tensor_data().data()), t->tensor_data().size()); + arg_buffers_[i] = xla::MakeUnique( + /*on_host_shape=*/shape, /*on_device_shape=*/shape, + client_->platform(), client_->default_device_ordinal()); + arg_buffers_[i]->set_buffer(dmem, /*index=*/{}); + arg_ptrs_[i] = arg_buffers_[i].get(); + } + } +} + +void XlaComputationLaunchContext::PopulateOutputs( + OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, + std::unique_ptr output) { + gpu::Stream* stream = + ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; + + // Computation output should always be a tuple. + if (VLOG_IS_ON(2)) { + VLOG(2) << "Result tuple shape: " << output->on_host_shape().DebugString(); + } + CHECK_EQ(ctx->num_outputs(), kernel->outputs.size()); + + // Copy XLA results to the OpOutputList. + int output_num = 0; + for (int i = 0; i < ctx->num_outputs(); ++i) { + AllocatorAttributes alloc_attrs = ctx->output_alloc_attr(i); + Allocator* allocator = ctx->device()->GetAllocator(alloc_attrs); + if (tensor_info_manager_ && !alloc_attrs.on_host()) { + allocator = tensor_info_manager_; + } + if (kernel->outputs[i].is_constant) { + // Output is a constant. + const Tensor& const_tensor = kernel->outputs[i].constant_value; + const size_t total_bytes = const_tensor.TotalBytes(); + if (stream && total_bytes > 0) { + // Copy host -> device. (Empty tensors don't have backing buffers.) + VLOG(1) << "Constant output tensor on device"; + Tensor* output_tensor; + TF_CHECK_OK( + ctx->allocate_output(i, const_tensor.shape(), &output_tensor)); + + const void* src_ptr = DMAHelper::base(&const_tensor); + void* dst_ptr = DMAHelper::base(output_tensor); + gpu::DeviceMemoryBase gpu_dst_ptr(dst_ptr, total_bytes); + stream->ThenMemcpy(&gpu_dst_ptr, src_ptr, total_bytes); + } else { + // No copy required. + ctx->set_output(i, const_tensor); + } + } else { + const TensorShape& shape = kernel->outputs[i].shape; + VLOG(2) << "Retval " << i << " shape " << shape.DebugString(); + + gpu::DeviceMemoryBase buffer = output->buffer({output_num}); + Tensor output_tensor = XlaTensorBuffer::MakeTensor( + ctx->expected_output_dtype(i), shape, buffer, allocator); + xla_allocator_->Release(buffer.opaque()); + + xla::Shape output_shape = xla::ShapeUtil::GetTupleElementShape( + output->on_device_shape(), output_num); + if (xla::ShapeUtil::IsTuple(output_shape)) { + CHECK(tensor_info_manager_); + XlaTensorInfo* tensor_info = + tensor_info_manager_->GetOrCreateTensorInfo(output_tensor); + tensor_info->set_shaped_buffer( + ExtractSubShapedBuffer(*output, output_num)); + } + ctx->set_output(i, output_tensor); + ++output_num; + } + + if (VLOG_IS_ON(3)) { + VLOG(3) << ctx->mutable_output(i)->DebugString(); + } + } + + // Apply variable updates, if any. + VLOG(2) << "Applying variable updates"; + for (int i = 0; i < kernel->resource_updates.size(); ++i) { + Allocator* allocator = ctx->device()->GetAllocator({}); + if (tensor_info_manager_) { + allocator = tensor_info_manager_; + } + const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i]; + OP_REQUIRES(ctx, + write.input_index >= 0 && write.input_index < ctx->num_inputs(), + errors::Internal("Invalid input index for variable write.")); + + gpu::DeviceMemoryBase buffer = output->buffer({output_num}); + + Var* variable = nullptr; + // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor, + // not a Tensor. + OP_REQUIRES_OK(ctx, LookupOrCreateResource( + ctx, HandleFromInput(ctx, write.input_index), + &variable, [this, ctx, &write](Var** ptr) { + *ptr = new Var(write.type); + return Status::OK(); + })); + + core::ScopedUnref s(variable); + + mutex_lock ml(*variable->mu()); + OP_REQUIRES(ctx, variable->tensor()->dtype() == write.type, + errors::Internal("Mismatched type in variable write")); + *variable->tensor() = + XlaTensorBuffer::MakeTensor(write.type, write.shape, buffer, allocator); + xla_allocator_->Release(buffer.opaque()); + + xla::Shape output_shape = xla::ShapeUtil::GetTupleElementShape( + output->on_device_shape(), output_num); + if (xla::ShapeUtil::IsTuple(output_shape)) { + CHECK(tensor_info_manager_); + XlaTensorInfo* tensor_info = + tensor_info_manager_->GetOrCreateTensorInfo(*variable->tensor()); + tensor_info->set_shaped_buffer( + ExtractSubShapedBuffer(*output, output_num)); + } + ++output_num; + } +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h new file mode 100644 index 0000000000000000000000000000000000000000..8694f6ce58b72ca188bf831528db30daf93b905d --- /dev/null +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -0,0 +1,149 @@ +/* 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. +==============================================================================*/ + +// Contains utilities for launching compiled XLA kernels for a KernelContext. + +#ifndef TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ +#define TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ + +#include "tensorflow/compiler/jit/xla_compilation_cache.h" +#include "tensorflow/compiler/jit/xla_tensor_info.h" +#include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/core/framework/allocation_description.pb.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/kernels/variable_ops.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { +class XlaAllocator; + +// Takes a snapshot of the values of resource variable arguments, which are +// the last `num_variables` arguments. We snapshot tensors that back +// resource variables since concurrent updates may modify the shape, and it is +// important that the shapes used for compilation match the true shapes of the +// buffers. +// +// Returns a map of TensorFlow argument index to resource variable. +std::map SnapshotResourceVariables(OpKernelContext* ctx, + int num_variables); + +// Adapter class that wraps a Tensorflow allocator as an XLA allocator. +// Assumes that the Tensorflow allocator permits asynchronous deallocation: +// see comment on `AllowsAsynchronousDeallocation()`. +class XlaAllocator : public xla::DeviceMemoryAllocator { + public: + XlaAllocator(const perftools::gputools::Platform* platform, + OpKernelContext* op_context); + ~XlaAllocator() override; + xla::StatusOr Allocate( + int device_ordinal, uint64 size, bool retry_on_failure) override; + Status Deallocate(int device_ordinal, + perftools::gputools::DeviceMemoryBase* mem) override; + + // Un-track 'ptr' - do not delete it on destruction. + void Release(void* ptr); + + // The Tensorflow BFC allocator used on GPU allows host-side deallocation + // before GPU execution takes place. Tensorflow uses the ordering of the main + // compute stream to enforce a happens-before relationship between a memory + // allocation and code that reuses the same memory. If Tensorflow adds + // support for multiple GPU streams or allocators with different ordering + // requirements, this code may need to change. + // (This attribute has no effect on CPU.) + bool AllowsAsynchronousDeallocation() const override { return true; } + + private: + OpKernelContext* const op_context_; + std::unordered_set allocated_; +}; + +// Helper class to perform the marshalling of TensorFlow inputs and outputs to +// ShapedBuffers suitable for passing to an XLA computation. +class XlaComputationLaunchContext { + public: + XlaComputationLaunchContext(int64 num_resource_args, xla::LocalClient* client, + XlaAllocator* xla_allocator, + XlaTensorInfoManager* tensor_info_manager); + + // Add all inputs within `ctx` as XLA arguments (returned by arguments()). + // `variables` is a map from TensorFlow argument number to resource variable. + void PopulateInputs(OpKernelContext* ctx, + const XlaCompiler::CompilationResult* kernel, + const std::map& variables); + + // Given the XLA output in `output`, populate all outputs of `ctx`. + void PopulateOutputs(OpKernelContext* ctx, + const XlaCompiler::CompilationResult* kernel, + std::unique_ptr output); + + // Return the argument list. Only valid after PopulateInputs() has been + // called. + const std::vector& arguments() const { return arg_ptrs_; } + + private: + int64 num_resource_args_; + xla::LocalClient* client_; + XlaAllocator* xla_allocator_; + XlaTensorInfoManager* tensor_info_manager_; + std::vector> arg_buffers_; + std::vector arg_ptrs_; +}; + +// A simple TensorBuffer implementation that allows us to create Tensors that +// take ownership of pre-allocated memory. +class XlaTensorBuffer : public TensorBuffer { + public: + XlaTensorBuffer(const void* ptr, size_t expected_size, size_t actual_size, + Allocator* allocator) + : expected_size_(expected_size), + actual_size_(actual_size), + allocator_(allocator) { + data_ = const_cast(ptr); + } + + ~XlaTensorBuffer() override { allocator_->DeallocateRaw(data_); } + + void* data() const override { return data_; } + size_t size() const override { return expected_size_; } + + TensorBuffer* root_buffer() override { return this; } + + void FillAllocationDescription(AllocationDescription* proto) const override { + proto->set_allocated_bytes(actual_size_); + } + + static Tensor MakeTensor(DataType dtype, const TensorShape& shape, + perftools::gputools::DeviceMemoryBase buffer, + Allocator* allocator) { + size_t expected_size = shape.num_elements() * DataTypeSize(dtype); + auto* tensor_buffer = new XlaTensorBuffer(buffer.opaque(), expected_size, + buffer.size(), allocator); + Tensor t(dtype, shape, tensor_buffer); + tensor_buffer->Unref(); + return t; + } + + private: + void* data_; + size_t expected_size_; + size_t actual_size_; + Allocator* allocator_; +}; + +} // namespace tensorflow + +#endif diff --git a/tensorflow/compiler/jit/xla_tensor_info.cc b/tensorflow/compiler/jit/xla_tensor_info.cc new file mode 100644 index 0000000000000000000000000000000000000000..0ce18c27cbe1d46eb61f8000506396fedc509e9c --- /dev/null +++ b/tensorflow/compiler/jit/xla_tensor_info.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/compiler/jit/xla_tensor_info.h" + +namespace tensorflow { + +const XlaTensorInfo* XlaTensorInfoManager::GetTensorInfo( + const void* device_ptr) const { + mutex_lock lock(lock_); + auto iterator = tensor_infos_.find(device_ptr); + return (iterator == tensor_infos_.end()) ? nullptr + : tensor_infos_.at(device_ptr).get(); +} + +XlaTensorInfo* XlaTensorInfoManager::GetOrCreateTensorInfo( + const void* device_ptr) { + mutex_lock lock(lock_); + auto iterator = tensor_infos_.find(device_ptr); + if (iterator != tensor_infos_.end()) { + return iterator->second.get(); + } + auto iterator_and_inserted = + tensor_infos_.emplace(device_ptr, MakeUnique()); + CHECK(iterator_and_inserted.second); + return iterator_and_inserted.first->second.get(); +} + +const XlaTensorInfo* XlaTensorInfoManager::GetTensorInfo(const Tensor& tensor) { + return GetTensorInfo(tensor.tensor_data().data()); +} + +XlaTensorInfo* XlaTensorInfoManager::GetOrCreateTensorInfo( + const Tensor& tensor) { + return GetOrCreateTensorInfo(tensor.tensor_data().data()); +} + +void XlaTensorInfoManager::DeallocateRaw(void* ptr) { + wrapped()->DeallocateRaw(ptr); + mutex_lock lock(lock_); + tensor_infos_.erase(ptr); +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_tensor_info.h b/tensorflow/compiler/jit/xla_tensor_info.h new file mode 100644 index 0000000000000000000000000000000000000000..0b0736bf0199f7693bd5510c9d4b87c24db0433f --- /dev/null +++ b/tensorflow/compiler/jit/xla_tensor_info.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_JIT_XLA_TENSOR_INFO_H_ +#define TENSORFLOW_COMPILER_JIT_XLA_TENSOR_INFO_H_ + +#include "tensorflow/compiler/xla/service/shaped_buffer.h" +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/device_base.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/mutex.h" + +namespace tensorflow { + +// Information about a tensor. The XlaTensorInfoManager can maintain one of +// these per device Tensor. +class XlaTensorInfo { + public: + XlaTensorInfo() {} + + // Some Tensors can have complex on-device shapes, including tuple shapes. To + // manage the memory for these tensors a ShapedBuffer may be required. + + // Return true if this TensorInfo contains a ShapedBuffer. + bool has_shaped_buffer() const { return shaped_buffer_ != nullptr; } + // Return the contained ShapedBuffer. + // REQUIRES: has_shaped_buffer() + const xla::ShapedBuffer& shaped_buffer() const { return *shaped_buffer_; } + // Mutates the TensorInfo to set the ShapedBuffer. + void set_shaped_buffer(xla::ShapedBuffer shaped_buffer) { + shaped_buffer_.reset(new xla::ShapedBuffer(std::move(shaped_buffer))); + } + + private: + // The optional contained ShapedBuffer. + std::unique_ptr shaped_buffer_; +}; + +// Manages XlaTensorInfo objects. This class is also an Allocator, so that +// XlaTensorInfo objects can be deleted when their Tensor is deallocated. +class XlaTensorInfoManager : public AllocatorWrapper { + public: + // Creates a new XlaTensorInfoManager, delegating all DeallocateRaw calls to + // allocator. + XlaTensorInfoManager(Allocator* allocator) : AllocatorWrapper(allocator) {} + + // Returns the XlaTensorInfo for the given device memory pointer or nullptr if + // none exists. + const XlaTensorInfo* GetTensorInfo(const void* device_ptr) const; + // Returns the XlaTensorInfo for the device memory pointer extracted from + // tensor or nullptr if none exists. + const XlaTensorInfo* GetTensorInfo(const Tensor& tensor); + + // Returns the XlaTensorInfo for the given device memory pointer, creating one + // if necessary. + XlaTensorInfo* GetOrCreateTensorInfo(const Tensor& tensor); + // Returns the XlaTensorInfo for the device memory pointer extracted from + // tensor, creating one if necessary. + XlaTensorInfo* GetOrCreateTensorInfo(const void* device_ptr); + + // Allocator interface + void DeallocateRaw(void* ptr) override; + + private: + mutable mutex lock_; + // The managed tensor infos. The mapped value is a unique_ptr so that returned + // references are stable over rehashes. + std::unordered_map> tensor_infos_ + GUARDED_BY(lock_); +}; +} // namespace tensorflow + +#endif diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 9e64f3e9a35c875036e03bc4a51fe02d81ef3ffb..85a2adab283c273af607b6b80fb4fd76f8dac2b2 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -98,7 +98,7 @@ tf_xla_py_test( tf_xla_py_test( name = "binary_ops_test", - size = "small", + size = "medium", srcs = ["binary_ops_test.py"], shard_count = 5, tags = [ @@ -255,6 +255,18 @@ 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", @@ -341,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", @@ -454,8 +479,9 @@ tf_xla_py_test( tf_xla_py_test( name = "reverse_sequence_op_test", - size = "small", + size = "medium", srcs = ["reverse_sequence_op_test.py"], + tags = ["optonly"], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -511,6 +537,7 @@ tf_xla_py_test( size = "medium", srcs = ["spacetobatch_op_test.py"], shard_count = 3, + tags = ["notsan"], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -614,6 +641,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", @@ -640,6 +668,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", @@ -764,6 +817,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 c95fb1c515242ca38369b11aa5e616b13624edf9..ba7b9bacd2b794c74409d517a9c05bfbb14a845f 100644 --- a/tensorflow/compiler/tests/binary_ops_test.py +++ b/tensorflow/compiler/tests/binary_ops_test.py @@ -71,7 +71,7 @@ class BinaryOpsTest(XLATestCase): expected=np.array([[[[False, True], [True, False]]]], dtype=dtype)) self._testBinary( - gen_math_ops._real_div, + gen_math_ops.real_div, np.array([3, 3, -1.5, -8, 44], dtype=dtype), np.array([2, -2, 7, -4, 0], dtype=dtype), expected=np.array( @@ -108,57 +108,57 @@ class BinaryOpsTest(XLATestCase): [0, np.pi / 4, np.pi / 2, np.pi * 3 / 4, np.pi], dtype=dtype)) self._testBinary( - gen_math_ops._reciprocal_grad, + gen_math_ops.reciprocal_grad, np.array([4, -3, -2, 1], dtype=dtype), np.array([5, -6, 7, -8], dtype=dtype), expected=np.array([-80, 54, -28, 8], dtype=dtype)) self._testBinary( - gen_math_ops._sigmoid_grad, + gen_math_ops.sigmoid_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array([-60, -36, -14, 0], dtype=dtype)) self._testBinary( - gen_math_ops._rsqrt_grad, + gen_math_ops.rsqrt_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array([-160, -81, -28, -4], dtype=dtype)) self._testBinary( - gen_math_ops._sqrt_grad, + gen_math_ops.sqrt_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array([0.625, 1, 1.75, 4], dtype=dtype)) self._testBinary( - gen_nn_ops._softplus_grad, + gen_nn_ops.softplus_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array( [3.97322869, 2.99258232, 1.99817801, 0.99966466], dtype=dtype)) self._testBinary( - gen_nn_ops._softsign_grad, + gen_nn_ops.softsign_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array( [0.11111111, 0.06122449, 0.03125, 0.01234568], dtype=dtype)) self._testBinary( - gen_math_ops._tanh_grad, + gen_math_ops.tanh_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array([-75, -48, -21, 0], dtype=dtype)) self._testBinary( - gen_nn_ops._elu_grad, + gen_nn_ops.elu_grad, np.array([1, 2, 3, 4, 5, 6], dtype=dtype), np.array([-.6, -.4, -.2, 0, .2, .4], dtype=dtype), expected=np.array([0.4, 1.2, 2.4, 4, 5, 6], dtype=dtype)) self._testBinary( - gen_nn_ops._selu_grad, + gen_nn_ops.selu_grad, np.array([1, 2, 3, 4, 5, 6], dtype=dtype), np.array([-.6, -.4, -.2, .2, .4, .6], dtype=dtype), expected=np.array( @@ -166,20 +166,20 @@ class BinaryOpsTest(XLATestCase): 4.202803949422, 5.2535049367774, 6.30420592413], dtype=dtype)) self._testBinary( - gen_nn_ops._relu_grad, + gen_nn_ops.relu_grad, np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype), np.array([0, 0, 0, 0, 0, 0.1, 0.3, 0.5, 0.7, 0.9], dtype=dtype), expected=np.array([0, 0, 0, 0, 0, 6, 7, 8, 9, 10], dtype=dtype)) self._testBinary( - gen_nn_ops._relu6_grad, + gen_nn_ops.relu6_grad, np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=dtype), np.array( [0, 0, 0, 0, 0, 0.1, 0.3, 0.5, 0.7, 0.9, 6.1, 10.0], dtype=dtype), expected=np.array([0, 0, 0, 0, 0, 6, 7, 8, 9, 10, 0, 0], dtype=dtype)) self._testBinary( - gen_nn_ops._softmax_cross_entropy_with_logits, + gen_nn_ops.softmax_cross_entropy_with_logits, np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=dtype), np.array([[0.1, 0.2, 0.3, 0.4], [0.4, 0.3, 0.2, 0.1]], dtype=dtype), expected=[ @@ -191,7 +191,7 @@ class BinaryOpsTest(XLATestCase): equality_test=self.ListsAreClose) self._testBinary( - gen_nn_ops._sparse_softmax_cross_entropy_with_logits, + gen_nn_ops.sparse_softmax_cross_entropy_with_logits, np.array([[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]], dtype=dtype), np.array([2, 1, 7], dtype=np.int32), @@ -207,7 +207,7 @@ class BinaryOpsTest(XLATestCase): def testIntOps(self): for dtype in self.int_types: self._testBinary( - gen_math_ops._truncate_div, + gen_math_ops.truncate_div, np.array([3, 3, -1, -9, -8], dtype=dtype), np.array([2, -2, 7, 2, -4], dtype=dtype), expected=np.array([1, -1, 0, -4, 2], dtype=dtype)) @@ -232,11 +232,16 @@ class BinaryOpsTest(XLATestCase): expected=np.right_shift(lhs, rhs)) if dtype in [np.int8, np.int16, np.int32, np.int64]: - lhs = np.array([-1, -5, -3, -14], dtype=dtype) - rhs = np.array([5, 0, 1, 11], dtype=dtype) - self._testBinary( - bitwise_ops.right_shift, lhs, rhs, - expected=np.right_shift(lhs, rhs)) + lhs = np.array([-1, -5, -3, -14, -2], dtype=dtype) + rhs = np.array([5, 0, 1, 11, 36], dtype=dtype) + # HLO has saturating shift behavior. + bits = np.ceil( + np.log(np.iinfo(dtype).max - np.iinfo(dtype).min) / np.log(2)) + expected = [ + np.right_shift(l, r) if r < bits else np.sign(l) + for l, r in zip(lhs, rhs) + ] + self._testBinary(bitwise_ops.right_shift, lhs, rhs, expected=expected) def testNumericOps(self): for dtype in self.numeric_types: @@ -255,12 +260,18 @@ class BinaryOpsTest(XLATestCase): np.array([[1], [2]], dtype=dtype), dtype(7), expected=np.array([[8], [9]], dtype=dtype)) + self._testBinary( + math_ops.add, + np.array([0xffffffff, 0xfffffffff, 1, 1], dtype=np.int64), + np.array([1, 1, 0xffffffff, 0xfffffffff], dtype=np.int64), + expected=np.array( + [1 << 32, 1 << 36, 1 << 32, 1 << 36], dtype=np.int64)) self._testBinary( math_ops.subtract, - np.array([1, 2], dtype=dtype), - np.array([10, 20], dtype=dtype), - expected=np.array([-9, -18], dtype=dtype)) + np.array([1, 2, 100], dtype=dtype), + np.array([10, 20, -1], dtype=dtype), + expected=np.array([-9, -18, 101], dtype=dtype)) self._testBinary( math_ops.subtract, dtype(5), @@ -369,7 +380,7 @@ class BinaryOpsTest(XLATestCase): expected=np.array([[[[False, True], [True, False]]]], dtype=dtype)) self._testBinary( - gen_math_ops._real_div, + gen_math_ops.real_div, np.array([3, 3j, -1.5j, -8, 2 + 3j, 2 + 4j], dtype=dtype), np.array([2, -2, 7j, -4j, 4 - 6j, 1 + 2j], dtype=dtype), expected=np.array( @@ -378,7 +389,7 @@ class BinaryOpsTest(XLATestCase): # Test inf/nan scenarios. self._testBinary( - gen_math_ops._real_div, + gen_math_ops.real_div, np.array([4 + 3j, 4, 3j, -4, -4j, 2 - 3j], dtype=dtype), np.array([0, 0, 0, 0, 0, 0], dtype=dtype), expected=np.array( @@ -418,19 +429,19 @@ class BinaryOpsTest(XLATestCase): lhs = np.array([4 + 2j, -3 - 1j, 2j, 1], dtype=dtype) rhs = np.array([5, -6j, 7 - 3j, -8j], dtype=dtype) self._testBinary( - gen_math_ops._reciprocal_grad, lhs, rhs, expected=-rhs * lhs * lhs) + gen_math_ops.reciprocal_grad, lhs, rhs, expected=-rhs * lhs * lhs) self._testBinary( - gen_math_ops._sigmoid_grad, lhs, rhs, expected=rhs * lhs * (1 - lhs)) + gen_math_ops.sigmoid_grad, lhs, rhs, expected=rhs * lhs * (1 - lhs)) self._testBinary( - gen_math_ops._rsqrt_grad, lhs, rhs, expected=lhs**3 * rhs / -2) + gen_math_ops.rsqrt_grad, lhs, rhs, expected=lhs**3 * rhs / -2) self._testBinary( - gen_math_ops._sqrt_grad, lhs, rhs, expected=rhs / (2 * lhs)) + gen_math_ops.sqrt_grad, lhs, rhs, expected=rhs / (2 * lhs)) self._testBinary( - gen_math_ops._tanh_grad, lhs, rhs, expected=rhs * (1 - lhs * lhs)) + gen_math_ops.tanh_grad, lhs, rhs, expected=rhs * (1 - lhs * lhs)) def testComplexMath(self): for dtype in self.complex_types: @@ -538,7 +549,7 @@ class BinaryOpsTest(XLATestCase): if dtype not in self.complex_types: # floordiv unsupported for complex. self._testBinary( - gen_math_ops._floor_div, + gen_math_ops.floor_div, np.array([3, 3, -1, -9, -8], dtype=dtype), np.array([2, -2, 7, 2, -4], dtype=dtype), expected=np.array([1, -2, -1, -5, 2], dtype=dtype)) @@ -554,12 +565,12 @@ class BinaryOpsTest(XLATestCase): def _testRemainder(self, dtype): """Test cases for remainder operators.""" self._testBinary( - gen_math_ops._floor_mod, + gen_math_ops.floor_mod, np.array([3, 3, -1, -8], dtype=dtype), np.array([2, -2, 7, -4], dtype=dtype), expected=np.array([1, -1, 6, 0], dtype=dtype)) self._testBinary( - gen_math_ops._truncate_mod, + gen_math_ops.truncate_mod, np.array([3, 3, -1, -8], dtype=dtype), np.array([2, -2, 7, -4], dtype=dtype), expected=np.array([1, 1, -1, 0], dtype=dtype)) @@ -668,6 +679,11 @@ class BinaryOpsTest(XLATestCase): np.array([[10], [7], [2]], dtype=np.float32), np.float32(7), expected=np.array([[False], [False], [True]], dtype=np.bool)) + self._testBinary( + less_op, + np.array([[10], [7], [2], [-1]], dtype=np.int64), + np.int64(7), + expected=np.array([[False], [False], [True], [True]], dtype=np.bool)) for less_equal_op in [math_ops.less_equal, (lambda x, y: x <= y)]: self._testBinary( @@ -686,6 +702,80 @@ class BinaryOpsTest(XLATestCase): np.float32(7), expected=np.array([[False], [True], [True]], dtype=np.bool)) + def testS64Comparisons(self): + for op in [(lambda x, y: x < y), (lambda x, y: x <= y), + (lambda x, y: x >= y), (lambda x, y: x > y)]: + lhs = np.array( + [ + np.int64(0x000000007FFFFFFF), + np.int64(0x000000007FFFFFFF), + np.int64(0x0000000080000000), + np.int64(0x0000000080000000), + np.int64(0x0000000080000001), + np.int64(0x00000000FFFF0000), + np.int64(0x00000000FFFF0000), + np.int64(0x00000000FFFFFFFE), + np.int64(0x00000000FFFFFFFF), + np.int64(0x00000000FFFFFFFF), + np.int64(0x0000000100000000), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(0x0000000200000002), + np.int64(-0x7FFFFFFF00000002), + np.int64(-0x7FFFFFFF00000002), + np.int64(-0x7FFFFFFF00000001), + np.int64(-0x7FFFFFFF00000001), + np.int64(-0x7FFFFFFF00000001), + np.int64(-0x7FFFFFFF00000001), + np.int64(0x7ffffffefff00010), + np.int64(0x7ffffffefff00010), + np.int64(-1), + np.int64(-1) + ], + dtype=np.int64) + rhs = np.array( + [ + np.int64(0x000000007FFFFFFE), + np.int64(0x000000007FFFFFFF), + np.int64(0x000000007FFFFFFF), + np.int64(0x0000000080000000), + np.int64(0x0000000080000001), + np.int64(0x00000000FFFF0000), + np.int64(0x00000000FFFF0001), + np.int64(0x00000000FFFFFFFF), + np.int64(0x00000000FFFFFFFE), + np.int64(0x00000000FFFFFFFF), + np.int64(0x00000000FFFFFFFF), + np.int64(0x0000000100000001), + np.int64(0x0000000100000002), + np.int64(0x0000000100000003), + np.int64(0x0000000200000001), + np.int64(0x0000000200000002), + np.int64(0x0000000200000003), + np.int64(0x0000000300000001), + np.int64(0x0000000300000002), + np.int64(0x0000000300000003), + np.int64(0x00000000FFFFFFFF), + np.int64(-0x7FFFFFFF00000001), + np.int64(0x00000000FFFFFFFE), + np.int64(0x00000000FFFFFFFF), + np.int64(-0x7FFFFFFF00000002), + np.int64(-0x7FFFFFFF00000001), + np.int64(0x00000000FFFFFFFF), + np.int64(-0x7FFFFFFF00000001), + np.int64(-2), + np.int64(-1) + ], + dtype=np.int64) + expected = np.array([op(l, r) for l, r in zip(lhs, rhs)], dtype=np.bool) + self._testBinary(op, lhs, rhs, expected=expected) + def testBroadcasting(self): """Tests broadcasting behavior of an operator.""" @@ -1045,6 +1135,20 @@ class BinaryOpsTest(XLATestCase): ], equality_test=self.ListsAreClose) + def splitvOp(x, y): # pylint: disable=invalid-name + return array_ops.split(value=y, num_or_size_splits=[2, 3], axis=x) + for axis in [1, -1]: + self._testBinary( + splitvOp, + np.int32(axis), + np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], + dtype=dtype), + expected=[ + np.array([[0, 1], [5, 6]], dtype=dtype), + np.array([[2, 3, 4], [7, 8, 9]], dtype=dtype), + ], + equality_test=self.ListsAreClose) + def testTile(self): for dtype in self.numeric_types: self._testBinary( @@ -1181,6 +1285,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/concat_ops_test.py b/tensorflow/compiler/tests/concat_ops_test.py index 81734082d9aab86f8bc763681265ef64ef32bd31..f10973e19f1945515b776cf86349445ed7334629 100644 --- a/tensorflow/compiler/tests/concat_ops_test.py +++ b/tensorflow/compiler/tests/concat_ops_test.py @@ -301,7 +301,7 @@ class ConcatOffsetTest(XLATestCase): s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5], dtypes.int32) s2 = constant_op.constant([2, 20, 5], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1, s2]) + off = gen_array_ops.concat_offset(cdim, [s0, s1, s2]) ans = sess.run(off) self.assertAllEqual(ans, [[0, 0, 0], [0, 3, 0], [0, 10, 0]]) 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/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 538fa8e8e570b83ed681ecc0501285520cabdecb..3bc41b7cfd72bec7572097f8c53eef314a4369f6 100644 --- a/tensorflow/compiler/tests/image_ops_test.py +++ b/tensorflow/compiler/tests/image_ops_test.py @@ -426,7 +426,7 @@ class ResizeBilinearTest(XLATestCase): with self.test_session() as sess, self.test_scope(): dtype = dtype or np.float32 grads = array_ops.placeholder(np.float32) - resized = gen_image_ops._resize_bilinear_grad( + resized = gen_image_ops.resize_bilinear_grad( grads, np.zeros([1, input_shape[0], input_shape[1], 1], dtype=dtype), align_corners=True) diff --git a/tensorflow/compiler/tests/lrn_ops_test.py b/tensorflow/compiler/tests/lrn_ops_test.py index 5d8d89224d4a778d84803811710bb095872e86b2..69bd8f7230d4394c45764d02a88fb0ec097c5756 100644 --- a/tensorflow/compiler/tests/lrn_ops_test.py +++ b/tensorflow/compiler/tests/lrn_ops_test.py @@ -115,11 +115,11 @@ class LRNTest(XLATestCase): out_image = constant_op.constant(out_image_vals, shape=shape) out_grads = constant_op.constant(out_grads_vals, shape=shape) with ops.device(CPU_DEVICE): - expected = gen_nn_ops._lrn_grad(out_grads, in_image, out_image, - depth_radius, bias, alpha, beta) + expected = gen_nn_ops.lrn_grad(out_grads, in_image, out_image, + depth_radius, bias, alpha, beta) with self.test_scope(): - actual = gen_nn_ops._lrn_grad(out_grads, in_image, out_image, - depth_radius, bias, alpha, beta) + actual = gen_nn_ops.lrn_grad(out_grads, in_image, out_image, + depth_radius, bias, alpha, beta) expected_val = expected.eval() actual_val = actual.eval() self.assertAllClose(actual_val, expected_val, rtol=1e-3) 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/pooling_ops_3d_test.py b/tensorflow/compiler/tests/pooling_ops_3d_test.py index eb48fe555a0b182ea7983cbd8c3b217d56350408..4eed903963a34a253ea5c409782d9a89a97a4fdf 100644 --- a/tensorflow/compiler/tests/pooling_ops_3d_test.py +++ b/tensorflow/compiler/tests/pooling_ops_3d_test.py @@ -33,7 +33,7 @@ from tensorflow.python.platform import test # MaxPoolGrad. def _AvgPoolGrad(inputs, outputs, output_gradients, ksize, strides, padding): del outputs # Unused by average-pooling gradients. - return gen_nn_ops._avg_pool3d_grad( + return gen_nn_ops.avg_pool3d_grad( inputs.get_shape().as_list(), output_gradients, ksize=ksize, @@ -263,7 +263,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradValidPadding1_1_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[1, 3, 3, 3, 1], ksize=[1, 1, 1], strides=[1, 1, 1], @@ -272,7 +272,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradValidPadding2_1_6_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[2, 3, 3, 6, 3], ksize=[2, 2, 2], strides=[1, 1, 1], @@ -281,7 +281,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradValidPadding2_1_7_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[2, 3, 5, 7, 3], ksize=[2, 2, 2], strides=[1, 1, 1], @@ -290,7 +290,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradValidPadding2_2_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[2, 2, 2, 2, 3], ksize=[2, 2, 2], strides=[2, 2, 2], @@ -299,7 +299,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradSamePadding1_1_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[2, 3, 2, 4, 1], ksize=[1, 1, 1], strides=[1, 1, 1], @@ -308,7 +308,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradSamePadding2_1_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[2, 3, 2, 4, 1], ksize=[2, 2, 2], strides=[1, 1, 1], @@ -317,7 +317,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradSamePadding2_2_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[2, 5, 2, 4, 3], ksize=[2, 2, 2], strides=[2, 2, 2], @@ -326,7 +326,7 @@ class Pooling3DTest(XLATestCase): def testMaxPoolGradSamePadding3_1_3d(self): self._VerifyGradient( nn_ops.max_pool3d, - gen_nn_ops._max_pool3d_grad, + gen_nn_ops.max_pool3d_grad, input_sizes=[1, 3, 3, 7, 1], ksize=[3, 3, 3], strides=[1, 1, 1], diff --git a/tensorflow/compiler/tests/pooling_ops_test.py b/tensorflow/compiler/tests/pooling_ops_test.py index 7c19a99c4eb4be3ca34b3ce949216e557b0a681d..fe270af3d636c0824621f36360ce9e7d14d8fc91 100644 --- a/tensorflow/compiler/tests/pooling_ops_test.py +++ b/tensorflow/compiler/tests/pooling_ops_test.py @@ -292,8 +292,15 @@ class PoolGradTest(XLATestCase): CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0" - def _VerifyOneTest(self, pool_func, pool_grad_func, input_sizes, ksize, - strides, padding, data_format): + def _VerifyOneTest(self, + pool_func, + pool_grad_func, + input_sizes, + ksize, + strides, + padding, + data_format, + pool_grad_grad_func=None): """Verifies the output values of the pooling gradient function. Args: @@ -304,9 +311,19 @@ class PoolGradTest(XLATestCase): strides: The stride dimensions padding: Padding type. data_format: The data format we use to run the pooling operation. + pool_grad_grad_func: Second-order gradient function, if available. """ total_size = np.prod(input_sizes) - x = np.arange(1, total_size + 1, dtype=np.float32).reshape(input_sizes) + # TODO(b/73062247): MaxPoolGradGrad can confuse gradients when x is equally + # maximal at 16 bits. Switch to np.random.randn when resolved. + x = np.arange(1, total_size + 1, dtype=np.float32) + x *= (np.random.randint(2, size=total_size) * 2 - 1) # Flip signs randomly + # Verify some specifically interesting values... + x[np.random.choice(total_size)] = np.inf + x[np.random.choice(total_size)] = -np.inf + # TODO(b/74222344): Fix nan handling for max pool grad. + # x[np.random.choice(total_size)] = np.nan + x = x.reshape(input_sizes) with self.test_session() as sess: # Use the forward pool function to compute some corresponding outputs # (needed for the CPU device, and we need the shape in both cases). @@ -323,6 +340,8 @@ class PoolGradTest(XLATestCase): output_gradient_vals = np.arange( 1, output_vals.size + 1, dtype=np.float32) output_gradient_vals = output_gradient_vals.reshape(output_vals.shape) + output_grad_grad_vals = np.arange(1, x.size + 1, dtype=np.float32) + output_grad_grad_vals = output_grad_grad_vals.reshape(x.shape) # Use the Tensorflow CPU pooling gradient to compute the expected input # gradients. @@ -342,18 +361,36 @@ class PoolGradTest(XLATestCase): {inputs: x, output_gradients: output_gradient_vals}) + output_grad_gradients = array_ops.placeholder( + dtypes.float32, shape=expected_input_gradient_vals.shape) + if pool_grad_grad_func is not None: + expected_grad_gradients = pool_grad_grad_func( + inputs, + outputs, + output_grad_gradients, + ksize=ksize, + strides=strides, + padding=padding, + data_format="NHWC") + expected_grad_gradients_vals = sess.run(expected_grad_gradients, { + inputs: x, + output_grad_gradients: output_grad_grad_vals + }) + # Run the gradient op on the XLA device with self.test_scope(): outputs = array_ops.placeholder(dtypes.float32, shape=output_vals.shape) xla_inputs = inputs xla_outputs = outputs xla_output_gradients = output_gradients + xla_output_grad_gradients = output_grad_gradients xla_ksize = ksize xla_strides = strides if data_format == "NCHW": xla_inputs = NHWCToNCHW(inputs) xla_outputs = NHWCToNCHW(outputs) xla_output_gradients = NHWCToNCHW(output_gradients) + xla_output_grad_gradients = NHWCToNCHW(output_grad_gradients) xla_ksize = NHWCToNCHW(ksize) xla_strides = NHWCToNCHW(strides) actual_input_gradients = pool_grad_func( @@ -366,22 +403,54 @@ class PoolGradTest(XLATestCase): data_format=data_format) if data_format == "NCHW": actual_input_gradients = NCHWToNHWC(actual_input_gradients) - actual = sess.run(actual_input_gradients, { + if pool_grad_grad_func is not None: + actual_grad_gradients = pool_grad_grad_func( + xla_inputs, + xla_outputs, + xla_output_grad_gradients, + ksize=xla_ksize, + strides=xla_strides, + padding=padding, + data_format=data_format) + if data_format == "NCHW": + actual_grad_gradients = NCHWToNHWC(actual_grad_gradients) + actual_input_gradients_vals = sess.run(actual_input_gradients, { inputs: x, outputs: output_vals, output_gradients: output_gradient_vals }) - # Compare the Tensorflow and XLA results. self.assertAllClose( - expected_input_gradient_vals.flatten(), - actual.flatten(), + expected_input_gradient_vals, + actual_input_gradients_vals, rtol=1e-4, atol=1e-6) - self.assertShapeEqual(actual, inputs) - - def _VerifyValues(self, pool_func, pool_grad_func, input_sizes, ksize, - strides, padding): + self.assertShapeEqual(actual_input_gradients_vals, inputs) + + if pool_grad_grad_func is not None: + actual_grad_gradients_vals = sess.run( + actual_grad_gradients, { + inputs: x, + outputs: output_vals, + output_grad_gradients: output_grad_grad_vals + }) + + # Compare the Tensorflow and XLA results. + self.assertAllClose( + expected_grad_gradients_vals, + actual_grad_gradients_vals, + rtol=1e-4, + atol=1e-6) + self.assertShapeEqual(actual_grad_gradients_vals, outputs) + + def _VerifyValues(self, + pool_func, + pool_grad_func, + input_sizes, + ksize, + strides, + padding, + pool_grad_grad_func=None): """Verifies the output values of the pooling function. Args: @@ -391,12 +460,20 @@ class PoolGradTest(XLATestCase): ksize: The kernel size dimensions strides: The stride dimensions padding: Padding type. + pool_grad_grad_func: Second-order gradient function, if available. """ for data_format in GetTestConfigs(): - self._VerifyOneTest(pool_func, pool_grad_func, input_sizes, ksize, - strides, padding, data_format) - - def _TestPooling(self, forward_op, backward_op): + self._VerifyOneTest( + pool_func, + pool_grad_func, + input_sizes, + ksize, + strides, + padding, + data_format, + pool_grad_grad_func=pool_grad_grad_func) + + def _TestPooling(self, forward_op, backward_op, pool_grad_grad_func=None): # VALID padding self._VerifyValues( forward_op, @@ -404,7 +481,8 @@ class PoolGradTest(XLATestCase): input_sizes=[1, 3, 3, 3], ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], - padding="VALID") + padding="VALID", + pool_grad_grad_func=pool_grad_grad_func) # SAME padding self._VerifyValues( @@ -413,7 +491,8 @@ class PoolGradTest(XLATestCase): input_sizes=[1, 2, 3, 3], ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], - padding="SAME") + padding="SAME", + pool_grad_grad_func=pool_grad_grad_func) # SAME padding, non square window self._VerifyValues( @@ -422,7 +501,8 @@ class PoolGradTest(XLATestCase): input_sizes=[1, 2, 2, 1], ksize=[1, 1, 2, 1], strides=[1, 1, 1, 1], - padding="SAME") + padding="SAME", + pool_grad_grad_func=pool_grad_grad_func) # VALID padding, uneven stride self._VerifyValues( @@ -431,14 +511,16 @@ class PoolGradTest(XLATestCase): input_sizes=[1, 4, 4, 1], ksize=[1, 2, 2, 1], strides=[1, 1, 2, 1], - padding="VALID") + padding="VALID", + pool_grad_grad_func=pool_grad_grad_func) self._VerifyValues( forward_op, backward_op, input_sizes=[1, 4, 4, 1], ksize=[1, 2, 2, 1], strides=[1, 2, 1, 1], - padding="VALID") + padding="VALID", + pool_grad_grad_func=pool_grad_grad_func) # SAME padding, size 4 input self._VerifyValues( @@ -447,7 +529,8 @@ class PoolGradTest(XLATestCase): input_sizes=[1, 4, 4, 4], ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], - padding="SAME") + padding="SAME", + pool_grad_grad_func=pool_grad_grad_func) # SAME padding, size 8 input self._VerifyValues( @@ -456,10 +539,14 @@ class PoolGradTest(XLATestCase): input_sizes=[1, 8, 8, 8], ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1], - padding="SAME") + padding="SAME", + pool_grad_grad_func=pool_grad_grad_func) def testMaxPool(self): - self._TestPooling(nn_ops.max_pool, gen_nn_ops._max_pool_grad) + self._TestPooling( + nn_ops.max_pool, + gen_nn_ops.max_pool_grad, + pool_grad_grad_func=gen_nn_ops.max_pool_grad_grad) def testAvgPool(self): # Wrapper around AvgPoolGrad that ignores extra arguments needed by @@ -467,7 +554,7 @@ class PoolGradTest(XLATestCase): def AvgPoolGrad(inputs, outputs, output_gradients, ksize, strides, padding, data_format): del outputs # Unused by average-pooling gradients. - return gen_nn_ops._avg_pool_grad( + return gen_nn_ops.avg_pool_grad( inputs.get_shape().as_list(), output_gradients, ksize=ksize, @@ -483,7 +570,7 @@ class PoolGradTest(XLATestCase): def testMaxPoolKernelSmallerThanStrideValid(self): self._VerifyValues( nn_ops.max_pool, - gen_nn_ops._max_pool_grad, + gen_nn_ops.max_pool_grad, input_sizes=[1, 7, 7, 1], ksize=[1, 2, 2, 1], strides=[1, 3, 3, 1], @@ -492,7 +579,7 @@ class PoolGradTest(XLATestCase): def testMaxPoolKernelSmallerThanStrideSame(self): self._VerifyValues( nn_ops.max_pool, - gen_nn_ops._max_pool_grad, + gen_nn_ops.max_pool_grad, input_sizes=[1, 3, 3, 1], ksize=[1, 1, 1, 1], strides=[1, 2, 2, 1], @@ -500,7 +587,7 @@ class PoolGradTest(XLATestCase): self._VerifyValues( nn_ops.max_pool, - gen_nn_ops._max_pool_grad, + gen_nn_ops.max_pool_grad, input_sizes=[1, 4, 4, 1], ksize=[1, 1, 1, 1], strides=[1, 2, 2, 1], diff --git a/tensorflow/compiler/tests/randomized_tests.cc b/tensorflow/compiler/tests/randomized_tests.cc index e72dd4eea9f127e1df96ab166103c4c16372adb6..e53efc3091d8935e745122af29abd7b8063b1d01 100644 --- a/tensorflow/compiler/tests/randomized_tests.cc +++ b/tensorflow/compiler/tests/randomized_tests.cc @@ -83,8 +83,8 @@ string LocalDeviceToFullDeviceName(const string& device) { return strings::StrCat("/job:localhost/replica:0/task:0/device:", device); } -constexpr std::array kAllXlaTypes = { - {DT_INT32, DT_FLOAT, DT_BOOL, DT_COMPLEX64}}; +constexpr std::array kAllXlaTypes = { + {DT_INT32, DT_FLOAT, DT_BOOL, DT_COMPLEX64, DT_INT64}}; // An OpTestBuilder is a graph builder class that takes as input an operator to // test, its inputs and attributes, and builds a graph that executes the diff --git a/tensorflow/compiler/tests/reduce_ops_test.py b/tensorflow/compiler/tests/reduce_ops_test.py index 965fdf684b973498d0b3c3cde17711cca7279705..2c084b04fa2f67ad0d86508109522d7bead206eb 100644 --- a/tensorflow/compiler/tests/reduce_ops_test.py +++ b/tensorflow/compiler/tests/reduce_ops_test.py @@ -18,6 +18,7 @@ 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 @@ -30,8 +31,13 @@ from tensorflow.python.platform import googletest class ReduceOpsTest(XLATestCase): - def _testReduction(self, tf_reduce_fn, np_reduce_fn, dtype, test_inputs, - rtol=1e-4, atol=1e-4): + def _testReduction(self, + tf_reduce_fn, + np_reduce_fn, + dtype, + test_inputs, + rtol=1e-4, + atol=1e-4): """Tests that the output of 'tf_reduce_fn' matches numpy's output.""" for test_input in test_inputs: @@ -41,16 +47,16 @@ class ReduceOpsTest(XLATestCase): index = array_ops.placeholder(dtypes.int32) out = tf_reduce_fn(a, index) result = sess.run(out, {a: test_input, index: [0]}) - self.assertAllClose(result, np_reduce_fn(test_input, axis=0), - rtol=rtol, atol=atol) + self.assertAllClose( + result, np_reduce_fn(test_input, axis=0), rtol=rtol, atol=atol) result = sess.run(out, {a: test_input, index: [1]}) - self.assertAllClose(result, np_reduce_fn(test_input, axis=1), - rtol=rtol, atol=atol) + self.assertAllClose( + result, np_reduce_fn(test_input, axis=1), rtol=rtol, atol=atol) result = sess.run(out, {a: test_input, index: [-1]}) - self.assertAllClose(result, np_reduce_fn(test_input, axis=1), - rtol=rtol, atol=atol) + self.assertAllClose( + result, np_reduce_fn(test_input, axis=1), rtol=rtol, atol=atol) with self.assertRaisesWithPredicateMatch( errors_impl.InvalidArgumentError, 'Invalid reduction dim'): @@ -60,7 +66,7 @@ class ReduceOpsTest(XLATestCase): errors_impl.InvalidArgumentError, 'Invalid reduction dim'): sess.run(out, {a: test_input, index: [2]}) - FLOAT_DATA = [ + REAL_DATA = [ np.zeros(shape=(2, 0)), np.zeros(shape=(0, 30)), np.arange(1, 7).reshape(2, 3), @@ -74,7 +80,7 @@ class ReduceOpsTest(XLATestCase): np.arange(-14, -2, dtype=np.float32).view(np.complex64).reshape(2, 3), np.arange(-4, 8, dtype=np.float32).view(np.complex64).reshape(2, 3), ] - NONEMPTY_FLOAT_DATA = [x for x in FLOAT_DATA if np.size(x) > 0] + NONEMPTY_REAL_DATA = [x for x in REAL_DATA if np.size(x) > 0] NONEMPTY_COMPLEX_DATA = [x for x in COMPLEX_DATA if np.size(x) > 0] BOOL_DATA = [ np.array([], dtype=np.bool).reshape(2, 0), @@ -83,8 +89,7 @@ class ReduceOpsTest(XLATestCase): ] def testReduceSumF32(self): - self._testReduction(math_ops.reduce_sum, np.sum, np.float32, - self.FLOAT_DATA) + self._testReduction(math_ops.reduce_sum, np.sum, np.float32, self.REAL_DATA) def testReduceSumC64(self): self._testReduction(math_ops.reduce_sum, np.sum, np.complex64, @@ -92,7 +97,7 @@ class ReduceOpsTest(XLATestCase): def testReduceProdF32(self): self._testReduction(math_ops.reduce_prod, np.prod, np.float32, - self.FLOAT_DATA) + self.REAL_DATA) def testReduceProdC64(self): self._testReduction(math_ops.reduce_prod, np.prod, np.complex64, @@ -100,31 +105,44 @@ class ReduceOpsTest(XLATestCase): def testReduceMin(self): - def reference_min(inp, axis): + def reference_min(dtype, inp, axis): """Wrapper around np.amin that returns +infinity for an empty input.""" if inp.shape[axis] == 0: - return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], float('inf')) + if np.issubdtype(dtype, np.floating): + return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], float('inf')) + return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], + np.iinfo(dtype).max) return np.amin(inp, axis) - self._testReduction(math_ops.reduce_min, reference_min, np.float32, - self.FLOAT_DATA) + for dtype in set(self.all_types).intersection( + [np.float32, np.int32, np.int64]): + self._testReduction(math_ops.reduce_min, + functools.partial(reference_min, dtype), dtype, + self.REAL_DATA) def testReduceMax(self): - def reference_max(inp, axis): + def reference_max(dtype, inp, axis): """Wrapper around np.amax that returns -infinity for an empty input.""" if inp.shape[axis] == 0: - return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], float('-inf')) + if np.issubdtype(dtype, np.floating): + return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], + float('-inf')) + return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], + np.iinfo(dtype).min) return np.amax(inp, axis) - self._testReduction(math_ops.reduce_max, reference_max, np.float32, - self.FLOAT_DATA) + for dtype in set(self.all_types).intersection( + [np.float32, np.int32, np.int64]): + self._testReduction(math_ops.reduce_max, + functools.partial(reference_max, dtype), dtype, + self.REAL_DATA) def testReduceMeanF32(self): # TODO(phawkins): mean on XLA currently returns 0 instead of NaN when # reducing across zero inputs. self._testReduction(math_ops.reduce_mean, np.mean, np.float32, - self.NONEMPTY_FLOAT_DATA) + self.NONEMPTY_REAL_DATA) def testReduceMeanC64(self): self._testReduction(math_ops.reduce_mean, np.mean, np.complex64, 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..4a9c0e7471f9cdb2a47b54705495d2dda9748890 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([6, 3, 0, 6], dtype=dtype), + self.UnsortedSegmentSum( + np.array([0, 1, 2, 3, 4, 5, 6], dtype=dtype), + np.array([3, -1, 0, 1, 0, -1, 3], dtype=np.int32), 4)) + def testUnsortedSegmentSum1DIndices2DDataDisjoint(self): for dtype in self.numeric_types: data = np.array( diff --git a/tensorflow/compiler/tests/slice_ops_test.py b/tensorflow/compiler/tests/slice_ops_test.py index a7cbfb04003c397212a35e16c6b23d7c2a18f7df..305ca0c6b78d3ef985deb38816f9388e7983906b 100644 --- a/tensorflow/compiler/tests/slice_ops_test.py +++ b/tensorflow/compiler/tests/slice_ops_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.compiler.tests.xla_test import XLATestCase from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.platform import googletest @@ -137,6 +138,34 @@ class StridedSliceTest(XLATestCase): self.assertAllEqual([6, 4], result) + def test2DDegenerate(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[2, 3]) + with self.test_scope(): + o = array_ops.strided_slice(i, [-1, 0], [0, 3]) + params = { + i: [[0, 1, 2], + [3, 4, 5]] + } + result = o.eval(feed_dict=params) + + self.assertEqual(tensor_shape.TensorShape((0, 3)), result.shape) + + def test2DDegenerateNegativeStride(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[2, 3]) + with self.test_scope(): + o = array_ops.strided_slice(i, [0, 0], [-1, 3], [-1, 1]) + params = { + i: [[0, 1, 2], + [3, 4, 5]] + } + result = o.eval(feed_dict=params) + + self.assertEqual(tensor_shape.TensorShape((0, 3)), result.shape) + def test3D(self): for dtype in self.numeric_types: with self.test_session(): diff --git a/tensorflow/compiler/tests/spacetobatch_op_test.py b/tensorflow/compiler/tests/spacetobatch_op_test.py index c013f4b50a4cf95be8028248c52b10b1c3be2bd3..92518aadc4bf5c601cfb4192c093799784b6aa72 100644 --- a/tensorflow/compiler/tests/spacetobatch_op_test.py +++ b/tensorflow/compiler/tests/spacetobatch_op_test.py @@ -75,11 +75,11 @@ class SpaceToBatchTest(XLATestCase): for dtype in self.float_types: # outputs = space_to_batch(inputs) placeholder = array_ops.placeholder(dtype) - x_tf = gen_array_ops._space_to_batch( + x_tf = gen_array_ops.space_to_batch( placeholder, paddings, block_size=block_size) self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs) # inputs = batch_to_space(outputs) - x_tf = gen_array_ops._batch_to_space( + x_tf = gen_array_ops.batch_to_space( placeholder, paddings, block_size=block_size) self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs) diff --git a/tensorflow/compiler/tests/stack_ops_test.py b/tensorflow/compiler/tests/stack_ops_test.py index 2b9c2279737ccee531d488d27ccdb0cafa1dc8fc..94342f9567ca71274609e63b0482d55637c98d51 100644 --- a/tensorflow/compiler/tests/stack_ops_test.py +++ b/tensorflow/compiler/tests/stack_ops_test.py @@ -34,33 +34,33 @@ class StackOpTest(XLATestCase): with self.test_session(), self.test_scope(): size = array_ops.placeholder(dtypes.int32) v = array_ops.placeholder(dtypes.float32) - h = gen_data_flow_ops._stack_v2(size, dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push_v2(h, v) + h = gen_data_flow_ops.stack_v2(size, dtypes.float32, stack_name="foo") + c = gen_data_flow_ops.stack_push_v2(h, v) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32) self.assertAllClose([[4.0, 5.0]], c1.eval({size: 5, v: [[4.0, 5.0]]})) def testStackPushPopSwap(self): with self.test_session(), self.test_scope(): a = np.arange(2000) x = array_ops.placeholder(dtypes.float32) - h = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push_v2(h, x, swap_memory=True) + h = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo") + c = gen_data_flow_ops.stack_push_v2(h, x, swap_memory=True) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32) self.assertAllClose(a, c1.eval({x: a})) def testMultiStack(self): with self.test_session(), self.test_scope(): v = array_ops.placeholder(dtypes.float32) - h1 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_push_v2(h1, v) + h1 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo") + c1 = gen_data_flow_ops.stack_push_v2(h1, v) with ops.control_dependencies([c1]): - c1 = gen_data_flow_ops._stack_pop_v2(h1, dtypes.float32) - h2 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="bar") - c2 = gen_data_flow_ops._stack_push_v2(h2, 5.0) + c1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32) + h2 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="bar") + c2 = gen_data_flow_ops.stack_push_v2(h2, 5.0) with ops.control_dependencies([c2]): - c2 = gen_data_flow_ops._stack_pop_v2(h2, dtypes.float32) + c2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32) r = c1 + c2 self.assertAllClose(9.0, r.eval({v: 4.0})) @@ -69,15 +69,15 @@ class StackOpTest(XLATestCase): with self.test_session() as sess, self.test_scope(): v1 = array_ops.placeholder(dtypes.float32) v2 = array_ops.placeholder(dtypes.float32) - h1 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") - h2 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") + h1 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo") + h2 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_push_v2(h1, v1) + c1 = gen_data_flow_ops.stack_push_v2(h1, v1) with ops.control_dependencies([c1]): - c2 = gen_data_flow_ops._stack_push_v2(h2, v2) + c2 = gen_data_flow_ops.stack_push_v2(h2, v2) with ops.control_dependencies([c2]): - pop1 = gen_data_flow_ops._stack_pop_v2(h1, dtypes.float32) - pop2 = gen_data_flow_ops._stack_pop_v2(h2, dtypes.float32) + pop1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32) + pop2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32) out1, out2 = sess.run([pop1, pop2], {v1: 4.0, v2: 5.0}) self.assertAllClose(out1, 4.0) @@ -86,17 +86,17 @@ class StackOpTest(XLATestCase): def testCloseStack(self): with self.test_session() as sess, self.test_scope(): size = array_ops.placeholder(dtypes.int32) - h = gen_data_flow_ops._stack_v2(size, dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_close_v2(h) + h = gen_data_flow_ops.stack_v2(size, dtypes.float32, stack_name="foo") + c1 = gen_data_flow_ops.stack_close_v2(h) sess.run(c1, {size: 5}) def testPushCloseStack(self): with self.test_session() as sess, self.test_scope(): v = array_ops.placeholder(dtypes.float32) - h = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push_v2(h, v) + h = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo") + c = gen_data_flow_ops.stack_push_v2(h, v) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_close_v2(h) + c1 = gen_data_flow_ops.stack_close_v2(h) sess.run(c1, {v: [[4.0, 5.0]]}) diff --git a/tensorflow/compiler/tests/tensor_array_ops_test.py b/tensorflow/compiler/tests/tensor_array_ops_test.py index a62925a1818da00cb0a9e82e1281db20fb38b208..7624d6e4b2e2ece6a61155743fc8b866f6903f32 100644 --- a/tensorflow/compiler/tests/tensor_array_ops_test.py +++ b/tensorflow/compiler/tests/tensor_array_ops_test.py @@ -338,7 +338,7 @@ class TensorArrayTest(xla_test.XLATestCase): w0 = ta.write(0, [[4.0, 5.0]]) # Test reading wrong datatype. - r0_bad = gen_data_flow_ops._tensor_array_read_v3( + r0_bad = gen_data_flow_ops.tensor_array_read_v3( handle=w0.handle, index=0, dtype=dtype2, flow_in=w0.flow) with self.assertRaisesOpError("TensorArray dtype is "): r0_bad.eval() diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py index 8e4b8a38336c5e8b2e10edc4c81502eeebb628d2..3d3e112f4821ea8e57ea9589a5b4433647ad294b 100644 --- a/tensorflow/compiler/tests/unary_ops_test.py +++ b/tensorflow/compiler/tests/unary_ops_test.py @@ -154,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/tests/xla_test.py b/tensorflow/compiler/tests/xla_test.py index 7e1f5c76ed65946363cc3c113ab1a9862f87b289..cc778f1c3c0098da5ab933f9b4674890a724d160 100644 --- a/tensorflow/compiler/tests/xla_test.py +++ b/tensorflow/compiler/tests/xla_test.py @@ -71,14 +71,14 @@ class XLATestCase(test.TestCase): self._all_types = set( [dtype.as_numpy_dtype for dtype in self._all_tf_types]) - self.int_types = set([dtype.as_numpy_dtype for dtype in self.int_tf_types]) + self._int_types = set([dtype.as_numpy_dtype for dtype in self.int_tf_types]) self._float_types = set( [dtype.as_numpy_dtype for dtype in self._float_tf_types]) self.complex_types = set([ dtype.as_numpy_dtype for dtype in self.complex_tf_types ]) - self._numeric_types = set( - self.int_types | self._float_types | self.complex_types) + self._numeric_types = set(self._int_types | self._float_types + | self.complex_types) # Parse the manifest file, if any, into a regex identifying tests to # disable @@ -130,6 +130,11 @@ class XLATestCase(test.TestCase): name = '{}.{}'.format(type(self).__name__, self._testMethodName) return self._float_tf_types - self._method_types_filter.get(name, set()) + @property + def int_types(self): + name = '{}.{}'.format(type(self).__name__, self._testMethodName) + return self._int_types - self._method_types_filter.get(name, set()) + @property def numeric_tf_types(self): name = '{}.{}'.format(type(self).__name__, self._testMethodName) diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 3c7dfef03dfb5d86dd63fd4aa84ad56081833035..eb20ca501c80b01c76198e1ad54173f1c601714d 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -58,6 +58,15 @@ xla_proto_library( ], ) +xla_proto_library( + name = "host_compute_metadata_proto", + srcs = ["host_compute_metadata.proto"], + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core:protos_all_cc", + ], +) + cc_library( name = "tf2xla", srcs = ["tf2xla.cc"], @@ -149,6 +158,7 @@ cc_library( ":common", ":dump_graph", ":functionalize_control_flow", + ":host_compute_metadata_proto", ":sharding_util", ":tf2xla_util", "//tensorflow/compiler/tf2xla/lib:util", @@ -312,6 +322,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..de1008803d69fefa415c7bdbe6c27a62e625b417 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.cc +++ b/tensorflow/compiler/tf2xla/const_analysis.cc @@ -37,7 +37,7 @@ Status BackwardsConstAnalysis(const Graph& g, }; Status status; - std::unordered_set must_be_const; + std::unordered_set must_be_const; auto visit = [&status, &metadata_ops, &must_be_const, compile_time_const_args](Node* node) { if (!status.ok()) return; @@ -55,8 +55,10 @@ Status BackwardsConstAnalysis(const Graph& g, compile_time_const_args->at(index) = true; return; } - for (Node* pred : node->in_nodes()) { - must_be_const.insert(pred); + for (const Edge* pred : node->in_edges()) { + if (!pred->IsControlEdge()) { + must_be_const.insert(pred->src()); + } } return; } @@ -64,7 +66,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/const_analysis_test.cc b/tensorflow/compiler/tf2xla/const_analysis_test.cc index 9d125f8d499863cfaa0e26b5b633ca02914d1b7d..992b12c06db5efc0ae54284d0ea77017c1c79aca 100644 --- a/tensorflow/compiler/tf2xla/const_analysis_test.cc +++ b/tensorflow/compiler/tf2xla/const_analysis_test.cc @@ -79,5 +79,24 @@ TEST(ConstAnalysisTest, TopologicalOrder) { } } +TEST(ConstAnalysisTest, DontFollowControlDependencies) { + Scope root = Scope::NewRootScope(); + + Output arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0); + Output arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1); + Output c1 = + ops::Const(root.WithOpName("c1").WithControlDependencies(arg0), 1, {1}); + Output add = ops::Add(root, arg1, c1); + Output reshape = ops::Reshape(root, arg1, add); + + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(root.ToGraph(&graph)); + + std::vector const_args(2, false); + TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args)); + + EXPECT_EQ(const_args, std::vector({false, true})); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc index 1d9e0fb33ee4a4229c78d116831e95391a5ac3f8..8b7beef83ec2ed0df780d6a9cb2a4bcf737d008b 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc @@ -285,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 @@ -427,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); + } } } } @@ -450,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); @@ -531,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(); } @@ -561,14 +583,16 @@ class FunctionalizeCond { // CondArgNode represents a input to the conditional and its corresponding // switch nodes. struct CondArgNode { - explicit CondArgNode(Node* input) : input(input) {} + explicit CondArgNode(Node* src, int src_output) + : src(src), src_output(src_output) {} string ToString() const { - return strings::StrCat("input=", input->name(), - " switches=", NodesToString(switch_nodes)); + return strings::StrCat("src=", src->name(), ":", src_output, + " switches=", NodesToString(switches)); } - Node* input; - std::vector switch_nodes; + Node* src; + int src_output; + std::vector switches; }; using CondArgNodes = std::vector; @@ -582,15 +606,23 @@ class FunctionalizeCond { int count; }; - struct PredicateSwitches { - explicit PredicateSwitches(Node* predicate) : predicate(predicate) {} + // Group of switch nodes that will be part of the same XlaIf. + struct SwitchCluster { + explicit SwitchCluster(const Edge* predicate_edge) + : predicate_edge(predicate_edge) {} + string ToString() const { + return strings::StrCat(name, " predicate=", predicate_edge->src()->name(), + " switches=", NodesToString(switches)); + } - Node* predicate; + string name; + const Edge* predicate_edge; std::vector switches; }; - FunctionalizeCond(Graph* graph, FunctionLibraryDefinition* library) - : library_(library), graph_(graph) {} + 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 @@ -601,40 +633,38 @@ class FunctionalizeCond { // frontier (the nodes where the cond ends). StatusOr, std::unordered_set>> - DetermineBranchMapAndFrontier(const std::vector& switches); + 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 std::vector& switch_nodes, - const std::vector& merge_nodes, - Node* predicate); + const SwitchCluster& switch_cluster, + const std::vector& switches); // Builds a XlaIfOp to replace the Switch-Graph-Merge cluster with. StatusOr BuildAndAddXlaIfOp(const CondArgNodes& cond_arg_nodes, - const std::vector& switch_nodes, - const std::vector& merge_nodes, - Node* predicate); + 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 CondArgNodes& cond_arg_nodes, - const std::vector& switch_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 CondArgNodes& cond_arg_nodes, Node* predicate, - Node* if_node); + Status AddInputEdges(const CondArgNodes& cond_arg_nodes, + const Edge* predicate_edge, 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_ (along with grouping predicates) in - // postorder. Dead switch nodes are skipped and removed from the graph. - std::vector DeterminePredicateSwitchOrder(); + // 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. @@ -657,6 +687,7 @@ class FunctionalizeCond { FunctionLibraryDefinition* library_; Graph* graph_; + bool dump_graphs_; }; bool IsDeadSwitch(const Node* node) { @@ -704,10 +735,13 @@ Status FunctionalizeCond::ValidateFrontier( ") in both Else and Then branch should be in Both."); } } - if (pending[kBoth].empty() && pending[kThenBranch].empty() && - pending[kElseBranch].empty()) { - return errors::Internal("Unexpected empty frontier for switch nodes"); - } + // 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(); } @@ -725,8 +759,8 @@ Status FunctionalizeCond::Join(const ForwardFlowNode& src_state, if (IsMerge(dst)) { dst_state->branch = Branch::kBoth; } else { - return errors::Internal("Illegal merge: ", src_state.ToString(), " with ", - dst_state->ToString(), " for ", + return errors::Internal("Illegal merge:\n", src_state.ToString(), + " with ", dst_state->ToString(), " for\n", dst->DebugString()); } } @@ -734,48 +768,217 @@ Status FunctionalizeCond::Join(const ForwardFlowNode& src_state, return Status::OK(); } -std::vector +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); } - std::vector predicate_switch_order; - if (switch_order.empty()) { - return predicate_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; + }; + + // 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()] = find_output_cluster(input); + UnionFind* cluster = entry_cluster[n->id()]; + 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()]); + } + } + } + + 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_); + } + + // 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::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 // node whose cluster is later in the topological order of clustered // switches). 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); + const Edge* pred_edge; + TF_CHECK_OK((*it)->input_edge(1, &pred_edge)); + // The predicate can be preceded by a identity node. Look through identity + // nodes to predicate. + while (pred_edge->src()->IsIdentity()) { + TF_CHECK_OK(pred_edge->src()->input_edge(0, &pred_edge)); + } + auto repr = std::make_pair(pred_edge->src(), clusters[(*it)->id()].Get()); + if (predicate_index.find(repr) == predicate_index.end()) { + predicate_index[repr] = switch_clusters.size(); + switch_clusters.emplace_back(pred_edge); + // 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_edge->src()->name(), "_", repr.second.representative, "_If"); } - predicate_switch_order[predicate_index[pred]].switches.push_back(*it); + switch_clusters[predicate_index[repr]].switches.push_back(*it); } - return predicate_switch_order; + + return switch_clusters; } StatusOr> @@ -823,10 +1026,10 @@ StatusOr< std::pair, std::unordered_set>> FunctionalizeCond::DetermineBranchMapAndFrontier( - const std::vector& switches) { + const SwitchCluster& switch_cluster) { std::unordered_map branch_map; std::unordered_set frontier; - std::vector stack = switches; + std::vector stack = switch_cluster.switches; std::vector visited(graph_->num_node_ids(), false); while (!stack.empty()) { Node* n = stack.back(); @@ -849,9 +1052,12 @@ FunctionalizeCond::DetermineBranchMapAndFrontier( 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)); + TF_RETURN_WITH_CONTEXT_IF_ERROR( + Join(ForwardFlowNode(Branch(index)), out, &ffn), " when joining ", + e->DebugString()); } else { - TF_RETURN_IF_ERROR(Join(branch_map[n], out, &ffn)); + TF_RETURN_WITH_CONTEXT_IF_ERROR(Join(branch_map[n], out, &ffn), + " when joining ", e->DebugString()); } if (IsMerge(out)) { if (out->in_edges().size() == ffn.count) { @@ -868,7 +1074,7 @@ FunctionalizeCond::DetermineBranchMapAndFrontier( } } - if (VLOG_IS_ON(2)) { + 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. @@ -880,41 +1086,50 @@ FunctionalizeCond::DetermineBranchMapAndFrontier( } Status FunctionalizeCond::FunctionalizeInternal() { - std::vector predicate_switch_order = - DeterminePredicateSwitchOrder(); + 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: " << NodesToString(ps.switches) << " (" - << ps.predicate->name() << ")"; + VLOG(3) << "Flow down from: " << ps.ToString(); std::unordered_map branch_map; std::unordered_set frontier; TF_ASSIGN_OR_RETURN(std::tie(branch_map, frontier), - DetermineBranchMapAndFrontier(ps.switches)); + DetermineBranchMapAndFrontier(ps)); - VLOG(2) << "FunctionalizeControlFlow (before XlaIf conversion): " - << dump_graph::DumpGraphToFile("functionalize_bc", *graph_); + if (dump_graphs_) + LOG(INFO) << "FunctionalizeControlFlow (before XlaIf conversion): " + << dump_graph::DumpGraphToFile("functionalize_bc", *graph_, + library_); TF_RETURN_IF_ERROR(ValidateFrontier(branch_map, frontier)); + struct Hash { + size_t operator()(const std::pair& item) const { + return Hash64Combine(hash()(item.first), + std::hash()(item.second)); + } + }; + // 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()); + std::unordered_map, int, Hash> input_index; 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); + const Edge* e; + TF_RETURN_IF_ERROR(switch_node->input_edge(0, &e)); + std::pair key = std::make_pair(e->src(), e->src_output()); + if (input_index.find(key) == input_index.end()) { + input_index[key] = cond_arg_nodes.size(); + cond_arg_nodes.emplace_back(key.first, key.second); } - cond_arg_nodes.at(input_index.at(in)).switch_nodes.push_back(switch_node); + cond_arg_nodes.at(input_index.at(key)).switches.push_back(switch_node); } std::vector merge_nodes(frontier.begin(), frontier.end()); std::sort(merge_nodes.begin(), merge_nodes.end(), NodeCmp()); @@ -923,9 +1138,8 @@ Status FunctionalizeCond::FunctionalizeInternal() { EnsureDominanceAndReturnNonDominatedControlNodes( branch_map, ps.switches)); - TF_ASSIGN_OR_RETURN( - Node * if_node, - ConvertToXlaIf(cond_arg_nodes, ps.switches, merge_nodes, ps.predicate)); + 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); } @@ -934,25 +1148,26 @@ Status FunctionalizeCond::FunctionalizeInternal() { graph_->RemoveNode(del_kv.first); } for (auto& kv : cond_arg_nodes) { - for (Node* node : kv.switch_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(); } StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( - const CondArgNodes& cond_arg_nodes, 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); + 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); @@ -962,12 +1177,9 @@ 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(cond_arg_nodes, 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()); - VLOG(4) << "FunctionalizeControlFlow (" << branch[i] << "): " - << dump_graph::DumpGraphToFile( - strings::StrCat("functionalize_", branch[i]), *body); FunctionDef body_fdef; TF_RETURN_IF_ERROR(GraphToFunctionDef(*body, body_name.name(), &body_fdef)); TF_RETURN_IF_ERROR(library_->AddFunctionDef(body_fdef)); @@ -979,7 +1191,7 @@ StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( DataTypeVector in_arg_types; for (auto& kv : cond_arg_nodes) { bool inserted = false; - for (const Node* arg : kv.switch_nodes) { + for (const Node* arg : kv.switches) { const Edge* in_edge; TF_RETURN_IF_ERROR(arg->input_edge(0, &in_edge)); if (in_edge->IsControlEdge()) { @@ -1006,10 +1218,12 @@ StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( builder.Attr("Tout", out_type); builder.Attr("Tcond", DT_BOOL); - builder.Device(predicate->assigned_device_name()); + builder.Device(switch_cluster.predicate_edge->src()->assigned_device_name()); // Conditional should be the first input ... - builder.Input( - NodeDefBuilder::NodeOut(predicate->name(), 0, predicate->output_type(0))); + builder.Input(NodeDefBuilder::NodeOut( + switch_cluster.predicate_edge->src()->name(), + switch_cluster.predicate_edge->src_output(), + switch_cluster.predicate_edge->src()->output_type(0))); // ... followed by the other inputs. builder.Input(inputs); @@ -1019,7 +1233,7 @@ StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( } Status FunctionalizeCond::ExtractBody(const CondArgNodes& cond_arg_nodes, - const std::vector& switch_nodes, + const std::vector& switches, const std::vector& merge_nodes, int input_edge, Graph* body) { VLOG(2) << "ExtractBody for " << NodesToString(merge_nodes) << " along edge " @@ -1029,7 +1243,7 @@ Status FunctionalizeCond::ExtractBody(const CondArgNodes& cond_arg_nodes, int arg_count = 0; for (auto& kv : cond_arg_nodes) { Node* arg_node = nullptr; - for (const auto* arg : kv.switch_nodes) { + 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++)); @@ -1053,8 +1267,7 @@ Status FunctionalizeCond::ExtractBody(const CondArgNodes& cond_arg_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 { @@ -1070,24 +1283,17 @@ Status FunctionalizeCond::ExtractBody(const CondArgNodes& cond_arg_nodes, } Status FunctionalizeCond::AddInputEdges(const CondArgNodes& cond_arg_nodes, - Node* predicate, Node* if_node) { + const Edge* predicate_edge, + Node* if_node) { VLOG(3) << "AddInputEdges for " << if_node->name(); int index = 0; - graph_->AddEdge(predicate, 0, if_node, index++); - for (auto& kv : cond_arg_nodes) { - bool inserted = false; - for (const Node* arg : kv.switch_nodes) { - 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; - } - } + graph_->AddEdge(predicate_edge->src(), predicate_edge->src_output(), if_node, + index++); + for (auto& arg : cond_arg_nodes) { + if (arg.src_output == Graph::kControlSlot) { + graph_->AddControlEdge(arg.src, if_node); + } else { + graph_->AddEdge(arg.src, arg.src_output, if_node, index++); } } return Status::OK(); @@ -1108,10 +1314,10 @@ Status FunctionalizeCond::AddOutputEdges(const std::vector& outputs, return errors::Unimplemented("Output of index (", edge->src_output(), ") of merge node ", node->name()); } - graph_->RemoveEdge(edge); int src_output = dst_input == Graph::kControlSlot ? Graph::kControlSlot : i; + graph_->RemoveEdge(edge); graph_->AddEdge(if_node, src_output, dst, dst_input); } } @@ -1119,16 +1325,17 @@ Status FunctionalizeCond::AddOutputEdges(const std::vector& outputs, } StatusOr FunctionalizeCond::ConvertToXlaIf( - const CondArgNodes& cond_arg_nodes, const std::vector& switch_nodes, - const std::vector& merge_nodes, Node* predicate) { - VLOG(1) << "ConvertToXlaIf for " << NodesToString(switch_nodes) << " -> " + 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(cond_arg_nodes, switch_nodes, merge_nodes, predicate)); - TF_RETURN_IF_ERROR(AddInputEdges(cond_arg_nodes, predicate, if_node)); + BuildAndAddXlaIfOp(cond_arg_nodes, switch_cluster, merge_nodes)); + TF_RETURN_IF_ERROR( + AddInputEdges(cond_arg_nodes, switch_cluster.predicate_edge, if_node)); TF_RETURN_IF_ERROR(AddOutputEdges(merge_nodes, if_node)); return if_node; @@ -1137,18 +1344,20 @@ StatusOr FunctionalizeCond::ConvertToXlaIf( 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. @@ -1160,7 +1369,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() : "---"); @@ -1228,7 +1438,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 44f7db5ffdfa08d27c6d0d56b3368fcb5f6e3a47..20179b67991d3d23d678cf1df2642e029ea037fd 100644 --- a/tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md +++ b/tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md @@ -3,6 +3,7 @@ Operator | Type Constraint ------------------------------------- | --------------- `Abs` | `T={double,float,int32,int64}` +`Acos` | `T={complex64,double,float,int32,int64}` `Acosh` | `T={complex64,double,float}` `Add` | `T={complex64,double,float,int32,int64}` `AddN` | `T={complex64,double,float,int32,int64,uint32,uint64}` @@ -15,10 +16,12 @@ Operator | Type Constraint `ApproximateEqual` | `T={complex64,double,float,int32,int64,uint32,uint64}` `ArgMax` | `Tidx={int32,int64}`
`output_type={int32,int64}`
`T={float}` `ArgMin` | `Tidx={int32,int64}`
`output_type={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` +`Asin` | `T={complex64,double,float,int32,int64}` `Asinh` | `T={complex64,double,float}` `AssignAddVariableOp` | `dtype={complex64,double,float,int32,int64,uint32,uint64}` `AssignSubVariableOp` | `dtype={complex64,double,float,int32,int64,uint32,uint64}` `AssignVariableOp` | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}` +`Atan` | `T={complex64,double,float,int32,int64}` `Atan2` | `T={double,float}` `Atanh` | `T={complex64,double,float}` `AvgPool` | `T={double,float}` @@ -71,9 +74,14 @@ 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}` +`ExtractImagePatches` | `T={double,float,int32,int64,uint32,uint64}` `FFT` | `FFT2D` | `FFT3D` | +`FakeQuantWithMinMaxArgs` | +`FakeQuantWithMinMaxArgsGradient` | +`FakeQuantWithMinMaxVars` | +`FakeQuantWithMinMaxVarsGradient` | `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}` @@ -83,6 +91,7 @@ Operator | Type Constraint `FusedBatchNormGradV2` | `U={float}`
`T={float}` `FusedBatchNormV2` | `U={float}`
`T={float}` `Gather` | `Tindices={int32,int64}`
`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}` +`GatherNd` | `Tindices={int32,int64}`
`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}` `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}` @@ -116,14 +125,20 @@ Operator | Type Constraint `LogicalNot` | `LogicalOr` | `MatMul` | `T={complex64,double,float}` +`MatrixBandPart` | `Tindex={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `MatrixDiag` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `MatrixDiagPart` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`MatrixSetDiag` | `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}` +`MaxPoolGradGrad` | `T={float}` +`MaxPoolGradGradV2` | `T={float}` +`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}` @@ -176,12 +191,14 @@ 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}` `Round` | `T={complex64,double,float,int32,int64}` `Rsqrt` | `T={complex64,double,float}` `RsqrtGrad` | `T={complex64,double,float}` +`ScatterNd` | `Tindices={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Select` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Selu` | `T={double,float}` `SeluGrad` | `T={double,float}` @@ -194,6 +211,7 @@ Operator | Type Constraint `Sinh` | `T={complex64,double,float}` `Size` | `out_type={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Slice` | `Index={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`Snapshot` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Softmax` | `T={double,float}` `SoftmaxCrossEntropyWithLogits` | `T={double,float}` `Softplus` | `T={double,float,int32,int64,uint32,uint64}` diff --git a/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md b/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md index eb1f891125383a2138b27367c350a3d4e16a060a..55f0538dba7c1941dfea88e0631cd299e51f76d0 100644 --- a/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md +++ b/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md @@ -3,6 +3,7 @@ Operator | Type Constraint ------------------------------------- | --------------- `Abs` | `T={double,float,int32,int64}` +`Acos` | `T={complex64,double,float,int32,int64}` `Acosh` | `T={complex64,double,float}` `Add` | `T={complex64,double,float,int32,int64}` `AddN` | `T={complex64,double,float,int32,int64,uint32,uint64}` @@ -15,10 +16,12 @@ Operator | Type Constraint `ApproximateEqual` | `T={complex64,double,float,int32,int64,uint32,uint64}` `ArgMax` | `Tidx={int32,int64}`
`output_type={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `ArgMin` | `Tidx={int32,int64}`
`output_type={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` +`Asin` | `T={complex64,double,float,int32,int64}` `Asinh` | `T={complex64,double,float}` `AssignAddVariableOp` | `dtype={complex64,double,float,int32,int64,uint32,uint64}` `AssignSubVariableOp` | `dtype={complex64,double,float,int32,int64,uint32,uint64}` `AssignVariableOp` | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}` +`Atan` | `T={complex64,double,float,int32,int64}` `Atan2` | `T={double,float}` `Atanh` | `T={complex64,double,float}` `AvgPool` | `T={double,float}` @@ -71,9 +74,14 @@ 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}` +`ExtractImagePatches` | `T={double,float,int32,int64,uint32,uint64}` `FFT` | `FFT2D` | `FFT3D` | +`FakeQuantWithMinMaxArgs` | +`FakeQuantWithMinMaxArgsGradient` | +`FakeQuantWithMinMaxVars` | +`FakeQuantWithMinMaxVarsGradient` | `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}` @@ -83,6 +91,7 @@ Operator | Type Constraint `FusedBatchNormGradV2` | `U={float}`
`T={float}` `FusedBatchNormV2` | `U={float}`
`T={float}` `Gather` | `Tindices={int32,int64}`
`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}` +`GatherNd` | `Tindices={int32,int64}`
`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}` `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}` @@ -116,14 +125,20 @@ Operator | Type Constraint `LogicalNot` | `LogicalOr` | `MatMul` | `T={complex64,double,float}` +`MatrixBandPart` | `Tindex={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `MatrixDiag` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `MatrixDiagPart` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`MatrixSetDiag` | `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}` +`MaxPoolGradGrad` | `T={float}` +`MaxPoolGradGradV2` | `T={float}` +`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}` @@ -173,12 +188,14 @@ 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}` `Round` | `T={complex64,double,float,int32,int64}` `Rsqrt` | `T={complex64,double,float}` `RsqrtGrad` | `T={complex64,double,float}` +`ScatterNd` | `Tindices={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Select` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Selu` | `T={double,float}` `SeluGrad` | `T={double,float}` @@ -191,6 +208,7 @@ Operator | Type Constraint `Sinh` | `T={complex64,double,float}` `Size` | `out_type={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Slice` | `Index={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`Snapshot` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Softmax` | `T={double,float}` `SoftmaxCrossEntropyWithLogits` | `T={double,float}` `Softplus` | `T={double,float,int32,int64,uint32,uint64}` diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index 02215b5112d37f726604da2c2caa4f804388d6e5..b20c1ffc7d8956f3f5530ee63e9b711a26439be5 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -60,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( @@ -132,11 +130,11 @@ Status GraphCompiler::Compile() { // Set up inputs from outputs of previous nodes. for (auto* e : n->in_edges()) { if (e->IsControlEdge()) continue; - Node* src = e->src(); + const Node* src = e->src(); 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()); 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/host_compute_metadata.proto b/tensorflow/compiler/tf2xla/host_compute_metadata.proto new file mode 100644 index 0000000000000000000000000000000000000000..43ab371a217e6c4521a160715104c96e3c8782c6 --- /dev/null +++ b/tensorflow/compiler/tf2xla/host_compute_metadata.proto @@ -0,0 +1,38 @@ +syntax = "proto3"; + +package tensorflow.tf2xla; +option cc_enable_arenas = true; +option java_outer_classname = "Tf2XlaProtos"; +option java_multiple_files = true; +option java_package = "org.tensorflow.tf2xla"; + +import "tensorflow/core/framework/tensor_shape.proto"; +import "tensorflow/core/framework/types.proto"; + +// TensorMetadata indicates the type and shape of a Tensor that is +// part of a host compute transfer. +message TensorMetadata { + DataType type = 1; + TensorShapeProto shape = 2; +} + +// HostTransferMetadata describes a transfer either from host to device +// or device to host. It has a key that is unique to the computation, +// and metadata about the list of tensors being transferred. +message HostTransferMetadata { + // The key used to identify this transfer. + string key = 1; + + // For each Tensor being transferred, its type and shape. + repeated TensorMetadata metadata = 2; +} + +// HostComputeMetadata describes all the sends and recvs +// from all host compute transfer ops in a computation. +message HostComputeMetadata { + // Metadata about each device_to_host transfer + repeated HostTransferMetadata device_to_host = 1; + + // Metadata about each host_to_device transfer + repeated HostTransferMetadata host_to_device = 2; +} diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index 84fa43f4fb9897120afe63a6b6b1595d2dddd19f..d2fa933cf9c085f92b2f442827a94d72938e4bb2 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -31,6 +31,8 @@ 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", @@ -43,6 +45,8 @@ 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", @@ -61,6 +65,7 @@ tf_kernel_library( "reverse_op.cc", "reverse_sequence_op.cc", "scan_ops.cc", + "scatter_nd_op.cc", "segment_reduction_ops.cc", "select_op.cc", "sendrecv_ops.cc", @@ -84,7 +89,6 @@ tf_kernel_library( "variable_ops.cc", ], hdrs = [ - "gather_op.h", "index_ops.h", "shape_util.h", ], @@ -94,12 +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", diff --git a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc index 344a2ab2b6835c518c41de6f7a30fb2a34d130d2..cbade79e85eed10ecb5ead7151ee778c86a0de37 100644 --- a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc @@ -159,7 +159,9 @@ class BatchToSpaceNDOp : public XlaOpKernel { block_shape, crops); } }; -REGISTER_XLA_OP(Name("BatchToSpaceND").CompileTimeConstInput("crops"), +REGISTER_XLA_OP(Name("BatchToSpaceND") + .CompileTimeConstInput("block_shape") + .CompileTimeConstInput("crops"), BatchToSpaceNDOp); class BatchToSpaceOp : public XlaOpKernel { 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/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index ffed38249416766850ba10f1069e706570b995fe..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,25 +26,38 @@ 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); + for (int64 i = 0; i < num_index_dims; ++i) { + slice_shape.set_dim(axis + i, 1); + } TensorShape loop_out_shape; loop_out_shape.AppendShape(input_shape_pre_axis); @@ -62,131 +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, which 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 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()] = 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/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/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc index d4fb5dd4e06c7c70591262c0d63a91c383a2a6e0..086a9491aa93ebfae99f296dd355ae2e322084ec 100644 --- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc @@ -525,5 +525,172 @@ class AvgPool3DGradOp : public AvgPoolGradOp { REGISTER_XLA_OP(Name("AvgPool3DGrad").CompileTimeConstInput("orig_input_shape"), AvgPool3DGradOp); +class MaxPoolGradGradOp : public XlaOpKernel { + public: + MaxPoolGradGradOp(OpKernelConstruction* ctx, int num_spatial_dims) + : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) { + 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(ctx, stride_.size() == num_dims(), + errors::InvalidArgument("Sliding window strides field must " + "specify ", + num_dims(), " dimensions")); + + const TensorShape tensor_in_shape = ctx->InputShape(0); + const TensorShape tensor_out_shape = ctx->InputShape(1); + const TensorShape out_backprop_shape = ctx->InputShape(2); + + // For maxpooling, tensor_in should have num_dims() dimensions. + OP_REQUIRES(ctx, tensor_in_shape.dims() == num_dims(), + errors::InvalidArgument("tensor_in must be ", num_dims(), + "-dimensional")); + OP_REQUIRES(ctx, tensor_out_shape.dims() == num_dims(), + errors::InvalidArgument("tensor_out must be ", num_dims(), + "-dimensional")); + // For maxpooling, out_backprop should have num_dims() dimensions. + OP_REQUIRES(ctx, out_backprop_shape.dims() == num_dims(), + errors::InvalidArgument("out_backprop must be ", num_dims(), + "-dimensional")); + + // What we want to compute: + // Given y = MaxPool(x), and xs_grad = MaxPoolGrad(x, y, ys_grad) + // MaxPoolGradGrad computes {ys_grad}_grad given x, y, and {xs_grad}_grad. + // + // In the regular TF op, this amounts to selecting for each window the + // incoming backprop value from xs_grad_grad that corresponds to the maximal + // value in the corresponding window of x. + // + // TODO(b/73062247): What we really want is a ReduceWindow with different + // arrays for index selection vs return value selection--a select-to-gather. + // + // Here, we implement a bitwise hack: we use the hi 16 bits of input for + // separate max pooling alongside each of the hi and lo 16 bits of + // out_backprop packed into 16 lo bits, which we then glue back together at + // the end to get a full 32 bits of gradient. + // + // This could select the wrong backprop value for two x values that are + // equally maximal up to the first 16 bits, in which case we are taking the + // latter. + // + // Note that in principle we could use 32 separate maxpools to recover each + // of 32 bits of the gradient while preserving 31 bits of input for the max + // pooling criteria; here, we just truncate to the first 16 bits of input. + + auto input = ctx->Input(0); + auto out_backprop = ctx->Input(2); + + auto b = ctx->builder(); + + auto sixteen = b->ConstantR0(16); + // in (f32) -> round to bf16 -> f32 for correct bitwidth -> 16-high-bit u32 + auto in_hi = b->BitcastConvertType( + b->ConvertElementType(b->ConvertElementType(input, xla::BF16), + xla::F32), + xla::U32); + auto bp_int = b->BitcastConvertType(out_backprop, xla::U32); + auto bp_hi = b->ShiftRightLogical(bp_int, sixteen); + auto bp_lo = b->ShiftRightLogical(b->ShiftLeft(bp_int, sixteen), sixteen); + auto in_hi_bp_hi = b->Add(in_hi, bp_hi); // Want an unsigned add. + auto in_hi_bp_lo = b->Add(in_hi, bp_lo); // Want an unsigned add. + + auto init_value = XlaHelpers::MinValue(b, DT_FLOAT); + // We will reduce by taking the maximal value up to 16 bits (ignoring the lo + // 16 bits of packed-in hi/lo backprop value). + auto rb = b->CreateSubBuilder("GreaterOrEqOf_ByFirst16Bits"); + { + // F32 parameters to satisfy lowering type restriction for reduce opcode. + const xla::Shape scalar = xla::ShapeUtil::MakeShape(xla::F32, {}); + auto lhs = rb->Parameter(0, scalar, "lhs"); + auto rhs = rb->Parameter(1, scalar, "rhs"); + auto sixteen = rb->ConstantR0(16); + auto lhs_criteria = rb->ShiftLeft( + rb->ShiftRightLogical(rb->BitcastConvertType(lhs, xla::S32), sixteen), + sixteen); + auto rhs_criteria = rb->ShiftLeft( + rb->ShiftRightLogical(rb->BitcastConvertType(rhs, xla::S32), sixteen), + sixteen); + // Must use a F32 comparison, because S32 would not work for negatives. + rb->Select(rb->Ge(rb->BitcastConvertType(lhs_criteria, xla::F32), + rb->BitcastConvertType(rhs_criteria, xla::F32)), + lhs, rhs); + } + auto reduce = rb->BuildAndNoteError(); + xla::Padding xla_padding = + (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame; + auto pooled_hi = + b->ReduceWindow(b->BitcastConvertType(in_hi_bp_hi, xla::F32), + init_value, reduce, ksize_, stride_, xla_padding); + auto pooled_lo = + b->ReduceWindow(b->BitcastConvertType(in_hi_bp_lo, xla::F32), + init_value, reduce, ksize_, stride_, xla_padding); + auto grads_hi = + b->ShiftLeft(b->BitcastConvertType(pooled_hi, xla::U32), sixteen); + auto grads_lo = b->ShiftRightLogical( + b->ShiftLeft(b->BitcastConvertType(pooled_lo, xla::U32), sixteen), + sixteen); + auto grads = b->Add(grads_hi, grads_lo); // Want an unsigned add. + + xla::PrimitiveType element_type; + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(2), &element_type)); + ctx->SetOutput(0, b->BitcastConvertType(grads, element_type)); + } + + protected: + const int num_spatial_dims_; + std::vector ksize_; + std::vector stride_; + Padding padding_; + TensorFormat data_format_ = FORMAT_NHWC; +}; + +class MaxPool2DGradGradOp : public MaxPoolGradGradOp { + public: + explicit MaxPool2DGradGradOp(OpKernelConstruction* ctx) + : MaxPoolGradGradOp(ctx, /*num_spatial_dims=*/2) { + string data_format; + OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format)); + OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_), + errors::InvalidArgument("Invalid data format")); + } +}; +REGISTER_XLA_OP(Name("MaxPoolGradGrad").TypeConstraint("T", DT_FLOAT), + MaxPool2DGradGradOp); +REGISTER_XLA_OP(Name("MaxPoolGradGradV2") + .TypeConstraint("T", DT_FLOAT) + .CompileTimeConstInput("ksize") + .CompileTimeConstInput("strides"), + MaxPool2DGradGradOp); + } // anonymous 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 c220edd588071ef262621784015d34cd475b2918..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,125 +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); - - // 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); - TensorShape loop_out_slice_shape(out_shape); - loop_out_slice_shape.set_dim(0, 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), - 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, 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(); - 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. - std::vector out_index_vals(out_shape.dims(), - zero); - out_index_vals[0] = 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(); - auto body = body_status.ConsumeValueOrDie(); - - auto gather_while = builder->While(cond, body, init); - return builder->GetTupleElement(gather_while, 3); -} - namespace { class UnsortedSegmentSum : public XlaOpKernel { @@ -153,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)); @@ -174,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/split_op.cc b/tensorflow/compiler/tf2xla/kernels/split_op.cc index 79c435c90a1f57250be90c2c2523bf3d7d231461..43c15e753805352875034dfd2c70a2a1ed9a4114 100644 --- a/tensorflow/compiler/tf2xla/kernels/split_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/split_op.cc @@ -111,27 +111,24 @@ class SplitVOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { const int32 num_split = num_outputs(); + const TensorShape input_shape = ctx->InputShape(0); const TensorShape index_shape = ctx->InputShape(2); - xla::Literal literal_index; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(2, &literal_index)); - int32 split_dim; - OP_REQUIRES(ctx, index_shape.dims() == 0, - errors::InvalidArgument("split_dim input to Split Op must be a " - "scalar")); - split_dim = literal_index.Get({}); + int64 split_dim_orig; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(2, &split_dim_orig)); + int64 split_dim = split_dim_orig < 0 ? split_dim_orig + input_shape.dims() + : split_dim_orig; + OP_REQUIRES(ctx, 0 <= split_dim && split_dim < input_shape.dims(), + errors::InvalidArgument("-input rank(-", input_shape.dims(), + ") <= split_dim < input rank (", + input_shape.dims(), "), but got ", + split_dim_orig)); xla::ComputationDataHandle input = ctx->Input(0); - const TensorShape input_shape = ctx->InputShape(0); OP_REQUIRES(ctx, input_shape.dims() > 0, errors::InvalidArgument("Can't split a 0 dimensional input")); - OP_REQUIRES( - ctx, 0 <= split_dim && split_dim < input_shape.dims(), - errors::InvalidArgument("0 <= split_dim < number of input dimensions (", - input_shape.dims(), "), but got ", split_dim)); - OP_REQUIRES( ctx, num_split > 0, errors::InvalidArgument( diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc index d77fb768ef4d124c403a1dc9b321c4f29571d806..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; @@ -119,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); } @@ -164,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); } @@ -208,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..6204aa4e27000fddec7f5b82b2198d37956f6aba 100644 --- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc @@ -77,13 +77,14 @@ class StridedSliceOp : public XlaOpKernel { for (int i = 0; i < begin.size(); ++i) { if (strides[i] > 0) { slice_begin.push_back(begin[i]); - slice_end.push_back(end[i]); + slice_end.push_back(std::max(end[i], begin[i])); slice_strides.push_back(strides[i]); } else { // Negative stride: swap begin and end, add 1 because the interval // is semi-open, and mark the dimension to be reversed. slice_begin.push_back(input_shape.dim_size(i) - begin[i] - 1); - slice_end.push_back(input_shape.dim_size(i) - end[i] - 1); + slice_end.push_back(std::max(input_shape.dim_size(i) - end[i] - 1, + input_shape.dim_size(i) - begin[i] - 1)); slice_strides.push_back(-strides[i]); dimensions_to_reverse.push_back(i); } @@ -231,6 +232,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 +254,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 +284,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 +319,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/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index a266e9013c41b88788dbc99849f01c09f3d61348..7cb47f908d4ff43f455f1e77c53cd3cc956579ee 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -50,18 +50,43 @@ 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)) +// = log(x + sqrt((x+1)*(x-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))))); + b->Log(b->Add(x, + b->Pow(b->Mul(b->Add(x, XlaHelpers::One(b, input_type(0))), + b->Sub(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 4a711e4d9b7aedb166a8a0ec9fe9ec2390f01b17..0ff1b65ae9179d506e453f98097cd88083eb2be7 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -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(); @@ -154,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 @@ -182,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"; @@ -292,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 d184f59e01788829d0ba97092c14d36e5188e4e8..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"], @@ -107,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/scatter.cc b/tensorflow/compiler/tf2xla/lib/scatter.cc new file mode 100644 index 0000000000000000000000000000000000000000..45699233ea8b2a75e3850098250307b95546cc28 --- /dev/null +++ b/tensorflow/compiler/tf2xla/lib/scatter.cc @@ -0,0 +1,192 @@ +/* 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)); + + // Make the index in bounds to prevent implementation defined behavior. + index = body_builder->Max(index, zero_index); + 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. + auto current_value = + body_builder->DynamicSlice(buffer, index, updates_slice_shape); + if (combiner) { + update = combiner(current_value, update, body_builder); + } + // Use the current value instead of the update if the index is out of + // bounds. + update = body_builder->Select(index_in_range, update, current_value); + // Apply the update. + buffer = body_builder->DynamicUpdateSlice(buffer, update, index); + + 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/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc index 9b7492f8cf6e86498d7e2f5d42e42ea978c664d8..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) { diff --git a/tensorflow/compiler/tf2xla/lib/util.h b/tensorflow/compiler/tf2xla/lib/util.h index 7f93102ee78bec60018814975a0badfeb7874aa6..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, 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 fcbd157c6191655865d5e250fdf71338780bc2a6..2c3cd658e0462368ac0b51938979b7a6815a7574 100644 --- a/tensorflow/compiler/tf2xla/literal_util.cc +++ b/tensorflow/compiler/tf2xla/literal_util.cc @@ -40,20 +40,20 @@ Status HostTensorToLiteral(const Tensor& host_tensor, xla::Literal* literal) { 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.untyped_data(); @@ -63,4 +63,12 @@ Status LiteralToHostTensor(const xla::Literal& literal, DataType target_type, 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_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index 69b265436bb19bbbdd9deb872f4097d4bac7ea52..7cdf4d1b3e03492f20629422887b19de30e7c0d1 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,23 +320,133 @@ Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr graph, return Status::OK(); } +// Builds the XLA computation. +// +// `retvals` is the list of retvals produced by _Retval operators, in index +// order. `variable_map` is a map from variable ID numbers to XlaOpContext +// variable states, generated by the symbolic evaluation. +// If `return_updated_values_for_all_resources` is true, all resources will be +// included in `resource_updates`, regardless of whether their value changed. +// Sets `*num_nonconst_outputs` to the number of outputs of the `computation`. +// Sets `*resource_updates` to a description of resources whose values are +// written by the computation; the variable writes are the last +// `resource_updates.size()` return values from the computation. Each entry in +// `resource_updates` is a (input_index, type) pair, where `input_index` is the +// index of a resource variable argument to the computation, and `type` is the +// type of the final output. +Status BuildComputation( + const std::vector& args, + const std::vector& arg_cores, + const std::vector& retvals, + const std::vector>& resources, + bool return_updated_values_for_all_resources, + xla::ComputationBuilder* builder, xla::Computation* computation, + int* num_computation_outputs, int* num_nonconst_outputs, + std::vector* resource_updates) { + std::vector elems; + elems.reserve(retvals.size()); + for (const XlaExpression& retval : retvals) { + if (!retval.has_constant_value()) { + elems.push_back(retval.handle()); + } + } + *num_nonconst_outputs = elems.size(); + + // Add return values for resources whose values have changed. + std::vector arg_resources; + arg_resources.reserve(resources.size()); + for (const auto& resource : resources) { + if (resource->arg_num() >= 0) { + arg_resources.push_back(resource.get()); + } + } + std::sort(arg_resources.begin(), arg_resources.end(), + [](const XlaResource* a, const XlaResource* b) { + return a->arg_num() < b->arg_num(); + }); + + // Attach a common operator name as metadata. This has no semantic effect — it + // merely makes the HLO graph more readable when visualized via TensorBoard, + // since TensorBoard forms groups out of operators with similar names. + xla::OpMetadata retval_metadata; + retval_metadata.set_op_name("XLA_Retvals"); + builder->SetOpMetadata(retval_metadata); + + for (const XlaResource* resource : arg_resources) { + const XlaCompiler::Argument& arg = args[resource->arg_num()]; + const int core = arg_cores[resource->arg_num()]; + DCHECK_LT(resource->arg_num(), arg_cores.size()); + bool modified = + resource->value().handle() != resource->initial_value().handle(); + // TensorArray gradients were modified if their values changed or there are + // any newly created gradients. + for (const auto& grad : resource->tensor_array_gradients()) { + modified = modified || + grad.second->value().handle() != + grad.second->initial_value().handle() || + arg.tensor_array_gradients.count(grad.first) == 0; + } + if (return_updated_values_for_all_resources || modified) { + resource_updates->emplace_back(); + 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); + } + + // Request that the value be returned on a specific core. + xla::ScopedShardingAssignment assign_sharding( + builder, core == -1 ? tensorflow::gtl::optional() + : xla::sharding_builder::AssignDevice(core)); + + xla::ComputationDataHandle handle; + TF_RETURN_IF_ERROR(resource->Pack(&handle, builder)); + + // Since we can't change the sharding metadata of as this point, + // create a tuple/get-tuple-element combination so that sharding + // assignment will be placed on this value, which will cause the resource + // update to be returned from the same device that provided the resource. + handle = builder->GetTupleElement(builder->Tuple({handle}), 0); + + elems.push_back(handle); + } + } + + *num_computation_outputs = elems.size(); + + // Builds the XLA computation. + builder->Tuple(elems); + builder->ClearOpMetadata(); + + xla::StatusOr computation_status = builder->Build(); + if (!computation_status.ok()) { + return computation_status.status(); + } + *computation = computation_status.ConsumeValueOrDie(); + return Status::OK(); +} + +} // namespace + // 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) { +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. - std::vector parameters, resources; - parameters.reserve(args.size()); + 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. @@ -290,18 +460,20 @@ Status BuildArguments(const Graph& graph, // 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); + 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: - parameters.push_back(i); + case XlaCompiler::Argument::kParameter: { + input_mapping->push_back(i); break; + } case XlaCompiler::Argument::kConstant: arg_expression.set_constant_value(arg.constant_value); break; @@ -312,19 +484,17 @@ Status BuildArguments(const Graph& graph, // Append parameters containing variable values after the other runtime // parameters. - parameters.insert(parameters.end(), resources.begin(), resources.end()); - if (parameters.empty()) { + input_mapping->insert(input_mapping->end(), resources.begin(), + resources.end()); + if (input_mapping->empty()) { return Status::OK(); } - std::vector arg_shapes; - arg_shapes.reserve(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]]; + 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. - arg_shapes.push_back(arg.shape); - (*input_mapping)[i] = parameters[i]; + TF_RETURN_IF_ERROR( + XLAShapeForArgument(args[(*input_mapping)[i]], &arg_shapes[i])); } if (use_tuple_arg) { @@ -353,14 +523,21 @@ Status BuildArguments(const Graph& graph, } } + // Attach a common operator name as metadata. This has no semantic effect — it + // merely makes the HLO graph more readable when visualized via TensorBoard, + // since TensorBoard forms groups out of operators with similar names. + xla::OpMetadata arg_metadata; + arg_metadata.set_op_name("XLA_Args"); + builder->SetOpMetadata(arg_metadata); + // Build parameter handles for non-constant arguments. - std::vector arg_handles(parameters.size()); + 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 : parameters) { + for (int64 parameter : *input_mapping) { const int core = (*arg_cores)[parameter]; const int root_device = 0; *tuple_sharding.add_tuple_shardings() = @@ -373,16 +550,16 @@ Status BuildArguments(const Graph& graph, } else { tuple = builder->Parameter(0, (*input_shapes)[0], "arg_tuple"); } - for (std::vector::size_type i = 0; i < parameters.size(); ++i) { - const int core = (*arg_cores)[parameters[i]]; + 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 < parameters.size(); ++i) { - const int core = (*arg_cores)[parameters[i]]; + 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)); @@ -391,21 +568,22 @@ Status BuildArguments(const Graph& graph, } } + builder->ClearOpMetadata(); + // 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]]; + 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 " << parameters[i]; - XlaExpression& arg_expression = (*arg_expressions)[parameters[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], - /*reset_initial_values=*/true, builder)); + 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; @@ -422,107 +600,6 @@ Status BuildArguments(const Graph& graph, return Status::OK(); } -// Builds the XLA computation. -// -// `retvals` is the list of retvals produced by _Retval operators, in index -// order. `variable_map` is a map from variable ID numbers to XlaOpContext -// variable states, generated by the symbolic evaluation. -// If `return_updated_values_for_all_resources` is true, all resources will be -// included in `resource_updates`, regardless of whether their value changed. -// Sets `*num_nonconst_outputs` to the number of outputs of the `computation`. -// Sets `*resource_updates` to a description of resources whose values are -// written by the computation; the variable writes are the last -// `resource_updates.size()` return values from the computation. Each entry in -// `resource_updates` is a (input_index, type) pair, where `input_index` is the -// index of a resource variable argument to the computation, and `type` is the -// type of the final output. -Status BuildComputation( - const std::vector& args, - const std::vector& arg_cores, - const std::vector& retvals, - const std::vector>& resources, - bool return_updated_values_for_all_resources, - xla::ComputationBuilder* builder, xla::Computation* computation, - int* num_computation_outputs, int* num_nonconst_outputs, - std::vector* resource_updates) { - std::vector elems; - elems.reserve(retvals.size()); - for (const XlaExpression& retval : retvals) { - if (!retval.has_constant_value()) { - elems.push_back(retval.handle()); - } - } - *num_nonconst_outputs = elems.size(); - - // Add return values for resources whose values have changed. - std::vector arg_resources; - arg_resources.reserve(resources.size()); - for (const auto& resource : resources) { - if (resource->arg_num() >= 0) { - arg_resources.push_back(resource.get()); - } - } - std::sort(arg_resources.begin(), arg_resources.end(), - [](const XlaResource* a, const XlaResource* b) { - return a->arg_num() < b->arg_num(); - }); - - for (const XlaResource* resource : arg_resources) { - const XlaCompiler::Argument& arg = args[resource->arg_num()]; - const int core = arg_cores[resource->arg_num()]; - DCHECK_LT(resource->arg_num(), arg_cores.size()); - bool modified = - resource->value().handle() != resource->initial_value().handle(); - // TensorArray gradients were modified if their values changed or there are - // any newly created gradients. - for (const auto& grad : resource->tensor_array_gradients()) { - modified = modified || - grad.second->value().handle() != - grad.second->initial_value().handle() || - arg.tensor_array_gradients.count(grad.first) == 0; - } - if (return_updated_values_for_all_resources || modified) { - resource_updates->emplace_back(); - XlaCompiler::ResourceUpdate& update = resource_updates->back(); - update.input_index = resource->arg_num(); - update.type = resource->type(); - update.modified = modified; - for (const auto& grad : resource->tensor_array_gradients()) { - update.tensor_array_gradients_accessed.insert(grad.first); - } - - // Request that the value be returned on a specific core. - xla::ScopedShardingAssignment assign_sharding( - builder, core == -1 ? tensorflow::gtl::optional() - : xla::sharding_builder::AssignDevice(core)); - - xla::ComputationDataHandle handle; - TF_RETURN_IF_ERROR(resource->Pack(&handle, builder)); - - // Since we can't change the sharding metadata of as this point, - // create a tuple/get-tuple-element combination so that sharding - // assignment will be placed on this value, which will cause the resource - // update to be returned from the same device that provided the resource. - handle = builder->GetTupleElement(builder->Tuple({handle}), 0); - - elems.push_back(handle); - } - } - - *num_computation_outputs = elems.size(); - - // Builds the XLA computation. - builder->Tuple(elems); - xla::StatusOr computation_status = builder->Build(); - if (!computation_status.ok()) { - return computation_status.status(); - } - *computation = computation_status.ConsumeValueOrDie(); - return Status::OK(); -} - -} // namespace - Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, string const& name, std::unique_ptr graph, @@ -547,7 +624,8 @@ 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; @@ -596,6 +674,14 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, VLOG(2) << "XLA output shape: " << xla::ShapeUtil::HumanString(result->xla_output_shape); + // Copy the host transfer metadata to the result. + for (const auto& send : host_compute_sends_) { + *result->host_compute_metadata.add_device_to_host() = send.second; + } + for (const auto& recv : host_compute_recvs_) { + *result->host_compute_metadata.add_host_to_device() = recv.second; + } + // Tensorflow expects a major-to-minor order of results. xla::LayoutUtil::SetToDefaultLayout(&result->xla_output_shape); @@ -616,13 +702,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(); } @@ -637,4 +716,59 @@ Status XlaCompiler::GetChannelHandle(const string& key, return Status::OK(); } +namespace { + +void SetTransfer(const string& key, gtl::ArraySlice types, + gtl::ArraySlice shapes, + tf2xla::HostTransferMetadata* transfer) { + transfer->set_key(key); + CHECK(types.size() == shapes.size()); + for (int i = 0; i < types.size(); ++i) { + tf2xla::TensorMetadata* metadata = transfer->add_metadata(); + metadata->set_type(types[i]); + shapes[i].AsProto(metadata->mutable_shape()); + } +} + +} // namespace + +Status XlaCompiler::SetDeviceToHostMetadata( + const string& key, gtl::ArraySlice types, + gtl::ArraySlice shapes) { + if (host_compute_sends_.find(key) != host_compute_sends_.end()) { + return errors::InvalidArgument( + "Duplicate calls to SetDeviceToHostMetadata with key ", key); + } + tf2xla::HostTransferMetadata& transfer = host_compute_sends_[key]; + SetTransfer(key, types, shapes, &transfer); + return Status::OK(); +} + +Status XlaCompiler::GetDeviceToHostShapes( + const string& key, std::vector* shapes) const { + const auto iter = host_compute_sends_.find(key); + if (iter == host_compute_sends_.end()) { + return errors::InvalidArgument( + "No host compute send shapes registered for key ", key); + } + shapes->clear(); + for (int i = 0; i < iter->second.metadata_size(); ++i) { + TensorShape shape(iter->second.metadata(i).shape()); + shapes->push_back(shape); + } + return Status::OK(); +} + +Status XlaCompiler::SetHostToDeviceMetadata( + const string& key, gtl::ArraySlice types, + gtl::ArraySlice shapes) { + if (host_compute_recvs_.find(key) != host_compute_sends_.end()) { + return errors::InvalidArgument( + "Duplicate calls to SetHostToDeviceMetadata with key ", key); + } + tf2xla::HostTransferMetadata& transfer = host_compute_recvs_[key]; + SetTransfer(key, types, shapes, &transfer); + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index 30d3c05ee9aa33accc0ad122901f70b0b6613104..5f1c631976cff77421eb1bfeab83ffa8c660c730 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_COMPILER_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_COMPILER_H_ +#include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h" #include "tensorflow/compiler/tf2xla/xla_compilation_device.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/core/common_runtime/device.h" @@ -29,6 +30,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 +108,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. @@ -175,8 +187,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.) @@ -204,6 +217,10 @@ class XlaCompiler { // containing both constant and non-constant results. std::vector outputs; + // TensorFlow shapes and types of sends/recvs from HostCompute Ops to their + // matching RecvAtHost/SendFromHost Ops in the outer graph. + tf2xla::HostComputeMetadata host_compute_metadata; + // Resources whose values were updated by the computation, ordered // by return value position. Resource updates follow the non-constant // results in the outputs of XLA computation. @@ -230,6 +247,12 @@ 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 @@ -266,11 +289,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. @@ -279,6 +301,22 @@ class XlaCompiler { // same XlaCompiler. Status GetChannelHandle(const string& key, xla::ChannelHandle* channel); + // Sets the shapes and types for the device to host transfer associated with + // 'key'. + Status SetDeviceToHostMetadata(const string& key, + gtl::ArraySlice types, + gtl::ArraySlice shapes); + + // Gets the shapes the device to host transfer associated with 'key'. + Status GetDeviceToHostShapes(const string& key, + std::vector* shapes) const; + + // Sets the shapes and types for the host to device transfer associated with + // 'key'. + Status SetHostToDeviceMetadata(const string& key, + gtl::ArraySlice types, + gtl::ArraySlice shapes); + const Options& options() const { return options_; } xla::Client* client() const { return options_.client; } FunctionLibraryRuntime* flib_runtime() const { return flib_runtime_; } @@ -291,6 +329,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; @@ -331,6 +380,9 @@ class XlaCompiler { std::unordered_map channels_; + std::unordered_map host_compute_sends_; + std::unordered_map host_compute_recvs_; + TF_DISALLOW_COPY_AND_ASSIGN(XlaCompiler); }; diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 7ebe4b75bc1e33e506624314b11163e36a2477de..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()); @@ -242,10 +243,10 @@ TEST_F(XlaCompilerTest, HasSaneErrorOnNonCompileTimeConstantInputToReshape) { 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()); @@ -281,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); @@ -373,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(); @@ -420,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(); @@ -472,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"}; @@ -540,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"}; @@ -574,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"}; @@ -689,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 77e24162676045b88dc8b62d2c6a4ecc1e738e96..f048662953e20b2a612271e2daeef6e370c4822a 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.cc +++ b/tensorflow/compiler/tf2xla/xla_helpers.cc @@ -135,58 +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 = 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 b->ConstantLiteral(literal); + return ::tensorflow::IntegerLiteral(b, type, value); } xla::ComputationDataHandle XlaHelpers::FloatLiteral(xla::ComputationBuilder* b, diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc index ee0aed672e1b264fee0a7f381c334400c55f3581..c4bb90d58755f16672ca7c6a6738065be6330485 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc @@ -286,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(); @@ -296,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(); } @@ -312,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(); } @@ -396,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 = @@ -405,7 +418,24 @@ 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 { diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.h b/tensorflow/compiler/tf2xla/xla_op_kernel.h index 6d3b6db2289d6c0b8f266062f9f3baca1145154a..4e4b97e0cec8d16b9b5686a779b1285906765dbd 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.h +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.h @@ -164,13 +164,18 @@ 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(const Status& s); diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc index 0dde6a986c61bdd5b0b2e6d7a16b29ab95be98ab..bbe808595d958346bd55bf8419306bf3de4cd1d0 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc @@ -255,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 c22fd37129c5344825631ecf422bbcf3434e4534..cd13db4d300bb5bba21a734173b6afb9223539d8 100644 --- a/tensorflow/compiler/xla/BUILD +++ b/tensorflow/compiler/xla/BUILD @@ -52,6 +52,7 @@ xla_proto_library( visibility = ["//visibility:public"], deps = [ ":xla_data_proto", + "//tensorflow/compiler/xla/service:hlo_proto", "//tensorflow/compiler/xla/service:session_proto", ], ) @@ -88,7 +89,6 @@ cc_library( visibility = [":friends"], deps = [ "//tensorflow/core:framework_lite", - "//tensorflow/core:lib", "//third_party/eigen3", ], ) @@ -373,7 +373,6 @@ tf_cc_test( cc_library( name = "array2d", - srcs = ["array2d.cc"], hdrs = ["array2d.h"], visibility = ["//visibility:public"], deps = [ diff --git a/tensorflow/compiler/xla/array.h b/tensorflow/compiler/xla/array.h index 71aa057cd3a1c273c0e851497a78f94ba37c778e..24b58bec11bd8d8b5c79ac84c5f43c509644b51d 100644 --- a/tensorflow/compiler/xla/array.h +++ b/tensorflow/compiler/xla/array.h @@ -121,6 +121,44 @@ class Array { CHECK(idx == num_elements()); } + // Creates a 1D array of a floating-point type (half, bfloat16, float, + // or double) from an initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::value) && + std::is_same::value>::type> + Array(std::initializer_list values) + : Array(ToInt64Vector({values.size()})) { + int64 idx = 0; + for (const auto& it1 : values) { + values_[idx] = static_cast(it1); + ++idx; + } + CHECK(idx == num_elements()); + } + + // Creates a 2D array of a floating-point type (half, bfloat16, float, + // or double) from an initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::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 +176,30 @@ class Array { CHECK(idx == num_elements()); } + // Creates a 3D array of a floating-point type (half, bfloat16, float, + // or double) from an initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::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 +220,34 @@ class Array { CHECK(idx == num_elements()); } + // Creates a 4D array of a floating-point type (half, bfloat16, float, + // or double) from an initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::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 +275,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..a17e81f44832f272fd93dce9f854042b4a84fde4 100644 --- a/tensorflow/compiler/xla/array2d.h +++ b/tensorflow/compiler/xla/array2d.h @@ -25,6 +25,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/array.h" +#include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/bits.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -52,6 +53,17 @@ class Array2D : public Array { Array2D(std::initializer_list> values) : Array(values) {} + // Creates an array of a floating-point type (half, bfloat16, float, + // or double) from the given nested initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::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); } @@ -86,9 +98,23 @@ class Array2D : public Array { // Returns a linspace-populated Array2D in the range [from, to] (inclusive) // with dimensions n1 x n2. -std::unique_ptr> MakeLinspaceArray2D(float from, float to, - int64 n1, int64 n2); - +template +std::unique_ptr> MakeLinspaceArray2D(double from, double to, + int64 n1, int64 n2) { + auto array = MakeUnique>(n1, n2); + int64 count = n1 * n2; + NativeT step = + static_cast((count > 1) ? (to - from) / (count - 1) : 0); + auto set = [&array, n1, n2](int64 index, NativeT value) { + (*array)(index / n2, index % n2) = value; + }; + for (int64 i = 0; i < count - 1; ++i) { + set(i, (static_cast(from) + + static_cast(i) * static_cast(step))); + } + set(count - 1, static_cast(to)); + return array; +} } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_ARRAY2D_H_ 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..0e9a0722ae43e1dc6ecddde9cbc3daf1db058840 100644 --- a/tensorflow/compiler/xla/array3d.h +++ b/tensorflow/compiler/xla/array3d.h @@ -57,6 +57,19 @@ class Array3D : public Array { values) : Array(values) {} + // Creates an array of a floating-point type (half, bfloat16, float, + // or double) from the given nested initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::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..a75fffc605aa0df3e1e2eeb6d3129718cbbba0e4 100644 --- a/tensorflow/compiler/xla/array4d.h +++ b/tensorflow/compiler/xla/array4d.h @@ -82,6 +82,19 @@ class Array4D : public Array { values) : Array(values) {} + // Creates an array of a floating-point type (half, bfloat16, float, + // or double) from the given nested initializer list of float values. + template ::value || + std::is_same::value || + std::is_same::value || + std::is_same::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/compile_only_client.cc b/tensorflow/compiler/xla/client/compile_only_client.cc index c7e2c4367b89ca2112022fa40449ae3ebe28463e..59662c95ac15e7c23790c5b5ff5d75a694613aeb 100644 --- a/tensorflow/compiler/xla/client/compile_only_client.cc +++ b/tensorflow/compiler/xla/client/compile_only_client.cc @@ -39,16 +39,15 @@ CompileOnlyClient::CompileAheadOfTime( return compiler_service_->CompileAheadOfTime(service_instances, options); } -int64 CompileOnlyClient::PointerSizeForTriple( - tensorflow::StringPiece target_triple) { - llvm::Triple triple(llvm::Triple::normalize( - llvm::StringRef(target_triple.data(), target_triple.size()))); - if (triple.isArch64Bit()) { +int64 CompileOnlyClient::PointerSizeForTriple(tensorflow::StringPiece triple) { + llvm::Triple llvm_triple( + llvm::Triple::normalize(llvm::StringRef(triple.data(), triple.size()))); + if (llvm_triple.isArch64Bit()) { return 8; - } else if (triple.isArch32Bit()) { + } else if (llvm_triple.isArch32Bit()) { return 4; } else { - CHECK(triple.isArch16Bit()); + CHECK(llvm_triple.isArch16Bit()); return 2; } } 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 46f2ed4836eda6bf6d5b68f2e29ac6888cd1749b..39d02f0863f78d4094f2cc4805f534713fb7e929 100644 --- a/tensorflow/compiler/xla/client/computation_builder.cc +++ b/tensorflow/compiler/xla/client/computation_builder.cc @@ -233,6 +233,26 @@ StatusOr> ComputationBuilder::GetShape( 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(); @@ -388,7 +408,7 @@ ComputationDataHandle ComputationBuilder::Reshape( ComputationDataHandle ComputationBuilder::Collapse( const ComputationDataHandle& operand, - tensorflow::gtl::ArraySlice dims_to_collapse) { + tensorflow::gtl::ArraySlice dimensions) { if (!first_error_.ok()) { return ComputationDataHandle(); } @@ -396,8 +416,8 @@ ComputationDataHandle ComputationBuilder::Collapse( // Don't support out-of-order collapse here. // Checks that the collapsed dimensions are in order and consecutive. for (tensorflow::gtl::ArraySlice::size_type i = 1; - i < dims_to_collapse.size(); ++i) { - if (dims_to_collapse[i] - 1 != dims_to_collapse[i - 1]) { + i < dimensions.size(); ++i) { + if (dimensions[i] - 1 != dimensions[i - 1]) { NoteError(InvalidArgument( "Collapsed dimensions are not in order and consecutive.")); return ComputationDataHandle(); @@ -414,9 +434,9 @@ ComputationDataHandle ComputationBuilder::Collapse( VLOG(3) << "original shape: " << ShapeUtil::HumanString(*original_shape); VLOG(3) << "dims to collapse: " - << tensorflow::str_util::Join(dims_to_collapse, ","); + << tensorflow::str_util::Join(dimensions, ","); - if (dims_to_collapse.size() <= 1) { + if (dimensions.size() <= 1) { // Not collapsing anything, trivially we can return the operand versus // enqueueing a trivial reshape. return operand; @@ -424,7 +444,7 @@ ComputationDataHandle ComputationBuilder::Collapse( std::vector new_sizes; for (int i = 0; i < ShapeUtil::Rank(*original_shape); ++i) { - if (i <= dims_to_collapse.front() || i > dims_to_collapse.back()) { + if (i <= dimensions.front() || i > dimensions.back()) { new_sizes.push_back(original_shape->dimensions(i)); } else { new_sizes.back() *= original_shape->dimensions(i); @@ -733,13 +753,13 @@ ComputationDataHandle ComputationBuilder::Infeed(const Shape& shape, } void ComputationBuilder::Outfeed(const ComputationDataHandle& operand, - const Shape& shape, + const Shape& shape_with_layout, const string& outfeed_config) { OpRequest op_request; OutfeedRequest* request = op_request.mutable_outfeed_request(); request->set_outfeed_config(outfeed_config); *request->mutable_operand() = operand; - *request->mutable_shape() = shape; + *request->mutable_shape() = shape_with_layout; RunOpAndNoteError(&op_request); } @@ -769,6 +789,20 @@ ComputationDataHandle ComputationBuilder::CustomCall( return RunOpAndParseResponse(&op_request); } +ComputationDataHandle ComputationBuilder::HostCompute( + tensorflow::gtl::ArraySlice operands, + const string& channel_name, int64 cost_estimate_ns, const Shape& shape) { + OpRequest op_request; + HostComputeRequest* request = op_request.mutable_host_compute_request(); + for (const ComputationDataHandle& operand : operands) { + *request->add_operands() = operand; + } + *request->mutable_shape() = shape; + request->set_channel_name(channel_name); + request->set_cost_estimate_ns(cost_estimate_ns); + return RunOpAndParseResponse(&op_request); +} + ComputationDataHandle ComputationBuilder::Complex( const ComputationDataHandle& real, const ComputationDataHandle& imag, tensorflow::gtl::ArraySlice broadcast_dimensions) { @@ -834,6 +868,14 @@ ComputationDataHandle ComputationBuilder::Or( return BinaryOp(BINOP_OR, lhs, rhs, broadcast_dimensions); } +// TODO(b/65209188): Create a dedicated lowering for Xor +ComputationDataHandle ComputationBuilder::Xor( + const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, + tensorflow::gtl::ArraySlice broadcast_dimensions) { + return Or(And(Not(lhs), rhs, broadcast_dimensions), + And(lhs, Not(rhs), broadcast_dimensions)); +} + ComputationDataHandle ComputationBuilder::Not( const ComputationDataHandle& operand) { return UnaryOp(UNOP_NOT, operand); @@ -1200,6 +1242,22 @@ ComputationDataHandle ComputationBuilder::While( return RunOpAndParseResponse(&op_request); } +ComputationDataHandle ComputationBuilder::Gather( + const ComputationDataHandle& input, + const ComputationDataHandle& gather_indices, + const GatherDimensionNumbers& dimension_numbers, + tensorflow::gtl::ArraySlice window_bounds) { + OpRequest op_request; + GatherRequest* gather_request = op_request.mutable_gather_request(); + *gather_request->mutable_input() = input; + *gather_request->mutable_gather_indices() = gather_indices; + *gather_request->mutable_dimension_numbers() = dimension_numbers; + for (int64 window_bound : window_bounds) { + gather_request->add_window_bounds(window_bound); + } + return RunOpAndParseResponse(&op_request); +} + ComputationDataHandle ComputationBuilder::Conditional( const ComputationDataHandle& predicate, const ComputationDataHandle& true_operand, @@ -1332,15 +1390,16 @@ ComputationDataHandle ComputationBuilder::BatchNormInference( ComputationDataHandle ComputationBuilder::BatchNormGrad( const ComputationDataHandle& operand, const ComputationDataHandle& scale, - const ComputationDataHandle& mean, const ComputationDataHandle& var, + const ComputationDataHandle& batch_mean, + const ComputationDataHandle& batch_var, const ComputationDataHandle& grad_output, float epsilon, int64 feature_index) { OpRequest op_request; 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_mean() = batch_mean; + *request->mutable_variance() = batch_var; *request->mutable_grad_output() = grad_output; request->set_epsilon(epsilon); request->set_feature_index(feature_index); diff --git a/tensorflow/compiler/xla/client/computation_builder.h b/tensorflow/compiler/xla/client/computation_builder.h index ea4cdb76673b1c99036224bcd754ce4fe1360945..2141ebc2065a1a80d2fe820a7b6fe15434c89e28 100644 --- a/tensorflow/compiler/xla/client/computation_builder.h +++ b/tensorflow/compiler/xla/client/computation_builder.h @@ -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, @@ -195,9 +198,8 @@ class ComputationBuilder { tensorflow::gtl::ArraySlice new_sizes); // Enqueues an operation onto the computation that collapses the operand, from - // minor to major order, then reshapes it into the shape with the given - // dimension sizes, also from major to minor. Conceptually, this is a limited - // form of "shape casting". + // first to last dimension (C order), then reshapes it to the given dimension + // sizes. Conceptually, this is a limited form of "shape casting". ComputationDataHandle Reshape(const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice new_sizes); @@ -443,6 +445,16 @@ class ComputationBuilder { tensorflow::gtl::ArraySlice operands, const Shape& shape); + // Enqueues a pseudo-op to represent host-side computation data-dependencies. + // During code generation, host send and receive operations will be generated + // to transfer |operands| to the host and a single result of |shape| back to + // the device. Host send/recv operations are emitted using |channel_name|. + // Dataflow dependencies and the |cost_estimate_ns| field may be used in HLO + // instruction scheduling. + ComputationDataHandle HostCompute( + tensorflow::gtl::ArraySlice operands, + const string& channel_name, int64 cost_estimate_ns, const Shape& shape); + // The following methods enqueue element-wise binary arithmetic operations // onto the computation. The shapes of the operands have to match unless one // of the operands is a scalar, or an explicit broadcast dimension is given @@ -500,6 +512,10 @@ class ComputationBuilder { const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions = {}); + ComputationDataHandle Xor( + const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, + tensorflow::gtl::ArraySlice broadcast_dimensions = {}); + ComputationDataHandle Not(const ComputationDataHandle& operand); ComputationDataHandle ShiftLeft( @@ -705,6 +721,13 @@ class ComputationBuilder { const int exponent_bits, const int mantissa_bits); + // Enqueues a Gather node onto the computation. + ComputationDataHandle Gather( + const ComputationDataHandle& input, + const ComputationDataHandle& gather_indices, + const GatherDimensionNumbers& dimension_numbers, + tensorflow::gtl::ArraySlice window_bounds); + // Enqueues a Send node onto the computation, to send the given operand to // a Recv instruction that shares the same channel handle. void Send(const ComputationDataHandle& operand, const ChannelHandle& handle); @@ -853,7 +876,7 @@ class ComputationBuilder { Window* window); // Internal helper method that does the building for an arbitrary unary op. - ComputationDataHandle UnaryOp(UnaryOperation binop, + ComputationDataHandle UnaryOp(UnaryOperation unop, const ComputationDataHandle& operand); // Internal helper method that does the building for an arbitrary binary op. 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 ef98dbb6403beedb0c08ab9a0fc9e7d4ee31ab3b..91396f055fe4a3ecbd436139be9470e2a35e1c63 100644 --- a/tensorflow/compiler/xla/client/local_client.cc +++ b/tensorflow/compiler/xla/client/local_client.cc @@ -172,7 +172,9 @@ StatusOr> LocalExecutable::Run( std::unique_ptr result, executable_->ExecuteOnStreamWrapper( &service_options, run_options.execution_profile(), arguments)); - return ScopedShapedBuffer::MakeScoped(result.get(), run_options.allocator()); + + return MakeUnique(std::move(*result), + run_options.allocator()); } StatusOr> LocalExecutable::ExecuteAndDump( diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index b52a30f5a0b92e0094e6b0de3241c10a5a909cad..de0ed13c43f87966c272102b2e9af9ff3be63aea 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -69,7 +69,7 @@ class LocalExecutable { // of the computation. tensorflow::Status ValidateExecutionOptions( const tensorflow::gtl::ArraySlice arguments, - const ExecutableRunOptions& options, const Backend& backend); + const ExecutableRunOptions& run_options, const Backend& backend); // Records the computation in a SessionModule proto with the arguments used to // invoke it, and the result. Enabled by flag: --tla_dump_executions_to. diff --git a/tensorflow/compiler/xla/client/xla_client/BUILD b/tensorflow/compiler/xla/client/xla_client/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..b912889e2627aa01e5a7441e71e6bf002916ba5e --- /dev/null +++ b/tensorflow/compiler/xla/client/xla_client/BUILD @@ -0,0 +1,77 @@ +# Description: +# The new XLA client libraries. +# +# This is NOT YET ready to use. + +licenses(["notice"]) # Apache 2.0 + +package(default_visibility = [":friends"]) + +package_group( + name = "friends", + includes = [ + "//tensorflow/compiler/xla:friends", + ], +) + +# Filegroup used to collect source files for dependency checking. +filegroup( + name = "c_srcs", + data = glob([ + "**/*.cc", + "**/*.h", + ]), +) + +load("//tensorflow:tensorflow.bzl", "tf_cc_test") + +# TODO(b/74197823): Replace computation_builder with xla_builder. +cc_library( + name = "xla_builder", + srcs = ["xla_builder.cc"], + hdrs = ["xla_builder.h"], + deps = [ + "//tensorflow/compiler/xla:literal_util", + "//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:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_proto", + "//tensorflow/compiler/xla/service:shape_inference", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "xla_builder_test", + srcs = ["xla_builder_test.cc"], + deps = [ + ":xla_builder", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/core:test", + ], +) + +# ----------------------------------------------------------------------------- + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/compiler/xla/client/xla_client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_client/xla_builder.cc new file mode 100644 index 0000000000000000000000000000000000000000..6328a4f350fc70efaa96102f8202fb00b88b51f2 --- /dev/null +++ b/tensorflow/compiler/xla/client/xla_client/xla_builder.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/compiler/xla/client/xla_client/xla_builder.h" + +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/shape_inference.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/mutex.h" + +namespace xla { + +using tensorflow::strings::StrCat; + +namespace { + +int64 GetUniqueId() { + static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); + static int64 built_counter = 0; + tensorflow::mutex_lock loc(mu); + const int64 id = built_counter++; + return id; +} + +// Returns true if an instruction with the given opcode can be the root of the +// computation. +bool CanBeRoot(HloOpcode opcode) { + switch (opcode) { + case HloOpcode::kSend: + case HloOpcode::kOutfeed: + case HloOpcode::kTrace: + return false; + default: + return true; + } +} + +void SetOpcode(HloInstructionProto* instr, HloOpcode opcode) { + instr->set_opcode(HloOpcodeString(opcode)); +} + +} // namespace + +StatusOr> XlaBuilder::GetShape(const XlaOp& op) const { + TF_ASSIGN_OR_RETURN(auto instr, LookUpInstruction(op)); + return MakeUnique(instr->shape()); +} + +StatusOr XlaOp::GetShape() const { + TF_RET_CHECK(builder_ != nullptr); + TF_ASSIGN_OR_RETURN(auto shape, builder_->GetShape(*this)); + return *shape; +} + +XlaBuilder::XlaBuilder(const string& computation_name) + : name_(computation_name) {} + +XlaBuilder::~XlaBuilder() {} + +void XlaBuilder::NoteError(const Status& error) { + CHECK(!error.ok()); + if (die_immediately_on_error_) { + LOG(FATAL) << "error building computation: " << error; + } + + if (first_error_.ok()) { + first_error_ = error; + first_error_backtrace_.CreateCurrent(/*skip_count=*/1); + } +} + +StatusOr XlaBuilder::Build() { + if (!first_error_.ok()) { + string backtrace; + first_error_backtrace_.Dump(tensorflow::DebugWriteToString, &backtrace); + return AppendStatus(first_error_, backtrace); + } + + HloComputationProto entry; + ProgramShape* program_shape = entry.mutable_program_shape(); + + entry.set_name(name_); + + // Not all instructions can be roots. Walk backwards from the last added + // instruction until a valid root is found. + for (int64 i = instructions_.size() - 1; i >= 0; i--) { + TF_ASSIGN_OR_RETURN(HloOpcode opcode, + StringToHloOpcode(instructions_[i].opcode())); + if (CanBeRoot(opcode)) { + entry.set_root_name(instructions_[i].name()); + *program_shape->mutable_result() = instructions_[i].shape(); + break; + } + } + if (entry.root_name().empty()) { + return FailedPrecondition("no root instruction was found"); + } + + // Check that the parameter numbers are continuous from 0, and add parameter + // shapes and names to the program shape. + const int64 param_count = parameter_numbers_.size(); + for (int64 i = 0; i < param_count; i++) { + program_shape->add_parameters(); + program_shape->add_parameter_names(); + } + for (const HloInstructionProto& instr : instructions_) { + // Parameter number uniqueness is guaranteed in XlaBuilder::Parameter(). So + // to verify continuity, we just need to verify that every parameter is in + // the right range. + if (instr.opcode() == HloOpcodeString(HloOpcode::kParameter)) { + const int64 index = instr.parameter_number(); + TF_RET_CHECK(index >= 0 && index < param_count) + << "invalid parameter number: " << index; + *program_shape->mutable_parameters(index) = instr.shape(); + *program_shape->mutable_parameter_names(index) = instr.name(); + } + } + + for (auto& instruction : instructions_) { + entry.add_instructions()->Swap(&instruction); + } + + const int64 id = GetUniqueId(); + entry.set_id(id); + XlaComputation computation(id); + HloModuleProto* module = computation.mutable_proto(); + module->set_name(entry.name()); + module->set_entry_computation_name(entry.name()); + *module->mutable_program_shape() = entry.program_shape(); + for (auto& e : embedded_) { + module->add_computations()->Swap(&e.second); + } + module->add_computations()->Swap(&entry); + + return std::move(computation); +} + +XlaOp XlaBuilder::Add(const XlaOp& lhs, const XlaOp& rhs, + tensorflow::gtl::ArraySlice broadcast_dimensions) { + auto op = [&]() -> StatusOr { + HloInstructionProto instr; + SetOpcode(&instr, HloOpcode::kAdd); + TF_ASSIGN_OR_RETURN(const auto* lhs_instr, LookUpInstruction(lhs)); + TF_ASSIGN_OR_RETURN(const auto* rhs_instr, LookUpInstruction(rhs)); + TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + ShapeInference::InferBinaryOpShape( + HloOpcode::kAdd, lhs_instr->shape(), + rhs_instr->shape(), broadcast_dimensions)); + instr.add_operand_names(lhs_instr->name()); + instr.add_operand_names(rhs_instr->name()); + return AddInstruction(std::move(instr)); + }; + return NoteErrorOrReturn(op()); +} + +XlaOp XlaBuilder::ConstantLiteral(const Literal& literal) { + HloInstructionProto instr; + SetOpcode(&instr, HloOpcode::kConstant); + *instr.mutable_shape() = literal.shape(); + *instr.mutable_literal() = literal.ToProto(); + return AddInstruction(std::move(instr)); +} + +XlaOp XlaBuilder::Call(const XlaComputation& computation, + tensorflow::gtl::ArraySlice operands) { + auto op = [&]() -> StatusOr { + HloInstructionProto instr; + SetOpcode(&instr, HloOpcode::kCall); + std::vector operand_shapes; + for (const auto& operand : operands) { + TF_ASSIGN_OR_RETURN(const auto* input, LookUpInstruction(operand)); + operand_shapes.push_back(&input->shape()); + } + TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + ShapeInference::InferCallShape( + operand_shapes, + /*to_apply=*/computation.GetProgramShape())); + + // Add input operands. + for (const auto& operand : operands) { + TF_ASSIGN_OR_RETURN(auto operand_instr, LookUpInstruction(operand)); + instr.add_operand_names(operand_instr->name()); + } + + // Add called computation. + *instr.add_called_computation_names() = computation.proto().name(); + for (const HloComputationProto& e : computation.proto().computations()) { + embedded_.insert({e.id(), e}); + } + + return AddInstruction(std::move(instr)); + }; + return NoteErrorOrReturn(op()); +} + +XlaOp XlaBuilder::Parameter(int64 parameter_number, const Shape& shape, + const string& name) { + auto op = [&]() -> StatusOr { + HloInstructionProto instr; + SetOpcode(&instr, HloOpcode::kParameter); + if (parameter_numbers_.find(parameter_number) != parameter_numbers_.end()) { + return InvalidArgument("parameter %lld already registered", + parameter_number); + } + parameter_numbers_.insert(parameter_number); + instr.set_parameter_number(parameter_number); + instr.set_name(name); + *instr.mutable_shape() = shape; + return AddInstruction(std::move(instr)); + }; + return NoteErrorOrReturn(op()); +} + +XlaOp XlaBuilder::AddInstruction(HloInstructionProto&& instr) { + const int64 handle = instructions_.size(); + if (instr.name().empty()) { + instr.set_name(StrCat(instr.opcode(), ".", handle)); + } else { + // Append the handle to make sure the name is unique. + instr.set_name(StrCat(instr.name(), ".", handle)); + } + instructions_.push_back(instr); + + XlaOp op(handle, this); + return op; +} + +StatusOr XlaBuilder::LookUpInstruction( + const XlaOp& op) const { + TF_RET_CHECK(op.builder_ == this); + if (op.handle() >= instructions_.size() || op.handle() < 0) { + return InvalidArgument("no XlaOp value %lld", op.handle()); + } + return &instructions_[op.handle()]; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/xla_client/xla_builder.h b/tensorflow/compiler/xla/client/xla_client/xla_builder.h new file mode 100644 index 0000000000000000000000000000000000000000..7632bd289d792ef487fb667de3cea335e06778bf --- /dev/null +++ b/tensorflow/compiler/xla/client/xla_client/xla_builder.h @@ -0,0 +1,216 @@ +/* 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. +==============================================================================*/ + +// TODO(b/74197823): Replace computation_builder.h with this file. +// +// This is NOT YET ready to use. + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_XLA_CLIENT_XLA_BUILDER_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_XLA_CLIENT_XLA_BUILDER_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatset.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/stacktrace.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +class XlaBuilder; + +// This represents an instruction that has been enqueued using the XlaBuilder. +// This is used to pass to subsequent computations that depends upon the +// instruction as an operand. +// +// TODO(b/74197823): Replace xla::ComputationDataHandle with this one. +class XlaOp { + public: + StatusOr GetShape() const; + + private: + XlaOp() : handle_(0), builder_(nullptr) {} + XlaOp(int64 handle, XlaBuilder* builder) + : handle_(handle), builder_(builder) {} + + int64 handle() const { return handle_; } + friend class XlaBuilder; + + int64 handle_; + XlaBuilder* builder_; // Not owned. +}; + +// The computation graph that the user builds up with the XlaBuilder. +// +// TODO(b/74197823): Replace xla::Computation with this one. +class XlaComputation { + public: + XlaComputation(const XlaComputation&) = delete; + XlaComputation& operator=(const XlaComputation&) = delete; + + XlaComputation(XlaComputation&& from) { *this = std::move(from); } + + XlaComputation& operator=(XlaComputation&& from) { + proto_ = std::move(from.proto()); + unique_id_ = from.unique_id_; + return *this; + } + + // Returns the "program shape" (parameter and return shapes) for this + // computation. + const ProgramShape& GetProgramShape() const { return proto_.program_shape(); } + + const HloModuleProto& proto() const { return proto_; } + + private: + // Creates a null Computation. + XlaComputation(const int64 unique_id) : unique_id_(unique_id) {} + HloModuleProto* mutable_proto() { return &proto_; } + friend class XlaBuilder; + + int64 unique_id_; + HloModuleProto proto_; +}; + +// A convenient interface for building up computations. +// +// Thread-compatible. +// +// TODO(b/74197823): Replace xla::ComputationBuilder with this one. +class XlaBuilder { + public: + // computation_name: name to use for the built computation. + XlaBuilder(const string& computation_name); + + XlaBuilder(const XlaBuilder&) = delete; + XlaBuilder& operator=(const XlaBuilder&) = delete; + + ~XlaBuilder(); + + // Returns the computation name. + const string& name() const { return name_; } + + // Sets the builder to a mode where it will die immediately when an error is + // encountered, rather than producing it in a deferred fashion when Build() is + // called (which is the default). + void set_die_immediately_on_error(bool enabled) { + die_immediately_on_error_ = enabled; + } + + // Enqueues an add instruction onto the computation. + XlaOp Add(const XlaOp& lhs, const XlaOp& rhs, + tensorflow::gtl::ArraySlice broadcast_dimensions = {}); + + // Enqueues a call instruction onto the computation. + XlaOp Call(const XlaComputation& computation, + tensorflow::gtl::ArraySlice operands); + + // Enqueues a "retrieve parameter value" instruction for a parameter that was + // passed to the computation. + XlaOp Parameter(int64 parameter_number, const Shape& shape, + const string& name); + + // Enqueues a constant with the value of the given literal onto the + // computation. + XlaOp ConstantLiteral(const Literal& literal); + + // Enqueues a constant onto the computation. Methods are templated on the + // native host type (NativeT) which corresponds to a specific XLA + // PrimitiveType as given in the following table: + // + // Native Type PrimitiveType + // ----------------------------- + // bool PRED + // int32 S32 + // int64 S64 + // uint32 U32 + // uint64 U64 + // float F32 + // double F64 + // + // Note: not all primitive types defined in xla_data.proto have a + // corresponding native type yet. + template + XlaOp ConstantR0(NativeT value); + + // Returns the shape of the given op. + StatusOr> GetShape(const XlaOp& op) const; + + // Builds the computation with the requested operations, or returns a non-ok + // status. + StatusOr Build(); + + private: + XlaOp AddInstruction(HloInstructionProto&& instr); + + // 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); + + XlaOp NoteErrorOrReturn(StatusOr&& op) { + if (!op.ok()) { + NoteError(op.status()); + return XlaOp(); + } + return op.ConsumeValueOrDie(); + } + + StatusOr LookUpInstruction(const XlaOp& op) const; + + string name_; // Name to use for the built computation. + + // The first error encountered while building the computation. + // This is OK until the first error is encountered. + Status first_error_; + + // The saved stack trace from the point at which the first error occurred. + tensorflow::SavedStackTrace first_error_backtrace_; + + // The instructions of this computation. + std::vector instructions_; + + // The embedded computations used by this computation. Each computation was + // the entry computation of some XlaComputation, the key is the unique id of + // that XlaComputation. + std::map embedded_; + + // The unique parameter numbers. + tensorflow::gtl::FlatSet parameter_numbers_; + + // Mode bit that indicates whether to die when a first error is encountered. + bool die_immediately_on_error_ = false; +}; + +template +XlaOp XlaBuilder::ConstantR0(NativeT value) { + return ConstantLiteral(*Literal::CreateR0(value)); +} + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_XLA_CLIENT_XLA_BUILDER_H_ diff --git a/tensorflow/compiler/xla/client/xla_client/xla_builder_test.cc b/tensorflow/compiler/xla/client/xla_client/xla_builder_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a400e4e78b044ae633a0135b0011d5267eacc115 --- /dev/null +++ b/tensorflow/compiler/xla/client/xla_client/xla_builder_test.cc @@ -0,0 +1,137 @@ +/* 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/xla_client/xla_builder.h" + +#include + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +namespace { + +namespace op = xla::testing::opcode_matchers; + +using ::testing::HasSubstr; + +// TODO(b/74197823): Move the tests to service/. +class XlaBuilderTest : public ::testing::Test { + protected: + StatusOr> BuildHloModule(XlaBuilder* b) { + TF_ASSIGN_OR_RETURN(XlaComputation computation, b->Build()); + const HloModuleProto& proto = computation.proto(); + TF_ASSIGN_OR_RETURN(const auto& config, + HloModule::CreateModuleConfigFromProto(proto)); + return HloModule::CreateFromProto(proto, config); + } + + // Returns the name of the test currently being run. + string TestName() const { + return ::testing::UnitTest::GetInstance()->current_test_info()->name(); + } +}; + +TEST_F(XlaBuilderTest, OnePlusTwo) { + XlaBuilder b(TestName()); + b.Add(b.ConstantR0(1.0), b.ConstantR0(2.0)); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Add(op::Constant(), op::Constant())); +} + +TEST_F(XlaBuilderTest, ParamPlusConstant) { + XlaBuilder b(TestName()); + auto x = b.Parameter(0, ShapeUtil::MakeShape(F32, {3, 5}), "x"); + b.Add(x, b.ConstantR0(1.0)); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Add(op::Parameter(), op::Constant())); +} + +TEST_F(XlaBuilderTest, ParamPlusParam) { + XlaBuilder b(TestName()); + const auto& x_shape = ShapeUtil::MakeShape(S32, {2, 4, 6}); + const auto& y_shape = ShapeUtil::MakeShape(S32, {2, 4}); + auto x = b.Parameter(0, x_shape, "x"); + auto y = b.Parameter(1, y_shape, "y"); + auto add = b.Add(x, y, /*broadcast_dimensions=*/{0, 1}); + + TF_ASSERT_OK_AND_ASSIGN(auto add_shape, add.GetShape()); + EXPECT_TRUE(ShapeUtil::Equal(add_shape, x_shape)); + + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Add(op::Parameter(0), op::Parameter(1))); +} + +TEST_F(XlaBuilderTest, XPlusX) { + XlaBuilder b(TestName()); + auto x = b.Parameter(0, ShapeUtil::MakeShape(S32, {1, 3, 5, 7}), "x"); + b.Add(x, x); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Add(op::Parameter(0), op::Parameter(0))); +} + +TEST_F(XlaBuilderTest, ShapeInferenceError) { + XlaBuilder b(TestName()); + auto x = b.Parameter(0, ShapeUtil::MakeShape(U32, {2, 4, 6}), "x"); + auto y = b.Parameter(1, ShapeUtil::MakeShape(U32, {2, 4}), "y"); + b.Add(x, y); + auto statusor = BuildHloModule(&b); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), HasSubstr("shape inference")); +} + +TEST_F(XlaBuilderTest, ParameterAlreadyRegistered) { + XlaBuilder b_call("add"); + b_call.Parameter(0, ShapeUtil::MakeShape(PRED, {}), "x"); + + XlaBuilder b(TestName()); + auto x = b.Parameter(0, ShapeUtil::MakeShape(PRED, {}), "x"); + auto y = b.Parameter(0, ShapeUtil::MakeShape(PRED, {}), "y"); + b.Add(x, y); + auto statusor = BuildHloModule(&b); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("parameter 0 already registered")); +} + +TEST_F(XlaBuilderTest, Call) { + XlaBuilder b_call("the_only_to_apply"); + auto p0 = b_call.Parameter(0, ShapeUtil::MakeShape(F32, {}), "p0"); + auto p1 = b_call.Parameter(1, ShapeUtil::MakeShape(F32, {}), "p1"); + b_call.Add(p0, p1); + TF_ASSERT_OK_AND_ASSIGN(auto call, b_call.Build()); + XlaBuilder b(TestName()); + auto x = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); + auto y = b.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y"); + auto one = b.ConstantR0(1); + auto two = b.ConstantR0(2); + b.Add(b.Call(call, {x, y}), b.Call(call, {one, two})); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Add(op::Call(op::Parameter(), op::Parameter()), + op::Call(op::Constant(), op::Constant()))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/index_util.h b/tensorflow/compiler/xla/index_util.h index 0b9188e8524d6f1367541496dc5a86a250a0d530..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( diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc index 89279b659c75ce4775581dfbfa8d830f54ae6fe8..0a24db046a390eb447bc3518476c3dd9897d973c 100644 --- a/tensorflow/compiler/xla/literal_util.cc +++ b/tensorflow/compiler/xla/literal_util.cc @@ -223,7 +223,7 @@ Status Literal::CopySliceFromInternal( Literal::StrideConfig stride_config(src_literal.shape(), shape(), copy_size); - auto copy_proc = [&](const std::vector& indexes) { + auto copy_proc = [&](tensorflow::gtl::ArraySlice indexes) { // Map from multi-dimensional index, to source index. std::transform(indexes.begin(), indexes.end(), src_base.begin(), src_indexes.begin(), std::plus()); @@ -234,7 +234,8 @@ Status Literal::CopySliceFromInternal( int64 src_index = linear_index(src_literal.shape(), src_indexes); int64 dest_index = linear_index(shape(), dest_indexes); - StridedCopy(data(), dest_index, stride_config.dest_stride, + // `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; @@ -247,6 +248,28 @@ Status Literal::CopySliceFromInternal( return Status::OK(); } +Status Literal::CopyElementFrom(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_index, + tensorflow::gtl::ArraySlice dest_index) { + DCHECK_EQ(shape().element_type(), src_literal.shape().element_type()); + const int64 src_linear_index = IndexUtil::MultidimensionalIndexToLinearIndex( + src_literal.shape(), src_index); + const int64 dest_linear_index = + IndexUtil::MultidimensionalIndexToLinearIndex(shape(), dest_index); + const int64 primitive_size = + ShapeUtil::ByteSizeOfPrimitiveType(shape().element_type()); + + char* dest_address = + static_cast(untyped_data()) + dest_linear_index * primitive_size; + const char* source_address = + static_cast(src_literal.untyped_data()) + + src_linear_index * primitive_size; + if (dest_address != source_address) { + memcpy(dest_address, source_address, primitive_size); + } + return Status::OK(); +} + std::vector Literal::DecomposeTuple() { CHECK(ShapeUtil::IsTuple(shape())); std::vector elements; @@ -342,7 +365,7 @@ Status Literal::Piece::CopyFrom(const Literal::Piece& src) { #undef COPY_ELEMENTS default: return Unimplemented( - "Unhandled primitive type %s", + "Copying a Literal object with element type %s is not implemented.", PrimitiveType_Name(subshape().element_type()).c_str()); } } @@ -490,7 +513,10 @@ Status Literal::CopySliceFrom(const Literal& src_literal, default: break; } - return Unimplemented("Unhandled primitive type %d", shape().element_type()); + return Unimplemented( + "Copying a slice from a Literal object with element type %d is not " + "implemented.", + shape().element_type()); } /* static */ Literal Literal::Zero(PrimitiveType primitive_type) { @@ -807,9 +833,10 @@ std::unique_ptr Literal::Slice( DimensionVector result_dimensions; for (int64 dnum = 0; dnum < ShapeUtil::Rank(shape()); ++dnum) { CHECK_GE(start_indices[dnum], 0); - CHECK_LE(limit_indices[dnum], shape().dimensions(dnum)); + CHECK_LE(limit_indices[dnum], shape().dimensions(dnum)) + << "dnum = " << dnum; int64 dimension = limit_indices[dnum] - start_indices[dnum]; - CHECK_GT(dimension, 0); + CHECK_GE(dimension, 0) << "dnum = " << dnum; result_dimensions.push_back(dimension); } const auto result_shape = @@ -996,6 +1023,36 @@ StatusOr Literal::GetIntegralAsS64( } } +Status Literal::SetIntegralAsS64(tensorflow::gtl::ArraySlice multi_index, + int64 value) { + CHECK(LayoutUtil::IsDenseArray(shape())); + switch (shape().element_type()) { + case PRED: + Set(multi_index, value); + break; + case U8: + Set(multi_index, value); + break; + case S32: + Set(multi_index, value); + break; + case S64: + Set(multi_index, value); + break; + case U32: + Set(multi_index, value); + break; + case U64: + Set(multi_index, value); + break; + default: + return FailedPrecondition( + "Array element type is not integral: %s", + PrimitiveType_Name(shape().element_type()).c_str()); + } + return Status::OK(); +} + tensorflow::gtl::ArraySlice Literal::GetSparseIndex( int64 sparse_element_number, const ShapeIndex& shape_index) const { const Piece& p = piece(shape_index); @@ -1008,6 +1065,49 @@ void Literal::SortSparseElements(const ShapeIndex& shape_index) { piece(shape_index).SortSparseElements(); } +Literal Literal::GetFirstScalarLiteral() const { + CHECK(ShapeUtil::IsArray(shape_)); + CHECK_GT(ShapeUtil::ElementsIn(shape_), 0); + switch (shape_.element_type()) { + case PRED: + return std::move(*Literal::CreateR0(GetFirstElement())); + // 8 bit types. + case S8: + return std::move(*Literal::CreateR0(GetFirstElement())); + case U8: + return std::move(*Literal::CreateR0(GetFirstElement())); + // 16 bit types. + case BF16: + return std::move( + *Literal::CreateR0(GetFirstElement())); + case F16: + return std::move(*Literal::CreateR0(GetFirstElement())); + case S16: + return std::move(*Literal::CreateR0(GetFirstElement())); + case U16: + return std::move(*Literal::CreateR0(GetFirstElement())); + // 32 bit types. + case F32: + return std::move(*Literal::CreateR0(GetFirstElement())); + case S32: + return std::move(*Literal::CreateR0(GetFirstElement())); + case U32: + return std::move(*Literal::CreateR0(GetFirstElement())); + // 64 bit types. + case C64: + return std::move( + *Literal::CreateR0(GetFirstElement())); + case F64: + return std::move(*Literal::CreateR0(GetFirstElement())); + case S64: + return std::move(*Literal::CreateR0(GetFirstElement())); + case U64: + return std::move(*Literal::CreateR0(GetFirstElement())); + default: + LOG(FATAL) << "Unhandled primitive type " << shape_.element_type(); + } +} + void Literal::Piece::SortSparseElements() { switch (subshape().element_type()) { case PRED: @@ -1257,11 +1357,17 @@ string Literal::ToString(bool print_layout) const { /* static */ std::unique_ptr Literal::MakeTupleOwned( std::vector> elements) { - std::vector element_ptrs; + std::vector element_shapes; + element_shapes.reserve(elements.size()); for (const auto& element : elements) { - element_ptrs.push_back(element.get()); + 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})); } - return MakeTuple(element_ptrs); + return literal; } void Literal::EachCellAsString( @@ -1344,8 +1450,8 @@ StatusOr> ConvertIfDestTypeMatches( return ConvertToC64(src_literal); // Other types are not yet supported. default: - return InvalidArgument( - "Unimplemented: Convert from type %s to type %s", + return Unimplemented( + "Converting from type %s to type %s is not implemented.", PrimitiveType_Name(src_literal.shape().element_type()).c_str(), PrimitiveType_Name(primitive_dest_type).c_str()); } @@ -1374,10 +1480,29 @@ StatusOr> Literal::Convert( #undef CONVERT_IF_DEST_TYPE_MATCHES // Other types are not yet supported. default: - return InvalidArgument("Unimplemented: Convert from type %s to type %s", - PrimitiveType_Name(shape().element_type()).c_str(), - PrimitiveType_Name(primitive_dest_type).c_str()); + return Unimplemented( + "Converting from type %s to type %s is not implemented.", + PrimitiveType_Name(shape().element_type()).c_str(), + PrimitiveType_Name(primitive_dest_type).c_str()); + } +} + +StatusOr> Literal::ConvertToShape( + const Shape& dest_shape) const { + if (!ShapeUtil::IsTuple(dest_shape)) { + return Convert(dest_shape.element_type()); + } + std::vector elements; + for (int i = 0; i < ShapeUtil::TupleElementCount(shape()); ++i) { + auto element = LiteralView::Create(*this, {i}); + TF_ASSIGN_OR_RETURN( + auto new_element, + element.ConvertToShape(ShapeUtil::GetSubshape(dest_shape, {i}))); + elements.push_back(std::move(*new_element)); } + auto converted = MakeUnique(); + *converted = Literal::MoveIntoTuple(&elements); + return std::move(converted); } template @@ -1564,6 +1689,92 @@ bool Literal::IsAllComplex(complex64 value) const { } } +bool Literal::IsAllFirst() const { + for (const auto& pair : pieces_) { + const Piece& piece = pair.second; + if (!ShapeUtil::IsArray(piece.subshape())) { + continue; + } + + // Empty shapes are not all the first element since there is no first + // element. + if (ShapeUtil::HasZeroElements(piece.subshape())) { + return false; + } + auto piece_is_all = [&]() { + switch (piece.subshape().element_type()) { + case PRED: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + // 8 bit types + case S8: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case U8: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + // 16 bit types + case BF16: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case F16: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case S16: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case U16: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + // 32 bit types + case F32: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case U32: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case S32: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + // 64 bit types + case C64: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case F64: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case S64: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + case U64: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } + default: + return false; + } + }; + + if (!piece_is_all()) { + return false; + } + } + return true; +} + bool Literal::IsZero(tensorflow::gtl::ArraySlice indices) const { CHECK(ShapeUtil::IsArray(shape())); switch (shape().element_type()) { diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h index 2b68b8f177dcae19f55a45f2f284ea10f7ce6e52..e24f5285d9a14cf26216e4a16c6d1e516afc413f 100644 --- a/tensorflow/compiler/xla/literal_util.h +++ b/tensorflow/compiler/xla/literal_util.h @@ -262,6 +262,11 @@ class Literal { tensorflow::gtl::ArraySlice dest_base, tensorflow::gtl::ArraySlice copy_size); + // Copies one element from src_literal[src_index] to (*this)[dest_index]. + Status CopyElementFrom(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_index, + tensorflow::gtl::ArraySlice dest_index); + // 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 @@ -333,6 +338,11 @@ class Literal { StatusOr> Convert( PrimitiveType primitive_dest_type) const; + // Converts this literal to the given shape. Returns an error is the + // conversion is not possible. + StatusOr> ConvertToShape( + const Shape& dest_shape) const; + // Creates a scalar literal value zero of the given primitive type. static Literal Zero(PrimitiveType primitive_type); @@ -451,6 +461,9 @@ class Literal { template NativeT GetFirstElement() const; + // Returns a literal scalar representing the first element. + Literal GetFirstScalarLiteral() const; + // As Get(), but determines the correct type and converts the value // into text. string GetAsString(tensorflow::gtl::ArraySlice multi_index, @@ -466,6 +479,11 @@ class Literal { StatusOr GetIntegralAsS64( tensorflow::gtl::ArraySlice multi_index) const; + // As Set(), but truncates `value` to the literal element type before storing. + // This literal must be an array. + Status SetIntegralAsS64(tensorflow::gtl::ArraySlice multi_index, + int64 value); + // Returns an identity matrix (rank 2) with the given row and column count. template static std::unique_ptr MakeIdentityR2(int64 size); @@ -485,6 +503,27 @@ class Literal { static std::unique_ptr MakeTupleOwned( std::vector> elements); + // 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; @@ -581,6 +620,9 @@ class Literal { // This literal must have a dense layout. bool IsAllComplex(complex64 value) const; + // Literal consists entirely of the first element of the literal. + bool IsAllFirst() const; + // Returns whether this literal is zero at the specified index. This literal // must be an array with a dense layout. bool IsZero(tensorflow::gtl::ArraySlice indices) const; @@ -1242,7 +1284,7 @@ Status Literal::Populate(const FnType& generator) { int64 minor_dimension_size = ShapeUtil::GetDimension(this_shape, stride_config.minor_dimension); - auto init_function = [&](const std::vector& indexes) { + auto init_function = [&](tensorflow::gtl::ArraySlice indexes) { const int64 index = IndexUtil::MultidimensionalIndexToLinearIndex(shape(), indexes); std::copy(indexes.begin(), indexes.end(), minor_scan_indexes.begin()); diff --git a/tensorflow/compiler/xla/literal_util_test.cc b/tensorflow/compiler/xla/literal_util_test.cc index b3583c2eb75de8297d5e7507430491f119bd4462..04e45f00491b0bef94f3c0af1c875b2d007194fd 100644 --- a/tensorflow/compiler/xla/literal_util_test.cc +++ b/tensorflow/compiler/xla/literal_util_test.cc @@ -30,6 +30,7 @@ limitations under the License. namespace xla { namespace { +using tensorflow::gtl::ArraySlice; using ::testing::ElementsAre; using ::testing::HasSubstr; @@ -214,11 +215,11 @@ TEST_F(LiteralUtilTest, CreateSparse) { 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)); + ArraySlice(expected_indices.data(), + expected_indices.num_elements())); + EXPECT_EQ( + ArraySlice(literal->data().data(), expected_values.size()), + ArraySlice(expected_values)); } TEST_F(LiteralUtilTest, LiteralR4F32ProjectedStringifies) { @@ -290,7 +291,7 @@ TEST_F(LiteralUtilTest, EachCellR2F32) { // clang-format on std::vector> seen; literal->EachCellAsString( - [&seen](tensorflow::gtl::ArraySlice indices, const string& value) { + [&seen](ArraySlice indices, const string& value) { seen.emplace_back(indices[0], indices[1], value); }); @@ -501,6 +502,24 @@ TEST_F(LiteralUtilTest, IsAllComplex) { ->IsAllComplex({8.0f, 9.0f})); } +TEST_F(LiteralUtilTest, IsAllFirst) { + // IsAllComplex always returns false when the literal is not complex. + EXPECT_FALSE(Literal::CreateR1({false, true})->IsAllFirst()); + EXPECT_TRUE(Literal::CreateR1({false, false})->IsAllFirst()); + EXPECT_FALSE(Literal::CreateR1({1, 1, 2})->IsAllFirst()); + EXPECT_TRUE(Literal::CreateR1({5, 5, 5, 5})->IsAllFirst()); + EXPECT_FALSE(Literal::CreateR1({1, 1, 2})->IsAllFirst()); + EXPECT_TRUE(Literal::CreateR1({5, 5, 5, 5})->IsAllFirst()); + EXPECT_FALSE(Literal::CreateR1({1, 1, 2})->IsAllFirst()); + EXPECT_TRUE(Literal::CreateR1({5, 5, 5, 5})->IsAllFirst()); + EXPECT_FALSE(Literal::CreateR1({1, 1, 2})->IsAllFirst()); + + complex64 c8_9 = {8, 9}; + complex64 c7_9 = {7, 9}; + EXPECT_TRUE(Literal::CreateR2({{c8_9}, {c8_9}})->IsAllFirst()); + EXPECT_FALSE(Literal::CreateR2({{c7_9}, {c8_9}})->IsAllFirst()); +} + TEST_F(LiteralUtilTest, IsZero) { auto scalar_zero = Literal::CreateR0(0.0f); auto scalar_one = Literal::CreateR0(1.0f); @@ -604,11 +623,10 @@ TEST_F(LiteralUtilTest, TransposeR4) { // clang-format on auto reshape = original->Transpose(/*permutation=*/{2, 3, 0, 1}); - reshape->EachCell( - [&](tensorflow::gtl::ArraySlice indices, float value) { - EXPECT_EQ(value, original->Get( - {indices[2], indices[3], indices[0], indices[1]})); - }); + reshape->EachCell([&](ArraySlice indices, float value) { + EXPECT_EQ(value, original->Get( + {indices[2], indices[3], indices[0], indices[1]})); + }); } TEST_F(LiteralUtilTest, TestR4RelayoutEquivalence) { @@ -845,7 +863,7 @@ TEST_F(LiteralUtilTest, CopySliceFrom) { const int64 zero_base[] = {0, 0, 0, 0}; const int64 step[] = {1, 1, 1, 1}; uint32 seqnr = 0; - auto init_proc = [&](const std::vector& indexes) { + auto init_proc = [&](ArraySlice indexes) { source->Set(indexes, ++seqnr); return true; }; @@ -861,7 +879,7 @@ TEST_F(LiteralUtilTest, CopySliceFrom) { std::vector source_indexes(TF_ARRAYSIZE(dimensions), 0); std::vector blank_indexes(TF_ARRAYSIZE(dimensions), 0); bool matched = true; - auto check_proc = [&](const std::vector& indexes) { + auto check_proc = [&](ArraySlice indexes) { std::copy(indexes.begin(), indexes.end(), source_indexes.begin()); std::transform(source_indexes.begin(), source_indexes.end(), src_base, source_indexes.begin(), std::plus()); @@ -1049,7 +1067,7 @@ TEST_F(LiteralUtilTest, Populate) { primitive_util::NativeToPrimitiveType(), data.dimensions, data.layout); auto literal = Literal::CreateFromShape(shape); - auto generator = [&](tensorflow::gtl::ArraySlice indexes) -> uint32 { + auto generator = [&](ArraySlice indexes) -> uint32 { // Offsets from linear index just to avoid R0 literals to be initialized // with zero. return IndexUtil::MultidimensionalIndexToLinearIndex(literal->shape(), @@ -1061,7 +1079,7 @@ TEST_F(LiteralUtilTest, Populate) { std::vector zero_base(data.dimensions.size(), 0); std::vector step(data.dimensions.size(), 1); bool matched = true; - auto check_function = [&](const std::vector& indexes) { + auto check_function = [&](ArraySlice indexes) { auto value = literal->Get(indexes); matched = matched && (value == generator(indexes)); return matched; @@ -1214,15 +1232,15 @@ TEST_F(LiteralUtilTest, ConvertIfTypesMatch) { EXPECT_EQ(*conv, *c64); EXPECT_EQ(s32->Convert(TUPLE).status().code(), - tensorflow::error::INVALID_ARGUMENT); + tensorflow::error::UNIMPLEMENTED); EXPECT_EQ(s32->Convert(S16).status().code(), - tensorflow::error::INVALID_ARGUMENT); + tensorflow::error::UNIMPLEMENTED); EXPECT_EQ(s32->Convert(U16).status().code(), - tensorflow::error::INVALID_ARGUMENT); + tensorflow::error::UNIMPLEMENTED); EXPECT_EQ(c64->Convert(F32).status().code(), - tensorflow::error::INVALID_ARGUMENT); + tensorflow::error::UNIMPLEMENTED); EXPECT_EQ(c64->Convert(S32).status().code(), - tensorflow::error::INVALID_ARGUMENT); + tensorflow::error::UNIMPLEMENTED); } TEST_F(LiteralUtilTest, CopyFromProto_Bool) { diff --git a/tensorflow/compiler/xla/map_util.h b/tensorflow/compiler/xla/map_util.h index 0ad0b9133075fdb29effe1be2c007a851f420a33..8db8c6f3de84a6c46625eadbb6b0f83d2262e5f7 100644 --- a/tensorflow/compiler/xla/map_util.h +++ b/tensorflow/compiler/xla/map_util.h @@ -65,6 +65,25 @@ MaybeFind(const Collection& collection, 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 diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc index 5772532b84820de29f9c303876ed128518c6a36e..b21ab3044fae7136071f50bdba6e74b799a309d5 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.cc +++ b/tensorflow/compiler/xla/python/local_computation_builder.cc @@ -98,15 +98,25 @@ const std::unique_ptr& LocalShapedBuffer::shaped_buffer() 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) { +LocalShapedBuffer* LocalShapedBuffer::FromLiteral( + const Literal& argument, + const tensorflow::gtl::optional& shape_with_layout) { LocalClient* client = GetOrCreateLocalClient(); - std::unique_ptr buf = - client - ->LiteralToShapedBuffer(argument, - /*device_ordinal=*/0, - client->backend().memory_allocator()) - .ConsumeValueOrDie(); + 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)); } @@ -120,7 +130,8 @@ CompiledLocalComputation::CompiledLocalComputation( : executable_(std::move(executable)) {} StatusOr> CompiledLocalComputation::Execute( - const std::vector& arguments) { + const std::vector& arguments, + const std::vector>& shapes_with_layout) { LocalClient* client = GetOrCreateLocalClient(); VLOG(1) << "Execution requested with " << GetReplicaCount() << " replicas."; @@ -133,7 +144,8 @@ StatusOr> CompiledLocalComputation::Execute( GetReplicaCount()); for (int replica = 0; replica < GetReplicaCount(); ++replica) { - pool.Schedule([this, client, replica, &arguments, &results] { + pool.Schedule([this, client, replica, &arguments, &shapes_with_layout, + &results] { StatusOr device_ordinal_status = client->ReplicaNumberToDeviceOrdinal(replica); if (!device_ordinal_status.ok()) { @@ -144,18 +156,28 @@ StatusOr> CompiledLocalComputation::Execute( 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 (const Literal& argument : arguments) { - StatusOr> pushed = - client->LiteralToShapedBuffer( - argument, device_ordinal, - client->backend().memory_allocator()); + 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()); } @@ -256,6 +278,12 @@ const Computation& LocalComputation::computation() const { 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_(GetOrCreateLocalClient(), computation_name) {} @@ -281,6 +309,11 @@ 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); } @@ -335,6 +368,12 @@ ComputationDataHandle LocalComputationBuilder::Slice( return builder_.Slice(operand, start_indices, limit_indices, strides); } +ComputationDataHandle LocalComputationBuilder::SliceInDim( + const ComputationDataHandle& operand, int64 start_index, int64 limit_index, + int64 stride, int64 dimno) { + return builder_.SliceInDim(operand, start_index, limit_index, stride, dimno); +} + ComputationDataHandle LocalComputationBuilder::DynamicSlice( const ComputationDataHandle& operand, const ComputationDataHandle& start_indices, @@ -367,12 +406,6 @@ LocalComputationBuilder::SelectAndScatterWithGeneralPadding( source, init_value, scatter.computation()); } -ComputationDataHandle LocalComputationBuilder::Select( - const ComputationDataHandle& pred, const ComputationDataHandle& on_true, - const ComputationDataHandle& on_false) { - return builder_.Select(pred, on_true, on_false); -} - ComputationDataHandle LocalComputationBuilder::Tuple( tensorflow::gtl::ArraySlice elements) { return builder_.Tuple(elements); @@ -388,6 +421,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, @@ -471,6 +510,17 @@ ComputationDataHandle LocalComputationBuilder::While( 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; \ @@ -487,6 +537,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) @@ -507,6 +566,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) @@ -523,6 +583,7 @@ _FORWARD_UNOP(Sort) #undef _FORWARD #undef _FORWARD_UNOP #undef _FORWARD_BINOP +#undef _FORWARD_TRIOP void DeleteLocalShapedBuffer(LocalShapedBuffer* local_shaped_buffer) { delete local_shaped_buffer; diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h index 6851c2644d097b04bb893f943692862e6b3b5f17..a7375c8965e9041226ffee08dab6ffafa25312af 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.h +++ b/tensorflow/compiler/xla/python/local_computation_builder.h @@ -59,7 +59,9 @@ StatusOr > TransferFromOutfeedLocalReplica( // client. class LocalShapedBuffer { public: - static LocalShapedBuffer* FromLiteral(const Literal& argument); + 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; @@ -77,8 +79,15 @@ class LocalShapedBuffer { class CompiledLocalComputation { public: CompiledLocalComputation(std::unique_ptr executable); + + // 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& arguments, + const std::vector >& shapes_with_layout); + LocalShapedBuffer* ExecuteWithShapedBuffers( tensorflow::gtl::ArraySlice argument_handles); @@ -93,11 +102,16 @@ class CompiledLocalComputation { class LocalComputation { public: 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: Computation computation_; }; @@ -124,6 +138,9 @@ class LocalComputationBuilder { 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, @@ -153,6 +170,10 @@ class LocalComputationBuilder { tensorflow::gtl::ArraySlice limit_indices, tensorflow::gtl::ArraySlice strides); + ComputationDataHandle SliceInDim(const ComputationDataHandle& operand, + int64 start_index, int64 limit_index, + int64 stride, int64 dimno); + ComputationDataHandle DynamicSlice( const ComputationDataHandle& operand, const ComputationDataHandle& start_indices, @@ -174,10 +195,6 @@ class LocalComputationBuilder { const ComputationDataHandle& source, const ComputationDataHandle& init_value, const LocalComputation& scatter); - ComputationDataHandle Select(const ComputationDataHandle& pred, - const ComputationDataHandle& on_true, - const ComputationDataHandle& on_false); - ComputationDataHandle Tuple( tensorflow::gtl::ArraySlice elements); @@ -187,6 +204,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, @@ -241,6 +262,12 @@ class LocalComputationBuilder { 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; @@ -254,6 +281,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) @@ -274,6 +309,7 @@ class LocalComputationBuilder { _FORWARD_UNOP(Exp) _FORWARD_UNOP(Floor) _FORWARD_UNOP(Ceil) + _FORWARD_UNOP(Round) _FORWARD_UNOP(Log) _FORWARD_UNOP(Sign) _FORWARD_UNOP(Cos) @@ -290,6 +326,7 @@ class LocalComputationBuilder { #undef _FORWARD #undef _FORWARD_UNOP #undef _FORWARD_BINOP +#undef _FORWARD_TRIOP private: ComputationBuilder builder_; diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i index 6a52a088ddefa5485ba34d321a6ded07cc616d02..b5354131c94930b75ea66036ddb61ecd3993414f 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.i +++ b/tensorflow/compiler/xla/python/local_computation_builder.i @@ -27,12 +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. @@ -55,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 @@ -199,6 +201,15 @@ tensorflow::ImportNumpy(); } } +%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( @@ -353,15 +364,31 @@ tensorflow::ImportNumpy(); // Shape %typemap(in) const Shape& (Shape temp) { - Status shape_status = numpy::CheckPyShapeInfo($input); - if (!shape_status.ok()) { - PyErr_SetString(PyExc_RuntimeError, shape_status.ToString().c_str()); + 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); } @@ -374,14 +401,37 @@ tensorflow::ImportNumpy(); const int size = PySequence_Size($input); for (int i = 0; i < size; ++i) { PyObject* o = PySequence_GetItem($input, i); - Status shape_status = numpy::CheckPyShapeInfo(o); - if (!shape_status.ok()) { - PyErr_SetString(PyExc_RuntimeError, shape_status.ToString().c_str()); - 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; } @@ -471,6 +521,135 @@ 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& @@ -653,6 +832,21 @@ tensorflow::ImportNumpy(); } 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; } } @@ -673,6 +867,7 @@ tensorflow::ImportNumpy(); %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; @@ -680,6 +875,7 @@ tensorflow::ImportNumpy(); %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; @@ -690,6 +886,7 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Collapse; %unignore xla::swig::LocalComputationBuilder::CrossReplicaSum; %unignore xla::swig::LocalComputationBuilder::Slice; +%unignore xla::swig::LocalComputationBuilder::SliceInDim; %unignore xla::swig::LocalComputationBuilder::DynamicSlice; %unignore xla::swig::LocalComputationBuilder::DynamicUpdateSlice; %unignore xla::swig::LocalComputationBuilder::ConcatInDim; @@ -701,6 +898,7 @@ 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; @@ -708,6 +906,7 @@ tensorflow::ImportNumpy(); %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; @@ -715,6 +914,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; @@ -730,6 +930,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; diff --git a/tensorflow/compiler/xla/python/numpy_bridge.cc b/tensorflow/compiler/xla/python/numpy_bridge.cc index 5c722623e318ece9eca6bdc8750195ce5fd5defb..3d87480728aab1d4ebbc71c6c7504d37cae5edaf 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.cc +++ b/tensorflow/compiler/xla/python/numpy_bridge.cc @@ -176,85 +176,107 @@ static string PyObjectCppRepr(PyObject* o) { return ExtractStringAndDecref(r); } -Status CheckPyShapeInfo(PyObject* o) { +StatusOr XlaShapeFromPyShape(PyObject* o) { auto error = [o](const string& prefix) { return InvalidArgument("%s; got %s", prefix.c_str(), PyObjectCppRepr(o).c_str()); }; - // The object is a tuple (a pair) - if (!PyTuple_Check(o)) { - return error("Shape record must be a tuple"); - } - if (PyTuple_Size(o) != 2) { - return error("Shape record tuple must be of length 2"); - } - // It has a first element, which is a numpy dtype object - PyObject* first = PyTuple_GetItem(o, 0); - if (first == nullptr) { - return error("Tuple has no item 0 (shape dtype)"); - } - if (first->ob_type != &PyArrayDescr_Type) { - return error( - "Shape record does not have a numpy dtype as its first element"); - } - const int np_type = NumpyTypenum(first); - if (!NumpyTypeIsValid(np_type)) { - return error("Shape record has an invalid integer dtype"); - } + 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; + }; - // 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 error("Tuple has no item 0 (shape dimensions)"); - } - if (!PyTuple_Check(second)) { - return error("Shape record does not have a tuple as its second element"); - } - 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) { - VLOG(3) << "element_type is tuple, checking member: " << i; - Status result = CheckPyShapeInfo(dimension); - if (!result.ok()) { - return AddStatus( - result, tensorflow::strings::StrCat("Validating tuple member ", i, - " of ", PyObjectCppRepr(o))); - } - } else if (!CheckPyIntOrLong(dimension)) { - return error("Non-tuple shape record has a non-integer dimension"); + 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; + }; - return Status::OK(); -} + 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"); + } + } + } + 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); } - return ShapeUtil::MakeShape(element_type, dimensions); } } diff --git a/tensorflow/compiler/xla/python/numpy_bridge.h b/tensorflow/compiler/xla/python/numpy_bridge.h index 6ff1c34cfc5e0323a6729bdfd5572239f4966211..adfcc3b8588dce01718bb19dea936bace483be4d 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.h +++ b/tensorflow/compiler/xla/python/numpy_bridge.h @@ -56,15 +56,11 @@ 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. -Status 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. diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py index a89e2643c8937db2b57273e373cf4601695f1a5d..90cda42f3227c80826ffbf4e5473647c2795544d 100644 --- a/tensorflow/compiler/xla/python/xla_client.py +++ b/tensorflow/compiler/xla/python/xla_client.py @@ -30,9 +30,9 @@ 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. +# 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 @@ -89,6 +89,7 @@ _UNARY_OPS = [ 'Abs', 'Exp', 'Floor', + 'Round', 'Ceil', 'Log', 'Sign', @@ -122,24 +123,34 @@ _BINARY_OPS = [ 'Pow', ] + 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.PRED: np.dtype('bool'), + xla_data_pb2.S8: np.dtype('int8'), + xla_data_pb2.S16: np.dtype('int16'), + xla_data_pb2.S32: np.dtype('int32'), + xla_data_pb2.S64: np.dtype('int64'), + xla_data_pb2.U8: np.dtype('uint8'), + xla_data_pb2.U16: np.dtype('uint16'), + xla_data_pb2.U32: np.dtype('uint32'), + xla_data_pb2.U64: np.dtype('uint64'), + xla_data_pb2.F16: np.dtype('float16'), + xla_data_pb2.F32: np.dtype('float32'), + xla_data_pb2.F64: np.dtype('float64'), + xla_data_pb2.C64: np.dtype('complex64'), 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() -} +DTYPE_TO_XLA_ELEMENT_TYPE = {str(dt): et + for et, dt in XLA_ELEMENT_TYPE_TO_DTYPE.items()} + + +def dtype_to_etype(dtype): + """Convenience function for reading DTYPE_TO_XLA_ELEMENT_TYPE.""" + return DTYPE_TO_XLA_ELEMENT_TYPE[str(np.dtype(dtype))] class LocalBuffer(object): @@ -155,9 +166,14 @@ class LocalBuffer(object): self._delete = c_api.DeleteLocalShapedBuffer @staticmethod - def from_py(npval): + def from_py(npval, layout_fn=None): npval = require_numpy_array_layout(npval) - return LocalBuffer(c_api.LocalShapedBuffer.FromLiteral(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() @@ -182,13 +198,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})'.format( - self.np_dtype, self._dimensions) + 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)] @@ -201,11 +227,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): @@ -222,23 +286,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 @@ -302,8 +353,7 @@ def transfer_from_outfeed(shape, replica_number=None): Returns: The literal value that is produced from the outfeed queue. """ - return c_api.TransferFromOutfeedLocalReplica( - _unwrap_shape(shape), replica_number or 0) + return c_api.TransferFromOutfeedLocalReplica(shape, replica_number or 0) class LocalComputation(object): @@ -320,28 +370,70 @@ 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=(), compile_options=None): + 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), compile_options), + self.c_local_computation.Compile(argument_shapes, compile_options), is_compiled=True) - def CompileWithExampleArguments(self, arguments=(), compile_options=None): + def CompileWithExampleArguments(self, + arguments=(), + compile_options=None, + layout_fn=None): return self.Compile( argument_shapes=[Shape.from_numpy(arg) for arg in arguments], - compile_options=compile_options) + 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.""" @@ -397,7 +489,7 @@ class ComputationBuilder(object): Returns: A ComputationDataHandle message. """ - return _wrap_data_handle(self._client.Infeed(_unwrap_shape(shape))) + return _wrap_data_handle(self._client.Infeed(shape)) def Outfeed(self, operand): """Enqueues an outfeed op onto the computation. @@ -406,7 +498,7 @@ class ComputationBuilder(object): outfeed queue for subsequent dequeue via the client API. """ self._client.Outfeed( - _unwrap_data_handle(operand), _unwrap_shape(self.GetShape(operand)), + _unwrap_data_handle(operand), self.GetShape(operand), ''.encode('utf-8')) def Constant(self, value): @@ -497,8 +589,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. @@ -558,6 +649,9 @@ 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() @@ -572,7 +666,7 @@ class ComputationBuilder(object): representing the configuration of the padding operation. Returns: - A ComputationDataHandle representing the added pad op. + A ComputationDataHandle representing the added Pad op. """ if not isinstance(padding_config, xla_data_pb2.PaddingConfig): padding_config = GetPaddingConfigFromTriples(padding_config) @@ -582,7 +676,20 @@ class ComputationBuilder(object): padding_config)) def Reshape(self, operand, dimensions, new_sizes): - """Reshape op.""" + """Enqueues a reshape op onto the computation. + + Args: + operand: ComputationDataHandle representing the array to be reshaped. + dimensions: sequence of integers encoding the order in which dimensions + are collapsed or None, in which case dimensions are flattened in order. + new_sizes: sequence of integers encoding the new dimension sizes (shape). + + Returns: + A ComputationDataHandle representing the added Reshape op. + """ + if dimensions is None: + ndim = len(self.GetShape(operand).dimensions()) + dimensions = tuple(range(ndim)) return _wrap_data_handle( self._client.Reshape( _unwrap_data_handle(operand), dimensions, new_sizes)) @@ -619,6 +726,13 @@ 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. @@ -681,11 +795,27 @@ class ComputationBuilder(object): strides = [1] * len(start_indices) return _wrap_data_handle( self._client.Slice( - _unwrap_data_handle(operand), - start_indices, - limit_indices, + _unwrap_data_handle(operand), start_indices, limit_indices, strides)) + def SliceInDim(self, operand, start_index, limit_index, stride, dimno): + """Enqueues a slice-in-dimension operation onto the computation. + + Args: + operand: ComputationDataHandle for the N dimensional array to be sliced. + start_index: an integer containing the start index of the slice. + limit_index: an integer containing the end index of the slice. + stride: an integer containing the stride size for the slice. + dimno: an integer indicating the dimension along which to slice. + + Returns: + A ComputationDataHandle representing the added Slice op. + """ + return _wrap_data_handle( + self._client.SliceInDim( + _unwrap_data_handle(operand), start_index, limit_index, stride, + dimno)) + def DynamicSlice(self, operand, start_indices, slice_sizes): """Enqueues a slice op with dynamic start indices onto the computation. @@ -838,8 +968,7 @@ class ComputationBuilder(object): shape = Shape(self.GetShape(mu).np_dtype, dims) return _wrap_data_handle( self._client.RngNormal( - _unwrap_data_handle(mu), _unwrap_data_handle(sigma), - _unwrap_shape(shape))) + _unwrap_data_handle(mu), _unwrap_data_handle(sigma), shape)) def RngUniform(self, a, b, dims): """Enqueues an RngUniform operation onto the computation. @@ -859,8 +988,7 @@ class ComputationBuilder(object): shape = Shape(self.GetShape(a).np_dtype, dims) return _wrap_data_handle( self._client.RngUniform( - _unwrap_data_handle(a), _unwrap_data_handle(b), - _unwrap_shape(shape))) + _unwrap_data_handle(a), _unwrap_data_handle(b), shape)) def While(self, cond, body, init): """Enqueues a While operation onto the computation. @@ -868,7 +996,7 @@ class ComputationBuilder(object): 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. """ @@ -877,11 +1005,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. @@ -992,7 +1167,7 @@ def initialize_replica_count(replica_count): Args: replica_count: number of replicas that are desired for set up during XLA - initalization. + initialization. Raises: A runtime exception if the XLA service has already been initialized. @@ -1018,3 +1193,13 @@ def GetPaddingConfigFromTriples(triples): 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 c0413b9bbc3b7f8b63e4cf7a8f24980322cffc47..4c16c1f8b07a28d8098e92e27f81a126ed9bdf0c 100644 --- a/tensorflow/compiler/xla/python/xla_client_test.py +++ b/tensorflow/compiler/xla/python/xla_client_test.py @@ -86,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): @@ -444,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") @@ -496,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]) @@ -699,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( @@ -714,6 +762,23 @@ class SingleOpTest(LocalComputationTest): [3, 2]) self._ExecuteAndCompareExact(c, expected=[[4, 5], [7, 8]]) + def testSliceInDim(self): + c = self._NewComputation() + c.SliceInDim( + c.Constant(NumpyArrayS32([[1, 2, 3], [4, 5, 6], [7, 8, 9]])), + start_index=1, + limit_index=2, + stride=1, + dimno=1) + self._ExecuteAndCompareExact(c, expected=[[2], [5], [8]]) + c.SliceInDim( + c.Constant(NumpyArrayS32([[1, 2, 3], [4, 5, 6], [7, 8, 9]])), + start_index=0, + limit_index=3, + stride=2, + dimno=0) + self._ExecuteAndCompareExact(c, expected=[[1, 2, 3], [7, 8, 9]]) + def testDynamicSlice(self): c = self._NewComputation() c.DynamicSlice( @@ -834,6 +899,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") @@ -974,6 +1046,14 @@ 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.]])), @@ -1172,6 +1252,28 @@ 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() diff --git a/tensorflow/compiler/xla/reference_util.cc b/tensorflow/compiler/xla/reference_util.cc index a9acdae380af5b7f9efb3d08302fc717108f5e40..ad3a28e11939d6259ebd75d544a950ba7abd741f 100644 --- a/tensorflow/compiler/xla/reference_util.cc +++ b/tensorflow/compiler/xla/reference_util.cc @@ -30,29 +30,23 @@ limitations under the License. namespace xla { -/* static */ std::unique_ptr> ReferenceUtil::TransposeArray2D( - const Array2D& operand) { - auto result = MakeUnique>(operand.width(), operand.height()); - for (int64 w = 0; w < operand.width(); ++w) { - for (int64 h = 0; h < operand.height(); ++h) { - (*result)(w, h) = operand(h, w); - } - } - - return result; -} - -/* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( - const Array2D& lhs, const Array2D& rhs) { +namespace { + +template +std::unique_ptr> MatmulArray2DImpl( + const Array2D& lhs, const Array2D& rhs, + const std::function& impl_fn) { CHECK_EQ(lhs.width(), rhs.height()); int m = lhs.height(); int n = rhs.width(); int k = lhs.width(); - auto result = MakeUnique>(m, n); + auto result = MakeUnique>(m, n); // Because Eigen is a header-oriented library, make sure that the Eigen code // is the same as the code used by the CPU backend (otherwise the linker will // randomly pick *some* definition). - __xla_cpu_runtime_EigenSingleThreadedMatMulF32( + impl_fn( /*run_options_ptr=*/nullptr, result->data(), rhs.data(), lhs.data(), n, m, k, /*transpose_lhs=*/0, @@ -60,22 +54,24 @@ namespace xla { return result; } +} // namespace + +/* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( + const Array2D& lhs, const Array2D& rhs) { + return MatmulArray2DImpl( + lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF16); +} + +/* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( + const Array2D& lhs, const Array2D& rhs) { + return MatmulArray2DImpl( + lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF32); +} + /* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( const Array2D& lhs, const Array2D& rhs) { - CHECK_EQ(lhs.width(), rhs.height()); - int m = lhs.height(); - int n = rhs.width(); - int k = lhs.width(); - auto result = MakeUnique>(m, n); - // Because Eigen is a header-oriented library, make sure that the Eigen code - // is the same as the code used by the CPU backend (otherwise the linker will - // randomly pick *some* definition). - __xla_cpu_runtime_EigenSingleThreadedMatMulF64( - /*run_options_ptr=*/nullptr, result->data(), rhs.data(), lhs.data(), n, m, - k, - /*transpose_lhs=*/0, - /*transpose_rhs=*/0); - return result; + return MatmulArray2DImpl( + lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF64); } /* static */ std::unique_ptr> ReferenceUtil::Array2DF32ToF64( @@ -188,18 +184,6 @@ ReferenceUtil::SeparableConvArray4D(const Array4D& input, return tensorflow::MathUtil::CeilOfRatio(unpadded_width, stride); } -/* static */ std::unique_ptr> -ReferenceUtil::ReduceWindow1DGeneric( - const tensorflow::gtl::ArraySlice& operand, float init, - const std::function& reduce_func, - const tensorflow::gtl::ArraySlice& window, - const tensorflow::gtl::ArraySlice& stride, Padding padding) { - std::vector dim_lengths{static_cast(operand.size())}; - return ReduceWindow1DGeneric( - operand, init, reduce_func, window, stride, - xla::MakePadding(dim_lengths, window, stride, padding)); -} - /* static */ std::unique_ptr> ReferenceUtil::ReduceWindow1DGeneric( const tensorflow::gtl::ArraySlice& operand, float init, @@ -239,23 +223,28 @@ ReferenceUtil::ReduceWindow1DAdd( const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, Padding padding) { const auto add_reduce = [](float arg1, float arg2) { return arg1 + arg2; }; - return ReduceWindow1DGeneric(operand, init, add_reduce, window, stride, - padding); + std::vector dim_lengths{static_cast(operand.size())}; + return ReduceWindow1DGeneric( + operand, init, add_reduce, window, stride, + xla::MakePadding(dim_lengths, window, stride, padding)); } -/* static */ std::unique_ptr> ReferenceUtil::ReduceWindow2DAdd( +/* static */ std::unique_ptr> +ReferenceUtil::ReduceWindow2DGeneric( const Array2D& operand, float init, + const std::function& reduce_func, const tensorflow::gtl::ArraySlice& window, - const tensorflow::gtl::ArraySlice& stride, Padding padding) { + const tensorflow::gtl::ArraySlice& stride, + const tensorflow::gtl::ArraySlice>& padding) { std::vector dim_lengths{operand.height(), operand.width()}; - 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) { + int64 padded_width = padding[i].first + dim_lengths[i] + padding[i].second; window_counts[i] = - WindowCount(dim_lengths[i], window[i], stride[i], padding); - pad_low[i] = padding_both[i].first; + window_util::StridedBound(padded_width, window[i], stride[i]); + pad_low[i] = padding[i].first; } auto result = MakeUnique>(window_counts[0], window_counts[1]); @@ -271,7 +260,7 @@ ReferenceUtil::ReduceWindow1DAdd( if (i0_base + i0_win >= 0 && i1_base + i1_win >= 0 && i0_base + i0_win < operand.n1() && i1_base + i1_win < operand.n2()) { - val += operand(i0_base + i0_win, i1_base + i1_win); + val = reduce_func(val, operand(i0_base + i0_win, i1_base + i1_win)); } } } @@ -281,6 +270,17 @@ ReferenceUtil::ReduceWindow1DAdd( return result; } +/* static */ std::unique_ptr> ReferenceUtil::ReduceWindow2DAdd( + const Array2D& operand, float init, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding) { + const auto add_reduce = [](float arg1, float arg2) { return arg1 + arg2; }; + std::vector dim_lengths{operand.height(), operand.width()}; + return ReduceWindow2DGeneric( + operand, init, add_reduce, window, stride, + xla::MakePadding(dim_lengths, window, stride, padding)); +} + /* static */ std::unique_ptr> ReferenceUtil::ReduceWindow3DAdd( const Array3D& operand, float init, const tensorflow::gtl::ArraySlice& window, @@ -472,7 +472,7 @@ ReferenceUtil::SelectAndScatter4DGePlus( i3_base + i3_win < operand.n4()) { float tmp = operand(i0_base + i0_win, i1_base + i1_win, i2_base + i2_win, i3_base + i3_win); - if (tmp >= val) { + if (tmp > val) { val = tmp; scatter_0 = i0_base + i0_win; scatter_1 = i1_base + i1_win; diff --git a/tensorflow/compiler/xla/reference_util.h b/tensorflow/compiler/xla/reference_util.h index 3ec96f2f38b8f91e1549419b60481327fa9bbd5f..28d6a8c3fe85fa4179bf2f41c82ad4eb93a045fe 100644 --- a/tensorflow/compiler/xla/reference_util.h +++ b/tensorflow/compiler/xla/reference_util.h @@ -39,10 +39,22 @@ namespace xla { class ReferenceUtil { public: // Returns the result of a transpose operation on the input matrix. - static std::unique_ptr> TransposeArray2D( - const Array2D& operand); + template + static std::unique_ptr> TransposeArray2D( + const Array2D& operand) { + auto result = MakeUnique>(operand.width(), operand.height()); + for (int64 w = 0; w < operand.width(); ++w) { + for (int64 h = 0; h < operand.height(); ++h) { + (*result)(w, h) = operand(h, w); + } + } + + return result; + } // Returns the result of a matrix multiply `lhs x rhs`. + static std::unique_ptr> MatmulArray2D( + const Array2D& lhs, const Array2D& rhs); static std::unique_ptr> MatmulArray2D( const Array2D& lhs, const Array2D& rhs); static std::unique_ptr> MatmulArray2D( @@ -187,9 +199,10 @@ class ReferenceUtil { const tensorflow::gtl::ArraySlice& operand, float init, const std::function& reduce_func, const tensorflow::gtl::ArraySlice& window, - const tensorflow::gtl::ArraySlice& stride, Padding padding); - static std::unique_ptr> ReduceWindow1DGeneric( - const tensorflow::gtl::ArraySlice& operand, float init, + const tensorflow::gtl::ArraySlice& stride, + const tensorflow::gtl::ArraySlice>& padding); + static std::unique_ptr> ReduceWindow2DGeneric( + const Array2D& operand, float init, const std::function& reduce_func, const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, @@ -215,6 +228,7 @@ class ReferenceUtil { // Performs select and scatter with Greater Than or equal as the select, plus // as the scatter, and Same Padding. + // TODO(b/74533103) Switch tests to evaluator and remove this implementation. static std::unique_ptr> SelectAndScatter4DGePlus( const Array4D& operand, const Array4D& source, float init, const tensorflow::gtl::ArraySlice& window, diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index 0f2d0a9e96e20007aa24a22832bdca4f0add372d..fba20c94cafea587bffcd766d1122d6327f32182 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -43,6 +43,118 @@ 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", + ":hlo_verifier", + "//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_propagation", + srcs = ["bfloat16_propagation.cc"], + hdrs = ["bfloat16_propagation.h"], + deps = [ + ":bfloat16_support", + ":hlo", + ":hlo_dataflow_analysis", + ":hlo_dce", + ":hlo_pass", + ":tuple_simplifier", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_tree", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "bfloat16_propagation_test", + srcs = ["bfloat16_propagation_test.cc"], + deps = [ + ":bfloat16_propagation", + ":bfloat16_support", + ":hlo", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep + ], +) + cc_library( name = "shape_inference", srcs = ["shape_inference.cc"], @@ -70,7 +182,8 @@ tf_cc_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep + "//tensorflow/core:lib", ], ) @@ -643,6 +756,7 @@ cc_library( hdrs = ["llvm_compiler.h"], deps = [ ":compiler", + "//tensorflow/core:lib_internal", "@llvm//:core", ], ) @@ -875,6 +989,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", "//tensorflow/core:lib", ], ) @@ -951,6 +1066,38 @@ tf_cc_test( ], ) +cc_library( + name = "hlo_module_group_metadata", + srcs = ["hlo_module_group_metadata.cc"], + hdrs = ["hlo_module_group_metadata.h"], + deps = [ + ":hlo", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + +cc_library( + name = "hlo_module_group_util", + srcs = ["hlo_module_group_util.cc"], + hdrs = ["hlo_module_group_util.h"], + deps = [ + ":hlo", + ":hlo_module_group_metadata", + ":hlo_reachability", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + cc_library( name = "hlo_scheduling", srcs = ["hlo_scheduling.cc"], @@ -1019,6 +1166,19 @@ tf_cc_test( ], ) +cc_library( + name = "hlo_creation_utils", + srcs = ["hlo_creation_utils.cc"], + hdrs = ["hlo_creation_utils.h"], + deps = [ + ":hlo", + ":shape_inference", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + ], +) + cc_library( name = "batchnorm_expander", srcs = ["batchnorm_expander.cc"], @@ -1027,7 +1187,6 @@ cc_library( ":hlo", ":hlo_pass", ":hlo_query", - ":shape_inference", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -1039,6 +1198,20 @@ cc_library( ], ) +cc_library( + name = "gather_expander", + srcs = ["gather_expander.cc"], + hdrs = ["gather_expander.h"], + deps = [ + ":hlo", + ":hlo_creation_utils", + ":hlo_pass", + ":while_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + ], +) + tf_cc_test( name = "batchnorm_expander_test", size = "small", @@ -1066,9 +1239,9 @@ cc_library( hdrs = ["algebraic_simplifier.h"], deps = [ ":hlo", + ":hlo_creation_utils", ":hlo_pass", ":hlo_query", - ":shape_inference", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -1102,6 +1275,41 @@ tf_cc_test( ], ) +cc_library( + name = "conditional_simplifier", + srcs = ["conditional_simplifier.cc"], + hdrs = ["conditional_simplifier.h"], + deps = [ + ":call_inliner", + ":hlo", + ":hlo_pass", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "conditional_simplifier_test", + srcs = ["conditional_simplifier_test.cc"], + deps = [ + ":conditional_simplifier", + ":hlo", + ":hlo_matchers", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + cc_library( name = "while_loop_simplifier", srcs = ["while_loop_simplifier.cc"], @@ -1124,6 +1332,7 @@ tf_cc_test( ":while_loop_simplifier", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/core:lib", "//tensorflow/core:test", ], ) @@ -2238,6 +2447,24 @@ cc_library( ":hlo", ":hlo_proto", "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:util", + ], +) + +tf_cc_test( + name = "hlo_proto_util_test", + srcs = ["hlo_proto_util_test.cc"], + deps = [ + ":hlo", + ":hlo_proto", + ":hlo_proto_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", ], ) @@ -2336,7 +2563,9 @@ cc_library( deps = [ ":call_inliner", ":hlo", + ":hlo_creation_utils", ":tuple_util", + "//tensorflow/core:lib", ], ) diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index ba82e822b216528c28536181059bc2417048de01..971c2935c835f6b10d1e4a7e0df42e5a7785c4ec 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -26,10 +26,10 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.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_creation_utils.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_query.h" -#include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -122,6 +122,8 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { Status HandleBitcast(HloInstruction* bitcast) override; + Status HandleBitcastConvert(HloInstruction* bitcast) override; + Status HandleBroadcast(HloInstruction* broadcast) override; Status HandleConcatenate(HloInstruction* concatenate) override; @@ -381,13 +383,9 @@ Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { !lhs->operand(0)->IsConstant() && lhs->operand(1)->IsConstant()) { auto* c1 = lhs->mutable_operand(1); auto* c2 = rhs; - TF_ASSIGN_OR_RETURN( - Shape sum_of_constants_shape, - ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, c1, c2)); - auto* sum_of_constants = - computation_->AddInstruction(HloInstruction::CreateBinary( - sum_of_constants_shape, HloOpcode::kAdd, c1, c2)); + TF_ASSIGN_OR_RETURN(auto* sum_of_constants, + MakeBinaryHlo(HloOpcode::kAdd, c1, c2)); return ReplaceWithNewInstruction( add, HloInstruction::CreateBinary(add->shape(), HloOpcode::kAdd, lhs->mutable_operand(0), @@ -411,6 +409,13 @@ Status AlgebraicSimplifierVisitor::HandleBitcast(HloInstruction* bitcast) { return Status::OK(); } +Status AlgebraicSimplifierVisitor::HandleBitcastConvert( + HloInstruction* bitcast) { + // Eliminate bitcast converts between same shape. + ReplaceInstructionIfSameShape(bitcast, bitcast->mutable_operand(0)); + return Status::OK(); +} + Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) { // If a copy feeds a copy, make it a single copy. if (copy->operand(0)->opcode() == HloOpcode::kCopy) { @@ -516,6 +521,18 @@ Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant) { return ReplaceInstruction( constant, BuildTupleConstant(computation_, constant->literal())); } + + // If a literal is all the same element replace it with a scalar broadcast. + if (ShapeUtil::ElementsIn(constant->shape()) > 1 && + constant->literal().IsAllFirst()) { + std::unique_ptr unique_scalar = + MakeUnique(constant->literal().GetFirstScalarLiteral()); + HloInstruction* scalar = computation_->AddInstruction( + HloInstruction::CreateConstant(std::move(unique_scalar))); + return ReplaceWithNewInstruction( + constant, + HloInstruction::CreateBroadcast(constant->shape(), scalar, {})); + } return Status::OK(); } @@ -619,32 +636,23 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { // (A / B) / (C / D) => (A / B)*(D / C) => (A * D) / (B * C) if (lhs->opcode() == HloOpcode::kDivide && rhs->opcode() == HloOpcode::kDivide) { - TF_ASSIGN_OR_RETURN( - const Shape a_times_d_shape, - ShapeInference::InferBinaryOpShape(HloOpcode::kMultiply, - lhs->operand(0), rhs->operand(1))); - auto a_times_d = computation_->AddInstruction(HloInstruction::CreateBinary( - a_times_d_shape, HloOpcode::kMultiply, lhs->mutable_operand(0), - rhs->mutable_operand(1))); - TF_ASSIGN_OR_RETURN( - const Shape b_times_c_shape, - ShapeInference::InferBinaryOpShape(HloOpcode::kMultiply, - lhs->operand(1), rhs->operand(0))); - auto b_times_c = computation_->AddInstruction(HloInstruction::CreateBinary( - b_times_c_shape, HloOpcode::kMultiply, lhs->mutable_operand(1), - rhs->mutable_operand(0))); - return ReplaceWithNewInstruction( - divide, HloInstruction::CreateBinary( - divide->shape(), HloOpcode::kDivide, a_times_d, b_times_c)); + TF_ASSIGN_OR_RETURN(auto a_times_d, MakeBinaryHlo(HloOpcode::kMultiply, + lhs->mutable_operand(0), + rhs->mutable_operand(1))); + TF_ASSIGN_OR_RETURN(auto b_times_c, MakeBinaryHlo(HloOpcode::kMultiply, + lhs->mutable_operand(1), + rhs->mutable_operand(0))); + TF_ASSIGN_OR_RETURN(auto new_divide, MakeBinaryHlo(HloOpcode::kDivide, + a_times_d, b_times_c)); + + return ReplaceInstruction(divide, new_divide); } // (A / B) / C => A / (B * C) if (lhs->opcode() == HloOpcode::kDivide) { - TF_ASSIGN_OR_RETURN(const Shape b_times_c_shape, - ShapeInference::InferBinaryOpShape( - HloOpcode::kMultiply, lhs->operand(1), rhs)); - auto b_times_c = computation_->AddInstruction(HloInstruction::CreateBinary( - b_times_c_shape, HloOpcode::kMultiply, lhs->mutable_operand(1), rhs)); + TF_ASSIGN_OR_RETURN( + auto b_times_c, + MakeBinaryHlo(HloOpcode::kMultiply, lhs->mutable_operand(1), rhs)); return ReplaceWithNewInstruction( divide, HloInstruction::CreateBinary(divide->shape(), HloOpcode::kDivide, @@ -653,11 +661,8 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { // A / (B / C) => (A*C) / B if (rhs->opcode() == HloOpcode::kDivide) { - TF_ASSIGN_OR_RETURN(const Shape a_times_c_shape, - ShapeInference::InferBinaryOpShape( - HloOpcode::kMultiply, lhs, rhs->operand(1))); - auto a_times_c = computation_->AddInstruction(HloInstruction::CreateBinary( - a_times_c_shape, HloOpcode::kMultiply, lhs, rhs->mutable_operand(1))); + TF_ASSIGN_OR_RETURN(auto a_times_c, MakeBinaryHlo(HloOpcode::kMultiply, lhs, + rhs->mutable_operand(1))); return ReplaceWithNewInstruction( divide, HloInstruction::CreateBinary(divide->shape(), HloOpcode::kDivide, @@ -1290,17 +1295,14 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { padding_dimension->set_edge_padding_high(0); } } - TF_ASSIGN_OR_RETURN(Shape nonzero_pad_shape, - ShapeInference::InferPadShape(pad->operand(0)->shape(), - pad->operand(1)->shape(), - nonzero_padding)); + + TF_ASSIGN_OR_RETURN(HloInstruction * nonzero_pad, + MakePadHlo(pad->mutable_operand(0), + pad->mutable_operand(1), nonzero_padding)); // Copy the layout from the original pad instructions. The new pad and the // slice instruction should all have the same layout. - TF_RETURN_IF_ERROR( - LayoutUtil::CopyLayoutBetweenShapes(pad->shape(), &nonzero_pad_shape)); - HloInstruction* nonzero_pad = computation_->AddInstruction( - HloInstruction::CreatePad(nonzero_pad_shape, pad->mutable_operand(0), - pad->mutable_operand(1), nonzero_padding)); + TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( + pad->shape(), nonzero_pad->mutable_shape())); // Second, construct the slice instruction to perform the negative padding. std::vector start_indices; @@ -1313,7 +1315,7 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { if (padding_dimension.edge_padding_low() < 0) { start = -1 * padding_dimension.edge_padding_low(); } - int64 end = nonzero_pad_shape.dimensions(i); + int64 end = nonzero_pad->shape().dimensions(i); if (padding_dimension.edge_padding_high() < 0) { end += padding_dimension.edge_padding_high(); } @@ -1322,16 +1324,14 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { strides.push_back(1); } - // Verify that the slice shape matches the pad shape. TF_ASSIGN_OR_RETURN( - Shape inferred_slice_shape, - ShapeInference::InferSliceShape(nonzero_pad_shape, start_indices, - end_indices, strides)); - TF_RET_CHECK(ShapeUtil::Compatible(inferred_slice_shape, pad->shape())); + HloInstruction * slice, + MakeSliceHlo(nonzero_pad, start_indices, end_indices, strides)); + + // Verify that the slice shape matches the pad shape. + TF_RET_CHECK(ShapeUtil::Compatible(slice->shape(), pad->shape())); - std::unique_ptr slice = HloInstruction::CreateSlice( - pad->shape(), nonzero_pad, start_indices, end_indices, strides); - return ReplaceWithNewInstruction(pad, std::move(slice)); + return ReplaceInstruction(pad, slice); } return Status::OK(); @@ -1604,6 +1604,14 @@ Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice( if (IsAll(start_indices, 0) && SameShape(dynamic_update_slice, update)) { return ReplaceInstruction(dynamic_update_slice, update); } + + // If any dimension of update is 0, elide the DynamicUpdateSlice. This + // optimization becomes invalid should we later prefer to warn about out of + // bound indices. + if (ShapeUtil::HasZeroElements(update->shape())) { + return ReplaceInstruction(dynamic_update_slice, + dynamic_update_slice->mutable_operand(0)); + } return Status::OK(); } @@ -1618,9 +1626,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) { diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc index e43ea50af45318adf2c95aa69b3e53a5225c5579..451294ef5d8367686d7fc22b7f5ebfde89d14d42 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc @@ -61,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); } @@ -83,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())); } @@ -110,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))); } @@ -133,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); } @@ -156,17 +152,47 @@ 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); } +TEST_F(AlgebraicSimplifierTest, ConstantToBroadcast) { + HloComputation::Builder builder(TestName()); + builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({3.14f, 3.14f, 3.14f}))); + + auto computation = module().AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Constant()); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_THAT(root, op::Broadcast(op::Constant())); + EXPECT_EQ(3.14f, root->operand(0)->literal().GetFirstElement()); +} + +TEST_F(AlgebraicSimplifierTest, ConstantNotToBroadcast) { + HloComputation::Builder builder(TestName()); + builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({3.14, 3.14, 4}))); + + auto computation = module().AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Constant()); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_FALSE(simplifier.Run(&module()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_THAT(root, op::Constant()); +} + // Test that A - 0 is simplified to A TEST_F(AlgebraicSimplifierTest, SubZero) { Shape r0f32 = ShapeUtil::MakeShape(F32, {}); @@ -178,13 +204,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); } @@ -200,13 +225,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))); } @@ -226,15 +250,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))); @@ -255,15 +278,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)); @@ -289,8 +311,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(), @@ -298,7 +319,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(), @@ -320,15 +341,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)))); @@ -349,15 +369,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)))); @@ -380,15 +399,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)))); @@ -411,12 +429,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))); @@ -438,11 +455,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))); } @@ -451,24 +467,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. @@ -482,13 +497,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); } @@ -504,13 +518,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); } @@ -529,13 +542,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); } @@ -554,13 +566,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); } @@ -579,13 +590,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); } @@ -607,13 +617,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)); } @@ -633,15 +642,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))); @@ -662,15 +670,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))); @@ -689,15 +696,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))); @@ -716,15 +722,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)); @@ -741,14 +746,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); } @@ -770,15 +774,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)); } @@ -795,14 +798,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()); @@ -820,14 +822,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()); @@ -849,14 +850,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); } @@ -872,14 +872,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)); } @@ -895,14 +894,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)); @@ -941,16 +939,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())); } @@ -969,7 +966,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"); @@ -980,20 +976,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())); } @@ -1014,14 +1010,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())); } @@ -1039,17 +1034,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. @@ -1060,14 +1054,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); } @@ -1081,14 +1074,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); } @@ -1102,14 +1094,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); } @@ -1132,8 +1123,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(), @@ -1141,7 +1131,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)); @@ -1163,15 +1153,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); } @@ -1188,14 +1177,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)); } @@ -1209,8 +1197,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}); @@ -1220,7 +1207,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)); @@ -1236,8 +1223,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}); @@ -1247,7 +1233,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); @@ -1268,14 +1254,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)); @@ -1314,8 +1299,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, @@ -1323,7 +1307,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( @@ -1344,8 +1328,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)); @@ -1353,7 +1336,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))); } @@ -1371,8 +1354,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)); @@ -1380,7 +1362,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)); @@ -1405,9 +1387,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 @@ -1424,14 +1405,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) { @@ -1448,14 +1429,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)); @@ -1475,14 +1455,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)); @@ -1501,15 +1480,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)); } @@ -1529,14 +1507,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)); } @@ -1554,14 +1531,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}), @@ -1576,17 +1552,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)); } @@ -1601,15 +1576,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)); } @@ -1623,15 +1597,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))); @@ -1646,15 +1619,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(), @@ -1670,15 +1642,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 = @@ -1696,15 +1667,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))); @@ -2410,12 +2380,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. @@ -2428,12 +2397,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())); } @@ -2453,11 +2421,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()); } @@ -2487,11 +2454,10 @@ 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())); } @@ -2554,15 +2520,16 @@ TEST_P(PadReduceWindowEffectiveBroadcastTest, DoIt) { 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( - ShapeUtil::MakeShape( - F32, decorate_spatials(param.reduce_window_spatials, 128, 2048)), - input, + pad_shape, input, builder.AddInstruction( HloInstruction::CreateConstant(Literal::CreateR0(0.0f))), padding)); - std::unique_ptr module = CreateNewModule(); HloComputation* add_computation = nullptr; { HloComputation::Builder builder(TestName() + ".add"); @@ -2573,24 +2540,24 @@ TEST_P(PadReduceWindowEffectiveBroadcastTest, DoIt) { 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()); } - TF_ASSERT_OK_AND_ASSIGN( - const Shape output_shape, - ShapeInference::InferPadShape(input_shape, ShapeUtil::MakeShape(F32, {}), - padding)); 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()); + 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( @@ -2667,11 +2634,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); @@ -2699,7 +2665,7 @@ struct DotOfConcatTestSpec { }; class DotOfConcatSimplificationTest - : public HloTestBase, + : public HloVerifiedTestBase, public ::testing::WithParamInterface {}; // Test that we transform @@ -2745,11 +2711,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( @@ -2790,17 +2755,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); @@ -2810,11 +2775,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)); @@ -2836,6 +2800,29 @@ DotOfConcatTestSpec kDotOfConcatTestSpecs[] = { {/*m=*/1, /*k=*/16, /*n=*/1}, // }; +// Test that DynamicUpdateSlice update param with any dimension equal to zero +// gets removed. +TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceZeroUpdate) { + HloComputation::Builder builder(TestName()); + const Shape dslice_shape = ShapeUtil::MakeShape(F32, {10}); + HloInstruction* const operand = builder.AddInstruction( + HloInstruction::CreateParameter(0, dslice_shape, "operand")); + const Shape update_shape = ShapeUtil::MakeShape(F32, {0}); + HloInstruction* const update = builder.AddInstruction( + HloInstruction::CreateParameter(1, update_shape, "update")); + HloInstruction* const start_indices = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0}))); + builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( + dslice_shape, operand, update, start_indices)); + const HloComputation* const computation = + module().AddEntryComputation(builder.Build()); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), operand); +} + INSTANTIATE_TEST_CASE_P(DotOfConcatSimplificationTestInstantiation, DotOfConcatSimplificationTest, ::testing::ValuesIn(kDotOfConcatTestSpecs)); diff --git a/tensorflow/compiler/xla/service/allocation_tracker.cc b/tensorflow/compiler/xla/service/allocation_tracker.cc index 4e80679c11dfdf7fdf8077a9f354139a4cab6803..4f819a743c48f30df8dde00ece72a0b4e1748802 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.cc +++ b/tensorflow/compiler/xla/service/allocation_tracker.cc @@ -34,40 +34,54 @@ StatusOr AllocationTracker::Register( std::unique_ptr shaped_buffer, const string& tag) { tensorflow::mutex_lock lock(mutex_); VLOG(2) << "Register"; - return RegisterInternal(std::move(shaped_buffer), tag); + std::vector> replicated_buffers; + replicated_buffers.emplace_back(std::move(shaped_buffer)); + return RegisterInternal(std::move(replicated_buffers), tag); +} + +StatusOr AllocationTracker::RegisterReplicatedBuffers( + std::vector> replicated_buffers, + const string& tag) { + tensorflow::mutex_lock lock(mutex_); + VLOG(2) << "RegisterReplicatedBuffers"; + return RegisterInternal(std::move(replicated_buffers), tag); } StatusOr AllocationTracker::RegisterInternal( - std::unique_ptr shaped_buffer, const string& tag) { + std::vector> replicated_buffers, + const string& tag) { VLOG(2) << "RegisterInternal(" - << "tag: \"" << tag << "\" " - << "shaped_buffer: " << *shaped_buffer; - if (shaped_buffer->platform() != backend_->platform()) { - return InvalidArgument( - "AllocationTracker for platform %s cannot register buffer from " - "platform %s", - backend_->platform()->Name().c_str(), - shaped_buffer->platform()->Name().c_str()); + << "tag: \"" << tag << "\" with " << replicated_buffers.size() + << " shaped_buffers."; + for (const auto& shaped_buffer : replicated_buffers) { + VLOG(2) << "shaped_buffer:" << *shaped_buffer; + if (shaped_buffer->platform() != backend_->platform()) { + return InvalidArgument( + "AllocationTracker for platform %s cannot register buffer from " + "platform %s", + backend_->platform()->Name().c_str(), + shaped_buffer->platform()->Name().c_str()); + } } int64 handle = next_handle_++; - std::vector shape_indices; - ShapeUtil::ForEachSubshape(shaped_buffer->on_device_shape(), - [this, &shape_indices](const Shape& /*subshape*/, - const ShapeIndex& index) { - shape_indices.push_back(index); - }); - for (const ShapeIndex& index : shape_indices) { - AddAllocationOrIncrementRefCount(shaped_buffer->buffer(index), - shaped_buffer->device_ordinal()); + for (auto& shaped_buffer : replicated_buffers) { + std::vector shape_indices; + ShapeUtil::ForEachSubshape(shaped_buffer->on_device_shape(), + [this, &shape_indices](const Shape& /*subshape*/, + const ShapeIndex& index) { + shape_indices.push_back(index); + }); + for (const ShapeIndex& index : shape_indices) { + AddAllocationOrIncrementRefCount(shaped_buffer->buffer(index), + shaped_buffer->device_ordinal()); + } + handle_to_shaped_buffers_[handle].emplace_back(std::move(shaped_buffer)); } + GlobalDataHandle result; result.set_handle(handle); - - handle_to_shaped_buffer_[handle] = std::move(shaped_buffer); - VLOG(2) << "handle: " << handle; - return result; } @@ -75,23 +89,35 @@ tensorflow::Status AllocationTracker::Unregister(const GlobalDataHandle& data) { tensorflow::mutex_lock lock(mutex_); VLOG(2) << "Unregister(" << "handle: " << data.handle() << ")"; - TF_ASSIGN_OR_RETURN(ShapedBuffer * shaped_buffer, ResolveInternal(data)); - std::vector shape_indices; - ShapeUtil::ForEachSubshape(shaped_buffer->on_device_shape(), - [this, &shape_indices](const Shape& /*subshape*/, - const ShapeIndex& index) { - shape_indices.push_back(index); - }); - for (const ShapeIndex& index : shape_indices) { - TF_RETURN_IF_ERROR(DecrementRefCount(shaped_buffer->buffer(index), - shaped_buffer->device_ordinal())); + TF_ASSIGN_OR_RETURN(std::vector replicated_buffers, + ResolveInternal(data)); + for (const auto& shaped_buffer : replicated_buffers) { + std::vector shape_indices; + ShapeUtil::ForEachSubshape(shaped_buffer->on_device_shape(), + [this, &shape_indices](const Shape& /*subshape*/, + const ShapeIndex& index) { + shape_indices.push_back(index); + }); + for (const ShapeIndex& index : shape_indices) { + TF_RETURN_IF_ERROR(DecrementRefCount(shaped_buffer->buffer(index), + shaped_buffer->device_ordinal())); + } } + return Reset(data); +} - // Keep a nullptr as a tombstone for unregistered handles. This enables better - // error messages. That is, "handle has been deallocated" versus "handle does - // not exist". - handle_to_shaped_buffer_.at(data.handle()).reset(); - +Status AllocationTracker::Reset(const GlobalDataHandle& data) { + // Keep a nullptr as a tombstone for unregistered handles. This enables + // better error messages. That is, "handle has been deallocated" versus + // "handle does not exist". + auto it = handle_to_shaped_buffers_.find(data.handle()); + if (it == handle_to_shaped_buffers_.end()) { + return NotFound("no allocation record for global data handle: %lld", + data.handle()); + } + for (auto& shaped_buffer : it->second) { + shaped_buffer.reset(); + } return tensorflow::Status::OK(); } @@ -99,7 +125,11 @@ StatusOr> AllocationTracker::DeconstructTuple( const GlobalDataHandle& data) { tensorflow::mutex_lock lock(mutex_); - TF_ASSIGN_OR_RETURN(ShapedBuffer * shaped_buffer, ResolveInternal(data)); + TF_ASSIGN_OR_RETURN(std::vector replicated_buffers, + ResolveInternal(data)); + // We only need to care about replica id 0 here, since the GlobalDataHandle is + // the same for all buffers across replicas. + const ShapedBuffer* shaped_buffer = replicated_buffers[0]; if (!ShapeUtil::IsTuple(shaped_buffer->on_host_shape())) { return InvalidArgument("global data handle %lld is not a tuple", data.handle()); @@ -109,7 +139,7 @@ StatusOr> AllocationTracker::DeconstructTuple( TF_RET_CHECK(ShapeUtil::IsTuple(shaped_buffer->on_device_shape())); if (ShapeUtil::IsNestedTuple(shaped_buffer->on_device_shape())) { - return Unimplemented("deconstructing nested tuples not yet supported"); + return Unimplemented("Deconstructing nested tuples is not implemented."); } std::vector element_handles; @@ -122,37 +152,55 @@ StatusOr> AllocationTracker::DeconstructTuple( shaped_buffer->platform(), shaped_buffer->device_ordinal()); element_buffer->set_buffer(shaped_buffer->buffer(/*index=*/{i}), /*index=*/{}); + std::vector> replicated_buffers; + replicated_buffers.emplace_back(std::move(element_buffer)); TF_ASSIGN_OR_RETURN( GlobalDataHandle element_handle, - RegisterInternal(std::move(element_buffer), "deconstructed tuple")); + RegisterInternal(std::move(replicated_buffers), "deconstructed tuple")); element_handles.push_back(element_handle); } return std::move(element_handles); } -StatusOr AllocationTracker::Resolve( +StatusOr> AllocationTracker::Resolve( const GlobalDataHandle& data) { tensorflow::mutex_lock lock(mutex_); return AllocationTracker::ResolveInternal(data); } -StatusOr AllocationTracker::ResolveInternal( +StatusOr AllocationTracker::ResolveForReplica( + const GlobalDataHandle& data, int replica_id) { + tensorflow::mutex_lock lock(mutex_); + TF_ASSIGN_OR_RETURN(std::vector replicated_buffers, + ResolveInternal(data)); + if (replica_id >= replicated_buffers.size()) { + return InvalidArgument( + "Requesting buffer for replica %d, but found buffers only for %lu " + "replicas.", + replica_id, replicated_buffers.size()); + } + return replicated_buffers[replica_id]; +} + +StatusOr> AllocationTracker::ResolveInternal( const GlobalDataHandle& data) { VLOG(2) << "resolve:" << data.handle(); - auto it = handle_to_shaped_buffer_.find(data.handle()); - if (it == handle_to_shaped_buffer_.end()) { + auto it = handle_to_shaped_buffers_.find(data.handle()); + if (it == handle_to_shaped_buffers_.end()) { return NotFound("no allocation record for global data handle: %lld", data.handle()); } - ShapedBuffer* shaped_buffer = it->second.get(); - - if (shaped_buffer == nullptr) { - return InvalidArgument("global data handle %lld was previously deallocated", - data.handle()); + std::vector replicated_buffers; + for (const auto& shaped_buffer : it->second) { + if (shaped_buffer == nullptr) { + return InvalidArgument( + "global data handle %lld was previously deallocated", data.handle()); + } + replicated_buffers.push_back(shaped_buffer.get()); } - return shaped_buffer; + return replicated_buffers; } void AllocationTracker::AddAllocationOrIncrementRefCount( diff --git a/tensorflow/compiler/xla/service/allocation_tracker.h b/tensorflow/compiler/xla/service/allocation_tracker.h index 807af8694972083d097604a67ee46d2f73d9545a..038aee8541b297d6f91fe2b3bce7455fd9a7084e 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.h +++ b/tensorflow/compiler/xla/service/allocation_tracker.h @@ -43,10 +43,17 @@ class AllocationTracker { AllocationTracker(Backend* backend) : backend_(backend), next_handle_(1) {} // Registers a shaped buffer of device memory, and returns a corresponding - // handle that can be used for talking to XLA clients. + // handle that can be used for talking to XLA clients. The given shaped buffer + // will be treated as the buffer corresponding to the only replica. StatusOr Register( std::unique_ptr shaped_buffer, const string& tag); + // Registers a vector of shaped buffers of device memory, one per replica, and + // returns a corresponding handle that can be used for talking to XLA clients. + StatusOr RegisterReplicatedBuffers( + std::vector> replicated_buffers, + const string& tag); + // Unregister the allocation for the given data handle. Status Unregister(const GlobalDataHandle& data); @@ -54,9 +61,17 @@ class AllocationTracker { StatusOr> DeconstructTuple( const GlobalDataHandle& Data); - // Resolve a handle from an XLA client to a shaped buffer, or provide an error - // status to say whether it was not found (or found, but found deallocated). - StatusOr Resolve(const GlobalDataHandle& data); + // Resolve a handle from an XLA client to a vector of shaped buffers, one per + // replica, or provide an error status to say whether any of those buffers + // were not found (or found, but found deallocated). + StatusOr> Resolve( + const GlobalDataHandle& data); + + // Resolves a handle from an XLA client and replica id to a shaped buffer, or + // provide an error status to say whether it was not found (or found, but + // found deallocated). + StatusOr ResolveForReplica(const GlobalDataHandle& data, + int replica_id); private: // Data structure encapsulating single memory allocation on the device. @@ -74,13 +89,17 @@ class AllocationTracker { // Internal helper which resolves the given GlobalDataHandle to a // ShapedBuffer. - StatusOr ResolveInternal(const GlobalDataHandle& data) - EXCLUSIVE_LOCKS_REQUIRED(mutex_); + StatusOr> ResolveInternal( + const GlobalDataHandle& data) EXCLUSIVE_LOCKS_REQUIRED(mutex_); - // Internal helper which registers a shaped buffer. + // Internal helper which registers a vector of shaped buffers, one per + // replica. StatusOr RegisterInternal( - std::unique_ptr shaped_buffer, const string& tag) - EXCLUSIVE_LOCKS_REQUIRED(mutex_); + std::vector> replicated_buffers, + const string& tag) EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + // Resets the shaped buffers corresponding to the given handle. + Status Reset(const GlobalDataHandle& data) EXCLUSIVE_LOCKS_REQUIRED(mutex_); // Adds the given device address to the allocation tracker, or if it already // exists, then increment it's reference count. @@ -111,9 +130,10 @@ class AllocationTracker { tensorflow::gtl::FlatMap opaque_to_allocation_map_ GUARDED_BY(mutex_); - // A map from data handle to ShapedBuffer. - tensorflow::gtl::FlatMap> - handle_to_shaped_buffer_ GUARDED_BY(mutex_); + // A map from data handle to a vector of shaped buffers that represent the + // buffers for different replicas. + tensorflow::gtl::FlatMap>> + handle_to_shaped_buffers_ GUARDED_BY(mutex_); TF_DISALLOW_COPY_AND_ASSIGN(AllocationTracker); }; diff --git a/tensorflow/compiler/xla/service/batchnorm_expander.cc b/tensorflow/compiler/xla/service/batchnorm_expander.cc index 27ddfd47aa3096afd3e245af1ac3cedd9b48ce4a..38086bd7e121847be6b6b69415cfe87814e7fc24 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander.cc @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_query.h" -#include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -153,6 +152,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormTraining( std::vector added_instructions; auto add = [&](std::unique_ptr inst) { HloInstruction* added_inst = computation_->AddInstruction(std::move(inst)); + added_inst->set_metadata(batch_norm->metadata()); added_instructions.push_back(added_inst); return added_inst; }; @@ -334,6 +334,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormInference( std::vector added_instructions; auto add = [&](std::unique_ptr inst) { HloInstruction* added_inst = computation_->AddInstruction(std::move(inst)); + added_inst->set_metadata(batch_norm->metadata()); added_instructions.push_back(added_inst); return added_inst; }; @@ -419,6 +420,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( std::vector added_instructions; auto add = [&](std::unique_ptr inst) { HloInstruction* added_inst = computation_->AddInstruction(std::move(inst)); + added_inst->set_metadata(batch_norm->metadata()); added_instructions.push_back(added_inst); return added_inst; }; 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..432448e9bbc7db30ed67a0130d52b060032362d5 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc @@ -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. +==============================================================================*/ + +#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::kSend || // + hlo->opcode() == HloOpcode::kSendDone || // + hlo->opcode() == HloOpcode::kRecv || // + hlo->opcode() == HloOpcode::kRecvDone || // + 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..14c54ddd135af024327f63418b410da1ed3c4fd4 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_normalization.cc @@ -0,0 +1,381 @@ +/* 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()); + int64 f32_count = 0; + int64 bf16_count = 0; + bool has_unsupported_bf16_operand = false; + bool has_unsupported_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) { + f32_count += 1; + } else if (operand_types[i] == BF16) { + bf16_count += 1; + if (!bfloat16_support_->SupportsBF16Operand(*crs, i)) { + has_unsupported_bf16_operand = true; + } + } + if (output_types[i] == F32) { + f32_count += 1; + } else if (output_types[i] == BF16) { + bf16_count += 1; + if (!bfloat16_support_->SupportsBF16Output(*crs)) { + has_unsupported_bf16_output = true; + } + } + } + + if (bf16_count == 0) { + return Status::OK(); + } + + auto should_convert_operand = [&](int64 i) { + if (operand_types[i] != BF16) { + return false; + } + if (!bfloat16_support_->SupportsBF16Operand(*crs, i)) { + return true; + } + if (bfloat16_support_->SupportsMixedPrecisions(*crs)) { + return false; + } + return has_unsupported_bf16_operand || has_unsupported_bf16_output || + f32_count > 0; + }; + + for (int64 i = 0; i < crs->operand_count(); ++i) { + if (should_convert_operand(i)) { + TF_RETURN_IF_ERROR(InsertConvertBeforeOperand(crs, i, F32, computation_)); + f32_count += 1; + bf16_count -= 1; + } + } + + if (!has_unsupported_bf16_output && + (bfloat16_support_->SupportsMixedPrecisions(*crs) || f32_count == 0 || + bf16_count == 0)) { + return Status::OK(); + } + + std::vector materialized_users = crs->users(); + 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)); + + // 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) { + int f32_count = 0; + int bf16_count = 1; + + for (int64 i = 0; i < hlo->operand_count(); ++i) { + if (hlo->operand(i)->shape().element_type() == F32) { + f32_count += 1; + } else if (hlo->operand(i)->shape().element_type() == BF16) { + bf16_count += 1; + } + } + + if (hlo->shape().element_type() == F32) { + f32_count += 1; + } else if (hlo->shape().element_type() == BF16) { + bf16_count += 1; + } + + std::vector bf16_called_comps; + for (auto* comp : hlo->called_computations()) { + bool comp_has_bf16 = false; + if (comp->root_instruction()->shape().element_type() == F32) { + f32_count += 1; + } else if (comp->root_instruction()->shape().element_type() == BF16) { + bf16_count += 1; + comp_has_bf16 = true; + } + for (auto* param : comp->parameter_instructions()) { + if (param->shape().element_type() == F32) { + f32_count += 1; + } else if (param->shape().element_type() == BF16) { + bf16_count += 1; + comp_has_bf16 = true; + } + } + if (comp_has_bf16) { + bf16_called_comps.push_back(comp); + } + } + + // Resolve unsupported BF16 operands. + for (int i = 0; i < hlo->operand_count(); ++i) { + if (hlo->operand(i)->shape().element_type() == BF16 && + !bfloat16_support_->SupportsBF16Operand(*hlo, i)) { + TF_RETURN_IF_ERROR(InsertConvertBeforeOperand(hlo, i, F32, computation_)); + bf16_count -= 1; + f32_count += 1; + } + } + + // Resolve unsupported BF16 output. + if (hlo->shape().element_type() == BF16 && + !bfloat16_support_->SupportsBF16Output(*hlo)) { + TF_RETURN_IF_ERROR( + ChangeOutputTypeThenInsertConvertBack(hlo, F32, computation_)); + bf16_count -= 1; + f32_count += 1; + } + + // Resolve unsupported mixed precision after resolving unsupported BF16 + // operands and output, because the numbers of BF16 operands/output and F32 + // operands/output may have changed. + if (bfloat16_support_->SupportsMixedPrecisions(*hlo) || bf16_count == 0 || + f32_count == 0) { + return Status::OK(); + } + // 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 = 0; i < hlo->operand_count(); ++i) { + if (hlo->operand(i)->shape().element_type() == BF16) { + continue; + } + if ((bfloat16_support_->EffectiveOperandPrecisionIsBF16(*hlo, i) || + bfloat16_support_->EffectiveOperandPrecisionIsOutputPrecision(*hlo, + i)) && + bfloat16_support_->SupportsBF16Operand(*hlo, i)) { + continue; + } + can_use_bf16 = false; + break; + } + if (can_use_bf16) { + for (int i = 0; i < hlo->operand_count(); ++i) { + if (hlo->operand(i)->shape().element_type() == F32) { + 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 = 0; i < hlo->operand_count(); ++i) { + if (hlo->operand(i)->shape().element_type() == BF16) { + TF_RETURN_IF_ERROR(InsertConvertBeforeOperand(hlo, i, F32, computation_)); + } + } + return ConvertCalledComputations(hlo, bf16_called_comps); +} + +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..1afaefd9df9c5771fb9e134ae9050f3abb00ea4a --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc @@ -0,0 +1,287 @@ +/* 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/service/hlo_verifier.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; + } + if (hlo.opcode() == HloOpcode::kDot) { + // Test that only the first operand of kDot supports BF16. + return operand_index == 0; + } + 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::kDot || 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()); + + HloVerifier verifier(/*allow_mixed_precision=*/true); + EXPECT_IS_OK(verifier.Run(module).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, {}); + + 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); +} + +// Tests that the normalization should not cause unsupported mixed precision due +// to resolving unsupported BF16 operand. +TEST_F(BFloat16NormalizationTest, DoNotAddUnsupportedMixedPrecision) { + auto builder = HloComputation::Builder(TestName()); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {4, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, bf16_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_shape, "b")); + + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_contracting_dimensions(1); + dot_dnums.add_rhs_contracting_dimensions(0); + HloInstruction* dot = builder.AddInstruction( + HloInstruction::CreateDot(bf16_shape, a, b, dot_dnums)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(Normalize(module.get())); + + EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kConvert); + EXPECT_EQ(dot->shape().element_type(), F32); + EXPECT_EQ(dot->operand(0)->shape().element_type(), F32); + EXPECT_EQ(dot->operand(0)->opcode(), HloOpcode::kConvert); + EXPECT_EQ(dot->operand(1)->shape().element_type(), F32); + EXPECT_EQ(dot->operand(1)->opcode(), HloOpcode::kConvert); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.cc b/tensorflow/compiler/xla/service/bfloat16_propagation.cc new file mode 100644 index 0000000000000000000000000000000000000000..531f36e8c5473ef684e654ed6b89c4d5ef04b401 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_propagation.cc @@ -0,0 +1,683 @@ +/* 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_propagation.h" + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/map_util.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_module.h" +#include "tensorflow/compiler/xla/service/tuple_simplifier.h" +#include "tensorflow/compiler/xla/shape_tree.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/lib/gtl/cleanup.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +BFloat16Propagation::BFloat16Propagation( + const BFloat16Support* bfloat16_support) + : bfloat16_support_(bfloat16_support) {} + +void BFloat16Propagation::DetermineAndMutateFusionComputationPrecision( + HloInstruction* fusion) { + CHECK_EQ(fusion->opcode(), HloOpcode::kFusion); + if (!bfloat16_support_->SupportsMixedPrecisions(*fusion)) { + return; + } + + // We are depending on the fusion node itself having already been analyzed + // for whether it can output BF16 and this has been adjusted in the output + // shape, and now we're looking to update the interior of the fusion node to + // match the new output shape, as well as recursively process the whole fusion + // node even if the output shape was not modified. + auto root = fusion->fused_instructions_computation()->root_instruction(); + + // Adjust root's element types according to the fusion's output shape. + ShapeUtil::ForEachMutableSubshape( + root->mutable_shape(), [&](Shape* subshape, const ShapeIndex& index) { + if (subshape->element_type() != F32) { + return; + } + if (ShapeUtil::GetSubshape(fusion->shape(), index).element_type() == + BF16) { + subshape->set_element_type(BF16); + changed_ = true; + VLOG(2) << "Fused root " << root->ToString() << " at shape index " + << index << " changed to BF16 precision for fusion " + << fusion->ToString(); + } + }); + + // Propagate BF16 in the fusion computation. + auto insts = + fusion->fused_instructions_computation()->MakeInstructionPostOrder(); + for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) { + DetermineAndMutateInstructionPrecision(*inst_it, /*skip_parameters=*/false); + } + computations_visited_in_mutation_pass_.insert( + fusion->fused_instructions_computation()); +} + +void BFloat16Propagation::DetermineAndMutateWhileComputationsPrecision( + HloInstruction* while_hlo) { + CHECK_EQ(while_hlo->opcode(), HloOpcode::kWhile); + + // We are depending on the while node itself having already been analyzed for + // whether it can output BF16 and this has been adjusted in the output shape, + // and now we're looking to update the body and condition computations to + // match the new output shape, as well as recursively process the whole while + // node even if the output shape was not modified. + HloComputation* body = while_hlo->while_body(); + auto body_root = body->root_instruction(); + HloComputation* condition = while_hlo->while_condition(); + + ShapeUtil::ForEachMutableSubshape( + body_root->mutable_shape(), + [this, while_hlo, body_root](Shape* subshape, const ShapeIndex& index) { + if (subshape->element_type() != F32) { + return; + } + if (ShapeUtil::GetSubshape(while_hlo->shape(), index).element_type() == + BF16) { + subshape->set_element_type(BF16); + changed_ = true; + VLOG(2) << "While body root " << body_root->ToString() + << " at shape index " << index + << " changed to BF16 precision for while " + << while_hlo->ToString(); + } + }); + + auto body_insts = body->MakeInstructionPostOrder(); + for (auto inst_it = body_insts.rbegin(); inst_it != body_insts.rend(); + ++inst_it) { + DetermineAndMutateInstructionPrecision(*inst_it, /*skip_parameters=*/false); + } + computations_visited_in_mutation_pass_.insert(body); + + auto condition_insts = condition->MakeInstructionPostOrder(); + for (auto inst_it = condition_insts.rbegin(); + inst_it != condition_insts.rend(); ++inst_it) { + DetermineAndMutateInstructionPrecision(*inst_it, /*skip_parameters=*/false); + } + computations_visited_in_mutation_pass_.insert(condition); +} + +bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo, + const ShapeIndex& index) const { + auto value_set = dataflow_->GetValueSet(&hlo, index); + for (const HloValue* value : value_set.values()) { + if (ContainsKey(values_that_must_be_kept_as_f32_, value)) { + return false; + } + if (value->shape().element_type() == BF16) { + continue; + } + for (const HloUse& use : value->uses()) { + if (!ContainsKey(instructions_visited_in_mutation_pass_, + use.instruction)) { + // We don't know yet whether use.instruction will consume BF16 since it + // hasn't been visited. Although we visit instructions in reverse + // topological order, this is still possible because there may be + // unvisited instruction that alias the same buffer. In this case, we + // aggressively skip this use, and if this causes inconsistency (e.g., + // one use is in BF16 but another use is in F32), it will be resolved at + // the end of the BFloat16Propagation pass. + continue; + } + // Any visited user that can accept BF16 has already been updated if + // necessary, e.g., the output has been changed to BF16 if it propagates + // precision, or a called computation's parameters have been changed to + // BF16 for fusions or whiles. + if (use.instruction->opcode() == HloOpcode::kFusion) { + const auto* fused_parameter = + use.instruction->fused_parameter(use.operand_number); + if (ShapeUtil::GetSubshape(fused_parameter->shape(), use.operand_index) + .element_type() != BF16) { + return false; + } + continue; + } else if (use.instruction->opcode() == HloOpcode::kWhile) { + const auto* cond_parameter = + use.instruction->while_condition()->parameter_instruction( + use.operand_number); + if (ShapeUtil::GetSubshape(cond_parameter->shape(), use.operand_index) + .element_type() != BF16) { + return false; + } + const auto* body_parameter = + use.instruction->while_body()->parameter_instruction( + use.operand_number); + if (ShapeUtil::GetSubshape(body_parameter->shape(), use.operand_index) + .element_type() != BF16) { + return false; + } + continue; + } + if (bfloat16_support_->EffectiveOperandPrecisionIsBF16( + *use.instruction, use.operand_number)) { + continue; + } + // If the op propagates precision and it outputs a BF16, then it's OK to + // supply BF16 also as the input. In the backward mutation pass, the users + // shapes should have already been processed. + PrimitiveType user_output_type = PRIMITIVE_TYPE_INVALID; + if (use.instruction->opcode() == HloOpcode::kTuple || + (use.instruction->opcode() == HloOpcode::kCrossReplicaSum && + ShapeUtil::IsTuple(use.instruction->shape()))) { + user_output_type = ShapeUtil::GetSubshape( + ShapeUtil::GetSubshape(use.instruction->shape(), + {use.operand_number}), + use.operand_index) + .element_type(); + } else { + user_output_type = use.instruction->shape().element_type(); + } + if (bfloat16_support_->EffectiveOperandPrecisionIsOutputPrecision( + *use.instruction, use.operand_number) && + user_output_type == BF16) { + continue; + } + return false; + } + } + return true; +} + +void BFloat16Propagation::DetermineAndMutateInstructionPrecision( + HloInstruction* hlo, bool skip_parameters) { + // We handle any fusion computation or while body/condition after the + // instruction is handled, because we need to know the output shape of a + // fusion or while before propagating inside its computations. + bool postpone_processing_called_computations = false; + auto cleaner = tensorflow::gtl::MakeCleanup( + [this, hlo, &postpone_processing_called_computations] { + if (!postpone_processing_called_computations) { + if (hlo->opcode() == HloOpcode::kFusion) { + DetermineAndMutateFusionComputationPrecision(hlo); + } else if (hlo->opcode() == HloOpcode::kWhile) { + DetermineAndMutateWhileComputationsPrecision(hlo); + } + } + instructions_visited_in_mutation_pass_.insert(hlo); + }); + + if (hlo->opcode() == HloOpcode::kWhile && + (caller_counts_[hlo->while_condition()] > 1 || + caller_counts_[hlo->while_body()] > 1)) { + postpone_processing_called_computations = true; + return; + } + + // Do not change precision for instructions related to entry and exit of a + // computation, and control flow, because this pass might break the interfaces + // or assumptions for them. + if (hlo->opcode() == HloOpcode::kInfeed || // + hlo->opcode() == HloOpcode::kOutfeed || // + hlo->opcode() == HloOpcode::kSend || // + hlo->opcode() == HloOpcode::kSendDone || // + hlo->opcode() == HloOpcode::kRecv || // + hlo->opcode() == HloOpcode::kRecvDone || // + hlo->opcode() == HloOpcode::kCustomCall || // + hlo->opcode() == HloOpcode::kCall || // + hlo->opcode() == HloOpcode::kConditional || // + (hlo->opcode() == HloOpcode::kParameter && skip_parameters)) { + return; + } + + // Prevent root instructions from having their output modified by recording + // all F32 output values as needing to stay as F32. + CHECK(hlo->parent() != nullptr); + if (hlo == hlo->parent()->root_instruction()) { + if (!hlo->parent()->IsFusionComputation()) { + ShapeUtil::ForEachSubshape(hlo->shape(), [&](const Shape& subshape, + const ShapeIndex& index) { + if (subshape.element_type() != F32) { + return; + } + for (const auto* value : dataflow_->GetValueSet(hlo, index).values()) { + // Since we use HloValues from the dataflow analysis, this can also + // affect HLO instructions beyond the root, e.g., if the root is a + // Tuple HLO, then its operands are also affected. + values_that_must_be_kept_as_f32_.insert(value); + } + }); + } + return; + } + + if (!ContainsKey(consider_using_bfloat16_, hlo)) { + return; + } + + if (!bfloat16_support_->SupportsBF16Output(*hlo)) { + return; + } + + ShapeUtil::ForEachMutableSubshape( + hlo->mutable_shape(), + [hlo, this](Shape* subshape, const ShapeIndex& index) { + if (subshape->element_type() == F32 && + AllUsersConsumeBF16(*hlo, index)) { + subshape->set_element_type(BF16); + changed_ = true; + VLOG(2) << "HloInstruction output at shape index " << index + << " changed to BF16 precision: " << hlo->ToString(); + } + }); +} + +bool BFloat16Propagation::InstructionIsCandidateForBF16Output( + HloInstruction* hlo) { + if (!bfloat16_support_->SupportsMixedPrecisions(*hlo) && + hlo->opcode() != HloOpcode::kTuple && + hlo->opcode() != HloOpcode::kGetTupleElement && + hlo->shape().element_type() != BF16) { + for (int64 i = 0; i < hlo->operand_count(); ++i) { + if (!bfloat16_support_->EffectiveOperandPrecisionIsOutputPrecision(*hlo, + i) || + !ContainsKey(consider_using_bfloat16_, hlo->operand(i))) { + return false; + } + } + } + return true; +} + +void BFloat16Propagation::AdjustCalledComputationParameters( + HloInstruction* hlo) { + auto adjust_computation = + [this, hlo](HloComputation* computation, + tensorflow::gtl::ArraySlice operands) { + // Adjust parameters. + CHECK_EQ(operands.size(), computation->num_parameters()); + for (int64 i = 0; i < operands.size(); ++i) { + auto parameter = computation->parameter_instruction(i); + ShapeUtil::ForEachMutableSubshape( + parameter->mutable_shape(), + [this, i, hlo, &operands, parameter](Shape* subshape, + const ShapeIndex& index) { + if (!ShapeUtil::IsLeafIndex(parameter->shape(), index)) { + return; + } + PrimitiveType operand_type = + ShapeUtil::GetSubshape(operands[i]->shape(), index) + .element_type(); + if (subshape->element_type() == operand_type) { + return; + } + CHECK(operand_type == F32 || operand_type == BF16); + subshape->set_element_type(operand_type); + changed_ = true; + VLOG(2) << "Called computation parameter " + << parameter->ToString() << " at shape index " << index + << " adjusted to match operand in HLO " + << hlo->ToString(); + }); + } + }; + + switch (hlo->opcode()) { + case HloOpcode::kFusion: + adjust_computation(hlo->fused_instructions_computation(), + hlo->operands()); + break; + case HloOpcode::kWhile: + adjust_computation(hlo->while_condition(), hlo->operands()); + adjust_computation(hlo->while_body(), hlo->operands()); + break; + default: + break; + } +} + +void BFloat16Propagation::AdjustCalledComputationRoot(HloInstruction* hlo) { + auto adjust_computation = [this, hlo](HloComputation* computation, + const Shape& output_shape) { + // Adjust root. + HloInstruction* root = computation->root_instruction(); + ShapeUtil::ForEachMutableSubshape( + root->mutable_shape(), [this, hlo, root, &output_shape]( + Shape* subshape, const ShapeIndex& index) { + if (!ShapeUtil::IsLeafIndex(hlo->shape(), index)) { + return; + } + const PrimitiveType output_type = + ShapeUtil::GetSubshape(output_shape, index).element_type(); + if (subshape->element_type() == output_type) { + return; + } + CHECK(output_type == F32 || output_type == BF16); + subshape->set_element_type(output_type); + // It's possible that output_type is F32, but the root instruction's + // type is BF16; e.g., a fusion node's output was changed to BF16 + // initially but then adjusted back to F32, and the fusion computation + // is now being adjusted after the fusion node. + if (output_type == F32) { + for (const auto* value : + dataflow_->GetValueSet(root, index).values()) { + // We rely on the fact that this adjustment works in reverse + // topological order so that called computation will be + // processed later. Adding the value to + // values_that_must_be_kept_as_f32_ will ensure the + // correctness of the adjustment for HLOs that will be + // processed later. + values_that_must_be_kept_as_f32_.insert(value); + } + } + changed_ = true; + VLOG(2) << "Called computation root " << root->ToString() + << " at shape index " << index + << " adjusted to match output shape of " << hlo->ToString(); + }); + }; + + switch (hlo->opcode()) { + case HloOpcode::kFusion: + adjust_computation(hlo->fused_instructions_computation(), hlo->shape()); + break; + case HloOpcode::kWhile: + adjust_computation(hlo->while_condition(), hlo->shape()); + adjust_computation(hlo->while_body(), hlo->shape()); + break; + default: + break; + } +} + +bool BFloat16Propagation::ResolveInconsistencyOfAliasingBuffersHelper( + HloComputation* computation, + tensorflow::gtl::FlatSet* visited_computations) { + bool parameter_changed = false; + auto insts = computation->MakeInstructionPostOrder(); + // Do the adjustment on each instruction in the computation in reverse + // topological order. + for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) { + auto hlo = *inst_it; + auto adjust_hlo_output = [this, hlo, ¶meter_changed]( + Shape* subshape, const ShapeIndex& index) { + if (subshape->element_type() != F32 && subshape->element_type() != BF16) { + return; + } + PrimitiveType type = BF16; + for (const auto* value : dataflow_->GetValueSet(hlo, index).values()) { + if (value->shape().element_type() == BF16) { + continue; + } + CHECK_EQ(value->shape().element_type(), F32); + type = F32; + break; + } + // It's possible that a user has been changed from BF16 to F32 + // during this final adjustment pass, so we need to check + // AllUsersConsumeBF16() again. + if (type == BF16 && !AllUsersConsumeBF16(*hlo, index)) { + type = F32; + } + if (type == F32) { + for (const auto* value : dataflow_->GetValueSet(hlo, index).values()) { + // We rely on the fact that this adjustment works in reverse + // topological order. Adding the value to + // values_that_must_be_kept_as_f32_ will ensure the correctness + // of the adjustment for HLOs that will be processed later. + values_that_must_be_kept_as_f32_.insert(value); + } + } + if (type != subshape->element_type()) { + subshape->set_element_type(type); + VLOG(2) << "HloInstruction output at shape index " << index + << " adjusted to " << *subshape << ": " << hlo->ToString(); + if (hlo->opcode() == HloOpcode::kParameter) { + parameter_changed = true; + } + } + }; + ShapeUtil::ForEachMutableSubshape(hlo->mutable_shape(), adjust_hlo_output); + AdjustCalledComputationRoot(hlo); + if (hlo->opcode() == HloOpcode::kWhile) { + // We need to run on the while body and condition repeatedly until a fixed + // point is reached, i.e., the parameters do not change any more. We may + // need more than one iteration because the while input and output alias + // each other, so changing one input parameter requires changing the + // corresponding output element and thus may transitively require changing + // another input parameter. A fixed point will be reached because the + // parameters can only be changed from BF16 to F32, not the other way + // around. + tensorflow::gtl::FlatSet visited_in_while; + while (ResolveInconsistencyOfAliasingBuffersHelper(hlo->while_condition(), + &visited_in_while) || + ResolveInconsistencyOfAliasingBuffersHelper(hlo->while_body(), + &visited_in_while)) { + visited_in_while.clear(); + ShapeUtil::ForEachMutableSubshape(hlo->mutable_shape(), + adjust_hlo_output); + AdjustCalledComputationRoot(hlo); + } + visited_computations->insert(visited_in_while.begin(), + visited_in_while.end()); + } + } + // Now adjust parameters of called computations. + for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) { + AdjustCalledComputationParameters(*inst_it); + } + return parameter_changed; +} + +Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers( + HloModule* module) { + std::list computations_topological_order = + module->MakeComputationPostOrder(); + tensorflow::gtl::FlatSet resolved; + for (auto comp_it = computations_topological_order.rbegin(); + comp_it != computations_topological_order.rend(); ++comp_it) { + if (ContainsKey(resolved, *comp_it)) { + continue; + } + ResolveInconsistencyOfAliasingBuffersHelper(*comp_it, &resolved); + } + + // We could have changed a fusion computation's root shape to have a different + // precision than the fusion node's output, if the fusion root does not + // define a buffer (e.g., a tuple). Now we add conversions after such fusion + // roots to make them match the fusion output. If the fusion output is a + // (possibly nested) tuple, we first create get-tuple-elements, then convert + // the unmatching leaf nodes, and finally create a new tuple as the fusion + // computation's root. If tuples and get-tuple-elements are created, we will + // run tuple simplifier and dead code elimination at the end (dead code is not + // allowed in fusion computation). E.g., + // + // (1) (2) (3) + // a b a b a b + // |\ | |\ | |\ | + // \ add -> |add -> | add + // \ | \ | convert | + // tuple tuple \ | + // / \ tuple + // gte gte + // | | + // convert | + // \ / + // tuple + // (1) a is F32 but tuple is BF16 + // (2) after adding conversion + // (3) after tuple simplifier and DCE. + bool needs_tuple_simplifier = false; + for (auto computation : computations_topological_order) { + auto insts = computation->MakeInstructionPostOrder(); + for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) { + auto hlo = *inst_it; + if (hlo->opcode() != HloOpcode::kFusion) { + continue; + } + auto fusion_computation = hlo->fused_instructions_computation(); + auto fusion_root = fusion_computation->root_instruction(); + if (ShapeUtil::Compatible(fusion_root->shape(), hlo->shape())) { + continue; + } + ShapeTree converted_outputs(hlo->shape()); + // Iterate through nodes in the shape tree in pre-order and initialize + // each non-root node with a corresponding get-tuple-element. For a leaf + // node, if its shape does not match the fusion output, create a + // conversion node to overwrite the node value. + for (auto it = converted_outputs.begin(); it != converted_outputs.end(); + ++it) { + ShapeIndex output_index = it->first; + HloInstruction*& output = it->second; + const Shape subshape = + ShapeUtil::GetSubshape(hlo->shape(), output_index); + if (output_index.empty()) { + output = fusion_root; + } else { + ShapeIndex parent_index = output_index; + parent_index.pop_back(); + output = fusion_computation->AddInstruction( + HloInstruction::CreateGetTupleElement( + subshape, converted_outputs.element(parent_index), + output_index.back())); + } + if (ShapeUtil::IsTuple(subshape)) { + continue; + } + if (!ShapeUtil::Compatible( + subshape, + ShapeUtil::GetSubshape(fusion_root->shape(), output_index))) { + output = fusion_computation->AddInstruction( + HloInstruction::CreateConvert(subshape, output)); + } + } + // Iterate through nodes in the shape tree in reverse pre-order and create + // a tuple instruction for each non-leaf node where the elements are the + // values of its child nodes. + for (auto it = converted_outputs.rbegin(); it != converted_outputs.rend(); + ++it) { + ShapeIndex output_index = it->first; + HloInstruction*& output = it->second; + const Shape& subshape = + ShapeUtil::GetSubshape(hlo->shape(), output_index); + if (!ShapeUtil::IsTuple(subshape)) { + continue; + } + std::vector elements( + ShapeUtil::TupleElementCount(subshape)); + ShapeIndex child_index = output_index; + for (int64 i = 0; i < elements.size(); ++i) { + child_index.push_back(i); + elements[i] = converted_outputs.element(child_index); + child_index.pop_back(); + } + output = fusion_computation->AddInstruction( + HloInstruction::CreateTuple(elements)); + } + fusion_computation->set_root_instruction(converted_outputs.element({})); + needs_tuple_simplifier |= ShapeUtil::IsTuple(hlo->shape()); + } + } + + // We may have converted some constants from F32 to BF16, so adjust the + // constant literals in such cases. We do this here instead of when the + // constant node's is changed because 1) the HloInstruction interface does not + // allow resetting the literal so we have to create a new kConstant + // instruction to replace the old one, which invalidates dataflow analysis, + // and 2) it's possible that a kConstant's output gets changed to BF16 at the + // beginning but later on adjusted back to F32, so converting literals here + // can avoid repeated conversions. + // + // TODO(b/73833576): Consider resetting literal in HloInstruction. + bool needs_dce = needs_tuple_simplifier; + for (auto computation : computations_topological_order) { + for (auto hlo : computation->MakeInstructionPostOrder()) { + if (hlo->opcode() != HloOpcode::kConstant) { + continue; + } + if (!ShapeUtil::Equal(hlo->literal().shape(), hlo->shape())) { + TF_ASSIGN_OR_RETURN(auto converted_literal, + hlo->literal().ConvertToShape(hlo->shape())); + auto new_constant = computation->AddInstruction( + HloInstruction::CreateConstant(std::move(converted_literal))); + TF_RETURN_IF_ERROR(hlo->ReplaceAllUsesWith(new_constant)); + needs_dce = true; + } + } + } + + if (needs_tuple_simplifier) { + TupleSimplifier tuple_simplifier; + TF_RETURN_IF_ERROR(tuple_simplifier.Run(module).status()); + } + if (needs_dce) { + HloDCE dce; + TF_RETURN_IF_ERROR(dce.Run(module).status()); + } + return Status::OK(); +} + +// The algorithm first does a forward pass (parameters to root) to determine a +// set of instructions to consider using bfloat16, then does a backward pass to +// determine the precisions of those instructions according to the need of +// their users. +StatusOr BFloat16Propagation::Run(HloModule* module) { + TF_ASSIGN_OR_RETURN(dataflow_, HloDataflowAnalysis::Run(*module)); + + std::list computations_topological_order = + module->MakeComputationPostOrder(); + // The first step is a forward pass (parameters to root), where we determine + // the potential candidate instructions to use bfloat16 in the outputs that + // are not likely to cause overhead from extra explicit conversions. This is + // done forwardly because we determine whether an HLO is a candidate partially + // based on whether its operands are candidates. + for (auto computation : computations_topological_order) { + for (auto inst : computation->MakeInstructionPostOrder()) { + if (InstructionIsCandidateForBF16Output(inst)) { + consider_using_bfloat16_.insert(inst); + } + } + } + + // The second step is a backward pass (root to parameters), where we modify + // the precisions of the instructions identified in the first step when + // feasible. This is done backwardly because we determine the precision of an + // HLO's output based on how it is later used. + // + // The precision of an instruction is determined by its users, so we do the + // propagation in reverse topological order. + for (auto comp_it = computations_topological_order.rbegin(); + comp_it != computations_topological_order.rend(); ++comp_it) { + if ((*comp_it)->IsFusionComputation()) { + // Fusion computations are handled when visiting the fusion instruction. + continue; + } + auto insts = (*comp_it)->MakeInstructionPostOrder(); + for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) { + DetermineAndMutateInstructionPrecision(*inst_it, + /*skip_parameters=*/true); + } + } + + if (!changed_) { + return false; + } + + // It's possible that an instruction does not define a buffer, but the + // defining instruction's shape has changed. So we need to adjust the output + // shapes of instructions according to the HLO values they refer to. + TF_RETURN_IF_ERROR(ResolveInconsistencyOfAliasingBuffers(module)); + return true; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.h b/tensorflow/compiler/xla/service/bfloat16_propagation.h new file mode 100644 index 0000000000000000000000000000000000000000..89a5ac5db1549877a135182ae8df57fa6bf9d579 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_propagation.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_PROPAGATION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_PROPAGATION_H_ + +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/service/bfloat16_support.h" +#include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// HLO pass which reduces the precision of some HLO instructions to BF16 +// according to the backend-specific BFloat16Support rule provided by the +// caller. +// +// This pass can be used to reduce instruction precision without affecting the +// numerical accuracy of the module, i.e., the final output of the module would +// be bitwise identical to that without this pass; this is possible if the +// backend already reduces precision to BF16 on some HLO instructions. +// +// This pass will not modify the signature of a computation, unless it is a +// fusion computation or its only caller is a while. +// +// !!! WARNING !!! This pass can introduce mixed precision in individual HLOs, +// which has two issues: +// +// 1) It does not guarantee to respect the passed-in BFloat16Support +// specification in terms of mixed precision, so the backend may not support an +// HLO that has mixed precision produced by this pass. To address this issue, +// run BFloat16Normalization with the same BFloat16Support after this pass. +// +// 2) In general, mixed precision may break the assumptions of some other HLO +// passes even if the specific backend supports the individual HLOs. Such +// assumptions include that there are no HLOs using mixed precision, or that the +// precision of an HLO's output is determined by its inputs. It should be used +// at the end of the HLO optimization pipeline but before +// BFloat16ConversionFolding. If other passes are needed after this pass, run +// BFloat16MixedPrecisionRemoval first to undo some of the changes made by this +// pass. +class BFloat16Propagation : public HloPassInterface { + public: + explicit BFloat16Propagation(const BFloat16Support* bfloat16_support); + + ~BFloat16Propagation() override = default; + + tensorflow::StringPiece name() const override { + return "bfloat16-propagation"; + } + + // Runs the pass on the given module. Returns whether the module was changed + // (precision reductions were added). + StatusOr Run(HloModule* module) override; + + private: + // *************************** + // Function called and state produced by the forward analysis pass (from + // parameters to root) that determines the candidate HLOs to use BF16 outputs. + + // Determines whether we should consider changing the precision of the given + // instruction in the forward pass. + bool InstructionIsCandidateForBF16Output(HloInstruction* hlo); + + // The set of instructions to consider using bfloat16, computed in the forward + // pass. + tensorflow::gtl::FlatSet consider_using_bfloat16_; + + // *************************** + // Functions called and state produced by the backward mutation pass (from + // root to parameters). + + // Determines the precision for the given instruction in the mutation pass. + void DetermineAndMutateInstructionPrecision(HloInstruction* hlo, + bool skip_parameters); + + // Special handling in the mutation pass for fusion computations. + // + // Precondition: hlo->opcode() == kFusion + void DetermineAndMutateFusionComputationPrecision(HloInstruction* fusion); + + // Special handling in the mutation pass for while computations. + // + // Precondition: hlo->opcode() == kWhile + void DetermineAndMutateWhileComputationsPrecision(HloInstruction* while_hlo); + + // The set of HloInstructions that have been visited in the mutation pass. + tensorflow::gtl::FlatSet + instructions_visited_in_mutation_pass_; + + // The set of HloComputations that have been visited in the mutation pass. + tensorflow::gtl::FlatSet + computations_visited_in_mutation_pass_; + + // *************************** + // Functions called by the final inconsistency resolving pass. + + // Adjusts the output shapes of HloInstructions such that if two + // HloInstructions have aliasing buffers in their outputs, they must have the + // same precision. + Status ResolveInconsistencyOfAliasingBuffers(HloModule* module); + + // Resolves inconsistency of aliasing buffers for the given computation, and + // recursively runs on a while instruction's condition and body until a fixed + // point is reached. + bool ResolveInconsistencyOfAliasingBuffersHelper( + HloComputation* computation, + tensorflow::gtl::FlatSet* visited_computations); + + // Makes the parameters of called computations match how they are called by + // the given HLO. + void AdjustCalledComputationParameters(HloInstruction* hlo); + + // Makes the root instructions of called computations match how they are used + // by the given HLO. + void AdjustCalledComputationRoot(HloInstruction* hlo); + + // *************************** + // Functions called and state used by two or more passes. + + // Returns whether all uses of the given HloInstruction can consume BF16 + // input. + bool AllUsersConsumeBF16(const HloInstruction& hlo, + const ShapeIndex& index) const; + + // The set of F32 HLO values that must be kept in F32. + tensorflow::gtl::FlatSet values_that_must_be_kept_as_f32_; + + // Mapping from each HloComputation to the number of callers to it in the + // module. Populated at the beginning of this pass. + tensorflow::gtl::FlatMap caller_counts_; + + const BFloat16Support* bfloat16_support_; + std::unique_ptr dataflow_; + + bool changed_ = false; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_PROPAGATION_H_ diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..5950b004b3da439c442eec6e5e09ea2307fcb018 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc @@ -0,0 +1,620 @@ +/* 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_propagation.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/tests/literal_test_util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +// A class specifying the BF16 support used to test the propagation pass. It +// specifies that BF16 and mixed precision are supported in all HloInstructions, +// and that kDot reduces its operands precision to BF16. +class TestBFloat16Support : public BFloat16Support { + public: + TestBFloat16Support() {} + ~TestBFloat16Support() override {} + + 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 true; + } + + bool EffectiveOperandPrecisionIsBF16(const HloInstruction& hlo, + int64 operand_index) const override { + return hlo.opcode() == HloOpcode::kDot; + } +}; + +class BFloat16PropagationTest : public HloTestBase { + protected: + // Runs the propagation pass on the given module, and returns whether the + // module is changed after this pass. + bool PropagatePrecision(HloModule* module) { + TestBFloat16Support bfloat16_support; + BFloat16Propagation propagation(&bfloat16_support); + StatusOr result = propagation.Run(module); + EXPECT_IS_OK(result.status()); + return result.ValueOrDie(); + } + + // Returns whether the given HloInstruction's output element type is BF16 or + // the only use of it is converting to BF16. + bool OutputsBF16(const HloInstruction* inst) { + if (inst->shape().element_type() == BF16) { + return true; + } + return inst->user_count() == 1 && + inst->users()[0]->opcode() == HloOpcode::kConvert && + inst->users()[0]->shape().element_type() == BF16; + } +}; + +// Tests that BF16 can propagate through select over non-tuple buffers, but not +// through add where reducing operand precision can affect the result. +TEST_F(BFloat16PropagationTest, PropagateThroughSelectButNotAdd) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* c = + builder.AddInstruction(HloInstruction::CreateParameter(2, shape, "c")); + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add0, b)); + HloInstruction* pred = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kEq, a, b)); + HloInstruction* sel = builder.AddInstruction( + HloInstruction::CreateTernary(shape, HloOpcode::kSelect, pred, c, add1)); + HloInstruction* xpose = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {4, 2}), sel, {1, 0})); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, xpose, a)); + HloInstruction* root = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, dot, dot)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), root); + EXPECT_TRUE(OutputsBF16(xpose)); + EXPECT_TRUE(OutputsBF16(sel)); + EXPECT_TRUE(OutputsBF16(add1)); + EXPECT_FALSE(OutputsBF16(add0)); + EXPECT_FALSE(OutputsBF16(a)); + EXPECT_FALSE(OutputsBF16(b)); + EXPECT_FALSE(OutputsBF16(c)); +} + +// Tests that if a constant is converted to BF16 then its literal must also be +// converted. +TEST_F(BFloat16PropagationTest, ConvertConstantLiteral) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); + Array2D array_a(4, 4); + array_a.FillUnique(1.0f); + Array2D array_b(4, 4); + array_b.FillUnique(10.0f); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateFromArray(array_a))); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateFromArray(array_b))); + HloInstruction* dot = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kDot, a, b)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), dot); + EXPECT_TRUE(OutputsBF16(dot->operand(0))); + EXPECT_TRUE(OutputsBF16(dot->operand(1))); + EXPECT_EQ(dot->operand(0)->opcode(), HloOpcode::kConstant); + EXPECT_EQ(dot->operand(1)->opcode(), HloOpcode::kConstant); + LiteralTestUtil::ExpectEqual( + dot->operand(0)->literal(), + *LiteralTestUtil::ConvertF32ToBF16(*Literal::CreateFromArray(array_a))); + LiteralTestUtil::ExpectEqual( + dot->operand(1)->literal(), + *LiteralTestUtil::ConvertF32ToBF16(*Literal::CreateFromArray(array_b))); +} + +// Tests that BF16 can be propagated through nested tuples. +TEST_F(BFloat16PropagationTest, PropagateThroughTuples) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, a)); + HloInstruction* add2 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, b, b)); + HloInstruction* xpose = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {4, 2}), add1, {1, 0})); + + HloInstruction* tuple0 = + builder.AddInstruction(HloInstruction::CreateTuple({add0, add1, add2})); + HloInstruction* tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({tuple0, xpose})); + + HloInstruction* lhs = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(xpose->shape(), tuple1, 1)); + HloInstruction* rhs = + builder.AddInstruction(HloInstruction::CreateGetTupleElement( + add0->shape(), + builder.AddInstruction(HloInstruction::CreateGetTupleElement( + tuple0->shape(), tuple1, 0)), + 0)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, lhs, rhs)); + + HloInstruction* output_tuple = + builder.AddInstruction(HloInstruction::CreateTuple({dot, add2})); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), output_tuple); + EXPECT_TRUE(OutputsBF16(xpose)); + EXPECT_TRUE(OutputsBF16(add0)); + EXPECT_TRUE(OutputsBF16(add1)); + EXPECT_FALSE(OutputsBF16(add2)); +} + +// Tests that even if an instruction does not define a buffer in its output, its +// shape must match the defining instruction. +TEST_F(BFloat16PropagationTest, SameValueReferencedTwice) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, a)); + + HloInstruction* lhs = builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {4, 2}), add1, {1, 0})); + + HloInstruction* tuple = + builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); + HloInstruction* rhs = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(add1->shape(), tuple, 1)); + + // lhs is the transpose of add1, and rhs is a get-tuple-element aliasing add1. + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, lhs, rhs)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), dot); + EXPECT_TRUE(OutputsBF16(add0)); + EXPECT_TRUE(OutputsBF16(add1)); + EXPECT_TRUE(OutputsBF16(lhs)); + // rhs is a get-tuple-element, which does not define a buffer, but its shape + // should also be adjusted accordingly. + EXPECT_TRUE(OutputsBF16(rhs)); +} + +// Tests that a non-fusion computation's root should not be changed. +TEST_F(BFloat16PropagationTest, DoNotChangeComputationRoot) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); + + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, add, add)); + + HloInstruction* tuple = + builder.AddInstruction(HloInstruction::CreateTuple({add, dot})); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_FALSE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), tuple); + EXPECT_FALSE(OutputsBF16(add)); +} + +// Tests that BF16 is propagated properly through fused computations. +TEST_F(BFloat16PropagationTest, PropagateThroughFusion) { + auto module = CreateNewModule(); + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, param)); + + auto builder_f0 = HloComputation::Builder("fusion0"); + HloInstruction* a_f0 = + builder_f0.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b_f0 = + builder_f0.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* tuple_f0 = + builder_f0.AddInstruction(HloInstruction::CreateTuple({a_f0, b_f0})); + auto comp_f0 = module->AddEmbeddedComputation(builder_f0.Build()); + auto fusion0 = builder.AddInstruction(HloInstruction::CreateFusion( + tuple_f0->shape(), HloInstruction::FusionKind::kCustom, {add, add}, + comp_f0)); + + auto builder_f1 = HloComputation::Builder("fusion1"); + HloInstruction* p_f1 = builder_f1.AddInstruction( + HloInstruction::CreateParameter(0, tuple_f0->shape(), "param")); + HloInstruction* a_f1 = builder_f1.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, p_f1, 0)); + HloInstruction* b_f1 = builder_f1.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, p_f1, 1)); + HloInstruction* dot = builder_f1.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, a_f1, b_f1)); + auto comp_f1 = module->AddEmbeddedComputation(builder_f1.Build()); + auto fusion1 = builder.AddInstruction(HloInstruction::CreateFusion( + dot->shape(), HloInstruction::FusionKind::kCustom, {fusion0}, comp_f1)); + + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), fusion1); + EXPECT_TRUE(OutputsBF16(add)); + EXPECT_TRUE(OutputsBF16(a_f0)); + EXPECT_TRUE(OutputsBF16(b_f0)); + EXPECT_TRUE(OutputsBF16(a_f1)); + EXPECT_TRUE(OutputsBF16(b_f1)); +} + +// Tests that if 1) the root instruction of a fusion is a tuple, 2) the fusion +// outputs are only used by a dot, and 3) one element of the tuple is used by +// an add in the fusion computation, then the propagation pass should create a +// convert in the fusion computation to keep the add's operand in F32 but change +// the fusion output to BF16. E.g., the following fusion computation +// (F32, F32) fusion_computation(F32 a, F32 b) +// = tuple(F32 a, F32 add(F32 a, F32 b)) +// will be changed to +// (BF16, BF16) fusion_computation(F32 a, F32 b) +// = tuple(BF16 convert(a), BF16 add(F32 a, F32 b)) +TEST_F(BFloat16PropagationTest, ConvertTupleFusionElementIfUsedByAdd) { + auto module = CreateNewModule(); + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); + + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, param)); + + auto builder_f = HloComputation::Builder("fusion0"); + HloInstruction* a_f = + builder_f.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b_f = + builder_f.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* add_f = builder_f.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a_f, b_f)); + HloInstruction* tuple_f = + builder_f.AddInstruction(HloInstruction::CreateTuple({a_f, add_f})); + auto comp_f = module->AddEmbeddedComputation(builder_f.Build()); + auto fusion = builder.AddInstruction(HloInstruction::CreateFusion( + tuple_f->shape(), HloInstruction::FusionKind::kCustom, {add, add}, + comp_f)); + + HloInstruction* gte0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, fusion, 0)); + HloInstruction* gte1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, fusion, 1)); + HloInstruction* dot = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kDot, gte0, gte1)); + + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), dot); + EXPECT_TRUE(OutputsBF16(gte0)); + EXPECT_TRUE(OutputsBF16(gte1)); + EXPECT_FALSE(OutputsBF16(a_f)); + EXPECT_FALSE(OutputsBF16(b_f)); + EXPECT_TRUE(OutputsBF16(add_f)); + auto new_fusion_root = comp_f->root_instruction(); + EXPECT_EQ(new_fusion_root->opcode(), HloOpcode::kTuple); + EXPECT_EQ(new_fusion_root->operand(1), add_f); + EXPECT_EQ(new_fusion_root->operand(0)->opcode(), HloOpcode::kConvert); + EXPECT_TRUE(OutputsBF16(new_fusion_root->operand(0))); +} + +// A select over tuples does not define the leaf buffers, so the types in +// on_true and on_false must match, so that as long as one of them is F32, the +// other must be F32 as well. +TEST_F(BFloat16PropagationTest, SelectOverTuples) { + auto module = CreateNewModule(); + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + HloInstruction* pred = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(PRED, {}), "pred")); + + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, param)); + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add0, param)); + HloInstruction* tuple0 = + builder.AddInstruction(HloInstruction::CreateTuple({param, add0})); + HloInstruction* tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({param, add1})); + HloInstruction* sel = builder.AddInstruction(HloInstruction::CreateTernary( + tuple0->shape(), HloOpcode::kSelect, pred, tuple0, tuple1)); + HloInstruction* gte0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, sel, 0)); + HloInstruction* gte1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, sel, 1)); + HloInstruction* xpose = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {4, 2}), gte0, {1, 0})); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, xpose, gte1)); + + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), dot); + EXPECT_FALSE(OutputsBF16(add0)); + EXPECT_FALSE(OutputsBF16(add1)); + EXPECT_FALSE(OutputsBF16(gte0)); + EXPECT_FALSE(OutputsBF16(gte1)); + EXPECT_TRUE(OutputsBF16(xpose)); +} + +// Tests that BF16 is propagated properly through while computations. +TEST_F(BFloat16PropagationTest, PropagateThroughWhile) { + auto module = CreateNewModule(); + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); + + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, shape, "param1")); + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + HloInstruction* tuple = + builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); + + auto builder_cond = HloComputation::Builder("cond"); + auto cond_param = builder_cond.AddInstruction( + HloInstruction::CreateParameter(0, tuple->shape(), "cond_param")); + auto cond_lhs = builder_cond.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, cond_param, 0)); + auto cond_rhs = builder_cond.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, cond_param, 1)); + // This add should prevent RHS from using BF16 + auto cond_add_rhs = builder_cond.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, cond_rhs, cond_rhs)); + auto cond_dot = builder_cond.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, cond_lhs, cond_add_rhs)); + builder_cond.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kGt, + builder_cond.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {}), cond_dot, {0, 0}, {1, 1}, {1, 1})), + builder_cond.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {}), cond_dot, {1, 1}, {2, 2}, {1, 1})))); + auto cond = module->AddEmbeddedComputation(builder_cond.Build()); + + auto builder_body = HloComputation::Builder("body"); + auto body_param = builder_body.AddInstruction( + HloInstruction::CreateParameter(0, tuple->shape(), "body_param")); + auto body_lhs = builder_body.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, body_param, 0)); + auto body_rhs = builder_body.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, body_param, 1)); + auto body_dot = builder_body.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kDot, body_lhs, body_rhs)); + builder_body.AddInstruction( + HloInstruction::CreateTuple({body_dot, body_rhs})); + auto body = module->AddEmbeddedComputation(builder_body.Build()); + + auto while_hlo = builder.AddInstruction( + HloInstruction::CreateWhile(tuple->shape(), cond, body, tuple)); + + auto lhs = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, while_hlo, 0)); + auto rhs = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, while_hlo, 1)); + auto dot = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kDot, lhs, rhs)); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), dot); + EXPECT_TRUE(OutputsBF16(lhs)); + EXPECT_FALSE(OutputsBF16(rhs)); + EXPECT_TRUE(OutputsBF16(body_dot)); + EXPECT_TRUE(OutputsBF16(body_lhs)); + EXPECT_FALSE(OutputsBF16(body_rhs)); + EXPECT_TRUE(OutputsBF16(cond_lhs)); + EXPECT_FALSE(OutputsBF16(cond_rhs)); + EXPECT_TRUE(OutputsBF16(add0)); + EXPECT_FALSE(OutputsBF16(add1)); +} + +// Tests that BF16 is not propagated through multiple whiles that invoke the +// same computation as long as one while prevents the propagation. +TEST_F(BFloat16PropagationTest, DoNotPropagateWhilesCallingSameComputation) { + auto module = CreateNewModule(); + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); + + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, shape, "param1")); + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + HloInstruction* add2 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + HloInstruction* add3 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + HloInstruction* tuple0 = + builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); + HloInstruction* tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({add2, add3})); + + // Condition computation for the first while. + auto builder_cond0 = HloComputation::Builder("cond0"); + auto cond0_param = builder_cond0.AddInstruction( + HloInstruction::CreateParameter(0, tuple0->shape(), "cond0_param")); + auto cond0_lhs = builder_cond0.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, cond0_param, 0)); + auto cond0_rhs = builder_cond0.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, cond0_param, 1)); + // This add should prevent RHS from using BF16 + auto cond0_add_rhs = + builder_cond0.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kAdd, cond0_rhs, cond0_rhs)); + auto cond0_dot = builder_cond0.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, cond0_lhs, cond0_add_rhs)); + builder_cond0.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kGt, + builder_cond0.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {}), cond0_dot, {0, 0}, {1, 1}, {1, 1})), + builder_cond0.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {}), cond0_dot, {1, 1}, {2, 2}, {1, 1})))); + auto cond0 = module->AddEmbeddedComputation(builder_cond0.Build()); + + // Condition computation for the second while. + auto builder_cond1 = HloComputation::Builder("cond1"); + auto cond1_param = builder_cond1.AddInstruction( + HloInstruction::CreateParameter(0, tuple1->shape(), "cond1_param")); + auto cond1_lhs = builder_cond1.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, cond1_param, 0)); + auto cond1_rhs = builder_cond1.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, cond1_param, 1)); + // This add should prevent LHS from using BF16 + auto cond1_add_lhs = + builder_cond1.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kAdd, cond1_lhs, cond1_lhs)); + auto cond1_dot = builder_cond1.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, cond1_add_lhs, cond1_rhs)); + builder_cond1.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kGt, + builder_cond1.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {}), cond1_dot, {0, 0}, {1, 1}, {1, 1})), + builder_cond1.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {}), cond1_dot, {1, 1}, {2, 2}, {1, 1})))); + auto cond1 = module->AddEmbeddedComputation(builder_cond1.Build()); + + // Body computation shared by both whiles. + auto builder_body = HloComputation::Builder("body"); + auto body_param = builder_body.AddInstruction( + HloInstruction::CreateParameter(0, tuple0->shape(), "body_param")); + auto body_lhs = builder_body.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, body_param, 0)); + auto body_rhs = builder_body.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, body_param, 1)); + auto body_dot = builder_body.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kDot, body_lhs, body_rhs)); + builder_body.AddInstruction( + HloInstruction::CreateTuple({body_dot, body_rhs})); + auto body = module->AddEmbeddedComputation(builder_body.Build()); + + auto while0 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple0->shape(), cond0, body, tuple0)); + auto while1 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple1->shape(), cond1, body, tuple1)); + + auto lhs = builder.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, + builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, while0, 0)), + builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, while0, 1)))); + auto rhs = builder.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, + builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, while1, 0)), + builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, while1, 1)))); + auto dot = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kDot, lhs, rhs)); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + EXPECT_FALSE(OutputsBF16(body_dot)); + EXPECT_FALSE(OutputsBF16(body_rhs)); + EXPECT_FALSE(OutputsBF16(body_lhs)); + EXPECT_FALSE(OutputsBF16(cond0_lhs)); + EXPECT_FALSE(OutputsBF16(cond0_rhs)); + EXPECT_FALSE(OutputsBF16(cond1_lhs)); + EXPECT_FALSE(OutputsBF16(cond1_rhs)); + EXPECT_TRUE(OutputsBF16(cond0_add_rhs)); + EXPECT_TRUE(OutputsBF16(cond1_add_lhs)); + EXPECT_EQ(computation->root_instruction(), dot); +} + +} // 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..07b4b14b5ec1bdbc01345091105df69368b0b2fb --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_support.cc @@ -0,0 +1,112 @@ +/* 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::kConvert: + 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..82c2745f444e4f9c544c78cb36dafc11f678518a --- /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 preserves 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 774b11478c6d2faf0eb5db29df3cfd3cc1e98d5b..dbe45e932cdeed00e959355d5b3199d2e858148f 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -45,6 +45,185 @@ using ::tensorflow::gtl::FlatMap; using ::tensorflow::gtl::FlatSet; using ::tensorflow::strings::Appendf; using ::tensorflow::strings::HumanReadableNumBytes; +using ::tensorflow::strings::Printf; +using ::tensorflow::strings::StrAppend; + +namespace { + +template +string ColocatedBufferSetsToString(const T& container, const char* title) { + string result; + StrAppend(&result, title, "\n"); + for (const auto& it : container) { + StrAppend(&result, "\t", it->ToString(), "\n"); + } + return result; +} + +// Walk the call graph of the HLO module and place each computation into either +// thread_local_computations or global_computations depending upon whether the +// computation requires thread-local allocations or global allocations. The +// elements in thread_local_computations and global_computations are in post +// order (if computation A has an instruction which calls computation B, then A +// will appear after B in the vector). +Status GatherComputationsByAllocationType( + const HloModule* module, + std::vector* thread_local_computations, + std::vector* global_computations) { + // Create a worklist of computations paired with whether the allocation must + // be thread-local. + std::deque> worklist; + worklist.push_back(std::make_pair(module->entry_computation(), + /*is_thread_local*/ false)); + + // Sets for quickly checking membership. Computations are returned in vectors + // for stable iteration. + FlatSet thread_local_set; + FlatSet global_set; + + while (!worklist.empty()) { + auto worklist_front = worklist.front(); + worklist.pop_front(); + const HloComputation* computation = worklist_front.first; + bool is_thread_local = worklist_front.second; + bool in_thread_local_set = thread_local_set.count(computation) > 0; + bool in_global_set = global_set.count(computation) > 0; + + // If the computation has already been added to the respective set, then + // nothing to do. + if ((is_thread_local && in_thread_local_set) || + (!is_thread_local && in_global_set)) { + continue; + } + + // If the computation has already been added to the other set this is an + // error condition because the global call to the computation (eg, + // while/call) may return a reference to one of the thread-local buffers to + // the calling computation which will become a dangling reference when the + // thread-local is deallocated with the call return. + if ((is_thread_local && in_global_set) || + (!is_thread_local && in_thread_local_set)) { + return InvalidArgument( + "computation %s has conflicting allocation requirements (global " + "and thread-local)", + computation->name().c_str()); + } + + if (is_thread_local) { + thread_local_set.insert(computation); + } else { + global_set.insert(computation); + } + + for (auto* instruction : computation->instructions()) { + for (HloComputation* subcomputation : + instruction->called_computations()) { + switch (instruction->opcode()) { + case HloOpcode::kCall: + case HloOpcode::kConditional: + case HloOpcode::kWhile: + // Call and while must be called from a computation with global + // allocations as they may return references to buffers inside the + // called computation which cannot be thread-local. + if (is_thread_local) { + return InvalidArgument( + "computation %s cannot contain call/while op because it " + "requires thread-local buffer allocations", + computation->name().c_str()); + } + worklist.push_back(std::make_pair(subcomputation, + false)); // Not thread local. + break; + case HloOpcode::kMap: + case HloOpcode::kReduce: + case HloOpcode::kReduceWindow: + case HloOpcode::kSelectAndScatter: + case HloOpcode::kFusion: + // Map/reduce etc computations are always thread-local. + worklist.push_back(std::make_pair(subcomputation, + true)); // Thread local. + break; + default: + return InternalError( + "Unexpected calling opcode: %s", + HloOpcodeString(instruction->opcode()).c_str()); + } + } + } + } + + // Add the computations to the vectors in post order. + for (auto* computation : module->MakeComputationPostOrder()) { + if (thread_local_set.count(computation) > 0) { + thread_local_computations->push_back(computation); + } else if (global_set.count(computation) > 0) { + global_computations->push_back(computation); + } + // If the computation is not reachable from the entry computation, then it + // will not appear in either thread_local_set or global_set. We don't bother + // assigning buffers for these. + } + return Status::OK(); +} + +// Checks that points-to set of 'instruction' is unambiguous and distinct +// (ensured by CopyInsertion), then adds the buffer from the points-to set at +// 'index' to 'colocated_set'. +const LogicalBuffer* AddBufferToColocatedSet( + const HloInstruction* instruction, const ShapeIndex& index, + const TuplePointsToAnalysis& points_to_analysis, + std::vector* colocated_set) { + // CopyInsertion ensures root points-to set is unambiguous and distinct. + const auto& points_to = points_to_analysis.GetPointsToSet(instruction); + DCHECK(!points_to.IsAmbiguous()); + colocated_set->push_back(points_to.element(index)[0]); + 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 size_t BufferAllocation::Slice::Hasher::operator()(Slice s) const { uint64 h = std::hash()(s.index()); @@ -93,6 +272,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_); @@ -110,27 +292,132 @@ BufferAllocationProto BufferAllocation::ToProto() const { return proto; } +std::pair> +BufferAllocation::ComputePeakMemoryLogicalBuffers() const { + if (HeapTraces().empty()) { + // Just return the largest LogicalBuffer in the allocation. + const LogicalBuffer* largest_buffer = nullptr; + int64 largest_size = 0; + for (const auto& pair : assigned_buffers()) { + const LogicalBuffer* buffer = pair.first; + int64 size = pair.second.size; + if (largest_buffer == nullptr) { + largest_buffer = buffer; + largest_size = size; + continue; + } + // Tie-break with LogicalBuffer::Id so the return value is stable relative + // to changing addresses. + if (size > largest_size || + ((size == largest_size) && (largest_buffer->id() > buffer->id()))) { + largest_buffer = buffer; + largest_size = size; + } + } + CHECK(largest_buffer != nullptr) + << "No logical buffers in allocation: " << ToString(); + return {largest_size, {largest_buffer}}; + } + + // Create a map from LogicalBuffer::Id to LogicalBuffer* for the logical + // buffers in this allocation. + tensorflow::gtl::FlatMap + id_to_buffer; + tensorflow::gtl::FlatMap buffer_sizes; + for (const auto& pair : assigned_buffers()) { + const LogicalBuffer* buffer = pair.first; + const OffsetSize& offset_size = pair.second; + id_to_buffer[buffer->id()] = buffer; + buffer_sizes[buffer] = offset_size.size; + } + + // Returns how much the given event increases the total size of live + // buffers. Can be negative. + auto memory_delta = [this, &id_to_buffer, &buffer_sizes]( + const HeapSimulatorTrace::Event& event) -> int64 { + const LogicalBuffer* buffer = id_to_buffer.at(event.buffer_id()); + const int64 buffer_size = buffer_sizes.at(buffer); + if (event.kind() == HeapSimulatorTrace::Event::ALLOC) { + return buffer_size; + } else if (event.kind() == HeapSimulatorTrace::Event::SHARE_WITH) { + // Sharing a buffer does not change the live set size for the purposes of + // the heap simulator. Even though the shared-with buffer may be smaller, + // the entire allocation remains live. + return 0; + } else if (event.kind() == HeapSimulatorTrace::Event::FREE) { + return -1 * buffer_size; + } + LOG(FATAL) << "Unknown event kind: " << event.kind(); + }; + + int64 total_max_live_size = 0; + std::vector live_buffers_vector; + for (const HeapSimulatorTrace& heap_trace : HeapTraces()) { + // First compute the size of the maximal live set. + int64 max_live_size = 0; + int64 live_size = 0; + for (const auto& event : heap_trace.events()) { + live_size += memory_delta(event); + if (max_live_size < live_size) { + max_live_size = live_size; + } + } + + // Next gather the set of logical buffers live at the earliest point of + // maximal live set size. + tensorflow::gtl::FlatSet live_buffers; + live_size = 0; + for (const auto& event : heap_trace.events()) { + const LogicalBuffer* buffer = id_to_buffer.at(event.buffer_id()); + if (event.kind() == HeapSimulatorTrace::Event::ALLOC) { + InsertOrDie(&live_buffers, buffer); + } else if (event.kind() == HeapSimulatorTrace::Event::SHARE_WITH) { + // Nothing to do. + } else if (event.kind() == HeapSimulatorTrace::Event::FREE) { + CHECK(ContainsKey(live_buffers, buffer)); + live_buffers.erase(buffer); + } + + live_size += memory_delta(event); + if (live_size == max_live_size) { + break; + } + } + CHECK_EQ(live_size, max_live_size); + total_max_live_size += max_live_size; + + live_buffers_vector.insert(live_buffers_vector.end(), live_buffers.begin(), + live_buffers.end()); + } + + // Stabily sort the live buffers. + std::sort(live_buffers_vector.begin(), live_buffers_vector.end(), + [](const LogicalBuffer* a, const LogicalBuffer* b) { + return a->id() < b->id(); + }); + return {total_max_live_size, live_buffers_vector}; +} + 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_) { @@ -142,12 +429,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; } @@ -345,6 +631,7 @@ void BufferAssignment::AddAssignment(BufferAllocation* allocation, // Combines allocations of temporary buffers of the same color into one big // BufferAllocation. void BufferAssignment::CombineTempAllocations() { + VLOG(1) << "CombineTempAllocations()"; FlatMap combined_allocation_map; @@ -366,11 +653,16 @@ void BufferAssignment::CombineTempAllocations() { if (combined_it == combined_allocation_map.end()) { // We have found the first temp allocation of this color. Collect // the other temp allocations of the same color into it. + VLOG(1) << "Combined temp allocation for color " << color + << " is: " << temp_allocation; combined_allocation_map.emplace(color, temp_allocation); continue; } auto* combined_allocation = &combined_it->second; + VLOG(1) << "Combined allocation absorbing temp allocation: " + << temp_allocation; + // Each temp allocation is placed end-to-end, accounting for alignment. // The offset of each buffer in the combined allocation is computed from // the base offset of the allocation. @@ -384,6 +676,10 @@ void BufferAssignment::CombineTempAllocations() { const int64 size = buffer_offset_size.second.size; combined_allocation->AddAssignment(*buffer, base + offset, size); } + if (!temp_allocation.HeapTraces().empty()) { + CHECK_EQ(temp_allocation.HeapTraces().size(), 1); + combined_allocation->AddHeapTrace(temp_allocation.HeapTraces().front()); + } } // Replace all existing temporary allocations with the new combined // allocations. @@ -513,123 +809,13 @@ BufferAssignmentProto BufferAssignment::ToProto() const { for (const BufferAllocation& allocation : Allocations()) { BufferAllocationProto proto_allocation = allocation.ToProto(); proto.add_buffer_allocations()->Swap(&proto_allocation); - } - for (const HeapSimulatorTrace& trace : heap_simulator_traces_) { - *proto.add_heap_simulator_traces() = trace; - } - return proto; -} - -namespace { - -// Walk the call graph of the HLO module and place each computation into either -// thread_local_computations or global_computations depending upon whether the -// computation requires thread-local allocations or global allocations. The -// elements in thread_local_computations and global_computations are in post -// order (if computation A has an instruction which calls computation B, then A -// will appear after B in the vector). -Status GatherComputationsByAllocationType( - const HloModule* module, - std::vector* thread_local_computations, - std::vector* global_computations) { - // Create a worklist of computations paired with whether the allocation must - // be thread-local. - std::deque> worklist; - worklist.push_back(std::make_pair(module->entry_computation(), - /*is_thread_local*/ false)); - - // Sets for quickly checking membership. Computations are returned in vectors - // for stable iteration. - FlatSet thread_local_set; - FlatSet global_set; - - while (!worklist.empty()) { - auto worklist_front = worklist.front(); - worklist.pop_front(); - const HloComputation* computation = worklist_front.first; - bool is_thread_local = worklist_front.second; - bool in_thread_local_set = thread_local_set.count(computation) > 0; - bool in_global_set = global_set.count(computation) > 0; - - // If the computation has already been added to the respective set, then - // nothing to do. - if ((is_thread_local && in_thread_local_set) || - (!is_thread_local && in_global_set)) { - continue; - } - - // If the computation has already been added to the other set this is an - // error condition because the global call to the computation (eg, - // while/call) may return a reference to one of the thread-local buffers to - // the calling computation which will become a dangling reference when the - // thread-local is deallocated with the call return. - if ((is_thread_local && in_global_set) || - (!is_thread_local && in_thread_local_set)) { - return InvalidArgument( - "computation %s has conflicting allocation requirements (global " - "and thread-local)", - computation->name().c_str()); - } - - if (is_thread_local) { - thread_local_set.insert(computation); - } else { - global_set.insert(computation); - } - - for (auto* instruction : computation->instructions()) { - for (HloComputation* subcomputation : - instruction->called_computations()) { - switch (instruction->opcode()) { - case HloOpcode::kCall: - case HloOpcode::kConditional: - case HloOpcode::kWhile: - // Call and while must be called from a computation with global - // allocations as they may return references to buffers inside the - // called computation which cannot be thread-local. - if (is_thread_local) { - return InvalidArgument( - "computation %s cannot contain call/while op because it " - "requires thread-local buffer allocations", - computation->name().c_str()); - } - worklist.push_back(std::make_pair(subcomputation, - false)); // Not thread local. - break; - case HloOpcode::kMap: - case HloOpcode::kReduce: - case HloOpcode::kReduceWindow: - case HloOpcode::kSelectAndScatter: - case HloOpcode::kFusion: - // Map/reduce etc computations are always thread-local. - worklist.push_back(std::make_pair(subcomputation, - true)); // Thread local. - break; - default: - return InternalError( - "Unexpected calling opcode: %s", - HloOpcodeString(instruction->opcode()).c_str()); - } - } - } - } - - // Add the computations to the vectors in post order. - for (auto* computation : module->MakeComputationPostOrder()) { - if (thread_local_set.count(computation) > 0) { - thread_local_computations->push_back(computation); - } else if (global_set.count(computation) > 0) { - global_computations->push_back(computation); + for (const HeapSimulatorTrace& heap_trace : allocation.HeapTraces()) { + *proto.add_heap_simulator_traces() = heap_trace; } - // If the computation is not reachable from the entry computation, then it - // will not appear in either thread_local_set or global_set. We don't bother - // assigning buffers for these. } - return Status::OK(); + return proto; } -} // namespace - /* static */ StatusOr> BufferAssigner::Run( const HloModule* module, std::unique_ptr hlo_ordering, @@ -840,7 +1026,7 @@ 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; @@ -1061,7 +1247,8 @@ void BufferAssigner::AssignBuffersFromHeapSimulator( assignment->AddAssignment(allocation, buffer, chunk.offset, chunk.size); } - assignment->heap_simulator_traces_.push_back(result.debug_trace); + VLOG(1) << "Ran heap simulation for allocation: " << allocation->ToString(); + allocation->AddHeapTrace(result.debug_trace); } // Adds the 'colocated_set' of buffers to 'colocated_buffer_sets', maintaining @@ -1082,7 +1269,8 @@ void BufferAssigner::AddSetToColocatedBufferSets( if (colocated_set.empty()) { return; } - + VLOG(5) << ColocatedBufferSetsToString(colocated_set, + "Adding colocated buffer set"); // Find existing sets that overlap with at least one buffer from the // colocated_set. The resulting 'overlap_set_indices' will have at most // colocated_buffer_sets->size() entries, and will be in increasing order. @@ -1090,6 +1278,10 @@ void BufferAssigner::AddSetToColocatedBufferSets( for (size_t index = 0; index < colocated_buffer_sets->size(); ++index) { for (const LogicalBuffer* buffer : colocated_set) { if ((*colocated_buffer_sets)[index].count(buffer) > 0) { + VLOG(5) << "Found overlap with existing set on buffer " + << buffer->ToString() << "\n" + << ColocatedBufferSetsToString((*colocated_buffer_sets)[index], + "Overlapping set"); overlap_set_indices.push_back(index); break; } @@ -1101,6 +1293,7 @@ void BufferAssigner::AddSetToColocatedBufferSets( colocated_buffer_sets->emplace_back(); colocated_buffer_sets->back().insert(colocated_set.begin(), colocated_set.end()); + VLOG(5) << "No overlap found, new group created"; return; } @@ -1112,6 +1305,8 @@ void BufferAssigner::AddSetToColocatedBufferSets( first->insert(overlap_set.begin(), overlap_set.end()); } first->insert(colocated_set.begin(), colocated_set.end()); + VLOG(5) << ColocatedBufferSetsToString( + *first, "Result of the colocated buffer set merging"); // Remove overlap sets that we just merged. The offset accounts for the fact // that as elements are erased, the indices need to be adjusted. Keep in mind @@ -1122,158 +1317,94 @@ 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; - } +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(); + }; - // 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; + // 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) { + // Do not merge if one of the sets includes live outs or entry parameters. + for (int64 key : {i, j}) { + for (auto& buffer : colocated_buffer_sets[key]) { + if (buffer_liveness.MaybeLiveOut(*buffer) || + is_entry_parameter(*buffer)) { + return true; + } + } } - // 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; + // Colocated sets satisfy the invariant that all buffers within a set have + // the same size. That means we need to check whether the size is the same + // between the two sets, but also that it's enough to look at just one + // buffer within each set. + if (buffer_size(**colocated_buffer_sets[i].begin()) != + buffer_size(**colocated_buffer_sets[j].begin())) { + return true; } - // 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]); + // Do not merge if some pair of buffers interferes with each other. + for (auto& buffer_a : colocated_buffer_sets[i]) { + for (auto& buffer_b : colocated_buffer_sets[j]) { + if (buffer_a->id() != buffer_b->id() && + buffer_liveness.MayInterfere(*buffer_a, *buffer_b)) { + return true; + } } } - 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; - } + return false; + }; - if (while_result_buffer->id() != buffer->id() && - buffer_liveness.MayInterfere(*while_result_buffer, *buffer)) { - return true; + // 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); } - - 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 { + // 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()); + } -// Checks that points-to set of 'instruction' is unambiguous and distinct -// (ensured by CopyInsertion), then adds the buffer from the points-to set at -// 'index' to 'colocated_set'. -const LogicalBuffer* AddBufferToColocatedSet( - const HloInstruction* instruction, const ShapeIndex& index, - const TuplePointsToAnalysis& points_to_analysis, - std::vector* colocated_set) { - // CopyInsertion ensures root points-to set is unambiguous and distinct. - const auto& points_to = points_to_analysis.GetPointsToSet(instruction); - DCHECK(!points_to.IsAmbiguous()); - colocated_set->push_back(points_to.element(index)[0]); - return colocated_set->back(); + VLOG(1) << "colocation sets count after coalescing:" + << colocated_buffer_sets.size(); + return new_colocated_buffer_sets; } -} // namespace - // Builds sets of buffers in 'colocated_buffer_sets' which should be colocated // in the same allocation (currently just supports kWhile, kCall, and // kConditional). @@ -1299,12 +1430,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, @@ -1317,10 +1447,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; @@ -1400,6 +1527,22 @@ void BufferAssigner::BuildColocatedBufferSets( } } } + + 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 @@ -1411,14 +1554,17 @@ 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; } } @@ -1439,7 +1585,8 @@ void BufferAssigner::AssignColocatedBufferSets( // 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 { diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h index 08a40bfeb2a2a78c25805308e73154c6cc667f21..3086d0e2ca0026547134285b8ceb357390fc7ece 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_; } @@ -185,6 +192,37 @@ class BufferAllocation { !is_thread_local(); } + // Add a heap trace which was used to assign slices to logical buffers in this + // allocation. A single BufferAllocation may include multiple heap traces + // in the case of the temporary block where there is a heap trace per + // computation. + void AddHeapTrace(const HeapSimulatorTrace& heap_trace) { + heap_traces_.push_back(heap_trace); + } + + // Return the set of heap traces used to assign slices to logical buffers in + // this allocation. + const std::vector HeapTraces() const { + return heap_traces_; + } + + // Compute and return the LogicalBuffers which are live at the point of peak + // memory usage for the given allocation. The point of peak memory usage is + // the point at which the total size of all live logical buffers is + // maximal. If peak memory is reached at multiple points, the set of logical + // buffers live at the earliest maximal point is returned. The vector is + // stabily asserted by LogicalBuffer::Index. + // + // The return value is a pair of total size of the logical buffers at peak, + // and the buffers themselves. + std::pair> + ComputePeakMemoryLogicalBuffers() const; + + // Get the number of bytes lost to fragmentation. This is equal to the + // difference between the size of the allocation and the size of the maximal + // live set. + int64 fragmentation_bytes() const { return fragmentation_bytes_; } + bool operator==(const BufferAllocation& other) const { return index_ == other.index_; } @@ -203,9 +241,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 +275,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_` @@ -244,6 +288,9 @@ class BufferAllocation { // Mapping from the set of buffers assigned to this allocation to their // logical offsets and sizes. tensorflow::gtl::FlatMap assigned_buffers_; + + int64 fragmentation_bytes_ = 0; + std::vector heap_traces_; }; // Add stream operators for nicer output of CHECK/RET_CHECK failures. @@ -428,7 +475,6 @@ class BufferAssignment { LogicalBuffer::AlignmentFunction color_alignment_; Stats stats_; - std::vector heap_simulator_traces_; TF_DISALLOW_COPY_AND_ASSIGN(BufferAssignment); }; @@ -528,15 +574,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..513a8785bbd52b0a3bfa3642bbfc62b1035ffb17 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc @@ -37,14 +37,16 @@ limitations under the License. #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/tools/parser/hlo_parser.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/macros.h" namespace xla { - namespace { +using ::testing::UnorderedElementsAre; + // DFS visitor that collects the instructions referenced by a computation // without descending into nested computations, i.e., only from the operands. class InstructionListVisitor : public DfsHloVisitorWithDefault { @@ -101,6 +103,22 @@ class BufferAssignmentTest : public HloTestBase { .ConsumeValueOrDie(); } + std::unique_ptr RunBufferAssignmentWithInstructionSequence( + HloModule* module, + tensorflow::gtl::ArraySlice instruction_sequence, + int64 alignment = 1) { + SequentialHloOrdering::HloModuleSequence module_sequence; + module_sequence[module->entry_computation()] = + std::vector(instruction_sequence.begin(), + instruction_sequence.end()); + return BufferAssigner::Run( + module, + xla::MakeUnique(module, module_sequence), + backend().compiler()->BufferSizeBytesFunction(), + [alignment](LogicalBuffer::Color) { return alignment; }) + .ConsumeValueOrDie(); + } + // Builds an x+1.0 computation to use in a Map. std::unique_ptr BuildMapComputationPlus1(const string& name) { auto builder = HloComputation::Builder(name); @@ -614,7 +632,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 = @@ -1370,7 +1388,7 @@ TEST_F(BufferAssignmentTest, AmbiguousBufferAsOutput) { auto element_slices = assignment->GetAllSlices(select, /*index=*/{0}); EXPECT_EQ(2, element_slices.size()); EXPECT_THAT(element_slices, - ::testing::UnorderedElementsAre( + UnorderedElementsAre( assignment->GetUniqueSlice(tuple_param0, /*index=*/{0}) .ConsumeValueOrDie(), assignment->GetUniqueSlice(tuple_param1, /*index=*/{0}) @@ -1473,6 +1491,98 @@ TEST_F(BufferAssignmentTest, OneTempAllocation) { } } +TEST_F(BufferAssignmentTest, TrivialPeakBuffers) { + // paramscalar ------- (mul) -- (add) -- (sub) + // / / / + // param0[100] -------/ / / + // / / + // param1[100] --------------/--------/ + auto builder = HloComputation::Builder(TestName()); + auto paramscalar = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "")); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32vec100_, "")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32vec100_, "")); + auto mul = builder.AddInstruction(HloInstruction::CreateBinary( + f32vec100_, HloOpcode::kMultiply, paramscalar, param0)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); + builder.AddInstruction(HloInstruction::CreateBinary( + f32vec100_, HloOpcode::kSubtract, add, param1)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + auto buffers = RunBufferAssignment(module.get()); + + // Trivially, the set of peak memory logical buffer(s) of an allocation with a + // single logical buffer should be exactly the logical buffer in that + // allocation. + const BufferAllocation& mul_buffer = GetTopLevelAllocation(*buffers, mul); + int64 peak_size; + std::vector peak_buffers; + + std::tie(peak_size, peak_buffers) = + mul_buffer.ComputePeakMemoryLogicalBuffers(); + EXPECT_EQ(peak_size, ShapeUtil::ByteSizeOf(f32vec100_)); + ASSERT_EQ(peak_buffers.size(), 1); + EXPECT_EQ(peak_buffers[0]->instruction(), mul); +} + +TEST_F(BufferAssignmentTest, PeakBuffers) { + // Compute the peak liveness buffers of the following sequence: + // + // %param = ... + // %log = log(%param) + // %rev = reverse(%log) + // %neg = neg(%param) + // %concat = concat(%rev, %neg) + // ROOT %root = slice(concat) + // + // In the temporary block, the set of live buffers at peak memory use should + // be {%rev, %neg, %concat}. This occurs right at the concat itself. + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32vec100_, "")); + auto log = builder.AddInstruction( + HloInstruction::CreateUnary(f32vec100_, HloOpcode::kLog, param)); + auto rev = builder.AddInstruction( + HloInstruction::CreateReverse(f32vec100_, log, {0})); + auto neg = builder.AddInstruction( + HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param)); + const Shape concat_shape = ShapeUtil::MakeShape(F32, {200}); + auto concat = builder.AddInstruction( + HloInstruction::CreateConcatenate(concat_shape, {rev, neg}, 0)); + // Make the root tiny so no interior nodes can share its buffer. + auto root = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {1}), concat, {0}, {1}, {1})); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + auto buffers = RunBufferAssignmentWithInstructionSequence( + module.get(), {param, log, rev, neg, concat, root}); + + // The temporary buffer should hold the 4 interior instructions. + const BufferAllocation& buffer = GetTopLevelAllocation(*buffers, concat); + EXPECT_FALSE(buffer.IsInputOrOutput()); + EXPECT_TRUE(buffer.IsPreallocatedTempBuffer()); + ASSERT_EQ(buffer.assigned_buffers().size(), 4); + + int64 peak_size; + std::vector peak_buffers; + std::tie(peak_size, peak_buffers) = buffer.ComputePeakMemoryLogicalBuffers(); + + // The peak live set should be concat and its inputs. + EXPECT_EQ(peak_size, ShapeUtil::ByteSizeOf(ShapeUtil::MakeShape(F32, {400}))); + ASSERT_EQ(peak_buffers.size(), 3); + std::vector peak_instructions; + for (const LogicalBuffer* logical_buffer : peak_buffers) { + peak_instructions.push_back(logical_buffer->instruction()); + } + EXPECT_THAT(peak_instructions, UnorderedElementsAre(rev, neg, concat)); +} + class WhileBufferAssignmentTest : public HloTestBase { protected: std::unique_ptr BuildWhileConditionComputation( @@ -1587,6 +1697,192 @@ TEST_F(WhileBufferAssignmentTest, TwoForwardWhileLoops) { assignment->GetUniqueSlice(while1, {1}).ConsumeValueOrDie()); } +// Tests that two colocated buffer sets are not merged if an entry parameter +// buffer belongs to either of the colocation sets (b/73267882). +// +// %param --> %while.0 --> %mul --> %while.1 --> %broadcast +// +// %while.0 body just forwards the init value, so the loop carried variable +// remains the constant, whereas %while.1 changes the loop carried variable. +TEST_F(WhileBufferAssignmentTest, ColocatedBufferWithEntryParameter) { + const Shape r0s32 = ShapeUtil::MakeShape(S32, {}); + + const char* module_str = R"( +HloModule test_module + +%cond.v0 { + %param = s32[] parameter(0) + ROOT %constant = pred[] constant(true) +} + +%cond.v1 { + %param.0 = s32[] parameter(0) + ROOT %constant.0 = pred[] constant(true) +} + +%body.v0 { + ROOT %param.1 = s32[] parameter(0) +} + +%body.v1 { + %param.2 = s32[] parameter(0) + ROOT add = s32[] add(%param.2, %param.2) +} + +ENTRY %test_module { + %param.3 = s32[] parameter(0) + %while.0 = s32[] while(%param.3), condition=%cond.v0, body=%body.v0 + %mul = s32[] multiply(%while.0, %while.0) + %while.1 = s32[] while(%mul), condition=%cond.v1, body=%body.v1 + ROOT %bcast = s32[1024,1024]{1,0} broadcast(s32[] %while.1), dimensions={} +})"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + tools::Parse(module_str)); + + // 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()); + + // Get the instructions in the module. + const HloInstruction* bcast = module->entry_computation()->root_instruction(); + const HloInstruction* param = + module->entry_computation()->parameter_instruction(0); + ASSERT_EQ(bcast->opcode(), HloOpcode::kBroadcast); + const HloInstruction* while1 = bcast->operand(0); + ASSERT_EQ(while1->opcode(), HloOpcode::kWhile); + const HloInstruction* while0 = while1->operand(0)->operand(0); + ASSERT_EQ(while0->opcode(), HloOpcode::kWhile); + + // Run buffer assignment. + auto assignment = RunBufferAssignment(module.get()); + TF_ASSERT_OK_AND_ASSIGN(auto slice_param, + assignment->GetUniqueSlice(param, {})); + TF_ASSERT_OK_AND_ASSIGN(auto slice_while0, + assignment->GetUniqueSlice(while0, {})); + TF_ASSERT_OK_AND_ASSIGN(auto slice_while1, + assignment->GetUniqueSlice(while1, {})); + + // The parameter slice is part of the while0's colocation set (init value), + // but not merged into the while1's colocation set. + EXPECT_EQ(slice_param, slice_while0); + EXPECT_NE(slice_param, slice_while1); +} + +// 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 e7749252ce44f0daf7016f72d80401695eaaacb9..37982aaef9eddd64ef6b57ad5a9cf8dd6a565097 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness.cc @@ -117,11 +117,12 @@ bool BufferLiveness::live_range_strictly_before(const LogicalBuffer& a, // 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. 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. + // 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() && + alias.instruction()->parent() == b.instruction()->parent()) { return false; } } diff --git a/tensorflow/compiler/xla/service/compile_only_service.cc b/tensorflow/compiler/xla/service/compile_only_service.cc index dab73596e1639eed62151197048ee8d29570b20a..6664496ab6c603c35c7dce923fcf94c54d1ce714 100644 --- a/tensorflow/compiler/xla/service/compile_only_service.cc +++ b/tensorflow/compiler/xla/service/compile_only_service.cc @@ -72,8 +72,7 @@ CompileOnlyService::CompileAheadOfTime( VersionedComputationHandle versioned_handle = user_computation->GetVersionedHandle(); - // TODO(b/63773457): Track DebugOptions in AotCompilationOptions. - DebugOptions debug_options = legacy_flags::GetDebugOptionsFromFlags(); + const DebugOptions& debug_options = options.debug_options(); // Dump computation proto state if flag is set. const string& directory_path = debug_options.xla_dump_computations_to(); diff --git a/tensorflow/compiler/xla/service/compiler.cc b/tensorflow/compiler/xla/service/compiler.cc index e2e9d2a0c048fec6c6ffbeef1223ae0e6aef50d1..0392d4af48a040c4a648f7bf9bf21a62ce03a990 100644 --- a/tensorflow/compiler/xla/service/compiler.cc +++ b/tensorflow/compiler/xla/service/compiler.cc @@ -86,4 +86,7 @@ Compiler::GetPlatformCompilers() { return compilers->at(platform->id()).get(); } +AotCompilationOptions::AotCompilationOptions() + : debug_options_(legacy_flags::GetDebugOptionsFromFlags()) {} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h index 74fd24edf88d44b2dfdc87556b0af43987e69e08..33e19efc72c6d30ccd7e0b3a13f664a4f42208bf 100644 --- a/tensorflow/compiler/xla/service/compiler.h +++ b/tensorflow/compiler/xla/service/compiler.h @@ -79,11 +79,15 @@ class AotCompilationOptions { device_allocator_ = device_allocator; } + const DebugOptions& debug_options() const { return debug_options_; } + DebugOptions* mutable_debug_options() { return &debug_options_; } + protected: - AotCompilationOptions() = default; + AotCompilationOptions(); private: DeviceMemoryAllocator* device_allocator_ = nullptr; + DebugOptions debug_options_; }; // Abstract compiler interface that is subclassed for compilation on a diff --git a/tensorflow/compiler/xla/service/conditional_simplifier.cc b/tensorflow/compiler/xla/service/conditional_simplifier.cc new file mode 100644 index 0000000000000000000000000000000000000000..f35de080853f7ec986565cb2df1050946ac3f244 --- /dev/null +++ b/tensorflow/compiler/xla/service/conditional_simplifier.cc @@ -0,0 +1,106 @@ +/* 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/conditional_simplifier.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/call_inliner.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/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" + +namespace xla { + +// Tries to replace a conditional with a call operation of the corresponding +// computation. If the given conditional has a constant predicate, tries to +// replace it with a call to its true/false computation as appropirate and then +// inline that computation. +// +// Returns true if it made a change to the graph. +static StatusOr TryRemoveConditional(HloInstruction* conditional) { + CHECK_EQ(conditional->opcode(), HloOpcode::kConditional); + // Do not remove conditionals that contain side-effecting instructions or + // have control predecessors/successors in either true/false computation. + if (!conditional->parent()->IsRemovable(conditional) || + conditional->HasSideEffect()) { + VLOG(2) << "Not attempting to remove conditional as it is not removable or " + "has side effect: " + << conditional->ToShortString(); + return false; + } + + if (conditional->operand(0)->opcode() != HloOpcode::kConstant) { + VLOG(2) << "Not attempting to remove conditional as its predicate is not a " + "compile-time constant: " + << conditional->ToShortString(); + return false; + } + + auto computation = conditional->parent(); + HloInstruction* call_op; + if (conditional->operand(0)->literal().Get({})) { + call_op = computation->AddInstruction(HloInstruction::CreateCall( + conditional->shape(), {conditional->mutable_operand(1)}, + conditional->true_computation())); + } else { + call_op = computation->AddInstruction(HloInstruction::CreateCall( + conditional->shape(), {conditional->mutable_operand(2)}, + conditional->false_computation())); + } + + TF_RETURN_IF_ERROR(computation->ReplaceInstruction(conditional, call_op)); + TF_RETURN_IF_ERROR(CallInliner::Inline(call_op).status()); + + return true; +} + +StatusOr ConditionalSimplifier::Run(HloModule* module) { + XLA_VLOG_LINES( + 3, "ConditionalSimplifier::Run(), before:\n" + module->ToString()); + bool changed = false; + + // Gather all the conditional ops in our module. We do this ahead of time so + // we don't have to worry about mutating the lists of computations or + // instructions as we iterate. + std::vector conditional_ops; + for (auto* comp : module->computations()) { + for (auto* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kConditional) { + conditional_ops.push_back(instr); + } + } + } + + for (HloInstruction* conditional_op : conditional_ops) { + TF_ASSIGN_OR_RETURN(bool result, TryRemoveConditional(conditional_op)); + changed |= result; + } + + XLA_VLOG_LINES(3, + "ConditionalSimplifier::Run(), after:\n" + module->ToString()); + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/convolution_folding.h b/tensorflow/compiler/xla/service/conditional_simplifier.h similarity index 64% rename from tensorflow/compiler/xla/service/gpu/convolution_folding.h rename to tensorflow/compiler/xla/service/conditional_simplifier.h index f9c898721f8dd6b8b7e74c82bb2085cc437eaad5..063261e26d06e21a297e8e3c405898a17221b7ca 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_folding.h +++ b/tensorflow/compiler/xla/service/conditional_simplifier.h @@ -13,25 +13,26 @@ 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_CONDITIONAL_SIMPLIFIER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CONDITIONAL_SIMPLIFIER_H_ #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/core/stringpiece.h" namespace xla { -namespace gpu { -class ConvolutionFolding : public HloPassInterface { +// HLO pass that removes kConditional with a constant predicate, replacing them +// with their true or false computation as appropriate. +class ConditionalSimplifier : public HloPassInterface { public: tensorflow::StringPiece name() const override { - return "convolution-folding"; + return "simplify-conditional"; } - StatusOr Run(HloModule* module) override; }; -} // namespace gpu } // namespace xla -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONVOLUTION_FOLDING_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CONDITIONAL_SIMPLIFIER_H_ diff --git a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..868348547d9f5cbdc7576c7fc0697d72c3a3e557 --- /dev/null +++ b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc @@ -0,0 +1,153 @@ +/* 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/conditional_simplifier.h" + +#include +#include + +#include "tensorflow/compiler/xla/literal_util.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_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#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/xla_data.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { + +namespace op = xla::testing::opcode_matchers; + +class ConditionalSimplifierTest : public HloVerifiedTestBase { + public: + // Makes a computation that contains a conditional with constant predicate. + HloComputation* MakeConditional(HloModule* module); +}; + +HloComputation* ConditionalSimplifierTest::MakeConditional(HloModule* module) { + HloComputation::Builder builder(TestName()); + + // true_computation returns param+1. + HloComputation* true_computation; + { + HloComputation::Builder true_computation_builder(TestName() + + ".true_computation"); + auto param = + true_computation_builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(S32, {}), "param")); + auto one = true_computation_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1))); + + true_computation_builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(S32, {}), HloOpcode::kAdd, param, one)); + + true_computation = + module->AddEmbeddedComputation(true_computation_builder.Build()); + } + + // false_computation returns param+42. + HloComputation* false_computation; + { + HloComputation::Builder false_computation_builder(TestName() + + ".false_computation"); + auto param = false_computation_builder.AddInstruction( + HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(S32, {}), + "param")); + auto forty_two = false_computation_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42))); + + false_computation_builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(S32, {}), HloOpcode::kAdd, param, forty_two)); + false_computation = + module->AddEmbeddedComputation(false_computation_builder.Build()); + } + + auto false_instrn = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto false_param = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(S32, {}), "false_param")); + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1))); + + builder.AddInstruction(HloInstruction::CreateConditional( + ShapeUtil::MakeShape(S32, {}), false_instrn, one, true_computation, + false_param, false_computation)); + + return module->AddEntryComputation(builder.Build()); +} + +TEST_F(ConditionalSimplifierTest, ConditionalGetsInlined) { + HloComputation* computation = MakeConditional(&module()); + ASSERT_TRUE(ConditionalSimplifier().Run(&module()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Parameter(), op::Constant())); +} + +TEST_F(ConditionalSimplifierTest, ConditionalWithControlDependency) { + HloComputation* computation = MakeConditional(&module()); + + auto* true_op = computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(true))); + TF_ASSERT_OK( + true_op->AddControlDependencyTo(computation->root_instruction())); + + EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); +} + +TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsSend) { + HloComputation* computation = MakeConditional(&module()); + auto* conditional = computation->root_instruction(); + ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); + + auto* true_computation = conditional->true_computation(); + auto* send = true_computation->AddInstruction(HloInstruction::CreateSend( + true_computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(true))), + /*channel_id=*/0)); + true_computation->AddInstruction(HloInstruction::CreateSendDone(send)); + EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); +} + +TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsRecv) { + HloComputation* computation = MakeConditional(&module()); + auto* conditional = computation->root_instruction(); + ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); + + auto* true_computation = conditional->true_computation(); + auto* recv = true_computation->AddInstruction(HloInstruction::CreateRecv( + ShapeUtil::MakeShape(F32, {1}), /*channel_id=*/0)); + true_computation->AddInstruction(HloInstruction::CreateRecvDone(recv)); + EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); +} + +TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsNonRemovableInstruction) { + HloComputation* computation = MakeConditional(&module()); + auto* conditional = computation->root_instruction(); + ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); + auto* false_computation = conditional->false_computation(); + false_computation->AddInstruction( + HloInstruction::CreateInfeed(ShapeUtil::MakeShape(F32, {1}), "config")); + EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index cd983bc03e993caed883916de01d75dffdbc4bab..40519ecc799c8f0343294ad88009820dbd8535e9 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -58,6 +58,46 @@ bool ValueIsReadOnly(const HloValue& value) { return IsConstantValue(value) || IsEntryParameterValue(value); } +// Data structure describing the action which should be taken on parts of a +// computation buffers, with respect to the adding of special case copies. +struct SpecialCaseCopyPolicy { + // Insert a copy if the same buffer is found at multiple indices within the + // output tuple. + bool copy_root_replicated_buffers = false; + // If true, insert a copy if a buffer coming from a constant or a parameter + // is found wihtin the output tuple. + bool copy_parameters_and_constants = false; +}; + +SpecialCaseCopyPolicy GetSpecialCaseCopyPolicy(const CallGraphNode& node, + HloModule* module, + HloComputation* computation) { + SpecialCaseCopyPolicy policy; + if (computation == module->entry_computation()) { + policy.copy_parameters_and_constants = true; + policy.copy_root_replicated_buffers = true; + } + for (const CallSite& site : node.caller_callsites()) { + // The AddCopiesForConditional() already adds copies, but the copy remover + // removes them, so we re-add them by returning the policy here. But really + // the copy remover should not be removing them. + if (site.instruction()->opcode() == HloOpcode::kConditional) { + policy.copy_parameters_and_constants = true; + policy.copy_root_replicated_buffers = true; + } + } + return policy; +} + +bool ShouldCopyRootValue(const HloValue& value, + const SpecialCaseCopyPolicy& policy) { + if (policy.copy_parameters_and_constants) { + return IsConstantValue(value) || + value.defining_instruction()->opcode() == HloOpcode::kParameter; + } + return false; +} + // Deep copy the given instructions 'from' and 'to' at the ShapeIndexes given in // 'indices_to_copy'. Add control edges from the respective kCopy instructions // in deep copy of 'from' to the respective kCopy instruction in the deep copy @@ -282,6 +322,29 @@ Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, return Status::OK(); } +// We add copies for all the indices of the true and false computaiton roots, +// in order to resolve interference. We later rely on the CopyRemover to drop +// the unnecessary ones. +Status AddCopiesForConditional(const HloAliasAnalysis& alias_analysis, + HloInstruction* conditional) { + VLOG(2) << "Adding copies for kConditional instruction " + << conditional->name(); + TF_RET_CHECK(conditional->opcode() == HloOpcode::kConditional); + + for (HloComputation* computation : + {conditional->true_computation(), conditional->false_computation()}) { + HloInstruction* root = computation->root_instruction(); + std::vector users = root->users(); + TF_ASSIGN_OR_RETURN(HloInstruction * deep_copy, + computation->DeepCopyInstruction(root)); + for (HloInstruction* user : users) { + TF_RETURN_IF_ERROR(root->ReplaceUseWith(user, deep_copy)); + } + computation->set_root_instruction(deep_copy); + } + return Status::OK(); +} + // Removes any control dependencies to or from the given instruction. Status StripControlDependenciesFrom(HloInstruction* instruction) { while (!instruction->control_successors().empty()) { @@ -309,6 +372,9 @@ Status AddCopiesToResolveInterference(HloModule* module) { for (HloInstruction* instruction : computation->instructions()) { if (instruction->opcode() == HloOpcode::kWhile) { TF_RETURN_IF_ERROR(AddCopiesForWhile(*alias_analysis, instruction)); + } else if (instruction->opcode() == HloOpcode::kConditional) { + TF_RETURN_IF_ERROR( + AddCopiesForConditional(*alias_analysis, instruction)); } } } @@ -557,6 +623,7 @@ class CopyRemover { auto is_live_range_before = [this](const ValueNode& a, const ValueNode& b) { + VLOG(3) << "Checking live range of " << *a.value << " WRT " << *b.value; if (LiveRangeBefore(a, b)) { VLOG(2) << " Live range of " << a.value->ToShortString() << " is before " << b.value->ToShortString(); @@ -571,7 +638,7 @@ class CopyRemover { VLOG(3) << copy->name() << " copies value " << src->value->ToShortString(); VLOG(3) << "Source buffer values: " << ValueListToString(src); - VLOG(3) << "Dest buffer values: " << ValueListToString(src); + VLOG(3) << "Dest buffer values: " << ValueListToString(dest); // A kCopy instruction copies an HLO value from a source buffer and // defines an HLO value in a destination buffer. Most generally, the @@ -729,7 +796,8 @@ class CopyRemover { // has a different operand (the operand of the elided copy). for (const HloUse* copy_use : copy_value_node->uses) { operand_node->uses.push_back(copy_use); - if (copy_use->instruction->opcode() == HloOpcode::kCopy) { + if (copy_use->instruction->opcode() == HloOpcode::kCopy && + ContainsKey(copy_map_, copy_use->instruction)) { copy_map_.at(copy_use->instruction).src = operand_node; } } @@ -746,16 +814,16 @@ class CopyRemover { // updated as copies are removed. bool LiveRangeBefore(const ValueNode& a, const ValueNode& b) { if (a.uses.empty()) { - VLOG(2) << "Empty uses"; + VLOG(2) << "Empty uses for " << *a.value; return ordering_.IsDefinedBefore(*a.value, *b.value); } for (const HloUse* use : a.uses) { - VLOG(2) << "use: " << *use; - VLOG(2) << "is before:" << *b.value; + VLOG(2) << "Checking use " << *use << " against " << *b.value; if (!ordering_.UseIsBeforeValueDefinition(*use, *b.value, dataflow_)) { - VLOG(2) << "Not before"; + VLOG(2) << "Use " << *use << " is NOT before " << *b.value; return false; } + VLOG(2) << "Use " << *use << " is before " << *b.value; } return true; } @@ -891,7 +959,6 @@ Status RemoveUnnecessaryCopies( CopyRemover copy_remover(*alias_analysis, ordering, module); XLA_VLOG_LINES(3, copy_remover.ToString()); - tensorflow::gtl::FlatSet existing_copies; for (HloComputation* computation : module->computations()) { for (HloInstruction* instruction : computation->instructions()) { if (instruction->opcode() == HloOpcode::kCopy && @@ -900,7 +967,6 @@ Status RemoveUnnecessaryCopies( } } } - return Status::OK(); } @@ -920,7 +986,7 @@ Status AddSpecialCaseCopies(const CallGraph& call_graph, HloModule* module) { // Identify which shape indices of which instructions need to be copied. Store // these results in 'instructions_to_copy'. - std::unordered_map> instructions_to_copy; + HloInstructionMap> instructions_to_copy; auto add_index_to_copy = [&instructions_to_copy](HloInstruction* instruction, const ShapeIndex& index) { auto it = instructions_to_copy.find(instruction); @@ -956,7 +1022,8 @@ Status AddSpecialCaseCopies(const CallGraph& call_graph, HloModule* module) { } TF_RET_CHECK(node.context() == CallContext::kSequential); - const bool is_entry = computation == module->entry_computation(); + SpecialCaseCopyPolicy policy = + GetSpecialCaseCopyPolicy(node, module, computation); HloInstruction* root = computation->root_instruction(); // Mark nondistinct/ambiguous indices. @@ -969,27 +1036,26 @@ Status AddSpecialCaseCopies(const CallGraph& call_graph, HloModule* module) { for (const HloBuffer* buffer : buffers_at_index) { buffer_seen_before |= !seen.insert(buffer).second; } - if (buffers_at_index.size() > 1 || (buffer_seen_before && is_entry)) { - VLOG(2) << "Index " << index << " of root of computation " + if (buffers_at_index.size() > 1 || + (buffer_seen_before && policy.copy_root_replicated_buffers)) { + VLOG(2) << "Index " << index << " of computation " << computation->name() << " (" << root->name() << ") has ambiguous or non-distinct buffer. Copying."; add_index_to_copy(root, index); } }); - // For entry instructions, mark any parameter or constant values. - if (is_entry) { - for (const auto& pair : - alias_analysis->dataflow_analysis().GetInstructionValueSet(root)) { - const ShapeIndex& index = pair.first; - const HloValueSet& value_set = pair.second; - for (const HloValue* value : value_set.values()) { - if (ValueIsReadOnly(*value)) { - VLOG(2) << "Root of entry computation (" << root->name() - << ") has constant or entry parameter value at index " - << index << ". Copying."; - add_index_to_copy(root, index); - } + for (const auto& pair : + alias_analysis->dataflow_analysis().GetInstructionValueSet(root)) { + const ShapeIndex& index = pair.first; + const HloValueSet& value_set = pair.second; + for (const HloValue* value : value_set.values()) { + if (ShouldCopyRootValue(*value, policy)) { + VLOG(2) << "Root of (" << root->name() << ") of computation(" + << computation->name() + << ") has constant or parameter value at index " << index + << ". Copying."; + add_index_to_copy(root, index); } } } @@ -1011,7 +1077,6 @@ Status AddSpecialCaseCopies(const CallGraph& call_graph, HloModule* module) { instruction->parent()->set_root_instruction(deep_copy); } } - return Status::OK(); } @@ -1155,7 +1220,7 @@ bool IsWhileBody(const HloComputation* computation, HloModule* module) { std::unique_ptr call_graph = CallGraph::Build(module); TF_ASSIGN_OR_RETURN(std::unique_ptr dataflow, - HloDataflowAnalysis::Run(module)); + HloDataflowAnalysis::Run(*module)); bool changed = false; diff --git a/tensorflow/compiler/xla/service/copy_insertion_test.cc b/tensorflow/compiler/xla/service/copy_insertion_test.cc index 128ee726ea6e4a8b63727fdc9762d865cee1c985..153f062d015e49db11c4c9ae0a2a61e76c020f02 100644 --- a/tensorflow/compiler/xla/service/copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/copy_insertion_test.cc @@ -1724,8 +1724,58 @@ 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"( diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index 2f0259163120dd5d62a5d1289deada8dc59c2c6c..093db020c03aa3685005b4a42b26dd032fcf0001 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -105,9 +105,11 @@ cc_library( "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:buffer_liveness", "//tensorflow/compiler/xla/service:call_inliner", + "//tensorflow/compiler/xla/service:conditional_simplifier", "//tensorflow/compiler/xla/service:dot_decomposer", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:flatten_call_graph", + "//tensorflow/compiler/xla/service:gather_expander", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_constant_folding", "//tensorflow/compiler/xla/service:hlo_cse", @@ -159,13 +161,11 @@ cc_library( 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_fp16", ":runtime_conv2d", ":runtime_fft", ":runtime_fork_join", @@ -185,6 +185,20 @@ cc_library( ] + ORC_JIT_MEMORY_MAPPER_TARGETS, ) +cc_library( + name = "runtime_fp16", + srcs = [ + "runtime_fp16.cc", + ], + hdrs = [ + "runtime_fp16.h", + ], + copts = runtime_copts(), + deps = [ + "//tensorflow/core:framework_lite", + ], +) + cc_library( name = "cpu_executable", srcs = ["cpu_executable.cc"], @@ -408,9 +422,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", @@ -430,43 +441,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 = [ @@ -497,6 +471,7 @@ cc_library( "llvm_ir_runtime.h", ], deps = [ + ":vector_support_library", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:lib", "@llvm//:core", @@ -541,7 +516,6 @@ cc_library( cc_library( name = "runtime_matvec", - srcs = ["runtime_matvec.cc"], hdrs = ["runtime_matvec.h"], copts = runtime_copts(), deps = [ @@ -798,6 +772,31 @@ cc_library( ], ) +tf_cc_test( + name = "parallel_task_assignment_test", + srcs = ["parallel_task_assignment_test.cc"], + deps = [ + ":cpu_executable", + ":parallel_task_assignment", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_layout", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:algebraic_simplifier", + "//tensorflow/compiler/xla/service:computation_layout", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + cc_library( name = "cpu_options", srcs = ["cpu_options.cc"], @@ -852,6 +851,7 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "@llvm//:core", + "@llvm//:support", ], ) diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc index 04b4a8c5c80eeefdbe10001ba5c462affbc9b21d..61b2da7a7dce7f6fba46a23cc8e5462a3899a18c 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: @@ -105,8 +93,8 @@ class FilteredPassManager : public llvm::legacy::PassManager { }; } // anonymous namespace -llvm::object::OwningBinary CompilerFunctor:: -operator()(llvm::Module& module) const { +std::unique_ptr CompilerFunctor::operator()( + llvm::Module& module) const { FilteredPassManager module_passes(disable_expensive_passes_); FilteredFunctionPassManager function_passes(&module, disable_expensive_passes_); @@ -169,112 +157,32 @@ operator()(llvm::Module& module) const { codegen_passes.run(module); // Construct ObjectFile from machine code buffer. - std::unique_ptr memory_buffer( + return std::unique_ptr( new llvm::ObjectMemoryBuffer(std::move(stream_buffer))); - llvm::Expected> - object_file_or_error = llvm::object::ObjectFile::createObjectFile( - memory_buffer->getMemBufferRef()); - CHECK(object_file_or_error); - - std::unique_ptr object_file = - std::move(object_file_or_error.get()); - if (VLOG_IS_ON(2)) { - StatusOr disassembly_status = - disassembler_->DisassembleObjectFile(*object_file); - if (disassembly_status.ok()) { - auto result = disassembly_status.ValueOrDie(); - XLA_VLOG_LINES(2, result.text); - VLOG(2) << "compiled code size: " << result.code_size_bytes << " bytes"; - } - } - - return llvm::object::OwningBinary( - 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 +190,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..c38b896c5019b48fd2a16a51abd59e12ebdb29eb 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,12 +43,11 @@ 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) {} // Compile a Module to an ObjectFile. - llvm::object::OwningBinary operator()( + std::unique_ptr operator()( llvm::Module& module) const; // NOLINT private: @@ -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 d13a97bcc9a84afb22556389b4cdcd985f58d445..0a966fd5a7c1ce2c4e367b26701c9186ab2ebf74 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -47,6 +47,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/call_inliner.h" +#include "tensorflow/compiler/xla/service/conditional_simplifier.h" #include "tensorflow/compiler/xla/service/cpu/compiler_functor.h" #include "tensorflow/compiler/xla/service/cpu/conv_canonicalization.h" #include "tensorflow/compiler/xla/service/cpu/cpu_copy_insertion.h" @@ -66,6 +67,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/dot_decomposer.h" #include "tensorflow/compiler/xla/service/flatten_call_graph.h" +#include "tensorflow/compiler/xla/service/gather_expander.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_constant_folding.h" @@ -260,6 +262,7 @@ Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) { /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true, /*use_fusion=*/false); + pipeline.AddPass(); pass.AddPass( /*is_layout_sensitive=*/false, [](const Shape&, const Shape&) { return false; }, @@ -275,6 +278,7 @@ Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) { pass.AddPass(); pass.AddPass(); pass.AddPass(); + pass.AddPass(); } pipeline.AddPass( [](const HloInstruction& dot, @@ -888,13 +892,11 @@ CpuCompiler::CompileAheadOfTime(std::vector> modules, 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); - llvm::object::OwningBinary object_file = + pre_optimization_ir_dump_hook, post_optimization_ir_dump_hook); + std::unique_ptr object_file = compiler_functor(llvm_module); - llvm::StringRef object_file_data_ref = object_file.getBinary()->getData(); - ObjectFileData object_file_data(object_file_data_ref.begin(), - object_file_data_ref.end()); + ObjectFileData object_file_data(object_file->getBufferStart(), + object_file->getBufferEnd()); BufferSizes buffer_sizes; for (const BufferAllocation& allocation : assignment->Allocations()) { diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc index 802d0a6fb46890b31d14b1fbf3b2e7d6520caccc..c053703c3524a47ee1de9681c1b986edbf109430 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -63,7 +63,7 @@ CpuExecutable::CpuExecutable( assignment_(std::move(assignment)) { // Resolve symbols in the constructor rather than at execution time to avoid // races because FindSymbol is not thread safe. - llvm::JITSymbol sym = jit_->FindSymbol(entry_function_name); + llvm::JITSymbol sym = jit_->FindCompiledSymbol(entry_function_name); // We expect to find the symbol provided with entry_function_name; otherwise // this is an internal error. CHECK(sym) << "Symbol " << entry_function_name << " not found."; diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc index 482e04052d5a914eab0e5bff2c7a83f3b698052f..0fc5a746bbbc7685ff5d4647111a750e7d7b1c19 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc @@ -30,7 +30,6 @@ bool CanBeLoopFused(const HloInstruction& hlo) { // These are the only ones we fuse since we rely on effective elemental IR // generation. return hlo.IsElementwise() || // - hlo.opcode() == HloOpcode::kBitcast || hlo.opcode() == HloOpcode::kBroadcast || hlo.opcode() == HloOpcode::kConcatenate || hlo.opcode() == HloOpcode::kDynamicSlice || diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc index 595c3f55b321f47e2312b93e0c238c7637495d77..6ed1cd31b18f6360bdd7fd41bd5be2e657b310a5 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc @@ -77,7 +77,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_Basic_1) { EXPECT_THAT(computation->root_instruction(), op::Fusion()); } -TEST_F(InstructionFusionTest, DotOperationFusion_Bitcast) { +TEST_F(InstructionFusionTest, DotOperationNoFusion_Bitcast) { HloComputation::Builder builder(TestName()); HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {2, 512, 2, 128}), "arg0")); @@ -94,8 +94,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_Bitcast) { auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); - EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Fusion()); + EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); } TEST_F(InstructionFusionTest, DotOperationFusion_Reshape) { @@ -244,35 +243,33 @@ class OpcodeFusionTest : public InstructionFusionTest { } }; -TEST_F(OpcodeFusionTest, Exponential_Bitcast_Negate) { +TEST_F(OpcodeFusionTest, Exponential_Reshape_Negate) { HloComputation::Builder builder(TestName()); Shape param_shape = ShapeUtil::MakeShape(F32, {1, 4}); Shape result_shape = ShapeUtil::MakeShape(F32, {4}); HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, param_shape, "param")); - // InstructionFusion::ShouldFuse() precludes fusing a bitcast whose operand - // is a parameter, so create an operand between the parameter and bitcast. HloInstruction* exp1 = builder.AddInstruction( HloInstruction::CreateUnary(param_shape, HloOpcode::kExp, param0)); - HloInstruction* bitcast2 = builder.AddInstruction( - HloInstruction::CreateUnary(result_shape, HloOpcode::kBitcast, exp1)); + HloInstruction* reshape2 = + builder.AddInstruction(HloInstruction::CreateReshape(result_shape, exp1)); builder.AddInstruction( - HloInstruction::CreateUnary(result_shape, HloOpcode::kNegate, bitcast2)); + HloInstruction::CreateUnary(result_shape, HloOpcode::kNegate, reshape2)); auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( - module.get(), {HloOpcode::kNegate, HloOpcode::kBitcast, HloOpcode::kExp, + module.get(), {HloOpcode::kNegate, HloOpcode::kReshape, HloOpcode::kExp, HloOpcode::kParameter}); } -TEST_F(OpcodeFusionTest, Broadcast_Bitcast_DynamicSlice_Tanh) { +TEST_F(OpcodeFusionTest, Broadcast_Reshape_DynamicSlice_Tanh) { HloComputation::Builder builder(TestName()); Shape param_shape = ShapeUtil::MakeShape(F32, {8}); Shape starts_shape = ShapeUtil::MakeShape(F32, {2}); Shape broadcast_shape = ShapeUtil::MakeShape(F32, {1, 8, 8}); - Shape bitcast_shape = ShapeUtil::MakeShape(F32, {8, 8}); + Shape reshape_shape = ShapeUtil::MakeShape(F32, {8, 8}); Shape dynamic_slice_shape = ShapeUtil::MakeShape(F32, {4, 4}); HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, param_shape, "param")); @@ -280,11 +277,11 @@ TEST_F(OpcodeFusionTest, Broadcast_Bitcast_DynamicSlice_Tanh) { HloInstruction::CreateParameter(1, starts_shape, "starts")); HloInstruction* broadcast2 = builder.AddInstruction( HloInstruction::CreateBroadcast(broadcast_shape, param0, {1})); - HloInstruction* bitcast3 = builder.AddInstruction(HloInstruction::CreateUnary( - bitcast_shape, HloOpcode::kBitcast, broadcast2)); + HloInstruction* reshape3 = builder.AddInstruction( + HloInstruction::CreateReshape(reshape_shape, broadcast2)); HloInstruction* dynamic_slice4 = builder.AddInstruction(HloInstruction::CreateDynamicSlice( - dynamic_slice_shape, bitcast3, param1, {4, 4})); + dynamic_slice_shape, reshape3, param1, {4, 4})); builder.AddInstruction(HloInstruction::CreateUnary( dynamic_slice_shape, HloOpcode::kTanh, dynamic_slice4)); @@ -293,7 +290,7 @@ TEST_F(OpcodeFusionTest, Broadcast_Bitcast_DynamicSlice_Tanh) { RunFusionAndCheckOpcodesWereFused( module.get(), - {HloOpcode::kTanh, HloOpcode::kDynamicSlice, HloOpcode::kBitcast, + {HloOpcode::kTanh, HloOpcode::kDynamicSlice, HloOpcode::kReshape, HloOpcode::kBroadcast, HloOpcode::kParameter, HloOpcode::kParameter}); } diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc index 1ef45dbec39a0880ebb123ba3fcd1fd6c89eb39a..9a3bd68c80c6e8bcdb231c63ba025d1f73619eb7 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc @@ -31,17 +31,25 @@ XfeedManager* GetXfeedManager() { return manager; } +extern const char* const kEigenMatMulF16SymbolName = + "__xla_cpu_runtime_EigenMatMulF16"; 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 kEigenSingleThreadedMatMulF16SymbolName = + "__xla_cpu_runtime_EigenSingleThreadedMatMulF16"; 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 3e1f08071119c938619d02777513e5b834077118..e61d6ea28b633398863357541e056ee887582f9c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h @@ -41,12 +41,16 @@ namespace runtime { // the actual symbol. // 2. When using ahead-of-time compilation, the linker can resolve the name // because it is a symbol in the cpu_runtime library. +extern const char* const kEigenMatMulF16SymbolName; 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 kEigenSingleThreadedMatMulF16SymbolName; 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.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.cc deleted file mode 100644 index b1c1142e8d988be2ca00809b4be505466071c72f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.cc +++ /dev/null @@ -1,43 +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_avx.h" - -#define EIGEN_USE_THREADS - -#include "third_party/eigen3/Eigen/Core" - -#ifdef TF_XLA_HAS_AVX -xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_ExpV8F32AVX( - xla::cpu::runtime::V8F32AVX x) { - return Eigen::internal::pexp(x); -} - -xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_LogV8F32AVX( - xla::cpu::runtime::V8F32AVX x) { - return Eigen::internal::plog(x); -} -#endif // TF_XLA_HAS_AVX - -namespace xla { -namespace cpu { -namespace runtime { - -const char *const kExpV8F32AVXSymbolName = "__xla_cpu_runtime_ExpV8F32AVX"; -const char *const kLogV8F32AVXSymbolName = "__xla_cpu_runtime_LogV8F32AVX"; - -} // namespace runtime -} // namespace cpu -} // namespace xla 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 e5c782f93f54dc9f8f76fce7e4735a60e8847583..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h +++ /dev/null @@ -1,60 +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" - -#if defined(__AVX__) -#include -#define TF_XLA_HAS_AVX -#endif - -namespace xla { -namespace cpu { -namespace runtime { - -extern const char *const kExpV8F32AVXSymbolName; -extern const char *const kLogV8F32AVXSymbolName; - -#ifdef TF_XLA_HAS_AVX -typedef __m256 V8F32AVX; -#endif -} // namespace runtime -} // namespace cpu -} // namespace xla - -extern "C" { - -#ifdef TF_XLA_HAS_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 8099b722f10ecb83f7cf6c58ba2abb783478b97f..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 TF_XLA_HAS_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 // TF_XLA_HAS_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 2f5d1a872aaf3868d6d27f88a4f05c778d45660f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h +++ /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. -==============================================================================*/ - -#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 -#define TF_XLA_HAS_NEON -#endif // __ARM_NEON__ - -namespace xla { -namespace cpu { -namespace runtime { - -extern const char *const kExpV4F32NEONSymbolName; -extern const char *const kLogV4F32NEONSymbolName; - -#ifdef TF_XLA_HAS_NEON -typedef float32x4_t V4F32NEON; -#endif // TF_XLA_HAS_NEON - -} // namespace runtime -} // namespace cpu -} // namespace xla - -extern "C" { - -#ifdef TF_XLA_HAS_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 // TF_XLA_HAS_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 d8ecf231cc8c859ac88e1ef1478f7107cd86a052..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 TF_XLA_HAS_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 // TF_XLA_HAS_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 aeb1eda23f76a6b5cb520b6673e0a011fa1130c7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h +++ /dev/null @@ -1,63 +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" - -// MSVC does not have __SSE4_1__ macro. Eigen enables EIGEN_VECTORIZE_SSE4_1 -// when __AVX__ is defined, we should do the same. -#if defined(__SSE4_1__) || (defined(_MSC_VER) && defined(__AVX__)) -#include -#define TF_XLA_HAS_SSE4_1 -#endif - -namespace xla { -namespace cpu { -namespace runtime { - -extern const char *const kExpV4F32SSESymbolName; -extern const char *const kLogV4F32SSESymbolName; - -#ifdef TF_XLA_HAS_SSE4_1 -typedef __m128 V4F32SSE; -#endif - -} // namespace runtime -} // namespace cpu -} // namespace xla - -extern "C" { - -#ifdef TF_XLA_HAS_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/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc index c9fc586b9a4c06eb9e1f111d8f9bd2f717990aab..8b1e20d79e90fcc32e985ffb855a1a10cdd2f2b9 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc @@ -549,7 +549,7 @@ DotOpEmitter::DotOpEmitter( 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, @@ -715,6 +715,11 @@ tensorflow::Status DotOpEmitter::Emit() { // which performs the sum-of-products (the reduction loop) before storing // the result in the output buffer. + // This routine assumes that the dot operation is not in a parallelized + // enclosing computation. + CHECK( + dot_.parent()->root_instruction()->outer_dimension_partitions().empty()); + const Shape& lhs_shape = lhs_array_.GetShape(); const Shape& rhs_shape = rhs_array_.GetShape(); @@ -919,6 +924,12 @@ tensorflow::Status DotOpEmitter::EmitCallToRuntime() { llvm::Type* float_type; const char* fn_name; switch (type) { + case F16: + fn_name = multi_threaded_eigen + ? runtime::kEigenMatMulF16SymbolName + : runtime::kEigenSingleThreadedMatMulF16SymbolName; + float_type = ir_builder_->getHalfTy(); + break; case F32: fn_name = multi_threaded_eigen ? runtime::kEigenMatMulF32SymbolName @@ -1051,7 +1062,8 @@ static bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, // The inputs and the output must // 1) be matrices with no padding, and // 2) have an allowed element type. - return output_shape.element_type() == F32 && + PrimitiveType output_primitive_type = output_shape.element_type(); + return (output_primitive_type == F32 || output_primitive_type == F16) && IsRank2WithNoPadding(lhs_shape) && IsRank2WithNoPadding(rhs_shape) && IsRank2WithNoPadding(output_shape); } diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index 71e81331897a8bb82438dd5160d2964cb88fd31f..3b8056d50500cac381a1c5ad6b05028476504a47 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -479,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) { @@ -522,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. @@ -625,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. @@ -802,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) { @@ -850,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(); @@ -929,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); @@ -957,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), @@ -1196,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 @@ -1334,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; @@ -1811,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) { @@ -1981,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) { @@ -1995,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()); } } @@ -2071,7 +2074,7 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( /*instruction=*/*root, /*operands=*/{lhs, rhs}, - /*supported_types=*/{F32})); + /*supported_types=*/{F16, F32})); llvm_ir::IrArray lhs_array(GetIrArrayFor(lhs)); llvm_ir::IrArray rhs_array(GetIrArrayFor(rhs)); 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/llvm_ir_runtime.cc b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc index 0336fa61312e5cd626ae38ddd29875bff256212a..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,27 +48,23 @@ 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); + 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, @@ -73,31 +75,230 @@ llvm::Function* EmitVectorF32TanhIfNeeded(llvm::Module* module, 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_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/parallel_cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc index cd997f07890cdc1d9a546ede58cc1d992b6416ae..07a9f0efcb64db4b2ff0c6518d4b48eee9a505e0 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc @@ -394,7 +394,7 @@ Status ParallelCpuExecutable::ExecuteComputeFunctions( for (auto& entry : *function_names_) { tensorflow::mutex_lock lock(jit_mutex_); HloInstruction* instruction = entry.first; - llvm::JITSymbol sym = jit_->FindSymbol(entry.second); + llvm::JITSymbol sym = jit_->FindCompiledSymbol(entry.second); TF_RET_CHECK(sym); InsertOrDie( &functions, instruction, diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc index deb21bf4ef5895cfdbec5c2449b6ce7b306a7008..38f16681597f298604fbca1139e0309b4409f0b1 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc @@ -130,22 +130,21 @@ int64 ParallelTaskAssignment::GetTargetParallelTaskCount( // *) Emit custom loops (kSelectAndScatter, FusionKind::kTransposeDot). // *) Tuple-shaped. // TODO(b/27458679) Parallelize instructions which are skipped here. - if (instruction->opcode() == HloOpcode::kParameter || - instruction->opcode() == HloOpcode::kConstant || - instruction->opcode() == HloOpcode::kCall || - instruction->opcode() == HloOpcode::kCustomCall || - instruction->opcode() == HloOpcode::kSelectAndScatter || - instruction->opcode() == HloOpcode::kGetTupleElement || - instruction->opcode() == HloOpcode::kBitcast || - instruction->opcode() == HloOpcode::kFft || - (instruction->opcode() == HloOpcode::kConvolution && + auto opcode = instruction->opcode(); + if (opcode == HloOpcode::kParameter || opcode == HloOpcode::kConstant || + opcode == HloOpcode::kCall || opcode == HloOpcode::kCustomCall || + opcode == HloOpcode::kDot || opcode == HloOpcode::kSelectAndScatter || + opcode == HloOpcode::kGetTupleElement || opcode == HloOpcode::kBitcast || + opcode == HloOpcode::kFft || + (opcode == HloOpcode::kConvolution && PotentiallyImplementedAsEigenConvolution(*instruction)) || PotentiallyImplementedAsEigenDot(*instruction) || - (instruction->opcode() == HloOpcode::kFusion && + (opcode == HloOpcode::kFusion && instruction->fusion_kind() != HloInstruction::FusionKind::kLoop) || ShapeUtil::IsTuple(instruction->shape())) { return 1; } + // Consult 'cost_model_' to compute target parallel task count. return cost_model_->GetParallelTaskCount(instruction); } diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..90191221eb70187290131dbe6037d17244043e4a --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/cpu/parallel_task_assignment.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_executable.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" +#include "tensorflow/core/lib/strings/str_util.h" + +namespace xla { +namespace { + +class ParallelTaskAssignmentTest : public HloVerifiedTestBase { + protected: + const HloCostAnalysis::ShapeSizeFunction shape_size_func_ = + cpu::CpuExecutable::ShapeSizeBytes; + + // Use any value larger than 2 since we only test whether a module is + // parallelized or not + const int max_parallelism_ = 10; +}; + +TEST_F(ParallelTaskAssignmentTest, DotOperationNotParallelized) { + const string hlo_string = R"( + HloModule TestTaskParallel_Dot + ENTRY Dot { + dot_lhs = f32[196614,2]{1,0} parameter(0) + dot_rhs = f32[2,1]{1,0} parameter(1) + ROOT dot = f32[196614,1]{1,0} dot(dot_lhs, dot_rhs), + lhs_contracting_dims={1}, rhs_contracting_dims={0} + } + )"; + + ParseAndVerifyModule(hlo_string); + TF_ASSERT_OK_AND_ASSIGN(bool changed, cpu::ParallelTaskAssigner( + max_parallelism_, shape_size_func_) + .Run(&module())); + EXPECT_FALSE(changed); +} + +TEST_F(ParallelTaskAssignmentTest, + FusedComputationWithDotOperationNotParallelized) { + const string hlo_string = R"( + HloModule TestTaskParallel_DotNestedInFusedComp + fused_computation.0 { + parameter.0 = f32[196614,2]{1,0} parameter(0) + parameter.0.1 = f32[2,1]{1,0} parameter(1) + parameter.0.2 = f32[196614,1]{1,0} parameter(2) + dot.0 = f32[196614,1]{1,0} dot(parameter.0, parameter.0.1), + lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT add.0 = f32[196614,1]{1,0} add(dot.0, parameter.0.2) + + } + ENTRY DotNestedInFusedComp { + parameter = f32[196614,2]{1,0} parameter(0) + parameter.1 = f32[2,1]{1,0} parameter(1) + parameter.2 = f32[196614,1]{1,0} parameter(2) + ROOT fusion = f32[196614,1]{1,0} fusion(parameter, parameter.1, + parameter.2), kind=kOutput, calls=fused_computation.0 + } + )"; + + ParseAndVerifyModule(hlo_string); + TF_ASSERT_OK_AND_ASSIGN(bool changed, cpu::ParallelTaskAssigner( + max_parallelism_, shape_size_func_) + .Run(&module())); + EXPECT_FALSE(changed); +} + +} // namespace +} // namespace xla 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..7337c907f5c83d608641b7382e75902e6f6c05d4 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_conv2d.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_conv2d.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_CONV2D_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_CONV2D_H_ +#include "third_party/eigen3/Eigen/Core" #include "tensorflow/core/platform/types.h" extern "C" { @@ -34,6 +35,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_fp16.cc b/tensorflow/compiler/xla/service/cpu/runtime_fp16.cc new file mode 100644 index 0000000000000000000000000000000000000000..af0275c8bd00c82220fbe116eb90d2692393713b --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_fp16.cc @@ -0,0 +1,133 @@ +/* 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/compiler/xla/service/cpu/runtime_fp16.h" +#include "tensorflow/core/platform/macros.h" + +namespace { +using tensorflow::uint16; +using tensorflow::uint32; + +// Helper class that lets us access the underlying bit representation +// of a float without breaking C++ strict aliasing. +class AliasedFloatInt { + public: + static_assert(sizeof(float) == sizeof(uint32), ""); + + static AliasedFloatInt FromFloat(float f) { + AliasedFloatInt value; + value.set_float(f); + return value; + } + + static AliasedFloatInt FromUInt(uint32 u) { + AliasedFloatInt value; + value.set_uint(u); + return value; + } + + void set_float(float f) { memcpy(&value_, &f, sizeof(f)); } + float as_float() const { + float f; + memcpy(&f, &value_, sizeof(f)); + return f; + } + + void set_uint(uint32 u) { value_ = u; } + uint32 as_uint() const { return value_; } + + private: + uint32 value_; +}; +} // namespace + +// __gnu_f2h_ieee and __gnu_h2f_ieee are marked as weak symbols so if XLA is +// built with compiler-rt (that also defines these symbols) we don't get a +// duplicate definition linker error. Making these symbols weak also ensures +// that the compiler-rt definitions "win", but that isn't essential. + +// Algorithm copied from Eigen. +uint16 TF_ATTRIBUTE_WEAK __gnu_f2h_ieee(float float_value) { + AliasedFloatInt f = AliasedFloatInt::FromFloat(float_value); + + const AliasedFloatInt f32infty = AliasedFloatInt::FromUInt(255 << 23); + const AliasedFloatInt f16max = AliasedFloatInt::FromUInt((127 + 16) << 23); + const AliasedFloatInt denorm_magic = + AliasedFloatInt::FromUInt(((127 - 15) + (23 - 10) + 1) << 23); + unsigned int sign_mask = 0x80000000u; + uint32 o = static_cast(0x0u); + + unsigned int sign = f.as_uint() & sign_mask; + f.set_uint(f.as_uint() ^ sign); + + // NOTE all the integer compares in this function can be safely + // compiled into signed compares since all operands are below + // 0x80000000. Important if you want fast straight SSE2 code + // (since there's no unsigned PCMPGTD). + + if (f.as_uint() >= + f16max.as_uint()) { // result is Inf or NaN (all exponent bits set) + o = (f.as_uint() > f32infty.as_uint()) ? 0x7e00 + : 0x7c00; // NaN->qNaN and Inf->Inf + } else { // (De)normalized number or zero + if (f.as_uint() < (113 << 23)) { // resulting FP16 is subnormal or zero + // use a magic value to align our 10 mantissa bits at the bottom of + // the float. as long as FP addition is round-to-nearest-even this + // just works. + f.set_float(f.as_float() + denorm_magic.as_float()); + + // and one integer subtract of the bias later, we have our final float! + o = static_cast(f.as_uint() - denorm_magic.as_uint()); + } else { + unsigned int mant_odd = + (f.as_uint() >> 13) & 1; // resulting mantissa is odd + + // update exponent, rounding bias part 1 + f.set_uint(f.as_uint() + (static_cast(15 - 127) << 23) + + 0xfff); + // rounding bias part 2 + f.set_uint(f.as_uint() + mant_odd); + // take the bits! + o = static_cast(f.as_uint() >> 13); + } + } + + o |= static_cast(sign >> 16); + return o; +} + +// Algorithm copied from Eigen. +float TF_ATTRIBUTE_WEAK __gnu_h2f_ieee(uint16 h) { + const AliasedFloatInt magic = AliasedFloatInt::FromUInt(113 << 23); + const unsigned int shifted_exp = 0x7c00 << 13; // exponent mask after shift + AliasedFloatInt o; + + o.set_uint((h & 0x7fff) << 13); // exponent/mantissa bits + unsigned int exp = shifted_exp & o.as_uint(); // just the exponent + o.set_uint(o.as_uint() + ((127 - 15) << 23)); // exponent adjust + + // handle exponent special cases + if (exp == shifted_exp) { // Inf/NaN? + o.set_uint(o.as_uint() + ((128 - 16) << 23)); // extra exp adjust + } else if (exp == 0) { // Zero/Denormal? + o.set_uint(o.as_uint() + (1 << 23)); // extra exp adjust + o.set_float(o.as_float() - magic.as_float()); // renormalize + } + + o.set_uint(o.as_uint() | (h & 0x8000) << 16); // sign bit + return o.as_float(); +} diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fp16.h b/tensorflow/compiler/xla/service/cpu/runtime_fp16.h new file mode 100644 index 0000000000000000000000000000000000000000..01d92d031904af99884c2583a8c7b5086b289d44 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_fp16.h @@ -0,0 +1,27 @@ +/* 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_RUNTIME_FP16_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FP16_H_ + +#include "tensorflow/core/platform/types.h" + +// Converts an F32 value to a F16. +extern "C" tensorflow::uint16 __gnu_f2h_ieee(float); + +// Converts an F16 value to a F32. +extern "C" float __gnu_h2f_ieee(tensorflow::uint16); + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FP16_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc b/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc index bff57d33ae23fbba8c664cbd18df77e4c35eb592..39b13183ff093611a42b3931d45f64eadb420622 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc @@ -63,30 +63,41 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, C.device(*run_options->intra_op_thread_pool()) = A.contract(B, dims); } +template +void MatMulImpl(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, + int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { + if (m == 1 || n == 1) { + // Despite being single threaded, this version of matrix * vector is faster. + xla::EigenMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); + } else { + MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); + } +} + } // namespace +void __xla_cpu_runtime_EigenMatMulF16(const void* run_options_ptr, + Eigen::half* out, Eigen::half* lhs, + Eigen::half* rhs, int64 m, int64 n, + int64 k, int32 transpose_lhs, + int32 transpose_rhs) { + MatMulImpl(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); +} + void __xla_cpu_runtime_EigenMatMulF32(const void* run_options_ptr, float* out, float* lhs, float* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - if (m == 1 || n == 1) { - // Despite being single threaded, this version of matrix * vector is faster. - xla::EigenMatVecF32(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); - } else { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); - } + MatMulImpl(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); } void __xla_cpu_runtime_EigenMatMulF64(const void* run_options_ptr, double* out, double* lhs, double* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - if (m == 1 || n == 1) { - // Despite being single threaded, this version of matrix * vector is faster. - xla::EigenMatVecF64(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); - } else { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); - } + MatMulImpl(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); } diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matmul.h b/tensorflow/compiler/xla/service/cpu/runtime_matmul.h index fdb644651dd5d0fa0345580f52ed0fb051672285..d96fe3d58bd5ffbad347e3ede3534d1d47be697a 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_matmul.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_matmul.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATMUL_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATMUL_H_ +#include "third_party/eigen3/Eigen/Core" #include "tensorflow/core/platform/types.h" extern "C" { @@ -25,6 +26,12 @@ extern "C" { // order. 'out' is a pointer to a buffer sufficiently large to hold the result // of the operation. Following standard nomenclature: lhs is m x k, // rhs is k x n, and out is m x n. +extern void __xla_cpu_runtime_EigenMatMulF16( + const void* /* xla::ExecutableRunOptions* */ run_options_ptr, + Eigen::half* out, Eigen::half* lhs, Eigen::half* rhs, tensorflow::int64 m, + tensorflow::int64 n, tensorflow::int64 k, tensorflow::int32 transpose_lhs, + tensorflow::int32 transpose_rhs); + extern void __xla_cpu_runtime_EigenMatMulF32( const void* /* xla::ExecutableRunOptions* */ run_options_ptr, float* out, float* lhs, float* rhs, tensorflow::int64 m, tensorflow::int64 n, diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matvec.cc b/tensorflow/compiler/xla/service/cpu/runtime_matvec.cc deleted file mode 100644 index 435820cdd36e2a906d9dfbe2555f4c0df623c729..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/runtime_matvec.cc +++ /dev/null @@ -1,110 +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 "third_party/eigen3/Eigen/Core" -#include "tensorflow/compiler/xla/service/cpu/runtime_matvec.h" - -using tensorflow::int32; -using tensorflow::int64; - -namespace { - -// Does mat * x or mat^T * x. -template -void MatVec(T* out_buf, T* mat_buf, T* x_buf, int64 rows, int64 cols, - int32 transpose) { - // Use an Eigen Matrix instead of a Tensor, as the GEMV from Matrix seems to - // be faster (b/30223679). See also: the matmul op kernel in TensorFlow, - // which implements the same optimization. - using Matrix = Eigen::Matrix; - using MatrixMap = Eigen::Map; - - using Vector = Eigen::Matrix; - using VectorMap = Eigen::Map; - - auto x = VectorMap(x_buf, cols); - auto out = VectorMap(out_buf, rows); - - int64 mat_rows = rows; - int64 mat_cols = cols; - - if (transpose) { - std::swap(mat_rows, mat_cols); - } - - auto mat = MatrixMap(mat_buf, mat_rows, mat_cols); - - if (transpose) { - out = mat.transpose() * x; - } else { - out = mat * x; - } -} - -// Converts matmul-style args to matvec. -template -void DispatchMatVec(T* out, T* lhs, T* rhs, int64 m, int64 n, int64 k, - int32 transpose_lhs, int32 transpose_rhs) { - // If the input is in the form x * A, where x is the vector, then bring A back - // over to the left hand side. We make use of the identity - // - // (x * A)^T = A^T * x^T - // - // We do not need to take the transpose of x or of the result since taking - // the transpose of a vector does not change the memory layout. - const int64 cols = k; - - T* mat; - T* vec; - int64 rows; - bool transpose_mat; - - bool is_mat_vec = (n == 1); - - if (is_mat_vec) { - mat = lhs; - vec = rhs; - rows = m; - transpose_mat = transpose_lhs; - } else { - mat = rhs; - vec = lhs; - rows = n; - transpose_mat = !transpose_rhs; - } - - MatVec(out, mat, vec, rows, cols, transpose_mat); -} - -} // namespace - -namespace xla { - -void EigenMatVecF32(float* out, float* lhs, float* rhs, int64 m, int64 n, - int64 k, int32 transpose_lhs, int32 transpose_rhs) { - assert((m == 1 || n == 1) && "not a matrix-vector multiply"); - DispatchMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); -} - -void EigenMatVecF64(double* out, double* lhs, double* rhs, int64 m, int64 n, - int64 k, int32 transpose_lhs, int32 transpose_rhs) { - assert((m == 1 || n == 1) && "not a matrix-vector multiply"); - DispatchMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); -} - -} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matvec.h b/tensorflow/compiler/xla/service/cpu/runtime_matvec.h index 1bd8dfb377acc1f7cfbe9a92773f87f0ef25de3a..70eb98c54169824e220d9287753c0849362eade6 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_matvec.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_matvec.h @@ -16,10 +16,86 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ +#include "third_party/eigen3/Eigen/Core" + #include "tensorflow/core/platform/types.h" namespace xla { +namespace detail { + +using tensorflow::int32; +using tensorflow::int64; + +// Does mat * x or mat^T * x. +template +void MatVec(T* out_buf, T* mat_buf, T* x_buf, int64 rows, int64 cols, + int32 transpose) { + // Use an Eigen Matrix instead of a Tensor, as the GEMV from Matrix seems to + // be faster (b/30223679). See also: the matmul op kernel in TensorFlow, + // which implements the same optimization. + using Matrix = Eigen::Matrix; + using MatrixMap = Eigen::Map; + + using Vector = Eigen::Matrix; + using VectorMap = Eigen::Map; + + auto x = VectorMap(x_buf, cols); + auto out = VectorMap(out_buf, rows); + + int64 mat_rows = rows; + int64 mat_cols = cols; + + if (transpose) { + std::swap(mat_rows, mat_cols); + } + + auto mat = MatrixMap(mat_buf, mat_rows, mat_cols); + + if (transpose) { + out = mat.transpose() * x; + } else { + out = mat * x; + } +} + +// Converts matmul-style args to matvec. +template +void DispatchMatVec(T* out, T* lhs, T* rhs, int64 m, int64 n, int64 k, + int32 transpose_lhs, int32 transpose_rhs) { + // If the input is in the form x * A, where x is the vector, then bring A back + // over to the left hand side. We make use of the identity + // + // (x * A)^T = A^T * x^T + // + // We do not need to take the transpose of x or of the result since taking + // the transpose of a vector does not change the memory layout. + const int64 cols = k; + + T* mat; + T* vec; + int64 rows; + bool transpose_mat; + + bool is_mat_vec = (n == 1); + + if (is_mat_vec) { + mat = lhs; + vec = rhs; + rows = m; + transpose_mat = transpose_lhs; + } else { + mat = rhs; + vec = lhs; + rows = n; + transpose_mat = !transpose_rhs; + } + + MatVec(out, mat, vec, rows, cols, transpose_mat); +} + +} // namespace detail + // Performs a matrix-vector multiplication using Eigen. 'lhs' and 'rhs' are // pointers to buffers containing input matrices in column-major order. 'out' is // a pointer to a buffer sufficiently large to hold the result of the @@ -30,15 +106,15 @@ namespace xla { // // TODO(b/64684907): Compare runtime performance of these functions with dot // simplification. -void EigenMatVecF32(float* out, float* lhs, float* rhs, tensorflow::int64 m, - tensorflow::int64 n, tensorflow::int64 k, - tensorflow::int32 transpose_lhs, - tensorflow::int32 transpose_rhs); - -void EigenMatVecF64(double* out, double* lhs, double* rhs, tensorflow::int64 m, - tensorflow::int64 n, tensorflow::int64 k, - tensorflow::int32 transpose_lhs, - tensorflow::int32 transpose_rhs); +template +void EigenMatVec(T* out, T* lhs, T* rhs, tensorflow::int64 m, + tensorflow::int64 n, tensorflow::int64 k, + tensorflow::int32 transpose_lhs, + tensorflow::int32 transpose_rhs) { + assert((m == 1 || n == 1) && "not a matrix-vector multiply"); + detail::DispatchMatVec(out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); +} } // namespace xla 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..44b201725b2c724f48c1a3f0373c41e76211e0c2 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h @@ -16,10 +16,25 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_SINGLE_THREADED_CONV2D_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_SINGLE_THREADED_CONV2D_H_ +#include "third_party/eigen3/Eigen/Core" #include "tensorflow/core/platform/types.h" 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/runtime_single_threaded_matmul.cc b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc index ee8eb081556d60fcf6537b1036a4a5825c4c7bf6..17303e2f0d34e531a3a56aa147608b949e0f43ae 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc @@ -57,26 +57,38 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, C = A.contract(B, dims); } +template +void SingleThreadedMatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, + int64 m, int64 n, int64 k, int32 transpose_lhs, + int32 transpose_rhs) { + if (m == 1 || n == 1) { + xla::EigenMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); + } else { + MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); + } +} + } // namespace +void __xla_cpu_runtime_EigenSingleThreadedMatMulF16( + const void* run_options_ptr, Eigen::half* out, Eigen::half* lhs, + Eigen::half* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, + int32 transpose_rhs) { + SingleThreadedMatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); +} + void __xla_cpu_runtime_EigenSingleThreadedMatMulF32( const void* run_options_ptr, float* out, float* lhs, float* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - if (m == 1 || n == 1) { - xla::EigenMatVecF32(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); - } else { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); - } + SingleThreadedMatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } void __xla_cpu_runtime_EigenSingleThreadedMatMulF64( const void* run_options_ptr, double* out, double* lhs, double* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - if (m == 1 || n == 1) { - xla::EigenMatVecF64(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); - } else { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); - } + SingleThreadedMatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } diff --git a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h index 029eb9514287d8c69cde2cfb06e0d56e78d6f165..82a1fcce594fa5b04f4fe459870991863c32a91a 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_SINGLE_THREADED_MATMUL_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_SINGLE_THREADED_MATMUL_H_ +#include "third_party/eigen3/Eigen/Core" #include "tensorflow/core/platform/types.h" extern "C" { @@ -25,6 +26,12 @@ extern "C" { // 'out' is a pointer to a buffer sufficiently large to hold the result of the // operation. Following standard nomenclature: lhs is m x k, rhs is k x n, and // out is m x n. +extern void __xla_cpu_runtime_EigenSingleThreadedMatMulF16( + const void* /* xla::ExecutableRunOptions* */ run_options_ptr, + Eigen::half* out, Eigen::half* lhs, Eigen::half* rhs, tensorflow::int64 m, + tensorflow::int64 n, tensorflow::int64 k, tensorflow::int32 transpose_lhs, + tensorflow::int32 transpose_rhs); + extern void __xla_cpu_runtime_EigenSingleThreadedMatMulF32( const void* /* xla::ExecutableRunOptions* */ run_options_ptr, float* out, float* lhs, float* rhs, tensorflow::int64 m, tensorflow::int64 n, diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc index de5e9b411905a37a7db7d05f51cca2802c1526ed..80c24eaccfc2a83f8f3f311d60860715668d0c08 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc @@ -21,20 +21,19 @@ limitations under the License. #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_fp16.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" @@ -46,36 +45,6 @@ namespace xla { namespace cpu { namespace { -// A simple SymbolResolver that delegates to the host dynamic linker. -class SimpleResolver : public llvm::LegacyJITSymbolResolver { - public: - explicit SimpleResolver(ExternalConstantPool* external_constant_pool) - : external_constant_pool_(external_constant_pool) {} - - llvm::JITSymbol findSymbol(const std::string& name) override { - 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::JITSymbol findSymbolInLogicalDylib(const std::string& name) override { - return nullptr; - } - - private: - ExternalConstantPool* external_constant_pool_; -}; - llvm::SmallVector DetectMachineAttributes() { llvm::SmallVector result; llvm::StringMap host_features; @@ -100,27 +69,6 @@ llvm::StringRef GetHostCpuName() { cpu_name.consume_back("-avx512"); return cpu_name; } - -CompilerFunctor::VectorIntrinsics GetAvailableIntrinsics() { - CompilerFunctor::VectorIntrinsics intrinsics; -#ifdef TF_XLA_HAS_SSE4_1 - intrinsics.sse_intrinsics = true; -#else - intrinsics.sse_intrinsics = false; -#endif -#ifdef TF_XLA_HAS_AVX - intrinsics.avx_intrinsics = true; -#else - intrinsics.avx_intrinsics = false; -#endif -#ifdef TF_XLA_HAS_NEON - intrinsics.neon_intrinsics = true; -#else - intrinsics.neon_intrinsics = false; -#endif - return intrinsics; -} - } // namespace SimpleOrcJIT::SimpleOrcJIT(const llvm::TargetOptions& target_options, @@ -139,49 +87,71 @@ 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 { + return this->ResolveRuntimeSymbol(name); + }, + [](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( - std::unique_ptr module) { - auto handle = cantFail(compile_layer_.addModule( - std::move(module), MakeUnique(external_constant_pool()))); - module_handles_.push_back(handle); - return handle; +llvm::JITSymbol SimpleOrcJIT::ResolveRuntimeSymbol(const std::string& name) { + 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; } -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)); +SimpleOrcJIT::VModuleKeyT SimpleOrcJIT::AddModule( + std::unique_ptr module) { + auto key = execution_session_.allocateVModule(); + cantFail(compile_layer_.addModule(key, std::move(module))); + module_keys_.push_back(key); + return key; } -llvm::JITSymbol SimpleOrcJIT::FindSymbol(const std::string& name) { - std::string mangled_name; - { - llvm::raw_string_ostream mangled_name_stream(mangled_name); - llvm::Mangler::getNameWithPrefix(mangled_name_stream, name, data_layout_); - } +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::FindCompiledSymbol(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, name, /*ExportedSymbolsOnly=*/true)) { return symbol; } @@ -208,29 +178,24 @@ 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(EigenMatMulF16); REGISTER_CPU_RUNTIME_SYMBOL(EigenMatMulF32); REGISTER_CPU_RUNTIME_SYMBOL(EigenMatMulF64); + REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedConvF16); REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedConvF32); + REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulF16); REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulF32); REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulF64); -#ifdef TF_XLA_HAS_NEON - REGISTER_CPU_RUNTIME_SYMBOL(ExpV4F32NEON); - REGISTER_CPU_RUNTIME_SYMBOL(LogV4F32NEON); -#endif -#ifdef TF_XLA_HAS_SSE4_1 - REGISTER_CPU_RUNTIME_SYMBOL(ExpV4F32SSE); - REGISTER_CPU_RUNTIME_SYMBOL(LogV4F32SSE); -#endif -#ifdef TF_XLA_HAS_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); + registry->Register("__gnu_f2h_ieee", reinterpret_cast(__gnu_f2h_ieee)); + registry->Register("__gnu_h2f_ieee", reinterpret_cast(__gnu_h2f_ieee)); + #undef REGISTER_CPU_RUNTIME_SYMBOL // Register both the f32 (float) and f64 (double) versions of a libm symbol. @@ -273,15 +238,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)); @@ -292,7 +257,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..aaeff2de8785b99d271f13b261c63118bcf7bd4a 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" @@ -44,11 +46,9 @@ namespace cpu { class SimpleOrcJIT { public: using ObjLayerT = llvm::orc::RTDyldObjectLinkingLayer; - using CompileFtor = - std::function( - llvm::Module&)>; + using CompileFtor = std::function; 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,16 +78,16 @@ 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. - llvm::JITSymbol FindSymbol(const std::string& name); + llvm::JITSymbol FindCompiledSymbol(const std::string& name); llvm::TargetMachine* target_machine() const { return target_machine_.get(); } @@ -96,10 +96,15 @@ class SimpleOrcJIT { } private: - std::vector module_handles_; + llvm::JITSymbol ResolveRuntimeSymbol(const std::string& name); + + 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/vector_support_library.cc b/tensorflow/compiler/xla/service/cpu/vector_support_library.cc index 128b465be239130918687d8e2ba0458684086ee1..cd1165e23812861ba9951546b7dd744529232196 100644 --- a/tensorflow/compiler/xla/service/cpu/vector_support_library.cc +++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.cc @@ -15,6 +15,7 @@ limitations under the License. #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" @@ -35,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); } @@ -50,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()) { @@ -93,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()); } @@ -102,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()); @@ -230,6 +370,9 @@ std::vector VectorSupportLibrary::ComputeHorizontalSums( std::vector VectorSupportLibrary::ComputeAvxOptimizedHorizontalSums( std::vector vectors, llvm::Value* init_values) { + // vectors are N llvm vector values, each with N elements. + int64 lane_width = vectors.size(); + while (vectors.size() != 2) { std::vector new_vectors; for (int i = 0; i < vectors.size(); i += 2) { @@ -250,10 +393,14 @@ VectorSupportLibrary::ComputeAvxOptimizedHorizontalSums( high = AddInternal(ExtractHighHalf(init_values), high); } + // `low` has the first `lane_width / 2` horizontal reductions, and `high` has + // the next `lane_width / 2` horizontal reductions. + std::vector results; - for (int i = 0; i < 8; i++) { + for (int i = 0; i < lane_width; i++) { llvm::Value* scalar_result = ir_builder()->CreateExtractElement( - i < 4 ? low : high, ir_builder()->getInt32(i % 4), name()); + i < (lane_width / 2) ? low : high, + ir_builder()->getInt32(i % (lane_width / 2)), name()); results.push_back(scalar_result); } diff --git a/tensorflow/compiler/xla/service/cpu/vector_support_library.h b/tensorflow/compiler/xla/service/cpu/vector_support_library.h index 8fbac2a6670f8ef18c00877a1566bd4ab896a7c8..6479bf76aab581ae3ec2923d98dab53720cab203 100644 --- a/tensorflow/compiler/xla/service/cpu/vector_support_library.h +++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.h @@ -26,6 +26,16 @@ limitations under the License. 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 { @@ -41,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, @@ -144,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 @@ -163,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_; diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index a803b3171f9afa6297553c5507c4f9aa45e420ab..56723e765048698baedc50ae7b189d0287ee56b8 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -190,6 +190,7 @@ class DfsHloVisitorBase { virtual Status HandleInfeed(HloInstructionPtr hlo) = 0; virtual Status HandleOutfeed(HloInstructionPtr hlo) = 0; + virtual Status HandleHostCompute(HloInstructionPtr hlo) = 0; virtual Status HandleRng(HloInstructionPtr hlo) = 0; virtual Status HandleReverse(HloInstructionPtr hlo) = 0; virtual Status HandleSort(HloInstructionPtr hlo) = 0; @@ -213,6 +214,7 @@ class DfsHloVisitorBase { virtual Status HandleSelectAndScatter(HloInstructionPtr hlo) = 0; virtual Status HandleWhile(HloInstructionPtr hlo) = 0; virtual Status HandleConditional(HloInstructionPtr hlo) = 0; + virtual Status HandleGather(HloInstructionPtr hlo) = 0; virtual Status HandlePad(HloInstructionPtr hlo) = 0; 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 170adb3d241b3648bc53f96dde9866f0b794f80a..ecda5288ee17a3856ce95f0caa327c3524fd180b 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -103,6 +103,9 @@ class DfsHloVisitorWithDefaultBase Status HandleOutfeed(HloInstructionPtr outfeed) override { return DefaultAction(outfeed); } + Status HandleHostCompute(HloInstructionPtr host_compute) override { + return DefaultAction(host_compute); + } Status HandleReverse(HloInstructionPtr reverse) override { return DefaultAction(reverse); } @@ -185,6 +188,9 @@ class DfsHloVisitorWithDefaultBase Status HandleSendDone(HloInstructionPtr send_done) override { return DefaultAction(send_done); } + Status HandleGather(HloInstructionPtr gather) override { + return DefaultAction(gather); + } // Invoked to inform the visitor that the traversal has completed, and that // the root was "root". diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index 9780bac16ec17eed2c1df64f01bcb753e26b46f0..b6a0903b0eeaa04d8bc1488378c148b2016c5d48 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -226,7 +226,7 @@ StatusOr ElementalIrEmitter::EmitIntegerUnaryOp( if (primitive_util::IsIntegralType(to_type)) { return ir_builder_->CreateIntCast( operand_value, llvm_ir::PrimitiveTypeToIrType(to_type, module_), - primitive_util::IsSignedIntegralType(to_type)); + primitive_util::IsSignedIntegralType(from_type)); } if (primitive_util::IsFloatingPointType(to_type)) { if (to_type == BF16) { @@ -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); @@ -1000,6 +1003,30 @@ StatusOr ElementalIrEmitter::EmitReducePrecision( ir_builder_); } +static llvm::Value* SaturateShiftIfNecessary(llvm::IRBuilder<>* ir_builder, + llvm::Value* lhs, llvm::Value* rhs, + llvm::Value* shift_result, + bool saturate_to_sign_bit) { + llvm::IntegerType* integer_type = + llvm::cast(lhs->getType()); + unsigned integer_bitsize = integer_type->getBitWidth(); + llvm::ConstantInt* integer_bitsize_constant = + llvm::ConstantInt::get(integer_type, integer_bitsize); + llvm::ConstantInt* zero = llvm::ConstantInt::get(integer_type, 0); + llvm::ConstantInt* minus_one = llvm::ConstantInt::get(integer_type, -1); + llvm::Value* saturated_value; + if (saturate_to_sign_bit) { + saturated_value = ir_builder->CreateSelect( + ir_builder->CreateICmpSLT(lhs, zero), minus_one, zero); + } else { + saturated_value = zero; + } + llvm::Value* shift_amt_in_range = + ir_builder->CreateICmpULT(rhs, integer_bitsize_constant, "shft.chk"); + return ir_builder->CreateSelect(shift_amt_in_range, shift_result, + saturated_value); +} + StatusOr ElementalIrEmitter::EmitIntegerBinaryOp( const HloInstruction* op, llvm::Value* lhs_value, llvm::Value* rhs_value, bool is_signed) const { @@ -1040,33 +1067,60 @@ 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: return ir_builder_->CreateOr(lhs_value, rhs_value); - case HloOpcode::kShiftLeft: - return ir_builder_->CreateShl(lhs_value, rhs_value); + + // Shifting out bits >= the number of bits in the type being shifted + // produces a poison value in LLVM which is basically "deferred undefined + // behavior" -- doing something observable with such a value precipitates + // UB. We replace the poison value with a constant to avoid this deferred + // UB. case HloOpcode::kShiftRightArithmetic: - return ir_builder_->CreateAShr(lhs_value, rhs_value); + return SaturateShiftIfNecessary( + ir_builder_, lhs_value, rhs_value, + ir_builder_->CreateAShr(lhs_value, rhs_value), + /*saturate_to_sign_bit=*/true); + case HloOpcode::kShiftLeft: + return SaturateShiftIfNecessary( + ir_builder_, lhs_value, rhs_value, + ir_builder_->CreateShl(lhs_value, rhs_value), + /*saturate_to_sign_bit=*/false); case HloOpcode::kShiftRightLogical: - return ir_builder_->CreateLShr(lhs_value, rhs_value); + return SaturateShiftIfNecessary( + ir_builder_, lhs_value, rhs_value, + ir_builder_->CreateLShr(lhs_value, rhs_value), + /*saturate_to_sign_bit=*/false); default: return Unimplemented("binary integer op '%s'", HloOpcodeString(op->opcode()).c_str()); } } +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 { @@ -1363,7 +1417,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]( @@ -1457,15 +1522,12 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( case HloOpcode::kBroadcast: return [this, hlo, &operand_to_generator]( const IrArray::Index& target_index) -> StatusOr { + const HloInstruction* operand = hlo->operand(0); // The `dimensions` member of the broadcast instruction maps from // input dimensions to output dimensions. - const HloInstruction* operand = hlo->operand(0); - int64 rank = ShapeUtil::Rank(operand->shape()); - IrArray::Index source_index(rank); - for (int64 i = 0; i < rank; ++i) { - source_index[i] = target_index[hlo->dimensions(i)]; - } - return operand_to_generator.at(operand)(source_index); + return operand_to_generator.at( + operand)(target_index.SourceIndexOfBroadcast( + hlo->shape(), operand->shape(), hlo->dimensions(), ir_builder_)); }; case HloOpcode::kSlice: return [this, hlo, &operand_to_generator]( @@ -1657,6 +1719,14 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( SetToFirstInsertPoint(if_data.after_block, ir_builder_); return ir_builder_->CreateLoad(ret_value_addr); }; + case HloOpcode::kBitcast: + CHECK_EQ(ShapeUtil::ElementsIn(hlo->shape()), + ShapeUtil::ElementsIn(hlo->operand(0)->shape())); + return [this, hlo, &operand_to_generator](const IrArray::Index& index) { + const HloInstruction* operand = hlo->operand(0); + return operand_to_generator.at(operand)(index.SourceIndexOfBitcast( + hlo->shape(), operand->shape(), ir_builder_)); + }; case HloOpcode::kReshape: CHECK_EQ(ShapeUtil::ElementsIn(hlo->shape()), ShapeUtil::ElementsIn(hlo->operand(0)->shape())); 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 90481c7a88f90edea5399ee44aee2d2c77fc115f..be92b1629a2d8dae57b315751bd4f7f9ccddf171 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/lib/strings/proto_serialization.h" #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/env.h" diff --git a/tensorflow/compiler/xla/service/gather_expander.cc b/tensorflow/compiler/xla/service/gather_expander.cc new file mode 100644 index 0000000000000000000000000000000000000000..58c62d8ce951fc99538cf17929722e90ee809122 --- /dev/null +++ b/tensorflow/compiler/xla/service/gather_expander.cc @@ -0,0 +1,359 @@ +/* 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/compiler/xla/service/gather_expander.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/while_util.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/util.h" + +namespace xla { +using tensorflow::gtl::ArraySlice; + +static StatusOr TransposeIndexVectorDimToLast( + HloInstruction* gather_indices, int64 index_vector_dim) { + const Shape& gather_indices_shape = gather_indices->shape(); + if (index_vector_dim == (gather_indices_shape.dimensions_size() - 1)) { + return gather_indices; + } + std::vector permutation; + permutation.reserve(gather_indices_shape.dimensions_size()); + for (int64 i = 0, e = gather_indices_shape.dimensions_size(); i < e; i++) { + if (i != index_vector_dim) { + permutation.push_back(i); + } + } + permutation.push_back(index_vector_dim); + return MakeTransposeHlo(gather_indices, permutation); +} + +// If the gather_indices holds scalar indices (i.e. gather_indices has rank N +// and index_vector_dim is N) then reshape it to have a trailing degenerate +// dimension. This makes the code for slicing out the index vector more +// uniform. +static StatusOr DeScalarizeGatherIndices( + HloInstruction* gather_indices, int64 index_vector_dim) { + const Shape& gather_indices_shape = gather_indices->shape(); + if (index_vector_dim != gather_indices_shape.dimensions_size()) { + return gather_indices; + } + + int64 last_index = gather_indices_shape.dimensions( + gather_indices_shape.dimensions_size() - 1); + return ExpandLastDimIntoNDims(gather_indices, {last_index, 1}); +} + +// Canonicalizes the gather_indices tensors so that we only have deal with some +// specific cases in the while loop that does the heavy lifting. +// +// See the "High Level Algorithm" section for a broader picture. +static StatusOr CanonicalizeGatherIndices( + HloInstruction* gather_indices, int64 index_vector_dim) { + // If gather_indices holds scalar indices, normalize it to hold index vectors + // of size 1. + TF_ASSIGN_OR_RETURN( + HloInstruction * descalarized_gather_indices, + DeScalarizeGatherIndices(gather_indices, index_vector_dim)); + + // Transpose the non-index-vector dimensions to the front. + TF_ASSIGN_OR_RETURN(HloInstruction * transposed_gather_indices, + TransposeIndexVectorDimToLast(descalarized_gather_indices, + index_vector_dim)); + + // If there is only one index (i.e. gather_indices has rank 1 and this gather + // is really just a dynamic slice) add a leading degenerate dimension for + // uniformity. Otherwise create a "collapsed" leading dimension that subsumes + // all of the non-index-vector dimensions. + const Shape& shape = transposed_gather_indices->shape(); + if (shape.dimensions_size() == 1) { + return ExpandFirstDimIntoNDims(gather_indices, {1, shape.dimensions(0)}); + } else { + return CollapseFirstNDims(transposed_gather_indices, + shape.dimensions_size() - 1); + } +} + +// Expands out the gather dimensions in the accumulator produced by the while +// loop. +static StatusOr ExpandGatherDimsInAccumulator( + const Shape& gather_indices_shape, HloInstruction* accumulator, + int64 index_vector_dim) { + std::vector output_gather_dim_bounds; + output_gather_dim_bounds.reserve(gather_indices_shape.dimensions_size()); + for (int64 i = 0, e = gather_indices_shape.dimensions_size(); i < e; i++) { + if (i != index_vector_dim) { + output_gather_dim_bounds.push_back(gather_indices_shape.dimensions(i)); + } + } + + if (output_gather_dim_bounds.empty()) { + // If output_gather_dim_bounds is empty we must be lowering a (effectively) + // dynamic-slice. + CHECK_EQ(accumulator->shape().dimensions(0), 1); + return CollapseFirstNDims(accumulator, 2); + } + + return ExpandFirstDimIntoNDims(accumulator, output_gather_dim_bounds); +} + +// Expand an index vector from the gather_indices tensor into a vector that can +// be used to dynamic-slice out of the gather operand. +static StatusOr ExpandIndexVectorIntoOperandSpace( + HloInstruction* index_vector, const GatherDimensionNumbers& dim_numbers, + int64 operand_rank) { + HloComputation* computation = index_vector->parent(); + const Shape& index_shape = index_vector->shape(); + HloInstruction* zero = + computation->AddInstruction(HloInstruction::CreateConstant( + Literal::CreateFromDimensions(index_shape.element_type(), {1}))); + + // We extract out individual components from the smaller index and concatenate + // them (interspersing zeros as needed) into the larger index. + std::vector expanded_index_components; + + for (int i = 0; i < operand_rank; i++) { + int64 index_vector_dim_index = + FindIndex(dim_numbers.gather_dims_to_operand_dims(), i); + if (index_vector_dim_index != + dim_numbers.gather_dims_to_operand_dims_size()) { + TF_ASSIGN_OR_RETURN( + HloInstruction * component_to_concat, + MakeSliceHlo(index_vector, /*start_indices=*/{index_vector_dim_index}, + /*limit_indices=*/{index_vector_dim_index + 1}, + /*strides=*/{1})); + expanded_index_components.push_back(component_to_concat); + } else { + expanded_index_components.push_back(zero); + } + } + + return MakeConcatHlo(expanded_index_components, /*dimension=*/0); +} + +// This generates the body of the while that implements the main data movement +// behavior of gather using dynamic-slice and dynamic-update-slice. +static StatusOr> GatherLoopBody( + const HloInstruction& gather, HloInstruction* induction_var, + const std::vector& incoming_loop_state) { + CHECK_EQ(incoming_loop_state.size(), 3); + HloInstruction* const operand = incoming_loop_state[0]; + HloInstruction* const gather_indices = incoming_loop_state[1]; + HloInstruction* const output_accumulator = incoming_loop_state[2]; + + int64 index_vector_size = gather_indices->shape().dimensions(1); + + TF_ASSIGN_OR_RETURN( + HloInstruction * induction_var_as_vector, + MakeBroadcastHlo(induction_var, /*broadcast_dimensions=*/{}, + /*result_shape_bounds=*/{1})); + + TF_ASSIGN_OR_RETURN( + HloInstruction * index_into_gather_indices, + PadVectorWithZeros(induction_var_as_vector, + /*zeros_to_prepend=*/0, /*zeros_to_append=*/1)); + + TF_ASSIGN_OR_RETURN( + HloInstruction * index_vector_2d, + MakeDynamicSliceHlo(gather_indices, index_into_gather_indices, + {1, index_vector_size})); + + TF_ASSIGN_OR_RETURN(HloInstruction * index_vector, + ElideDegenerateDims(index_vector_2d, {0})); + + TF_ASSIGN_OR_RETURN(HloInstruction * gathered_slice_start, + ExpandIndexVectorIntoOperandSpace( + index_vector, gather.gather_dimension_numbers(), + operand->shape().dimensions_size())); + + TF_ASSIGN_OR_RETURN(HloInstruction * gathered_slice, + MakeDynamicSliceHlo(operand, gathered_slice_start, + gather.gather_window_bounds())); + + TF_ASSIGN_OR_RETURN( + HloInstruction * gathered_slice_for_update, + ExpandFirstDimIntoNDims(gathered_slice, + {1, gathered_slice->shape().dimensions(0)})); + + TF_ASSIGN_OR_RETURN( + HloInstruction * index_vector_into_accumulator, + PadVectorWithZeros( + induction_var_as_vector, /*zeros_to_prepend=*/0, + /*zeros_to_append=*/gathered_slice->shape().dimensions_size())); + + TF_ASSIGN_OR_RETURN( + HloInstruction * updated_accumulator, + MakeDynamicUpdateSliceHlo(output_accumulator, gathered_slice_for_update, + index_vector_into_accumulator)); + + // New loop state -- only the accumulator has changed. The + // WhileUtil::MakeCountedLoop functions takes care of the induction variable + // and the while loop exit condition. + return StatusOr>{ + {operand, gather_indices, updated_accumulator}}; +} + +static StatusOr CreateGatherLoopAccumulatorInitValue( + HloComputation* computation, PrimitiveType element_type, + ArraySlice window_bounds, int64 gather_loop_trip_count) { + std::vector accumulator_state_shape_dims; + accumulator_state_shape_dims.reserve(1 + window_bounds.size()); + accumulator_state_shape_dims.push_back(gather_loop_trip_count); + c_copy(window_bounds, std::back_inserter(accumulator_state_shape_dims)); + return BroadcastZeros(computation, element_type, + accumulator_state_shape_dims); +} + +static StatusOr ElideWindowDimsFromAccumulator( + HloInstruction* accumulator, const GatherDimensionNumbers& dim_numbers) { + std::vector dims_to_elide; + dims_to_elide.reserve(dim_numbers.elided_window_dims_size()); + for (int64 elided_window_dim : dim_numbers.elided_window_dims()) { + dims_to_elide.push_back(elided_window_dim + 1); + } + + return ElideDegenerateDims(accumulator, dims_to_elide); +} + +// `accumulator` is almost the tensor the gather operation would have produced, +// except that it has the dimensions in the wrong order -- the gather dimensions +// are the major dimensions and the window dimensions are the minor dimensions. +// Fix this up with a transpose. +static StatusOr PermuteGatherAndWindowDims( + HloInstruction* accumulator, ArraySlice output_window_dims, + int64 output_rank) { + std::vector permutation; + permutation.reserve(output_rank); + + int64 gather_idx_counter = 0; + int64 window_idx_counter = output_rank - output_window_dims.size(); + for (int64 i = 0; i < output_rank; i++) { + bool is_window_dim = c_binary_search(output_window_dims, i); + if (is_window_dim) { + permutation.push_back(window_idx_counter++); + } else { + permutation.push_back(gather_idx_counter++); + } + } + + return MakeTransposeHlo(accumulator, permutation); +} + +// High Level Algorithm +// +// We follow the following steps in sequence: +// +// 1. We canonicalize the gather_indices tensor such that it has rank +// 2 (i.e. is a matrix) where each row is an index vector into the +// operand. +// 2. We iterate over the set of indices in the canonicalized +// gather_indices tensor using a while loop, accumulating slices +// of the operand tensor into an accumulator using +// DynamicUpdateSlice. +// 3. The accumulator result from the while loop from (2) is then +// reshaped to split out all the individual gather dimensions and +// then transposed to give the final result. +// +// As an example, if we started with the following operation: +// +// HloModule TensorFlowGatherMultipleBatchDims +// +// ENTRY main { +// operand = s32[3,3] parameter(0) +// indices = s32[2,2] parameter(1) +// ROOT gather = s32[2,3,2] gather(operand, indices), +// output_window_dims={1}, +// elided_window_dims={1}, +// gather_dims_to_operand_dims={1}, +// index_vector_dim=2, +// window_bounds={3, 1} +// } +// +// We'd first reshape indices to s32[4,1], where each row is an index +// into operand. We'd then run a loop to slice out 4 tensors of shape +// [3,1] out of operand into an accumulator of shape [4,3,1]. We then +// reshape this result to [2,2,3] and finally transpose it to [2,3,2]. + +StatusOr GatherExpander::ExpandGather( + HloInstruction* gather_instr) { + HloComputation* computation = gather_instr->parent(); + HloInstruction* operand = gather_instr->mutable_operand(0); + HloInstruction* gather_indices = gather_instr->mutable_operand(1); + const Shape& output_shape = gather_instr->shape(); + int64 output_rank = output_shape.dimensions_size(); + + const GatherDimensionNumbers& dim_numbers = + gather_instr->gather_dimension_numbers(); + + TF_ASSIGN_OR_RETURN(HloInstruction * canonical_gather_indices, + CanonicalizeGatherIndices( + gather_indices, dim_numbers.index_vector_dim())); + + const int64 gather_loop_trip_count = + canonical_gather_indices->shape().dimensions(0); + + TF_ASSIGN_OR_RETURN( + HloInstruction * accumulator_init, + CreateGatherLoopAccumulatorInitValue( + computation, output_shape.element_type(), + gather_instr->gather_window_bounds(), gather_loop_trip_count)); + + StatusOr> gather_loop_result_or_error = + WhileUtil::MakeCountedLoop( + computation, gather_loop_trip_count, + {operand, canonical_gather_indices, accumulator_init}, + [&](HloInstruction* indvar, + const std::vector& loop_state) { + return GatherLoopBody(*gather_instr, indvar, loop_state); + }); + + TF_ASSIGN_OR_RETURN(std::vector gather_loop_result, + gather_loop_result_or_error); + + HloInstruction* accumulator_result = gather_loop_result.back(); + TF_ASSIGN_OR_RETURN( + HloInstruction * accumulator_with_window_dims_elided, + ElideWindowDimsFromAccumulator(accumulator_result, dim_numbers)); + + TF_ASSIGN_OR_RETURN( + HloInstruction * accumulator_with_output_gather_dims_decanonicalized, + ExpandGatherDimsInAccumulator(gather_indices->shape(), + accumulator_with_window_dims_elided, + dim_numbers.index_vector_dim())); + + return PermuteGatherAndWindowDims( + accumulator_with_output_gather_dims_decanonicalized, + AsInt64Slice(dim_numbers.output_window_dims()), output_rank); +} + +StatusOr GatherExpander::Run(HloModule* module) { + std::vector gather_instrs; + for (HloComputation* computation : module->MakeNonfusionComputations()) { + c_copy_if(computation->instructions(), std::back_inserter(gather_instrs), + [](HloInstruction* inst) { + return inst->opcode() == HloOpcode::kGather; + }); + } + + for (HloInstruction* inst : gather_instrs) { + TF_ASSIGN_OR_RETURN(HloInstruction * expanded_root, ExpandGather(inst)); + TF_RETURN_IF_ERROR(inst->parent()->ReplaceInstruction(inst, expanded_root)); + } + + return !gather_instrs.empty(); +} +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gather_expander.h b/tensorflow/compiler/xla/service/gather_expander.h new file mode 100644 index 0000000000000000000000000000000000000000..c1fc8574da99fff223c7dbb570b4533f76905b9a --- /dev/null +++ b/tensorflow/compiler/xla/service/gather_expander.h @@ -0,0 +1,37 @@ +/* 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_GATHER_EXPANDER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GATHER_EXPANDER_H_ + +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// This pass rewrites gather operations into (roughly) while loops of dynamic +// slices. This lets backends that don't support gather directly to +// nevertheless have a minimum level of support. +class GatherExpander : public HloPassInterface { + public: + tensorflow::StringPiece name() const override { return "gather_expander"; } + StatusOr Run(HloModule* module) override; + + private: + StatusOr ExpandGather(HloInstruction* gather_instr); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GATHER_EXPANDER_H_ diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.cc b/tensorflow/compiler/xla/service/generic_transfer_manager.cc index 78dc0ad4fcd167c93f19d0c2b18ea72d666897ef..a99e2b7794a399047fb5a77a140bd333214e3f23 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.cc @@ -38,14 +38,7 @@ namespace xla { GenericTransferManager::GenericTransferManager(se::Platform::Id platform_id, size_t pointer_size) - : platform_id_(platform_id), pointer_size_(pointer_size) { - // We currently only support kHostPlatformId for CPU, kCudaPlatformId for - // GPU and kInterpreterPlatformId for Interpreter. Before supporting other - // platforms, we need to test this transfer manager on them. - CHECK(platform_id_ == se::host::kHostPlatformId || - platform_id_ == se::interpreter::kInterpreterPlatformId || - platform_id_ == se::cuda::kCudaPlatformId); -} + : platform_id_(platform_id), pointer_size_(pointer_size) {} se::Platform::Id GenericTransferManager::PlatformId() const { return platform_id_; diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 80c2eed1097e4a7dcbf29b9b1c02fb9964983368..a3b7e10ae8df080879ce98b02b83f246bb19204b 100644 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -129,8 +129,11 @@ 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", @@ -262,6 +265,7 @@ cc_library( ], deps = [ ":buffer_allocations", + ":cudnn_convolution_runner", ":infeed_manager", ":ir_emission_utils", ":partition_assignment", @@ -309,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", @@ -325,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", ], ) @@ -358,6 +397,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", ], ) @@ -398,8 +438,10 @@ tf_cc_test( ":fusion_merger", ":instruction_fusion", "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", ], ) @@ -413,6 +455,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo_creation_utils", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/compiler/xla/service:shape_inference", ], @@ -446,7 +489,8 @@ 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", @@ -470,9 +514,11 @@ cc_library( "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:buffer_liveness", "//tensorflow/compiler/xla/service:call_inliner", + "//tensorflow/compiler/xla/service:conditional_simplifier", "//tensorflow/compiler/xla/service:dot_decomposer", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:flatten_call_graph", + "//tensorflow/compiler/xla/service:gather_expander", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_constant_folding", "//tensorflow/compiler/xla/service:hlo_cse", diff --git a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc index ed78fef4113bd9f7048ca3c8c2d4e38c5ec4762a..2029c303d47e9a62135b003c3bd9be6f8b3438d4 100644 --- a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc +++ b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc @@ -98,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/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/cudnn_convolution_rewriter.h b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h new file mode 100644 index 0000000000000000000000000000000000000000..0c0578d88840fed1d77f7456c9acef27dec380f5 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.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_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" + +namespace xla { +namespace gpu { + +// 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 "cudnn-convolution-rewriter"; + } + + StatusOr Run(HloModule* module) override; +}; + +} // namespace gpu +} // namespace xla + +#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/fusion_merger.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc index c137fbc97e29e24edb3603c611a5c8f093bc62a6..3cd30b754c3242f00c704de1afab2282ed827b41 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc @@ -45,6 +45,7 @@ void MaybeResolveTupleElements(HloInstruction* instruction, // Returns the bytes read by fusion parameter 'param', by returning the byte // size of 'param' shape (or the cumulative byte sizes of all leaf tuple // elements if 'param' is tuple-shaped). +// // In the special case where all users of 'param' (or all users of a leaf // tuple element if 'param' is tuple-shaped) are Slice instructions, the size // of each slice instruction is accumulated instead, to give a more accurate @@ -63,11 +64,10 @@ double CalculateBytesReadByFusionParameter(HloInstruction* param) { // Slice for a more accurate estimate of bytes read. double bytes = 0.0; for (auto& instruction : instructions) { - if (std::all_of(instruction->users().begin(), instruction->users().end(), - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kSlice || - instruction->opcode() == HloOpcode::kDynamicSlice; - })) { + if (c_all_of(instruction->users(), [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kSlice || + instruction->opcode() == HloOpcode::kDynamicSlice; + })) { // All users are slice: accumulate bytes of all user slice instructions. for (auto& user : instruction->users()) { bytes += ShapeUtil::ByteSizeOf(user->shape()); @@ -199,6 +199,7 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { ++total_visited_; // Skip 'fusion' instruction if there are no users into which we can merge. if (fusion->users().empty()) { + VLOG(3) << "Not merging " << fusion->name() << ": Has no users."; ++num_fail_no_users_; return Status::OK(); } @@ -208,24 +209,27 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { // Input fusion instructions need to be rooted at a particular HLO (e.g. // kReduce), so they shouldn't be further fused either. if (fusion->fusion_kind() != HloInstruction::FusionKind::kLoop) { + VLOG(3) << "Not merging " << fusion->name() << ": Is not loop fusion."; ++num_fail_not_loop_fusion_; return Status::OK(); } // Skip multiple output fusion. It's not yet supported. if (fusion->IsMultiOutputFusion()) { + VLOG(3) << "Not merging " << fusion->name() << ": Is multi-output fusion."; ++num_fail_not_loop_fusion_; return Status::OK(); } // Skip 'fusion' instruction if we cannot merge into all of its users. // Merging into all users enables the removal of 'fusion' from the // computation. - if (!std::all_of(fusion->users().begin(), fusion->users().end(), - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kFusion && - instruction->fusion_kind() == - HloInstruction::FusionKind::kLoop; - })) { + if (!c_all_of(fusion->users(), [](const HloInstruction* user) { + return user->opcode() == HloOpcode::kFusion && + (user->fusion_kind() == HloInstruction::FusionKind::kLoop || + user->fusion_kind() == HloInstruction::FusionKind::kInput); + })) { + VLOG(3) << "Not merging " << fusion->name() + << ": Some of its users are not loop/input fusion kernels."; ++num_fail_merge_all_users_; return Status::OK(); } @@ -233,18 +237,17 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { // Skip 'fusion' instruction if any of its fused instructions are expensive. // This is done to avoid the duplication of expensive instructions, which // would occur if 'fusion' were merged into multiple users. + // // If 'fusion' has just one user, then an earlier fusion pass chose not to // fuse this producer/comsumer pair (likely because of expensive instruction // re-use by the consumer), and so we honor that choice here as well. - if (!std::all_of(fusion->fused_instructions().begin(), - fusion->fused_instructions().end(), - [](const HloInstruction* instruction) { - if (instruction->opcode() != HloOpcode::kParameter && - GpuInstructionFusion::IsExpensive(*instruction)) { - return false; - } - return true; - })) { + if (c_any_of(fusion->fused_instructions(), + [](const HloInstruction* instruction) { + return instruction->opcode() != HloOpcode::kParameter && + GpuInstructionFusion::IsExpensive(*instruction); + })) { + VLOG(3) << "Not merging " << fusion->name() + << ": Contains one or more expensive instructions."; ++num_fail_expensive_fused_instruction_; return Status::OK(); } @@ -253,6 +256,8 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { // exceeds the threshold value. if (CalculateFlopsToBytesRatio(fusion) > FusionMerger::GetThresholdFlopsToBytesRatio()) { + VLOG(3) << "Not merging " << fusion->name() + << ": flops-to-bytes ratio is not favorable."; ++num_fail_flops_to_byte_ratio_; return Status::OK(); } @@ -265,6 +270,9 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { const double merged_to_current_bytes_ratio = merged_bytes_transferred / std::max(1.0, current_bytes_transferred); if (merged_to_current_bytes_ratio > 1.10) { + VLOG(3) << "Not merging " << fusion->name() + << ": merged-to-current-bytes-ratio of " + << merged_to_current_bytes_ratio << " is not favorable."; ++num_fail_net_bytes_transferred_ratio_; return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc index deef5966b80d1b7f16e9982eed9ac5c7131e9d73..2217776c7d5a5f92c520d56222988f80401be9e4 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc @@ -16,257 +16,21 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/fusion_merger.h" #include "tensorflow/compiler/xla/service/gpu/instruction_fusion.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" namespace xla { namespace gpu { namespace { -class FusionMergerTest : public HloTestBase { - protected: - FusionMergerTest() : module_(CreateNewModule()) {} - - // Builds the following computation: - // - // Param - // / | \ - // / | \ - // OnesVec GTE(0) GTE(1) GTE(2) - // \ / \ / - // Add Add OnesVec - // \ / \ / - // \ Add Mul OnesVec - // \ | | / - // \ Mul Add - // \ | / - // \ | / - // Tuple - // - HloComputation* BuildComputation0() { - auto builder = HloComputation::Builder(TestName() + ".Computation0"); - // Create param instruction to access computation state. - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape3_, "param")); - - // Create GetTupleElement instructions for each tuple element. - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, param, 1)); - auto gte2 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, param, 2)); - - // Create const vector of ones to be used in element-wise computations. - auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( - Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); - - // Create simple fusable computation for tuple element 0 (wont get merged). - auto out0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, one_vec, gte0)); - - // Create fusable computation which is dependent on second and third tuple - // elements (will initially be fused on its own). - auto add1 = builder.AddInstruction( - HloInstruction::CreateBinary(data_shape_, HloOpcode::kAdd, gte1, gte2)); - - // Create two sub-computations, both of which are users of 'add1'. - - // First sub-computation: out1 = Mul(Add(add1, one_vec), one_vec) - auto add2 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, add1, one_vec)); - auto out1 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, add2, one_vec)); - - // Second sub-computation: out2 = Add(Mul(add1, one_vec), one_vec) - auto mul0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, add1, one_vec)); - auto out2 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, mul0, one_vec)); - - // Create output Tuple. - builder.AddInstruction(HloInstruction::CreateTuple({out0, out1, out2})); - return module_->AddEntryComputation(builder.Build()); - } - - // Builds the following computation: - // - // Param - // / \ - // GTE(0) GTE(1) - // | | \ / - // | | Mul - // \ \ | - // \ Mul - // \ | - // OnesVec Mul OnesVec - // \ / \ / - // OnesVec Add Mul OnesVec - // \ | | / - // Mul Add - // \ / - // \ / - // Tuple - // - HloComputation* BuildComputation1() { - auto builder = HloComputation::Builder(TestName() + ".Computation1"); - Shape tuple_shape2_ = ShapeUtil::MakeTupleShape({data_shape_, data_shape_}); - // Create param instruction to access computation state. - auto state = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape2_, "state")); - - // Create shared sub-computation (will initially be fused on its own). - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, state, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, state, 2)); - // Calculate the flops we need to generate for this shared computation - // to exceed the threshold flops_to_bytes_ratio. - // Note that bytes transferred is multiplied by 3 because there are two - // operands and one output of size 'data_shape_'. - const int64 flops_needed = FusionMerger::GetThresholdFlopsToBytesRatio() * - ShapeUtil::ByteSizeOf(data_shape_) * 3; - const int64 vec_elements = ShapeUtil::ElementsIn(data_shape_); - const int64 iters = (flops_needed + vec_elements - 1) / vec_elements; - - auto mul0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, gte0, gte1)); - for (int i = 0; i < iters; ++i) { - mul0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, gte0, mul0)); - } - - // Create two sub-computations, both of which are users of 'mul0'. - auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( - Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); - - // First sub-computation: out0 = Mul(Add(mul0, one_vec), one_vec) - auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, mul0, one_vec)); - auto out0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, add0, one_vec)); - - // Second sub-computation: out1 = Add(Mul(mul0, one_vec), one_vec) - auto mul1 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, mul0, one_vec)); - auto out1 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, mul1, one_vec)); - - // Create output Tuple. - builder.AddInstruction(HloInstruction::CreateTuple({out0, out1})); - return module_->AddEntryComputation(builder.Build()); - } - - // Builds the following computation: - // - // Param - // / | | \ - // / | | \ - // / | | \ - // GTE(0) GTE(1) GTE(2) GTE(3) - // \ / / / - // Add / / - // \ / / - // Add / - // \ / - // \ / - // OnesVec Add OnesVec - // \ / \ / - // OnesVec Add Mul OnesVec - // \ | | / - // Mul Add - // \ / - // \ / - // Tuple - // - HloComputation* BuildComputation2(bool add_extra_input) { - auto builder = HloComputation::Builder(TestName() + ".Computation2"); - Shape state_shape = add_extra_input ? tuple_shape4_ : tuple_shape3_; - // Create param instruction to access computation state. - auto state = builder.AddInstruction( - HloInstruction::CreateParameter(0, state_shape, "state")); - - // Create GetTupleElement instructions for each tuple element. - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, state, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, state, 1)); - auto gte2 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, state, 2)); - - // Create shared fusable computation that reduces its operands. - auto reduce0 = builder.AddInstruction( - HloInstruction::CreateBinary(data_shape_, HloOpcode::kAdd, gte0, gte1)); - auto reduce_out = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, reduce0, gte2)); - if (add_extra_input) { - auto gte3 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, state, 3)); - reduce_out = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, reduce_out, gte3)); - } +namespace op = xla::testing::opcode_matchers; - // Create two fusable sub-computations which are dependent on shared - // computation 'reduce_out'. - auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( - Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); - - // First sub-computation: out0 = Mul(Add(reduce_out, one_vec), one_vec) - auto add2 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, reduce_out, one_vec)); - auto out0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, add2, one_vec)); - - // Second sub-computation: out1 = Add(Mul(reduce_out, one_vec), one_vec) - auto mul0 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kMultiply, reduce_out, one_vec)); - auto out1 = builder.AddInstruction(HloInstruction::CreateBinary( - data_shape_, HloOpcode::kAdd, mul0, one_vec)); - - // Create output Tuple. - builder.AddInstruction(HloInstruction::CreateTuple({out0, out1})); - return module_->AddEntryComputation(builder.Build()); - } - - Shape data_shape_ = ShapeUtil::MakeShape(F32, {4}); - Shape tuple_shape2_ = ShapeUtil::MakeTupleShape({data_shape_, data_shape_}); - Shape tuple_shape3_ = - ShapeUtil::MakeTupleShape({data_shape_, data_shape_, data_shape_}); - Shape tuple_shape4_ = ShapeUtil::MakeTupleShape( - {data_shape_, data_shape_, data_shape_, data_shape_}); - - std::unique_ptr module_; -}; +class FusionMergerTest : public HloTestBase {}; // Tests that we can merge a fusion instruction that is below threshold. // -// Original computation: -// -// Param -// / | \ -// / | \ -// OnesVec GTE(0) GTE(1) GTE(2) -// \ / \ / -// Add Add OnesVec -// \ / \ / -// \ Add Mul OnesVec -// \ | | / -// \ Mul Add -// \ | / -// \ | / -// Tuple -// -// Computation after fusion passes: -// -// Param -// / \ -// Fusion3 Fusion2 -// | / \ -// \ Fusion0 Fusion1 -// \ | / -// \ | / -// Tuple -// // Computation after fusion merger pass (Fusion2 is merged into Fusion0 and // Fusion1): // Param @@ -276,19 +40,50 @@ class FusionMergerTest : public HloTestBase { // Tuple // TEST_F(FusionMergerTest, MergeSharedFusionInstruction) { - auto computation = BuildComputation0(); - // Run standard fusion passes. - EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) - .Run(module_.get()) - .ValueOrDie()); - EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module_.get()) - .ValueOrDie()); - // Run fusion merger pass, which should merge the shared fusion instruction - // into its two users. - EXPECT_TRUE(FusionMerger().Run(module_.get()).ValueOrDie()); - - auto* root = computation->root_instruction(); + auto module = tools::Parse(R"( +HloModule MergeSharedFusionInstruction + +comp.3 { + constant.param_0 = f32[4]{0} parameter(0) + param.param_1.2 = (f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(1) + get-tuple-element.6 = f32[4]{0} get-tuple-element(param.param_1.2), index=0 + ROOT add.7 = f32[4]{0} add(constant.param_0, get-tuple-element.6) +} + +comp.2 { + param.param_1.1 = (f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(0) + get-tuple-element.4 = f32[4]{0} get-tuple-element(param.param_1.1), index=1 + get-tuple-element.5 = f32[4]{0} get-tuple-element(param.param_1.1), index=2 + ROOT add.6 = f32[4]{0} add(get-tuple-element.4, get-tuple-element.5) +} + +comp.1 { + add.1.param_1.1 = f32[4]{0} parameter(1) + constant.param_1.3 = f32[4]{0} parameter(0) + add.5 = f32[4]{0} add(add.1.param_1.1, constant.param_1.3) + ROOT multiply.3 = f32[4]{0} multiply(add.5, constant.param_1.3) +} + +comp { + add.1.param_1 = f32[4]{0} parameter(1) + constant.param_1.1 = f32[4]{0} parameter(0) + multiply.2 = f32[4]{0} multiply(add.1.param_1, constant.param_1.1) + ROOT add.4 = f32[4]{0} add(multiply.2, constant.param_1.1) +} + +ENTRY MergeSharedFusionInstruction.Computation0 { + constant = f32[4]{0} constant({1, 1, 1, 1}) + param = (f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(0) + fusion.3 = f32[4]{0} fusion(constant, param), kind=kLoop, calls=comp.3 + fusion.4 = f32[4]{0} fusion(param), kind=kLoop, calls=comp.2 + fusion.5 = f32[4]{0} fusion(constant, fusion.4), kind=kLoop, calls=comp.1 + fusion.6 = f32[4]{0} fusion(constant, fusion.4), kind=kLoop, calls=comp + ROOT tuple = (f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(fusion.3, fusion.5, fusion.6) +})") + .ValueOrDie(); + EXPECT_TRUE(FusionMerger().Run(module.get()).ValueOrDie()); + + auto* root = module->entry_computation()->root_instruction(); EXPECT_EQ(HloOpcode::kTuple, root->opcode()); // Check operand 0 (not merged). Should have 4 instructions. auto* operand0 = root->operand(0); @@ -307,156 +102,188 @@ TEST_F(FusionMergerTest, MergeSharedFusionInstruction) { // Tests that we do not merge a fusion instruction that above flops to bytes // threshold. // -// Original computation: -// -// Param -// / \ -// GTE(0) GTE(1) -// | | \ / -// | | Mul -// \ \ | -// \ Mul -// \ | -// OnesVec Mul OnesVec -// \ / \ / -// OnesVec Add Mul OnesVec -// \ | | / -// Mul Add -// \ / -// \ / -// Tuple -// -// Computation after fusion passes and fusion merger pass (Fusion2 is not -// merged because it exceeds the threshold flops to bytes ratio). -// -// Param -// | -// Fusion2 -// / \ -// Fusion0 Fusion1 -// \ / -// Tuple -// +// Fusion2 is not merged because it exceeds the threshold flops-to-bytes ratio. TEST_F(FusionMergerTest, FlopsToBytesRatioThresholdExceeded) { - BuildComputation1(); - // Run standard fusion passes. - EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) - .Run(module_.get()) - .ValueOrDie()); - EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module_.get()) - .ValueOrDie()); + auto module = tools::Parse(R"( +HloModule FlopsToBytesRatioThresholdExceeded + +comp.2 { + state.param_1.1 = (f32[4]{0}, f32[4]{0}) parameter(0) + get-tuple-element.3 = f32[4]{0} get-tuple-element(state.param_1.1), index=0 + get-tuple-element.4 = f32[4]{0} get-tuple-element(state.param_1.1), index=2 + multiply.29 = f32[4]{0} multiply(get-tuple-element.3, get-tuple-element.4) + multiply.30 = f32[4]{0} multiply(get-tuple-element.3, multiply.29) + multiply.31 = f32[4]{0} multiply(get-tuple-element.3, multiply.30) + multiply.32 = f32[4]{0} multiply(get-tuple-element.3, multiply.31) + multiply.33 = f32[4]{0} multiply(get-tuple-element.3, multiply.32) + multiply.34 = f32[4]{0} multiply(get-tuple-element.3, multiply.33) + multiply.35 = f32[4]{0} multiply(get-tuple-element.3, multiply.34) + multiply.36 = f32[4]{0} multiply(get-tuple-element.3, multiply.35) + multiply.37 = f32[4]{0} multiply(get-tuple-element.3, multiply.36) + multiply.38 = f32[4]{0} multiply(get-tuple-element.3, multiply.37) + multiply.39 = f32[4]{0} multiply(get-tuple-element.3, multiply.38) + multiply.40 = f32[4]{0} multiply(get-tuple-element.3, multiply.39) + ROOT multiply.41 = f32[4]{0} multiply(get-tuple-element.3, multiply.40) +} + +comp.1 { + multiply.12.param_1.1 = f32[4]{0} parameter(1) + constant.param_1.3 = f32[4]{0} parameter(0) + add.3 = f32[4]{0} add(multiply.12.param_1.1, constant.param_1.3) + ROOT multiply.16 = f32[4]{0} multiply(add.3, constant.param_1.3) +} + +comp { + multiply.12.param_1 = f32[4]{0} parameter(1) + constant.param_1.1 = f32[4]{0} parameter(0) + multiply.15 = f32[4]{0} multiply(multiply.12.param_1, constant.param_1.1) + ROOT add.2 = f32[4]{0} add(multiply.15, constant.param_1.1) +} + +ENTRY FlopsToBytesRatioThresholdExceeded.Computation1 { + constant = f32[4]{0} constant({1, 1, 1, 1}) + state = (f32[4]{0}, f32[4]{0}) parameter(0) + fusion.2 = f32[4]{0} fusion(state), kind=kLoop, calls=comp.2 + fusion.3 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp.1 + fusion.4 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp + ROOT tuple = (f32[4]{0}, f32[4]{0}) tuple(fusion.3, fusion.4) +})") + .ValueOrDie(); // Run fusion merger pass, which should detect that the flops/bytes of the // shared fusion instruction exceeds the threshold ratio, and therefore // cannot be merged with other fusion instructions. - EXPECT_FALSE(FusionMerger().Run(module_.get()).ValueOrDie()); + EXPECT_FALSE(FusionMerger().Run(module.get()).ValueOrDie()); } // Tests that threshold for bytes transferred if merged is exceeded. // -// Original computation: -// -// Param -// / | | \ -// / | | \ -// / | | \ -// GTE(0) GTE(1) GTE(2) GTE(3) -// \ / / / -// Add / / -// \ / / -// Add / -// \ / -// \ / -// OnesVec Add OnesVec -// \ / \ / -// OnesVec Add Mul OnesVec -// \ | | / -// Mul Add -// \ / -// \ / -// Tuple -// -// Computation after fusion passes and fusion merger pass. Fusion2 is not -// merged because it exceeds the threshold bytes transferred. This is because -// the bytes read by Fusion2 (when replicated if the instruction is merged -// into Fusion0 and Fusion1) would exceed the bytes transferred threshold. -// -// Param -// | -// Fusion2 -// / \ -// Fusion0 Fusion1 -// \ / -// Tuple -// +// Fusion2 is not merged because it exceeds the threshold bytes transferred. +// This is because the bytes read by Fusion2 (when replicated if the instruction +// is merged into Fusion0 and Fusion1) would exceed the bytes transferred +// threshold. TEST_F(FusionMergerTest, BytesTransferredThresholdExeceeded) { - BuildComputation2(/*add_extra_input=*/true); - // Run standard fusion passes. - EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) - .Run(module_.get()) - .ValueOrDie()); - EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module_.get()) - .ValueOrDie()); + auto module = tools::Parse(R"( +HloModule BytesTransferredThresholdExeceeded + +comp.2 { + state.param_1.1 = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(0) + get-tuple-element.7 = f32[4]{0} get-tuple-element(state.param_1.1), index=0 + get-tuple-element.8 = f32[4]{0} get-tuple-element(state.param_1.1), index=1 + add.9 = f32[4]{0} add(get-tuple-element.7, get-tuple-element.8) + get-tuple-element.9 = f32[4]{0} get-tuple-element(state.param_1.1), index=2 + add.10 = f32[4]{0} add(add.9, get-tuple-element.9) + get-tuple-element.10 = f32[4]{0} get-tuple-element(state.param_1.1), index=3 + ROOT add.11 = f32[4]{0} add(add.10, get-tuple-element.10) +} + +comp.1 { + add.2.param_1.1 = f32[4]{0} parameter(1) + constant.param_1.3 = f32[4]{0} parameter(0) + add.6 = f32[4]{0} add(add.2.param_1.1, constant.param_1.3) + ROOT multiply.3 = f32[4]{0} multiply(add.6, constant.param_1.3) +} + +comp { + add.2.param_1 = f32[4]{0} parameter(1) + constant.param_1.1 = f32[4]{0} parameter(0) + multiply.2 = f32[4]{0} multiply(add.2.param_1, constant.param_1.1) + ROOT add.5 = f32[4]{0} add(multiply.2, constant.param_1.1) +} + +ENTRY BytesTransferredThresholdExeceeded.Computation2 { + constant = f32[4]{0} constant({1, 1, 1, 1}) + state = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(0) + fusion.2 = f32[4]{0} fusion(state), kind=kLoop, calls=comp.2 + fusion.3 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp.1 + fusion.4 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp + ROOT tuple = (f32[4]{0}, f32[4]{0}) tuple(fusion.3, fusion.4) +})") + .ValueOrDie(); // Run fusion merger pass, which should detect that the net bytes transferred // (if merged) would increase. - EXPECT_FALSE(FusionMerger().Run(module_.get()).ValueOrDie()); + EXPECT_FALSE(FusionMerger().Run(module.get()).ValueOrDie()); } // Tests that threshold for bytes transferred if merged is not exceeded. // -// Original computation: -// -// Param -// / | \ -// / | \ -// / | \ -// GTE(0) GTE(1) GTE(2) -// \ / / -// Add / -// \ / -// OnesVec Add OnesVec -// \ / \ / -// OnesVec Add Mul OnesVec -// \ / \ / -// Mul Add -// \ / -// \ / -// Tuple -// -// Computation after fusion passes: -// -// Param -// | -// Fusion2 -// / \ -// Fusion0 Fusion1 -// \ / -// Tuple -// -// Computation after fusion merger pass (Fusion2 is merged into Fusion0 and -// Fusion1, because bytes read from Param by Fusion2 is reduced for this test -// which makes the merge operation into its operand below the bytes -// transferred threshold. -// -// Param -// / \ -// Fusion0 Fusion1 -// \ / -// Tuple -// +// Fusion2 is merged into Fusion0 and Fusion1, because bytes read from Param by +// Fusion2 is reduced for this test which makes the merge operation into its +// operand below the bytes transferred threshold. TEST_F(FusionMergerTest, BytesTransferredThresholdNotExeceeded) { - BuildComputation2(/*add_extra_input=*/false); - // Run standard fusion passes. - EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) - .Run(module_.get()) - .ValueOrDie()); - EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module_.get()) - .ValueOrDie()); + auto module = tools::Parse(R"( +HloModule BytesTransferredThresholdNotExeceeded + +comp.2 { + state.param_1.1 = (f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(0) + get-tuple-element.5 = f32[4]{0} get-tuple-element(state.param_1.1), index=0 + get-tuple-element.6 = f32[4]{0} get-tuple-element(state.param_1.1), index=1 + add.7 = f32[4]{0} add(get-tuple-element.5, get-tuple-element.6) + get-tuple-element.7 = f32[4]{0} get-tuple-element(state.param_1.1), index=2 + ROOT add.8 = f32[4]{0} add(add.7, get-tuple-element.7) +} + +comp.1 { + add.1.param_1.1 = f32[4]{0} parameter(1) + constant.param_1.3 = f32[4]{0} parameter(0) + add.5 = f32[4]{0} add(add.1.param_1.1, constant.param_1.3) + ROOT multiply.3 = f32[4]{0} multiply(add.5, constant.param_1.3) +} + +comp { + add.1.param_1 = f32[4]{0} parameter(1) + constant.param_1.1 = f32[4]{0} parameter(0) + multiply.2 = f32[4]{0} multiply(add.1.param_1, constant.param_1.1) + ROOT add.4 = f32[4]{0} add(multiply.2, constant.param_1.1) +} + +ENTRY BytesTransferredThresholdNotExeceeded.Computation2 { + constant = f32[4]{0} constant({1, 1, 1, 1}) + state = (f32[4]{0}, f32[4]{0}, f32[4]{0}) parameter(0) + fusion.2 = f32[4]{0} fusion(state), kind=kLoop, calls=comp.2 + fusion.3 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp.1 + fusion.4 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp + ROOT tuple = (f32[4]{0}, f32[4]{0}) tuple(fusion.3, fusion.4) +})") + .ValueOrDie(); // Run fusion merger pass, which should detect that the net bytes transferred // (if merged) would not increase. - EXPECT_TRUE(FusionMerger().Run(module_.get()).ValueOrDie()); + EXPECT_TRUE(FusionMerger().Run(module.get()).ValueOrDie()); +} + +// Check that we're willing to merge f1_computation into f2_computation, even +// though f2 is an input fusion node. +TEST_F(FusionMergerTest, WillMergeIntoInputFusion) { + auto module = tools::Parse(R"( + HloModule m + + f1_computation { + f1_p0 = f32[10]{0} parameter(0) + ROOT f1_root = f32[10]{0} add(f1_p0, f1_p0) + } + + add_computation { + add_lhs = f32[] parameter(0) + add_rhs = f32[] parameter(1) + ROOT add_root = f32[] add(add_lhs, add_rhs) + } + + f2_computation { + f2_p0 = f32[10]{0} parameter(0) + f2_mul = f32[10]{0} multiply(f2_p0, f2_p0) + f2_zero = f32[] constant(0) + ROOT f2_root = f32[] reduce(f2_mul, f2_zero), dimensions={0}, + to_apply=add_computation + } + + ENTRY entry { + p0 = f32[10]{0} parameter(0) + f1 = f32[10]{0} fusion(p0), kind=kLoop, calls=f1_computation + ROOT f2 = f32[] fusion(f1), kind=kInput, calls=f2_computation + })") + .ValueOrDie(); + EXPECT_TRUE(FusionMerger().Run(module.get()).ValueOrDie()); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Fusion(op::Parameter())); } } // namespace diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc index 8e3aebbc12b5e6d746700956b9743bc94db50167..38668ff455a44c7ef99b57b750f1a3b18a90bd2c 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc @@ -49,7 +49,7 @@ struct MatrixDescriptor { // rhs_matrix, and stores the result to output_matrix. template bool DoGemm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, - MatrixDescriptor output_matrix, se::Stream* stream) { + MatrixDescriptor output_matrix, double alpha, se::Stream* stream) { DCHECK(!output_matrix.transpose); se::DeviceMemory lhs_data(lhs_matrix.data); @@ -65,7 +65,7 @@ bool DoGemm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, return stream ->ThenBlasGemm( lhs_transpose, rhs_transpose, output_matrix.num_rows, - output_matrix.num_cols, /*size of reduce dim=*/k, /*alpha=*/1.0, + output_matrix.num_cols, /*size of reduce dim=*/k, /*alpha=*/alpha, lhs_data, /*leading dim of LHS=*/lhs_matrix.num_rows, rhs_data, /*leading dim of RHS=*/rhs_matrix.num_rows, /*beta=*/0.0, &output_data, /*leading dim of output=*/output_matrix.num_rows) @@ -89,7 +89,7 @@ bool DoGemm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, template bool DoGemmWithAlgorithm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, - MatrixDescriptor output_matrix, + MatrixDescriptor output_matrix, double alpha, se::blas::ComputationType computation_type, se::blas::AlgorithmType algorithm, se::Stream* stream, se::blas::ProfileResult* output_profile_result) { @@ -108,11 +108,13 @@ bool DoGemmWithAlgorithm(MatrixDescriptor lhs_matrix, return stream ->ThenBlasGemmWithAlgorithm( lhs_transpose, rhs_transpose, output_matrix.num_rows, - output_matrix.num_cols, /*size of reduce dim=*/k, /*alpha=*/1.0, - lhs_data, /*leading dim of LHS=*/lhs_matrix.num_rows, rhs_data, - /*leading dim of RHS=*/rhs_matrix.num_rows, /*beta=*/0.0, - &output_data, /*leading dim of output=*/output_matrix.num_rows, - computation_type, algorithm, output_profile_result) + output_matrix.num_cols, /*size of reduce dim=*/k, + /*alpha=*/static_cast(alpha), lhs_data, + /*leading dim of LHS=*/lhs_matrix.num_rows, rhs_data, + /*leading dim of RHS=*/rhs_matrix.num_rows, + /*beta=*/static_cast(0.0f), &output_data, + /*leading dim of output=*/output_matrix.num_rows, computation_type, + algorithm, output_profile_result) .ok(); } @@ -125,8 +127,8 @@ bool DoGemmWithAlgorithm(MatrixDescriptor lhs_matrix, template StatusOr DoGemmAutotune( MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, - MatrixDescriptor output_matrix, se::blas::ComputationType computation_type, - se::Stream* stream) { + MatrixDescriptor output_matrix, double alpha, + se::blas::ComputationType computation_type, se::Stream* stream) { std::vector algorithms; CHECK(stream->parent()->GetBlasGemmAlgorithms(&algorithms)); @@ -137,9 +139,9 @@ StatusOr DoGemmAutotune( // for all algorithms if we're targeting < sm_50. But because we pass a // non-null ProfileResult, DoGemmWithAlgorithm should always return true, // and the actual success-ness is returned in ProfileResult::is_valid. - DCHECK(DoGemmWithAlgorithm(lhs_matrix, rhs_matrix, output_matrix, - computation_type, algorithm, stream, - &profile_result)); + CHECK(DoGemmWithAlgorithm(lhs_matrix, rhs_matrix, output_matrix, + alpha, computation_type, algorithm, + stream, &profile_result)); if (profile_result.is_valid() && profile_result.elapsed_time_in_ms() < best_result.elapsed_time_in_ms()) { @@ -161,6 +163,8 @@ StatusOr DoGemmAutotune( // DoGemm/DoGemmWithAlgorithm/DoGemmAutotune. auto GetGemmFn(PrimitiveType type) -> decltype(&DoGemm) { switch (type) { + case F16: + return &DoGemm; case F32: return &DoGemm; case F64: @@ -172,6 +176,8 @@ auto GetGemmFn(PrimitiveType type) -> decltype(&DoGemm) { auto GetGemmWithAlgorithmFn(PrimitiveType type) -> decltype(&DoGemmWithAlgorithm) { switch (type) { + case F16: + return &DoGemmWithAlgorithm; case F32: return &DoGemmWithAlgorithm; case F64: @@ -182,6 +188,8 @@ auto GetGemmWithAlgorithmFn(PrimitiveType type) } auto GetGemmAutotuneFn(PrimitiveType type) -> decltype(&DoGemmAutotune) { switch (type) { + case F16: + return &DoGemmAutotune; case F32: return &DoGemmAutotune; case F64: @@ -196,6 +204,10 @@ auto GetGemmAutotuneFn(PrimitiveType type) -> decltype(&DoGemmAutotune) { // separately from the precision of the inputs and result. se::blas::ComputationType GetBlasComputationType(PrimitiveType type) { switch (type) { + case F16: + // Use F32 as computation type for F16 as we currently only implement the + // cuDNN pseudo half configuration for half precision. + return se::blas::ComputationType::kF32; case F32: return se::blas::ComputationType::kF32; case F64: @@ -212,7 +224,8 @@ GemmThunk::GemmThunk(const BufferAllocation::Slice& lhs_buffer, const BufferAllocation::Slice& output_buffer, const Shape& lhs_shape, const Shape& rhs_shape, const Shape& output_shape, bool transpose_lhs, - bool transpose_rhs, const HloInstruction* hlo_instruction) + bool transpose_rhs, double alpha, + const HloInstruction* hlo_instruction) : Thunk(Kind::kGemm, hlo_instruction), lhs_buffer_(lhs_buffer), rhs_buffer_(rhs_buffer), @@ -221,7 +234,8 @@ GemmThunk::GemmThunk(const BufferAllocation::Slice& lhs_buffer, rhs_shape_(rhs_shape), output_shape_(output_shape), transpose_lhs_(transpose_lhs), - transpose_rhs_(transpose_rhs) {} + transpose_rhs_(transpose_rhs), + alpha_(alpha) {} tensorflow::Status GemmThunk::ExecuteOnStream( const BufferAllocations& buffer_allocations, se::Stream* stream) { @@ -290,7 +304,7 @@ tensorflow::Status GemmThunk::ExecuteOnStream( if (autotune_it == autotune_results_.end()) { StatusOr best_algorithm = GetGemmAutotuneFn(element_type)(lhs_matrix, rhs_matrix, output_matrix, - computation_type, stream); + alpha_, computation_type, stream); autotune_it = autotune_results_.insert({device_name, best_algorithm}).first; @@ -311,12 +325,15 @@ tensorflow::Status GemmThunk::ExecuteOnStream( VLOG(2) << "Using algorithm " << algorithm << " chosen by autotuning on GemmThunk " << this; return GetGemmWithAlgorithmFn(element_type)( - lhs_matrix, rhs_matrix, output_matrix, computation_type, algorithm, - stream, + lhs_matrix, rhs_matrix, output_matrix, alpha_, computation_type, + algorithm, stream, /*output_profile_result=*/nullptr); } + + // Autotune will fail when CUDA 8 and GPU sm_50 or older are used. + // Use the older Gemm API in this case. return GetGemmFn(element_type)(lhs_matrix, rhs_matrix, output_matrix, - stream); + alpha_, stream); }; bool launch_ok; diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h index 8c6a1f51a8a09ef78950dfe7e89994a3fe247f49..df3edcefef898d465cd5ddc53e5d06a966a31f88 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h @@ -34,15 +34,16 @@ namespace gpu { // This is thread-compatible. class GemmThunk : public Thunk { public: - // Constructs a thunk that computes "output = lhs rhs" using BLAS gemm. - // transpose_lhs and transpose_rhs indicate whether gemm should transpose the - // lhs and rhs operand. hlo_instruction is as in Thunk. + // Constructs a thunk that computes "output = (lhs rhs) * alpha" using + // BLAS gemm. transpose_lhs and transpose_rhs indicate whether gemm should + // transpose the lhs and rhs operand. hlo_instruction is as in Thunk. alpha is + // a constant. GemmThunk(const BufferAllocation::Slice& lhs_buffer, const BufferAllocation::Slice& rhs_buffer, const BufferAllocation::Slice& output_buffer, const Shape& lhs_shape, const Shape& rhs_shape, const Shape& output_shape, bool transpose_lhs, bool transpose_rhs, - const HloInstruction* hlo_instruction); + double alpha, const HloInstruction* hlo_instruction); GemmThunk(const GemmThunk&) = delete; GemmThunk& operator=(const GemmThunk&) = delete; @@ -72,6 +73,7 @@ class GemmThunk : public Thunk { const bool transpose_lhs_; const bool transpose_rhs_; + const double alpha_; // Maps device names (StreamExecutor::DeviceDescription::name()) to autotune // results. The map's value is the best algorithm we've found for this thunk diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc index 07543d42e3e47c2acb6bf908399bb88c3ae353b3..07be2a0cf90c326af6e41764e79950db546e43e4 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc @@ -33,10 +33,13 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/call_inliner.h" +#include "tensorflow/compiler/xla/service/conditional_simplifier.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/gather_expander.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" @@ -46,8 +49,8 @@ limitations under the License. #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" @@ -127,7 +130,9 @@ string GetLibdeviceDir(const string& config_cuda_data_dir) { } // Runs optimization passes on the given HLO module. -tensorflow::Status OptimizeHloModule(HloModule* hlo_module) { +tensorflow::Status OptimizeHloModule(HloModule* hlo_module, + se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) { { HloPassPipeline pipeline("optimization"); pipeline.AddInvariantChecker(); @@ -143,6 +148,7 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module) { // most ops. pipeline.AddPass(BF16, F32); pipeline.AddPass(); + { auto& pass = pipeline.AddPass>("simplification"); @@ -160,6 +166,9 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module) { /*rewrite_grad_op=*/true, /*use_fusion=*/false); + // Rewrite gather ops into smaller ones. + pass.AddPass(); + // BatchNormExpander can create zero-sized ops, so zero-sized HLO // elimination has to come after that pass. pipeline.AddPass(); @@ -172,8 +181,9 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module) { pass.AddPass(); pass.AddPass(); pass.AddPass(); + pass.AddPass(); } - pipeline.AddPass(); + pipeline.AddPass( [](const HloInstruction& dot, const TransposeFolding::OperandIndices& candidate_operands) { @@ -185,6 +195,74 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module) { 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()); + } + + { + HloPassPipeline pipeline("layout_assignment"); + 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>( + /*is_layout_sensitive=*/true, + /*valid_bitcast_callback=*/[](const Shape&, const Shape&) { + return true; + }); + pipeline.AddPass(/*is_layout_sensitive=*/true); + TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status()); + } + { HloPassFix fusion("fusion"); fusion.AddInvariantChecker(); @@ -212,9 +290,7 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module) { // 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, se::StreamExecutor* stream_exec, - DeviceMemoryAllocator* /*device_allocator*/) { +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 @@ -222,15 +298,7 @@ tensorflow::Status PrepareHloModuleForIrEmitting( // the parameter. HloPassPipeline pipeline("GPU-ir-emit-prepare"); pipeline.AddInvariantChecker(); - pipeline.AddPass(); - 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>( - /*is_layout_sensitive=*/true, - [](const Shape&, const Shape&) { return true; }); - pipeline.AddPass(/*is_layout_sensitive=*/true); + // Copy insertion should be performed immediately before IR emission to avoid // inserting unnecessary copies (later pass adds an instruction which // materializes the value) or missing a necessary copy (later pass removes an @@ -417,7 +485,8 @@ StatusOr> GpuCompiler::RunHloPasses( XLA_SCOPED_LOGGING_TIMER("GpuCompiler::RunHloPasses"); Tracing::TraceMe annotation("HLO Transforms", module->name(), /*is_expensive=*/true); - TF_RETURN_IF_ERROR(OptimizeHloModule(module.get())); + TF_RETURN_IF_ERROR( + OptimizeHloModule(module.get(), stream_exec, device_allocator)); return std::move(module); } @@ -428,8 +497,7 @@ StatusOr> GpuCompiler::RunBackend( TF_RET_CHECK(stream_exec != nullptr); - TF_RETURN_IF_ERROR(PrepareHloModuleForIrEmitting(module.get(), stream_exec, - device_allocator)); + TF_RETURN_IF_ERROR(PrepareHloModuleForIrEmitting(module.get())); llvm::LLVMContext llvm_context; std::string buffer; @@ -464,8 +532,9 @@ StatusOr> GpuCompiler::RunBackend( /*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_optimized_hlo_proto_to = @@ -602,6 +671,8 @@ StatusOr> GpuCompiler::RunBackend( if (module->config().hlo_profiling_enabled()) { HloCostAnalysis cost_analysis(ShapeSizeBytesFunction()); + cost_analysis.set_bytes_per_second( + stream_exec->GetDeviceDescription().memory_bandwidth()); TF_RETURN_IF_ERROR(module->entry_computation()->Accept(&cost_analysis)); profile_index_map = MakeUnique(*module); profile_printer = diff --git a/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc b/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc index e3b493c6630d061c00dc6c67bdaecdb2e5d68533..9db85bc788bde46c890a46ce9b0902ddce3f5675 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc @@ -49,7 +49,7 @@ StatusOr GpuCopyInsertion::Run(HloModule* module) { TF_ASSIGN_OR_RETURN(bool changed, generic_copy_insertion.Run(module)); TF_ASSIGN_OR_RETURN(std::unique_ptr dataflow, - HloDataflowAnalysis::Run(module)); + HloDataflowAnalysis::Run(*module)); // Make sure all operands of a library call are in memory instead of constants // in IR. @@ -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) { diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index f5d67b9ea9498df3f023ea9a694a63b468c5be18..04b37d913e0bc8f8226057f107da05fd1e675010 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -46,12 +46,14 @@ namespace { class HloExecutionProfiler { public: // If profiling is enabled, start an execution timer running. - explicit HloExecutionProfiler(bool do_profile, HloExecutionProfile* profile, - se::Stream* stream, - const HloComputation* computation) + explicit HloExecutionProfiler( + bool do_profile, HloExecutionProfile* profile, se::Stream* stream, + const std::vector::SmartPtr>& sub_streams, + const HloComputation* computation) : do_profile_(do_profile), profile_(profile), stream_(stream), + sub_streams_(sub_streams), computation_(computation) { if (do_profile_) { clock_rate_ghz_ = @@ -70,6 +72,7 @@ class HloExecutionProfiler { CHECK(!finished_execution_) << "Call FinishExecution only once!"; finished_execution_ = true; if (do_profile_) { + stream_->ThenWaitFor(&sub_streams_); stream_->ThenStopTimer(execution_timer_.get()); stream_->BlockHostUntilDone().IgnoreError(); profile_->set_total_cycles_executed( @@ -88,6 +91,7 @@ class HloExecutionProfiler { // that the hlo_instruction took to execute in the profile. void FinishOperation(const HloInstruction* hlo_instruction) { if (do_profile_) { + stream_->ThenWaitFor(&sub_streams_); stream_->ThenStopTimer(per_op_timer_.get()); stream_->BlockHostUntilDone().IgnoreError(); profile_->SetCyclesTakenBy( @@ -100,6 +104,7 @@ class HloExecutionProfiler { double clock_rate_ghz_; HloExecutionProfile* profile_; se::Stream* stream_; + const std::vector::SmartPtr>& sub_streams_; const HloComputation* computation_; std::unique_ptr execution_timer_; std::unique_ptr per_op_timer_; @@ -147,13 +152,9 @@ Status GpuExecutable::ExecuteThunks( 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; + sub_streams.reserve(thunk_schedule_->StreamCount() - 1); while (sub_streams.size() + 1 < thunk_schedule_->StreamCount()) { sub_streams.emplace_back(); TF_ASSIGN_OR_RETURN( @@ -161,6 +162,10 @@ Status GpuExecutable::ExecuteThunks( run_options->BorrowStream(main_stream->parent()->device_ordinal())); } + HloExecutionProfiler profiler(do_profile, hlo_execution_profile, main_stream, + sub_streams, hlo_module_->entry_computation()); + uint64 start_micros = tensorflow::Env::Default()->NowMicros(); + // The next event enqueued on stream N must not run until the thunk at // last_blocking_thunk_for_stream[N] completes. std::map last_blocking_thunk_for_stream; @@ -262,9 +267,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_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/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.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc index b5962f069bf499c913bd5479f263a7cb77c00555..85ecbe8fdb34700ca738b99ddd9ea615afc35da3 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc @@ -25,13 +25,19 @@ namespace gpu { namespace { bool IsFusile(const HloInstruction& hlo) { + // Don't fuse get-tuple-element on GPU: We can, but it's slower than not + // fusing. We never generate kernels for unfused GTEs. Instead, if an + // unfused GTE is an input to a kernel (including a fusion kernel), we + // compute the address of the GTE at the top of the kernel. Often we know the + // address of the GTE result statically, so we can do this without chasing any + // pointers. return (hlo.IsElementwise() && hlo.operand_count() > 0) || + hlo.opcode() == HloOpcode::kBitcast || hlo.opcode() == HloOpcode::kBroadcast || hlo.opcode() == HloOpcode::kConcatenate || hlo.opcode() == HloOpcode::kDynamicSlice || hlo.opcode() == HloOpcode::kDynamicUpdateSlice || hlo.opcode() == HloOpcode::kFusion || - hlo.opcode() == HloOpcode::kGetTupleElement || hlo.opcode() == HloOpcode::kPad || hlo.opcode() == HloOpcode::kReduce || hlo.opcode() == HloOpcode::kReduceWindow || @@ -46,6 +52,34 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer, int64 operand_index) { HloInstruction* producer = consumer->mutable_operand(operand_index); + // Check if we can use output fusion for (A @ B) * alpha + if (producer->opcode() == HloOpcode::kDot) { + if (consumer->opcode() == HloOpcode::kMultiply) { + CHECK_EQ(consumer->operand_count(), 2); + int64 other_operand_index = 1 - operand_index; + const HloInstruction* alpha = consumer->operand(other_operand_index); + if (alpha->opcode() == HloOpcode::kConstant && + ShapeUtil::IsScalar(alpha->shape())) { + return true; + } + } + } + + // Only allow to fuse transpose into an output fusion. + if (consumer->opcode() == HloOpcode::kFusion && + consumer->fusion_kind() == HloInstruction::FusionKind::kOutput) { + if (producer->opcode() != HloOpcode::kTranspose) { + return false; + } + // Check that the transpose is the operand of a dot. + auto producer_operand_index = consumer->operand_index(producer); + auto fused_parameter = consumer->fused_parameter(producer_operand_index); + const std::vector& fused_parameter_users = + fused_parameter->users(); + return (fused_parameter_users.size() == 1 && + fused_parameter_users[0]->opcode() == HloOpcode::kDot); + } + // Output fusion is not currently supported on GPUs. if (producer->opcode() == HloOpcode::kFusion) { return false; @@ -70,17 +104,6 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer, return false; } - // We may need to know original operand layout to emit input fusion, and so - // far, we merely use the layout of an operand of the fusion node, which means - // we must fuse only elementwise operations. This restriction should be lifted - // later if we need to fuse other operations, e.g. transpose, for performance. - if ((IsReductionToVector(*consumer) || - (HloOpcode::kFusion == consumer->opcode() && - HloInstruction::FusionKind::kInput == consumer->fusion_kind())) && - !producer->IsElementwise()) { - return false; - } - // Cost condition: not fuse (simple, expensive producers) and (consumers who // reuse operand elements). if (producer->opcode() != HloOpcode::kFusion && @@ -98,6 +121,9 @@ HloInstruction::FusionKind GpuInstructionFusion::ChooseKind( if (IsReductionToVector(*consumer)) { return HloInstruction::FusionKind::kInput; } + if (producer->opcode() == HloOpcode::kDot) { + return HloInstruction::FusionKind::kOutput; + } if (HloOpcode::kFusion == consumer->opcode()) { return consumer->fusion_kind(); } diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc index 1d47ffde4331868cbc8a8afb2d01b11e77a7fab0..4b231c449f8f101127b4d30bfff20c69d8cef5c1 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" namespace op = xla::testing::opcode_matchers; @@ -137,73 +138,119 @@ 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)); +// Tests that broadcasts fused into a fusion with a reduce root. +TEST_F(InstructionFusionTest, BroadcastIntoReduce) { + auto module = tools::Parse(R"( + HloModule test_module + + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + + ENTRY BroadcastIntoReduce { + constant = f32[] constant(1) + broadcast = f32[16,16,16,16]{3,2,1,0} broadcast(constant), dimensions={} + constant.1 = f32[] constant(0) + ROOT reduce = f32[] reduce(broadcast, constant.1), dimensions={0,1,2,3}, + to_apply=add + })") + .ValueOrDie(); + + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Fusion()); + EXPECT_THAT(root->fused_expression_root(), + op::Reduce(op::Broadcast(op::Parameter()), op::Parameter())); +} + +TEST_F(InstructionFusionTest, BitcastIntoAdd) { + auto module = tools::Parse(R"( + HloModule test_module + + ENTRY BroadcastIntoAdd { + p0 = f32[4,1,1]{2,1,0} parameter(0) + p1 = f32[4,1]{1,0} parameter(1) + bitcast = f32[4,1]{1,0} bitcast(p0) + ROOT add = f32[4,1] add(bitcast, p1) + })") + .ValueOrDie(); + + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Fusion()); + EXPECT_THAT(root->fused_expression_root(), + op::Add(op::Bitcast(op::Parameter()), op::Parameter())); +} + +TEST_F(InstructionFusionTest, AddIntoBitcast) { + auto module = tools::Parse(R"( + HloModule test_module + + ENTRY BroadcastIntoAdd { + p0 = f32[4,1,1]{2,1,0} parameter(0) + p1 = f32[4,1]{1,0} parameter(1) + add = f32[4,1] add(p0, p1) + ROOT bitcast = f32[4,1,1] bitcast(add) + })") + .ValueOrDie(); + + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Fusion()); + EXPECT_THAT(root->fused_expression_root(), + op::Bitcast(op::Add(op::Parameter(), op::Parameter()))); +} + +TEST_F(InstructionFusionTest, DontFuseGTE) { + auto module = tools::Parse(R"( + HloModule test_module + ENTRY DontFuseGTE { + p0 = (f32[10], f32[10]) parameter(0) + gte0 = f32[10] get-tuple-element(p0), index=0 + gte1 = f32[10] get-tuple-element(p0), index=1 + ROOT add = f32[10] add(gte0, gte1) + })") + .ValueOrDie(); - 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}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({data_shape, data_shape}); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "param")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape, param, 1)); - builder.AddInstruction( - HloInstruction::CreateBinary(data_shape, HloOpcode::kAdd, gte0, gte1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); +TEST_F(InstructionFusionTest, DotOutputFusion) { + auto module = tools::Parse(R"( + HloModule test_module + ENTRY OutputFusion { + constant = f32[] constant(3) + p0 = f32[4,3]{1,0} parameter(0) + p1 = f32[4,3]{1,0} parameter(1) + transpose = f32[3,4]{1,0} transpose(p1), dimensions={1, 0} + dot = f32[4,4]{1,0} dot(p0, transpose) + ROOT mul = f32[4,4] multiply(constant, dot) + })") + .ValueOrDie(); + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) .Run(module.get()) .ValueOrDie()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(HloOpcode::kFusion, root->opcode()); - HloInstruction* fused_root = root->fused_expression_root(); - EXPECT_EQ(HloOpcode::kAdd, fused_root->opcode()); - // Check that operands of 'fused_root' are GTE. - EXPECT_EQ(HloOpcode::kGetTupleElement, fused_root->operand(0)->opcode()); - EXPECT_EQ(HloOpcode::kGetTupleElement, fused_root->operand(1)->opcode()); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Fusion()); + EXPECT_THAT( + root->fused_expression_root(), + op::Multiply(op::Parameter(), + op::Dot(op::Parameter(), op::Transpose(op::Parameter())))); } } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index 76566a9e3dbbc936ff90fe3f440ede14bf4e5233..32413f975a40c1abc334b16e81097bb44f56a44a 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -49,8 +49,10 @@ bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, // The inputs and the output must // 1) be matrices with no padding and a non-zero number of elements, // 2) have an allowed element type. - bool type_is_allowed = (output_shape.element_type() == F32 || - output_shape.element_type() == F64); + PrimitiveType output_primitive_type = output_shape.element_type(); + bool type_is_allowed = + (output_primitive_type == F16 || output_primitive_type == F32 || + output_primitive_type == F64); return type_is_allowed && IsRank2WithNoPadding(lhs_shape) && IsRank2WithNoPadding(rhs_shape) && IsRank2WithNoPadding(output_shape) && @@ -87,41 +89,17 @@ bool ImplementedAsGemm(const HloInstruction& hlo) { return true; } - 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 (hlo.opcode() == HloOpcode::kFusion && + hlo.fusion_kind() == HloInstruction::FusionKind::kOutput && + hlo.fused_expression_root()->opcode() == HloOpcode::kMultiply) { + // Try to find the dot inside the output fusion node. + const HloInstruction* dot = hlo.fused_expression_root()->operand(0); + if (dot->opcode() != HloOpcode::kDot) { + dot = hlo.fused_expression_root()->operand(1); } - - if (window_util::HasWindowReversal(hlo.window())) { - return false; + if (dot->opcode() == HloOpcode::kDot) { + return ImplementedAsGemm(*dot); } - - 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; @@ -144,9 +122,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 23b72c3f71dacf2be02a0719c07c7e6e88abd00c..1e0db2821a2c212d0f212ae94ab69231bc6053ea 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include "tensorflow/core/platform/logging.h" // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" @@ -27,6 +28,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" @@ -436,6 +439,32 @@ Status IrEmitter::HandleSelect(HloInstruction* select) { return IrEmitter::DefaultAction(select); } +namespace { +llvm::Value* Real(llvm::Value* x, llvm::IRBuilder<>* ir_builder) { + return ir_builder->CreateExtractValue(x, {0}); +} + +llvm::Value* Imag(llvm::Value* x, llvm::IRBuilder<>* ir_builder) { + return ir_builder->CreateExtractValue(x, {1}); +} + +std::pair MultiplyComplex( + llvm::Value* lhs_value, llvm::Value* rhs_value, + llvm::IRBuilder<>* ir_builder) { + llvm::Value* lhs_real = Real(lhs_value, ir_builder); + llvm::Value* lhs_imag = Imag(lhs_value, ir_builder); + llvm::Value* rhs_real = Real(rhs_value, ir_builder); + llvm::Value* rhs_imag = Imag(rhs_value, ir_builder); + llvm::Value* real_result1 = ir_builder->CreateFMul(lhs_real, rhs_real); + llvm::Value* real_result2 = ir_builder->CreateFMul(lhs_imag, rhs_imag); + llvm::Value* real_result = ir_builder->CreateFSub(real_result1, real_result2); + llvm::Value* imag_result1 = ir_builder->CreateFMul(lhs_real, rhs_imag); + llvm::Value* imag_result2 = ir_builder->CreateFMul(lhs_imag, rhs_real); + llvm::Value* imag_result = ir_builder->CreateFAdd(imag_result1, imag_result2); + return {real_result, imag_result}; +} +} // namespace + Status IrEmitter::HandleDot(HloInstruction* dot) { auto lhs_instruction = dot->operand(0); auto rhs_instruction = dot->operand(1); @@ -454,21 +483,10 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { rhs_array.EmitReadArrayElement(/*index=*/{}, &ir_builder_); llvm::Value* result; if (ShapeUtil::ElementIsComplex(lhs_shape)) { - auto real = [&](llvm::Value* x) { - return ir_builder_.CreateExtractValue(x, {0}); - }; - auto imag = [&](llvm::Value* x) { - return ir_builder_.CreateExtractValue(x, {1}); - }; - llvm::Value* real_result = ir_builder_.CreateFSub( - ir_builder_.CreateFMul(real(lhs_value), real(rhs_value)), - ir_builder_.CreateFMul(imag(lhs_value), imag(rhs_value))); - llvm::Value* imag_result = ir_builder_.CreateFAdd( - ir_builder_.CreateFMul(real(lhs_value), imag(rhs_value)), - ir_builder_.CreateFMul(imag(lhs_value), real(rhs_value))); + auto value = MultiplyComplex(lhs_value, rhs_value, &ir_builder_); result = llvm::ConstantAggregateZero::get(lhs_array.GetElementLlvmType()); - result = ir_builder_.CreateInsertValue(result, real_result, {0}); - result = ir_builder_.CreateInsertValue(result, imag_result, {1}); + result = ir_builder_.CreateInsertValue(result, value.first, {0}); + result = ir_builder_.CreateInsertValue(result, value.second, {1}); } else { result = ir_builder_.CreateFMul(lhs_value, rhs_value); } @@ -546,20 +564,13 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { llvm::Value* accum = ir_builder_.CreateLoad(accum_address); llvm::Value* updated_accum; if (ShapeUtil::ElementIsComplex(lhs_shape)) { -#define REAL(x) ir_builder_.CreateExtractValue(x, {0}) -#define IMAG(x) ir_builder_.CreateExtractValue(x, {1}) - llvm::Value* product_real = ir_builder_.CreateFSub( - ir_builder_.CreateFMul(REAL(lhs_element), REAL(rhs_element)), - ir_builder_.CreateFMul(IMAG(lhs_element), IMAG(rhs_element))); - llvm::Value* product_imag = ir_builder_.CreateFAdd( - ir_builder_.CreateFMul(REAL(lhs_element), IMAG(rhs_element)), - ir_builder_.CreateFMul(IMAG(lhs_element), REAL(rhs_element))); - updated_accum = ir_builder_.CreateInsertValue( - accum, ir_builder_.CreateFAdd(REAL(accum), product_real), {0}); - updated_accum = ir_builder_.CreateInsertValue( - updated_accum, ir_builder_.CreateFAdd(IMAG(accum), product_imag), {1}); -#undef IMAG -#undef REAL + auto value = MultiplyComplex(lhs_element, rhs_element, &ir_builder_); + llvm::Value* accum_real = Real(accum, &ir_builder_); + llvm::Value* real_sum = ir_builder_.CreateFAdd(accum_real, value.first); + updated_accum = ir_builder_.CreateInsertValue(accum, real_sum, {0}); + llvm::Value* accum_imag = Imag(accum, &ir_builder_); + llvm::Value* imag_sum = ir_builder_.CreateFAdd(accum_imag, value.second); + updated_accum = ir_builder_.CreateInsertValue(updated_accum, imag_sum, {1}); } else { llvm::Value* product = ir_builder_.CreateFMul(lhs_element, rhs_element); updated_accum = ir_builder_.CreateFAdd(accum, product); @@ -615,8 +626,7 @@ Status IrEmitter::HandleFft(HloInstruction* fft) { 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) { @@ -710,11 +720,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) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.h b/tensorflow/compiler/xla/service/gpu/ir_emitter.h index 3aa178410f05aef3630a4bd83b9651f6c1aac79b..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,19 +47,28 @@ 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; @@ -217,202 +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 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. - 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); - - // 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 ConvolutionThunk that calls DNN to implement `inst`. - std::unique_ptr BuildConvolutionThunk(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_; -}; - -// 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 fc8783e753d3819ee7a35b2ad660a25eafc42f76..2381d7a7d59ba2777e711138779b4493b8037f3d 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" @@ -32,6 +34,7 @@ limitations under the License. #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" @@ -39,7 +42,6 @@ limitations under the License. #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" @@ -73,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. @@ -135,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, @@ -152,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( @@ -175,13 +217,15 @@ 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( @@ -190,43 +234,32 @@ llvm::Function* IrEmitterUnnested::BuildKernelPrototype( // 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( /*Result=*/llvm::Type::getVoidTy(context), - std::vector(num_escaped_hlos + 1, - ir_builder_.getInt8PtrTy()), + 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); - } + // 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; - // 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); - } - kernel->addParamAttr(temp_buffer_arg_no, llvm::Attribute::NoAlias); - - // All arguments to a kernel must be aligned to kCudaMallocAlignBytes. - for (int64 i = 0; i < kernel->arg_size(); ++i) { + kernel->addDereferenceableAttr(arg_no + 1, alloc->size()); kernel->addParamAttr( - i, llvm::Attribute::get(context, llvm::Attribute::Alignment, - kCudaMallocAlignBytes)); + 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()))); + } } // TODO(b/65380986): Investigate if adding fast math flags for generated @@ -243,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)); @@ -278,10 +310,6 @@ Status IrEmitterUnnested::HandleConditional(HloInstruction* conditional) { } 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); } @@ -380,6 +408,76 @@ 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); } @@ -419,39 +517,6 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { TF_RETURN_IF_ERROR(root->Accept(&fused_emitter)); Shape input_shape = root->operand(0)->shape(); - // EmitReductionToVector requires the input shape to have a layout, but - // fused instructions don't have one. So we determine its layout from - // the fusion's operands. The choice of the layout only affects - // performance but not correctness. - auto choose_input_layout = []( - tensorflow::gtl::ArraySlice operands, - Shape* input_shape) -> Status { - // Prefer the layout of an operand whose shape is compatible with - // input_shape. - for (const HloInstruction* operand : operands) { - if (ShapeUtil::Compatible(*input_shape, operand->shape())) { - return LayoutUtil::CopyLayoutBetweenShapes(operand->shape(), - input_shape); - } - } - // If no operand has a compatible shape, prefer an operand that has - // the same rank at least. - for (const HloInstruction* operand : operands) { - if (ShapeUtil::Rank(*input_shape) == - ShapeUtil::Rank(operand->shape())) { - // Do not use CopyLayoutBetweenShapes because input_shape and - // operand->shape() may be incompatible. - *input_shape->mutable_layout() = operand->shape().layout(); - return Status::OK(); - } - } - // When all the above fails, which is rare, set the default layout. - LayoutUtil::SetToDefaultLayout(input_shape); - return Status::OK(); - }; - TF_RETURN_IF_ERROR( - choose_input_layout(fusion->operands(), &input_shape)); - return EmitReductionToVector( root, input_shape, fused_emitter.GetGenerator(root->operand(0)), fused_emitter.GetGenerator(root->operand(1)), root->dimensions(), @@ -500,10 +565,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); } @@ -805,7 +866,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(); } @@ -1599,24 +1661,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( @@ -1658,8 +1720,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 @@ -1868,62 +1929,207 @@ 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; +} + +Status IrEmitterUnnested::HandleGather(HloInstruction* gather) { + // TODO(b/72710576): Gather is not implemented on GPUs + return Unimplemented("Gather is not implemented on GPUs."); } 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; + } + } - // 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())); + // 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); } - // Create a KernelThunk that launches the kernel that implements "inst". - return MakeUnique(io_buffers, - llvm_ir::AsString(kernel->getName()), inst); + // 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())); + } + + return MakeUnique(buffers, llvm_ir::AsString(kernel->getName()), + inst); } std::unique_ptr IrEmitterUnnested::BuildHostToDeviceCopyThunk( @@ -1982,82 +2188,68 @@ std::unique_ptr IrEmitterUnnested::BuildGemmThunk( inst->shape(), // The shape of the output. false, // Do not transpose LHS. false, // Do not transpose RHS. + 1.0, // alpha. inst); } if (inst->opcode() == HloOpcode::kFusion) { - const HloInstruction* dot = inst->fused_expression_root(); - DCHECK(dot->opcode() == HloOpcode::kDot); - const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0)); - const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1)); - DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter && - rhs_parameter->opcode() == HloOpcode::kParameter); - const HloInstruction* lhs = - inst->operand(lhs_parameter->parameter_number()); - const HloInstruction* rhs = - inst->operand(rhs_parameter->parameter_number()); - - return MakeUnique( - GetAllocationSlice(*lhs), // The buffer assigned to LHS. - GetAllocationSlice(*rhs), // The buffer assigned to RHS. - GetAllocationSlice(*inst), // The output buffer. - lhs->shape(), // The shape of LHS. - rhs->shape(), // The shape of RHS. - inst->shape(), // The shape of the output. - dot->operand(0)->IsRank2Transpose(), // Transpose LHS. - dot->operand(1)->IsRank2Transpose(), // Trasnpose RHS. - inst); + if (inst->fusion_kind() == HloInstruction::FusionKind::kOutput) { + const HloInstruction* mul = inst->fused_expression_root(); + const HloInstruction* dot = mul->operand(0); + const HloInstruction* alpha = mul->operand(1); + if (dot->opcode() != HloOpcode::kDot) { + std::swap(dot, alpha); + } + DCHECK(dot->opcode() == HloOpcode::kDot); + const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0)); + const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1)); + DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter && + rhs_parameter->opcode() == HloOpcode::kParameter); + const HloInstruction* lhs = + inst->operand(lhs_parameter->parameter_number()); + const HloInstruction* rhs = + inst->operand(rhs_parameter->parameter_number()); + + return MakeUnique( + GetAllocationSlice(*lhs), // The buffer assigned to LHS. + GetAllocationSlice(*rhs), // The buffer assigned to RHS. + GetAllocationSlice(*mul), // The output buffer. + lhs->shape(), // The shape of LHS. + rhs->shape(), // The shape of RHS. + inst->shape(), // The shape of the output. + dot->operand(0)->IsRank2Transpose(), // Transpose LHS. + dot->operand(1)->IsRank2Transpose(), // Transpose RHS. + alpha->literal().Get({0}), // alpha. + inst); + } else { + const HloInstruction* dot = inst->fused_expression_root(); + DCHECK(dot->opcode() == HloOpcode::kDot); + const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0)); + const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1)); + DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter && + rhs_parameter->opcode() == HloOpcode::kParameter); + const HloInstruction* lhs = + inst->operand(lhs_parameter->parameter_number()); + const HloInstruction* rhs = + inst->operand(rhs_parameter->parameter_number()); + + return MakeUnique( + GetAllocationSlice(*lhs), // The buffer assigned to LHS. + GetAllocationSlice(*rhs), // The buffer assigned to RHS. + GetAllocationSlice(*inst), // The output buffer. + lhs->shape(), // The shape of LHS. + rhs->shape(), // The shape of RHS. + inst->shape(), // The shape of the output. + dot->operand(0)->IsRank2Transpose(), // Transpose LHS. + dot->operand(1)->IsRank2Transpose(), // Transpose RHS. + 1.0, // Alpha. + inst); + } } LOG(FATAL) << "Cannot build a GemmThunk for " << inst->ToString(); } -std::unique_ptr IrEmitterUnnested::BuildConvolutionThunk( - 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"; - } -} - std::unique_ptr IrEmitterUnnested::BuildFftThunk( const HloInstruction* inst) { const HloInstruction* operand = inst->operand(0); @@ -2259,6 +2451,8 @@ std::unique_ptr IrEmitterUnnested::BuildConditionalThunk( 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..b83a2337e2decd9d4fba3d40fcf33f131fca8a3c --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h @@ -0,0 +1,206 @@ +/* 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 HandleGather(HloInstruction* gather) 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 cfabae791d26d0eb49826085ad7ad166a19109a1..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(); } diff --git a/tensorflow/compiler/xla/service/gpu/pad_insertion.cc b/tensorflow/compiler/xla/service/gpu/pad_insertion.cc index 2923a79af0a559b08a2126162130a83801d024f8..7bda4e2fcd469bd430e5ef1846251c8504225383 100644 --- a/tensorflow/compiler/xla/service/gpu/pad_insertion.cc +++ b/tensorflow/compiler/xla/service/gpu/pad_insertion.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.h" #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/window_util.h" @@ -27,7 +28,7 @@ namespace gpu { namespace { bool IsForwardConvolutionCanonical(const HloInstruction& conv) { - CHECK_EQ(HloOpcode::kConvolution, conv.opcode()); + CHECK_EQ(conv.custom_call_target(), kCudnnConvForwardCallTarget); return window_util::HasSymmetricPadding(conv.window()) && !window_util::HasNegativePadding(conv.window()) && !window_util::HasDilation(conv.window()); @@ -47,6 +48,12 @@ HloInstruction* MaybePaddedAndSlicedInput( 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) { @@ -62,13 +69,7 @@ HloInstruction* MaybePaddedAndSlicedInput( HloInstruction* padding = computation->AddInstruction(HloInstruction::CreateConstant( MakeUnique(Literal::Zero(element_type)))); - input = computation->AddInstruction(HloInstruction::CreatePad( - ShapeInference::InferPadShape( - /*operand_shape=*/input->shape(), - /*padding_value_shape=*/ShapeUtil::MakeShape(element_type, {}), - padding_config) - .ConsumeValueOrDie(), - input, padding, padding_config)); + input = MakePadHlo(input, padding, padding_config).ValueOrDie(); } if (window_util::HasNegativePadding(conv_window)) { @@ -91,11 +92,8 @@ HloInstruction* MaybePaddedAndSlicedInput( std::max(0LL, -conv_window.dimensions(i).padding_high()); } - input = computation->AddInstruction(HloInstruction::CreateSlice( - ShapeInference::InferSliceShape(input->shape(), start_indices, - limit_indices, strides) - .ConsumeValueOrDie(), - input, start_indices, limit_indices, strides)); + input = + MakeSliceHlo(input, start_indices, limit_indices, strides).ValueOrDie(); } return input; @@ -128,13 +126,7 @@ HloInstruction* MaybePaddedKernel(const Window& conv_window, HloInstruction* padding = computation->AddInstruction(HloInstruction::CreateConstant( MakeUnique(Literal::Zero(element_type)))); - return computation->AddInstruction(HloInstruction::CreatePad( - ShapeInference::InferPadShape( - /*operand_shape=*/kernel->shape(), - /*padding_value_shape=*/ShapeUtil::MakeShape(element_type, {}), - padding_config) - .ConsumeValueOrDie(), - kernel, padding, padding_config)); + return MakePadHlo(kernel, padding, padding_config).ValueOrDie(); } } // namespace @@ -167,14 +159,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,6 +185,8 @@ void IncreasePaddingHighBy(int64 delta, WindowDimension* window_dim) { bool PadInsertion::CanonicalizeBackwardFilterConvolution( HloInstruction* backward_conv) { + CHECK_EQ(backward_conv->custom_call_target(), + kCudnnConvBackwardFilterCallTarget); if (window_util::HasSymmetricPadding(backward_conv->window())) { return false; } @@ -202,15 +199,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 +215,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 +226,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. @@ -254,25 +238,14 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( computation->AddInstruction(HloInstruction::CreateConstant( MakeUnique(Literal::Zero(input->shape().element_type())))); HloInstruction* padded_input = - computation->AddInstruction(HloInstruction::CreatePad( - ShapeInference::InferPadShape(input->shape(), padding->shape(), - input_padding_config) - .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); + MakePadHlo(input, padding, input_padding_config).ValueOrDie(); + + // 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 " @@ -289,14 +262,15 @@ bool PadInsertion::CanonicalizeBackwardInputConvolution( 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 +289,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 +348,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 +374,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.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/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc index cde5877e29f36abc61c5417ce960e2c7699e2749..3dd4c4a0794e5c41b877078c4e69c6c9584ce6c0 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.cc +++ b/tensorflow/compiler/xla/service/heap_simulator.cc @@ -27,38 +27,6 @@ namespace xla { using tensorflow::gtl::FlatMap; using tensorflow::gtl::FlatSet; -namespace { - -// Returns the set of buffers that may be sources of all operands of the given -// instruction. The returned buffers are guaranteed to have no duplicates, and -// to be sorted in a deterministic order. -std::vector UniqueOperandSourceBuffers( - const HloInstruction* instruction, - const TuplePointsToAnalysis& points_to_analysis) { - std::vector buffers; - for (const HloInstruction* operand : instruction->operands()) { - points_to_analysis.GetPointsToSet(operand).ForEachElement( - [&](const ShapeIndex& /*index*/, - const PointsToSet::BufferList& points_to) { - buffers.insert(buffers.end(), points_to.begin(), points_to.end()); - }); - } - - // Sort and then remove duplicates from buffers. - std::sort(buffers.begin(), buffers.end(), - [](const LogicalBuffer* a, const LogicalBuffer* b) { - return a->id() < b->id(); - }); - buffers.erase(std::unique(buffers.begin(), buffers.end(), - [](const LogicalBuffer* a, const LogicalBuffer* b) { - return a->id() == b->id(); - }), - buffers.end()); - return buffers; -} - -} // namespace - /*static*/ StatusOr HeapSimulator::Run( std::unique_ptr algorithm, const HloModule& module, @@ -93,6 +61,7 @@ Status HeapSimulator::RunComputation( const HloComputation& computation, const std::vector& instruction_sequence, const TuplePointsToAnalysis& points_to_analysis) { + VLOG(3) << "Computation:\n" << computation.ToString(); // The goal here is to minimize memory usage, assuming the given sequential // ordering of instructions. The strategy is to walk through the instruction // sequence, calling Alloc and Free on the underlying heap algorithm. The @@ -101,7 +70,51 @@ Status HeapSimulator::RunComputation( // 'live_buffers' tracks the liveness of each buffer that we assign, by // associating it with a set of HloInstructions that need to be visited. When // the set becomes empty, the buffer is no longer used, and can be freed. + // 'used_buffers' is the reverse map - it tracks which buffers were used by an + // instruction, so that we can remove the instructions from a buffer's live + // set after they are visited. FlatMap> live_buffers; + FlatMap> used_buffers; + auto add_user_to_buffer = [this, &live_buffers, &used_buffers]( + const HloInstruction* user, + const LogicalBuffer* buffer) { + if (!IgnoreBuffer(buffer)) { + VLOG(4) << " Adding user " << user->name() << " to buffer " + << buffer->ToString(); + live_buffers[buffer].insert(user); + used_buffers[user].insert(buffer); + } + }; + + // Initialize live_buffers for each buffer that we're going to assign. The + // set of instructions that need to be visited contains all users of all + // aliases, that is, all users of all instructions that have the buffer + // contained in their points-to set. + for (const HloInstruction* instruction : instruction_sequence) { + const PointsToSet& points_to = + points_to_analysis.GetPointsToSet(instruction); + const PointsToSet::BufferSet& buffer_set = points_to.CreateFlattenedSet(); + for (const HloInstruction* user : instruction->users()) { + if (user->opcode() != HloOpcode::kGetTupleElement) { + for (const LogicalBuffer* buffer : buffer_set) { + add_user_to_buffer(user, buffer); + } + } else { + // A GetTupleElement doesn't need to keep all of its operand's buffers + // alive. It only needs the buffers that relate to the element its + // extracting, and the tuple it's extracting from, but not the buffers + // for the other elements. + for (const LogicalBuffer* buffer : points_to.element({})) { + add_user_to_buffer(user, buffer); + } + const PointsToSet& gte_points_to = + points_to_analysis.GetPointsToSet(user); + for (const LogicalBuffer* buffer : gte_points_to.CreateFlattenedSet()) { + add_user_to_buffer(user, buffer); + } + } + } + } const HloInstruction* root = computation.root_instruction(); auto output_source_buffers = @@ -114,34 +127,17 @@ Status HeapSimulator::RunComputation( buffers_defined_by_instruction = points_to_analysis.GetBuffersDefinedByInstruction(instruction); - // Initialize live_buffers for each buffer that we're going to assign. The - // set of instructions that need to be visited contains all users of all - // aliases. The alias itself is not necessary; if it has users, the users - // are necessarily scheduled after the alias. And if it has no users, it is - // either a dead value or an output, both of which are handled below. - // - // We ignore control dependencies here. The reasoning is that the control - // dependencies have already been accounted for in the ordering of the given - // 'instruction_sequence', and should not otherwise artificially extend the - // lifetime of buffers that aren't already connected by a data dependency. + VLOG(3) << "Instruction: " << instruction->ToString(); + for (const LogicalBuffer* buffer : buffers_defined_by_instruction) { + VLOG(4) << " Defines: " << buffer->ToString() + << (IgnoreBuffer(buffer) ? " (Ignored)" : ""); + } + dead_buffers_to_free.clear(); for (const LogicalBuffer* buffer : buffers_defined_by_instruction) { if (IgnoreBuffer(buffer)) { continue; } - FlatSet* live_set = nullptr; - for (const BufferAlias& alias : - points_to_analysis.GetBufferAliases(*buffer)) { - const std::vector& users = - alias.instruction()->users(); - if (!users.empty()) { - if (live_set == nullptr) { - live_set = &live_buffers[buffer]; - } - live_set->insert(users.begin(), users.end()); - } - } - // Add a nullptr sentry to ensure entry parameters and output source // buffers are not freed until the very end. const bool entry_parameter = @@ -165,11 +161,12 @@ Status HeapSimulator::RunComputation( // have no instructions left to visit are moved from live_buffers to // operand_buffers_to_free. operand_buffers_to_free.clear(); - for (const LogicalBuffer* operand_buffer : - UniqueOperandSourceBuffers(instruction, points_to_analysis)) { + for (const LogicalBuffer* operand_buffer : used_buffers[instruction]) { if (IgnoreBuffer(operand_buffer)) { continue; } + VLOG(4) << " Removing user " << instruction->name() << " from buffer " + << operand_buffer->ToString(); auto it = live_buffers.find(operand_buffer); FlatSet* live_set = &it->second; live_set->erase(instruction); @@ -178,6 +175,11 @@ Status HeapSimulator::RunComputation( operand_buffers_to_free.push_back(operand_buffer); } } + // Sort to get a deterministic iteration order. + std::sort(operand_buffers_to_free.begin(), operand_buffers_to_free.end(), + [](const LogicalBuffer* x, const LogicalBuffer* y) { + return x->id() < y->id(); + }); // Allocate buffers defined by this instruction. This is the latest point // that we can allocate; right before the buffer is first used. This must @@ -203,6 +205,8 @@ Status HeapSimulator::RunComputation( CanShareOperandBufferWithUser( operand_buffer->instruction(), operand_buffer->index(), buffer->instruction(), buffer->index(), points_to_analysis)) { + VLOG(3) << " Sharing: " << buffer->ToString() << " with " + << operand_buffer->ToString(); ShareBuffer(buffer, operand_buffer, instruction); shared = true; break; @@ -211,6 +215,7 @@ Status HeapSimulator::RunComputation( } if (!shared) { + VLOG(3) << " Allocating: " << buffer->ToString(); Alloc(buffer, instruction); } } @@ -225,6 +230,7 @@ Status HeapSimulator::RunComputation( // sub-computations will never be run concurrently. if (module_sequence_ != nullptr) { if (instruction->opcode() == HloOpcode::kCall || + instruction->opcode() == HloOpcode::kConditional || instruction->opcode() == HloOpcode::kWhile) { for (const HloComputation* called_computation : instruction->called_computations()) { @@ -243,20 +249,34 @@ Status HeapSimulator::RunComputation( // Free buffers that are no longer live. This is the earliest point that we // can de-allocate; right after the last use of the buffer. for (const LogicalBuffer* buffer : dead_buffers_to_free) { + VLOG(3) << " Freeing dead: " << buffer->ToString(); Free(buffer, instruction); } for (const LogicalBuffer* buffer : operand_buffers_to_free) { + VLOG(3) << " Freeing operand: " << buffer->ToString(); Free(buffer, instruction); } } // Any remaining live buffers must be entry parameters or output source - // buffers, which had a nullptr sentry added. Free them now. + // buffers, which had a nullptr sentry added. Free them now, in a + // deterministic order. + std::vector to_free; + to_free.reserve(live_buffers.size()); for (const auto& buffer_pending : live_buffers) { const LogicalBuffer* buffer = buffer_pending.first; const FlatSet& pending = buffer_pending.second; CHECK_EQ(pending.size(), 1) << *buffer; CHECK(*pending.begin() == nullptr) << *buffer; + to_free.push_back(buffer); + } + + std::sort(to_free.begin(), to_free.end(), + [](const LogicalBuffer* x, const LogicalBuffer* y) { + return x->id() < y->id(); + }); + for (const LogicalBuffer* buffer : to_free) { + VLOG(3) << "Freeing pending: " << buffer->ToString(); Free(buffer, root); } diff --git a/tensorflow/compiler/xla/service/heap_simulator_test.cc b/tensorflow/compiler/xla/service/heap_simulator_test.cc index 387b649a731ebcbfd8307807469f39f22d192b06..688a271712ac243666ba4ff02932aa4f7f7ed21c 100644 --- a/tensorflow/compiler/xla/service/heap_simulator_test.cc +++ b/tensorflow/compiler/xla/service/heap_simulator_test.cc @@ -410,6 +410,56 @@ TEST_F(HeapSimulatorTest, MultiplyDotDotTuple) { }); } +TEST_F(HeapSimulatorTest, IndependentTupleElements) { + auto builder = HloComputation::Builder(TestName()); + auto paramA = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32scalar_, "paramA")); + auto paramB = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32scalar_, "paramB")); + auto mul = builder.AddInstruction(HloInstruction::CreateBinary( + f32scalar_, HloOpcode::kMultiply, paramA, paramB)); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + f32scalar_, HloOpcode::kAdd, paramA, paramB)); + auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({mul, add})); + auto element0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(f32scalar_, tuple, 0)); + auto broadcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(f32vec4_, element0, {0})); + auto sub = builder.AddInstruction(HloInstruction::CreateBinary( + f32scalar_, HloOpcode::kSubtract, paramA, paramB)); + auto element1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(f32scalar_, tuple, 1)); + auto output = builder.AddInstruction( + HloInstruction::CreateTuple({broadcast, sub, element1})); + + HeapSimulatorTracker tracker(TestName(), builder.Build(), + {paramA, paramB, mul, add, tuple, element0, + broadcast, sub, element1, output}); + tracker.ExpectCallSequence({ + {kAlloc, tracker.BufferAt(paramA, {})}, + {kAlloc, tracker.BufferAt(paramB, {})}, + {kAlloc, tracker.BufferAt(mul, {})}, + {kAlloc, tracker.BufferAt(add, {})}, + {kAlloc, tracker.BufferAt(tuple, {})}, + {kAlloc, tracker.BufferAt(broadcast, {})}, + // The mul can be freed right after the broadcast happens, even though + // The other GetTupleElement is still alive. + {kFree, tracker.BufferAt(mul, {})}, + {kAlloc, tracker.BufferAt(sub, {})}, + // The temporary tuple is now dead. + {kFree, tracker.BufferAt(tuple, {})}, + {kAlloc, tracker.BufferAt(output, {})}, + // All params and outputs are freed at the end. + {kFree, tracker.BufferAt(paramA, {})}, + {kFree, tracker.BufferAt(paramB, {})}, + {kFree, tracker.BufferAt(add, {})}, + {kFree, tracker.BufferAt(broadcast, {})}, + {kFree, tracker.BufferAt(sub, {})}, + {kFree, tracker.BufferAt(output, {})}, + {kFinish, nullptr}, + }); +} + TEST_F(HeapSimulatorTest, WholeModule) { HeapSimulatorTracker tracker(TestName()); diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index 0e9a852788e978f79fa6f6c802f855a4c476583f..bf903d6a390fe2951d33942dfc2e124868c9fdb5 100644 --- a/tensorflow/compiler/xla/service/hlo.proto +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -129,6 +129,13 @@ message HloInstructionProto { // FFT length. repeated int64 fft_length = 32; + + // Gather dimension numbers. + xla.GatherDimensionNumbers gather_dimension_numbers = 33; + repeated int64 gather_window_bounds = 34; + + // The id of this instruction. + int64 id = 35; } // Serialization of HloComputation. @@ -141,6 +148,12 @@ message HloComputationProto { // The name of the root of the computation. string root_name = 3; + + // The program shape (with layout) of this computation. + xla.ProgramShape program_shape = 4; + + // The id of this computation. + int64 id = 5; } // Serialization of HloModule. @@ -151,6 +164,12 @@ message HloModuleProto { // The array of computations is always in a valid dependency order, where // callees appear before their callers. repeated HloComputationProto computations = 3; + + // The program shape (with layout) of the entry computation. + xla.ProgramShape program_shape = 4; + + // The id of this module. + int64 id = 5; } // Serialization of HloOrdering. @@ -200,6 +219,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_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc index 6d2a3aa5b531650a658502531e050702ffbd3760..a88283ed9a6459b4fa9310e160b59c77d51f1027 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc @@ -171,24 +171,21 @@ class BufferValueMap { return value_to_buffer_number_.at(&value); } - // Compute and return a vector of buffers that the given value must be - // contained in due to HLO aliasing rules. - std::vector ComputeAliasedBuffers(const HloValue& value) { + void ComputeWhileAliasedBuffers(const HloValue& value, + std::vector* aliased_buffers) { + VLOG(3) << "Compute kWhile aliases"; // Value is init of a while (use is while). - std::vector aliased_buffers; for (const HloUse& use : value.uses()) { - VLOG(2) << "use of value " << value.ToShortString() << ": " << use; if (use.instruction->opcode() == HloOpcode::kWhile) { // Determine the while value that this shares a buffer with. const HloValue& while_value = dataflow_.GetUniqueValueAt(use.instruction, use.operand_index); - aliased_buffers.push_back(GetBufferForValue(while_value)); + aliased_buffers->push_back(GetBufferForValue(while_value)); VLOG(3) << " value is init value to a while; must share buffer with " "while value " << while_value.ToShortString(); } } - // Value is a parameter of a while body/condition. if (value.defining_instruction()->opcode() == HloOpcode::kParameter) { const HloComputation* computation = @@ -205,11 +202,10 @@ class BufferValueMap { VLOG(3) << " value is parameter value of the body or condition of a " "while; must share buffer with while value " << while_value.ToShortString(); - aliased_buffers.push_back(GetBufferForValue(while_value)); + aliased_buffers->push_back(GetBufferForValue(while_value)); } } } - // Value is the root of a while body. for (const HloPosition& position : value.positions()) { const HloComputation* computation = position.instruction->parent(); @@ -224,27 +220,71 @@ class BufferValueMap { const HloValue& while_value = dataflow_.GetUniqueValueAt( callsite.instruction(), position.index); - VLOG(3) << " value is root the body computation of a while; must " - "share buffer with while value " + VLOG(3) << " value @ " << position << " is root of " + << callsite.instruction()->name() + << "; body root and while value root must share buffer " + "among them : " << while_value.ToShortString(); - aliased_buffers.push_back(GetBufferForValue(while_value)); + aliased_buffers->push_back(GetBufferForValue(while_value)); } } } } - // Value is the output of the while instruction itself. if (value.defining_instruction()->opcode() == HloOpcode::kWhile) { VLOG(3) << " value is output of a while instruction"; - aliased_buffers.push_back(GetBufferForValue(value)); + aliased_buffers->push_back(GetBufferForValue(value)); + } + } + + void ComputeConditionalAliasedBuffers( + const HloValue& value, std::vector* aliased_buffers) { + VLOG(3) << "Compute kConditional aliases"; + // Aliases the buffers of the true/false computations roots, with the one of + // the conditional. + for (const HloPosition& position : value.positions()) { + const HloComputation* computation = position.instruction->parent(); + const CallGraphNode& call_graph_node = + dataflow_.call_graph().GetNode(computation); + if (position.instruction == computation->root_instruction()) { + for (const CallSite& callsite : call_graph_node.caller_callsites()) { + if (callsite.instruction()->opcode() == HloOpcode::kConditional) { + // Call graph must have been flattened. + CHECK_EQ(call_graph_node.caller_callsites().size(), 1); + + const HloValue& cond_value = dataflow_.GetUniqueValueAt( + callsite.instruction(), position.index); + VLOG(3) + << " value @ " << position << " is root of " + << callsite.instruction()->name() + << "; true/false branch roots must share buffer among them : " + << cond_value.ToShortString(); + aliased_buffers->push_back(GetBufferForValue(cond_value)); + } + } + } + } + // Value is the output of the conditional instruction itself. + if (value.defining_instruction()->opcode() == HloOpcode::kConditional) { + VLOG(3) << " value is output of a conditional instruction"; + aliased_buffers->push_back(GetBufferForValue(value)); } + } + // Compute and return a vector of buffers that the given value must be + // contained in due to HLO aliasing rules. + std::vector ComputeAliasedBuffers(const HloValue& value) { + for (const HloUse& use : value.uses()) { + VLOG(2) << "Use of value " << value.ToShortString() << ": " << use; + } + std::vector aliased_buffers; + ComputeWhileAliasedBuffers(value, &aliased_buffers); + ComputeConditionalAliasedBuffers(value, &aliased_buffers); // Uniquify aliased buffers. std::sort(aliased_buffers.begin(), aliased_buffers.end()); aliased_buffers.erase( std::unique(aliased_buffers.begin(), aliased_buffers.end()), aliased_buffers.end()); - return aliased_buffers; } @@ -419,7 +459,7 @@ StatusOr> HloAliasAnalysis::Run( auto alias_analysis = WrapUnique(new HloAliasAnalysis(module)); TF_ASSIGN_OR_RETURN( alias_analysis->dataflow_analysis_, - HloDataflowAnalysis::Run(module, /*ssa_form=*/true, + HloDataflowAnalysis::Run(*module, /*ssa_form=*/true, /*bitcast_defines_value=*/false)); BufferValueMap buffer_map(alias_analysis->dataflow_analysis()); diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index a63affa06caf75f1ccab084bd114e39ba7c91a38..f99c7cf5e495eaf83e0dda859ef31a7487bc6ffe 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -399,6 +399,7 @@ HloComputationProto HloComputation::ToProto() const { proto.add_instructions()->Swap(&instruction_proto); } proto.set_root_name(root_instruction()->name()); + *proto.mutable_program_shape() = ComputeProgramShape(); return proto; } @@ -461,20 +462,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) { @@ -523,13 +510,14 @@ StatusOr HloComputation::DeepCopyInstruction( "Can't deep copy instruction %s: instruction is not in computation %s", instruction->name().c_str(), name().c_str()); } - if (indices_to_copy != nullptr && !ShapeUtil::Compatible(instruction->shape(), indices_to_copy->shape())) { return FailedPrecondition( "Can't deep copy instruction %s: given shape tree of indices to copy " - "has incompatible shape", - instruction->name().c_str()); + "has incompatible shapes: %s vs. %s", + instruction->name().c_str(), + ShapeUtil::HumanString(instruction->shape()).c_str(), + ShapeUtil::HumanString(indices_to_copy->shape()).c_str()); } ShapeIndex index; @@ -545,7 +533,6 @@ ProgramShape HloComputation::ComputeProgramShape() const { } *program_shape.mutable_result() = root_instruction_->shape(); - LayoutUtil::ClearLayout(&program_shape); return program_shape; } @@ -577,8 +564,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..dd9d346999f0eae448d74628278c802ccd3f51b4 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -77,6 +77,14 @@ class HloComputation { return last_added_instruction_; } + Status ForEachInstruction( + const std::function& func) const { + for (const auto& instruction : instructions_) { + TF_RETURN_IF_ERROR(func(instruction.get())); + } + return Status::OK(); + } + private: const string name_; HloInstruction* last_added_instruction_; @@ -224,15 +232,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 @@ -249,7 +248,7 @@ class HloComputation { ShapeTree* copies_added = nullptr); // Computes and returns the ProgramShape of this computation (shape of - // parameters and result without layout). + // parameters and result with layout). ProgramShape ComputeProgramShape() const; // Return whether `*this` and `other` are functionally equivalent. diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding.cc b/tensorflow/compiler/xla/service/hlo_constant_folding.cc index 53450991b6fad5b9651d9d23b55c908e6b68e5dd..35ecd4428d0dfde2de445ea34472d2c78148c6c9 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding.cc +++ b/tensorflow/compiler/xla/service/hlo_constant_folding.cc @@ -35,7 +35,10 @@ limitations under the License. namespace xla { StatusOr HloConstantFolding::Run(HloModule* module) { - auto evaluator = MakeUnique(); + // Limit the constant folding to 0 iterations to skip folding loops. This + // retains the behavior from before while loop support in HloEvaluator and may + // be revised. + auto evaluator = MakeUnique(/*max_loop_iterations=*/0); XLA_VLOG_LINES(2, "HloConstantFolding::Run(), before:\n" + module->ToString()); diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index cd54eb74d18d0be714b5b56fc8ae0dfa55ff31a0..4ec2ef27bf59b0c877ec38e55ef5c12debeec227 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -229,6 +229,10 @@ Status HloCostAnalysis::HandleOutfeed(const HloInstruction*) { return Status::OK(); } +Status HloCostAnalysis::HandleHostCompute(const HloInstruction*) { + return Status::OK(); +} + Status HloCostAnalysis::HandleMap(const HloInstruction* map) { // Compute properties of the mapped function. TF_ASSIGN_OR_RETURN(const Properties sub_properties, @@ -469,7 +473,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) { @@ -523,6 +533,11 @@ Status HloCostAnalysis::HandleConditional(const HloInstruction* conditional) { return Status::OK(); } +Status HloCostAnalysis::HandleGather(const HloInstruction* gather) { + // Gather does not issue any flops. + return Status::OK(); +} + Status HloCostAnalysis::FinishVisit(const HloInstruction*) { return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.h b/tensorflow/compiler/xla/service/hlo_cost_analysis.h index e5783539e5436f09fa58bf7889118380ee90fea0..d17678d20f2a23fd98d18b77d5fb25853901a789 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -71,6 +71,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleCrossReplicaSum(const HloInstruction* crs) override; Status HandleInfeed(const HloInstruction* infeed) override; Status HandleOutfeed(const HloInstruction* outfeed) override; + Status HandleHostCompute(const HloInstruction* host_compute) override; Status HandleRng(const HloInstruction* random) override; Status HandleReverse(const HloInstruction* reverse) override; Status HandleSort(const HloInstruction* sort) override; @@ -99,6 +100,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleTranspose(const HloInstruction* transpose) override; Status HandleWhile(const HloInstruction* xla_while) override; Status HandleConditional(const HloInstruction* conditional) override; + Status HandleGather(const HloInstruction* gather) override; Status FinishVisit(const HloInstruction* root) override; Status Preprocess(const HloInstruction* hlo) override; diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.cc b/tensorflow/compiler/xla/service/hlo_creation_utils.cc new file mode 100644 index 0000000000000000000000000000000000000000..fbe71f8d5bf6daa7437a143e3cad44c8c54e697f --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc @@ -0,0 +1,295 @@ +/* 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_creation_utils.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/shape_inference.h" +#include "tensorflow/compiler/xla/util.h" + +namespace xla { +using tensorflow::gtl::ArraySlice; +using tensorflow::strings::StrCat; + +StatusOr MakeBinaryHlo(HloOpcode opcode, HloInstruction* lhs, + HloInstruction* rhs) { + HloComputation* computation = lhs->parent(); + CHECK_EQ(computation, rhs->parent()); + TF_ASSIGN_OR_RETURN(Shape binary_op_shape, + ShapeInference::InferBinaryOpShape(opcode, lhs, rhs)); + return computation->AddInstruction( + HloInstruction::CreateBinary(binary_op_shape, opcode, lhs, rhs)); +} + +StatusOr MakePadHlo(HloInstruction* operand, + HloInstruction* padding_value, + const PaddingConfig& padding_config) { + HloComputation* computation = operand->parent(); + CHECK_EQ(computation, padding_value->parent()); + TF_ASSIGN_OR_RETURN( + Shape pad_shape, + ShapeInference::InferPadShape(operand->shape(), padding_value->shape(), + padding_config)); + return computation->AddInstruction(HloInstruction::CreatePad( + pad_shape, operand, padding_value, padding_config)); +} + +StatusOr MakeSliceHlo(HloInstruction* operand, + ArraySlice start_indices, + ArraySlice limit_indices, + ArraySlice strides) { + HloComputation* computation = operand->parent(); + TF_ASSIGN_OR_RETURN(Shape slice_shape, ShapeInference::InferSliceShape( + operand->shape(), start_indices, + limit_indices, strides)); + return computation->AddInstruction(HloInstruction::CreateSlice( + slice_shape, operand, start_indices, limit_indices, strides)); +} + +StatusOr MakeConvolveHlo( + HloInstruction* lhs, HloInstruction* rhs, const Window& window, + const ConvolutionDimensionNumbers& dimension_numbers) { + HloComputation* computation = lhs->parent(); + CHECK_EQ(computation, rhs->parent()); + TF_ASSIGN_OR_RETURN(Shape convolve_shape, ShapeInference::InferConvolveShape( + lhs->shape(), rhs->shape(), + window, dimension_numbers)); + return computation->AddInstruction(HloInstruction::CreateConvolve( + convolve_shape, lhs, rhs, window, dimension_numbers)); +} + +StatusOr MakeTransposeHlo(HloInstruction* operand, + ArraySlice dimensions) { + HloComputation* computation = operand->parent(); + TF_ASSIGN_OR_RETURN( + Shape transpose_shape, + ShapeInference::InferTransposeShape(operand->shape(), dimensions)); + return computation->AddInstruction( + HloInstruction::CreateTranspose(transpose_shape, operand, dimensions)); +} + +StatusOr MakeReshapeHlo(const Shape& result_shape, + HloInstruction* operand) { + HloComputation* computation = operand->parent(); + return computation->AddInstruction( + HloInstruction::CreateReshape(result_shape, operand)); +} + +StatusOr MakeReshapeHlo( + ArraySlice result_shape_dim_bounds, HloInstruction* operand) { + Shape new_shape = ShapeUtil::MakeShape(operand->shape().element_type(), + result_shape_dim_bounds); + return MakeReshapeHlo(new_shape, operand); +} + +StatusOr MakeDynamicSliceHlo(HloInstruction* operand, + HloInstruction* start_indices, + ArraySlice slice_sizes) { + HloComputation* computation = operand->parent(); + CHECK_EQ(computation, start_indices->parent()); + TF_ASSIGN_OR_RETURN( + Shape dynamic_slice_shape, + ShapeInference::InferDynamicSliceShape( + operand->shape(), start_indices->shape(), slice_sizes)); + return computation->AddInstruction(HloInstruction::CreateDynamicSlice( + dynamic_slice_shape, operand, start_indices, slice_sizes)); +} + +StatusOr MakeDynamicUpdateSliceHlo( + HloInstruction* operand, HloInstruction* update, + HloInstruction* start_indices) { + HloComputation* computation = operand->parent(); + CHECK_EQ(computation, update->parent()); + CHECK_EQ(computation, start_indices->parent()); + TF_ASSIGN_OR_RETURN( + Shape dynamic_update_slice_shape, + ShapeInference::InferDynamicUpdateSliceShape( + operand->shape(), update->shape(), start_indices->shape())); + return computation->AddInstruction(HloInstruction::CreateDynamicUpdateSlice( + dynamic_update_slice_shape, operand, update, start_indices)); +} + +StatusOr MakeBroadcastHlo( + HloInstruction* operand, ArraySlice broadcast_dimensions, + ArraySlice result_shape_bounds) { + HloComputation* computation = operand->parent(); + Shape broadcast_shape = ShapeUtil::MakeShape(operand->shape().element_type(), + result_shape_bounds); + + return computation->AddInstruction(HloInstruction::CreateBroadcast( + broadcast_shape, operand, broadcast_dimensions)); +} + +StatusOr MakeGetTupleElementHlo(HloInstruction* operand, + int64 index) { + HloComputation* computation = operand->parent(); + + TF_ASSIGN_OR_RETURN( + Shape gte_shape, + ShapeInference::InferGetTupleElementShape(operand->shape(), index)); + return computation->AddInstruction( + HloInstruction::CreateGetTupleElement(gte_shape, operand, index)); +} + +StatusOr MakeConcatHlo(ArraySlice operands, + int64 dimension) { + CHECK_GT(operands.size(), 0); + + HloComputation* computation = operands[0]->parent(); + CHECK(c_all_of(operands, [&](HloInstruction* instr) { + return instr->parent() == computation; + })); + + std::vector operand_shapes; + c_transform(operands, std::back_inserter(operand_shapes), + [](HloInstruction* instr) { return &instr->shape(); }); + + TF_ASSIGN_OR_RETURN(Shape concat_shape, ShapeInference::InferConcatOpShape( + operand_shapes, dimension)); + return computation->AddInstruction( + HloInstruction::CreateConcatenate(concat_shape, operands, dimension)); +} + +StatusOr CollapseFirstNDims(HloInstruction* operand, int64 n) { + const Shape& operand_shape = operand->shape(); + CHECK_GE(operand_shape.dimensions_size(), n); + int64 new_shape_leading_bound = 1; + for (int64 i = 0; i < n; i++) { + new_shape_leading_bound *= operand_shape.dimensions(i); + } + + std::vector new_shape_dims; + new_shape_dims.reserve(operand_shape.dimensions_size() - n + 1); + new_shape_dims.push_back(new_shape_leading_bound); + + std::copy(operand_shape.dimensions().begin() + n, + operand_shape.dimensions().end(), + std::back_inserter(new_shape_dims)); + + Shape output_shape = + ShapeUtil::MakeShape(operand_shape.element_type(), new_shape_dims); + + return MakeReshapeHlo(output_shape, operand); +} + +StatusOr ExpandFirstDimIntoNDims( + HloInstruction* operand, ArraySlice expanded_dims) { + CHECK_GT(operand->shape().dimensions_size(), 0); + CHECK_EQ(operand->shape().dimensions(0), Product(expanded_dims)); + + std::vector expanded_shape_dim_bounds; + expanded_shape_dim_bounds.reserve(expanded_dims.size() + + operand->shape().dimensions_size() - 1); + c_copy(expanded_dims, std::back_inserter(expanded_shape_dim_bounds)); + std::copy(operand->shape().dimensions().begin() + 1, + operand->shape().dimensions().end(), + std::back_inserter(expanded_shape_dim_bounds)); + Shape new_shape = ShapeUtil::MakeShape(operand->shape().element_type(), + expanded_shape_dim_bounds); + return MakeReshapeHlo(new_shape, operand); +} + +StatusOr ExpandLastDimIntoNDims( + HloInstruction* operand, ArraySlice expanded_dims) { + CHECK_GT(operand->shape().dimensions_size(), 0); + CHECK_EQ(operand->shape().dimensions(operand->shape().dimensions_size() - 1), + Product(expanded_dims)); + + std::vector expanded_shape_dim_bounds; + expanded_shape_dim_bounds.reserve(expanded_dims.size() + + operand->shape().dimensions_size() - 1); + std::copy(operand->shape().dimensions().begin(), + operand->shape().dimensions().end() - 1, + std::back_inserter(expanded_shape_dim_bounds)); + c_copy(expanded_dims, std::back_inserter(expanded_shape_dim_bounds)); + Shape new_shape = ShapeUtil::MakeShape(operand->shape().element_type(), + expanded_shape_dim_bounds); + return MakeReshapeHlo(new_shape, operand); +} + +StatusOr ElideDegenerateDims(HloInstruction* operand, + ArraySlice dims_to_elide) { + CHECK(c_is_sorted(dims_to_elide)); + + const Shape& input_shape = operand->shape(); + // First accumulate in reverse + std::vector new_shape_dim_bounds; + new_shape_dim_bounds.reserve(input_shape.dimensions_size() - + dims_to_elide.size()); + int64 dims_to_elide_idx = dims_to_elide.size() - 1; + for (int64 i = input_shape.dimensions_size() - 1; i >= 0; i--) { + if (dims_to_elide_idx >= 0 && i == dims_to_elide[dims_to_elide_idx]) { + CHECK_EQ(input_shape.dimensions(i), 1); + dims_to_elide_idx--; + } else { + new_shape_dim_bounds.push_back(input_shape.dimensions(i)); + } + } + + c_reverse(new_shape_dim_bounds); + Shape output_shape = + ShapeUtil::MakeShape(input_shape.element_type(), new_shape_dim_bounds); + return MakeReshapeHlo(output_shape, operand); +} + +StatusOr PadVectorWithZeros(HloInstruction* operand, + int64 zeros_to_prepend, + int64 zeros_to_append) { + HloComputation* computation = operand->parent(); + CHECK_EQ(operand->shape().dimensions_size(), 1); + PaddingConfig padding_config; + PaddingConfig::PaddingConfigDimension padding_config_dim; + padding_config_dim.set_edge_padding_low(zeros_to_prepend); + padding_config_dim.set_edge_padding_high(zeros_to_append); + *padding_config.add_dimensions() = padding_config_dim; + + HloInstruction* zero = + computation->AddInstruction(HloInstruction::CreateConstant( + MakeUnique(Literal::Zero(operand->shape().element_type())))); + return MakePadHlo(operand, zero, padding_config); +} + +StatusOr BroadcastZeros( + HloComputation* computation, PrimitiveType element_type, + ArraySlice broadcast_dimensions) { + HloInstruction* zero = + computation->AddInstruction(HloInstruction::CreateConstant( + MakeUnique(Literal::Zero(element_type)))); + return MakeBroadcastHlo(zero, /*broadcast_dimensions=*/{}, + /*result_shape_bounds=*/broadcast_dimensions); +} + +StatusOr> CreateComputationWithSignature( + ArraySlice domain, const Shape& range, + tensorflow::StringPiece name) { + HloComputation::Builder b(name.ToString()); + int64 param_idx = 0; + for (const Shape* param_shape : domain) { + b.AddInstruction(HloInstruction::CreateParameter( + param_idx, *param_shape, StrCat("param.", param_idx))); + param_idx++; + } + + // We can't change the root type of a computation once it is created so create + // a dummy root instruction to give the computation the right root shape. In + // the future we may want to use a (recursive) broadcast here to avoid + // creating large constants. + b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateFromShape(range))); + + return b.Build(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.h b/tensorflow/compiler/xla/service/hlo_creation_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..6032ebab74464a7e4668e98dc638e76320854013 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.h @@ -0,0 +1,163 @@ +/* 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_HLO_CREATION_UTILS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_CREATION_UTILS_H_ + +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// Some lightweight utilities intended to make HLO instruction creation more +// ergonomic. We don't have a complete set of helpers yet -- I expect we'll +// expand this interface as needed on an ad-hoc basis. + +// Creates a binary HLO instruction and adds it to the computation containing +// `lhs` and `rhs` (`lhs` and `rhs` must be in the same computation). +StatusOr MakeBinaryHlo(HloOpcode opcode, HloInstruction* lhs, + HloInstruction* rhs); + +// Creates a pad HLO instruction and adds it to the computation containing +// `operand` and `padding_value` (`operand` and `padding_value` must be in the +// same computation). +StatusOr MakePadHlo(HloInstruction* operand, + HloInstruction* padding_value, + const PaddingConfig& padding_config); + +// Creates a slice HLO instruction and adds it to the computation containing +// `operand`. +StatusOr MakeSliceHlo( + HloInstruction* operand, tensorflow::gtl::ArraySlice start_indices, + tensorflow::gtl::ArraySlice limit_indices, + tensorflow::gtl::ArraySlice strides); + +// Creates a convolution HLO instruction and adds it to the computation +// containing `lhs` and `rhs` (`lhs` and `rhs` must be in the same computation). +StatusOr MakeConvolveHlo( + HloInstruction* lhs, HloInstruction* rhs, const Window& window, + const ConvolutionDimensionNumbers& dimension_numbers); + +// Creates a transpose HLO instruction and adds it to the computation containing +// `operand`. +StatusOr MakeTransposeHlo( + HloInstruction* operand, tensorflow::gtl::ArraySlice dimensions); + +// Creates a reshape HLO instruction and adds it to the computation containing +// `operand`. +StatusOr MakeReshapeHlo(const Shape& result_shape, + HloInstruction* operand); + +StatusOr MakeReshapeHlo( + tensorflow::gtl::ArraySlice result_shape_dim_bounds, + HloInstruction* operand); + +// Creates a dynamic-slice HLO instruction and adds it to the computation +// containing `operand` and `start_indices` (`operand` and `start_indices` must +// be in the same computation). +StatusOr MakeDynamicSliceHlo( + HloInstruction* operand, HloInstruction* start_indices, + tensorflow::gtl::ArraySlice slice_sizes); + +// Creates a dynamic-update-slice HLO instruction and adds it to the computation +// containing `operand`, `update` and `start_indices` (`operand`, `update` and +// `start_indices` must be in the same computation). +StatusOr MakeDynamicUpdateSliceHlo( + HloInstruction* operand, HloInstruction* update, + HloInstruction* start_indices); + +// Creates a broadcast HLO instruction and adds it to the computation containing +// `operand`. +StatusOr MakeBroadcastHlo( + HloInstruction* operand, + tensorflow::gtl::ArraySlice broadcast_dimensions, + tensorflow::gtl::ArraySlice result_shape_bounds); + +// Creates a GetTupleElement HLO instruction and adds it to the computation +// containing `operand`. +StatusOr MakeGetTupleElementHlo(HloInstruction* operand, + int64 index); + +// Creates a Concatenate HLO instruction and adds it to the computation +// containing `operands` (`operands` must be non-empty and every element must be +// contained in the same computation). +StatusOr MakeConcatHlo( + tensorflow::gtl::ArraySlice operands, int64 dimension); + +// ----------------------------------------------------------------------------- +// Some other miscellaneous helpers to generate common HLO patterns. All of +// these add all the instructions they generate into the computation containing +// their operand(s). + +// Collapses (via reshape) the first N (logical) dimensions of `operand` into a +// single leading dimension. `operand` must have rank > n. +// +// For instance if `operand` has shape f32[7,8,9] and n is 2 then the output is +// the `operand` reshaped to [56,9]. +StatusOr CollapseFirstNDims(HloInstruction* operand, int64 n); + +// Expands (via reshape) the first (logical) dimension of `operand` into a +// sequence of `expanded_dims` dimensions. `operand` must at least be of rank 1 +// and the number of elements in its first dimension must be equal to the +// product of `expanded_dims`. +// +// For instance if `operand` has shape f32[200,9,7] and expanded_dims is +// {2,5,20} the result is `operand` reshaped to [2,5,20,9,7]. +StatusOr ExpandFirstDimIntoNDims( + HloInstruction* operand, tensorflow::gtl::ArraySlice expanded_dims); + +// Expands (via reshape) the last (logical) dimension of `operand` into a +// sequence of `expanded_dims` dimensions. `operand` must at least be of rank 1 +// and the number of elements in its last dimension must be equal to the +// product of `expanded_dims`. +// +// For instance if `operand` has shape f32[9,7,200] and expanded_dims is +// {2,5,20} the result is `operand` reshaped to [9,7,2,5,20]. +StatusOr ExpandLastDimIntoNDims( + HloInstruction* operand, tensorflow::gtl::ArraySlice expanded_dims); + +// Elides (via reshape) a set of degenerate dimensions (dimensions containing +// exactly one element), `dims_to_elide` from `operand`. Every dimension in +// `dims_to_elide` must be a degenerate dimension. `dims_to_elide` must be +// sorted and not contain duplicates. +// +// For example if `operand` is of shape f32[19,1,20,1,7,1,9] and dims_to_elide +// is {1,5} then the result is `operand` reshaped to [19,20,1,7,9]. +StatusOr ElideDegenerateDims( + HloInstruction* operand, tensorflow::gtl::ArraySlice dims_to_elide); + +// Pads `operand` (which must have rank 1) with `zeros_to_prepend` zeros in the +// front and `zeros_to_append` zeros in the back. +StatusOr PadVectorWithZeros(HloInstruction* operand, + int64 zeros_to_prepend, + int64 zeros_to_append); + +// Broadcasts a zero value of type `element_type` into a tensor with element +// type `element_type` and dimension bounds `broadcast_dimensions`. The +// broadcast instruction is emitted into `computation`. +StatusOr BroadcastZeros( + HloComputation* computation, PrimitiveType element_type, + tensorflow::gtl::ArraySlice broadcast_dimensions); + +// Creates a HLO computation that takes arguments of type `domain` and produces +// a value of type `range`. +StatusOr> CreateComputationWithSignature( + tensorflow::gtl::ArraySlice domain, const Shape& range, + tensorflow::StringPiece name); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_CREATION_UTILS_H_ diff --git a/tensorflow/compiler/xla/service/hlo_cse.cc b/tensorflow/compiler/xla/service/hlo_cse.cc index 7feda2b3b040de1f0a14303ce1adcd21c6624c8b..279edd4ba8772a9c576f76f554de8ec68631b953 100644 --- a/tensorflow/compiler/xla/service/hlo_cse.cc +++ b/tensorflow/compiler/xla/service/hlo_cse.cc @@ -119,9 +119,8 @@ StatusOr HloCSE::Run(HloModule* module) { equivalent_instructions; for (HloInstruction* user : operand->users()) { if (user != instruction && - user->Identical(*instruction, eq_instructions, eq_computations) && - (!is_layout_sensitive_ || - ShapeUtil::Equal(user->shape(), instruction->shape()))) { + 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 d25fc5d7418ae40c7167f88d6172906482a58925..0c37a8d75f38dabaad886cc9d4adce8ab29ddf18 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc @@ -38,12 +38,12 @@ namespace xla { using ::tensorflow::strings::StrAppend; using ::tensorflow::strings::StrCat; -HloDataflowAnalysis::HloDataflowAnalysis(HloModule* module, bool ssa_form, +HloDataflowAnalysis::HloDataflowAnalysis(const HloModule& module, bool ssa_form, bool bitcast_defines_value) : module_(module), ssa_form_(ssa_form), bitcast_defines_value_(bitcast_defines_value), - call_graph_(CallGraph::Build(module)) {} + call_graph_(CallGraph::Build(&module)) {} bool HloDataflowAnalysis::ValueIsDefinedAt(const HloInstruction* instruction, const ShapeIndex& index) const { @@ -115,9 +115,9 @@ void HloDataflowAnalysis::DeleteMarkedValues() { } string HloDataflowAnalysis::ToString() const { - string out = StrCat("HloDataflowAnalysis, module ", module_->name(), "\n"); + string out = StrCat("HloDataflowAnalysis, module ", module_.name(), "\n"); StrAppend(&out, " Instruction value sets:\n"); - for (const HloComputation* computation : module_->computations()) { + for (const HloComputation* computation : module_.computations()) { for (const HloInstruction* instruction : computation->instructions()) { StrAppend(&out, " ", instruction->name(), ":\n"); if (ShapeUtil::IsTuple(instruction->shape())) { @@ -368,11 +368,11 @@ bool HloDataflowAnalysis::UpdateConditionalValueSet( conditional->true_computation()->root_instruction()), &GetInstructionValueSet( conditional->false_computation()->root_instruction())}; - // A phi-node is not defined for a kConditional instruction even though it - // represents a join point. This is because the current approach is to define - // a phi-node only for kWhile to account for the dataflow through back-edges - // and deal with the ambiguity in other cases. - return GetInstructionValueSet(conditional).AssignUnionOf(inputs); + if (ssa_form_) { + return Phi(conditional, inputs); + } else { + return GetInstructionValueSet(conditional).AssignUnionOf(inputs); + } } bool HloDataflowAnalysis::UpdateCopyValueSet(HloInstruction* copy) { @@ -585,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 (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(); @@ -608,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. @@ -625,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()) { @@ -636,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)); } } @@ -652,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)); } @@ -678,7 +686,7 @@ InstructionValueSet& HloDataflowAnalysis::GetInstructionValueSet( } Status HloDataflowAnalysis::InitializeInstructionValueSets() { - for (const HloComputation* computation : module_->computations()) { + for (const HloComputation* computation : module_.computations()) { const CallGraphNode& call_graph_node = call_graph_->GetNode(computation); for (HloInstruction* instruction : computation->instructions()) { // Create an empty shape tree. @@ -779,9 +787,9 @@ Status HloDataflowAnalysis::InitializeInstructionValueSets() { /* static */ StatusOr> HloDataflowAnalysis::Run( - HloModule* module, bool ssa_form, bool bitcast_defines_value) { - VLOG(1) << "HloDataflowAnalysis::Run on module " << module->name(); - XLA_VLOG_LINES(2, module->ToString()); + const HloModule& module, bool ssa_form, bool bitcast_defines_value) { + VLOG(1) << "HloDataflowAnalysis::Run on module " << module.name(); + XLA_VLOG_LINES(2, module.ToString()); auto dataflow_analysis = WrapUnique( new HloDataflowAnalysis(module, ssa_form, bitcast_defines_value)); @@ -798,7 +806,7 @@ StatusOr> HloDataflowAnalysis::Run( // lookup is faster. std::vector> value_positions( dataflow_analysis->next_value_id_); - for (const HloComputation* computation : module->computations()) { + for (const HloComputation* computation : module.computations()) { for (HloInstruction* instruction : computation->instructions()) { for (const auto& pair : dataflow_analysis->GetInstructionValueSet(instruction)) { @@ -850,7 +858,7 @@ Status HloDataflowAnalysis::Verify() const { // For each value in each value set, verify that the value set's position // appears in the value's positions(). - for (const auto& computation : module_->computations()) { + for (const auto& computation : module_.computations()) { for (const auto& instruction : computation->instructions()) { for (const auto& pair : GetInstructionValueSet(instruction)) { const ShapeIndex& index = pair.first; diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h index 89d318188f0855c7924836a51cfe98d531e08cb4..7b8a74b096ff48733717e78ada5bb56a28caed72 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h @@ -60,7 +60,7 @@ class HloDataflowAnalysis { // a new HLO value in the analysis. If false then Bitcast forwards the // value of its operand. static StatusOr> Run( - HloModule* module, bool ssa_form = false, + const HloModule& module, bool ssa_form = false, bool bitcast_defines_value = false); // Returns true if 'instruction' defines an HLO value at the given shape index @@ -119,7 +119,7 @@ class HloDataflowAnalysis { string ToString() const; protected: - HloDataflowAnalysis(HloModule* module, bool ssa_form, + HloDataflowAnalysis(const HloModule& module, bool ssa_form, bool bitcast_defines_value = false); // Returns a new HloValue defined at the given instruction and shape index. @@ -180,7 +180,7 @@ class HloDataflowAnalysis { // Verify various invariants of the dataflow analysis. Status Verify() const; - HloModule* const module_; + const HloModule& module_; const bool ssa_form_; const bool bitcast_defines_value_; diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc index e714b2567fd1b3eab607a19f0bb7e3288150dc64..07f69b8e1339fed636e4eb54791941b85e09fd17 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc @@ -50,7 +50,7 @@ class HloDataflowAnalysisTest : public HloTestBase, bool bitcast_defines_value = false) { hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before dataflow analysis"); analysis_ = - HloDataflowAnalysis::Run(module_.get(), ssa_form, bitcast_defines_value) + HloDataflowAnalysis::Run(*module_, ssa_form, bitcast_defines_value) .ConsumeValueOrDie(); return *analysis_; } @@ -1602,11 +1602,17 @@ TEST_P(HloDataflowAnalysisTest, ConditionalWithIdentity) { EXPECT_THAT(analysis.GetValueDefinedAt(constant2).uses(), ElementsAre(HloUse{conditional, 2, {}})); - EXPECT_EQ(analysis.values().size(), 3); - EXPECT_FALSE(analysis.ValueIsDefinedAt(conditional)); - EXPECT_THAT(HloValuesAt(conditional), - UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), - analysis.GetValueDefinedAt(constant2))); + bool ssa_form = GetParam(); + if (ssa_form) { + EXPECT_EQ(analysis.values().size(), 4); + EXPECT_TRUE(analysis.ValueIsDefinedAt(conditional)); + } else { + EXPECT_EQ(analysis.values().size(), 3); + EXPECT_FALSE(analysis.ValueIsDefinedAt(conditional)); + EXPECT_THAT(HloValuesAt(conditional), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2))); + } } TEST_P(HloDataflowAnalysisTest, ConditionalTakingTupleOperand) { @@ -1713,11 +1719,17 @@ TEST_P(HloDataflowAnalysisTest, ConditionalTakingTupleOperand) { HloUse{true_x, 0, {}}, HloUse{true_y, 0, {}}, HloUse{false_x, 0, {}}, HloUse{false_y, 0, {}})); - EXPECT_EQ(analysis.values().size(), 6); - EXPECT_FALSE(analysis.ValueIsDefinedAt(conditional)); - EXPECT_THAT(HloValuesAt(conditional), - UnorderedElementsAre(analysis.GetValueDefinedAt(add), - analysis.GetValueDefinedAt(sub))); + bool ssa_form = GetParam(); + if (ssa_form) { + EXPECT_EQ(analysis.values().size(), 7); + EXPECT_TRUE(analysis.ValueIsDefinedAt(conditional)); + } else { + EXPECT_EQ(analysis.values().size(), 6); + EXPECT_FALSE(analysis.ValueIsDefinedAt(conditional)); + EXPECT_THAT(HloValuesAt(conditional), + UnorderedElementsAre(analysis.GetValueDefinedAt(add), + analysis.GetValueDefinedAt(sub))); + } } TEST_P(HloDataflowAnalysisTest, NestedConditionals) { @@ -1834,20 +1846,27 @@ TEST_P(HloDataflowAnalysisTest, NestedConditionals) { EXPECT_EQ(analysis.GetUniqueValueAt(false_operand_cond), analysis.GetValueDefinedAt(constant2)); - EXPECT_EQ(analysis.values().size(), 9); - EXPECT_FALSE(analysis.ValueIsDefinedAt(inner_conditional)); - EXPECT_FALSE(analysis.ValueIsDefinedAt(conditional)); - EXPECT_THAT( - HloValuesAt(inner_conditional), - UnorderedElementsAre( - analysis.GetValueDefinedAt(computation1->root_instruction()), - analysis.GetValueDefinedAt(computation2->root_instruction()))); - EXPECT_THAT( - HloValuesAt(conditional), - UnorderedElementsAre( - analysis.GetValueDefinedAt(computation1->root_instruction()), - analysis.GetValueDefinedAt(computation2->root_instruction()), - analysis.GetValueDefinedAt(computation3->root_instruction()))); + bool ssa_form = GetParam(); + if (ssa_form) { + EXPECT_EQ(analysis.values().size(), 11); + EXPECT_TRUE(analysis.ValueIsDefinedAt(inner_conditional)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(conditional)); + } else { + EXPECT_EQ(analysis.values().size(), 9); + EXPECT_FALSE(analysis.ValueIsDefinedAt(inner_conditional)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(conditional)); + EXPECT_THAT( + HloValuesAt(inner_conditional), + UnorderedElementsAre( + analysis.GetValueDefinedAt(computation1->root_instruction()), + analysis.GetValueDefinedAt(computation2->root_instruction()))); + EXPECT_THAT( + HloValuesAt(conditional), + UnorderedElementsAre( + analysis.GetValueDefinedAt(computation1->root_instruction()), + analysis.GetValueDefinedAt(computation2->root_instruction()), + analysis.GetValueDefinedAt(computation3->root_instruction()))); + } } INSTANTIATE_TEST_CASE_P(HloDataflowAnalysisInstantiation, diff --git a/tensorflow/compiler/xla/service/hlo_dce.cc b/tensorflow/compiler/xla/service/hlo_dce.cc index 1e5f0f797a13fd7e7ce1cc934387a274a74153bc..fcd723af146e2227b8661b1a4993f1338f7de389 100644 --- a/tensorflow/compiler/xla/service/hlo_dce.cc +++ b/tensorflow/compiler/xla/service/hlo_dce.cc @@ -40,7 +40,7 @@ StatusOr HloDCE::Run(HloModule* module) { VLOG(2) << "Before dce:"; XLA_VLOG_LINES(2, module->ToString()); - for (auto* computation : module->MakeNonfusionComputations()) { + for (auto* computation : module->MakeComputationPostOrder()) { std::unordered_set live_instructions; TF_RETURN_IF_ERROR(computation->root_instruction()->Accept( [&live_instructions](HloInstruction* instruction) { diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc index ab604064d5e34a1e21f3209cc4d27fc05a9ca057..693004d364114b1a25ce6b6791092665c861d13f 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc @@ -34,8 +34,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_query.h" #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/status.h" -#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/window_util.h" @@ -53,12 +51,22 @@ namespace xla { namespace { +using tensorflow::gtl::ArraySlice; +using tensorflow::gtl::FlatSet; +using tensorflow::gtl::optional; + template struct is_complex_t : public std::false_type {}; template <> struct is_complex_t : public std::true_type {}; +template +struct is_complex64_t : public std::false_type {}; + +template <> +struct is_complex64_t : public std::true_type {}; + template StatusOr> Compare(const Shape& shape, HloOpcode opcode, const Literal& lhs_literal, @@ -101,11 +109,10 @@ StatusOr> Compare(const Shape& shape, HloOpcode opcode, } auto result = Literal::CreateFromShape(shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { - return compare_op(lhs_literal.Get(multi_index), - rhs_literal.Get(multi_index)); - })); + TF_RETURN_IF_ERROR(result->Populate([&](ArraySlice multi_index) { + return compare_op(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index)); + })); return std::move(result); } @@ -132,11 +139,10 @@ StatusOr> Compare( } auto result = Literal::CreateFromShape(shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { - return compare_op(lhs_literal.Get(multi_index), - rhs_literal.Get(multi_index)); - })); + TF_RETURN_IF_ERROR(result->Populate([&](ArraySlice multi_index) { + return compare_op(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index)); + })); return std::move(result); } @@ -161,8 +167,8 @@ StatusOr> ElementWiseUnaryOpImpl( auto result = Literal::CreateFromShape(shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice multi_index) { return unary_op(operand_literal.Get(multi_index)); })); return std::move(result); @@ -174,7 +180,7 @@ StatusOr> ElementWiseUnaryOpImpl( // 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 ArraySlice& window_count_index, const std::function&)>& f) { const int64 rank = ShapeUtil::Rank(base_shape); DimensionVector window_index(rank); @@ -250,17 +256,37 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { template < typename NativeT, - typename std::enable_if::value || - is_complex_t::value>::type* = nullptr> + typename std::enable_if::value>::type* = nullptr> Status HandleAbs(HloInstruction* abs) { TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs], - ElementWiseUnaryOp(abs, [](ElementwiseT elem_operand) { + ElementWiseUnaryOp(abs, [](NativeT elem_operand) { return std::abs(elem_operand); })); return Status::OK(); } + template < + typename NativeT, + typename std::enable_if::value>::type* = nullptr> + Status HandleAbs(HloInstruction* abs) { + const Literal& operand_literal = + parent_->GetEvaluatedLiteralFor(abs->operand(0)); + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[abs], + (ElementWiseUnaryOpImpl( + abs, [](NativeT elem_operand) { return std::abs(elem_operand); }, + operand_literal))); + + return Status::OK(); + } + Status HandleAbs(HloInstruction* abs) override { + // If the operand is of C64 type, the return type of abs will be F32. + // However, ElementwiseT would still be the return type, F32, and thus + // specifying the ElementwiseT explicitly as C64 is needed below. + if (abs->operand(0)->shape().element_type() == C64) { + return HandleAbs(abs); + } return HandleAbs(abs); } @@ -308,13 +334,12 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { operand_to_broadcast.shape().dimensions(i)); } - return output->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { - for (int64 i = 0; i < broadcast->dimensions().size(); ++i) { - broadcast_indices[i] = multi_index[broadcast->dimensions(i)]; - } - return operand_to_broadcast.Get(broadcast_indices); - }); + return output->Populate([&](ArraySlice multi_index) { + for (int64 i = 0; i < broadcast->dimensions().size(); ++i) { + broadcast_indices[i] = multi_index[broadcast->dimensions(i)]; + } + return operand_to_broadcast.Get(broadcast_indices); + }); } template < @@ -588,14 +613,25 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } - template < - typename NativeT, - typename std::enable_if::value>::type* = nullptr> + template ::value>::type* = + nullptr> Status HandleMaximum(HloInstruction* maximum) { TF_ASSIGN_OR_RETURN( parent_->evaluated_[maximum], ElementWiseBinaryOp(maximum, [](ElementwiseT lhs, ElementwiseT rhs) { - return std::fmax(lhs, rhs); + return std::max(lhs, rhs); + })); + return Status::OK(); + } + + template ::value>::type* = nullptr> + Status HandleMaximum(HloInstruction* maximum) { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[maximum], + ElementWiseBinaryOp(maximum, [](ElementwiseT lhs, ElementwiseT rhs) { + return ((lhs >= rhs) || std::isnan(lhs)) ? lhs : rhs; })); return Status::OK(); } @@ -611,18 +647,30 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return HandleMaximum(maximum); } - template < - typename NativeT, - typename std::enable_if::value>::type* = nullptr> + template ::value>::type* = + nullptr> Status HandleMinimum(HloInstruction* minimum) { TF_ASSIGN_OR_RETURN(parent_->evaluated_[minimum], ElementWiseBinaryOp(minimum, [](ElementwiseT lhs_el, ElementwiseT rhs_el) { - return std::fmin(lhs_el, rhs_el); + return std::min(lhs_el, rhs_el); })); return Status::OK(); } + template ::value>::type* = nullptr> + Status HandleMinimum(HloInstruction* minimum) { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[minimum], + ElementWiseBinaryOp(minimum, [](ElementwiseT lhs_el, + ElementwiseT rhs_el) { + return ((lhs_el <= rhs_el) || std::isnan(lhs_el)) ? lhs_el : rhs_el; + })); + return Status::OK(); + } + template < typename NativeT, typename std::enable_if::value>::type* = nullptr> @@ -742,7 +790,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN( parent_->evaluated_[shl], ElementWiseBinaryOp(shl, [](NativeT lhs_elem, NativeT rhs_elem) { - return lhs_elem << rhs_elem; + return IsShiftOutOfBounds(rhs_elem) ? 0 + : (lhs_elem << rhs_elem); })); return Status::OK(); } @@ -767,8 +816,12 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN( parent_->evaluated_[shr], ElementWiseBinaryOp(shr, [](NativeT lhs_elem, NativeT rhs_elem) { - return static_cast(static_cast(lhs_elem) >> - rhs_elem); + SignedT lhs_signed = static_cast(lhs_elem); + if (IsShiftOutOfBounds(rhs_elem)) { + return lhs_signed < 0 ? static_cast(-1) : 0; + } else { + return lhs_signed >> rhs_elem; + } })); return Status::OK(); } @@ -794,6 +847,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 (IsShiftOutOfBounds(rhs_elem)) { + return static_cast(0); + } return static_cast(static_cast(lhs_elem) >> rhs_elem); })); @@ -818,7 +875,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { Status HandleClamp(HloInstruction* clamp) { std::function clamp_op = [](ElementwiseT low, ElementwiseT value, ElementwiseT high) { - return std::fmax(low, std::fmin(value, high)); + return std::fmin(high, std::fmax(value, low)); }; TF_ASSIGN_OR_RETURN( parent_->evaluated_[clamp], @@ -839,6 +896,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { } Status HandleSelect(HloInstruction* select) override { + CHECK(!ShapeUtil::IsScalar(select->operand(0)->shape())); CHECK(!ShapeUtil::IsTuple(select->shape())); std::function select_op = [](bool pred, ReturnT on_true, ReturnT on_false) { @@ -869,8 +927,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); auto result = Literal::CreateFromShape(result_shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice out_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice out_index) { std::vector from_index(out_index.begin(), out_index.end()); for (const int64 dim : reverse_dimensions) { from_index[dim] = result_shape.dimensions(dim) - 1 - out_index[dim]; @@ -945,7 +1003,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { DimensionVector rhs_index(rhs_rank); DimensionVector rhs_spatial_index(dnums.kernel_spatial_dimensions_size()); - auto func = [&](tensorflow::gtl::ArraySlice out_index) { + auto func = [&](ArraySlice out_index) { ElementwiseT result_val = static_cast(0); std::fill(lhs_index.begin(), lhs_index.end(), 0); @@ -1028,55 +1086,118 @@ 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()); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + + 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; + 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([&](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)) * @@ -1112,9 +1233,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { parent_->GetEvaluatedLiteralFor(pad->operand(1)).Get({}); auto result = Literal::CreateFromShape(pad->shape()); TF_RETURN_IF_ERROR(result->Populate( - [&scalar](tensorflow::gtl::ArraySlice multi_index) { - return scalar; - })); + [&scalar](ArraySlice multi_index) { return scalar; })); const Literal& evaluated_operand = parent_->GetEvaluatedLiteralFor(pad->operand(0)); @@ -1127,7 +1246,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { // corresponding index of the resulting padded literal. const PaddingConfig& pad_config = pad->padding_config(); - auto func = [&](const std::vector& input_index) { + auto func = [&](ArraySlice input_index) { for (auto i = 0; i < input_index.size(); ++i) { // Interior padding occurs logically before edge padding, so in the case // of negative edge padding elements are removed from the @@ -1277,9 +1396,9 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { auto result = Literal::CreateFromShape(map->shape()); - HloEvaluator embedded_evaluator; - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + HloEvaluator embedded_evaluator(parent_->max_loop_iterations_); + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice multi_index) { std::vector> arg_literals; arg_literals.reserve(operands.size()); @@ -1339,6 +1458,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; @@ -1364,7 +1488,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { Status HandleReduce(HloInstruction* reduce) override { auto arg = reduce->operand(0); auto init_value = reduce->operand(1); - tensorflow::gtl::ArraySlice dimensions(reduce->dimensions()); + ArraySlice dimensions(reduce->dimensions()); HloComputation* function = reduce->to_apply(); TF_RET_CHECK(ShapeUtil::Rank(reduce->shape()) == ShapeUtil::Rank(arg->shape()) - dimensions.size()); @@ -1407,10 +1531,10 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { } } - HloEvaluator embedded_evaluator; + HloEvaluator embedded_evaluator(parent_->max_loop_iterations_); // For each resulting dimension, calculate and assign computed value. - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice multi_index) { ReturnT result_val = init_scalar; std::vector base(arg_dimensions.size()); @@ -1418,7 +1542,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { base[result_to_arg_index[i]] = multi_index[i]; } - auto func = [&](const std::vector& input_index) { + auto func = [&](ArraySlice input_index) { auto curr_val = arg_literal.Get(input_index); // Evaluate computation with specified literal operands. @@ -1464,9 +1588,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { // Initialize result array with the init value. TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice output_index) { - return init_scalar; - })); + [&](ArraySlice output_index) { return init_scalar; })); std::vector window_dimension_sizes; for (const auto& window_dimension : window.dimensions()) { @@ -1483,7 +1605,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { int64 rank = ShapeUtil::Rank(operand_literal.shape()); - HloEvaluator embedded_evaluator; + HloEvaluator embedded_evaluator(parent_->max_loop_iterations_); DimensionVector source_index(rank); std::fill(source_index.begin(), source_index.end(), 0); @@ -1499,8 +1621,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { // 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; + optional selected_val; + optional> selected_index; IterateThroughWindow( window_shape, window, operand_literal.shape(), source_index, @@ -1515,11 +1637,11 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { Literal::CreateR0(*selected_val); const std::vector args = { - curr_val_literal.get(), selected_val_literal.get()}; + selected_val_literal.get(), curr_val_literal.get()}; std::unique_ptr computed_result = embedded_evaluator.Evaluate(*select, args) .ConsumeValueOrDie(); - bool selected = computed_result->Get({}); + bool selected = !computed_result->Get({}); if (selected) { selected_val = curr_val; selected_index = operand_index; @@ -1594,10 +1716,10 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { DimensionVector window_index(window.dimensions_size()); DimensionVector operand_index(ShapeUtil::Rank(operand_literal.shape())); - HloEvaluator embedded_evaluator; + HloEvaluator embedded_evaluator(parent_->max_loop_iterations_); // For each resulting dimension, calculate and assign computed value. - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice output_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice output_index) { ReturnT result_val = init_scalar; std::fill(window_index.begin(), window_index.end(), 0); @@ -1647,7 +1769,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { const int64 rank = ShapeUtil::Rank(operand->shape()); const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); - auto func = [&](tensorflow::gtl::ArraySlice out_index) { + auto func = [&](ArraySlice out_index) { DimensionVector operand_index(rank); for (int64 i = 0; i < rank; ++i) { operand_index[i] = @@ -1828,8 +1950,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { std::vector operand_indices(start.size()); auto result = Literal::CreateFromShape(result_shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice multi_index) { for (int64 i = 0; i < operand_indices.size(); ++i) { CHECK_GE(multi_index[i] + start[i], 0); // Mod is only used here to be consistent with the existing @@ -1849,17 +1971,26 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { StatusOr> DynamicUpdateSlice( const Literal& operand_literal, const Literal& update_literal, const Literal& start_indices_literal) { - auto start_indices_typed = start_indices_literal.data(); - const std::vector start(start_indices_typed.begin(), - start_indices_typed.end()); - auto result = operand_literal.CloneToUnique(); - std::vector result_index(ShapeUtil::Rank(result->shape()), 0); + auto start_indices_typed = start_indices_literal.data(); + const auto rank = ShapeUtil::Rank(result->shape()); + std::vector start(rank, 0); + for (int64 i = 0; i < rank; ++i) { + // All other implementations currently wrap-around the index, so this + // should do so as well. + start[i] = (start_indices_typed[i] % result->shape().dimensions(i)); + start[i] += (start[i] < 0) * result->shape().dimensions(i); + } + std::vector result_index(rank, 0); - auto func = [&](const std::vector& update_index) { + auto func = [&](ArraySlice update_index) { std::transform(update_index.begin(), update_index.end(), start.begin(), result_index.begin(), std::plus()); - + // Same as above, wrap-around only to match other implementations' + // semantics. + std::transform(result_index.begin(), result_index.end(), + result->shape().dimensions().begin(), result_index.begin(), + std::modulus()); result->Set(result_index, update_literal.Get(update_index)); return true; @@ -1912,8 +2043,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { auto result = Literal::CreateFromShape(shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice multi_index) { return ConvertBinaryFunction(binary_op)( lhs_literal.Get(multi_index), rhs_literal.Get(multi_index)); @@ -1950,8 +2081,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { auto result = Literal::CreateFromShape(shape); - TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + TF_RETURN_IF_ERROR( + result->Populate([&](ArraySlice multi_index) { return ternary_op(lhs_literal.Get(multi_index), rhs_literal.Get(multi_index), ehs_literal.Get(multi_index)); @@ -1960,26 +2091,35 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return std::move(result); } + template + static bool IsShiftOutOfBounds(NativeT rhs) { + typedef typename std::make_unsigned::type UnsignedT; + UnsignedT lhs_size_unsigned = sizeof(NativeT) * CHAR_BIT; + UnsignedT rhs_unsigned = static_cast(rhs); + return rhs_unsigned >= lhs_size_unsigned; + } + HloEvaluator* parent_; }; // class HloEvaluator::TypedVisitor -HloEvaluator::HloEvaluator() { +HloEvaluator::HloEvaluator(int64 max_loop_iterations) + : max_loop_iterations_(max_loop_iterations) { typed_visitors_[PRED] = MakeUnique>(this); typed_visitors_[U8] = MakeUnique>(this); typed_visitors_[U16] = MakeUnique([](HloInstruction*) { - return Unimplemented("HloEvaluator: unhandled primitive type: U16."); + return Unimplemented( + "HloEvaluator::TypedVisitor: unhandled primitive type: U16."); }); typed_visitors_[U32] = MakeUnique>(this); typed_visitors_[U64] = MakeUnique>(this); typed_visitors_[S8] = MakeUnique>(this); typed_visitors_[S16] = MakeUnique([](HloInstruction*) { - return Unimplemented("HloEvaluator: unhandled primitive type: S16."); + return Unimplemented( + "HloEvaluator::TypedVisitor: unhandled primitive type: S16."); }); 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); @@ -1989,18 +2129,20 @@ HloEvaluator::HloEvaluator() { // elementwise computations to be done in F32 and do BF16<->F32 conversion // around the input and the output of the computations. typed_visitors_[BF16] = MakeUnique>(this); + typed_visitors_[TUPLE] = MakeUnique([](HloInstruction*) { - return Unimplemented("HloEvaluator: unhandled primitive type: TUPLE."); + return Unimplemented( + "HloEvaluator::TypedVistor: unhandled primitive type: TUPLE."); }); typed_visitors_[OPAQUE] = MakeUnique([](HloInstruction*) { - return Unimplemented("HloEvaluator: unhandled primitive type: OPAQUE."); + return Unimplemented( + "HloEvaluator::TypedVisitor: unhandled primitive type: OPAQUE."); }); } template StatusOr> HloEvaluator::Evaluate( - const HloModule& module, - tensorflow::gtl::ArraySlice arg_literals) { + const HloModule& module, ArraySlice arg_literals) { XLA_VLOG_LINES(2, "HloEvaluator::Evaluate module:\n" + module.ToString()); evaluated_.clear(); @@ -2017,8 +2159,8 @@ StatusOr> HloEvaluator::Evaluate( template StatusOr> HloEvaluator::Evaluate( - const HloComputation& computation, - tensorflow::gtl::ArraySlice arg_literals) { + const HloComputation& computation, ArraySlice arg_literals) { + CHECK(computation.parent() != nullptr); XLA_VLOG_LINES( 2, "HloEvaluator::Evaluate computation:\n" + computation.ToString()); @@ -2034,8 +2176,7 @@ StatusOr> HloEvaluator::Evaluate( template StatusOr> HloEvaluator::Evaluate( - HloInstruction* instruction, - tensorflow::gtl::ArraySlice arg_literals) { + HloInstruction* instruction, ArraySlice arg_literals) { TF_RET_CHECK(hlo_query::AllOperandsAreParametersOrConstants(*instruction)); TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(instruction->shape())); @@ -2160,8 +2301,7 @@ Status HloEvaluator::HandleTranspose(HloInstruction* transpose) { } Status HloEvaluator::HandleConcatenate(HloInstruction* concatenate) { - tensorflow::gtl::ArraySlice operands( - concatenate->operands()); + ArraySlice operands(concatenate->operands()); // The result concatenate dimension is going to be the sum of all // concatenate dimensions of the operands taking part of the operation. const Shape& reference_shape = operands[0]->shape(); @@ -2333,6 +2473,349 @@ Status HloEvaluator::HandleTuple(HloInstruction* tuple) { return Status::OK(); } +// Returns an ShapeUtil::IndexIterationSpace that iterates over the output +// gather dimensions while keeping the rest of the output dimensions clamped to +// 0. +ShapeUtil::IndexIterationSpace IterationSpaceForOutputGatherIndices( + const Shape& output_shape, const GatherDimensionNumbers& dim_numbers) { + int64 output_rank = output_shape.dimensions_size(); + std::vector index_base(output_rank, 0); + std::vector index_count; + index_count.reserve(output_rank); + for (int64 i = 0; i < output_rank; i++) { + bool is_output_gather_dim = + !c_binary_search(dim_numbers.output_window_dims(), i); + index_count.push_back(is_output_gather_dim ? output_shape.dimensions(i) + : 1); + } + + return {std::move(index_base), std::move(index_count), + std::vector(output_rank, 1)}; +} + +// Return an ShapeUtil::IndexIterationSpace that iterates over the output window +// dimensions while keeping the rest of the output dimensions clamped to 0. +ShapeUtil::IndexIterationSpace IterationSpaceForOutputWindowIndices( + int64 output_rank, ArraySlice window_bounds, + const GatherDimensionNumbers& dim_numbers) { + std::vector index_base(output_rank, 0); + std::vector index_count(output_rank, 1); + int64 window_bounds_idx = 0; + for (int64 i = 0; i < output_rank; i++) { + bool is_output_window_dim = + c_binary_search(dim_numbers.output_window_dims(), i); + if (is_output_window_dim) { + while (c_binary_search(dim_numbers.elided_window_dims(), + window_bounds_idx)) { + window_bounds_idx++; + } + index_count[i] = window_bounds[window_bounds_idx++]; + } + } + + return {std::move(index_base), std::move(index_count), + std::vector(output_rank, 1)}; +} + +// This functor computes the contribution of gather_indices to an input index +// corresponding to an output index. That is, given an output index I, it picks +// out the gather output indices in I and uses them to look up a gather index, +// G, from the gather indices tensor, and expands G into the input space +// according to gather_dims_to_operand_dims. +class OutputGatherIndexToInputIndex { + public: + // The constructor does some setup work that is amortized across all + // iterations. + explicit OutputGatherIndexToInputIndex( + const GatherDimensionNumbers* dim_numbers, const Shape& input_shape, + const Shape& output_shape, const Literal* gather_indices) + : dim_numbers_(*dim_numbers), gather_indices_(*gather_indices) { + for (int64 i = 0; i < output_shape.dimensions_size(); i++) { + output_dim_is_gather_dims_.push_back( + !c_binary_search(dim_numbers_.output_window_dims(), i)); + } + + for (int64 i = 0; i < input_shape.dimensions_size(); i++) { + int64 index_of_input_dim_in_index_vector = + std::distance(dim_numbers_.gather_dims_to_operand_dims().begin(), + c_find(dim_numbers_.gather_dims_to_operand_dims(), i)); + if (index_of_input_dim_in_index_vector == + dim_numbers_.gather_dims_to_operand_dims_size()) { + input_dim_value_to_index_vector_.push_back(-1); + } else { + input_dim_value_to_index_vector_.push_back( + index_of_input_dim_in_index_vector); + } + } + + index_vector_index_.resize(gather_indices_.shape().dimensions_size()); + input_index_.resize(input_shape.dimensions_size()); + int64 index_vector_size = + gather_indices_.shape().dimensions(dim_numbers_.index_vector_dim()); + index_vector_.resize(index_vector_size); + } + + // Returns the contribution of gather_indices to the input index corresponding + // to output_index. See gather_inner_loop_body. + // + // This is conceptually a stateless transformation from output_index to the + // gather input index, but: + // + // - Instead of allocating memory to represent the gather input index on + // every invocation we reuse the same storage for the result + // (input_index_), mutating it in place. + // - Instead of allocating buffers for temporary values like + // index_vector_index_ and index_vector on every invocation, we reuse the + // same storage for all invocations. + // + // This returns an arrayslice into memory owned by the class. + StatusOr> operator()(ArraySlice output_index) { + PropagateOutputIndexGatherDimsToIndexVectorIndex(output_index); + TF_RETURN_IF_ERROR(FetchIndexVector()); + PropagateIndexVectorToInputIndex(); + return ArraySlice(input_index_); + } + + private: + // Propagates the gather index dimensions from the output index into + // index_vector_index_ by mutating index_vector_index_ in place. Does not + // update the dim_numbers.index_vector_dim() dimension -- that's the dimension + // we iterate over in FetchIndexVector. + void PropagateOutputIndexGatherDimsToIndexVectorIndex( + ArraySlice output_index) { + int64 index_vector_index_i = 0; + for (int64 i = 0, e = output_index.size(); i < e; i++) { + if (!output_dim_is_gather_dims_[i]) { + continue; + } + + if (index_vector_index_i == dim_numbers_.index_vector_dim()) { + index_vector_index_i++; + } + + index_vector_index_[index_vector_index_i++] = output_index[i]; + } + } + + // Populates index_vector_ by iterating over gather_indices_ according to + // index_vector_index_. + Status FetchIndexVector() { + int64 index_vector_dim = dim_numbers_.index_vector_dim(); + for (int64 i = 0, e = index_vector_.size(); i < e; i++) { + index_vector_index_[index_vector_dim] = i; + TF_ASSIGN_OR_RETURN(index_vector_[i], gather_indices_.GetIntegralAsS64( + index_vector_index_)); + } + return Status::OK(); + } + + // Populates input_index_. + void PropagateIndexVectorToInputIndex() { + for (int64 i = 0, e = input_index_.size(); i < e; i++) { + if (input_dim_value_to_index_vector_[i] != -1) { + input_index_[i] = index_vector_[input_dim_value_to_index_vector_[i]]; + } + + // If input_dim_value_to_index_vector_[i] == -1 then input_index_[i] + // remains 0, as set by the constructor. + } + } + + // input_dim_value_to_index_vector_[i] tells us how to compute dimension i of + // the input index from the index vector. See + // PropagateIndexVectorToInputIndex. + std::vector input_dim_value_to_index_vector_; + + // output_dim_is_gather_dims_[i] is true iff the output index i is a gather + // dimension. + std::vector output_dim_is_gather_dims_; + + // The buffer into which we construct an index into gather_indices_ to fetch + // the index vector. + std::vector index_vector_index_; + + // The index vector fetched from gather_indices_. + std::vector index_vector_; + + // The result computed by this functor. operator() returns an ArraySlice into + // this vector. + std::vector input_index_; + + const GatherDimensionNumbers& dim_numbers_; + const Literal& gather_indices_; +}; + +// This functor computes the contribution of the window indices in an output +// index to an input index. That is, given an output index I it picks out the +// output window indices in I and expands it into a window index into the input +// shape. +class OutputWindowIndexToInputIndex { + public: + // The constructor does some setup work that is amortized across all + // iterations. + explicit OutputWindowIndexToInputIndex( + const GatherDimensionNumbers& dim_numbers, const Shape& input_shape, + const Shape& output_shape) { + std::vector window_index_to_output_index; + int64 output_index_count = 0; + for (int64 i = 0; i < output_shape.dimensions_size(); i++) { + if (c_binary_search(dim_numbers.output_window_dims(), i)) { + window_index_to_output_index.push_back(output_index_count++); + } else { + output_index_count++; + } + } + + int64 window_dim_count = 0; + for (int64 i = 0; i < input_shape.dimensions_size(); i++) { + if (c_binary_search(dim_numbers.elided_window_dims(), i)) { + input_dim_value_to_output_index_.push_back(-1); + } else { + input_dim_value_to_output_index_.push_back( + window_index_to_output_index[window_dim_count++]); + } + } + + input_index_.resize(input_shape.dimensions_size()); + } + + // Returns the contribution of the window indices to the input index + // corresponding to output_index. See gather_inner_loop_body. + // + // This is conceptually a stateless transformation from output_index to the + // window input index, but instead of allocating memory to represent the + // gather input index on every invocation we reuse the same storage for the + // result (input_index_), mutating it in place. + // + // This returns an arrayslice into memory owned by the class. + StatusOr> operator()(ArraySlice output_index) { + PropagateOutputIndexWindowDimsToInputIndex(output_index); + return ArraySlice(input_index_); + } + + private: + // Propagates window dimensions from the output index to input_index_ by + // mutating input_index_ in place. + void PropagateOutputIndexWindowDimsToInputIndex( + ArraySlice output_index) { + for (int64 i = 0, e = input_index_.size(); i < e; i++) { + if (input_dim_value_to_output_index_[i] != -1) { + input_index_[i] = output_index[input_dim_value_to_output_index_[i]]; + } + + // If input_dim_value_to_index_vector_[i] == -1 then input_index_[i] + // remains 0, as set by the constructor. + } + } + + // input_dim_value_to_index_vector_[i] tells us how to compute dimension i of + // the input index from the output index. See + // PropagateOutputIndexToInputIndex. + std::vector input_dim_value_to_output_index_; + + // The result computed by this functor. operator() returns an ArraySlice into + // this vector. + std::vector input_index_; +}; + +// Rehapes the gather indices input to have a trailing degenerate `1` dimension +// if necessary. Hands over the ownership of the newly created literal (if +// there is one) to `reshaped_gather_indices`. +static StatusOr> ReshapedGatherIndices( + int64 index_vector_dim, const Literal& gather_indices, + std::unique_ptr* reshaped_gather_indices) { + if (gather_indices.shape().dimensions_size() != index_vector_dim) { + return std::cref(gather_indices); + } + + std::vector new_shape(gather_indices.shape().dimensions().begin(), + gather_indices.shape().dimensions().end()); + new_shape.push_back(1); + TF_ASSIGN_OR_RETURN(*reshaped_gather_indices, + gather_indices.Reshape(new_shape)); + return std::cref(**reshaped_gather_indices); +} + +Status HloEvaluator::HandleGather(HloInstruction* gather) { + std::unique_ptr result = Literal::CreateFromShape(gather->shape()); + const Shape& shape = gather->shape(); + const GatherDimensionNumbers& dim_numbers = + gather->gather_dimension_numbers(); + const Literal& operand = GetEvaluatedLiteralFor(gather->operand(0)); + std::unique_ptr reshaped_gather_indices; + TF_ASSIGN_OR_RETURN( + const Literal& gather_indices, + ReshapedGatherIndices(dim_numbers.index_vector_dim(), + GetEvaluatedLiteralFor(gather->operand(1)), + &reshaped_gather_indices)); + + // We iterate over the gather dimensions in the output shape in an outer loop + // nest, and iterate over the window dimensions in the output shape in an + // inner loop nest. + + ShapeUtil::IndexIterationSpace gather_indices_iteration_space = + IterationSpaceForOutputGatherIndices(shape, dim_numbers); + ShapeUtil::IndexIterationSpace window_indices_iteration_space = + IterationSpaceForOutputWindowIndices( + shape.dimensions_size(), gather->gather_window_bounds(), dim_numbers); + + // Scratch buffers that hold an index in the output shape and the + // corresponding index in the input shape. + std::vector input_index(operand.shape().dimensions_size()); + std::vector output_index(gather->shape().dimensions_size()); + + OutputGatherIndexToInputIndex output_gather_index_to_input_index( + &gather->gather_dimension_numbers(), /*input_shape=*/operand.shape(), + /*output_shape=*/shape, &gather_indices); + OutputWindowIndexToInputIndex output_window_index_to_input_index( + gather->gather_dimension_numbers(), /*input_shape=*/operand.shape(), + /*output_shape=*/shape); + + const Shape& operand_shape = operand.shape(); + + auto gather_inner_loop_body = + [&](ArraySlice output_window_index, + ArraySlice input_gather_index, + ArraySlice output_gather_index) -> StatusOr { + TF_ASSIGN_OR_RETURN( + ArraySlice input_window_index, + output_window_index_to_input_index(output_window_index)); + for (int i = 0, e = output_index.size(); i < e; i++) { + output_index[i] = output_gather_index[i] + output_window_index[i]; + DCHECK_LT(output_index[i], shape.dimensions(i)); + } + for (int i = 0, e = input_index.size(); i < e; i++) { + // TODO(b/74360564): We should implement whatever out of bounds behavior + // we decide for dynamic-slice here as well. + input_index[i] = (input_gather_index[i] + input_window_index[i]) % + operand_shape.dimensions(i); + if (input_index[i] < 0) { + input_index[i] += operand_shape.dimensions(i); + } + } + TF_RETURN_IF_ERROR( + result->CopyElementFrom(operand, input_index, output_index)); + return true; + }; + + auto gather_outer_loop_body = + [&](ArraySlice output_gather_index) -> StatusOr { + TF_ASSIGN_OR_RETURN( + ArraySlice input_gather_index, + output_gather_index_to_input_index(output_gather_index)); + TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus( + shape, window_indices_iteration_space, + std::bind(gather_inner_loop_body, std::placeholders::_1, + input_gather_index, output_gather_index))); + return true; + }; + + TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus( + shape, gather_indices_iteration_space, gather_outer_loop_body)); + evaluated_[gather] = std::move(result); + return Status::OK(); +} + Status HloEvaluator::HandleGetTupleElement(HloInstruction* get_tuple_element) { const auto result_shape = get_tuple_element->shape(); const int64 index = get_tuple_element->tuple_index(); @@ -2363,6 +2846,135 @@ Status HloEvaluator::HandleCopy(HloInstruction* copy) { return Status::OK(); } +Status HloEvaluator::HandleCall(HloInstruction* call) { + auto* computation = call->to_apply(); + auto operands = call->operands(); + + std::vector arg_literals; + arg_literals.reserve(operands.size()); + for (auto operand : operands) { + const Literal& arg_literal = GetEvaluatedLiteralFor(operand); + arg_literals.push_back(&arg_literal); + } + + HloEvaluator embedded_evaluator; + std::unique_ptr result = + embedded_evaluator.Evaluate(*computation, arg_literals) + .ConsumeValueOrDie(); + + evaluated_[call] = std::move(result); + return Status::OK(); +} + +Status HloEvaluator::HandleFusion(HloInstruction* fusion) { + // Attach cloned computation to an empty HLO module so the existing ones are + // not modified. + HloModule empty_hlo_module("EmptyModuleForFusion"); + auto cloned_fused_computation = + fusion->fused_instructions_computation()->Clone( + /*suffix=*/"clone_with_layout", &empty_hlo_module); + for (auto* instruction : cloned_fused_computation->instructions()) { + LayoutUtil::SetToDefaultLayout(instruction->mutable_shape()); + } + auto readded_computation = + empty_hlo_module.AddEntryComputation(std::move(cloned_fused_computation)); + + auto operands = fusion->operands(); + std::vector arg_literals; + arg_literals.reserve(operands.size()); + for (auto operand : operands) { + const Literal& arg_literal = GetEvaluatedLiteralFor(operand); + arg_literals.push_back(&arg_literal); + } + + HloEvaluator embedded_evaluator; + std::unique_ptr result = + embedded_evaluator + .Evaluate(*readded_computation, arg_literals) + .ConsumeValueOrDie(); + + evaluated_[fusion] = std::move(result); + return Status::OK(); +} + +Status HloEvaluator::HandleConditional(HloInstruction* conditional) { + const auto& pred = GetEvaluatedLiteralFor(conditional->operand(0)); + const auto& true_computation_arg = + GetEvaluatedLiteralFor(conditional->operand(1)); + const auto& false_computation_arg = + GetEvaluatedLiteralFor(conditional->operand(2)); + + auto* true_computation = conditional->true_computation(); + auto* false_computation = conditional->false_computation(); + + auto result = Literal::CreateFromShape(conditional->shape()); + HloEvaluator embedded_evaluator; + if (pred.Get({})) { + result = embedded_evaluator + .Evaluate(*true_computation, + {&true_computation_arg}) + .ConsumeValueOrDie(); + } else { + result = embedded_evaluator + .Evaluate(*false_computation, + {&false_computation_arg}) + .ConsumeValueOrDie(); + } + + evaluated_[conditional] = std::move(result); + return Status::OK(); +} + +Status HloEvaluator::HandleSelect(HloInstruction* select) { + const auto& pred = GetEvaluatedLiteralFor(select->operand(0)); + const auto& on_true = GetEvaluatedLiteralFor(select->operand(1)); + const auto& on_false = GetEvaluatedLiteralFor(select->operand(2)); + + // If predicate is of scalar type, no element-wise selection would be needed. + // This would also handle output array of tuple types as the DefaultAction + // would go through the TypedVisitor which doesn't handle tuples. + if (ShapeUtil::IsScalar(pred.shape())) { + if (pred.Get({})) { + evaluated_[select] = on_true.CloneToUnique(); + } else { + evaluated_[select] = on_false.CloneToUnique(); + } + return Status::OK(); + } + + return DefaultAction(select); +} + +Status HloEvaluator::HandleWhile(HloInstruction* while_hlo) { + HloComputation* cond_comp = while_hlo->while_condition(); + HloComputation* body_comp = while_hlo->while_body(); + // Initialize the loop carried valued with the input to the While instruction. + auto lcv = GetEvaluatedLiteralFor(while_hlo->operand(0)).CloneToUnique(); + bool keep_going = true; + int64 iteration_count = 0; + HloEvaluator cond_evaluator(max_loop_iterations_); + HloEvaluator loop_body_evaluator(max_loop_iterations_); + while (keep_going) { + if (max_loop_iterations_ >= 0 && iteration_count++ > max_loop_iterations_) { + return InvalidArgument("Loop %s exceeded loop iteration limit (%lld).", + while_hlo->name().c_str(), max_loop_iterations_); + } + TF_ASSIGN_OR_RETURN(auto cond_val, cond_evaluator.Evaluate( + *cond_comp, {lcv.get()})); + keep_going = cond_val->GetFirstElement(); + if (keep_going) { + TF_ASSIGN_OR_RETURN(auto body_val, loop_body_evaluator.Evaluate( + *body_comp, {lcv.get()})); + VLOG(3) << "Loop iteration result: " << body_val->ToString(); + lcv = std::move(body_val); + cond_evaluator.ResetVisitStates(); + loop_body_evaluator.ResetVisitStates(); + } + } + evaluated_[while_hlo] = std::move(lcv); + return Status::OK(); +} + Status HloEvaluator::Preprocess(HloInstruction* hlo) { VLOG(2) << "About to visit HLO: " << hlo->ToString(); return Status::OK(); @@ -2376,28 +2988,27 @@ Status HloEvaluator::Postprocess(HloInstruction* hlo) { // Explicit instantiation of templatized Evaluate* methods. // -template StatusOr> HloEvaluator::Evaluate< - const Literal*>(const HloModule& module, - tensorflow::gtl::ArraySlice arg_literals); +template StatusOr> +HloEvaluator::Evaluate(const HloModule& module, + ArraySlice arg_literals); template StatusOr> HloEvaluator::Evaluate>( - const HloModule& module, - tensorflow::gtl::ArraySlice> arg_literals); + const HloModule& module, ArraySlice> arg_literals); -template StatusOr> HloEvaluator::Evaluate< - const Literal*>(const HloComputation& computation, - tensorflow::gtl::ArraySlice arg_literals); +template StatusOr> +HloEvaluator::Evaluate(const HloComputation& computation, + ArraySlice arg_literals); template StatusOr> HloEvaluator::Evaluate>( const HloComputation& computation, - tensorflow::gtl::ArraySlice> arg_literals); + ArraySlice> arg_literals); -template StatusOr> HloEvaluator::Evaluate< - const Literal*>(HloInstruction* instruction, - tensorflow::gtl::ArraySlice arg_literals); +template StatusOr> +HloEvaluator::Evaluate(HloInstruction* instruction, + ArraySlice arg_literals); template StatusOr> HloEvaluator::Evaluate>( HloInstruction* instruction, - tensorflow::gtl::ArraySlice> arg_literals); + ArraySlice> arg_literals); } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h index 3b2b697e492a78a06a4e5ae6bf056ff8676f2ff5..c0dcee0c3e382f74de72a2b89f39e06f042e2b80 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator.h @@ -36,7 +36,10 @@ namespace xla { // This class is not thread-safe. class HloEvaluator : public DfsHloVisitorWithDefault { public: - HloEvaluator(); + // Only evaluate up to max_loop_iterations per while-loop execution if + // specified. + explicit HloEvaluator(int64 max_loop_iterations = -1); + // Evaluates an HLO module and an array of pointers to literals. // Returns the evaluated result as a literal if successful. // Precondition: The indices of arg_literals correspond to the parameter @@ -149,10 +152,22 @@ class HloEvaluator : public DfsHloVisitorWithDefault { Status HandleTuple(HloInstruction* tuple) override; + Status HandleGather(HloInstruction* gather) override; + Status HandleGetTupleElement(HloInstruction* get_tuple_element) override; Status HandleCopy(HloInstruction* copy) override; + Status HandleConditional(HloInstruction* conditional) override; + + Status HandleCall(HloInstruction* call) override; + + Status HandleFusion(HloInstruction* fusion) override; + + Status HandleWhile(HloInstruction* while_hlo) override; + + Status HandleSelect(HloInstruction* select) override; + private: // Returns the already-evaluated literal result for the instruction. // A Constant instruction is considered evaluated and its literal will be @@ -190,6 +205,9 @@ class HloEvaluator : public DfsHloVisitorWithDefault { // Must be cleared for each evaluation. std::vector arg_literals_; + // Max loop iterations to execute with no maximum if negative. + int64 max_loop_iterations_; + TF_DISALLOW_COPY_AND_ASSIGN(HloEvaluator); }; diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc index 97765d65909cee192f65069777f8f195081603b2..685cacd7f74c00789296dee16f0a6a94c35a4393 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc @@ -1729,6 +1729,207 @@ TEST_P(HloEvaluatorTest, EvaluateWithSubstitutionsWithConstantOperand) { *result.ValueOrDie()); } +TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherV1) { + const char* hlo_text = R"( +HloModule TensorFlowGatherV1 + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[2,3] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0}, + gather_dims_to_operand_dims={0}, + index_vector_dim=1, + window_bounds={1, 3} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR1({0, 2}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR2({{1, 2, 3}, {7, 8, 9}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherV2) { + const char* hlo_text = R"( +HloModule TensorFlowGatherV2 + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[3,2] gather(operand, indices), + output_window_dims={0}, + elided_window_dims={1}, + gather_dims_to_operand_dims={1}, + index_vector_dim=1, + window_bounds={3, 1} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR1({0, 2}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR2({{1, 3}, {4, 6}, {7, 9}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherMultipleBatchDims) { + const char* hlo_text = R"( +HloModule TensorFlowGatherMultipleBatchDims + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,3,2] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={1}, + gather_dims_to_operand_dims={1}, + index_vector_dim=2, + window_bounds={3, 1} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{0, 2}, {2, 1}}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR3( + {{{1, 3}, {4, 6}, {7, 9}}, {{3, 2}, {6, 5}, {9, 8}}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherNd) { + const char* hlo_text = R"( +HloModule TensorFlowGatherNd + +ENTRY main { + operand = s32[3,3,2] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,2] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0,1}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=1, + window_bounds={1,1,2} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // + {{-4, 4}, {-5, 5}, {-6, 6}}, // + {{-7, 7}, {-8, 8}, {-9, 9}}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{0, 0}, {1, 0}}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR2({{-1, 1}, {-4, 4}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, + EvaluateGather_TensorFlowGatherNdNonDefaultIndexVectorDim) { + const char* hlo_text = R"( +HloModule TensorFlowGatherNd + +ENTRY main { + operand = s32[3,3,2] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,2] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0,1}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=0, + window_bounds={1,1,2} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // + {{-4, 4}, {-5, 5}, {-6, 6}}, // + {{-7, 7}, {-8, 8}, {-9, 9}}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{0, 0}, {1, 0}}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR2({{-2, 2}, {-1, 1}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, EvaluateGather_DynamicSlice) { + const char* hlo_text = R"( +HloModule DynamicSlice + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[1,1] gather(operand, indices), + output_window_dims={0,1}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=0, + window_bounds={1,1} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR1({1, 1}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR2({{5}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, EvaluateGather_BatchDynamicSlice) { + const char* hlo_text = R"( +HloModule BatchDynamicSlice + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,1,1] gather(operand, indices), + output_window_dims={1,2}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=0, + window_bounds={1,1} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{2, 1}, {1, 1}}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR3({{{8}}, {{5}}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + +TEST_P(HloEvaluatorTest, EvaluateGather_ZeroDimBounds) { + const char* hlo_text = R"( +HloModule TensorFlowGatherV1 + +ENTRY main { + operand = s32[3,0] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[2,0] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0}, + gather_dims_to_operand_dims={0}, + index_vector_dim=1, + window_bounds={1, 0} +} +)"; + ParseAndVerifyModule(hlo_text); + std::unique_ptr operand = Literal::CreateR2({{}, {}, {}}); + std::unique_ptr gather_indices = Literal::CreateR1({0, 2}); + LiteralTestUtil::ExpectEqual( + *Literal::CreateR2({{}, {}}), + *Evaluate({operand.get(), gather_indices.get()})); +} + INSTANTIATE_TEST_CASE_P(HloEvaluatorTest_Instantiation, HloEvaluatorTest, ::testing::ValuesIn(use_bf16_params)); diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index 44fcd36370dcd0cf77601aa1cd2b92810947bd5f..1dc72355cf179e996caab4d6b52068dc99d02244 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -157,52 +157,60 @@ enum ColorScheme { kDashedBorder, }; +// Graphviz attributes/colors that make up a color scheme. +struct NodeColors { + const char* style; + const char* fill_color; + const char* stroke_color; + const char* font_color; +}; + +NodeColors NodeColorsForScheme(ColorScheme color) { + switch (color) { + case kBlue: + return NodeColors{"filled", "#bbdefb", "#8aacc8", "black"}; + case kBrown: + return NodeColors{"filled", "#bcaaa4", "#8c7b75", "black"}; + case kDarkBlue: + return NodeColors{"filled", "#1565c0", "#003c8f", "white"}; + case kDarkGreen: + return NodeColors{"filled", "#2e7d32", "#005005", "white"}; + case kDarkRed: + return NodeColors{"filled", "#b71c1c", "#7f0000", "white"}; + case kGray: + return NodeColors{"filled", "#cfd8dc", "#9ea7aa", "black"}; + case kGreen: + return NodeColors{"filled", "#c8e6c9", "#97b498", "black"}; + case kOrange: + return NodeColors{"filled", "#ffe0b2", "#cbae82", "black"}; + case kPurple: + return NodeColors{"filled", "#e1bee7", "#af8eb5", "black"}; + case kRed: + return NodeColors{"filled", "#ffcdd2", "#cb9ca1", "black"}; + case kWhite: + return NodeColors{"filled", "white", "black", "black"}; + case kYellow: + return NodeColors{"filled", "#fff9c4", "#cbc693", "black"}; + case kDashedBorder: + // "filled,dashed" looks the same as "dashed", since we have a white + // background. But we use "filled,dashed" so that when you hover over + // any part of the node (not just the text inside the node), our css + // :hover rule is triggered. + return NodeColors{"filled,dashed", "white", "#757575", "#757575"}; + } +} + // Given a ColorScheme, returns an attribute string for a node of that color. // Sets the node's style and fill/stroke/text colors. // // Colors are from https://material.io/color. string NodeColorAttributes(ColorScheme color) { - using std::make_tuple; - - const char *style, *fill_color, *stroke_color, *font_color; - std::tie(style, fill_color, stroke_color, font_color) = [color] { - switch (color) { - case kBlue: - return make_tuple("filled", "#bbdefb", "#8aacc8", "black"); - case kBrown: - return make_tuple("filled", "#bcaaa4", "#8c7b75", "black"); - case kDarkBlue: - return make_tuple("filled", "#1565c0", "#003c8f", "white"); - case kDarkGreen: - return make_tuple("filled", "#2e7d32", "#005005", "white"); - case kDarkRed: - return make_tuple("filled", "#b71c1c", "#7f0000", "white"); - case kGray: - return make_tuple("filled", "#cfd8dc", "#9ea7aa", "black"); - case kGreen: - return make_tuple("filled", "#c8e6c9", "#97b498", "black"); - case kOrange: - return make_tuple("filled", "#ffe0b2", "#cbae82", "black"); - case kPurple: - return make_tuple("filled", "#e1bee7", "#af8eb5", "black"); - case kRed: - return make_tuple("filled", "#ffcdd2", "#cb9ca1", "black"); - case kWhite: - return make_tuple("filled", "white", "black", "black"); - case kYellow: - return make_tuple("filled", "#fff9c4", "#cbc693", "black"); - case kDashedBorder: - // "filled,dashed" looks the same as "dashed", since we have a white - // background. But we use "filled,dashed" so that when you hover over - // any part of the node (not just the text inside the node), our css - // :hover rule is triggered. - return make_tuple("filled,dashed", "white", "#757575", "#757575"); - } - }(); + NodeColors node_colors = NodeColorsForScheme(color); return Printf( - R"(style="%s", fontcolor="%s", color="%s", fillcolor="%s")", style, - font_color, stroke_color, fill_color); + R"(style="%s", fontcolor="%s", color="%s", fillcolor="%s")", + node_colors.style, node_colors.font_color, node_colors.stroke_color, + node_colors.fill_color); } // Replaces <> with <>, so that this string is safe(er) for use in a @@ -604,11 +612,21 @@ tooltip = " "; StrAppend(&subcomp_label, "
", extra_info); } - // Subcomputation's fill/stroke color is light/dark red/gray, depending on - // whether or not the subcomputation's fusion node is highlighted. bool highlight = filter_.Highlight(parent_instr); - const char* fillcolor = highlight ? "#ffcdd2" : "#f5f5f5"; - const char* strokecolor = highlight ? "#b71c1c" : "#c2c2c2"; + const char* fillcolor; + const char* strokecolor; + if (debug_options_.xla_hlo_graph_sharding_color() && !highlight) { + // Use the sharding color, if the node isn't highlighted. + NodeColors node_colors = + NodeColorsForScheme(GetInstructionColor(parent_instr)); + fillcolor = node_colors.fill_color; + strokecolor = node_colors.stroke_color; + } else { + // Subcomputation's fill/stroke color is light/dark red/gray, depending on + // whether or not the subcomputation's fusion node is highlighted. + fillcolor = highlight ? "#ffcdd2" : "#f5f5f5"; + strokecolor = highlight ? "#b71c1c" : "#c2c2c2"; + } style = Printf(R"(style="rounded,filled,bold"; fillcolor="%s"; color="%s;")", fillcolor, strokecolor); @@ -782,6 +800,14 @@ string HloDotDumper::GetInstructionNodeInlinedOperands( auto stringify_constant = [](const HloInstruction* constant) { const auto& shape = constant->shape(); + // If the shape has a dimension of size zero, print it as e.g. + // "{} (f32[42, 0, 10])". The alternative, calling Literal::ToString(), + // enumerates all of its empty dimensions (e.g. "{ { {}, {} }, ..."), which + // is just noise. + if (ShapeUtil::HasZeroElements(shape)) { + return Printf("{} (%s)", ShapeUtil::HumanString(constant->shape())); + } + // Print the literal value of constants with <= K elements. optional elem_count; if (!ShapeUtil::IsOpaque(shape) && !ShapeUtil::IsTuple(shape)) { @@ -940,6 +966,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kConcatenate: case HloOpcode::kCopy: case HloOpcode::kDynamicSlice: + case HloOpcode::kGather: case HloOpcode::kPad: case HloOpcode::kReshape: case HloOpcode::kReverse: @@ -988,6 +1015,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kCall: case HloOpcode::kConditional: case HloOpcode::kCustomCall: + case HloOpcode::kHostCompute: case HloOpcode::kWhile: return kDarkGreen; case HloOpcode::kConstant: diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index a889c35aeb297bd118c40ced2dd9539957dce67a..d33add23d07b52cb56e4b212a29b415259af7694 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -182,6 +182,7 @@ StatusOr> HloInstruction::CreateFromProto( /* static */ std::unique_ptr HloInstruction::CreateGetTupleElement(const Shape& shape, HloInstruction* operand, int64 index) { + CHECK(ShapeUtil::IsTuple(operand->shape())); auto instruction = WrapUnique(new HloInstruction(HloOpcode::kGetTupleElement, shape)); instruction->tuple_index_ = index; @@ -763,16 +764,13 @@ HloInstruction::CreateBroadcastSequence( return instruction; } -// We put the fusion kind into the instruction's name for transpose-dot and -// backward-conv fusions, since those fusions are really just describing a type -// of dot/conv rather than generating a novel computation. +// 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"; - case HloInstruction::FusionKind::kConvBackwardInput: - case HloInstruction::FusionKind::kConvBackwardFilter: - return "conv_fusion"; default: return "fusion"; } @@ -804,16 +802,20 @@ static string FusionNodeName(HloInstruction::FusionKind fusion_kind) { 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; +HloInstruction* HloInstruction::AddFusionOperand(HloInstruction* new_operand) { + CHECK_EQ(opcode(), HloOpcode::kFusion); + CHECK_EQ(operand_count(), + fused_instructions_computation()->parameter_instructions().size()); + const int64 param_no = operand_count(); + // Name the parameter after the instruction it represents in the outer + // (non-fusion) computation. + string param_name = StrCat(new_operand->name(), ".param_", param_no); + HloInstruction* fused_parameter = + fused_instructions_computation()->AddParameter( + HloInstruction::CreateParameter(param_no, new_operand->shape(), + param_name)); + AppendOperand(new_operand); + return fused_parameter; } void HloInstruction::MergeFusionInstruction( @@ -1008,13 +1010,7 @@ HloInstruction* HloInstruction::CloneAndFuseInternal( // Clone's operand was not already an operand of the fusion // instruction. Add it as an operand and add a corresponding fused // parameter instruction. - int64 param_no = fused_parameters.size(); - // Name the parameter after the instruction it represents in the outer - // (non-fusion) computation. - string param_name = StrCat(operand->name(), ".param_", param_no); - fused_param = fused_instructions_computation()->AddParameter( - CreateParameter(param_no, operand->shape(), param_name)); - AppendOperand(operand); + fused_param = AddFusionOperand(operand); } TF_CHECK_OK(clone->ReplaceOperandWith(operand_num, fused_param)); } @@ -1099,6 +1095,7 @@ bool HloInstruction::HasSideEffect() const { case HloOpcode::kInfeed: case HloOpcode::kOutfeed: case HloOpcode::kTrace: + case HloOpcode::kHostCompute: return true; default: { // Check if any of the called computations has a side effect. @@ -1136,6 +1133,19 @@ bool HloInstruction::HasSideEffect() const { return instruction; } +/* static */ std::unique_ptr HloInstruction::CreateHostCompute( + const Shape& shape, tensorflow::gtl::ArraySlice operands, + tensorflow::StringPiece channel_name, const int64 cost_estimate_ns) { + std::unique_ptr instruction = + WrapUnique(new HloInstruction(HloOpcode::kHostCompute, shape)); + for (auto operand : operands) { + instruction->AppendOperand(operand); + } + instruction->channel_name_ = channel_name.ToString(); + instruction->cost_estimate_ns_ = cost_estimate_ns; + return instruction; +} + /* static */ std::unique_ptr HloInstruction::CreateTuple( tensorflow::gtl::ArraySlice elements) { std::vector element_shapes; @@ -1146,6 +1156,40 @@ bool HloInstruction::HasSideEffect() const { return CreateVariadic(tuple_shape, HloOpcode::kTuple, elements); } +/* static */ std::unique_ptr HloInstruction::CreateGather( + const Shape& shape, HloInstruction* operand, HloInstruction* gather_indices, + const GatherDimensionNumbers& gather_dim_numbers, + tensorflow::gtl::ArraySlice window_bounds) { + std::unique_ptr instruction = + WrapUnique(new HloInstruction(HloOpcode::kGather, shape)); + instruction->AppendOperand(operand); + instruction->AppendOperand(gather_indices); + instruction->gather_dimension_numbers_ = + MakeUnique(gather_dim_numbers); + c_copy(window_bounds, std::back_inserter(instruction->gather_window_bounds_)); + return instruction; +} + +/* static */ GatherDimensionNumbers HloInstruction::MakeGatherDimNumbers( + tensorflow::gtl::ArraySlice output_window_dims, + tensorflow::gtl::ArraySlice elided_window_dims, + tensorflow::gtl::ArraySlice gather_dims_to_operand_dims, + int64 index_vector_dim) { + GatherDimensionNumbers gather_dim_numbers; + for (int64 output_window_dim : output_window_dims) { + gather_dim_numbers.add_output_window_dims(output_window_dim); + } + for (int64 elided_window_dim : elided_window_dims) { + gather_dim_numbers.add_elided_window_dims(elided_window_dim); + } + for (int64 gather_dim_to_input_dim : gather_dims_to_operand_dims) { + gather_dim_numbers.add_gather_dims_to_operand_dims(gather_dim_to_input_dim); + } + + gather_dim_numbers.set_index_vector_dim(index_vector_dim); + return gather_dim_numbers; +} + std::unique_ptr HloInstruction::CloneWithNewOperands( const Shape& shape, tensorflow::gtl::ArraySlice new_operands, @@ -1227,6 +1271,10 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kCustomCall: clone = CreateCustomCall(shape, new_operands, custom_call_target_); break; + case HloOpcode::kHostCompute: + clone = CreateHostCompute(shape, new_operands, channel_name_, + cost_estimate_ns_); + break; case HloOpcode::kConcatenate: clone = CreateConcatenate(shape, new_operands, dimensions(0)); break; @@ -1376,12 +1424,19 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( break; case HloOpcode::kRecv: CHECK_EQ(new_operands.size(), 0); - clone = CreateRecv(shape, channel_id()); + // The shape is a tuple, but CreateRecv() wants the raw data shape. + clone = + CreateRecv(ShapeUtil::GetTupleElementShape(shape, 0), channel_id()); break; case HloOpcode::kRecvDone: CHECK_EQ(new_operands.size(), 1); clone = CreateRecvDone(new_operands[0]); break; + case HloOpcode::kGather: + CHECK_EQ(new_operands.size(), 2); + clone = CreateGather(shape, new_operands[0], new_operands[1], + *gather_dimension_numbers_, gather_window_bounds_); + break; case HloOpcode::kTrace: LOG(FATAL) << "Not yet implemented, clone: " << HloOpcodeString(opcode_); } @@ -1627,7 +1682,8 @@ bool HloInstruction::HasConstantOperand() const { bool HloInstruction::IdenticalSlowPath( const HloInstruction& other, const std::function& - eq_computations) const { + 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 @@ -1675,8 +1731,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: @@ -1686,7 +1746,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: @@ -1720,6 +1780,11 @@ bool HloInstruction::IdenticalSlowPath( return protobuf_util::ProtobufEquals(dot_dimension_numbers(), other.dot_dimension_numbers()); + case HloOpcode::kGather: + return protobuf_util::ProtobufEquals(gather_dimension_numbers(), + other.gather_dimension_numbers()) && + gather_window_bounds() == other.gather_window_bounds(); + // FFT has various types & lengths. case HloOpcode::kFft: return fft_type() == other.fft_type() && @@ -1742,18 +1807,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(); @@ -1767,10 +1832,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()); @@ -1790,6 +1855,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kRecvDone: case HloOpcode::kSend: case HloOpcode::kSendDone: + case HloOpcode::kHostCompute: return false; } } @@ -1815,7 +1881,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()); @@ -1838,8 +1905,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; @@ -2149,6 +2216,11 @@ std::vector HloInstruction::ExtraAttributesToString( if (dot_dimension_numbers_ != nullptr) { extra.push_back(DotDimensionNumbersToString()); } + if (gather_dimension_numbers_ != nullptr) { + extra.push_back(GatherDimensionNumbersToString()); + extra.push_back( + StrCat("window_bounds={", Join(gather_window_bounds(), ","), "}")); + } if (opcode() == HloOpcode::kFft) { extra.push_back(StrCat("fft_type=", FftType_Name(fft_type()))); extra.push_back(StrCat("fft_length={", Join(fft_length(), ","), "}")); @@ -2280,6 +2352,14 @@ HloInstructionProto HloInstruction::ToProto() const { if (dot_dimension_numbers_ != nullptr) { *proto.mutable_dot_dimension_numbers() = *dot_dimension_numbers_; } + if (gather_dimension_numbers_ != nullptr) { + *proto.mutable_gather_dimension_numbers() = *gather_dimension_numbers_; + } + if (opcode() == HloOpcode::kGather) { + for (int64 bound : gather_window_bounds()) { + proto.add_gather_window_bounds(bound); + } + } for (int i = 0; i < slice_starts_.size(); ++i) { auto* slice_dimension = proto.add_slice_dimensions(); slice_dimension->set_start(slice_starts_[i]); @@ -2318,7 +2398,7 @@ string HloInstruction::ToCategory() const { return "data formatting"; } - auto conv_category = [&] { + if (opcode() == HloOpcode::kConvolution) { string category = "convolution"; if (window_util::HasBaseDilation(window())) { category += " base-dilated"; @@ -2327,10 +2407,6 @@ string HloInstruction::ToCategory() const { category += " window-dilated"; } return category; - }; - - if (opcode() == HloOpcode::kConvolution) { - return conv_category(); } // Give transpose-dot and backwards-conv fusions the categories "dot" and @@ -2348,9 +2424,6 @@ string HloInstruction::ToCategory() const { return "output fusion"; case FusionKind::kTransposeDot: return "dot"; - case FusionKind::kConvBackwardFilter: - case FusionKind::kConvBackwardInput: - return conv_category(); case FusionKind::kCustom: return "custom fusion"; } @@ -2581,6 +2654,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleInfeed(this); case HloOpcode::kOutfeed: return visitor->HandleOutfeed(this); + case HloOpcode::kHostCompute: + return visitor->HandleHostCompute(this); case HloOpcode::kRng: return visitor->HandleRng(this); case HloOpcode::kWhile: @@ -2601,13 +2676,17 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleSend(this); case HloOpcode::kSendDone: return visitor->HandleSendDone(this); + case HloOpcode::kGather: + return visitor->HandleGather(this); // These opcodes are not handled here. case HloOpcode::kTrace: break; } - return Unimplemented("unhandled HloOpcode for DfsHloVisitor: %s", - HloOpcodeString(opcode_).c_str()); + return InternalError( + "Unhandled HloOpcode for DfsHloVisitor: %s. This should not happen - " + "please file a bug for XLA.", + HloOpcodeString(opcode_).c_str()); } // Explicit instantiations. @@ -3125,10 +3204,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"; } @@ -3148,12 +3223,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; } @@ -3261,7 +3330,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; } @@ -3288,6 +3363,26 @@ string HloInstruction::DotDimensionNumbersToString() const { return Join(result, ", "); } +string HloInstruction::GatherDimensionNumbersToString() const { + CHECK_NE(gather_dimension_numbers_.get(), nullptr); + string output_window_dims = + StrCat("output_window_dims={", + Join(gather_dimension_numbers_->output_window_dims(), ","), "}"); + string elided_window_dims = + StrCat("elided_window_dims={", + Join(gather_dimension_numbers_->elided_window_dims(), ","), "}"); + string gather_dims_to_operand_dims = StrCat( + "gather_dims_to_operand_dims={", + Join(gather_dimension_numbers_->gather_dims_to_operand_dims(), ","), "}"); + string index_vector_dim = StrCat( + "index_vector_dim=", gather_dimension_numbers_->index_vector_dim()); + + return Join>( + {output_window_dims, elided_window_dims, gather_dims_to_operand_dims, + index_vector_dim}, + ", "); +} + bool HloInstruction::CouldBeBitcast() const { switch (opcode_) { case HloOpcode::kTranspose: diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 5e89dc79bea81e650331e320f7836fdde90b2a53..e4c86214c2014095b2e171ff10691e1221574cb7 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -162,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(); @@ -454,6 +451,12 @@ class HloInstruction { HloInstruction* true_computation_arg, HloComputation* true_computation, HloInstruction* false_computation_arg, HloComputation* false_computation); + static std::unique_ptr CreateGather( + const Shape& shape, HloInstruction* operand, + HloInstruction* gather_indices, + const GatherDimensionNumbers& gather_dim_numbers, + tensorflow::gtl::ArraySlice window_bounds); + // Creates a fusion instruction. A fusion instruction contains one or more // fused instructions forming an expression with a single root // "fused_root". Additional instructions can be added to the fusion @@ -466,14 +469,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( @@ -486,6 +481,12 @@ class HloInstruction { const Shape& shape, tensorflow::gtl::ArraySlice operands, tensorflow::StringPiece custom_call_target); + // Creates a HostCompute instruction, which records host-side control and + // data dependencies for use in instruction scheduling. + static std::unique_ptr CreateHostCompute( + const Shape& shape, tensorflow::gtl::ArraySlice operands, + tensorflow::StringPiece channel_name, const int64 cost_estimate_ns); + // Creates a tuple instruction with the given elements. This is a convenience // wrapper around CreateVariadic. static std::unique_ptr CreateTuple( @@ -497,6 +498,13 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, tensorflow::gtl::ArraySlice dimensions); + // Creates an instance of GatherDimensionNumbers. + static GatherDimensionNumbers MakeGatherDimNumbers( + tensorflow::gtl::ArraySlice output_window_dims, + tensorflow::gtl::ArraySlice elided_window_dims, + tensorflow::gtl::ArraySlice gather_dims_to_operand_dims, + int64 index_vector_dim); + // Returns the opcode for this instruction. HloOpcode opcode() const { return opcode_; } @@ -565,27 +573,33 @@ 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, const std::function& eq_operands = std::equal_to(), const std::function& - eq_computations = std::equal_to()) const { + 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. + // 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) { @@ -594,7 +608,7 @@ class HloInstruction { } } - return IdenticalSlowPath(other, eq_computations); + return IdenticalSlowPath(other, eq_computations, eq_shapes); } // Returns whether the instruction has a constant operand. @@ -772,6 +786,10 @@ class HloInstruction { // // (We express the default options using an overload rather than a default // param because gdb ignores default params, but does resolve overloads.) + // + // TODO(b/73348663): Make ToString() adaptive to the size of the string by + // default, backing off on providing full information for very large strings, + // or provide a different name for a ToString-like function that does that. string ToString() const { return ToString(HloPrintOptions()); } string ToString(const HloPrintOptions& options) const; @@ -807,6 +825,12 @@ class HloInstruction { // Precondition: opcode() == HloOpcode::kSend or HloOpcode::kRecv int64 channel_id() const { return channel_id_; } + // Returns the channel name associated with the instruction. The name is + // used to identify host Send/Recv operations. + // + // Precondition: opcode() == HloOpcode::kHostCompute + string channel_name() const { return channel_name_; } + // Returns feature_index field associated with the instruction. The index // represents the index of the feature dimension. // @@ -885,8 +909,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 { @@ -919,6 +943,9 @@ class HloInstruction { // Return true if this operator has a sharding assigned. bool has_sharding() const { return sharding_ != nullptr; } + // Adds a new operand the fusion instruction. + HloInstruction* AddFusionOperand(HloInstruction* new_operand); + // Merges the fused instructions from 'instruction_to_merge' into the // fused instruction set of 'this', updating operands as necessary. // @@ -1052,13 +1079,23 @@ 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_; @@ -1081,6 +1118,19 @@ class HloInstruction { // Returns the dump string of the dot dimension numbers. string DotDimensionNumbersToString() const; + const GatherDimensionNumbers& gather_dimension_numbers() const { + CHECK(gather_dimension_numbers_ != nullptr); + return *gather_dimension_numbers_; + } + + tensorflow::gtl::ArraySlice gather_window_bounds() const { + CHECK_EQ(opcode(), HloOpcode::kGather); + return gather_window_bounds_; + } + + // Returns the dump string of the gather dimension numbers. + string GatherDimensionNumbersToString() const; + // Returns the random distribution for this rng node. // // Precondition: opcode() == HloOpcode::kRng @@ -1233,10 +1283,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, const std::function& - eq_computations) const; + eq_computations, + const std::function& eq_shapes) const; // Creates an n-ary elementwise operation. static std::unique_ptr CreateNary( @@ -1341,6 +1395,9 @@ class HloInstruction { // Describes the dimension numbers used for a dot. std::unique_ptr dot_dimension_numbers_; + std::unique_ptr gather_dimension_numbers_; + std::vector gather_window_bounds_; + // Describes FFT type for an FFT instruction. FftType fft_type_ = FftType::FFT; @@ -1379,6 +1436,12 @@ class HloInstruction { // Name of a global symbol to call, only present for kCustomCall. string custom_call_target_; + // Name to use for host send/recv channels, only present for kHostCompute. + string channel_name_; + + // Estimate of the duration of a host computation in nanoseconds. + int64 cost_estimate_ns_; + // Computations called by this instruction. std::vector called_computations_; diff --git a/tensorflow/compiler/xla/service/hlo_instruction_test.cc b/tensorflow/compiler/xla/service/hlo_instruction_test.cc index 1038ab555567aa654342d59e02efaf844f2b95ba..f2980d309d01fdf3b3e601bc260a0ad0895b3064 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction_test.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction_test.cc @@ -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) { @@ -1182,5 +1271,77 @@ TEST_F(HloInstructionTest, Stringification) { "true_computation=%TransposeDot, false_computation=%TransposeDot"); } +TEST_F(HloInstructionTest, StringifyGather_0) { + Shape input_tensor_shape = ShapeUtil::MakeShape(F32, {50, 49, 48, 47, 46}); + Shape gather_indices_tensor_shape = + ShapeUtil::MakeShape(S64, {10, 9, 8, 7, 5}); + Shape gather_result_shape = + ShapeUtil::MakeShape(F32, {10, 9, 8, 7, 30, 29, 28, 27, 26}); + + HloComputation::Builder builder("Gather"); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_tensor_shape, "input_tensor")); + HloInstruction* gather_indices = + builder.AddInstruction(HloInstruction::CreateParameter( + 1, gather_indices_tensor_shape, "gather_indices")); + + HloInstruction* gather_instruction = + builder.AddInstruction(HloInstruction::CreateGather( + gather_result_shape, input, gather_indices, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26})); + + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); + + EXPECT_EQ(gather_instruction->ToString(), + "%gather = f32[10,9,8,7,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} " + "gather(f32[50,49,48,47,46]{4,3,2,1,0} %input_tensor, " + "s64[10,9,8,7,5]{4,3,2,1,0} %gather_indices), " + "output_window_dims={4,5,6,7,8}, elided_window_dims={}, " + "gather_dims_to_operand_dims={0,1,2,3,4}, " + "index_vector_dim=4, window_bounds={30,29,28,27,26}"); +} + +TEST_F(HloInstructionTest, StringifyGather_1) { + Shape input_tensor_shape = ShapeUtil::MakeShape(F32, {50, 49, 48, 47, 46}); + Shape gather_indices_tensor_shape = + ShapeUtil::MakeShape(S64, {10, 9, 5, 7, 6}); + Shape gather_result_shape = + ShapeUtil::MakeShape(F32, {10, 9, 7, 6, 30, 29, 28, 27, 26}); + + HloComputation::Builder builder("Gather"); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_tensor_shape, "input_tensor")); + HloInstruction* gather_indices = + builder.AddInstruction(HloInstruction::CreateParameter( + 1, gather_indices_tensor_shape, "gather_indices")); + + HloInstruction* gather_instruction = + builder.AddInstruction(HloInstruction::CreateGather( + gather_result_shape, input, gather_indices, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/2), + /*window_bounds=*/{30, 29, 28, 27, 26})); + + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); + + EXPECT_EQ(gather_instruction->ToString(), + "%gather = f32[10,9,7,6,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} " + "gather(f32[50,49,48,47,46]{4,3,2,1,0} %input_tensor, " + "s64[10,9,5,7,6]{4,3,2,1,0} %gather_indices), " + "output_window_dims={4,5,6,7,8}, elided_window_dims={}, " + "gather_dims_to_operand_dims={0,1,2,3,4}, " + "index_vector_dim=2, window_bounds={30,29,28,27,26}"); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index 60270b0595dcfca8f1fcea5ab0914428880f35b5..cdea3d597824d155241a544d226aa18d3b0b0274 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -145,6 +145,21 @@ void HloModule::ReplaceComputations( } break; } + case HloOpcode::kConditional: { + HloComputation* new_true_computation = + tensorflow::gtl::FindWithDefault( + replacements, instruction->true_computation(), nullptr); + if (new_true_computation != nullptr) { + instruction->set_true_computation(new_true_computation); + } + HloComputation* new_false_computation = + tensorflow::gtl::FindWithDefault( + replacements, instruction->false_computation(), nullptr); + if (new_false_computation != nullptr) { + instruction->set_false_computation(new_false_computation); + } + break; + } case HloOpcode::kSelectAndScatter: { HloComputation* new_select = tensorflow::gtl::FindWithDefault( replacements, instruction->select(), nullptr); @@ -198,74 +213,23 @@ HloModuleProto HloModule::ToProto() const { continue; } HloComputationProto computation_proto = computation->ToProto(); + if (computation->name() == entry_computation_->name()) { + *proto.mutable_program_shape() = computation_proto.program_shape(); + } proto.add_computations()->Swap(&computation_proto); } return proto; } -namespace { - -// Construct a ProgramShape matching the shape of the parameters and root of the -// given module's entry computation. -StatusOr ProgramShapeFromProto(const HloModuleProto& module) { - const HloComputationProto* entry_computation = nullptr; - for (const HloComputationProto& computation : module.computations()) { - if (computation.name() == module.entry_computation_name()) { - entry_computation = &computation; - break; - } - } - TF_RET_CHECK(entry_computation != nullptr) - << "No computation with entry computation name" - << module.entry_computation_name(); - - tensorflow::gtl::FlatMap> parameters; - const HloInstructionProto* root = nullptr; - for (const HloInstructionProto& instruction : - entry_computation->instructions()) { - if (instruction.name() == entry_computation->root_name()) { - TF_RET_CHECK(root == nullptr) << "Entry computation has more than " - "one instruction with (root) name " - << instruction.name(); - root = &instruction; - } - if (instruction.opcode() == HloOpcodeString(HloOpcode::kParameter)) { - TF_RET_CHECK(!ContainsKey(parameters, instruction.parameter_number())) - << "Entry computation has more than one parameter instruction " - "with parameter number " - << instruction.parameter_number(); - parameters[instruction.parameter_number()] = {instruction.name(), - &instruction.shape()}; - } - } - TF_RET_CHECK(root != nullptr) - << "Entry computation is missing root instruction named " - << entry_computation->root_name(); - - ProgramShape program_shape; - *program_shape.mutable_result() = root->shape(); - for (int64 i = 0; i < parameters.size(); ++i) { - TF_RET_CHECK(ContainsKey(parameters, i)) - << "Entry computation missing parameter number " << i; - const string& name = parameters.at(i).first; - const Shape& shape = *parameters.at(i).second; - *program_shape.add_parameters() = shape; - program_shape.add_parameter_names(name); - } - - return std::move(program_shape); -} - -} // namespace - /* static */ StatusOr> HloModule::CreateFromProto( const HloModuleProto& proto, const HloModuleConfig& module_config, const VersionedComputationHandle& entry_computation_handle) { // The ProgramShape in the passed in module config must match the shapes of // the entry parameters and root. - TF_ASSIGN_OR_RETURN(ProgramShape expected_program_shape, - ProgramShapeFromProto(proto)); + TF_RET_CHECK(proto.has_program_shape()) + << "No program shape found in the proto"; + const auto& expected_program_shape = proto.program_shape(); TF_RET_CHECK(expected_program_shape.parameters_size() == module_config.entry_computation_layout().parameter_count()); for (int i = 0; i < expected_program_shape.parameters_size(); ++i) { @@ -339,8 +303,9 @@ StatusOr> HloModule::CreateFromProto( /* static */ StatusOr HloModule::CreateModuleConfigFromProto( const HloModuleProto& module) { - TF_ASSIGN_OR_RETURN(ProgramShape program_shape, - ProgramShapeFromProto(module)); + TF_RET_CHECK(module.has_program_shape()) + << "No program shape found in the proto"; + const auto& program_shape = module.program_shape(); HloModuleConfig module_config(program_shape); @@ -563,6 +528,18 @@ std::unique_ptr HloModule::Clone(const string& suffix) const { return module; } +HloComputation* HloModule::DeepCloneComputation(HloComputation* computation) { + HloComputation* clone = AddEmbeddedComputation(computation->Clone("", this)); + TF_CHECK_OK( + clone->root_instruction()->Accept([this](HloInstruction* instruction) { + instruction->ReplaceCalledComputations([this](HloComputation* callee) { + return DeepCloneComputation(callee); + }); + return Status::OK(); + })); + return clone; +} + uint64 HloModule::RandomNew64() const { tensorflow::mutex_lock l(rng_mutex_); return rng_(); diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index 4bfe8d89ce0a285de6d05d4867aaa6b266d78d12..755bbd359f7b95e7f3f3cbee1b46df85908202c6 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -85,6 +85,10 @@ class HloModule { // Returns a deep copy of this module including all computations. std::unique_ptr Clone(const string& suffix = "clone") const; + // Performs a deep clone of the computation, by recursively cloning all + // the called computations as well. + HloComputation* DeepCloneComputation(HloComputation* computation); + // Return a pointer to the entry computation of the module.. const HloComputation* entry_computation() const { CHECK_NE(nullptr, entry_computation_); @@ -99,7 +103,7 @@ class HloModule { return config_.mutable_entry_computation_layout(); } - ComputationLayout entry_computation_layout() const { + const ComputationLayout& entry_computation_layout() const { return config_.entry_computation_layout(); } @@ -183,11 +187,6 @@ class HloModule { // Returns a randomly generated uint64. uint64 RandomNew64() const; - // Returns the unique name for a computation in this module. - string GetUniqueCompuationName(const string& prefix) { - return computation_name_uniquer_.GetUniqueName(prefix); - } - // Returns the NameUniquer for uniquing instruction names in this module. NameUniquer& instruction_name_uniquer() { return instruction_name_uniquer_; } diff --git a/tensorflow/compiler/xla/service/hlo_module_config.cc b/tensorflow/compiler/xla/service/hlo_module_config.cc index 822e2f1f53e5ee460b88c2241ecf7f6b91ef608b..4205b0402cb8b2c31141d65be652cd84c22e7262 100644 --- a/tensorflow/compiler/xla/service/hlo_module_config.cc +++ b/tensorflow/compiler/xla/service/hlo_module_config.cc @@ -40,7 +40,7 @@ void HloModuleConfig::SetDefaultComputationLayout( string HloModuleConfig::compilation_cache_key() const { string key = - tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled_); + tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled()); StrAppend(&key, "::("); std::vector params; for (const ShapeLayout& param_layout : diff --git a/tensorflow/compiler/xla/service/hlo_module_config.h b/tensorflow/compiler/xla/service/hlo_module_config.h index a5ee895e48448fbb8fa3879dc1b6764c1f9f6966..586a03d412681cacdd780f48e77baf4cd4c51415 100644 --- a/tensorflow/compiler/xla/service/hlo_module_config.h +++ b/tensorflow/compiler/xla/service/hlo_module_config.h @@ -63,9 +63,19 @@ class HloModuleConfig { return &(*entry_computation_layout_); } - // Sets/returns whether to enable HLO-level profiling. - bool hlo_profiling_enabled() const { return hlo_profiling_enabled_; } - void enable_hlo_profiling(bool enabled) { hlo_profiling_enabled_ = enabled; } + // Returns whether to enable HLO-level profiling. + bool hlo_profiling_enabled() const { + return debug_options_.xla_hlo_profile(); + } + + // Sets/returns whether this is a "host module". Host modules are used to + // record the data- and control-flow dependencies of host side computation + // that communicates with compiled code. They are used for analysis and + // scheduling purposes, but no code is generated. + bool is_host_module() const { return is_host_module_; } + void set_is_host_module(bool is_host_module) { + is_host_module_ = is_host_module; + } // Sets/returns the module seed set during execution. void set_seed(uint64 seed) { seed_ = seed; } @@ -101,8 +111,8 @@ class HloModuleConfig { tensorflow::gtl::optional entry_computation_layout_; - // Whether to enable HLO-level profiling. - bool hlo_profiling_enabled_ = false; + // Whether this is a 'host module'. + bool is_host_module_ = false; // Module/graph-level seed handle. uint64 seed_ = 0; diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc new file mode 100644 index 0000000000000000000000000000000000000000..fa5dcb0b369d17c70c64c67b9f11640c93fb4278 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -0,0 +1,350 @@ +/* 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_module_group_metadata.h" + +#include +#include + +#include "tensorflow/compiler/xla/ptr_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/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +string HloModuleGroupMetadata::TrackedInstruction::ToString() const { + string repr = + (instruction_ != nullptr) ? instruction_->ToShortString() : "NULL"; + switch (kind_) { + case ComputationKind::kInvalid: + repr += ":INVALID"; + break; + case ComputationKind::kWhileCondition: + repr += ":WHILE_CONDITION"; + break; + case ComputationKind::kWhileBody: + repr += ":WHILE_BODY"; + break; + case ComputationKind::kConditionalTrue: + repr += ":CONDITIONAL_TRUE"; + break; + case ComputationKind::kConditionalFalse: + repr += ":CONDITIONAL_FALSE"; + break; + } + return repr; +} + +/* static */ StatusOr> +HloModuleGroupMetadata::Build(const std::vector& modules) { + auto metadata = absl::make_unique(modules); + TF_RETURN_IF_ERROR(metadata->Build()); + return std::move(metadata); +} + +Status HloModuleGroupMetadata::Build() { + TF_RETURN_IF_ERROR(RecordInstructions()); + TF_RETURN_IF_ERROR(VerifyChannelInstructions()); + + // Record all companion while instructions. + const auto visitor = [this](HloInstruction* hlo) -> Status { + // We only need to process if the instruction is within the computation + // of a companion instruction, like in the condition or body computation + // of a While. + const TrackedInstruction* tracked = GetTrackedInstruction(hlo->parent()); + if (tracked == nullptr) { + return Status::OK(); + } + // Add the parent computation of this channel instruction and its peer + // computation (both must be while computations) as companions. + if (IsChannelInstruction(hlo)) { + HloComputation* peer_computation = PeerComputation(hlo); + const TrackedInstruction* peer_tracked = + GetTrackedInstruction(peer_computation); + TF_RET_CHECK(peer_tracked != nullptr) + << "Peer instruction is not a possible companion"; + TF_RET_CHECK(*tracked == *peer_tracked) + << "Peer instruction does not match the computation kind"; + TF_RETURN_IF_ERROR( + AddCompanion(tracked->instruction(), peer_tracked->instruction())); + } + + // Add the parents of companion instructions (they must be all of the same + // kind of instructions, opcode wise) as companions. + if (IsCompanionInstruction(hlo)) { + for (HloInstruction* companion : Companions(hlo)) { + const TrackedInstruction* companion_tracked = + GetTrackedInstruction(companion->parent()); + TF_RET_CHECK(companion_tracked != nullptr); + TF_RET_CHECK(*tracked == *companion_tracked); + TF_RETURN_IF_ERROR(AddCompanion(tracked->instruction(), + companion_tracked->instruction())); + } + } + return Status::OK(); + }; + + // Visit the computations in postorder so that the companion information grows + // from inner computations to outer ones. + for (HloModule* module : modules_) { + for (HloComputation* computation : module->MakeComputationPostOrder()) { + TF_RETURN_IF_ERROR(computation->Accept(visitor)); + } + } + return Status::OK(); +} + +bool HloModuleGroupMetadata::IsChannelInstruction( + const HloInstruction* instruction) const { + switch (instruction->opcode()) { + case HloOpcode::kSend: + case HloOpcode::kRecv: + case HloOpcode::kSendDone: + case HloOpcode::kRecvDone: + return true; + default: + return false; + } +} + +bool HloModuleGroupMetadata::IsCompanionInstruction(HloInstruction* hlo) const { + return companion_set_index_.count(hlo) > 0; +} + +bool HloModuleGroupMetadata::InstructionCommunicates( + HloInstruction* hlo) const { + return IsChannelInstruction(hlo) || IsCompanionInstruction(hlo); +} + +const HloModuleGroupMetadata::Channel& HloModuleGroupMetadata::GetChannel( + int64 channel_id) const { + CHECK(channel_id_map_.find(channel_id) != channel_id_map_.end()); + return channels_[channel_id_map_.at(channel_id)]; +} + +HloComputation* HloModuleGroupMetadata::PeerComputation( + const HloInstruction* instruction) const { + CHECK(IsChannelInstruction(instruction)); + const Channel& channel = GetChannel(instruction->channel_id()); + switch (instruction->opcode()) { + case HloOpcode::kSend: + case HloOpcode::kSendDone: + return channel.recv->parent(); + case HloOpcode::kRecv: + case HloOpcode::kRecvDone: + return channel.send->parent(); + default: + LOG(FATAL) << "opcode not supported"; + } +} + +std::vector +HloModuleGroupMetadata::GetCompanionsPath(const HloInstruction* hlo) const { + std::vector path; + const HloComputation* parent = hlo->parent(); + const TrackedInstruction* companion; + while ((companion = GetTrackedInstruction(parent)) != nullptr) { + parent = companion->instruction()->parent(); + path.push_back(*companion); + } + return path; +} + +bool HloModuleGroupMetadata::CheckCompanionPathsCompatibility( + const std::vector& path0, + const std::vector& path1) const { + if (path0.size() != path1.size()) { + VLOG(5) << "Companion path size do not match: " << path0.size() + << " != " << path1.size(); + return false; + } + for (int64 i = 0; i < path0.size(); ++i) { + if (path0[i] != path1[i]) { + VLOG(5) << "Companion instructions at path index " << i + << " do not have the same opcode: " << path0[i].ToString() + << " vs " << path1[i].ToString(); + return false; + } + } + return true; +} + +int64 HloModuleGroupMetadata::GetModuleId(const HloModule* module) const { + for (int64 i = 0; i < modules_.size(); ++i) { + if (modules_[i] == module) { + return i; + } + } + LOG(FATAL) << "unknown module"; +} + +Status HloModuleGroupMetadata::RecordInstructions() { + const auto visitor = [this](HloInstruction* hlo) -> Status { + if (hlo->opcode() == HloOpcode::kWhile) { + tracked_instructions_[hlo->while_condition()] = + TrackedInstruction(hlo, ComputationKind::kWhileCondition); + tracked_instructions_[hlo->while_body()] = + TrackedInstruction(hlo, ComputationKind::kWhileBody); + } else if (hlo->opcode() == HloOpcode::kConditional) { + tracked_instructions_[hlo->true_computation()] = + TrackedInstruction(hlo, ComputationKind::kConditionalTrue); + tracked_instructions_[hlo->false_computation()] = + TrackedInstruction(hlo, ComputationKind::kConditionalFalse); + } + if (!IsChannelInstruction(hlo)) { + return Status::OK(); + } + + // Add a new channel if needed. + if (channel_id_map_.find(hlo->channel_id()) == channel_id_map_.end()) { + channels_.emplace_back(); + channels_.back().id = hlo->channel_id(); + channel_id_map_[hlo->channel_id()] = channels_.size() - 1; + max_channel_id_ = std::max(max_channel_id_, hlo->channel_id()); + } + Channel& channel = channels_[channel_id_map_[hlo->channel_id()]]; + + if (hlo->opcode() == HloOpcode::kSend) { + TF_RET_CHECK(channel.send == nullptr) + << "channel id " << hlo->channel_id() + << " is used by multiple send instructions"; + channel.send = hlo; + } + if (hlo->opcode() == HloOpcode::kRecv) { + TF_RET_CHECK(channel.recv == nullptr) + << "channel id " << hlo->channel_id() + << " is used by multiple recv instructions"; + channel.recv = hlo; + } + if (hlo->opcode() == HloOpcode::kSendDone) { + TF_RET_CHECK(channel.send_done == nullptr) + << "channel id " << hlo->channel_id() + << " is used by multiple send-done instructions"; + channel.send_done = hlo; + } + if (hlo->opcode() == HloOpcode::kRecvDone) { + TF_RET_CHECK(channel.recv_done == nullptr) + << "channel id " << hlo->channel_id() + << " is used by multiple recv-done instructions"; + channel.recv_done = hlo; + } + return Status::OK(); + }; + + for (HloModule* module : modules_) { + for (auto* computation : module->computations()) { + TF_RETURN_IF_ERROR(computation->Accept(visitor)); + } + } + return Status::OK(); +} + +Status HloModuleGroupMetadata::AddCompanion(HloInstruction* instruction1, + HloInstruction* instruction2) { + TF_RET_CHECK(instruction1->opcode() == HloOpcode::kWhile || + instruction1->opcode() == HloOpcode::kConditional); + VLOG(2) << "adding as companions:" << instruction1->ToString() << " and " + << instruction2->ToString(); + + if (!ContainsKey(companion_set_index_, instruction1) && + !ContainsKey(companion_set_index_, instruction2)) { + companion_sets_.push_back( + absl::make_unique>()); + auto companion_set = companion_sets_.back().get(); + companion_set->insert(instruction1); + companion_set->insert(instruction2); + companion_set_index_[instruction1] = companion_sets_.size() - 1; + companion_set_index_[instruction2] = companion_sets_.size() - 1; + } else if (!ContainsKey(companion_set_index_, instruction1)) { + companion_sets_[companion_set_index_[instruction2]]->insert(instruction1); + companion_set_index_[instruction1] = companion_set_index_[instruction2]; + } else if (!ContainsKey(companion_set_index_, instruction2)) { + companion_sets_[companion_set_index_[instruction1]]->insert(instruction2); + companion_set_index_[instruction2] = companion_set_index_[instruction1]; + } else if (companion_set_index_[instruction1] != + companion_set_index_[instruction2]) { + companion_sets_[companion_set_index_[instruction1]]->insert( + Companions(instruction2).begin(), Companions(instruction2).end()); + int64 index_to_remove = companion_set_index_[instruction2]; + for (HloInstruction* hlo : Companions(instruction2)) { + companion_set_index_[hlo] = companion_set_index_[instruction1]; + } + companion_sets_.erase(companion_sets_.begin() + index_to_remove); + } + return Status::OK(); +} + +Status HloModuleGroupMetadata::VerifyChannelInstructions() { + for (const Channel& channel : channels_) { + if (channel.send == nullptr) { + return FailedPrecondition("missing send for id : %lld", channel.id); + } + if (channel.recv == nullptr) { + return FailedPrecondition("missing recv for id : %lld", channel.id); + } + if (channel.send_done == nullptr) { + return FailedPrecondition("missing send-done for id : %lld", channel.id); + } + if (channel.recv_done == nullptr) { + return FailedPrecondition("missing recv-done for id : %lld", channel.id); + } + } + + // Check if the shapes match for each channel. + for (const Channel& channel : channels_) { + const Shape& send_shape = channel.send->operand(0)->shape(); + const Shape& recv_shape = channel.recv_done->shape(); + if (!ShapeUtil::Compatible(send_shape, recv_shape)) { + return FailedPrecondition("send/recv shapes do not match"); + } + } + + // Check if channel instructions are used only in allowed computations. + const auto allowed = [this](HloInstruction* hlo) { + HloComputation* computation = hlo->parent(); + const HloModule* module = computation->parent(); + if (module->entry_computation() == computation || + tracked_instructions_.count(computation) > 0) { + return true; + } + return false; + }; + for (const Channel& channel : channels_) { + if (!allowed(channel.send) || !allowed(channel.send_done) || + !allowed(channel.recv) || !allowed(channel.recv_done)) { + return FailedPrecondition("channel is used in disallowed computation"); + } + } + // Check if the nest levels match for each channel. + for (const Channel& channel : channels_) { + std::vector path = GetCompanionsPath(channel.send); + if (!CheckCompanionPathsCompatibility( + path, GetCompanionsPath(channel.send_done)) || + !CheckCompanionPathsCompatibility(path, + GetCompanionsPath(channel.recv)) || + !CheckCompanionPathsCompatibility( + path, GetCompanionsPath(channel.recv_done))) { + return FailedPrecondition( + "Nest companion paths do not match for channel %lld", channel.id); + } + } + return Status::OK(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h new file mode 100644 index 0000000000000000000000000000000000000000..c48a7ab0b59269474f7406ef24a249355528e085 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h @@ -0,0 +1,239 @@ +/* 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_HLO_MODULE_GROUP_METADATA_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_GROUP_METADATA_H_ + +#include +#include +#include +#include +#include + +#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/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// Class for bookkeeping the information on the given modules, in particular on +// the interaction between computations. +// +// Companion instructions are one of the information collected as we build the +// metadata. For example, for each While instruction, companion instructions +// refer to a set of While instructions in other computations that communicate +// with each other. +// In the example below with 3 modules, {While_0, While_2, While_5}, {While_1, +// While_4}, {While_3, While_6} are companion sets. +// +// +// While_0() { While_2() { While_5() { +// While_1() { Send(0) } While_3() { Send(1) } While_6() { Recv(1) } +// } While_4() { Recv(0) } +// } +// +// Companion instructions are used to detect cycles in the graph and also for +// global scheduling. +class HloModuleGroupMetadata { + public: + // The kind of companion computation a given instruction can be within. + enum class ComputationKind { + kInvalid, + kWhileCondition, + kWhileBody, + kConditionalTrue, + kConditionalFalse, + }; + + // Tracks the instruction mapped to a given computation, and the computation + // kind. + // For example, a body computation of a while instruction, will generate a + // TrackedInstruction with instruction being the while instruction, and + // kind being ComputationKind::kWhileBody. + class TrackedInstruction { + public: + TrackedInstruction() = default; + TrackedInstruction(HloInstruction* instruction, ComputationKind kind) + : instruction_(instruction), kind_(kind) {} + + bool operator==(const TrackedInstruction& rhs) const { + return instruction_->opcode() == rhs.instruction_->opcode() && + kind_ == rhs.kind_; + } + bool operator!=(const TrackedInstruction& rhs) const { + return !operator==(rhs); + } + + HloInstruction* instruction() const { return instruction_; } + + string ToString() const; + + private: + HloInstruction* instruction_ = nullptr; + ComputationKind kind_ = ComputationKind::kInvalid; + }; + + // Represents a channel and the 4 instructions that form the channel. + struct Channel { + int64 id = -1; + HloInstruction* send = nullptr; + HloInstruction* recv = nullptr; + HloInstruction* send_done = nullptr; + HloInstruction* recv_done = nullptr; + }; + + explicit HloModuleGroupMetadata(const std::vector& modules) + : modules_(modules) {} + + ~HloModuleGroupMetadata() = default; + + // Build and return the metadata for the given modules. + static StatusOr> Build( + const std::vector& modules); + + // Returns true if the instruction is one of the 4 channel instructions (Send, + // Recv, SendDone, RecvDone). + bool IsChannelInstruction(const HloInstruction* instruction) const; + + // Returns true if the instruction is a companion instruction. See the class + // comment above on companion instructions. + bool IsCompanionInstruction(HloInstruction* hlo) const; + + // Returns true if the instruction is either a channel instruction or a + // companion instruction. + bool InstructionCommunicates(HloInstruction* hlo) const; + + // Returns the Channel instance for the given channel id. + const Channel& GetChannel(int64 channel_id) const; + + // Returns the computation that contains the peer channel instructions for + // the given instruction. + // + // Precondition: IsChannelInstruction(instruction) is true. + HloComputation* PeerComputation(const HloInstruction* instruction) const; + + // Returns the path of the nested companion instructions, in terms of HLO + // instructions. The path goes from inner to outer companions. + // The returned path does not include the input hlo instruction, in case it + // is a companion instruction. + std::vector GetCompanionsPath( + const HloInstruction* hlo) const; + + // Checks whether two companion paths (as returned by the GetCompanionsPath() + // API) are compatible. The two paths are compatible if the sequence of + // opcodes, and the companion kinds, of the two paths matches. + bool CheckCompanionPathsCompatibility( + const std::vector& path0, + const std::vector& path1) const; + + // Returns the unique integer for each module. The returned id is the index of + // the module in the module vector. + int64 GetModuleId(const HloModule* module) const; + + // Returns the companion instructions for the given instruction. + // + // Precondition: IsCompanionWhile(instruction) is true. + const std::unordered_set& Companions( + HloInstruction* instruction) const { + CHECK_EQ(companion_set_index_.count(instruction), 1); + return companion_set(companion_set_index_.at(instruction)); + } + + // Returns the companion set at the given index. + const std::unordered_set& companion_set(int64 index) const { + CHECK_LT(index, companion_sets_.size()); + return *companion_sets_[index]; + } + + // Returns the companion set index of the given instruction. + int64 companion_set_index(HloInstruction* instruction) const { + return companion_set_index_.at(instruction); + } + + // Returns the list of all companion sets in the HLO module group. + const std::vector>>& + companion_sets() const { + return companion_sets_; + } + + // Returns all channels in the module group. + const std::vector& channels() const { return channels_; } + + // Returns the maximum channel id used in the module group. + int64 max_channel_id() const { return max_channel_id_; } + + private: + Status Build(); + + // Record all channel instructions and While instructions. + Status RecordInstructions(); + + // Verifies the given HloModules are well-formed and follow the specification, + // in particular with respect to using channel instructions. + // + // * Each channel has all 4 instructions (Send, Recv, SendDone, RecvDone). + // * The shape of channel instructions match. + // * The nest level of channel instructions match. + // * Channel instructions are used in allowed computations; i.e., in the + // entry computation of the module or condition/body of While computations. + // + // TODO(b/62064342): Currently, HloModuleGroupScheduler checks if there is a + // cycle in the graph, but it would be good to verify here. + Status VerifyChannelInstructions(); + + // Adds metadata that the given two instructions are companions. + Status AddCompanion(HloInstruction* instruction1, + HloInstruction* instruction2); + + // Retrieves a pointer to the stored TrackedInstruction associated with a + // tracked computation, or nullptr in case such computation is not tracked. + const TrackedInstruction* GetTrackedInstruction( + const HloComputation* computation) const { + auto it = tracked_instructions_.find(computation); + return it != tracked_instructions_.end() ? &it->second : nullptr; + } + + // List of all companion instructions sets in the module. + std::vector>> + companion_sets_; + + // Map from each companion while instruction to the index into companion_set_. + tensorflow::gtl::FlatMap companion_set_index_; + + // Map from computation to the instruction using it (a kWhile, kConditional). + tensorflow::gtl::FlatMap + tracked_instructions_; + + // All channels in the module. + std::vector channels_; + + // Map from channel ids to the index in channels_. + tensorflow::gtl::FlatMap channel_id_map_; + + // The maximum channel id used in the module group. + int64 max_channel_id_ = -1; + + // The modules that this metadata was built from. + const std::vector& modules_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_GROUP_METADATA_H_ diff --git a/tensorflow/compiler/xla/service/hlo_module_group_util.cc b/tensorflow/compiler/xla/service/hlo_module_group_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..289c96b0a7b90c5f8a122cd3fc327a5762099106 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_module_group_util.cc @@ -0,0 +1,316 @@ +/* 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_module_group_util.h" + +#include +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.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/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +std::vector HloModuleGroupUtil::GlobalPredecessors( + HloInstruction* instruction) { + std::vector predecessors; + + // Adds to the unique predecessors list and also add companion instructions + // if the given predecessor has those. + auto add_unique_predecessor = [&](HloInstruction* predecessor) { + if (std::find(predecessors.begin(), predecessors.end(), predecessor) != + predecessors.end()) { + return; + } + if (!metadata_.IsCompanionInstruction(predecessor)) { + predecessors.push_back(predecessor); + return; + } + for (HloInstruction* companion : metadata_.Companions(predecessor)) { + predecessors.push_back(companion); + } + }; + + // If the given instruction is a companion instruction, we need to find the + // predecessors of all of its companion instructions. + std::vector instruction_group; + if (metadata_.IsCompanionInstruction(instruction)) { + for (HloInstruction* companion : metadata_.Companions(instruction)) { + instruction_group.push_back(companion); + } + } else { + instruction_group.push_back(instruction); + } + + for (HloInstruction* hlo : instruction_group) { + for (HloInstruction* operand : hlo->operands()) { + add_unique_predecessor(operand); + } + for (HloInstruction* control_predecessor : hlo->control_predecessors()) { + add_unique_predecessor(control_predecessor); + } + } + if (instruction->opcode() == HloOpcode::kRecvDone) { + // Send is a remote predecessor of RecvDone. + HloInstruction* send = metadata_.GetChannel(instruction->channel_id()).send; + add_unique_predecessor(send); + } + if (instruction->opcode() == HloOpcode::kSend) { + // Recv is a remote predecessor of Send. + HloInstruction* recv_done = + metadata_.GetChannel(instruction->channel_id()).recv_done; + CHECK(recv_done->opcode() == HloOpcode::kRecvDone); + CHECK_EQ(recv_done->operand_count(), 1); + HloInstruction* recv = recv_done->mutable_operand(0); + add_unique_predecessor(recv); + } + return predecessors; +} + +std::vector HloModuleGroupUtil::GlobalSuccessors( + HloInstruction* instruction) { + std::vector successors; + + // Adds to the unique successors list and also add companion instructions + // if the given successor has those. + auto add_unique_successor = [&](HloInstruction* successor) { + if (std::find(successors.begin(), successors.end(), successor) != + successors.end()) { + return; + } + if (!metadata_.IsCompanionInstruction(successor)) { + successors.push_back(successor); + return; + } + for (HloInstruction* companion : metadata_.Companions(successor)) { + successors.push_back(companion); + } + }; + + // If the given instruction is a companion instruction, we need to find the + // successors of all of its companion instructions. + std::vector instruction_group; + if (metadata_.IsCompanionInstruction(instruction)) { + for (HloInstruction* companion : metadata_.Companions(instruction)) { + instruction_group.push_back(companion); + } + } else { + instruction_group.push_back(instruction); + } + + for (HloInstruction* hlo : instruction_group) { + for (HloInstruction* user : hlo->users()) { + add_unique_successor(user); + } + for (HloInstruction* control_successor : hlo->control_successors()) { + add_unique_successor(control_successor); + } + } + if (instruction->opcode() == HloOpcode::kRecv) { + // Send is a remote successor of Recv. + const HloInstruction* recv_done = instruction->users().front(); + CHECK(recv_done->opcode() == HloOpcode::kRecvDone); + HloInstruction* send = metadata_.GetChannel(instruction->channel_id()).send; + add_unique_successor(send); + } + if (instruction->opcode() == HloOpcode::kSend) { + // RecvDone is a remote successor of Send. + HloInstruction* recv_done = + metadata_.GetChannel(instruction->channel_id()).recv_done; + add_unique_successor(recv_done); + } + return successors; +} + +std::vector HloModuleGroupUtil::RootInstructions( + tensorflow::gtl::ArraySlice computations) { + std::vector roots; + for (HloComputation* computation : computations) { + for (HloInstruction* instruction : computation->instructions()) { + if (GlobalSuccessors(instruction).empty()) { + roots.push_back(instruction); + } + } + } + return roots; +} + +Status HloModuleGroupUtil::VisitTopologicalOrder( + VisitStates* visit_state, const VisitFunction& visit_function, + HloInstruction* root) { + // Stack of HLO instructions visited in DFS order. + std::stack stack; + stack.push(root); + + while (!stack.empty()) { + HloInstruction* hlo = stack.top(); + + // Find the instruction group of the currently visited instruction. The + // instruction group represents all companion instructions of the + // current instruction, and are considered to be a single entity for the + // purpose of the traversal (i.e., they must always be in the same visit + // state). + std::vector instruction_group; + if (metadata_.IsCompanionInstruction(hlo)) { + for (HloInstruction* companion : metadata_.Companions(hlo)) { + instruction_group.push_back(companion); + } + } else { + instruction_group.push_back(hlo); + } + + if ((*visit_state)[hlo] == VisitState::kVisited) { + // All instructions in the group must be in the same state. + for (HloInstruction* instruction : instruction_group) { + TF_RET_CHECK((*visit_state)[instruction] == VisitState::kVisited); + } + stack.pop(); + continue; + } + + if ((*visit_state)[hlo] == VisitState::kVisiting) { + TF_RETURN_IF_ERROR(visit_function(hlo, instruction_group)); + + // Set the visit state of all instructions in the group to kVisited. + for (HloInstruction* instruction : instruction_group) { + TF_RET_CHECK((*visit_state)[instruction] == VisitState::kVisiting); + (*visit_state)[instruction] = VisitState::kVisited; + } + stack.pop(); + continue; + } + + // Set the visit state of all instructions in the group to kVisiting. + for (HloInstruction* instruction : instruction_group) { + TF_RET_CHECK((*visit_state)[instruction] == VisitState::kNotVisited) + << instruction->ToString(); + (*visit_state)[instruction] = VisitState::kVisiting; + } + + // For each instruction in the group, visit its predecessors (operands, + // control predecessors and remote predecessors). + for (HloInstruction* instruction : instruction_group) { + for (HloInstruction* predecessor : GlobalPredecessors(instruction)) { + // Visiting a node that is already being visited implies that there is + // a cycle. Generate an error with the list of instructions in the + // cycle. + if ((*visit_state)[predecessor] == VisitState::kVisiting) { + string cyclic_instructions; + for (const auto& state : *visit_state) { + if (state.second == VisitState::kVisiting) { + tensorflow::strings::StrAppend(&cyclic_instructions, + state.first->ToString(), "\n"); + } + } + // TODO(b/64305524): Improve the error message to print out the + // instructions in a deterministic order that forms the cycle. + return FailedPrecondition( + "Cross-computation cycle detected via communicating nodes. The " + "cycle contains the node %s. The cycle is found among the " + "following nodes. Note that the order of the nodes is arbitrary " + "and that the list may include nodes that are not part of the " + "cycle.\n%s", + predecessor->ToString().c_str(), cyclic_instructions.c_str()); + } + stack.push(predecessor); + } + } + } + + return Status::OK(); +} + +Status HloModuleGroupUtil::VerifyComputations( + tensorflow::gtl::ArraySlice computations) { + auto visit_function = + [&](HloInstruction* instruction, + const std::vector& instruction_group) { + return Status::OK(); + }; + int64 instructions_count = 0; + VisitStates visit_states; + for (HloComputation* computation : computations) { + // Visit all instructions, and not just from the root instruction of the + // computation. This allows us to detect dead cycles (i.e., cycles that + // are not reachable from the root) or to enforce an order for the + // communication instructions that are not reachable from any roots. + for (HloInstruction* instruction : computation->instructions()) { + TF_RETURN_IF_ERROR( + VisitTopologicalOrder(&visit_states, visit_function, instruction)); + } + instructions_count += computation->instruction_count(); + } + + // Check if all instructions are visited and are in the visited state. + TF_RET_CHECK(visit_states.size() == instructions_count); + for (auto& state : visit_states) { + TF_RET_CHECK(state.second == VisitState::kVisited); + } + + return Status::OK(); +} + +StatusOr> +HloModuleGroupUtil::ComputeReachability( + tensorflow::gtl::ArraySlice computations) { + std::list post_order; + auto visit_function = + [&](HloInstruction* instruction, + const std::vector& instruction_group) { + post_order.insert(post_order.end(), instruction_group.begin(), + instruction_group.end()); + return Status::OK(); + }; + HloModuleGroupUtil::VisitStates visit_states; + for (HloInstruction* root : RootInstructions(computations)) { + TF_RETURN_IF_ERROR( + VisitTopologicalOrder(&visit_states, visit_function, root)); + } + auto reachability = absl::make_unique(post_order); + for (HloInstruction* hlo : post_order) { + reachability->SetReachabilityToUnion(GlobalPredecessors(hlo), hlo); + } + return std::move(reachability); +} + +void HloModuleGroupUtil::UpdateReachabilityThroughInstruction( + HloInstruction* instruction, HloReachabilityMap* reachability_map) { + std::queue worklist; + worklist.push(instruction); + + while (!worklist.empty()) { + HloInstruction* item = worklist.front(); + worklist.pop(); + if (reachability_map->SetReachabilityToUnion(GlobalPredecessors(item), + item)) { + for (HloInstruction* successor : GlobalSuccessors(item)) { + worklist.push(successor); + } + } + } +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_group_util.h b/tensorflow/compiler/xla/service/hlo_module_group_util.h new file mode 100644 index 0000000000000000000000000000000000000000..c25ca1aff50b288f3ac3885cbed53e7ba9768430 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_module_group_util.h @@ -0,0 +1,117 @@ +/* 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_HLO_MODULE_GROUP_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_GROUP_UTIL_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module_group_metadata.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatmap.h" + +namespace xla { + +// Collection of utilities for handling HloModuleGroups. +class HloModuleGroupUtil { + public: + explicit HloModuleGroupUtil(const HloModuleGroupMetadata& metadata) + : metadata_(metadata) {} + + // Returns all unique predecessors of the instruction. This includes: + // * predecessors in the same computation: operands and control predecessors + // * Recv is a predecessor of Send + // * Send is a predecessor of RecvDone + // * predecessors of companions (if the instruction is a companion while) + // * predecessors' companions (for any predecessor that is a companion while) + std::vector GlobalPredecessors(HloInstruction* instruction); + + // Returns all unique successors of the instruction. This includes: + // * successors in the same computation: users and control successors + // * Send is a successor of Recv + // * RecvDone is a predecessor of Send + // * successors of companions (if the instruction is a companion while) + // * successors' companions (for any successor that is a companion while) + std::vector GlobalSuccessors(HloInstruction* instruction); + + // Returns the root instructions of the computations. + std::vector RootInstructions( + tensorflow::gtl::ArraySlice computations); + + // Visit state of each instruction during DFS traversal. + enum VisitState { + kNotVisited = 0, + kVisiting, + kVisited, + }; + + // Function called on each instruction group during the DFS traversal. See the + // comment for VisitTopologicalOrder()). + using VisitFunction = std::function& instruction_group)>; + + // Given the hlo instruction as the root, recursively visits all its + // predecessor instructions in DFS order to visit nodes in topological order. + // + // Note that the DFS traversal does not only visit nodes in the same + // computation (parent of the root instruction), but also visits nodes in + // different computations connected via communication instructions. During the + // traversal, companion While instructions (see the class comment in + // HloModuleGroupMetadata) are treated as a single instruction (called + // instruction group, which contains only a single instruction if the visiting + // node is not a companion while) -- visiting one of the instructions in the + // group effectively visits all other instructions in the group, and then all + // predecessor instructions of the group are visited. + // + // * visit_state: map from each instruction to its visit state. + // * visit_function: function called when each instruction group. + // * root: the root instruction of the traversal. + using VisitStates = tensorflow::gtl::FlatMap; + Status VisitTopologicalOrder(VisitStates* visit_state, + const VisitFunction& visit_function, + HloInstruction* root); + + // Verifies that the computations are well-formed (e.g., no cycles). + Status VerifyComputations( + tensorflow::gtl::ArraySlice computations); + + // Below Reachability utils resemble those in HloComputation, except that + // they can handle instructions across multiple computations. + // + // Creates the reachability map for the instructions in the computations. + StatusOr> ComputeReachability( + tensorflow::gtl::ArraySlice computations); + + // Updates the reachability of the given instruction, taking the global + // predeccessorss and successors into account. + void UpdateReachabilityThroughInstruction( + HloInstruction* instruction, HloReachabilityMap* reachability_map); + + private: + const HloModuleGroupMetadata& metadata_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_GROUP_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index 3d64523a79fc50638fdf378b5d521a5cd4482b90..af24604c39b554f146793594958f373999844b4c 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -76,9 +76,11 @@ namespace xla { V(kFft, "fft") \ V(kFloor, "floor") \ V(kFusion, "fusion", kHloOpcodeIsVariadic) \ + V(kGather, "gather") \ V(kGe, "greater-than-or-equal-to", kHloOpcodeIsComparison) \ V(kGetTupleElement, "get-tuple-element") \ V(kGt, "greater-than", kHloOpcodeIsComparison) \ + V(kHostCompute, "host-compute") \ V(kImag, "imag") \ V(kInfeed, "infeed") \ V(kIsFinite, "is-finite") \ diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc index 68e3c9618c1fe9daacb0aee3ee98862c8b9e4bc4..e89d94bede6c437ca1131a1b1b0098390d58c0d9 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering.cc @@ -66,6 +66,28 @@ bool HloOrdering::ExecutesBefore(const HloInstruction* a, } } + // If the common ancestor is a conditional instruction, even though the true + // and false computations are not really ordered per-se, we define the true + // computation to be ordered before the false one. + // This ensures that buffers can still be shared among the two computations + // as they will forcibly have disjoint liveness. + if (a_ancestor == b_ancestor && + a_ancestor->opcode() == HloOpcode::kConditional) { + const HloComputation* true_computation = a_ancestor->true_computation(); + const HloComputation* false_computation = a_ancestor->false_computation(); + if (call_graph_->InstructionIsNestedIn(a, true_computation) && + call_graph_->InstructionIsNestedIn(b, false_computation)) { + return true; + } + // If 'b' is the conditional ancestor, and 'a' is within the true or false + // computations, 'a' executes before 'b'. + if (b == a_ancestor && + (call_graph_->InstructionIsNestedIn(a, true_computation) || + call_graph_->InstructionIsNestedIn(a, false_computation))) { + return true; + } + } + return ExecutesBeforeInSameComputation(a_ancestor, b_ancestor); } @@ -118,7 +140,18 @@ bool HloOrdering::IsDefinedBefore(const HloValue& a, const HloValue& b) const { b.defining_instruction()->while_condition()))) { return true; } - + // If 'b' is a conditional phi and 'a' is in the true or false computation, + // then 'a' executes before 'b'. + if (b.is_phi() && + b.defining_instruction()->opcode() == HloOpcode::kConditional && + (call_graph_->InstructionIsNestedIn( + a.defining_instruction(), + b.defining_instruction()->true_computation()) || + call_graph_->InstructionIsNestedIn( + a.defining_instruction(), + b.defining_instruction()->false_computation()))) { + return true; + } return ExecutesBefore(a.defining_instruction(), b.defining_instruction()); } @@ -186,6 +219,22 @@ bool HloOrdering::UseIsBeforeValueDefinition( } } + if (use.instruction->opcode() == HloOpcode::kConditional) { + const HloInstruction* conditional = use.instruction; + if (call_graph_->InstructionIsNestedIn(value.defining_instruction(), + conditional->true_computation())) { + VLOG(4) << " use is conditional " << use.instruction->name() + << " and def is in TRUE computation"; + return true; + } + if (call_graph_->InstructionIsNestedIn(value.defining_instruction(), + conditional->false_computation())) { + VLOG(4) << " use is conditional " << use.instruction->name() + << " and def is in FALSE computation"; + return true; + } + } + VLOG(4) << " use is not before value"; return false; } @@ -196,18 +245,17 @@ bool HloOrdering::LiveRangeStrictlyBefore( VLOG(4) << "LiveRangeStrictlyBefore(a = " << a.ToShortString() << ", b = " << b.ToShortString() << ")"; if (!IsDefinedBefore(a, b)) { - VLOG(4) << "a not defined before b"; + VLOG(4) << a << " not defined before " << b; return false; } - // All uses of 'a' must be before 'b' is defined. for (const HloUse& use : a.uses()) { if (!UseIsBeforeValueDefinition(use, b, dataflow)) { - VLOG(4) << "use of a (" << use << ") not before b is defined"; + VLOG(4) << "use of " << a << " (" << use << ") not before " << b + << " is defined"; return false; } } - return true; } diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc index aba66114de649ce7667ae77174e9c4073b010b90..441d790f0e96c8f26e1540b79e2dd83a4a0b9ba5 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering_test.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc @@ -262,8 +262,8 @@ TEST_F(HloOrderingTest, ValuesInWhileComputations) { scalar_shape, HloOpcode::kAdd, constant, xla_while)); module->AddEntryComputation(builder.Build()); - TF_ASSERT_OK_AND_ASSIGN( - auto dataflow, HloDataflowAnalysis::Run(module.get(), /*ssa_form=*/true)); + TF_ASSERT_OK_AND_ASSIGN(auto dataflow, + HloDataflowAnalysis::Run(*module, /*ssa_form=*/true)); DependencyHloOrdering ordering(module.get()); // Init value is defined before the while, but live range is not before the @@ -362,5 +362,66 @@ ENTRY while.v11 { ordering.ToString(); // Shouldn't crash. } +TEST_F(HloOrderingTest, ConditionalInstructionOrdering) { + const char* module_str = R"( +HloModule test_conditional_module + +true_branch { + param.1 = (s32[], s32[]) parameter(0) + get-tuple-element.1 = s32[] get-tuple-element(param.1), index=0 + get-tuple-element.2 = s32[] get-tuple-element(param.1), index=1 + add.1 = s32[] add(get-tuple-element.1, get-tuple-element.2) + ROOT tuple.1 = (s32[], s32[]) tuple(add.1, get-tuple-element.1) +} + +false_branch { + param.2 = (s32[], s32[]) parameter(0) + get-tuple-element.3 = s32[] get-tuple-element(param.2), index=0 + get-tuple-element.4 = s32[] get-tuple-element(param.2), index=1 + add.2 = s32[] add(get-tuple-element.3, get-tuple-element.4) + ROOT tuple.2 = (s32[], s32[]) tuple(add.2, get-tuple-element.4) +} + +ENTRY root { + param.3 = (pred[], (s32[], s32[])) parameter(0) + pred.1 = pred[] get-tuple-element(param.3), index=0 + cond_arg.1 = (s32[], s32[]) get-tuple-element(param.3), index=1 + conditional = (s32[], s32[]) conditional(pred.1, cond_arg.1, cond_arg.1), true_computation=true_branch, false_computation=false_branch + cond_res.1 = s32[] get-tuple-element(conditional), index=0 + cond_res.2 = s32[] get-tuple-element(conditional), index=1 + add.3 = s32[] add(cond_res.1, cond_res.2) + ROOT result = (s32[], s32[], s32[]) tuple(add.3, cond_res.1, cond_res.2) +})"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + tools::Parse(module_str)); + TF_ASSERT_OK_AND_ASSIGN(auto dataflow, + HloDataflowAnalysis::Run(*module, /*ssa_form=*/true)); + DependencyHloOrdering ordering(module.get()); + + // Even though the true and false branches has no ordering, since they do not + // interfere (as they are mutually exclusive), we define the true computation + // to be before the false one. + // Similarly, any instruction in the true or false branches are considered + // before the conditional instruction. The roots are effectively "at the same + // time" WRT the conditional, but they are Phi-ed anyway. + HloInstruction* add_1 = FindInstruction(module.get(), "add.1"); + HloInstruction* add_2 = FindInstruction(module.get(), "add.2"); + HloInstruction* add_3 = FindInstruction(module.get(), "add.3"); + HloInstruction* conditional = FindInstruction(module.get(), "conditional"); + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(add_1), + dataflow->GetValueDefinedAt(add_2))); + EXPECT_TRUE( + ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(add_2), + dataflow->GetValueDefinedAt(conditional))); + EXPECT_TRUE( + ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(add_1), + dataflow->GetValueDefinedAt(conditional))); + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(add_1), + dataflow->GetValueDefinedAt(add_3))); + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(add_2), + dataflow->GetValueDefinedAt(add_3))); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.cc b/tensorflow/compiler/xla/service/hlo_proto_util.cc index 78e6a101c10a1e812e3e2631d520139fd0bc425c..3460679558d185d1e022660d9a1d23176d0d96bf 100644 --- a/tensorflow/compiler/xla/service/hlo_proto_util.cc +++ b/tensorflow/compiler/xla/service/hlo_proto_util.cc @@ -15,6 +15,10 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_proto_util.h" +#include + +#include "tensorflow/compiler/xla/util.h" + namespace xla { HloProto MakeHloProto(const HloModule& module, @@ -35,4 +39,35 @@ HloProto MakeHloProto(const HloModule& module) { return proto; } +StatusOr> EntryComputationParameterShapes( + const HloProto& hlo_proto) { + if (!hlo_proto.has_hlo_module()) { + return NotFound("HloProto missing HloModuleProto."); + } + if (!hlo_proto.hlo_module().has_program_shape()) { + return NotFound("HloProto missing program shape."); + } + + std::vector parameter_shapes; + const auto& program_shape = hlo_proto.hlo_module().program_shape(); + for (const Shape& shape : program_shape.parameters()) { + parameter_shapes.push_back(&shape); + } + return parameter_shapes; +} + +StatusOr EntryComputationOutputShape(const HloProto& hlo_proto) { + if (!hlo_proto.has_hlo_module()) { + return NotFound("HloProto missing HloModuleProto."); + } + if (!hlo_proto.hlo_module().has_program_shape()) { + return NotFound("HloProto missing program shape."); + } + if (!hlo_proto.hlo_module().program_shape().has_result()) { + return NotFound("HloProto missing result in its program shape"); + } + + return &hlo_proto.hlo_module().program_shape().result(); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.h b/tensorflow/compiler/xla/service/hlo_proto_util.h index 320288fdb9aa0810b306b1d78bd1ff4cfc366ed2..3d9c375cd5d26f92cf8316f78789daf4fc08c927 100644 --- a/tensorflow/compiler/xla/service/hlo_proto_util.h +++ b/tensorflow/compiler/xla/service/hlo_proto_util.h @@ -35,6 +35,15 @@ HloProto MakeHloProto(const HloModule& module, // will not be included in the output. HloProto MakeHloProto(const HloModule& module); +// Returns the shapes of the parameters of the entry computation. Shape pointers +// refer to shapes inside of the given HloProto. +StatusOr> EntryComputationParameterShapes( + const HloProto& hlo_proto); + +// Returns the shape of the output of the entry computation. The shape pointer +// refers to the output shape inside of the given HloProto. +StatusOr EntryComputationOutputShape(const HloProto& hlo_proto); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROTO_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/hlo_proto_util_test.cc b/tensorflow/compiler/xla/service/hlo_proto_util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b9cca138703c8fa61aadf69dd7304a215a9f4be2 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_proto_util_test.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/hlo_proto_util.h" + +#include "tensorflow/compiler/xla/service/hlo.pb.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/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/strings/str_util.h" + +namespace xla { +namespace { + +class HloProtoUtilTest : public ::testing::Test {}; + +TEST_F(HloProtoUtilTest, ParamsAndOutputShapeMissingModule) { + HloProto hlo_proto; + + auto status = EntryComputationParameterShapes(hlo_proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + ::testing::HasSubstr("missing HloModuleProto")); +} + +TEST_F(HloProtoUtilTest, MissingProgramShape) { + HloProto hlo_proto; + HloModuleProto* module = hlo_proto.mutable_hlo_module(); + module->set_name("entry"); + + auto status = EntryComputationParameterShapes(hlo_proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + ::testing::HasSubstr("missing program shape")); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc index c6b4dc0368d92fd477decdfb38045f74f8696803..98b8d34be1f331aaeac94e952deeae1e76379861 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc @@ -60,6 +60,7 @@ bool IsRematerializable(const HloInstruction* instruction) { switch (instruction->opcode()) { case HloOpcode::kCall: case HloOpcode::kConstant: + case HloOpcode::kConditional: case HloOpcode::kCrossReplicaSum: case HloOpcode::kCustomCall: case HloOpcode::kParameter: diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index 41b079eb799d06321a31f7d7ae0630dc8d58c46b..e5b1c2efa3fc25d23531df298e125521c002dba1 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -110,7 +110,7 @@ HloRunner::HloRunner(se::Platform* platform) { HloRunner::~HloRunner() {} -StatusOr> HloRunner::ExecuteInternal( +StatusOr> HloRunner::Execute( std::unique_ptr module, const tensorflow::gtl::ArraySlice arguments, bool run_hlo_passes) { @@ -158,8 +158,8 @@ StatusOr> HloRunner::ExecuteInternal( TF_ASSIGN_OR_RETURN( std::unique_ptr result, - executable->ExecuteOnStream(&service_run_options, argument_buffer_ptrs, - /*hlo_execution_profile=*/nullptr)); + executable->ExecuteOnStreamWrapper( + &service_run_options, /*profile=*/nullptr, argument_buffer_ptrs)); // Create a ScopedShapedBuffer of the result to manage deallocation. This will // deallocate all the device memory when it goes out of scope. diff --git a/tensorflow/compiler/xla/service/hlo_runner.h b/tensorflow/compiler/xla/service/hlo_runner.h index cbaebc68bee708090b8ccb2eae19b556c4d6d453..06ce22a5b9fc7b3d6c10857c84196094c0eed303 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.h +++ b/tensorflow/compiler/xla/service/hlo_runner.h @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" @@ -64,17 +65,27 @@ class HloRunner { const std::string& filename, const DebugOptions& debug_options); // Executes the given module with given literals as input and returns the - // result as a Literal. The LiteralPtr type accepts Literal* or - // std::unique_ptr. + // result as a Literal. // // If run_hlo_passes is false, the module will be executed without Hlo // optimization. - template StatusOr> Execute( std::unique_ptr module, - const tensorflow::gtl::ArraySlice arguments, + const tensorflow::gtl::ArraySlice arguments, bool run_hlo_passes = true); + StatusOr> Execute( + std::unique_ptr module, + const tensorflow::gtl::ArraySlice> arguments, + bool run_hlo_passes = true) { + // Construct a vector of plain pointers for the arguments. + std::vector argument_pointers; + c_transform( + arguments, std::back_inserter(argument_pointers), + [](const std::unique_ptr& literal) { return literal.get(); }); + return Execute(std::move(module), argument_pointers, run_hlo_passes); + } + // If backend is not created in the constructor, creates and returns the // default backend. If creation fails, crashes the program. // @@ -83,11 +94,6 @@ class HloRunner { Backend& backend(); private: - StatusOr> ExecuteInternal( - std::unique_ptr module, - const tensorflow::gtl::ArraySlice arguments, - bool run_hlo_passes = true); - struct EigenThreadPoolWrapper; std::unique_ptr thread_pool_wrapper_; @@ -95,19 +101,6 @@ class HloRunner { std::unique_ptr backend_; }; -template -StatusOr> HloRunner::Execute( - std::unique_ptr module, - const tensorflow::gtl::ArraySlice arguments, - bool run_hlo_passes) { - // Construct a vector of plain pointers for the arguments. - std::vector argument_pointers; - for (const auto& argument : arguments) { - argument_pointers.push_back(&*argument); - } - return ExecuteInternal(std::move(module), argument_pointers, run_hlo_passes); -} - } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_RUNNER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_scheduling.cc b/tensorflow/compiler/xla/service/hlo_scheduling.cc index 2594c29efd717b3bead34d326c28c7efdf093c50..da448ed71ab470e0c4d72e234bf1f1087d3ea7b4 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,7 +151,7 @@ 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 + // 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*> @@ -205,7 +206,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 +219,48 @@ class ListScheduler { } } - std::list ready_list; + // Use a multimap to sort ReadyListEntry according to their priority. + std::multimap ready_queue; + + // Map of ready instructions to their iterators in ready_queue. + tensorflow::gtl::FlatMap::iterator> + ready_instructions; + + auto add_to_ready_queue = [&](HloInstruction* inst) { + auto entry = MakeReadyListEntry(inst); + auto it = ready_queue.emplace(GetPriority(entry), std::move(entry)); + ready_instructions[inst] = it; + }; + 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); + auto best_it = ready_queue.end(); + --best_it; + const HloInstruction* best = best_it->second.instruction; + ready_queue.erase(best_it); + 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 +268,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 +279,31 @@ 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 update them in the + // ready queue, so that they can appear earlier. + if (adjust_ready_queue) { + for (HloInstruction* operand : best->operands()) { + for (HloInstruction* operand_user : operand->users()) { + auto ready_instructions_it = ready_instructions.find(operand_user); + if (ready_instructions_it == ready_instructions.end()) { + continue; + } + auto ready_queue_it = ready_instructions_it->second; + auto& entry = ready_queue_it->second; + Priority new_priority = GetPriority(entry); + if (new_priority == ready_queue_it->first) { + continue; + } + // Create a new entry in ready_queue, then update + // ready_instructions[operand_user] to refer to the new entry. + ready_instructions_it->second = + ready_queue.emplace(new_priority, std::move(entry)); + // Remove the old entry in ready_queue. + ready_queue.erase(ready_queue_it); + } + } + } } CHECK_EQ(schedule.size(), computation_.instruction_count()); CHECK_EQ(scheduled_instructions_.size(), computation_.instruction_count()); @@ -280,15 +316,17 @@ 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. + // LogicalBuffer. We rely on iterator stability in this map, and that the map + // entries are std::pair's. std::unordered_map unscheduled_use_count_; // Set of instructions which have been scheduled. - std::unordered_set scheduled_instructions_; + tensorflow::gtl::FlatSet scheduled_instructions_; }; int64 SumLogicalBufferSizes( @@ -310,6 +348,7 @@ StatusOr> RunDFSMemoryScheduler( // simply users-1 for each instruction. By subtracting 1, we're saying that // instructions with no users or a single user don't count; instructions with // lots of fan-out will be visited earlier. + int64 cumulative_total_size = 0; tensorflow::gtl::FlatMap extra_users; tensorflow::gtl::FlatMap total_sizes; for (const HloInstruction* hlo : computation.MakeInstructionPostOrder()) { @@ -319,14 +358,17 @@ StatusOr> RunDFSMemoryScheduler( continue; } extra_users[hlo] = hlo->users().empty() ? 0 : hlo->users().size() - 1; - total_sizes[hlo] = SumLogicalBufferSizes( + int64 logical_buffer_size = SumLogicalBufferSizes( points_to_analysis.GetBuffersDefinedByInstruction(hlo), size_function); + total_sizes[hlo] = logical_buffer_size; + cumulative_total_size += logical_buffer_size; tensorflow::gtl::FlatSet unique_operands( hlo->operands().begin(), hlo->operands().end()); for (const HloInstruction* operand : unique_operands) { extra_users[hlo] += extra_users[operand]; total_sizes[hlo] += total_sizes[operand]; } + total_sizes[hlo] = std::min(total_sizes[hlo], cumulative_total_size); } CHECK_EQ(extra_users.size(), computation.instruction_count()); CHECK_EQ(total_sizes.size(), computation.instruction_count()); diff --git a/tensorflow/compiler/xla/service/hlo_sharding.cc b/tensorflow/compiler/xla/service/hlo_sharding.cc index 447c2446668253c932b44b51b2db22bfd47f9957..aa9ff89e983aa5d35a18906afca1c6e8eeaefa06 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding.cc +++ b/tensorflow/compiler/xla/service/hlo_sharding.cc @@ -183,6 +183,10 @@ Status HloSharding::ValidateTuple(const Shape& shape, int64 num_devices) const { // shape tree. ShapeTree shape_tree = GetAsShapeTree(shape); for (const auto& index_to_sharding : shape_tree.leaves()) { + if (index_to_sharding.first.empty()) { + // An empty tuple has a ShapeTree with a single leaf at the empty index. + continue; + } Status status = index_to_sharding.second.ValidateNonTuple( ShapeUtil::GetSubshape(shape, index_to_sharding.first), num_devices); if (!status.ok()) { @@ -222,7 +226,7 @@ Status HloSharding::ValidateNonTuple(const Shape& shape, Status status = Status::OK(); std::set seen_cores; tile_assignment_.Each( - [&](tensorflow::gtl::ArraySlice indices, uint32 core) { + [&](tensorflow::gtl::ArraySlice indices, int32 core) { // Don't overwrite a bad status, so we report the first error. if (status.ok()) { if (core >= num_devices) { @@ -344,4 +348,30 @@ OpSharding HloSharding::ToProto() const { return result; } +HloSharding HloSharding::TransformShardedTileShape( + const Shape& new_shape, + const std::function& transform) const { + CHECK(!IsTuple()); + if (IsTileMaximal()) { + return *this; + } + CHECK_EQ(ShapeUtil::Rank(new_shape), ShapeUtil::Rank(tile_shape())); + Shape new_tile_shape; + new_tile_shape.set_element_type(tile_shape().element_type()); + for (int64 i = 0; i < ShapeUtil::Rank(new_shape); ++i) { + int64 dim; + if (tile_assignment().dim(i) == 1) { + dim = new_shape.dimensions(i); + } else if (transform) { + dim = transform(i, tile_shape().dimensions(i)); + } else { + dim = tile_shape().dimensions(i); + } + new_tile_shape.add_dimensions(dim); + } + TF_CHECK_OK( + LayoutUtil::CopyLayoutBetweenShapes(tile_shape_, &new_tile_shape)); + return HloSharding::Tile(new_tile_shape, tile_assignment()); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_sharding.h b/tensorflow/compiler/xla/service/hlo_sharding.h index 7263198385cf0c84b1dac1e15177dcac99adaafb..e715dff9a0b8fcc2301a1581919dba384206923c 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding.h +++ b/tensorflow/compiler/xla/service/hlo_sharding.h @@ -207,6 +207,19 @@ class HloSharding { // REQUIRES: !IsReplicated() && !IsTuple() const Array& tile_assignment() const { return tile_assignment_; } + // Return a new sharding that can apply to the given new shape. + // If this sharding is tile-maximal, the returned sharding will be the same as + // this sharding. If this sharding is not tile-maximal, the returned + // sharding's tile size will differ: + // - Non-sharded dimensions will be adapted to be the same as `new_shape`; + // tile_dimension(i) = new_shape.dimensions(i); + // - Sharded dimensions will be kept the same unless `transform` is supplied + // in which case tile_dimension(i) = transform(i, tile_dimension(i)); + // REQUIRES: !IsTuple(). + HloSharding TransformShardedTileShape( + const Shape& new_shape, + const std::function& transform = nullptr) const; + private: HloSharding() : replicated_(true), diff --git a/tensorflow/compiler/xla/service/hlo_sharding_test.cc b/tensorflow/compiler/xla/service/hlo_sharding_test.cc index 0c7487b3ac77ff181d44dd55ebcf2608feaf02ea..07fc4687cc1c0518b3ab2a86c62464fc54082a01 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding_test.cc +++ b/tensorflow/compiler/xla/service/hlo_sharding_test.cc @@ -269,5 +269,18 @@ TEST_F(HloShardingTest, Hash) { } } +TEST_F(HloShardingTest, TransformShardedTileShapeTest) { + HloSharding sharding = + HloSharding::Tile(ShapeUtil::MakeShape(F32, {3, 5, 7, 11}), + Array4D({{{{0, 1}, {2, 3}}}})); + HloSharding result = sharding.TransformShardedTileShape( + ShapeUtil::MakeShape(F32, {13, 15, 17, 19}), + [](int dim, int value) { return dim * 111; }); + HloSharding expected = + HloSharding::Tile(ShapeUtil::MakeShape(F32, {13, 15, 222, 333}), + Array4D({{{{0, 1}, {2, 3}}}})); + EXPECT_EQ(result, expected); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 04d4656546684063d3d6532e443ad7995c6ea8db..8c875698eb1992719d504d272ca338b05b60e36b 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -13,6 +13,8 @@ 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/status_macros.h" #include "tensorflow/core/lib/core/errors.h" @@ -123,6 +125,10 @@ Status ShapeVerifier::HandleOutfeed(HloInstruction* outfeed) { return CheckShape(outfeed, ShapeUtil::MakeNil()); } +Status ShapeVerifier::HandleHostCompute(HloInstruction*) { + return tensorflow::Status::OK(); +} + Status ShapeVerifier::HandleRng(HloInstruction*) { return tensorflow::Status::OK(); } @@ -164,6 +170,8 @@ 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; @@ -178,6 +186,8 @@ Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { } 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(); @@ -359,13 +369,130 @@ Status ShapeVerifier::HandleBatchNormGrad(HloInstruction* batch_norm_grad) { batch_norm_grad->feature_index())); } +namespace { + +// 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 Status::OK(); +} + +} // namespace + +Status ShapeVerifier::HandleGather(HloInstruction* gather) { + return CheckShape( + gather, + ShapeInference::InferGatherShape( + gather->operand(0)->shape(), gather->operand(1)->shape(), + gather->gather_dimension_numbers(), gather->gather_window_bounds())); +} + Status ShapeVerifier::CheckShape(const HloInstruction* instruction, - const Shape& expected_shape) { - if (!ShapeUtil::Compatible(instruction->shape(), expected_shape)) { + 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); + } + } + if (!compatible) { return InvalidArgument( "Expected instruction to have shape compatible with %s, actual " "shape is %s:\n%s", - ShapeUtil::HumanString(expected_shape).c_str(), + ShapeUtil::HumanString(inferred_shape).c_str(), ShapeUtil::HumanString(instruction->shape()).c_str(), instruction->ToString().c_str()); } @@ -373,14 +500,14 @@ Status ShapeVerifier::CheckShape(const HloInstruction* instruction, } Status ShapeVerifier::CheckShape(const HloInstruction* instruction, - const StatusOr& expected_shape_status) { - if (!expected_shape_status.ok()) { - Status s = expected_shape_status.status(); + 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, expected_shape_status.ValueOrDie()); + return CheckShape(instruction, inferred_shape_status.ValueOrDie()); } Status ShapeVerifier::CheckUnaryShape(const HloInstruction* instruction) { @@ -635,11 +762,14 @@ StatusOr HloVerifier::Run(HloModule* module) { } else if (instruction->opcode() == HloOpcode::kBroadcast) { // If you see this failure then someone has confused the difference // between the HLO broadcast op, and the UserComputation broadcast - // op. See https://groups.google.com/forum/#!topic/xla-dev/9LqijHmTt_I + // op. See https://groups.google.com/forum/#!topic/xla-dev/9LqijHmTt_I // or ComputationLowerer::Visit() TF_RET_CHECK(instruction->dimensions().size() == ShapeUtil::Rank(instruction->operand(0)->shape())) - << "Broadcast HLO has invalid number of dimensions."; + << "Broadcast HLO (" << instruction->ToShortString() + << ") has invalid number of dimensions: " + << instruction->dimensions().size() + << " != " << ShapeUtil::Rank(instruction->operand(0)->shape()); } else if (instruction->opcode() == HloOpcode::kWhile) { auto* while_cond = instruction->while_condition(); auto* while_body = instruction->while_body(); diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index 26d53dec1e52f0bf19d6a8af998c56db8a850518..1dd7ec3c51e18dcfe89bd478de87798ba3858119 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -27,6 +27,10 @@ namespace xla { // 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; @@ -56,6 +60,7 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleFusion(HloInstruction*) override; Status HandleCall(HloInstruction* call) override; Status HandleCustomCall(HloInstruction*) override; + Status HandleHostCompute(HloInstruction*) override; Status HandleSlice(HloInstruction* slice) override; Status HandleDynamicSlice(HloInstruction* dynamic_slice) override; Status HandleDynamicUpdateSlice( @@ -75,20 +80,21 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleBatchNormInference( HloInstruction* batch_norm_inference) override; Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override; + Status HandleGather(HloInstruction* gather) override; Status FinishVisit(HloInstruction*) override { return tensorflow::Status::OK(); } protected: - // Check the instruction's shape against the given expected shape and return - // an appropriate error if there is a mismatch. + // 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& expected_shape); + const Shape& inferred_shape); // Overload which takes a StatusOr to reduce boilerplate in the caller. Status CheckShape(const HloInstruction* instruction, - const StatusOr& expected_shape_status); + const StatusOr& inferred_shape_status); // Check a unary (binary, etc) instruction's shape against the inferred shape. Status CheckUnaryShape(const HloInstruction* instruction); @@ -99,19 +105,32 @@ class ShapeVerifier : public DfsHloVisitor { // 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: + using ShapeVerifierFactory = std::function()>; + // Uses standard shape inference. explicit HloVerifier() - : shape_verifier_factory_([] { return MakeUnique(); }) {} + : 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( - std::function()> shape_verifier_factory) + explicit HloVerifier(ShapeVerifierFactory shape_verifier_factory) : shape_verifier_factory_(std::move(shape_verifier_factory)) {} ~HloVerifier() override = default; @@ -129,7 +148,7 @@ class HloVerifier : public HloPassInterface { // 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. - std::function()> shape_verifier_factory_; + ShapeVerifierFactory shape_verifier_factory_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index 90e1f0acdc4cdeda280dabaab2df66b181d0f407..d69ad80bdb4d2eab2d34228be026d7bc0b76efc0 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -102,6 +102,8 @@ namespace xla { case HloOpcode::kExp: case HloOpcode::kFft: case HloOpcode::kFusion: + case HloOpcode::kGather: + case HloOpcode::kHostCompute: case HloOpcode::kLog: case HloOpcode::kMap: case HloOpcode::kParameter: @@ -300,7 +302,7 @@ StatusOr InstructionFusion::Run(HloModule* module) { // Consider each operand of this instruction for fusion into this // instruction. We want to consider the operands in a particular order to - // avoid created duplicate instruction clones in the fusion instruction. + // avoid creating duplicate instruction clones in the fusion instruction. // For example, consider the following expression: // // A = ... @@ -375,7 +377,7 @@ StatusOr InstructionFusion::Run(HloModule* module) { changed = true; if (operand->user_count() == 0) { - // Operand is now dead. Remove from post order by setting it's + // Operand is now dead. Remove from post order by setting its // location to nullptr. post_order[FindOrDie(post_order_index, operand)] = nullptr; post_order_index.erase(operand); diff --git a/tensorflow/compiler/xla/service/interpreter/BUILD b/tensorflow/compiler/xla/service/interpreter/BUILD index 0819ab3b90b2360c6b0b2afaa89f322afe566eb3..0db3863f2428cf0c9a66a928d54f774e39a18539 100644 --- a/tensorflow/compiler/xla/service/interpreter/BUILD +++ b/tensorflow/compiler/xla/service/interpreter/BUILD @@ -63,10 +63,7 @@ cc_library( name = "platform_id", srcs = ["platform_id.cc"], hdrs = ["platform_id.h"], - deps = [ - "@nsync//:nsync_headers", - "//tensorflow/core:stream_executor_headers_lib", - ] + if_static( + deps = ["//tensorflow/core:stream_executor_headers_lib"] + if_static( ["@protobuf_archive//:protobuf"], ["@protobuf_archive//:protobuf_headers"], ), diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc index c83880e03090bc316b029ae9c902e27f3d8d826e..9171e859c6f84ceef9664aa1eb90a07c87dfab40 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.cc +++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc @@ -44,28 +44,12 @@ 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(); } diff --git a/tensorflow/compiler/xla/service/interpreter/executable.cc b/tensorflow/compiler/xla/service/interpreter/executable.cc index 0cb9b5d8107cd8bf468b07d5fe2a22930d9e8b8c..883063d0f075f5b0d79edc01bcd27a7c579272f4 100644 --- a/tensorflow/compiler/xla/service/interpreter/executable.cc +++ b/tensorflow/compiler/xla/service/interpreter/executable.cc @@ -93,7 +93,7 @@ StatusOr> InterpreterExecutable::ExecuteOnStream( TF_ASSIGN_OR_RETURN(std::unique_ptr result, transfer_manager->AllocateShapedBuffer( result_literal->shape(), run_options->allocator(), - run_options->device_ordinal())); + executor->device_ordinal())); TF_RETURN_IF_ERROR(transfer_manager->TransferLiteralToDevice( executor, *result_literal, *result)); diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 5413b95cfb6aad464da27c7b4aeaed5011e16393..39f9120e552f014dd2759bff2892157402d9c47a 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.cc +++ b/tensorflow/compiler/xla/service/layout_assignment.cc @@ -53,6 +53,83 @@ limitations under the License. namespace xla { +// For now moving only one API here, but we should have a single top level +// anonymous namespace, instead of three or four spread all over this file. +namespace { + +// Creates and returns a copy of the given instruction with a different +// layout. Tuple-shaped instructions will be deep-copied, and the last Tuple +// instruction producing the copy is returned. +StatusOr CreateCopyWithNewLayout( + const Shape& shape_with_layout, HloInstruction* instruction) { + TF_RET_CHECK(LayoutUtil::HasLayout(shape_with_layout)); + DCHECK(ShapeUtil::Compatible(shape_with_layout, instruction->shape())) + << ShapeUtil::HumanString(shape_with_layout) << " " + << ShapeUtil::HumanString(instruction->shape()) + << " instruction: " << instruction->ToString(); + + if (ShapeUtil::IsTuple(instruction->shape())) { + // Deep-copy tuples. + std::vector element_copies; + for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); + ++i) { + HloInstruction* gte = instruction->parent()->AddInstruction( + HloInstruction::CreateGetTupleElement( + ShapeUtil::GetSubshape(instruction->shape(), {i}), instruction, + i)); + + // Recurse to copy each elements. + TF_ASSIGN_OR_RETURN( + HloInstruction * element_copy, + CreateCopyWithNewLayout( + ShapeUtil::GetSubshape(shape_with_layout, {i}), gte)); + element_copies.push_back(element_copy); + } + // Gather element copies into a tuple with a new Tuple instruction. + HloInstruction* tuple_copy = instruction->parent()->AddInstruction( + HloInstruction::CreateTuple(element_copies)); + LayoutUtil::ClearLayout(tuple_copy->mutable_shape()); + TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( + shape_with_layout, tuple_copy->mutable_shape())); + return tuple_copy; + } else if (ShapeUtil::IsArray(instruction->shape())) { + HloInstruction* copy = + instruction->parent()->AddInstruction(HloInstruction::CreateUnary( + instruction->shape(), HloOpcode::kCopy, instruction)); + LayoutUtil::ClearLayout(copy->mutable_shape()); + TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( + shape_with_layout, copy->mutable_shape())); + + return copy; + } else { + return FailedPrecondition( + "Can only copy array and tuple shaped instructions"); + } +} + +// Creates a copy of the given operand if the operand's layout does not match +// the given layout. This copy replaces the use in the given instruction. Tuple +// operands will be deep-copied. +Status CopyOperandIfLayoutsDiffer(const ShapeLayout& operand_layout, + HloInstruction* instruction, + int64 operand_no) { + HloInstruction* operand = instruction->mutable_operand(operand_no); + TF_RET_CHECK(operand_layout.LayoutIsSet()); + TF_RET_CHECK(LayoutUtil::HasLayout(operand->shape())); + + if (ShapeUtil::Equal(operand_layout.shape(), operand->shape())) { + // Operand layout already matches our constraint. Nothing to do. + return Status::OK(); + } + + TF_ASSIGN_OR_RETURN(HloInstruction * operand_copy, + CreateCopyWithNewLayout(operand_layout.shape(), operand)); + + return instruction->ReplaceOperandWith(operand_no, operand_copy); +} + +} // namespace + std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint) { out << constraint.ToString(); @@ -61,8 +138,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,8 +151,8 @@ 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) { @@ -115,17 +192,34 @@ LayoutConstraints::LayoutConstraints( } } +PointsToSet::BufferSet* LayoutConstraints::GetBufferSet( + const HloInstruction* instruction) const { + auto it = buffer_sets_cache_.find(instruction); + if (it != buffer_sets_cache_.end()) { + return it->second.get(); + } + auto& buffer_set = + buffer_sets_cache_ + .emplace(instruction, MakeUnique()) + .first->second; + const auto& points_to_set = points_to_analysis_.GetPointsToSet(instruction); + points_to_set.ForEachElement( + [&buffer_set](const ShapeIndex& /*index*/, + const PointsToSet::BufferList& buffers) { + buffer_set->insert(buffers.begin(), buffers.end()); + }); + return buffer_set.get(); +} + bool LayoutConstraints::OperandBufferForwarded( const HloInstruction* instruction, int64 operand_no) const { // The operand is potentially forwarded if the intersection of points-to sets // of the operand and the instruction is non-empty. - auto output_buffers = - points_to_analysis_.GetPointsToSet(instruction).CreateFlattenedSet(); - auto operand_buffers = - points_to_analysis_.GetPointsToSet(instruction->operand(operand_no)) - .CreateFlattenedSet(); - for (const LogicalBuffer* output_buffer : output_buffers) { - if (operand_buffers.count(output_buffer) > 0) { + PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); + PointsToSet::BufferSet* operand_buffers = + GetBufferSet(instruction->operand(operand_no)); + for (const LogicalBuffer* output_buffer : *output_buffers) { + if (operand_buffers->count(output_buffer) > 0) { return true; } } @@ -134,7 +228,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); @@ -171,10 +265,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); @@ -188,7 +283,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); @@ -226,12 +322,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); @@ -240,16 +336,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); @@ -267,14 +364,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); @@ -290,8 +388,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 = @@ -300,7 +398,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(); } @@ -394,8 +492,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. @@ -434,8 +531,8 @@ Status LayoutAssignment::AddMandatoryConstraints( {0})); Shape new_shape = channel_constraints->LayoutShapeForChannel( recv_buffer_shape, instruction->channel_id()); - TF_RETURN_IF_ERROR(constraints->SetBufferLayout( - new_shape.layout(), *buffer, /*mandatory=*/true)); + TF_RETURN_IF_ERROR( + constraints->SetBufferLayout(new_shape.layout(), *buffer)); } } } @@ -457,7 +554,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 @@ -508,7 +605,36 @@ 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, + body_layout.result_shape(), instruction, 0)); + } else if (instruction->opcode() == HloOpcode::kConditional) { + // The layout of the true and false computations must match, and must + // be the layout of the kConditional instruction. + TF_RET_CHECK(instruction->operand_count() == 3); + + HloComputation* true_computation = instruction->true_computation(); + HloComputation* false_computation = instruction->false_computation(); + const HloInstruction* true_operand = instruction->operand(1); + const HloInstruction* false_operand = instruction->operand(2); + + TF_RET_CHECK(true_computation->num_parameters() == 1); + TF_RET_CHECK(false_computation->num_parameters() == 1); + ComputationLayout& true_computation_layout = + FindOrDie(computation_layouts_, true_computation); + ComputationLayout& false_computation_layout = + FindOrDie(computation_layouts_, false_computation); + + DCHECK(ShapeUtil::Compatible(true_operand->shape(), + true_computation_layout.parameter_shape(0))); + DCHECK(ShapeUtil::Compatible( + false_operand->shape(), false_computation_layout.parameter_shape(0))); + + TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( + true_computation_layout.result_shape(), instruction)); + TF_RETURN_IF_ERROR(constraints->SetOperandLayout( + true_computation_layout.parameter_shape(0), instruction, 1, + /*mandatory=*/true)); + TF_RETURN_IF_ERROR(constraints->SetOperandLayout( + false_computation_layout.parameter_shape(0), instruction, 2, /*mandatory=*/true)); } else if (instruction->opcode() == HloOpcode::kCustomCall) { if (!CustomCallRequiresMajorFirstLayout(instruction)) { @@ -533,7 +659,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)); } } } @@ -596,6 +722,33 @@ Status CheckWhileLayout(HloInstruction* while_inst, return Status::OK(); } +Status CheckConditionalLayout( + HloInstruction* instruction, + const ComputationLayout& true_computation_layout, + const ComputationLayout& false_computation_layout) { + HloComputation* true_computation = instruction->true_computation(); + HloComputation* false_computation = instruction->false_computation(); + const HloInstruction* true_operand = instruction->operand(1); + const HloInstruction* false_operand = instruction->operand(2); + + TF_RET_CHECK(true_computation_layout.result_layout() == + false_computation_layout.result_layout()); + TF_RET_CHECK(true_computation_layout.result_layout().MatchesLayoutInShape( + instruction->shape())); + TF_RET_CHECK(true_computation_layout.result_layout().MatchesLayoutInShape( + true_computation->root_instruction()->shape())); + TF_RET_CHECK(false_computation_layout.result_layout().MatchesLayoutInShape( + instruction->shape())); + TF_RET_CHECK(false_computation_layout.result_layout().MatchesLayoutInShape( + false_computation->root_instruction()->shape())); + TF_RET_CHECK(true_computation_layout.parameter_layout(0).MatchesLayoutInShape( + true_operand->shape())); + TF_RET_CHECK( + false_computation_layout.parameter_layout(0).MatchesLayoutInShape( + false_operand->shape())); + return Status::OK(); +} + // Fusion parameters must match the layout of the fusion instructions operands, // and the root of the fusion expression must match the layout of the fusion // instruction. @@ -708,6 +861,13 @@ Status LayoutAssignment::CheckLayouts(HloModule* module) { FindOrDie(computation_layouts_, instruction->while_condition()), FindOrDie(computation_layouts_, instruction->while_body()))); break; + case HloOpcode::kConditional: + TF_RETURN_IF_ERROR(CheckConditionalLayout( + instruction, + FindOrDie(computation_layouts_, instruction->true_computation()), + FindOrDie(computation_layouts_, + instruction->false_computation()))); + break; default: break; } @@ -907,7 +1067,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(); @@ -1159,77 +1323,6 @@ StatusOr InferArrayLayout( return *first_buffer_layout; } -// Creates and returns a copy of the given instruction with a different -// layout. Tuple-shaped instructions will be deep-copied, and the last Tuple -// instruction producing the copy is returned. -StatusOr CreateCopyWithNewLayout( - const Shape& shape_with_layout, HloInstruction* instruction) { - TF_RET_CHECK(LayoutUtil::HasLayout(shape_with_layout)); - DCHECK(ShapeUtil::Compatible(shape_with_layout, instruction->shape())) - << ShapeUtil::HumanString(shape_with_layout) << " " - << ShapeUtil::HumanString(instruction->shape()) - << " instruction: " << instruction->ToString(); - - if (ShapeUtil::IsTuple(instruction->shape())) { - // Deep-copy tuples. - std::vector element_copies; - for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); - ++i) { - HloInstruction* gte = instruction->parent()->AddInstruction( - HloInstruction::CreateGetTupleElement( - ShapeUtil::GetSubshape(instruction->shape(), {i}), instruction, - i)); - - // Recurse to copy each elements. - TF_ASSIGN_OR_RETURN( - HloInstruction * element_copy, - CreateCopyWithNewLayout( - ShapeUtil::GetSubshape(shape_with_layout, {i}), gte)); - element_copies.push_back(element_copy); - } - // Gather element copies into a tuple with a new Tuple instruction. - HloInstruction* tuple_copy = instruction->parent()->AddInstruction( - HloInstruction::CreateTuple(element_copies)); - LayoutUtil::ClearLayout(tuple_copy->mutable_shape()); - TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( - shape_with_layout, tuple_copy->mutable_shape())); - return tuple_copy; - } else if (ShapeUtil::IsArray(instruction->shape())) { - HloInstruction* copy = - instruction->parent()->AddInstruction(HloInstruction::CreateUnary( - instruction->shape(), HloOpcode::kCopy, instruction)); - LayoutUtil::ClearLayout(copy->mutable_shape()); - TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( - shape_with_layout, copy->mutable_shape())); - - return copy; - } else { - return FailedPrecondition( - "Can only copy array and tuple shaped instructions"); - } -} - -// Creates a copy of the given operand if the operand's layout does not match -// the given layout. This copy replaces the use in the given instruction. Tuple -// operands will be deep-copied. -Status CopyOperandIfLayoutsDiffer(const ShapeLayout& operand_layout, - HloInstruction* instruction, - int64 operand_no) { - HloInstruction* operand = instruction->mutable_operand(operand_no); - TF_RET_CHECK(operand_layout.LayoutIsSet()); - TF_RET_CHECK(LayoutUtil::HasLayout(operand->shape())); - - if (ShapeUtil::Equal(operand_layout.shape(), operand->shape())) { - // Operand layout already matches our constraint. Nothing to do. - return Status::OK(); - } - - TF_ASSIGN_OR_RETURN(HloInstruction * operand_copy, - CreateCopyWithNewLayout(operand_layout.shape(), operand)); - - return instruction->ReplaceOperandWith(operand_no, operand_copy); -} - // For fusion instructions, set the layout of each fused parameter instruction // to match the layout of its corresponding fusion instruction operand. Also, // set the layout of the fused root to match the layout of the fusion @@ -1390,7 +1483,7 @@ Status LayoutAssignment::RunOnComputation( // 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 @@ -1455,7 +1548,12 @@ 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->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 @@ -1463,17 +1561,21 @@ StatusOr LayoutAssignment::Run(HloModule* module) { // 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::kBitcast) { + // bitcasts are inherently layout sensitive and so a bitcast instruction + // present in the IR before layout assignment is a bug. + return InternalError( + "Unexpected bitcast operation seen during layout assignment: %s.", + instruction->ToString().c_str()); + } 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 if (computation->IsFusionComputation()) { - continue; } else { ComputationLayout computation_layout(computation->ComputeProgramShape()); // Setting all embedded computations to the default layout is potentially diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index 6bfae2998609c0482b91368f1891ce1e8e43fa23..680f88048a1f0cd5ede7991640003ef407d4facf 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -38,6 +38,7 @@ limitations under the License. #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/flatmap.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -46,7 +47,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 +56,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 +71,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 +92,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 +112,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 +165,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 @@ -190,6 +200,11 @@ class LayoutConstraints { string ToString() const; private: + // Find a bufferset in the bufferset cache. This is useful since we can + // currently create the flattened buffer set for the same instruction many + // times, which is often slow. + PointsToSet::BufferSet* GetBufferSet(const HloInstruction* instruction) const; + // The set of BufferLayoutConstraints applied to the computation. std::unordered_map buffer_constraints_; @@ -212,6 +227,10 @@ class LayoutConstraints { // Array-shaped buffers which have not yet been constrained. std::set unconstrained_buffer_ids_; + mutable tensorflow::gtl::FlatMap> + buffer_sets_cache_; + HloComputation* computation_; }; @@ -384,7 +403,6 @@ 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 @@ -392,6 +410,7 @@ class LayoutAssignment : public HloPassInterface { // handled before their caller instructions so the layouts of caller // instructions can be set to match the computation. std::map computation_layouts_; + ChannelLayoutConstraints* channel_layout_constraints_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc index e269a13459f1146f1d2952870399827d9e705e38..4b1c9bad41de8030cf14bc6d1c0db21b9c56c3bf 100644 --- a/tensorflow/compiler/xla/service/layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc @@ -590,6 +590,85 @@ TEST_F(LayoutAssignmentTest, TransposeToBitcastToUser) { transpose->shape(), {2, 3, 0, 1})); } +// TransposeIsBitcast shouldn't be called without layout information. +TEST_F(LayoutAssignmentTest, TransposeIsBitcastFail) { + auto builder = HloComputation::Builder(TestName()); + Shape input_shape = ShapeUtil::MakeShape(F32, {2, 2, 2}); + Shape input_shape_with_layout(input_shape); + *input_shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({2, 1, 0}); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape_with_layout, "param")); + auto hlo = builder.AddInstruction( + HloInstruction::CreateTranspose(input_shape, param, {0, 2, 1})); + // Clear the default layout assigned to the instruction. + LayoutUtil::ClearLayout(hlo->mutable_shape()); + EXPECT_DEATH(ShapeUtil::TransposeIsBitcast(hlo->operand(0)->shape(), + hlo->shape(), hlo->dimensions()), + "LayoutUtil::HasLayout"); +} + +// ReshapeIsBitcast shouldn't be called without layout information. +TEST_F(LayoutAssignmentTest, ReshapeIsBitcastFail) { + auto builder = HloComputation::Builder(TestName()); + Shape input_shape = ShapeUtil::MakeShape(F32, {2, 2, 2}); + Shape input_shape_with_layout(input_shape); + *input_shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({2, 1, 0}); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape_with_layout, "param")); + auto hlo = + builder.AddInstruction(HloInstruction::CreateReshape(input_shape, param)); + // Clear the default layout assigned to the instruction. + LayoutUtil::ClearLayout(hlo->mutable_shape()); + EXPECT_DEATH( + ShapeUtil::ReshapeIsBitcast(hlo->operand(0)->shape(), hlo->shape()), + "LayoutUtil::HasLayout"); +} + +// Check that the computation below doesn't crash the compiler. +// +// Within a fusion computation, only the parameters and result get assigned a +// layout. When we run the algebraic simplifier on this computation post layout +// assignment, it should not call TransposeIsBitcast on the `transpose` node +// inside the fusion computation as TransposeIsBitcast checks both input_shape +// and output_shape have layouts. +TEST_F(LayoutAssignmentTest, TransposeWithinFusionDoesNotCrash) { + const char* module_str = R"( + HloModule test_module + + fused_computation { + param_1 = f32[2,2,2]{2,1,0} parameter(1) + transpose = f32[2,2,2]{2,1,0} transpose(param_1), dimensions={0,2,1} + reduce_1 = f32[] parameter(0) + broadcast_1 = f32[2,2,2]{2,1,0} broadcast(reduce_1), dimensions={} + ROOT divide_1 = f32[2,2,2]{2,1,0} divide(transpose, broadcast_1) + } + + ENTRY entry_computation { + fusion.1 = f32[2,2,2]{2,1,0} parameter(1) + reduce.1 = f32[] parameter(0) + fusion.2 = f32[2,2,2]{2,1,0} fusion(reduce.1, fusion.1), kind=kLoop, calls=fused_computation + ROOT tuple.1 = (f32[2,2,2]{2,1,0}) tuple(fusion.2) + } + )"; + + auto module = tools::Parse(module_str).ValueOrDie(); + + module = + backend() + .compiler() + ->RunHloPasses(std::move(module), backend().default_stream_executor(), + /*device_allocator=*/nullptr) + .ConsumeValueOrDie(); + + EXPECT_EQ( + ::tensorflow::Status::OK(), + backend() + .compiler() + ->RunBackend(std::move(module), backend().default_stream_executor(), + /*device_allocator=*/nullptr) + .status()); +} + // 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) { @@ -629,33 +708,113 @@ TEST_F(LayoutAssignmentTest, GTEInheritsLayoutFromOperand) { 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; + auto layout_of = [&](tensorflow::StringPiece name) { + return FindInstruction(module.get(), name) + ->shape() + .layout() + .minor_to_major(); }; - 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(), + EXPECT_THAT(layout_of("gte0"), ElementsAre(0, 1, 2)); + EXPECT_THAT(layout_of("gte1a"), ElementsAre(1, 2, 0)); + EXPECT_THAT(layout_of("gte1b"), ElementsAre(2, 0, 1)); + EXPECT_THAT(layout_of("fresult"), ElementsAre(2, 1, 0)); + EXPECT_THAT(FindInstruction(module.get(), "gte1") + ->shape() + .tuple_shapes(0) + .layout() + .minor_to_major(), ElementsAre(1, 2, 0)); - EXPECT_THAT(fused_instr("gte1b")->shape().layout().minor_to_major(), + EXPECT_THAT(FindInstruction(module.get(), "gte1") + ->shape() + .tuple_shapes(1) + .layout() + .minor_to_major(), ElementsAre(2, 0, 1)); - EXPECT_THAT(fused_instr("fresult")->shape().layout().minor_to_major(), - ElementsAre(2, 1, 0)); +} + +TEST_F(LayoutAssignmentTest, ConditionalAsymmetricLayout) { + auto builder = HloComputation::Builder(TestName()); + auto module = CreateNewModule(); + Shape shape = ShapeUtil::MakeShape(F32, {128, 8}); + Shape tshape = ShapeUtil::MakeTupleShape({shape, shape}); + Shape result_tshape = ShapeUtil::MakeTupleShape({shape}); + + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, shape, "param1")); + auto pred = builder.AddInstruction(HloInstruction::CreateParameter( + 2, ShapeUtil::MakeShape(PRED, {}), "param2")); + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({param0, param1})); + + auto true_builder = HloComputation::Builder(TestName() + "_TrueBranch"); + { + auto param = true_builder.AddInstruction( + HloInstruction::CreateParameter(0, tshape, "param")); + auto gte0 = true_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, param, 0)); + auto gte1 = true_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, param, 1)); + auto add = true_builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, gte0, gte1)); + true_builder.AddInstruction(HloInstruction::CreateTuple({add})); + } + HloComputation* true_computation = + module->AddEmbeddedComputation(true_builder.Build()); + + auto false_builder = HloComputation::Builder(TestName() + "_FalseBranch"); + { + Shape xshape = ShapeUtil::MakeShapeWithLayout(F32, {128, 8}, {0, 1}); + false_builder.AddInstruction( + HloInstruction::CreateParameter(0, tshape, "param")); + // Using infeed as layout assignment does not mess up with it. + auto infeed = + false_builder.AddInstruction(HloInstruction::CreateInfeed(xshape, "")); + false_builder.AddInstruction(HloInstruction::CreateTuple({infeed})); + } + HloComputation* false_computation = + module->AddEmbeddedComputation(false_builder.Build()); + builder.AddInstruction(HloInstruction::CreateConditional( + result_tshape, pred, tuple, true_computation, tuple, false_computation)); + + HloComputation* computation = module->AddEntryComputation(builder.Build()); + ComputationLayout computation_layout(computation->ComputeProgramShape()); + + AssignLayouts(module.get(), &computation_layout); + + const HloInstruction* true_root = true_computation->root_instruction(); + const HloInstruction* false_root = false_computation->root_instruction(); + EXPECT_THAT(true_root->opcode(), HloOpcode::kTuple); + EXPECT_THAT(false_root->opcode(), HloOpcode::kTuple); + + const HloInstruction* true_result = true_root->operand(0); + const HloInstruction* false_result = false_root->operand(0); + EXPECT_TRUE(LayoutUtil::Equal(true_result->shape().layout(), + false_result->shape().layout())); + EXPECT_THAT(false_result->opcode(), HloOpcode::kCopy); +} + +TEST_F(LayoutAssignmentTest, InternalErrorOnBitcast) { + auto builder = HloComputation::Builder(TestName()); + auto constant0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR2WithLayout( + {{1.0, 2.0}, {3.0, 4.0}}, LayoutUtil::MakeLayout({0, 1})))); + builder.AddInstruction(HloInstruction::CreateUnary( + constant0->shape(), HloOpcode::kBitcast, constant0)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + ComputationLayout computation_layout( + module->entry_computation()->ComputeProgramShape()); + LayoutAssignment layout_assignment(&computation_layout); + Status error_status = layout_assignment.Run(module.get()).status(); + EXPECT_FALSE(error_status.ok()); + EXPECT_THAT( + error_status.error_message(), + ::testing::HasSubstr( + "Unexpected bitcast operation seen during layout assignment")); } } // namespace diff --git a/tensorflow/compiler/xla/service/liveness_util_test.cc b/tensorflow/compiler/xla/service/liveness_util_test.cc index 2c2a02f6375343d67dfb155bbb03729ff6e490d2..f8b309488eeb5391b1cad5db760934ec1f7e3521 100644 --- a/tensorflow/compiler/xla/service/liveness_util_test.cc +++ b/tensorflow/compiler/xla/service/liveness_util_test.cc @@ -35,8 +35,7 @@ class PointsToAnalysisTestBase : public HloTestBase { CHECK_NOTNULL(module_.get()); points_to_analysis_ = TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); - dataflow_analysis_ = - HloDataflowAnalysis::Run(module_.get()).ConsumeValueOrDie(); + dataflow_analysis_ = HloDataflowAnalysis::Run(*module_).ConsumeValueOrDie(); } void BuildModuleAndRunAnalysis(std::unique_ptr computation) { diff --git a/tensorflow/compiler/xla/service/llvm_compiler.cc b/tensorflow/compiler/xla/service/llvm_compiler.cc index f98fc0400a7d827a29dcddc5eecf9a4a01e76590..911b243fe28a5baf8a4b8ed752b892265f5388ac 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.cc +++ b/tensorflow/compiler/xla/service/llvm_compiler.cc @@ -14,12 +14,29 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/llvm_compiler.h" +#include "tensorflow/core/platform/denormal.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, DeviceMemoryAllocator* device_allocator) { + // Tensorflow tries to enable the following behaviors in all its threads: + // + // - Denormals are zero (DAZ): roughly, operations treat denormal floats as + // zero. + // - Flush denormals to zero (FTZ): roughly, operations produce zero instead + // of denormal floats. + // + // In theory enabling these shouldn't matter since the compiler should ideally + // not leak its environment into generated code, but we turn off DAZ and FTZ + // to get some defense-in-depth. + tensorflow::port::ScopedDontFlushDenormal dont_flush_denormals; + std::vector> result; for (size_t i = 0; i < modules.size(); i++) { if (stream_execs[i].size() != 1) { diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index ffc78bd5cfac3df1001d8125327607c85169ae92..37261ed1e665ebed9685751161a412ad114a9e96 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/BUILD +++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD @@ -54,6 +54,7 @@ cc_library( "@llvm//:core", "@llvm//:support", "@llvm//:target", + "@llvm//:transform_utils", ], ) 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 6384c7f46f5ebbedaeda232b40095611a5d738a4..3312a888443233139841ce7a5e3173f907605e1d 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc @@ -29,18 +29,13 @@ limitations under the License. namespace xla { namespace llvm_ir { -IrArray::Index::Index(llvm::Value* linear, const Shape& shape, - llvm::IRBuilder<>* ir_builder) - : multidim_(ShapeUtil::Rank(shape)), - linear_(linear), - layout_(shape.layout()), - dims_(shape.dimensions().begin(), shape.dimensions().end()) { - CHECK(LayoutUtil::HasLayout(shape)) - << "Shape " << ShapeUtil::HumanStringWithLayout(shape) - << " should have a layout."; +static void Delinearize(std::vector* multidim, + llvm::Value* linear, const Shape& shape, + llvm::IRBuilder<>* ir_builder) { int64 divisor = 1; - for (int64 i = 0; i < layout_.minor_to_major_size(); ++i) { - int64 dimension = layout_.minor_to_major(i); + const Layout& layout = shape.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); // If i is not the last dimension, compute @@ -54,16 +49,28 @@ IrArray::Index::Index(llvm::Value* linear, const Shape& shape, // 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( + 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; + (*multidim)[dimension] = quot; } divisor *= size_of_current_dimension; } } +IrArray::Index::Index(llvm::Value* linear, const Shape& shape, + llvm::IRBuilder<>* ir_builder) + : multidim_(ShapeUtil::Rank(shape)), + linear_(linear), + layout_(shape.layout()), + dims_(shape.dimensions().begin(), shape.dimensions().end()) { + CHECK(LayoutUtil::HasLayout(shape)) + << "Shape " << ShapeUtil::HumanStringWithLayout(shape) + << " should have a layout."; + Delinearize(&multidim_, linear, shape, ir_builder); +} + IrArray::Index::Index(tensorflow::gtl::ArraySlice multidim, llvm::Value* linear, const Shape& shape) : multidim_(multidim.begin(), multidim.end()), @@ -83,7 +90,6 @@ IrArray::Index::Index(tensorflow::gtl::ArraySlice multidim, dims_(shape.dimensions().begin(), shape.dimensions().end()) { CHECK_EQ(shape.dimensions_size(), multidim.size()); CHECK(LayoutUtil::HasLayout(shape)); - linear_ = Linearize(AsInt64Slice(shape.dimensions()), ir_builder); } IrArray::IrArray(llvm::Value* base_ptr, const Shape& shape) @@ -106,16 +112,13 @@ IrArray::IrArray(llvm::Value* base_ptr, const Shape& shape) } } -// Returns whether given linear index valid on given shape. +// Returns whether the given linear index is valid on the given shape. bool IrArray::Index::LinearValidOnShape(const Shape& a) const { - auto b = ShapeUtil::MakeShape(PRED /* irrelevant */, dims_); + auto b = ShapeUtil::MakeShape(a.element_type(), dims_); *b.mutable_layout() = layout_; return linear_ != nullptr && - ContainersEqual( - ShapeUtil::StripDegenerateDimensions(a).dimensions(), - ShapeUtil::StripDegenerateDimensions(b).dimensions()) && - LayoutUtil::Equal(ShapeUtil::StripDegenerateDimensions(a).layout(), - ShapeUtil::StripDegenerateDimensions(b).layout()); + ShapeUtil::ElementsIn(a) == ShapeUtil::ElementsIn(b) && + ShapeUtil::ReshapeIsBitcast(a, b); } IrArray::Index IrArray::Index::SourceIndexOfReshape( @@ -160,7 +163,8 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape( } } - if (linear() != nullptr && + if (linear() != nullptr && LayoutUtil::HasLayout(input_shape) && + LayoutUtil::HasLayout(output_shape) && ShapeUtil::ReshapeIsBitcast(input_shape, output_shape)) { return Index(source_multidim_index, linear(), input_shape); } @@ -195,13 +199,111 @@ IrArray::Index IrArray::Index::SourceIndexOfTranspose( llvm::IRBuilder<>* builder) const { std::vector operand_multidim_index = Permute(dimension_mapping, multidim()); - if (linear() != nullptr && + + if (linear() != nullptr && LayoutUtil::HasLayout(operand_shape) && + LayoutUtil::HasLayout(shape) && ShapeUtil::TransposeIsBitcast(operand_shape, shape, dimension_mapping)) { return Index(operand_multidim_index, linear(), operand_shape); } + return Index(operand_multidim_index); } +IrArray::Index IrArray::Index::SourceIndexOfBitcast( + const Shape& shape, const Shape& operand_shape, + llvm::IRBuilder<>* builder) const { + CHECK(LayoutUtil::HasLayout(shape) && LayoutUtil::HasLayout(operand_shape)); + // In case the bitcast is just a reshape, we can use SourceIndexOfReshape() + // instead. This will reuse linear() if possible, so we don't have to build a + // new 'linear_index'. + if (ShapeUtil::ReshapeIsBitcast(operand_shape, shape)) { + return SourceIndexOfReshape(shape, operand_shape, builder); + } + + // First linearize the index coming from the output of the bitcast. We want + // the physical index of the element in the buffer. This is like Linearize, + // but takes the layout into account. + int64 scale = 1; + llvm::Value* linear_index = builder->getInt64(0); + for (auto dimension : LayoutUtil::MinorToMajor(shape)) { + linear_index = builder->CreateAdd( + linear_index, + builder->CreateMul(multidim_[dimension], builder->getInt64(scale), "", + /*HasNUW=*/true, /*HasNSW=*/true), + "", /*HasNUW=*/true, /*HasNSW=*/true); + scale *= shape.dimensions(dimension); + } + + // Now delinearize it for the input of the bitcast. + std::vector multi_index(operand_shape.dimensions_size()); + Delinearize(&multi_index, linear_index, operand_shape, builder); + + return Index(multi_index, linear_index, operand_shape); +} + +IrArray::Index IrArray::Index::SourceIndexOfBroadcast( + const Shape& shape, const Shape& operand_shape, + tensorflow::gtl::ArraySlice dimension_mapping, + llvm::IRBuilder<>* builder) const { + int64 rank = ShapeUtil::Rank(operand_shape); + std::vector source_index(rank); + for (int64 i = 0; i < rank; ++i) { + source_index[i] = multidim_[dimension_mapping[i]]; + } + if (linear_ == nullptr || !LayoutUtil::HasLayout(operand_shape) || + !LayoutUtil::HasLayout(shape)) { + return Index(source_index); + } + // High-level idea: we can reuse the linear index if the broadcasted + // dimensions are contiguous, and this part of the operation is a bitcast. + // The other dimensions can be masked out with a div and a mod operation. + std::vector logical_to_physical = + LayoutUtil::MakeLogicalToPhysical(shape.layout()); + int64 output_rank = ShapeUtil::Rank(shape); + // The minimum physical dimension that is broadcasted. + int64 min_broadcasted_dimension = output_rank; + // The maximum physical dimension that is broadcasted. + int64 max_broadcasted_dimension = -1; + for (int64 i = 0; i < rank; ++i) { + int64 physical_dim = logical_to_physical[dimension_mapping[i]]; + min_broadcasted_dimension = + std::min(min_broadcasted_dimension, physical_dim); + max_broadcasted_dimension = + std::max(max_broadcasted_dimension, physical_dim); + } + bool contiguous_broadcast_dimensions = + max_broadcasted_dimension - min_broadcasted_dimension == rank - 1; + if (!contiguous_broadcast_dimensions) { + return Index(source_index); + } + // Check if the mapped dimensions are a bitcast. + std::vector operand_logical_to_physical = + LayoutUtil::MakeLogicalToPhysical(operand_shape.layout()); + for (int64 i = 0; i < rank; ++i) { + if (operand_logical_to_physical[i] != + logical_to_physical[dimension_mapping[i]] - min_broadcasted_dimension) { + return Index(source_index); + } + } + llvm::Value* linear = linear_; + int64 divisor = 1; + for (int64 i = max_broadcasted_dimension + 1; i < output_rank; ++i) { + divisor *= shape.dimensions(LayoutUtil::Major(shape.layout(), i)); + } + if (divisor > 1) { + linear = builder->CreateUDiv(linear, builder->getInt64(divisor)); + } + if (min_broadcasted_dimension > 0) { + int64 mod = 1; + for (int64 i = min_broadcasted_dimension; i <= max_broadcasted_dimension; + ++i) { + mod *= shape.dimensions(LayoutUtil::Major(shape.layout(), i)); + } + linear = builder->CreateURem(linear, builder->getInt64(mod)); + } + return Index(source_index, linear, operand_shape); +} + llvm::Value* IrArray::Index::Linearize( tensorflow::gtl::ArraySlice dimensions, llvm::IRBuilder<>* builder) const { diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h index 387d4629125cbb791840e943013188d14159908a..06cfb2a36c56c5fdece7140e469379f8394111fa 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h @@ -76,8 +76,7 @@ class IrArray { llvm::IRBuilder<>* ir_builder); // Constructs an index from the given multi-dimensional index and the shape - // that it indexes into. Also, computes the linear index according to - // "shape". + // that it indexes into. // // Precondition: "shape" has a layout. Index(tensorflow::gtl::ArraySlice multidim, @@ -134,6 +133,18 @@ class IrArray { tensorflow::gtl::ArraySlice dimension_mapping, llvm::IRBuilder<>* builder) const; + // Given that "this" is the target index of a bitcast from `operand_shape` + // to `shape`, returns the source index. + Index SourceIndexOfBitcast(const Shape& shape, const Shape& operand_shape, + llvm::IRBuilder<>* builder) const; + + // Given that "this" is the target index of a broadcast from `operand_shape` + // to `shape` with the given dimension mapping, returns the source index. + Index SourceIndexOfBroadcast( + const Shape& shape, const Shape& operand_shape, + tensorflow::gtl::ArraySlice dimension_mapping, + llvm::IRBuilder<>* builder) const; + // Linearizes the index into the given shape, i.e. reshapes it to rank-1 and // returns the index into the sole dimension 0 of the new shape. llvm::Value* Linearize(tensorflow::gtl::ArraySlice dimensions, diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc index 8d1e6338e189a055ac20f09961a783b52600866d..2a282f3be79f847a6569416794d1a2a3fcd69148 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); @@ -94,8 +106,10 @@ llvm::Value* EmitFloatMax(llvm::Value* lhs_value, llvm::Value* rhs_value, auto cmp = ir_builder->CreateFCmpUGE(lhs_value, rhs_value); return ir_builder->CreateSelect(cmp, lhs_value, rhs_value); } else { - return EmitCallToIntrinsic(llvm::Intrinsic::maxnum, {lhs_value, rhs_value}, - {lhs_value->getType()}, ir_builder); + auto cmp_ge = ir_builder->CreateFCmpOGE(lhs_value, rhs_value); + auto lhs_is_nan = ir_builder->CreateFCmpUNE(lhs_value, lhs_value); + auto sel_lhs = ir_builder->CreateOr(cmp_ge, lhs_is_nan); + return ir_builder->CreateSelect(sel_lhs, lhs_value, rhs_value); } } @@ -105,8 +119,10 @@ llvm::Value* EmitFloatMin(llvm::Value* lhs_value, llvm::Value* rhs_value, auto cmp = ir_builder->CreateFCmpULE(lhs_value, rhs_value); return ir_builder->CreateSelect(cmp, lhs_value, rhs_value); } else { - return EmitCallToIntrinsic(llvm::Intrinsic::minnum, {lhs_value, rhs_value}, - {lhs_value->getType()}, ir_builder); + auto cmp_le = ir_builder->CreateFCmpOLE(lhs_value, rhs_value); + auto lhs_is_nan = ir_builder->CreateFCmpUNE(lhs_value, lhs_value); + auto sel_lhs = ir_builder->CreateOr(cmp_le, lhs_is_nan); + return ir_builder->CreateSelect(sel_lhs, lhs_value, rhs_value); } } @@ -672,6 +688,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) { @@ -686,13 +715,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, 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/service.cc b/tensorflow/compiler/xla/service/service.cc index a57b7e571788f0d3e3ffda31eb94aa81ee85b574..0becc9d8f8ed22b2d7174b76ce775efec4b646f5 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -44,6 +44,7 @@ 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/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -231,10 +232,14 @@ tensorflow::Status Service::ValidateResultShapeWithLayout( return ShapeUtil::ValidateShape(shape_with_layout); } -StatusOr> Service::ResolveAndValidateArguments( +StatusOr>> +Service::ResolveAndValidateArguments( tensorflow::gtl::ArraySlice arguments, - int device_ordinal) { - std::vector shaped_buffers; + tensorflow::gtl::ArraySlice + stream_executors) { + CHECK_EQ(options_.number_of_replicas(), stream_executors.size()); + std::vector> replicated_arguments; + replicated_arguments.resize(options_.number_of_replicas()); for (size_t i = 0; i < arguments.size(); ++i) { auto buffer_status = allocation_tracker_.Resolve(*arguments[i]); if (!buffer_status.ok()) { @@ -242,22 +247,25 @@ StatusOr> Service::ResolveAndValidateArguments( StrCat(buffer_status.status().error_message(), ", ", "failed to resolve allocation for parameter ", i)); } - const ShapedBuffer* shaped_buffer = buffer_status.ValueOrDie(); - - // Verify allocation is same platform and device as the execution. - if (shaped_buffer->platform() != execute_backend_->platform() || - shaped_buffer->device_ordinal() != device_ordinal) { - return InvalidArgument( - "argument %lu is on device %s:%d but computation will be executed " - "on device %s", - i, shaped_buffer->platform()->Name().c_str(), - shaped_buffer->device_ordinal(), - execute_backend_->device_name(device_ordinal).c_str()); + auto replicated_buffers = buffer_status.ValueOrDie(); + CHECK_EQ(options_.number_of_replicas(), replicated_buffers.size()); + for (int replica = 0; replica < options_.number_of_replicas(); ++replica) { + const ShapedBuffer* shaped_buffer = replicated_buffers[replica]; + int replica_device_ordinal = stream_executors[replica]->device_ordinal(); + // Verify allocation is same platform and device as the execution. + if (shaped_buffer->platform() != execute_backend_->platform() || + shaped_buffer->device_ordinal() != replica_device_ordinal) { + return InvalidArgument( + "argument %lu is on device %s:%d but computation will be executed " + "on device %s", + i, shaped_buffer->platform()->Name().c_str(), + shaped_buffer->device_ordinal(), + execute_backend_->device_name(replica_device_ordinal).c_str()); + } + replicated_arguments[replica].push_back(shaped_buffer); } - - shaped_buffers.push_back(shaped_buffer); } - return shaped_buffers; + return replicated_arguments; } StatusOr> Service::CreateModuleConfig( @@ -306,8 +314,6 @@ StatusOr> Service::CreateModuleConfig( if (execution_options != nullptr) { config->set_seed(execution_options->seed()); config->set_debug_options(execution_options->debug_options()); - config->enable_hlo_profiling( - execution_options->debug_options().xla_hlo_profile()); } else { config->set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); } @@ -489,7 +495,8 @@ StatusOr> Service::BuildAndCacheExecutable( StatusOr> Service::ExecuteParallelAndRegisterResult( tensorflow::gtl::ArraySlice executables, - tensorflow::gtl::ArraySlice> arguments, + tensorflow::gtl::ArraySlice>> + arguments, Backend* backend, tensorflow::gtl::ArraySlice device_handles, tensorflow::gtl::ArraySlice result_tags, ExecutionProfile* profile) { @@ -512,6 +519,8 @@ Service::ExecuteParallelAndRegisterResult( for (int64 i = 0; i < executables.size(); i++) { // Stream executors for the replicas of the current computation. TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handles[i])); + CHECK_EQ(replicas.size(), arguments[i].size()); + std::vector> result_buffers; for (int64 replica = 0; replica < replicas.size(); ++replica) { TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, backend->BorrowStream(replicas[replica])); @@ -544,23 +553,20 @@ Service::ExecuteParallelAndRegisterResult( backend->StreamBorrower()); // Asynchronously launch the computation. - TF_ASSIGN_OR_RETURN( - std::unique_ptr result, - executables[i]->ExecuteAsyncOnStream(&run_options, arguments[i])); + TF_ASSIGN_OR_RETURN(std::unique_ptr result, + executables[i]->ExecuteAsyncOnStream( + &run_options, arguments[i][replica])); if (replica == 0 && profile != nullptr) { streams.back()->ThenStopTimer(timers.back().get()); } - // All replicas share the same device address for the result allocation, - // so only one of the replicas need to register the result handle. - if (replica == 0) { - TF_ASSIGN_OR_RETURN( - GlobalDataHandle handle, - allocation_tracker_.Register(std::move(result), result_tags[i])); - result_handles.push_back(handle); - } + result_buffers.emplace_back(std::move(result)); } + TF_ASSIGN_OR_RETURN(GlobalDataHandle handle, + allocation_tracker_.RegisterReplicatedBuffers( + std::move(result_buffers), result_tags[i])); + result_handles.push_back(handle); } // Wait for all executions to complete. @@ -626,9 +632,9 @@ Service::ExecuteParallelAndRegisterResult( StatusOr Service::ExecuteAndRegisterResult( Executable* executable, - const tensorflow::gtl::ArraySlice arguments, - Backend* backend, perftools::gputools::StreamExecutor* executor, - const string& result_tag, ExecutionProfile* profile) { + const tensorflow::gtl::ArraySlice> + arguments, + Backend* backend, const string& result_tag, ExecutionProfile* profile) { // Set up streams. std::vector::SmartPtr> streams; @@ -661,21 +667,26 @@ StatusOr Service::ExecuteAndRegisterResult( backend->inter_op_thread_pool()); } - std::unique_ptr result; if (options_.number_of_replicas() == 1) { - TF_ASSIGN_OR_RETURN(result, executable->ExecuteOnStreamWrapper( - &run_options[0], profile, arguments)); - } else { - // TODO(b/69985541): Support profiling also on this path. - std::vector> - repeated_arguments(options_.number_of_replicas(), arguments); - - TF_ASSIGN_OR_RETURN(auto results, executable->ExecuteOnStreams( - run_options, repeated_arguments)); - TF_RET_CHECK(!results.empty()); - result = std::move(results[0]); + TF_ASSIGN_OR_RETURN( + auto result, executable->ExecuteOnStreamWrapper(&run_options[0], + profile, arguments[0])); + return allocation_tracker_.Register(std::move(result), result_tag); + } + + // TODO(b/69985541): Support profiling also on this path. + + std::vector> + replicated_arguments; + for (const auto& arg : arguments) { + replicated_arguments.emplace_back(arg); } - return allocation_tracker_.Register(std::move(result), result_tag); + + TF_ASSIGN_OR_RETURN(auto results, executable->ExecuteOnStreams( + run_options, replicated_arguments)); + TF_RET_CHECK(!results.empty()); + return allocation_tracker_.RegisterReplicatedBuffers(std::move(results), + result_tag); } tensorflow::Status Service::SetReturnValue(const SetReturnValueRequest* arg, @@ -689,7 +700,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, ExecuteParallelResponse* result) { VLOG(1) << "running execute-parallel request: " << arg->ShortDebugString(); - std::vector> all_arguments; + std::vector>> all_arguments; std::vector> all_executors; std::vector versioned_handles; std::vector> module_configs; @@ -717,6 +728,14 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, return FailedPrecondition( "device handles must be given to execute parallel computations"); } + if (arg->requests_size() > 1 && + execution_options.device_handles_size() > 1) { + return InvalidArgument( + "Parallel requests with multiple device handles is not supported. " + "Found %d parallel requests, with request %lld containing %d device " + "handles.", + arg->requests_size(), i, execution_options.device_handles_size()); + } std::vector executors; for (const auto& device_handle : execution_options.device_handles()) { TF_ASSIGN_OR_RETURN(auto replicas, @@ -746,22 +765,26 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, // In the case of partitioned computations, assume all arguments go on the // zeroth core. TF_ASSIGN_OR_RETURN( - std::vector arguments, - ResolveAndValidateArguments(request.arguments(), - executors[0]->device_ordinal())); + auto replicas, + Replicas(*execute_backend_, execution_options.device_handles(0))); + TF_ASSIGN_OR_RETURN( + std::vector> replicated_arguments, + ResolveAndValidateArguments(request.arguments(), replicas)); // Create an HloModuleConfig object for the computation, given the shape of - // the program and the argument allocations. + // the program and the argument allocations. Here, we care only about the + // shapes of the arguments, so, it is sufficient to use the arguments of + // replica 0. TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(*program_shape, arguments, + CreateModuleConfig(*program_shape, replicated_arguments.front(), request.execution_options(), *user_computation)); VLOG(3) << "ExecuteParallel created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); // Adds to the vectors to build and execute the computations after the loop. - all_arguments.push_back(arguments); - all_arguments.insert(all_arguments.end(), executors.size() - 1, {}); + all_arguments.push_back(replicated_arguments); + all_arguments.insert(all_arguments.end(), executors.size() - 1, {{}}); versioned_handles.push_back(versioned_handle); module_configs.push_back(std::move(module_config)); computation_names.insert(computation_names.end(), executors.size(), @@ -860,15 +883,18 @@ tensorflow::Status Service::Execute(const ExecuteRequest* arg, std::shared_ptr program_shape, user_computation->ComputeProgramShape(versioned_handle.version)); + TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, + SingleComputationDeviceHandle())); TF_ASSIGN_OR_RETURN( - std::vector arguments, - ResolveAndValidateArguments(arg->arguments(), - execute_backend_->default_device_ordinal())); + std::vector> replicated_arguments, + ResolveAndValidateArguments(arg->arguments(), replicas)); + // Since we care only about the shapes of the arguments, it is sufficient to + // use the arguments of replica 0. TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(*program_shape, arguments, arg->execution_options(), - *user_computation)); + CreateModuleConfig(*program_shape, replicated_arguments.front(), + arg->execution_options(), *user_computation)); VLOG(3) << "Execute created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); @@ -884,20 +910,21 @@ tensorflow::Status Service::Execute(const ExecuteRequest* arg, executable->session_module()->set_execution_platform( execute_backend_->platform()->Name()); TF_RETURN_IF_ERROR(RecordArguments( - arguments, execute_backend_->default_stream_executor(), + replicated_arguments.front(), + execute_backend_->default_stream_executor(), execute_backend_->transfer_manager(), executable->session_module())); } TF_ASSIGN_OR_RETURN( *result->mutable_output(), ExecuteAndRegisterResult( - executable.get(), arguments, execute_backend_.get(), - execute_backend_->default_stream_executor(), + executable.get(), replicated_arguments, execute_backend_.get(), "result of " + user_computation->name(), result->mutable_profile())); if (executable->dumping()) { - TF_ASSIGN_OR_RETURN(const ShapedBuffer* result_buffer, - allocation_tracker_.Resolve(result->output())); + TF_ASSIGN_OR_RETURN( + const ShapedBuffer* result_buffer, + allocation_tracker_.ResolveForReplica(result->output(), 0)); TF_RETURN_IF_ERROR(RecordResult( *result_buffer, execute_backend_->default_stream_executor(), execute_backend_->transfer_manager(), executable->session_module())); @@ -908,6 +935,11 @@ tensorflow::Status Service::Execute(const ExecuteRequest* arg, return tensorflow::Status::OK(); } +tensorflow::Status Service::ExecuteGraph(const ExecuteGraphRequest* /*arg*/, + ExecuteResponse* /*result*/) { + return Unimplemented("execute-graph is not yet implemented"); +} + tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, ExecuteAsyncResponse* result) { VLOG(1) << "running execute-async request: " << arg->ShortDebugString(); @@ -925,15 +957,17 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, std::shared_ptr program_shape, user_computation->ComputeProgramShape(versioned_handle.version)); + TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, + SingleComputationDeviceHandle())); + TF_RET_CHECK(!replicas.empty()); TF_ASSIGN_OR_RETURN( - std::vector arguments, - ResolveAndValidateArguments(arg->arguments(), - execute_backend_->default_device_ordinal())); + std::vector> replicated_arguments, + ResolveAndValidateArguments(arg->arguments(), replicas)); TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(*program_shape, arguments, arg->execution_options(), - *user_computation)); + CreateModuleConfig(*program_shape, replicated_arguments.front(), + arg->execution_options(), *user_computation)); VLOG(3) << "ExecuteAsync created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); @@ -946,21 +980,17 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, versioned_handle, std::move(module_config), execute_backend_.get(), execute_backend_->default_stream_executor(), &profile)); - TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, - SingleComputationDeviceHandle())); - TF_RET_CHECK(!replicas.empty()); - // Set up streams. std::vector::SmartPtr> streams; - for (se::StreamExecutor* executor : replicas) { TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, execute_backend_->BorrowStream(executor)); streams.push_back(std::move(stream)); } - std::unique_ptr result_buffer; - for (const Pool::SmartPtr& stream : streams) { + std::vector> result_buffers; + for (size_t i = 0; i < streams.size(); ++i) { + const auto& stream = streams[i]; ExecutableRunOptions options; options.set_stream(stream.get()); options.set_allocator(execute_backend_->memory_allocator()); @@ -971,20 +1001,17 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, ServiceExecutableRunOptions service_options( options, execute_backend_->StreamBorrower()); - TF_ASSIGN_OR_RETURN( - std::unique_ptr this_result_buffer, - executable->ExecuteAsyncOnStream(&service_options, arguments)); + TF_ASSIGN_OR_RETURN(std::unique_ptr this_result_buffer, + executable->ExecuteAsyncOnStream( + &service_options, replicated_arguments[i])); - // Take the first result. - if (result_buffer == nullptr) { - result_buffer = std::move(this_result_buffer); - } + result_buffers.emplace_back(std::move(this_result_buffer)); } TF_ASSIGN_OR_RETURN( GlobalDataHandle output, - allocation_tracker_.Register(std::move(result_buffer), - "result of " + user_computation->name())); + allocation_tracker_.RegisterReplicatedBuffers( + std::move(result_buffers), "result of " + user_computation->name())); *result->mutable_execution() = execution_tracker_.Register( execute_backend_.get(), std::move(streams), profile, output); @@ -1012,7 +1039,7 @@ tensorflow::Status Service::WaitForExecution(const WaitForExecutionRequest* arg, tensorflow::Status Service::TransferToClient(const TransferToClientRequest* arg, TransferToClientResponse* result) { TF_ASSIGN_OR_RETURN(const ShapedBuffer* shaped_buffer, - allocation_tracker_.Resolve(arg->data())); + allocation_tracker_.ResolveForReplica(arg->data(), 0)); const Shape* return_shape; if (arg->has_shape_with_layout()) { @@ -1073,37 +1100,24 @@ tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, replicas, Replicas(*execute_backend_, SingleComputationDeviceHandle())); } - // All memory allocation is done on the first replica. The allocations in all - // other replicas mirror the firsts'. - int master_device_ordinal = replicas[0]->device_ordinal(); - TF_ASSIGN_OR_RETURN( - std::unique_ptr shaped_buffer, - execute_backend_->transfer_manager()->AllocateShapedBuffer( - shape, execute_backend_->memory_allocator(), master_device_ordinal)); - - // Transfer the data to the replicas. + // Allocate memory in each replica and transfer the data to all replicas. + std::vector> replicated_buffers; for (se::StreamExecutor* executor : replicas) { - if (executor->device_ordinal() == master_device_ordinal) { - TF_RETURN_IF_ERROR( - execute_backend_->transfer_manager()->TransferLiteralToDevice( - 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 - // TransferLiteralToDevice verifies that the device ordinal of the shaped - // buffer matches that of the executor. - std::unique_ptr clone = - CloneShapedBufferOnDevice(*shaped_buffer, executor->device_ordinal()); - TF_RETURN_IF_ERROR( - execute_backend_->transfer_manager()->TransferLiteralToDevice( - executor, *literal, *clone)); - } + TF_ASSIGN_OR_RETURN( + std::unique_ptr shaped_buffer, + execute_backend_->transfer_manager()->AllocateShapedBuffer( + shape, execute_backend_->memory_allocator(), + executor->device_ordinal())); + TF_RETURN_IF_ERROR( + execute_backend_->transfer_manager()->TransferLiteralToDevice( + executor, *literal, *shaped_buffer)); + replicated_buffers.emplace_back(std::move(shaped_buffer)); } - TF_ASSIGN_OR_RETURN( - *result->mutable_data(), - allocation_tracker_.Register(std::move(shaped_buffer), - StrCat("TransferToServer literal of shape ", - ShapeUtil::HumanString(shape)))); + TF_ASSIGN_OR_RETURN(*result->mutable_data(), + allocation_tracker_.RegisterReplicatedBuffers( + std::move(replicated_buffers), + StrCat("TransferToServer literal of shape ", + ShapeUtil::HumanString(shape)))); return tensorflow::Status::OK(); } @@ -1286,7 +1300,7 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, tensorflow::Status Service::GetShape(const GetShapeRequest* arg, GetShapeResponse* result) { TF_ASSIGN_OR_RETURN(const ShapedBuffer* buffer, - allocation_tracker_.Resolve(arg->data())); + allocation_tracker_.ResolveForReplica(arg->data(), 0)); *result->mutable_shape() = buffer->on_host_shape(); return tensorflow::Status::OK(); } @@ -1445,6 +1459,9 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { case OpRequest::kFftRequest: handle_status = computation->AddFftInstruction(arg->fft_request()); break; + case OpRequest::kGatherRequest: + handle_status = computation->AddGatherInstruction(arg->gather_request()); + break; case OpRequest::kGetTupleElementRequest: handle_status = computation->AddGetTupleElementInstruction( arg->get_tuple_element_request()); @@ -1453,9 +1470,13 @@ 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::kHostComputeRequest: + handle_status = + computation->AddHostComputeInstruction(arg->host_compute_request()); + break; case OpRequest::kMapRequest: { TF_ASSIGN_OR_RETURN( UserComputation * to_apply, @@ -1548,8 +1569,10 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { case OpRequest::kSendRequest: { TF_RETURN_IF_ERROR( channel_tracker_.RegisterSend(arg->send_request().channel_handle())); - TF_RETURN_IF_ERROR(computation->AddSendInstruction(arg->send_request())); - return tensorflow::Status::OK(); + // Send does not return a value, but we need a handle to be able to + // set OpMetadata and OpSharding (device assignment). + handle_status = computation->AddSendInstruction(arg->send_request()); + break; } case OpRequest::kRecvRequest: { TF_RETURN_IF_ERROR( diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h index 6ce241971156599aaa25aea1b0caac0e1bd5379c..96352d9096e6aeeb33f84c7b6fc42c28820e5e84 100644 --- a/tensorflow/compiler/xla/service/service.h +++ b/tensorflow/compiler/xla/service/service.h @@ -112,6 +112,12 @@ class Service : public ServiceInterface { tensorflow::Status Execute(const ExecuteRequest* arg, ExecuteResponse* result) override; + // Executes a computation with the provided global data passed as + // immutable arguments. The request contains the whole computation graph. + // Returns global data output and execution timing. + tensorflow::Status ExecuteGraph(const ExecuteGraphRequest* arg, + ExecuteResponse* result) override; + // Executes one or more computations in parallel with the provided global data // passed as immutable arguments. Returns global data output for each // computation. @@ -265,11 +271,14 @@ class Service : public ServiceInterface { static StatusOr> CreateComputeConstantBackend(); // Resolves the given argument handles in the allocation tracker and returns - // the corresponding allocations. The function also verifies that each - // allocation matches the execution platform and device ordinal. - StatusOr> ResolveAndValidateArguments( + // the corresponding allocations for every replica. The function also verifies + // that each allocation matches the execution platform and device ordinal of + // the corresponding replica. + StatusOr>> + ResolveAndValidateArguments( tensorflow::gtl::ArraySlice arguments, - int device_ordinal); + tensorflow::gtl::ArraySlice + stream_executors); // Create a Hlo module config for the given program shape and arguments. // execution_options is optional; if not given a default is used. @@ -314,16 +323,17 @@ class Service : public ServiceInterface { // ExecutionProfile object which will be filled in with profile data. StatusOr ExecuteAndRegisterResult( Executable* executable, - const tensorflow::gtl::ArraySlice arguments, - Backend* backend, perftools::gputools::StreamExecutor* executor, - const string& result_tag, ExecutionProfile* profile); + const tensorflow::gtl::ArraySlice> + arguments, + Backend* backend, const string& result_tag, ExecutionProfile* profile); // Runs the given executables with the given arguments and register the result // from each executable in the allocation tracker. The handles of the result // from the tracker are returned. StatusOr> ExecuteParallelAndRegisterResult( tensorflow::gtl::ArraySlice executables, - tensorflow::gtl::ArraySlice> arguments, + tensorflow::gtl::ArraySlice>> + arguments, Backend* backend, tensorflow::gtl::ArraySlice device_handles, tensorflow::gtl::ArraySlice result_tags, diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index 4ba6da6ccc44be8f3c70d2af80b30f0b2e388c2a..8c8bd6d73ad41db7d609ac91c7bdfc4703f364e1 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -169,11 +169,11 @@ bool AllUnique(tensorflow::gtl::ArraySlice slice) { tensorflow::Status ExpectNotTupleOrOpaque(const Shape& shape, tensorflow::StringPiece op_type) { if (ShapeUtil::IsTuple(shape)) { - return InvalidArgument("Expected non-tuple argument for %s. Got: %s", + return InvalidArgument("Expected non-tuple argument for %s, but got %s.", op_type.ToString().c_str(), ShapeUtil::HumanString(shape).c_str()); } else if (ShapeUtil::IsOpaque(shape)) { - return InvalidArgument("Expected non-opaque argument for %s. Got: %s", + return InvalidArgument("Expected non-opaque argument for %s, but got %s.", op_type.ToString().c_str(), ShapeUtil::HumanString(shape).c_str()); } else { @@ -193,8 +193,10 @@ tensorflow::Status VerifyReducerShape(const ProgramShape& reducer_shape, const Shape& accumulator_shape = reducer_shape.result(); if (ShapeUtil::Rank(accumulator_shape) != 0) { - return Unimplemented( - "Reduction function currently must have rank-0 result."); + return InvalidArgument( + "Reduction function must have rank 0 (rank %lld reduction function " + "given).", + ShapeUtil::Rank(accumulator_shape)); } // Check that the accumulator can be passed in as the first argument. @@ -209,7 +211,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", @@ -220,8 +223,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", @@ -231,10 +234,11 @@ 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", + "Reduction function's second parameter shape must " + "match the result shape, but got %s vs %s.", ShapeUtil::HumanString(reducer_shape.parameters(1)).c_str(), ShapeUtil::HumanString(accumulator_shape).c_str()); } @@ -256,29 +260,29 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, for (int64 i = 0; i < window.dimensions_size(); ++i) { const auto& dim = window.dimensions(i); if (dim.size() <= 0) { - return InvalidArgument("Window has a non-positive dimension. Window: %s", + return InvalidArgument("Window %s has a non-positive dimension.", window.DebugString().c_str()); } if (dim.stride() <= 0) { - return InvalidArgument("Window has a non-positive stride. Window: %s", + return InvalidArgument("Window %s has a non-positive stride.", window.DebugString().c_str()); } if (!allow_negative_padding && dim.padding_low() < 0) { - return InvalidArgument("Window has a negative low padding. Window: %s", + return InvalidArgument("Window %s has a negative low padding.", window.DebugString().c_str()); } if (!allow_negative_padding && dim.padding_high() < 0) { - return InvalidArgument("Window has a negative high padding. Window: %s", + return InvalidArgument("Window %s has a negative high padding.", window.DebugString().c_str()); } if (dim.base_dilation() < 1) { return InvalidArgument( - "Window has a non-positive base area dilation factor. Window: %s", + "Window %s has a non-positive base area dilation factor.", window.DebugString().c_str()); } if (dim.window_dilation() < 1) { return InvalidArgument( - "Window has a non-positive window dilation factor. Window: %s", + "Window %s has a non-positive window dilation factor.", window.DebugString().c_str()); } @@ -318,8 +322,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, case UNOP_CEIL: if (!ShapeUtil::ElementIsFloating(arg)) { return InvalidArgument( - "expected element type in shape to be floating for floor/ceil " - "operation; got %s", + "Expected element type in shape to be floating for floor/ceil " + "operation; got %s.", PrimitiveType_Name(arg.element_type()).c_str()); } return arg; @@ -331,8 +335,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, if (!ShapeUtil::ElementIsFloating(arg) && !ShapeUtil::ElementIsComplex(arg)) { return InvalidArgument( - "expected element type in shape to be floating or complex for " - "sin/cos/exp/log/tanh operation; got %s", + "Expected element type in shape to be floating or complex for " + "sin/cos/exp/log/tanh operation; got %s.", PrimitiveType_Name(arg.element_type()).c_str()); } return arg; @@ -340,8 +344,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, case UNOP_IMAG: if (!ShapeUtil::ElementIsComplex(arg)) { return InvalidArgument( - "expected element type in shape to be complex for real/imag " - "operation; got %s", + "Expected element type in shape to be complex for real/imag " + "operation; got %s.", PrimitiveType_Name(arg.element_type()).c_str()); } return ShapeUtil::ChangeElementType(arg, F32); @@ -361,8 +365,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, if (arg.element_type() != PRED && !primitive_util::IsIntegralType(arg.element_type())) { return InvalidArgument( - "expected pred or an integral element type in argument to not " - "operation; got %s", + "Expected pred or an integral element type in argument to Not " + "operation; got %s.", PrimitiveType_Name(arg.element_type()).c_str()); } return arg; @@ -370,8 +374,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, case UNOP_IS_FINITE: if (!ShapeUtil::ElementIsFloating(arg)) { return InvalidArgument( - "expected element type in shape to be floating point for IsFinite " - "operation; got %s", + "Expected element type in shape to be floating point for IsFinite " + "operation; got %s.", PrimitiveType_Name(arg.element_type()).c_str()); } return ShapeUtil::ChangeElementType(arg, PRED); @@ -387,31 +391,33 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, tensorflow::gtl::ArraySlice arg_shapes, const int64 dimension) { if (arg_shapes.empty()) { - return InvalidArgument("Concatenate expects at least one argument"); + return InvalidArgument("Concatenate expects at least one argument."); } if (dimension < 0 || dimension >= ShapeUtil::Rank(*arg_shapes[0])) { - return InvalidArgument("dimension to concatenate along out of bounds: %lld", + return InvalidArgument("Concatenate dimension out of bounds: %lld.", 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)) { return InvalidArgument( "Cannot concatenate arrays with different ranks: %lld (%s) vs %lld " - "(%s)", + "(%s).", ShapeUtil::Rank(*arg_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", + "Cannot concatenate arrays with different element types: %s vs %s.", PrimitiveType_Name(arg_shape->element_type()).c_str(), PrimitiveType_Name(shape->element_type()).c_str()); } @@ -424,13 +430,14 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, // concatenating. } return InvalidArgument( - "cannot concatenate arrays that differ in dimensions other than " + "Cannot concatenate arrays that differ in dimensions other than " "the one being concatenated (the other array dimensions must be " - "the same): %s vs %s in dimension %lld", + "the same): %s vs %s in dimension %lld.", ShapeUtil::HumanString(*arg_shape).c_str(), ShapeUtil::HumanString(*shape).c_str(), dimension); } } + element_type = ShapeUtil::HigherPrecisionElementType(*shape, *arg_shape); } std::vector new_dimensions(arg_shape->dimensions().begin(), @@ -438,7 +445,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( @@ -447,7 +454,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, if (primitive_util::IsComplexType(old_element_type) && !primitive_util::IsComplexType(new_element_type)) { return Unimplemented( - "Unsupported conversion from complex to real type: %s => %s", + "Conversion from complex to real type %s => %s is not implemented.", ShapeUtil::HumanString(operand_shape).c_str(), PrimitiveType_Name(new_element_type).c_str()); } @@ -456,7 +463,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, // future, by recursing into the tuple elements to check all sub-conversions // are valid. For now we just reject them, though. return InvalidArgument( - "cannot convert from or to tuple type; requested conversion: %s => %s", + "Convert does not allow tuples, so cannot convert from %s to %s.", ShapeUtil::HumanString(operand_shape).c_str(), PrimitiveType_Name(new_element_type).c_str()); } @@ -469,24 +476,23 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, auto old_element_type = operand_shape.element_type(); if (primitive_util::IsComplexType(old_element_type) != primitive_util::IsComplexType(new_element_type)) { - return Unimplemented( - "Unsupported conversion between real and complex types: %s => %s", - ShapeUtil::HumanString(operand_shape).c_str(), - PrimitiveType_Name(new_element_type).c_str()); + return InvalidArgument("Conversion from complex to real type %s => %s.", + ShapeUtil::HumanString(operand_shape).c_str(), + PrimitiveType_Name(new_element_type).c_str()); } if (ShapeUtil::IsTuple(operand_shape) || new_element_type == TUPLE) { // Note: we may want to support tuple conversions via this operation in the // future, by recursing into the tuple elements to check all sub-conversions // are valid. For now we just reject them, though. return InvalidArgument( - "cannot convert from or to tuple type; requested conversion: %s => %s", + "Cannot convert from or to tuple type; requested conversion: %s => %s.", ShapeUtil::HumanString(operand_shape).c_str(), PrimitiveType_Name(new_element_type).c_str()); } if (primitive_util::BitWidth(old_element_type) != primitive_util::BitWidth(new_element_type)) { return InvalidArgument( - "cannot bitcast types with different bit-widths: %s => %s", + "Cannot bitcast types with different bit-widths: %s => %s.", PrimitiveType_Name(old_element_type).c_str(), PrimitiveType_Name(new_element_type).c_str()); } @@ -499,20 +505,20 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, const int mantissa_bits) { if (!ShapeUtil::ElementIsFloating(operand_shape)) { return InvalidArgument( - "expected element type in shape to be floating point for " - "ReducePrecision operation; got %s", + "Expected element type in shape to be floating point for " + "ReducePrecision operation; got %s.", PrimitiveType_Name(operand_shape.element_type()).c_str()); } if (exponent_bits < 1) { // One exponent bit is necessary to distinguish 0 from infinity. Having // no exponent bits doesn't produce a sensible number, so we require at // least one. - return InvalidArgument("expected exponent_bits >= 1; got %d", + return InvalidArgument("Expected exponent_bits >= 1; got %d.", exponent_bits); } if (mantissa_bits < 0) { // A number with no mantissa bits is still meaningful, however. - return InvalidArgument("expected non-negative mantissa_bits; got %d", + return InvalidArgument("Expected non-negative mantissa_bits; got %d.", mantissa_bits); } return operand_shape; @@ -523,22 +529,23 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, const PaddingConfig& padding_config) { if (ShapeUtil::IsTuple(operand_shape)) { return InvalidArgument( - "pad operation does not support tuple-shape operands"); + "Pad operation does not support tuple-shape operands."); } if (!ShapeUtil::IsScalar(padding_value_shape)) { return InvalidArgument( - "pad operation does not support non-scalar padding values"); + "Pad operation does not support non-scalar padding values."); } if (ShapeUtil::Rank(operand_shape) != padding_config.dimensions_size()) { return InvalidArgument( "The rank of the operand and the padding configuration do not match: " - "%s vs %s", + "%s vs %s.", 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"); + "The element types of the operands to Pad do not match."); } std::vector dimensions(ShapeUtil::Rank(operand_shape)); for (int64 i = 0; i < operand_shape.dimensions_size(); ++i) { @@ -548,7 +555,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: @@ -597,7 +606,7 @@ Status ValidateDotDimensionNumbers( lhs_batch_dimensions) || !dims_in_range(ShapeUtil::Rank(rhs), rhs_contracting_dimensions, rhs_batch_dimensions)) { - return InvalidArgument("A dimension number is out of range in dot: %s", + return InvalidArgument("A dimension number is out of range in Dot: %s.", dimension_numbers.DebugString().c_str()); } @@ -615,7 +624,7 @@ Status ValidateDotDimensionNumbers( if (!dims_unique(lhs_contracting_dimensions, lhs_batch_dimensions) || !dims_unique(rhs_contracting_dimensions, rhs_batch_dimensions)) { - return InvalidArgument("A dimension number is not unique in dot: %s", + return InvalidArgument("A dimension number is not unique in Dot: %s.", dimension_numbers.DebugString().c_str()); } @@ -633,8 +642,7 @@ Status ValidateDotDimensionNumbers( rhs_non_contracting_non_batch_dims < 0 || rhs_non_contracting_non_batch_dims > 1) { return InvalidArgument( - "batch and contracting dimension number mismatch " - "with rank "); + "Batch and contracting dimension number mismatch with rank."); } // Check that batch dimension numbers are ordered before all others, and @@ -646,7 +654,7 @@ Status ValidateDotDimensionNumbers( !std::equal(batch_dim_numbers.begin(), batch_dim_numbers.end(), rhs_batch_dimensions.begin())) { return InvalidArgument( - "batch dimension numbers must precede non-batch dimensions and be" + "Batch dimension numbers must precede non-batch dimensions and be" "monotonically increasing."); } @@ -663,22 +671,22 @@ Status ValidateDotDimensionNumbers( auto fail = [lhs, rhs](const string& addendum) -> Status { string message = tensorflow::strings::Printf( - "cannot infer shape for dot operation: %s %s", + "Cannot infer shape for dot operation: %s %s.", ShapeUtil::HumanString(lhs).c_str(), ShapeUtil::HumanString(rhs).c_str()); if (!addendum.empty()) { - message += ": " + addendum; + message += " " + addendum; } return InvalidArgument("%s", message.c_str()); }; // Check if both element types are the same. - if (lhs.element_type() != rhs.element_type()) { - return fail("element types do not match"); + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) { + return fail("Element types do not match."); } if ((ShapeUtil::Rank(lhs) < 1) || (ShapeUtil::Rank(rhs) < 1)) { - return fail("dot only supports rank 1 or above."); + return fail("Dot only supports rank 1 or above."); } // Validate basic properties of dot dimension numbers. @@ -688,7 +696,7 @@ Status ValidateDotDimensionNumbers( if (dimension_numbers.lhs_contracting_dimensions_size() != dimension_numbers.rhs_contracting_dimensions_size() || dimension_numbers.lhs_contracting_dimensions_size() != 1) { - return fail("must specify one contracting dimension for both lhs and rhs."); + return fail("Must specify one contracting dimension for both lhs and rhs."); } // Check that contracting dimension sizes match. @@ -698,13 +706,13 @@ Status ValidateDotDimensionNumbers( dimension_numbers.rhs_contracting_dimensions(0); if (lhs.dimensions(lhs_contracting_dimension) != rhs.dimensions(rhs_contracting_dimension)) { - return fail("contracting dimension sizes do not match."); + return fail("Contracting dimension sizes do not match."); } // Check that number of batch dimensions match. if (dimension_numbers.lhs_batch_dimensions_size() != dimension_numbers.rhs_batch_dimensions_size()) { - return fail("must the same number of batch dimensions for lhs and rhs."); + return fail("Must the same number of batch dimensions for lhs and rhs."); } // Check that batch dimension numbers and sizes match. @@ -713,7 +721,7 @@ Status ValidateDotDimensionNumbers( dimension_numbers.rhs_batch_dimensions(i) || lhs.dimensions(dimension_numbers.lhs_batch_dimensions(i)) != rhs.dimensions(dimension_numbers.rhs_batch_dimensions(i))) { - return fail("batch dimension numbers and sizes must match for lhs/rhs."); + return fail("Batch dimension numbers and sizes must match for lhs/rhs."); } } @@ -736,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); @@ -761,13 +770,15 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } else if (rhs.dimensions(i) == 1) { output_dimensions[i] = lhs.dimensions(i); } else { - return InvalidArgument("binary op %s with incompatible shapes: %s and %s", - BinaryOperation_Name(operation).c_str(), - ShapeUtil::HumanString(lhs).c_str(), - ShapeUtil::HumanString(rhs).c_str()); + return InvalidArgument( + "Binary op %s with incompatible shapes: %s and %s.", + BinaryOperation_Name(operation).c_str(), + ShapeUtil::HumanString(lhs).c_str(), + 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( @@ -778,15 +789,15 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // Reject "magic" inference for binops on different shapes, requiring // the user to provide an explicit broadcast dimension in this case. // See b/25177275 for more details. - return InvalidArgument("automatic shape inference not supported: %s and %s", + return InvalidArgument("Automatic shape inference not supported: %s and %s", ShapeUtil::HumanString(smaller_shape).c_str(), ShapeUtil::HumanString(larger_shape).c_str()); } else if (broadcast_dimensions.size() != ShapeUtil::Rank(smaller_shape)) { return InvalidArgument( - "size of broadcast_dimensions has to match lower-rank operand's " + "Size of broadcast_dimensions has to match lower-rank operand's " "rank; " " lower-rank operand's rank is %lld, size of broadcast_dimensions is " - "%zu", + "%zu.", ShapeUtil::Rank(smaller_shape), broadcast_dimensions.size()); } @@ -829,18 +840,20 @@ 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); if (dimension_to_match < 0) { return InvalidArgument( - "broadcast dimension number (%lld) cannot be negative", + "Broadcast dimension number (%lld) cannot be negative.", dimension_to_match); } if (dimension_to_match >= larger_shape.dimensions_size()) { return InvalidArgument( - "broadcast dimension number (%lld) too large; higher-rank " - "operand has rank %d", + "Broadcast dimension number (%lld) too large; higher-rank " + "operand has rank %d.", dimension_to_match, larger_shape.dimensions_size()); } int64 small_dimension_size = smaller_shape.dimensions(i); @@ -851,7 +864,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (small_dimension_size != large_dimension_size && small_dimension_size != 1 && large_dimension_size != 1) { return InvalidArgument( - "broadcast dimension %d mismatch: %lld != %lld; %s and %s", i, + "Broadcast dimension %d mismatch: %lld != %lld; %s and %s.", i, small_dimension_size, large_dimension_size, ShapeUtil::HumanString(smaller_shape).c_str(), ShapeUtil::HumanString(larger_shape).c_str()); @@ -860,7 +873,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // order. if (i > 0 && broadcast_dimensions.at(i - 1) >= dimension_to_match) { return InvalidArgument( - "broadcast dimensions order is wrong: %lld comes after %lld", + "Broadcast dimensions order is wrong: %lld comes after %lld.", dimension_to_match, broadcast_dimensions.at(i - 1)); } @@ -878,9 +891,9 @@ 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", + "Binary op %s with different element types: %s and %s.", BinaryOperation_Name(operation).c_str(), ShapeUtil::HumanString(lhs).c_str(), ShapeUtil::HumanString(rhs).c_str()); @@ -892,15 +905,16 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (!broadcast_dimensions.empty() && broadcast_dimensions != identity_dims) { return InvalidArgument( - "broadcast dimensions field must either be not set or be the " - "identity on binary operations with operands of the same rank"); + "Broadcast dimensions field must either be not set or be the " + "identity on binary operations with operands of the same rank."); } } - 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)) { @@ -930,6 +944,13 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( rhs->shape(), /*broadcast_dimensions=*/{}); } +/* static */ StatusOr ShapeInference::InferBinaryOpShape( + HloOpcode opcode, const Shape& lhs, const Shape& rhs, + tensorflow::gtl::ArraySlice broadcast_dimensions) { + return InferBinaryOpShape(OpcodeToBinaryOperation(opcode), lhs, rhs, + broadcast_dimensions); +} + /* static */ StatusOr ShapeInference::InferBinaryOpShape( BinaryOperation operation, const Shape& lhs, const Shape& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions) { @@ -966,17 +987,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( case BINOP_COMPLEX: { if (!ShapeUtil::ElementIsFloating(lhs)) { return InvalidArgument( - "expected element type in shape to be floating for complex compose " - "operation; got %s", + "Expected element type in shape to be floating for complex compose " + "operation; got %s.", PrimitiveType_Name(lhs.element_type()).c_str()); } 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"); + return Unimplemented("Complex component type is not implemented."); } } case BINOP_AND: @@ -984,8 +1005,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (lhs.element_type() != PRED && !primitive_util::IsIntegralType(lhs.element_type())) { return InvalidArgument( - "expected pred or integral type in argument to and/or operation; " - "got %s", + "Expected pred or integral type in argument to and/or operation; " + "got %s.", PrimitiveType_Name(lhs.element_type()).c_str()); } return InferElementwiseBinaryOpShape(operation, lhs, rhs, @@ -1003,7 +1024,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } default: return Unimplemented( - "not yet implemented; infer binary op shape: %s; lhs: %s; rhs: %s", + "Binary op shape inference: %s; lhs: %s; rhs: %s is not implemented.", BinaryOperation_Name(operation).c_str(), lhs.ShortDebugString().c_str(), rhs.ShortDebugString().c_str()); } @@ -1028,7 +1049,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( case TRIOP_SELECT: return InferSelectShape(lhs, rhs, ehs); default: - return InvalidArgument("unknown operation %s", + return InvalidArgument("Unknown operation %s.", TernaryOperation_Name(operation).c_str()); } } @@ -1059,7 +1080,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return result; } default: - return InvalidArgument("unknown operation %s", + return InvalidArgument("Unknown operation %s.", VariadicOperation_Name(operation).c_str()); } } @@ -1069,7 +1090,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const ProgramShape& to_apply, tensorflow::gtl::ArraySlice dimensions) { if (arg_shapes.empty()) { - return InvalidArgument("Map expects at least one argument"); + return InvalidArgument("Map expects at least one argument."); } // All arguments must have the same shape. @@ -1078,12 +1099,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; } @@ -1099,7 +1121,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } return InvalidArgument( "Map operation requires all operands to have the same shape; got: " - "%s", + "%s.", Join(pieces, ", ").c_str()); } @@ -1108,7 +1130,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (dimensions.size() != arg_shape->dimensions_size()) { return InvalidArgument( "Map applied to a subset of dimensions currently not supported: " - "arg_dimension_size: %d, requested_map_dimensions_size: %zu", + "arg_dimension_size: %d, requested_map_dimensions_size: %zu.", arg_shape->dimensions_size(), dimensions.size()); } @@ -1116,7 +1138,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( for (int i = 0; i < dimensions.size(); ++i) { if (dimensions[i] != i) { return InvalidArgument( - "Map requires monotonically increasing dimension numbers, found: %s ", + "Map requires monotonically increasing dimension numbers; got: %s.", Join(dimensions, ", ").c_str()); } } @@ -1125,7 +1147,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (arg_shapes.size() != to_apply.parameters_size()) { return InvalidArgument( "Map applied function arity must match number of arguments; got: " - "arity: %d, arguments: %zu", + "arity: %d, arguments: %zu.", to_apply.parameters_size(), arg_shapes.size()); } @@ -1133,8 +1155,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const Shape& output_shape = to_apply.result(); if (!ShapeUtil::IsScalar(output_shape)) { return InvalidArgument( - "mapped computation's result has to be a scalar; " - "got: %s", + "Mapped computation's result has to be a scalar; got: %s.", ShapeUtil::HumanString(output_shape).c_str()); } @@ -1143,15 +1164,16 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (!ShapeUtil::IsScalar(parameter_shape)) { return InvalidArgument( - "mapped computation's parameter has to be a scalar; " - "got parameter %d shape: %s", + "Mapped computation's parameter has to be a scalar; " + "got parameter %d shape: %s.", 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", + "Mapped computation's parameter type has to match argument element " + "type; got parameter %d shape: %s, argument shape: %s.", i, ShapeUtil::HumanString(parameter_shape).c_str(), ShapeUtil::HumanString(*arg_shape).c_str()); } @@ -1182,21 +1204,21 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Expected feature_index of batch-norm-training to be " "smaller than the rank of operand_shape; " - "got feature_index %lld, and rank %lld", + "got feature_index %lld, and rank %lld.", feature_index, ShapeUtil::Rank(operand_shape)); } if (feature_index < 0) { return InvalidArgument( "Expected feature_index of batch-norm-training to " - "be a non-negative number, got %lld", + "be a non-negative number, got %lld.", feature_index); } if (ShapeUtil::Rank(operand_shape) < 1) { return InvalidArgument( "Expected the rank of operand to " - "batch-norm-training to be at least 1; got %lld", + "batch-norm-training to be at least 1; got %lld.", ShapeUtil::Rank(operand_shape)); } @@ -1217,24 +1239,26 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (!ShapeUtil::ElementIsFloating(operand_shape)) { return InvalidArgument( "The operand to batch-norm-training must have a floating point " - "element type, but the shape is %s", + "element type, but the shape is %s.", 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 " - "and the shape of operand is %s", + "and the shape of operand is %s.", PrimitiveType_Name(offset_shape.element_type()).c_str(), 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 " - "and the shape of operand is %s", + "and the shape of operand is %s.", PrimitiveType_Name(scale_shape.element_type()).c_str(), PrimitiveType_Name(operand_shape.element_type()).c_str()); } @@ -1247,7 +1271,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of offset factor should be the same as feature count," "but the size of offset factor is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(offset_shape, 0), feature_count); } @@ -1255,7 +1279,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of scale factor should be the same as feature count," "but the size of scale factor is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(scale_shape, 0), feature_count); } @@ -1290,21 +1314,21 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Expected feature_index of batch-norm-inference to be " "smaller than the rank of operand_shape; " - "got feature_index %lld, and rank %lld", + "got feature_index %lld, and rank %lld.", feature_index, ShapeUtil::Rank(operand_shape)); } if (feature_index < 0) { return InvalidArgument( "Expected feature_index of batch-norm-inference to " - "be a non-negative number, got %lld", + "be a non-negative number, got %lld.", feature_index); } if (ShapeUtil::Rank(operand_shape) < 1) { return InvalidArgument( "Expected the rank of operand to " - "batch-norm-inference to be at least 1; got %lld", + "batch-norm-inference to be at least 1; got %lld.", ShapeUtil::Rank(operand_shape)); } @@ -1325,46 +1349,50 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (!ShapeUtil::ElementIsFloating(operand_shape)) { return InvalidArgument( "The operand to batch-norm-inference must have a floating point " - "element type, but the shape is %s", + "element type, but the shape is %s.", 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, " "but the shape of offset factor is %s " - "and the shape of operand is %s", + "and the shape of operand is %s.", PrimitiveType_Name(offset_shape.element_type()).c_str(), 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, " "but the shape of scale factor is %s " - "and the shape of operand is %s", + "and the shape of operand is %s.", PrimitiveType_Name(scale_shape.element_type()).c_str(), 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, " "but the shape of mean is %s " - "and the shape of operand is %s", + "and the shape of operand is %s.", PrimitiveType_Name(mean_shape.element_type()).c_str(), 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, " "but the shape of variance is %s " - "and the shape of operand is %s", + "and the shape of operand is %s.", PrimitiveType_Name(mean_shape.element_type()).c_str(), PrimitiveType_Name(variance_shape.element_type()).c_str()); } @@ -1377,7 +1405,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of offset factor should be the same as feature count," "but the size of offset factor is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(offset_shape, 0), feature_count); } @@ -1385,7 +1413,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of scale factor should be the same as feature count," "but the size of scale factor is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(scale_shape, 0), feature_count); } @@ -1393,7 +1421,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of mean should be the same as feature count," "but the size of mean is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(mean_shape, 0), feature_count); } @@ -1401,7 +1429,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of variance should be the same as feature count," "but the size of variance is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(variance_shape, 0), feature_count); } @@ -1434,7 +1462,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Expected feature_index of batch-norm-grad to be " "smaller than the rank of operand_shape; " - "got feature_index %lld, and rank %lld", + "got feature_index %lld, and rank %lld.", feature_index, ShapeUtil::Rank(operand_shape)); } @@ -1442,7 +1470,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Expected operand_shape of batch-norm-grad to have the same rank as" " output_grad_shape; got rank(oprand_shape) %lld, and" - " rank(output_grad_shape) %lld", + " rank(output_grad_shape) %lld.", ShapeUtil::Rank(operand_shape), ShapeUtil::Rank(output_grad_shape)); } @@ -1470,49 +1498,53 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (!ShapeUtil::ElementIsFloating(operand_shape)) { return InvalidArgument( "The operand to batch-norm-grad must have a floating point " - "element type, but the shape is %s", + "element type, but the shape is %s.", PrimitiveType_Name(operand_shape.element_type()).c_str()); } if (!ShapeUtil::ElementIsFloating(output_grad_shape)) { return InvalidArgument( "The output_grad to batch-norm-grad must have a floating point " - "element type, but the shape is %s", + "element type, but the shape is %s.", 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 " - "and the element type of operand is %s", + "and the element type of operand is %s.", PrimitiveType_Name(output_grad_shape.element_type()).c_str(), 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 " - "and the element type of operand is %s", + "and the element type of operand is %s.", PrimitiveType_Name(scale_shape.element_type()).c_str(), 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 " - "and the element type of operand is %s", + "and the element type of operand is %s.", PrimitiveType_Name(mean_shape.element_type()).c_str(), 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 " - "and the element type of operand is %s", + "and the element type of operand is %s.", PrimitiveType_Name(mean_shape.element_type()).c_str(), PrimitiveType_Name(operand_shape.element_type()).c_str()); } @@ -1526,7 +1558,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of mean should be the same as feature count," "but the size of offset factor is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(mean_shape, 0), feature_count); } @@ -1534,7 +1566,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of scale factor should be the same as feature count," "but the size of scale factor is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(scale_shape, 0), feature_count); } @@ -1542,7 +1574,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The size of variance should be the same as feature count," "but the size of variance is %lld " - "and the feature count is %lld", + "and the feature count is %lld.", ShapeUtil::GetDimension(var_shape, 0), feature_count); } @@ -1553,7 +1585,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "The bounds of operand shape should be the same as output_grad's," "but the bound of operand_shape at dimension %lld is %lld " - "and the bound of output_grad_shape is %lld", + "and the bound of output_grad_shape is %lld.", i, ShapeUtil::GetDimension(operand_shape, i), ShapeUtil::GetDimension(output_grad_shape, i)); } @@ -1569,9 +1601,9 @@ 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", + "Convolution with different element types: %s and %s.", ShapeUtil::HumanString(lhs).c_str(), ShapeUtil::HumanString(rhs).c_str()); } @@ -1587,21 +1619,19 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (window.dimensions_size() != num_spatial_dims) { return InvalidArgument( "Window must have same number of dimensions as dimension numbers.\n" - "Window: %s\nDimension numbers: %s", + "Window: %s\nDimension numbers: %s.", window.DebugString().c_str(), dnums.DebugString().c_str()); } const int num_dims = num_spatial_dims + 2; if (ShapeUtil::Rank(lhs) != num_dims) { return InvalidArgument( - "The LHS argument to a convolution should have rank %d.\n" - "lhs: %s", + "The LHS argument to a convolution should have rank %d; lhs: %s.", num_dims, ShapeUtil::HumanString(lhs).c_str()); } if (ShapeUtil::Rank(rhs) != num_dims) { return InvalidArgument( - "The RHS argument to a convolution should have rank %d.\n" - "lhs: %s", + "The RHS argument to a convolution should have rank %d; lhs: %s.", num_dims, ShapeUtil::HumanString(lhs).c_str()); } TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(lhs)); @@ -1638,26 +1668,26 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( !std::all_of(window_dnums.begin(), window_dnums.end(), in_range) || !std::all_of(output_dnums.begin(), output_dnums.end(), in_range)) { return InvalidArgument( - "A dimension number is out of range in convolution: %s", + "A dimension number is out of range in convolution: %s.", dnums.DebugString().c_str()); } if (input_dnums != expected_dnums) { return InvalidArgument( "Input dimensions of convolution must contain each dimension exactly " - "once: %s", + "once: %s.", dnums.DebugString().c_str()); } if (window_dnums != expected_dnums) { return InvalidArgument( "Window dimensions of convolution must contain each dimension exactly " - "once: %s", + "once: %s.", dnums.DebugString().c_str()); } if (output_dnums != expected_dnums) { return InvalidArgument( "Output dimensions of convolution must contain each dimension exactly " - "once: %s", + "once: %s.", dnums.DebugString().c_str()); } @@ -1681,7 +1711,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Expected LHS feature dimension (value %lld) to match RHS " "input feature dimension (value %lld); got (%s, %s)\n" - "Dimension numbers: {%s}", + "Dimension numbers: {%s}.", input_features, kernel_input_features, ShapeUtil::HumanString(lhs).c_str(), ShapeUtil::HumanString(rhs).c_str(), dnums.DebugString().c_str()); @@ -1695,7 +1725,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( "Window dimensions do not match RHS shape:\n\t" "RHS shape: %s\n\t" "Window: {%s}\n\t" - "Dimension numbers: {%s}", + "Dimension numbers: {%s}.", ShapeUtil::HumanString(rhs).c_str(), window.ShortDebugString().c_str(), dnums.ShortDebugString().c_str()); } @@ -1714,8 +1744,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( dimensions[dnums.output_spatial_dimensions(i)] = window_output_shape.dimensions(i); } - - return ShapeUtil::MakeShape(lhs.element_type(), dimensions); + return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), + dimensions); } /* static */ StatusOr ShapeInference::InferFftShape( @@ -1723,8 +1753,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( 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); + return InvalidArgument("FFT only supports ranks 1-3; got %lld.", fft_rank); } #define RET_CHECK_RANK(x) \ if (x.dimensions_size() < fft_rank) { \ @@ -1737,7 +1766,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( case FFT: case IFFT: if (in.element_type() != C64) { - return InvalidArgument("%s requires C64 input type, found %s", + return InvalidArgument("%s requires C64 input type, found %s.", FftType_Name(fft_type).c_str(), PrimitiveType_Name(in.element_type()).c_str()); } @@ -1745,7 +1774,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return in; case RFFT: { if (in.element_type() != F32) { - return InvalidArgument("RFFT requires F32 input type, found %s", + return InvalidArgument("RFFT requires F32 input type, found %s.", PrimitiveType_Name(in.element_type()).c_str()); } RET_CHECK_RANK(in); @@ -1754,7 +1783,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( fft_length[i]) { return InvalidArgument( "RFFT requires innermost dimensions match fft_length but " - "dimension %lld is %lld and should be %lld", + "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]); @@ -1767,7 +1796,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } case IRFFT: { if (in.element_type() != C64) { - return InvalidArgument("IRFFT requires C64 input type, found %s", + return InvalidArgument("IRFFT requires C64 input type, found %s.", PrimitiveType_Name(in.element_type()).c_str()); } RET_CHECK_RANK(in); @@ -1777,7 +1806,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( fft_length[i]) { return InvalidArgument( "IRFFT requires all but one innermost dimensions match " - "fft_length, but dimension %lld is %lld and should be %lld", + "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]); @@ -1787,7 +1816,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( 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", + "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); } @@ -1825,8 +1854,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( for (int64 dimension : dimensions_to_reduce) { if (dimension >= ShapeUtil::Rank(arg) || dimension < 0) { return InvalidArgument( - "attempting to reduce out-of-bounds dimension %lld in shape %s", - dimension, ShapeUtil::HumanString(arg).c_str()); + "Reducing out-of-bounds dimension %lld in shape %s.", dimension, + ShapeUtil::HumanString(arg).c_str()); } } TF_RETURN_IF_ERROR( @@ -1866,30 +1895,30 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // Check if the select function has a proper shape of (T,T) -> PRED. if (select_shape.parameters_size() != 2) { return InvalidArgument( - "select function must take 2 parameters, but " + "Select function must take 2 parameters, but " "takes %d parameter(s).", select_shape.parameters_size()); } const Shape& select_result_shape = select_shape.result(); if (!ShapeUtil::Compatible(select_result_shape, ShapeUtil::MakeShape(PRED, {}))) { - return Unimplemented("select function must have rank-0 PRED result."); + return InvalidArgument("Select function must have rank-0 PRED result."); } 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", + "Select function's first parameter shape currently must " + "match the operand element shape, but 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", + "Select function's second parameter shape currently must " + "match the operand element shape, but got %s vs %s.", ShapeUtil::HumanString(select_shape.parameters(1)).c_str(), ShapeUtil::HumanString(operand_element_shape).c_str()); } @@ -1903,10 +1932,11 @@ 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)", + "Source shape does not match the shape of window-reduced operand: " + "source(%s), window-reduced operand(%s).", ShapeUtil::HumanString(source_shape).c_str(), ShapeUtil::HumanString(window_result_shape).c_str()); } @@ -1920,7 +1950,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( auto error = [&](const string& message) { return InvalidArgument( "%s in slice operation; argument shape: %s; starts: {%s}; limits: " - "{%s}; strides: {%s}", + "{%s}; strides: {%s}.", message.c_str(), ShapeUtil::HumanString(arg).c_str(), Join(starts, ",").c_str(), Join(limits, ",").c_str(), Join(strides, ",").c_str()); @@ -1943,7 +1973,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (starts.size() != ShapeUtil::Rank(arg)) { return InvalidArgument( - "slice index count does not match argument rank: %zu vs %lld", + "Slice index count does not match argument rank: %zu vs %lld.", starts.size(), ShapeUtil::Rank(arg)); } @@ -1953,7 +1983,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( int64 limit_index = limits[dimension]; int64 stride = strides[dimension]; if (start_index < 0) { - return InvalidArgument("negative start index to slice: %lld", + return InvalidArgument("Negative start index to slice: %lld.", start_index); } if (limit_index > arg.dimensions(dimension)) { @@ -1973,7 +2003,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( limit_index, start_index)); } if (stride <= 0) { - return InvalidArgument("stride (%lld) must be positive", stride); + return InvalidArgument("Stride (%lld) must be positive.", stride); } sizes.push_back((limit_index - start_index + stride - 1) / stride); } @@ -1997,20 +2027,20 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (ShapeUtil::Rank(start_indices_shape) != 1) { return InvalidArgument( - "dynamic slice start indices of rank %lld must be rank1.", + "Dynamic slice start indices of rank %lld must be rank1.", ShapeUtil::Rank(start_indices_shape)); } if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { return InvalidArgument( - "dynamic slice start indices must be of integral type."); + "Dynamic slice start indices must be of integral type."); } const int64 start_num_dims = start_indices_shape.dimensions(0); if (ShapeUtil::Rank(operand_shape) != start_num_dims) { return InvalidArgument( - "dynamic slice start number of dimensions %lld (%s) must match rank " - "%lld of slice input (%s)", + "Dynamic slice start number of dimensions %lld (%s) must match rank " + "%lld of slice input (%s).", start_num_dims, ShapeUtil::HumanString(start_indices_shape).c_str(), ShapeUtil::Rank(operand_shape), ShapeUtil::HumanString(operand_shape).c_str()); @@ -2018,7 +2048,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (slice_sizes.size() != ShapeUtil::Rank(operand_shape)) { return InvalidArgument( - "dynamic slice index count does not match argument rank: %zu vs %lld", + "Dynamic slice index count does not match argument rank: %zu vs %lld.", slice_sizes.size(), ShapeUtil::Rank(operand_shape)); } @@ -2026,12 +2056,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const int64 input_dim_size = operand_shape.dimensions(dim); const int64 slice_dim_size = slice_sizes[dim]; if (slice_dim_size < 0) { - return InvalidArgument("negative size index to dynamic slice: %lld", + return InvalidArgument("Negative size index to dynamic slice: %lld.", slice_dim_size); } if (slice_dim_size > input_dim_size) { return InvalidArgument( - "slice dim size %lld greater than dynamic slice dimension: %lld", + "Slice dim size %lld greater than dynamic slice dimension: %lld.", slice_dim_size, input_dim_size); } VLOG(2) << tensorflow::strings::Printf("slice_sizes[%lld] = %lld", dim, @@ -2060,20 +2090,20 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (ShapeUtil::Rank(start_indices_shape) != 1) { return InvalidArgument( - "dynamic update slice start indices of rank %lld must be rank1.", + "Dynamic update slice start indices of rank %lld must be rank1.", ShapeUtil::Rank(start_indices_shape)); } if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { return InvalidArgument( - "dynamic update slice start indices must be of integral type."); + "Dynamic update slice start indices must be of integral type."); } const int64 start_num_dims = start_indices_shape.dimensions(0); if (ShapeUtil::Rank(operand_shape) != start_num_dims) { return InvalidArgument( - "dynamic slice start number of dimensions %lld (%s) must match rank " - "%lld of slice input (%s)", + "Dynamic update slice start number of dimensions %lld (%s) must match " + "rank %lld of slice input (%s).", start_num_dims, ShapeUtil::HumanString(start_indices_shape).c_str(), ShapeUtil::Rank(operand_shape), ShapeUtil::HumanString(operand_shape).c_str()); @@ -2081,15 +2111,16 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (ShapeUtil::Rank(update_shape) != ShapeUtil::Rank(operand_shape)) { return InvalidArgument( - "dynamic update slice update rank does not match argument rank: " - "%lld vs %lld", + "Dynamic update slice update rank does not match argument rank: " + "%lld vs %lld.", 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", + "Dynamic update slice update element type does not match argument. " + "operand.element_type: %s vs update.element_type: %s.", PrimitiveType_Name(operand_shape.element_type()).c_str(), PrimitiveType_Name(update_shape.element_type()).c_str()); } @@ -2099,12 +2130,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const int64 update_dim_size = update_shape.dimensions(dim); if (update_dim_size < 0) { return InvalidArgument( - "size index %lld to dynamic update slice must be >= 0", + "Size index %lld to dynamic update slice must be >= 0.", update_dim_size); } if (update_dim_size > input_dim_size) { return InvalidArgument( - "update dim size %lld greater than dynamic slice dimension: %lld", + "Update dim size %lld greater than dynamic slice dimension: %lld.", update_dim_size, input_dim_size); } VLOG(2) << tensorflow::strings::Printf("update_sizes[%lld] = %lld", dim, @@ -2124,7 +2155,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( for (int64 dimension : dimensions) { if (dimension >= ShapeUtil::Rank(operand_shape) || dimension < 0) { return InvalidArgument( - "one of the reverse dimensions (%lld) is out-of-bounds in shape %s", + "One of the reverse dimensions (%lld) is out-of-bounds in shape %s.", dimension, ShapeUtil::HumanString(operand_shape).c_str()); } } @@ -2135,14 +2166,14 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const Shape& arg, int64 index) { if (!ShapeUtil::IsTuple(arg)) { return InvalidArgument( - "cannot infer shape: attempting to index into non-tuple: %s", + "Cannot infer shape: attempting to index into non-tuple: %s.", ShapeUtil::HumanString(arg).c_str()); } if (index >= arg.tuple_shapes_size()) { return InvalidArgument( - "cannot infer shape: attempt to index out of tuple bounds: %lld " - ">= %d in shape %s", + "Cannot infer shape: attempt to index out of tuple bounds: %lld " + ">= %d in shape %s.", index, arg.tuple_shapes_size(), ShapeUtil::HumanString(arg).c_str()); } @@ -2154,17 +2185,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const Shape& init) { // Check the number of parameters for given computations. if (condition.parameters_size() != 1) { - return InvalidArgument("condition must take 1 arguments; got %d", + return InvalidArgument("Condition must take 1 arguments; got %d.", condition.parameters_size()); } if (body.parameters_size() != 1) { - return InvalidArgument("body must take 1 arguments; got %d", + return InvalidArgument("Body must take 1 arguments; got %d.", body.parameters_size()); } auto shape_string = [&]() { return tensorflow::strings::Printf( - "condition: %s; body: %s; init: %s", + "Condition: %s; body: %s; init: %s.", ShapeUtil::HumanString(condition).c_str(), ShapeUtil::HumanString(body).c_str(), ShapeUtil::HumanString(init).c_str()); @@ -2172,15 +2203,15 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // Check the shapes of computation parameters and return types. if (!ShapeUtil::ShapeIs(condition.result(), PRED, {})) { - return InvalidArgument("condition must return a boolean; got %s", + return InvalidArgument("Condition must return a boolean; got %s.", shape_string().c_str()); } if (!ShapeUtil::Compatible(body.result(), condition.parameters(0)) || !ShapeUtil::Compatible(body.result(), body.parameters(0)) || !ShapeUtil::Compatible(body.result(), init)) { return InvalidArgument( - "the parameter of condition and body, the result of the body, and init " - "must all have the same shape; got %s", + "The parameter of condition and body, the result of the body, and init " + "must all have the same shape; got %s.", shape_string().c_str()); } @@ -2192,7 +2223,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const Shape& false_operand, const ProgramShape& true_computation, const ProgramShape& false_computation) { if (!ShapeUtil::ShapeIs(predicate, PRED, {})) { - return InvalidArgument("predicate must be a boolean; got %s.", + return InvalidArgument("Predicate must be a boolean; got %s.", ShapeUtil::HumanString(predicate).c_str()); } @@ -2275,8 +2306,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (ShapeUtil::ElementsIn(operand) != ShapeUtil::ElementsIn(inferred_shape)) { return InvalidArgument( - "reshape operation has mismatched element counts: from=%lld (%s) " - "to=%lld (%s)", + "Reshape operation has mismatched element counts: from=%lld (%s) " + "to=%lld (%s).", ShapeUtil::ElementsIn(operand), ShapeUtil::HumanString(operand).c_str(), ShapeUtil::ElementsIn(inferred_shape), ShapeUtil::HumanString(inferred_shape).c_str()); @@ -2322,28 +2353,30 @@ 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)) { - return InvalidArgument("clamp op with different operand types: %s, %s, %s", + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(min, operand) || + !ShapeUtil::SameElementTypeIgnoringFpPrecision(max, operand)) { + return InvalidArgument("Clamp 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( - "not yet implemented: %s, %s %s", min.ShortDebugString().c_str(), + "%s, %s %s is not implemented.", min.ShortDebugString().c_str(), max.ShortDebugString().c_str(), operand.ShortDebugString().c_str()); } @@ -2352,26 +2385,36 @@ 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", + "Operands to select must be the same shape; got %s and %s.", ShapeUtil::HumanString(on_true).c_str(), ShapeUtil::HumanString(on_false).c_str()); } if (pred.element_type() != PRED) { return InvalidArgument( - "select's pred operand must have PRED element type; got %s", + "Select's pred operand must have PRED element type; got %s.", ShapeUtil::HumanString(pred).c_str()); } - if (ShapeUtil::SameDimensions(pred, on_true) || ShapeUtil::Rank(pred) == 0) { + if (ShapeUtil::CompatibleIgnoringElementType(pred, on_true) || + ShapeUtil::Rank(pred) == 0) { // 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 " - " different from the other operands: %s", + return InvalidArgument( + "Select operation with non-scalar predicate with dimensionality " + " different from the other operands: %s.", ShapeUtil::HumanString(pred).c_str()); } } @@ -2389,7 +2432,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Call applied function arity must match number of arguments; got: " "arity: %d, arguments: %zu; computation signature: %s; argument " - "shapes: [%s]", + "shapes: [%s].", to_apply.parameters_size(), arg_shapes.size(), computation_signature.c_str(), argument_shapes.c_str()); } @@ -2401,7 +2444,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (!ShapeUtil::Compatible(arg_shape, param_shape)) { return InvalidArgument( "Call parameter must match argument; got parameter %d shape: %s, " - "argument shape: %s", + "argument shape: %s.", i, ShapeUtil::HumanString(param_shape).c_str(), ShapeUtil::HumanString(arg_shape).c_str()); } @@ -2410,4 +2453,209 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return to_apply.result(); } +static Status ValidateGatherDimensionNumbers( + const Shape& input_shape, + tensorflow::gtl::ArraySlice gather_indices_shape, + const GatherDimensionNumbers& dim_numbers) { + if (!c_is_sorted(dim_numbers.output_window_dims())) { + return InvalidArgument( + "Output window dimensions in gather op must be ascending; got: %s.", + Join(dim_numbers.output_window_dims(), ", ").c_str()); + } + + if (c_adjacent_find(dim_numbers.output_window_dims()) != + dim_numbers.output_window_dims().end()) { + return InvalidArgument( + "Output window dimensions in gather op must not repeat; got: %s.", + Join(dim_numbers.output_window_dims(), ", ").c_str()); + } + + const int64 output_window_dim_count = dim_numbers.output_window_dims_size(); + const int64 output_shape_rank = + output_window_dim_count + gather_indices_shape.size() - 1; + + for (int i = 0; i < dim_numbers.output_window_dims_size(); ++i) { + int64 window_index = dim_numbers.output_window_dims(i); + if (window_index < 0 || window_index >= output_shape_rank) { + return InvalidArgument( + "Window index %d in gather op is out of bounds; got %lld, but should " + "have been in [0,%lld).", + i, window_index, output_shape_rank); + } + } + + if (dim_numbers.gather_dims_to_operand_dims_size() != + gather_indices_shape[dim_numbers.index_vector_dim()]) { + return InvalidArgument( + "Gather op has %d elements in gather_dims_to_operand_dims and the " + "bound of dimension index_vector_dim=%lld of gather_indices is " + "%lld. These two numbers must be equal.", + dim_numbers.gather_dims_to_operand_dims_size(), + dim_numbers.index_vector_dim(), + gather_indices_shape[dim_numbers.index_vector_dim()]); + } + + for (int i = 0; i < dim_numbers.gather_dims_to_operand_dims_size(); i++) { + int64 gather_dim_to_input_dim = dim_numbers.gather_dims_to_operand_dims(i); + if (gather_dim_to_input_dim < 0 || + gather_dim_to_input_dim >= input_shape.dimensions_size()) { + return InvalidArgument( + "Invalid gather_dims_to_operand_dims mapping; domain is [0, %d), " + "got: %d->%lld.", + input_shape.dimensions_size(), i, gather_dim_to_input_dim); + } + } + + std::vector sorted_gather_dims_to_operand_dims( + dim_numbers.gather_dims_to_operand_dims().begin(), + dim_numbers.gather_dims_to_operand_dims().end()); + + c_sort(sorted_gather_dims_to_operand_dims); + + if (c_adjacent_find(sorted_gather_dims_to_operand_dims) != + sorted_gather_dims_to_operand_dims.end()) { + return InvalidArgument( + "Repeated dimensions are not allowed in gather_dims_to_operand_dims; " + "got: %s.", + Join(dim_numbers.gather_dims_to_operand_dims(), ", ").c_str()); + } + + for (int64 elided_dim : dim_numbers.elided_window_dims()) { + if (elided_dim < 0 || elided_dim >= input_shape.dimensions_size()) { + return InvalidArgument( + "Invalid elided_window_dims set in gather op; valid range is [0, " + "%d), got: %lld.", + input_shape.dimensions_size(), elided_dim); + } + } + + if (!c_is_sorted(dim_numbers.elided_window_dims())) { + return InvalidArgument( + "elided_window_dims in gather op must be sorted; got: %s", + Join(dim_numbers.elided_window_dims(), ", ").c_str()); + } + + if (c_adjacent_find(dim_numbers.elided_window_dims()) != + dim_numbers.elided_window_dims().end()) { + return InvalidArgument( + "Repeated dimensions not allowed in elided_window_dims in gather op; " + "got: %s.", + Join(dim_numbers.elided_window_dims(), ", ").c_str()); + } + + return Status::OK(); +} + +/*static*/ StatusOr ShapeInference::InferGatherShape( + const Shape& input_shape, const Shape& gather_indices_shape, + const GatherDimensionNumbers& gather_dim_numbers, + tensorflow::gtl::ArraySlice window_bounds) { + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(input_shape, "input tensor operand gather op")); + TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque( + gather_indices_shape, "gather indices operand of gather op")); + + if (!ShapeUtil::ElementIsIntegral(gather_indices_shape)) { + return InvalidArgument( + "Gather indices parameter must be an integral tensor; got %s.", + ShapeUtil::HumanString(gather_indices_shape).c_str()); + } + + // We implicitly reshape gather indices of shape P[A,B,C] to P[A,B,C,1] if + // index_vector_dim is rank(P). The bounds of this expanded shape is + // stored in expanded_gather_indices_shape. + + if (gather_indices_shape.dimensions_size() < + gather_dim_numbers.index_vector_dim() || + gather_dim_numbers.index_vector_dim() < 0) { + return InvalidArgument( + "Gather index leaf dimension must be within [0, rank(gather_indices) + " + "1). rank(gather_indices) is %d and gather index leaf dimension is " + "%lld.", + gather_indices_shape.dimensions_size(), + gather_dim_numbers.index_vector_dim()); + } + + std::vector expanded_gather_indices_shape; + expanded_gather_indices_shape.reserve(gather_indices_shape.dimensions_size()); + c_copy(gather_indices_shape.dimensions(), + std::back_inserter(expanded_gather_indices_shape)); + if (expanded_gather_indices_shape.size() == + gather_dim_numbers.index_vector_dim()) { + expanded_gather_indices_shape.push_back(1); + } + + TF_RETURN_IF_ERROR(ValidateGatherDimensionNumbers( + input_shape, expanded_gather_indices_shape, gather_dim_numbers)); + + if (window_bounds.size() != input_shape.dimensions_size()) { + return InvalidArgument( + "Gather op must have one window bound for every input dimension; got: " + "len(window_bounds)=%lu, input_shape.rank=%d.", + window_bounds.size(), input_shape.dimensions_size()); + } + + if (window_bounds.size() != + gather_dim_numbers.output_window_dims_size() + + gather_dim_numbers.elided_window_dims_size()) { + return InvalidArgument( + "All components of the window index in a gather op must either be a " + "output window index or explicitly elided; got len(window_bounds)=%lu, " + "output_window_bounds=%s, elided_window_bounds=%s.", + window_bounds.size(), + Join(gather_dim_numbers.output_window_dims(), ",").c_str(), + Join(gather_dim_numbers.elided_window_dims(), ",").c_str()); + } + + for (int i = 0; i < window_bounds.size(); i++) { + int64 window_bound = window_bounds[i]; + int64 corresponding_input_bound = input_shape.dimensions(i); + if (window_bound < 0 || window_bound > corresponding_input_bound) { + return InvalidArgument( + "Window bound at index %d in gather op is out of range, must be " + "within " + "[0, %lld), got %lld.", + i, corresponding_input_bound + 1, window_bound); + } + } + + for (int i = 0; i < gather_dim_numbers.elided_window_dims_size(); i++) { + if (window_bounds[gather_dim_numbers.elided_window_dims(i)] != 1) { + return InvalidArgument( + "Gather op can only elide window indices with bound 1, but bound is " + "%lld for index %lld at position %d.", + window_bounds[gather_dim_numbers.elided_window_dims(i)], + gather_dim_numbers.elided_window_dims(i), i); + } + } + + int64 result_rank = gather_dim_numbers.output_window_dims_size() + + (expanded_gather_indices_shape.size() - 1); + int64 window_dims_seen = 0; + int64 gather_dims_seen = 0; + std::vector output_dim_bounds; + output_dim_bounds.reserve(result_rank); + for (int64 i = 0; i < result_rank; i++) { + int64 current_bound; + bool is_window_index = + c_binary_search(gather_dim_numbers.output_window_dims(), i); + if (is_window_index) { + while (c_binary_search(gather_dim_numbers.elided_window_dims(), + window_dims_seen)) { + window_dims_seen++; + } + current_bound = window_bounds[window_dims_seen++]; + } else { + if (gather_dims_seen == gather_dim_numbers.index_vector_dim()) { + gather_dims_seen++; + } + current_bound = expanded_gather_indices_shape[gather_dims_seen++]; + } + + output_dim_bounds.push_back(current_bound); + } + + return ShapeUtil::MakeShape(input_shape.element_type(), output_dim_bounds); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h index b39151ebbc19f5d0b702a80da5069f58c8dfb07d..085fdac60c6de161c457dff672175e82f4f4da51 100644 --- a/tensorflow/compiler/xla/service/shape_inference.h +++ b/tensorflow/compiler/xla/service/shape_inference.h @@ -37,6 +37,11 @@ namespace xla { // the expected result type for computations that are built up via the API -- // the shape that results from an operation is inferred. Some methods have // overloads for inferring shape at the HLO level. +// +// TODO(b/73352135): Shape inference does not issue very good error messages, in +// part because HloInstruction::ToString() is not available since shape +// inference runs before the HloInstruction object is created. We need a +// solution for this. class ShapeInference { public: // Infers the shape produced by applying the given unary operation to the @@ -51,6 +56,9 @@ class ShapeInference { static StatusOr InferBinaryOpShape( BinaryOperation operation, const Shape& lhs, const Shape& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions); + static StatusOr InferBinaryOpShape( + HloOpcode opcode, const Shape& lhs, const Shape& rhs, + tensorflow::gtl::ArraySlice broadcast_dimensions); static StatusOr InferBinaryOpShape(HloOpcode opcode, const HloInstruction* lhs, const HloInstruction* rhs); @@ -248,6 +256,14 @@ class ShapeInference { const Shape& lhs, const Shape& rhs, const DotDimensionNumbers& dimension_numbers); + // Helper that infers the shape of the tensor produced by a gather operation + // with the given input shape, gather indices shape and gather dimension + // numbers. + static StatusOr InferGatherShape( + const Shape& input_shape, const Shape& gather_indices_shape, + const GatherDimensionNumbers& gather_dim_numbers, + tensorflow::gtl::ArraySlice window_bounds); + private: // Helper that infers the shape produced by performing an element-wise binary // operation with the given LHS and RHS shapes. diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc index 026c021165785bd3945d6a846dae446ad45da9b7..0e61994a786b53a295ef9c9c2287b28fbf754d9b 100644 --- a/tensorflow/compiler/xla/service/shape_inference_test.cc +++ b/tensorflow/compiler/xla/service/shape_inference_test.cc @@ -18,15 +18,16 @@ limitations under the License. #include #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" - #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #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 { +using ::tensorflow::gtl::ArraySlice; using ::testing::ContainsRegex; using ::testing::HasSubstr; @@ -134,7 +135,7 @@ TEST_F(ShapeInferenceTest, SelectBadShapes) { TernaryOperation::TRIOP_SELECT, pred_, matrix_64_48_, matrix_32_64_); ASSERT_FALSE(inferred_status_error1.ok()); ASSERT_THAT(inferred_status_error1.status().error_message(), - HasSubstr("operands to select must be the same shape")); + HasSubstr("Operands to select must be the same shape")); auto inferred_status_error2 = ShapeInference::InferTernaryOpShape( TernaryOperation::TRIOP_SELECT, s32_, matrix_64_48_, matrix_64_48_); @@ -339,7 +340,7 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSourceShape) { init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); ASSERT_THAT(inferred_status_fail.status().error_message(), - HasSubstr("source shape does not match")); + HasSubstr("Source shape does not match")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape1) { @@ -350,7 +351,7 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape1) { init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); ASSERT_THAT(inferred_status_fail.status().error_message(), - HasSubstr("select function must take 2 parameters")); + HasSubstr("Select function must take 2 parameters")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape2) { @@ -361,7 +362,7 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape2) { init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); ASSERT_THAT(inferred_status_fail.status().error_message(), - HasSubstr("select function must have rank-0 PRED")); + HasSubstr("Select function must have rank-0 PRED")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape3) { @@ -372,7 +373,7 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape3) { init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); ASSERT_THAT(inferred_status_fail.status().error_message(), - HasSubstr("select function's first parameter")); + HasSubstr("Select function's first parameter")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape4) { @@ -383,7 +384,7 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape4) { init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); ASSERT_THAT(inferred_status_fail.status().error_message(), - HasSubstr("select function's second parameter")); + HasSubstr("Select function's second parameter")); } TEST_F(ShapeInferenceTest, Convolve) { @@ -905,7 +906,7 @@ TEST_F(ShapeInferenceTest, ScalarDotVector) { ShapeInference::InferDotOpShape(f32_, vector_32_, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("dot only supports rank")); + HasSubstr("Dot only supports rank")); } // 3D 2D: error @@ -917,7 +918,7 @@ TEST_F(ShapeInferenceTest, DotWithRankHigherThanTwo) { ShapeUtil::MakeShape(F32, {32, 32, 32}), matrix_32_64_, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("batch and contracting dimension number mismatch")); + HasSubstr("Batch and contracting dimension number mismatch")); } // vector vector -> scalar @@ -1023,7 +1024,7 @@ TEST_F(ShapeInferenceTest, DotWithTwoContractingDimsFails) { ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("must specify one contracting dimension for both " + HasSubstr("Must specify one contracting dimension for both " "lhs and rhs")); } @@ -1043,7 +1044,7 @@ TEST_F(ShapeInferenceTest, DotWithMisatchedBatchDimSizesFails) { ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("batch dimension numbers and sizes must match")); + HasSubstr("Batch dimension numbers and sizes must match")); } // BatchMatMul with different batch dimension numbers fails. @@ -1062,7 +1063,7 @@ TEST_F(ShapeInferenceTest, DotWithMisatchedBatchDimNumbersFails) { ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("batch dimension numbers must precede non-batch")); + HasSubstr("Batch dimension numbers must precede non-batch")); } // BatchMatMul with out-of-range dimension numbers fails. @@ -1165,42 +1166,42 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastBadDimension) { BinaryOperation::BINOP_ADD, tensor, vec8, {}); ASSERT_FALSE(inferred_status_error1.ok()); ASSERT_THAT(inferred_status_error1.status().error_message(), - HasSubstr("automatic")); + HasSubstr("Automatic")); // broadcast_dimension out of bounds for tensor's rank auto inferred_status_error2 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, vec8, {3}); ASSERT_FALSE(inferred_status_error2.ok()); ASSERT_THAT(inferred_status_error2.status().error_message(), - ContainsRegex("broadcast dimension number .* too large")); + ContainsRegex("Broadcast dimension number .* too large")); // broadcast_dimension doesn't match corresponding dimension auto inferred_status_error3 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, vec8, {0}); ASSERT_FALSE(inferred_status_error3.ok()); ASSERT_THAT(inferred_status_error3.status().error_message(), - HasSubstr("broadcast dimension 0 mismatch")); + HasSubstr("Broadcast dimension 0 mismatch")); // broadcast_dimensions list too long auto inferred_status_error4 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, matrix8_4, {0, 1, 2}); ASSERT_FALSE(inferred_status_error4.ok()); ASSERT_THAT(inferred_status_error4.status().error_message(), - HasSubstr("size of broadcast_dimensions has to match")); + HasSubstr("broadcast_dimensions has to match")); // there's a dimension above the rank of the tensor auto inferred_status_error5 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, matrix8_4, {3, 0}); ASSERT_FALSE(inferred_status_error5.ok()); ASSERT_THAT(inferred_status_error5.status().error_message(), - ContainsRegex("broadcast dimension number .* too large")); + ContainsRegex("dimension number .* too large")); // broadcasting dimensions don't match in this order auto inferred_status_error6 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, matrix8_4, {2, 1}); ASSERT_FALSE(inferred_status_error6.ok()); ASSERT_THAT(inferred_status_error6.status().error_message(), - HasSubstr("broadcast dimension 0 mismatch")); + HasSubstr("dimension 0 mismatch")); // The following two tests make sure that broadcasting dimensions are listed // in a proper (strictly increasing) order, even if the lower-rank array @@ -1209,13 +1210,13 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastBadDimension) { BinaryOperation::BINOP_ADD, tensor8_8_8, matrix8_8, {0, 0}); ASSERT_FALSE(inferred_status_error7.ok()); ASSERT_THAT(inferred_status_error7.status().error_message(), - HasSubstr("broadcast dimensions order is wrong")); + HasSubstr("dimensions order is wrong")); auto inferred_status_error8 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor8_8_8, matrix8_8, {1, 0}); ASSERT_FALSE(inferred_status_error8.ok()); ASSERT_THAT(inferred_status_error8.status().error_message(), - HasSubstr("broadcast dimensions order is wrong")); + HasSubstr("dimensions order is wrong")); } // Tests for the while instruction with proper shapes. @@ -1241,7 +1242,7 @@ TEST_F(ShapeInferenceTest, WhileWithBadShapes) { ShapeInference::InferWhileShape(bad_shape_1, body, result_shape); ASSERT_FALSE(inferred_status_error1.ok()); ASSERT_THAT(inferred_status_error1.status().error_message(), - HasSubstr("condition must take 1 arguments")); + HasSubstr("Condition must take 1 arguments")); auto bad_shape_2 = ShapeUtil::MakeProgramShape({s32_, result_shape}, result_shape); @@ -1249,14 +1250,14 @@ TEST_F(ShapeInferenceTest, WhileWithBadShapes) { ShapeInference::InferWhileShape(cond, bad_shape_2, result_shape); ASSERT_FALSE(inferred_status_error2.ok()); ASSERT_THAT(inferred_status_error2.status().error_message(), - HasSubstr("body must take 1 arguments")); + HasSubstr("Body must take 1 arguments")); auto bad_shape_3 = ShapeUtil::MakeProgramShape({result_shape}, s32_); auto inferred_status_error3 = ShapeInference::InferWhileShape(bad_shape_3, body, result_shape); ASSERT_FALSE(inferred_status_error3.ok()); ASSERT_THAT(inferred_status_error3.status().error_message(), - HasSubstr("condition must return a boolean")); + HasSubstr("Condition must return a boolean")); auto bad_shape_4 = ShapeUtil::MakeProgramShape({result_shape}, vector_32_); auto inferred_status_error4 = @@ -1300,13 +1301,13 @@ TEST_F(ShapeInferenceTest, ConcatenateWithBadShapes) { ShapeInference::InferConcatOpShape({&vector_32_}, /*dimension=*/-1); ASSERT_FALSE(inferred_status_error2.ok()); ASSERT_THAT(inferred_status_error2.status().error_message(), - HasSubstr("dimension to concatenate along out of bounds: -1")); + HasSubstr("dimension out of bounds: -1")); auto inferred_status_error3 = ShapeInference::InferConcatOpShape({&vector_32_}, /*dimension=*/1); ASSERT_FALSE(inferred_status_error3.ok()); ASSERT_THAT(inferred_status_error3.status().error_message(), - HasSubstr("dimension to concatenate along out of bounds: 1")); + HasSubstr("dimension out of bounds: 1")); Shape tuple = ShapeUtil::MakeTupleShape({vector_32_}); auto inferred_status_error4 = ShapeInference::InferConcatOpShape( @@ -1314,21 +1315,20 @@ TEST_F(ShapeInferenceTest, ConcatenateWithBadShapes) { ASSERT_FALSE(inferred_status_error4.ok()); ASSERT_THAT( inferred_status_error4.status().error_message(), - HasSubstr("Expected non-tuple argument for operand of concatenation.")); + HasSubstr("Expected non-tuple argument for operand of concatenation")); const Shape vector_s32 = ShapeUtil::MakeShape(S32, {32}); auto inferred_status_error5 = ShapeInference::InferConcatOpShape( {&vector_32_, &vector_s32}, /*dimension=*/0); ASSERT_FALSE(inferred_status_error5.ok()); - ASSERT_THAT( - inferred_status_error5.status().error_message(), - HasSubstr("cannot concatenate arrays with different element types")); + ASSERT_THAT(inferred_status_error5.status().error_message(), + HasSubstr("concatenate arrays with different element types")); auto inferred_status_error6 = ShapeInference::InferConcatOpShape( {&matrix_32_48_, &matrix_32_64_}, /*dimension=*/0); ASSERT_FALSE(inferred_status_error6.ok()); ASSERT_THAT(inferred_status_error6.status().error_message(), - HasSubstr("cannot concatenate arrays that differ in " + HasSubstr("concatenate arrays that differ in " "dimensions other than the one being " "concatenated")); } @@ -1466,7 +1466,7 @@ TEST_F(ShapeInferenceTest, Conditional) { ShapeUtil::MakeProgramShape({vector_64_}, f32_)); EXPECT_FALSE(inferred_status_error0.ok()); EXPECT_THAT(inferred_status_error0.status().error_message(), - HasSubstr("predicate must be a boolean")); + HasSubstr("Predicate must be a boolean")); auto inferred_status_error1 = ShapeInference::InferConditionalShape( pred_, ShapeUtil::MakeTupleShape({f32_, vector_32_}), matrix_32_48_, @@ -1527,5 +1527,458 @@ TEST_F(ShapeInferenceTest, BadSlice) { << statusor.status(); } +class GatherShapeInferenceTest : public ShapeInferenceTest { + protected: + const Shape s64_scalar_ = ShapeUtil::MakeShape(S64, {}); + const Shape s64_vector_5_ = ShapeUtil::MakeShape(S64, {5}); + const Shape s64_vector_32_ = ShapeUtil::MakeShape(S64, {32}); + const Shape s64_4d_tensor_10_9_8_7_1_ = + ShapeUtil::MakeShape(S64, {10, 9, 8, 7, 1}); + const Shape s64_4d_tensor_10_9_8_7_5_ = + ShapeUtil::MakeShape(S64, {10, 9, 8, 7, 5}); + const Shape s64_4d_tensor_5_10_9_7_6_ = + ShapeUtil::MakeShape(S64, {5, 10, 9, 7, 6}); + const Shape s64_4d_tensor_10_9_5_7_6_ = + ShapeUtil::MakeShape(S64, {10, 9, 5, 7, 6}); + const Shape f32_5d_tensor_50_49_48_47_46_ = + ShapeUtil::MakeShape(F32, {50, 49, 48, 47, 46}); + const Shape tuple_shape_ = ShapeUtil::MakeTupleShape( + {s64_4d_tensor_10_9_8_7_1_, s64_4d_tensor_10_9_8_7_1_}); +}; + +TEST_F(GatherShapeInferenceTest, TensorFlowGather) { + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape(matrix_64_48_, s64_vector_32_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{0}, + /*elided_window_dims=*/{1}, + /*gather_dims_to_operand_dims=*/{1}, + /*index_vector_dim=*/1), + /*window_bounds=*/{64, 1})); + EXPECT_TRUE( + ShapeUtil::Equal(gather_shape, ShapeUtil::MakeShape(F32, {64, 32}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, TensorFlowGatherV2) { + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape(matrix_64_48_, s64_vector_32_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{1}, + /*elided_window_dims=*/{0}, + /*gather_dims_to_operand_dims=*/{0}, + /*index_vector_dim=*/1), + /*window_bounds=*/{1, 48})); + EXPECT_TRUE( + ShapeUtil::Equal(gather_shape, ShapeUtil::MakeShape(F32, {32, 48}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, TensorFlowGatherNd) { + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape(matrix_64_48_, s64_4d_tensor_10_9_8_7_1_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4}, + /*elided_window_dims=*/{0}, + /*gather_dims_to_operand_dims=*/{0}, + /*index_vector_dim=*/4), + /*window_bounds=*/{1, 48})); + EXPECT_TRUE(ShapeUtil::Equal(gather_shape, + ShapeUtil::MakeShape(F32, {10, 9, 8, 7, 48}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, TensorFlowBatchDynamicSlice) { + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26})); + EXPECT_TRUE(ShapeUtil::Equal( + gather_shape, + ShapeUtil::MakeShape(F32, {10, 9, 8, 7, 30, 29, 28, 27, 26}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, NonDefaultGatherIndicesLeafDim_A) { + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_5_7_6_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/2), + /*window_bounds=*/{30, 29, 28, 27, 26})); + + EXPECT_TRUE(ShapeUtil::Equal( + gather_shape, + ShapeUtil::MakeShape(F32, {10, 9, 7, 6, 30, 29, 28, 27, 26}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, NonDefaultGatherIndicesLeafDim_B) { + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_5_10_9_7_6_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/0), + /*window_bounds=*/{30, 29, 28, 27, 26})); + + EXPECT_TRUE(ShapeUtil::Equal( + gather_shape, + ShapeUtil::MakeShape(F32, {10, 9, 7, 6, 30, 29, 28, 27, 26}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, NoOutputGatherDims) { + // This is equivalent to a dynamic slice. + TF_ASSERT_OK_AND_ASSIGN( + Shape gather_shape, + ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_vector_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{0, 1, 2, 3, 4}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/0), + /*window_bounds=*/{30, 29, 28, 27, 26})); + + EXPECT_TRUE(ShapeUtil::Equal(gather_shape, + ShapeUtil::MakeShape(F32, {30, 29, 28, 27, 26}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, ScalarGatherIndices) { + // The gather indices "tensor" is a scalar S here that's used to slice out + // [S,0,0,0,0]..[S,30,29,28,27] into a [30,29,28,27] shaped result. + TF_ASSERT_OK_AND_ASSIGN(Shape gather_shape, + ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_scalar_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{0, 1, 2, 3}, + /*elided_window_dims=*/{0}, + /*gather_dims_to_operand_dims=*/{0}, + /*index_vector_dim=*/0), + /*window_bounds=*/{1, 30, 29, 28, 27})); + + EXPECT_TRUE(ShapeUtil::Equal(gather_shape, + ShapeUtil::MakeShape(F32, {30, 29, 28, 27}))) + << ShapeUtil::HumanString(gather_shape); +} + +TEST_F(GatherShapeInferenceTest, TupleShapedTensorInput) { + StatusOr statusor = ShapeInference::InferGatherShape( + tuple_shape_, s64_vector_32_, + HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0}, + /*elided_window_dims=*/{1}, + /*gather_dims_to_operand_dims=*/{1}, + /*index_vector_dim=*/1), + /*window_bounds=*/{64, 1}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Expected non-tuple argument for input")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, TupleShapedGatherIndicesInput) { + StatusOr statusor = ShapeInference::InferGatherShape( + s64_vector_32_, tuple_shape_, + HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0}, + /*elided_window_dims=*/{1}, + /*gather_dims_to_operand_dims=*/{1}, + /*index_vector_dim=*/0), + /*window_bounds=*/{64, 1}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Expected non-tuple argument for gather indices")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, FloatingPointGatherIndicesInput) { + StatusOr statusor = ShapeInference::InferGatherShape( + s64_vector_32_, vector_32_, + HloInstruction::MakeGatherDimNumbers(/*output_window_dims=*/{0}, + /*elided_window_dims=*/{1}, + /*gather_dims_to_operand_dims=*/{1}, + /*index_vector_dim=*/0), + /*window_bounds=*/{64, 1}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Gather indices parameter must be an integral tensor")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_NonAscendingWindowIndices) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 8, 7}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr("Output window dimensions in gather op must be ascending")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_RepeatedWindowIndices) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 7}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr("Output window dimensions in gather op must not repeat")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_WindowIndexOutOfBounds) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 99, 100, 101}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Window index 2 in gather op is out of bounds")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_WindowIndexBarelyOutOfBounds) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 9}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Window index 4 in gather op is out of bounds")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_MismatchingElidedWindowDims) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{4}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr("All components of the window index in a gather op must either " + "be a output window index or explicitly elided")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_OutOfBoundsWindowToInputMapping) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{0, 1, 2, 3, 19}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Invalid elided_window_dims set in gather op; valid " + "range is [0, 5), got: 19")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_RepeatedWindowToInputMapping) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{0, 1, 2, 3, 3}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr( + "Repeated dimensions not allowed in elided_window_dims in gather op")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_MismatchingGatherToInputMapping) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr("Gather op has 4 elements in gather_dims_to_operand_dims and " + "the bound of dimension index_vector_dim=4 of " + "gather_indices is 5. These two numbers must be equal.")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_OutOfBoundsGatherToInputMapping) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 7}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr("Invalid gather_dims_to_operand_dims mapping; domain is " + "[0, 5), got: 4->7")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_RepeatedGatherToInputMapping) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 3}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr( + "Repeated dimensions are not allowed in gather_dims_to_operand_dims")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_NonAscendingElidedWindowDims) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{2, 1}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{1, 1, 28, 27, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("elided_window_dims in gather op must be sorted")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, InvalidGatherDimNumbers_WindowBoundsTooLarge) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7}, + /*elided_window_dims=*/{2}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 1, 300, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Window bound at index 3 in gather op is out of range, " + "must be within [0, 48), got 300")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_MismatchingNumberOfWindowBounds) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 26}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr( + "Gather op must have one window bound for every input dimension")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, + InvalidGatherDimNumbers_WindowBoundsNot1ForElidedDim) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_8_7_5_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7}, + /*elided_window_dims=*/{1}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/4), + /*window_bounds=*/{30, 29, 28, 26, 20}); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Gather op can only elide window indices with bound 1, " + "but bound is 29 for index 1 at position 0")) + << statusor.status(); +} + +TEST_F(GatherShapeInferenceTest, OutOfBoundsGatherIndicesLeafDim) { + StatusOr statusor = ShapeInference::InferGatherShape( + f32_5d_tensor_50_49_48_47_46_, s64_4d_tensor_10_9_5_7_6_, + HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/{4, 5, 6, 7, 8}, + /*elided_window_dims=*/{}, + /*gather_dims_to_operand_dims=*/{0, 1, 2, 3, 4}, + /*index_vector_dim=*/32), + /*window_bounds=*/{30, 29, 28, 27, 26}); + + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("Gather index leaf dimension must be within [0, " + "rank(gather_indices) + 1)")) + << 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/user_computation.cc b/tensorflow/compiler/xla/service/user_computation.cc index 2ea6507900e712200ce43e9b63577a4967381fdf..0dca30a804005c6f536aca5b54af24eb08d4560b 100644 --- a/tensorflow/compiler/xla/service/user_computation.cc +++ b/tensorflow/compiler/xla/service/user_computation.cc @@ -226,7 +226,8 @@ StatusOr UserComputation::AddParameterInstruction( return handle; } -Status UserComputation::AddSendInstruction(const SendRequest& send_request) { +StatusOr UserComputation::AddSendInstruction( + const SendRequest& send_request) { tensorflow::mutex_lock lock(mutex_); // Check if the operand of the instruction is valid. @@ -244,7 +245,7 @@ Status UserComputation::AddSendInstruction(const SendRequest& send_request) { VLOG(1) << "AddSendInstruction (" << GetVersionedHandleInternal() << "), data handle " << handle.handle() << ": " << send_request.ShortDebugString(); - return Status::OK(); + return handle; } StatusOr UserComputation::AddRecvInstruction( @@ -315,6 +316,36 @@ StatusOr UserComputation::AddConstantInstruction( return handle; } +StatusOr UserComputation::AddGatherInstruction( + const GatherRequest& gather_request) { + tensorflow::mutex_lock lock(mutex_); + + TF_ASSIGN_OR_RETURN(const OperationRequest* input_request, + LookUpRequest(gather_request.input())); + TF_ASSIGN_OR_RETURN(const OperationRequest* gather_indices_request, + LookUpRequest(gather_request.gather_indices())); + + TF_ASSIGN_OR_RETURN( + Shape shape, + ShapeInference::InferGatherShape( + input_request->output_shape(), gather_indices_request->output_shape(), + gather_request.dimension_numbers(), + AsInt64Slice(gather_request.window_bounds()))); + + 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_gather_request() = gather_request; + + VLOG(1) << "AddGatherInstruction (" << GetVersionedHandleInternal() + << "), data handle " << handle.handle() << ": " + << gather_request.ShortDebugString(); + return handle; +} + StatusOr UserComputation::AddGetTupleElementInstruction( const GetTupleElementRequest& get_tuple_element_request) { tensorflow::mutex_lock lock(mutex_); @@ -1185,7 +1216,7 @@ StatusOr UserComputation::AddInfeedInstruction( return handle; } -Status UserComputation::AddOutfeedInstruction( +StatusOr UserComputation::AddOutfeedInstruction( const OutfeedRequest& outfeed_request) { tensorflow::mutex_lock lock(mutex_); @@ -1197,8 +1228,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()]; @@ -1209,7 +1238,7 @@ Status UserComputation::AddOutfeedInstruction( VLOG(1) << "AddOutfeedInstruction (" << GetVersionedHandleInternal() << "), data handle " << handle.handle() << ": " << outfeed_request.ShortDebugString(); - return Status::OK(); + return handle; } StatusOr UserComputation::AddCallInstruction( @@ -1278,6 +1307,28 @@ StatusOr UserComputation::AddCustomCallInstruction( return handle; } +StatusOr UserComputation::AddHostComputeInstruction( + const HostComputeRequest& host_compute_request) { + tensorflow::mutex_lock lock(mutex_); + + for (const ComputationDataHandle& handle : host_compute_request.operands()) { + TF_RETURN_IF_ERROR(LookUpRequest(handle).status()); + } + + ComputationDataHandle handle = CreateComputationDataHandle(); + OperationRequest& request = + (*session_computation_.mutable_requests())[handle.handle()]; + *request.mutable_output_handle() = handle; + *request.mutable_output_shape() = host_compute_request.shape(); + *request.mutable_request()->mutable_host_compute_request() = + host_compute_request; + + VLOG(1) << "AddHostComputeInstruction (" << GetVersionedHandleInternal() + << "), data handle " << handle.handle() << ": " + << host_compute_request.ShortDebugString(); + return handle; +} + StatusOr UserComputation::AddDotInstruction( const DotRequest& dot_request) { tensorflow::mutex_lock lock(mutex_); @@ -1715,6 +1766,11 @@ void PureFunctionalVisitor(const SessionComputation& session_computation, break; } + case OpRequest::kHostComputeRequest: { + *is_functional = false; + break; + } + case OpRequest::kCallRequest: { const CallRequest& call_request = request.request().call_request(); for (const ComputationDataHandle& handle : call_request.operands()) { @@ -1993,12 +2049,25 @@ void PureFunctionalVisitor(const SessionComputation& session_computation, break; } + case OpRequest::kGatherRequest: { + PureFunctionalVisitor(session_computation, + request.request().gather_request().input(), + num_parameters, visited, is_functional); + PureFunctionalVisitor(session_computation, + request.request().gather_request().gather_indices(), + num_parameters, visited, is_functional); + break; + } + case OpRequest::OP_NOT_SET: LOG(FATAL) << "OperationRequest doesn't contain a request"; default: LOG(FATAL) << "Unexpected request type: " << request.request().op_case(); } + if (!*is_functional) { + VLOG(1) << "Non-functional: " << request.request().DebugString(); + } visited->insert(handle.handle()); } @@ -2642,6 +2711,15 @@ static void ForEachOperand( break; } + case OpRequest::kHostComputeRequest: { + const HostComputeRequest& hc_request = + request.request().host_compute_request(); + for (const ComputationDataHandle& operand : hc_request.operands()) { + apply(operand); + } + break; + } + case OpRequest::kDotRequest: { const DotRequest& dot_request = request.request().dot_request(); apply(dot_request.rhs()); @@ -2683,6 +2761,13 @@ static void ForEachOperand( break; } + case OpRequest::kGatherRequest: { + const GatherRequest& gather_request = request.request().gather_request(); + apply(gather_request.input()); + apply(gather_request.gather_indices()); + break; + } + case OpRequest::OP_NOT_SET: LOG(FATAL) << "OperationRequest doesn't contain a request"; @@ -3230,20 +3315,23 @@ void ComputationLowerer::Visit( HloInstruction* rhs = lookup_instruction(ternary_op_request.rhs()); HloInstruction* ehs = lookup_instruction(ternary_op_request.ehs()); auto hlo_opcode = TernaryOperationToHloOpcode(ternary_op_request.triop()); - - if (debug_options_.xla_eliminate_hlo_implicit_broadcast()) { - if (!ShapeUtil::SameDimensions(request.output_shape(), lhs->shape())) { + if (debug_options_.xla_eliminate_hlo_implicit_broadcast() && + !ShapeUtil::IsTuple(request.output_shape())) { + if (!ShapeUtil::IsTuple(lhs->shape()) && + !ShapeUtil::SameDimensions(request.output_shape(), lhs->shape())) { // lhs side is being implicitly broadcast. Change to explicit. lhs = ImplicitBroadcastToExplicitBroadcast(lhs, request.output_shape()); } - if (!ShapeUtil::SameDimensions(request.output_shape(), rhs->shape())) { + if (!ShapeUtil::IsTuple(rhs->shape()) && + !ShapeUtil::SameDimensions(request.output_shape(), rhs->shape())) { rhs = ImplicitBroadcastToExplicitBroadcast(rhs, request.output_shape()); } - if (!ShapeUtil::SameDimensions(request.output_shape(), ehs->shape())) { + if (!ShapeUtil::IsTuple(ehs->shape()) && + !ShapeUtil::SameDimensions(request.output_shape(), ehs->shape())) { ehs = ImplicitBroadcastToExplicitBroadcast(ehs, request.output_shape()); } @@ -3298,6 +3386,22 @@ void ComputationLowerer::Visit( break; } + case OpRequest::kHostComputeRequest: { + const HostComputeRequest& host_compute_request = + request.request().host_compute_request(); + std::vector operands; + for (const ComputationDataHandle& operand : + host_compute_request.operands()) { + operands.push_back(lookup_instruction(operand)); + } + auto output_shape = host_compute_request.shape(); + auto channel_name = host_compute_request.channel_name(); + auto cost_estimate_ns = host_compute_request.cost_estimate_ns(); + hlo_instruction = add_instruction(HloInstruction::CreateHostCompute( + output_shape, operands, channel_name, cost_estimate_ns)); + break; + } + case OpRequest::kUnaryOpRequest: { const UnaryOpRequest& unary_op_request = request.request().unary_op_request(); @@ -3400,6 +3504,20 @@ void ComputationLowerer::Visit( break; } + case OpRequest::kGatherRequest: { + const GatherRequest& gather_request = request.request().gather_request(); + HloInstruction* input_operand = + lookup_instruction(gather_request.input()); + HloInstruction* gather_indices_operand = + lookup_instruction(gather_request.gather_indices()); + std::vector window_bounds; + c_copy(gather_request.window_bounds(), std::back_inserter(window_bounds)); + hlo_instruction = add_instruction(HloInstruction::CreateGather( + request.output_shape(), input_operand, gather_indices_operand, + gather_request.dimension_numbers(), window_bounds)); + break; + } + case OpRequest::OP_NOT_SET: LOG(FATAL) << "OperationRequest doesn't contain a request"; diff --git a/tensorflow/compiler/xla/service/user_computation.h b/tensorflow/compiler/xla/service/user_computation.h index 4f92e58877a1d06728fdd250744ca2ce7b57d9ad..5544c868fe905c1ca7e6cab32738440add2e3b4f 100644 --- a/tensorflow/compiler/xla/service/user_computation.h +++ b/tensorflow/compiler/xla/service/user_computation.h @@ -146,7 +146,12 @@ 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 host compute instruction onto this user computation. + StatusOr AddHostComputeInstruction( + const HostComputeRequest& host_compute_request); // Enqueues a call instruction onto this user computation. StatusOr AddCallInstruction( @@ -231,12 +236,17 @@ class UserComputation { const UserComputation& false_computation); // Enqueues a Send instruction onto this user computation. - Status AddSendInstruction(const SendRequest& send_request); + StatusOr AddSendInstruction( + const SendRequest& send_request); // Enqueues a Recv instruction onto this user computation. StatusOr AddRecvInstruction( const RecvRequest& recv_request); + // Enqueues a Gather instruction onto this user computation. + StatusOr AddGatherInstruction( + const GatherRequest& gather_request); + // Returns the user-provided name of this user computation, which is provided // via the XLA computation-building API. const string& name() const { return name_; } 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 index a5f9b01f011ce04f1114c74391a967c62f015221..3ef0cdff6751258e4489ce350deb0931fdf69ef9 100644 --- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc @@ -106,20 +106,12 @@ static bool NotWorthHoistingIndividually(const HloInstruction& instruction) { case HloOpcode::kBitcast: case HloOpcode::kBroadcast: case HloOpcode::kConstant: + case HloOpcode::kReshape: case HloOpcode::kReverse: case HloOpcode::kSlice: + case HloOpcode::kTranspose: 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()); } } diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc index 87a7f86f4ec9844de3e350d7774093dd6248dd83..ec05a74e286c89dd8db5ae07580e461938d7c087 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/while_loop_simplifier.h" #include "tensorflow/compiler/xla/service/call_inliner.h" #include "tensorflow/compiler/xla/service/hlo_evaluator.h" +#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" @@ -212,7 +213,7 @@ static optional GetLoopTripCount(HloInstruction* while_op) { // Now that we know the index of the induction variable, we can we can try to // compute how many times the loop executes. Start by computing the induction // variable's initial value. - HloEvaluator evaluator; + HloEvaluator evaluator(/*max_loop_iterations=*/0); auto* while_init = while_op->mutable_operand(0); auto* indvar_init = while_init->mutable_operand(*indvar_tuple_idx); StatusOr> indvar_init_result = @@ -564,9 +565,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(); @@ -603,6 +606,78 @@ static StatusOr TryRemoveWhileLoop(HloInstruction* while_op) { return false; } +static StatusOr TryPropagateConstant(HloInstruction* while_op) { + auto while_init = while_op->operand(0); + if (while_init->opcode() != HloOpcode::kTuple) { + return false; + } + + auto while_body = while_op->while_body(); + auto while_body_root = while_body->root_instruction(); + if (while_body_root->opcode() != HloOpcode::kTuple) { + return false; + } + + auto while_body_param = while_body->parameter_instruction(0); + const HloInstruction::InstructionVector& root_operands = + while_body_root->operands(); + + // Find the loop invariant tuple elements with scalar constant init value and + // build a map from the tuple element index to the constant value. Limit this + // to scalar constant values because propagating array constants can regress + // performance by forcing us to copy constants. + tensorflow::gtl::FlatMap index_to_constant; + 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_param && + ShapeUtil::IsScalar(instr->shape())) { + auto tuple_element = while_init->operand(i); + if (tuple_element->IsConstant()) { + VLOG(3) << "Found loop invariant tuple element " << i << " " + << tuple_element->ToString(); + index_to_constant[i] = tuple_element; + } + } + } + + if (index_to_constant.empty()) { + return false; + } + + // Replace the use of each constant tuple element in the loop_condition and + // loop_body with the corresponding constant value. + auto propagate_constant = [&](HloComputation* computation) -> StatusOr { + HloInstruction* param = computation->parameter_instruction(0); + bool changed = false; + for (auto instr : param->users()) { + // Since only a while-loop with a tuple result reaches here, we can safely + // assume that `param` is a tuple and the first operand of the + // GetTupleElement instruction is a use of `param`. + if (instr->opcode() == HloOpcode::kGetTupleElement) { + VLOG(3) << "tuple index " << instr->tuple_index() << " " + << instr->ToString(); + auto iter = index_to_constant.find(instr->tuple_index()); + if (iter != index_to_constant.end()) { + const HloInstruction* hlo_constant = (*iter).second; + VLOG(3) << "Replace use of " << instr->ToString() << " with " + << hlo_constant->ToString(); + TF_RETURN_IF_ERROR(instr->ReplaceAllUsesWith( + computation->AddInstruction(hlo_constant->Clone()))); + changed = true; + } + } + } + return changed; + }; + + TF_ASSIGN_OR_RETURN(bool changed_cond, + propagate_constant(while_op->while_condition())); + TF_ASSIGN_OR_RETURN(bool changed_body, propagate_constant(while_body)); + + return changed_cond || changed_body; +} + StatusOr WhileLoopSimplifier::Run(HloModule* module) { XLA_VLOG_LINES(3, "WhileLoopSimplifier::Run(), before:\n" + module->ToString()); @@ -633,7 +708,11 @@ StatusOr WhileLoopSimplifier::Run(HloModule* module) { continue; } - StatusOr result = TryRemoveWhileLoop(while_op); + StatusOr result = TryPropagateConstant(while_op); + TF_RETURN_IF_ERROR(result.status()); + changed |= result.ValueOrDie(); + + result = TryRemoveWhileLoop(while_op); TF_RETURN_IF_ERROR(result.status()); if (result.ValueOrDie()) { changed = true; diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc index c5183f8d3aee99696ed4114c3f7e451888222137..619e87caa5b6d0f6ec3c3b1489b0d4f50ef29963 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc @@ -19,6 +19,7 @@ limitations under the License. #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" +#include "tensorflow/core/lib/strings/str_util.h" namespace xla { namespace { @@ -26,112 +27,140 @@ namespace { namespace op = xla::testing::opcode_matchers; class WhileLoopSimplifierTest : public HloVerifiedTestBase { - public: - // Makes a computation that contains a loop that runs num_iters times. - HloComputation* MakeSimpleLoop(int num_iters, HloModule* module); - - // 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); + protected: + // Makes an HloModule that contains a loop with `num_iters` iteration. + void MakeModuleWithSimpleLoop(int num_iters); + + // Similar to MakeModuleWithSimpleLoop except that the loop bound is passed to + // the loop-condition through an element of a tuple which is the + // loop-condition parameter. + void MakeModuleWithSimpleLoopTupleElementLoopBound(int num_iters); }; -HloComputation* WhileLoopSimplifierTest::MakeSimpleLoop(int num_iters, - HloModule* module) { - HloComputation::Builder builder(TestName()); - - auto loop_iter_init = builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(42))); - auto loop_data_init = builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR1({0, 1, 2}))); - auto loop_init = builder.AddInstruction( - HloInstruction::CreateTuple({loop_iter_init, loop_data_init})); - - HloComputation* condition; - { - HloComputation::Builder cond_builder(TestName() + ".condition"); - auto loop_var = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - auto loop_induction_var = - cond_builder.AddInstruction(HloInstruction::CreateGetTupleElement( - ShapeUtil::MakeShape(S32, {}), loop_var, 0)); - auto limit = cond_builder.AddInstruction(HloInstruction::CreateConstant( - Literal::CreateR0(42 + num_iters))); - cond_builder.AddInstruction(HloInstruction::CreateBinary( - ShapeUtil::MakeShape(PRED, {}), HloOpcode::kLt, loop_induction_var, - limit)); - condition = module->AddEmbeddedComputation(cond_builder.Build()); +void WhileLoopSimplifierTest::MakeModuleWithSimpleLoop(int num_iters) { + string hlo_string_template = R"( + HloModule SimpleLoop + SimpleLoop.body { + loop_var.1 = (s32[], s32[3]{0}) parameter(0) + get-tuple-element.1 = s32[] get-tuple-element(loop_var.1), index=0 + constant.1 = s32[] constant(1) + add = s32[] add(get-tuple-element.1, constant.1) + get-tuple-element.2 = s32[3]{0} get-tuple-element(loop_var.1), index=1 + multiply = s32[3]{0} multiply(get-tuple-element.2, get-tuple-element.2) + ROOT tuple = (s32[], s32[3]{0}) tuple(add, multiply) } - - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - auto loop_var = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - auto loop_induction_var = - body_builder.AddInstruction(HloInstruction::CreateGetTupleElement( - ShapeUtil::MakeShape(S32, {}), loop_var, 0)); - auto new_loop_induction_var = - body_builder.AddInstruction(HloInstruction::CreateBinary( - loop_induction_var->shape(), HloOpcode::kAdd, loop_induction_var, - body_builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(1))))); - auto loop_data = - body_builder.AddInstruction(HloInstruction::CreateGetTupleElement( - loop_data_init->shape(), loop_var, 1)); - auto new_loop_data = - body_builder.AddInstruction(HloInstruction::CreateBinary( - loop_data_init->shape(), HloOpcode::kMultiply, loop_data, - loop_data)); - body_builder.AddInstruction( - HloInstruction::CreateTuple({new_loop_induction_var, new_loop_data})); - body = module->AddEmbeddedComputation(body_builder.Build()); + SimpleLoop.condition { + loop_var.2 = (s32[], s32[3]{0}) parameter(0) + get-tuple-element.3 = s32[] get-tuple-element(loop_var.2), index=0 + constant.2 = s32[] constant({{LOOP_BOUND}}) + ROOT less-than = pred[] less-than(get-tuple-element.3, constant.2) } + ENTRY SimpleLoop { + constant.3 = s32[] constant(42) + constant.4 = s32[3]{0} constant({0, 1, 2}) + tuple.1 = (s32[], s32[3]{0}) tuple(constant.3, constant.4) + ROOT while = (s32[], s32[3]{0}) while(tuple.1), condition= + SimpleLoop.condition, body=SimpleLoop.body + } + )"; - builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - - return module->AddEntryComputation(builder.Build()); + string hlo_string = tensorflow::str_util::StringReplace( + hlo_string_template, "{{LOOP_BOUND}}", + tensorflow::strings::StrCat(42 + num_iters), + /*replace_all=*/true); + ParseAndVerifyModule(hlo_string); } -HloComputation* WhileLoopSimplifierTest::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()); +void WhileLoopSimplifierTest::MakeModuleWithSimpleLoopTupleElementLoopBound( + int num_iters) { + string hlo_string_template = R"( + HloModule SimpleLoopWithIndirectLoopBound + SimpleLoopWithIndirectLoopBound.body { + loop_var.1 = (s32[], s32[3]{0}, s32[]) parameter(0) + get-tuple-element.1 = s32[] get-tuple-element(loop_var.1), index=0 + constant.1 = s32[] constant(1) + add = s32[] add(get-tuple-element.1, constant.1) + get-tuple-element.2 = s32[3]{0} get-tuple-element(loop_var.1), index=1 + multiply = s32[3]{0} multiply(get-tuple-element.2, get-tuple-element.2) + limit = s32[] get-tuple-element(loop_var.1), index=2 + ROOT tuple = (s32[], s32[3]{0}, s32[]) tuple(add, multiply, limit) + } + SimpleLoopWithIndirectLoopBound.condition { + loop_var.2 = (s32[], s32[3]{0}, s32[]) parameter(0) + get-tuple-element.3 = s32[] get-tuple-element(loop_var.2), index=0 + get-tuple-element.4 = s32[] get-tuple-element(loop_var.2), index=2 + ROOT less-than = pred[] less-than(get-tuple-element.3, get-tuple-element.4) + } + ENTRY SimpleLoopWithIndirectLoopBound { + constant.3 = s32[] constant(42) + constant.4 = s32[3]{0} constant({0, 1, 2}) + constant.2 = s32[] constant({{LOOP_BOUND}}) + tuple.1 = (s32[], s32[3]{0}, s32[]) tuple(constant.3, constant.4, + constant.2) + ROOT while = (s32[], s32[3]{0}, s32[]) while(tuple.1), + condition=SimpleLoopWithIndirectLoopBound.condition, + body=SimpleLoopWithIndirectLoopBound.body + } + )"; + + string hlo_string = tensorflow::str_util::StringReplace( + hlo_string_template, "{{LOOP_BOUND}}", + tensorflow::strings::StrCat(42 + num_iters), + /*replace_all=*/true); + ParseAndVerifyModule(hlo_string); } -TEST_F(WhileLoopSimplifierTest, WhileLoopWithZeroIterations) { - HloComputation* computation = MakeSimpleLoop(/*num_iters=*/0, &module()); - ASSERT_TRUE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), +TEST_F(WhileLoopSimplifierTest, LoopWithZeroIterationSimiplified) { + MakeModuleWithSimpleLoop(/*num_iters=*/0); + HloModule* the_module = &module(); + ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_THAT(the_module->entry_computation()->root_instruction(), op::Tuple(op::Constant(), op::Constant())); } -TEST_F(WhileLoopSimplifierTest, WhileLoopWithOneIteration) { - HloComputation* computation = MakeSimpleLoop(/*num_iters=*/1, &module()); - ASSERT_TRUE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), +TEST_F(WhileLoopSimplifierTest, + LoopWithZeroIterationTupleElementLoopBoundSimplified) { + MakeModuleWithSimpleLoopTupleElementLoopBound(/*num_iters=*/0); + HloModule* the_module = &module(); + ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_THAT(the_module->entry_computation()->root_instruction(), + op::Tuple(op::Constant(), op::Constant(), op::Constant())); +} + +TEST_F(WhileLoopSimplifierTest, LoopWithOneIterationSimplified) { + MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloModule* the_module = &module(); + ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_THAT(the_module->entry_computation()->root_instruction(), op::Tuple(op::Add(), op::Multiply())); } -TEST_F(WhileLoopSimplifierTest, WhileLoopWithTwoIterations) { - MakeSimpleLoop(/*num_iters=*/2, &module()); +TEST_F(WhileLoopSimplifierTest, + LoopWithOneIterationTupleELementLoopBoundSimplified) { + MakeModuleWithSimpleLoopTupleElementLoopBound(/*num_iters=*/1); + HloModule* the_module = &module(); + ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_THAT(the_module->entry_computation()->root_instruction(), + op::Tuple(op::Add(), op::Multiply(), op::Constant())); +} + +TEST_F(WhileLoopSimplifierTest, LoopWithTwoIterationsNotSimplified) { + MakeModuleWithSimpleLoop(/*num_iters=*/2); EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } -TEST_F(WhileLoopSimplifierTest, WhileLoopWithControlDependency) { - HloComputation* computation = MakeSimpleLoop(/*num_iters=*/1, &module()); +TEST_F(WhileLoopSimplifierTest, + LoopWithControlDependencySimplifiedDependencyPreserved) { + MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloModule* the_module = &module(); + HloComputation* computation = the_module->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* true_op = while_op->while_body()->AddInstruction( HloInstruction::CreateConstant(Literal::CreateR0(true))); TF_ASSERT_OK(true_op->AddControlDependencyTo( while_op->while_body()->root_instruction())); - ASSERT_TRUE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); EXPECT_THAT(computation->root_instruction()->control_predecessors(), ElementsAre(op::Constant())) << computation->ToString(); @@ -139,8 +168,10 @@ TEST_F(WhileLoopSimplifierTest, WhileLoopWithControlDependency) { // Loops that contain send/recv nodes can't be simplified; the loop structure // around send/recv nodes must be preserved. -TEST_F(WhileLoopSimplifierTest, NotRemovedIfContainsSend) { - HloComputation* computation = MakeSimpleLoop(/*num_iters=*/1, &module()); +TEST_F(WhileLoopSimplifierTest, LoopWithSendNotSimplified) { + MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloModule* the_module = &module(); + HloComputation* computation = the_module->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* while_body = while_op->while_body(); @@ -149,11 +180,13 @@ TEST_F(WhileLoopSimplifierTest, NotRemovedIfContainsSend) { HloInstruction::CreateConstant(Literal::CreateR0(true))), /*channel_id=*/0)); while_body->AddInstruction(HloInstruction::CreateSendDone(send)); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); } -TEST_F(WhileLoopSimplifierTest, NotRemovedIfContainsRecv) { - HloComputation* computation = MakeSimpleLoop(/*num_iters=*/1, &module()); +TEST_F(WhileLoopSimplifierTest, LoopWithRecvNotSimplified) { + MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloModule* the_module = &module(); + HloComputation* computation = the_module->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* while_body = while_op->while_body(); @@ -161,247 +194,217 @@ TEST_F(WhileLoopSimplifierTest, NotRemovedIfContainsRecv) { HloInstruction::CreateRecv(ShapeUtil::MakeShape(F32, {1}), /*channel_id=*/0)); while_body->AddInstruction(HloInstruction::CreateRecvDone(recv)); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); } // The limitation on not being able to simplify loops that contain infeeds (and // other non-removable instructions) isn't fundamental -- it just stems from the // fact that our infrastructure sees simplifying such a loop as tantamount to // removing the non-removable instruction. -TEST_F(WhileLoopSimplifierTest, NotRemovedIfContainsNonRemovableInstruction) { - HloComputation* computation = MakeSimpleLoop(/*num_iters=*/1, &module()); +TEST_F(WhileLoopSimplifierTest, LoopWithInfeedNotSimplified) { + MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloModule* the_module = &module(); + HloComputation* computation = the_module->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* while_body = while_op->while_body(); while_body->AddInstruction( HloInstruction::CreateInfeed(ShapeUtil::MakeShape(F32, {1}), "config")); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); } -// Check that we don't crash when given a loop whose shape is not a tuple. -TEST_F(WhileLoopSimplifierTest, IgnoreNonTupleShapedLoop) { - HloComputation::Builder builder(TestName()); - auto loop_init = builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(42))); - - HloComputation* condition; - { - HloComputation::Builder cond_builder(TestName() + ".condition"); - auto param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - cond_builder.AddInstruction(HloInstruction::CreateBinary( - ShapeUtil::MakeShape(PRED, {}), HloOpcode::kLt, param, - cond_builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(100))))); - condition = module().AddEmbeddedComputation(cond_builder.Build()); +// A non-tuple shaped loop shouldn't be simplified or crash the compiler. +TEST_F(WhileLoopSimplifierTest, NonTupleShapedLoopNotSimplified) { + const string hlo_string = R"( + HloModule NonTupleShapedLoop + NonTupleShapedLoop.body { + loop_var.1 = s32[] parameter(0) + constant.1 = s32[] constant(-1) + ROOT add = s32[] add(s32[] loop_var.1, s32[] constant.1) + } + NonTupleShapedLoop.condition { + loop_var = s32[] parameter(0) + constant = s32[] constant(100) + ROOT less-than = pred[] less-than(s32[] loop_var, s32[] constant) + } + ENTRY INonTupleShapedLoop { + constant.2 = s32[] constant(42) + ROOT while = s32[] while(s32[] constant.2), + condition=NonTupleShapedLoop.condition, + body=NonTupleShapedLoop.body } + )"; - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - auto param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - body_builder.AddInstruction(HloInstruction::CreateBinary( - ShapeUtil::MakeShape(S32, {}), HloOpcode::kAdd, param, - body_builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(-1))))); - body = module().AddEmbeddedComputation(body_builder.Build()); - } - - builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - - module().AddEntryComputation(builder.Build()); + ParseAndVerifyModule(hlo_string); EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } -// Construct a loop where we swap the tuple elements in each iteration. -// Although the tuple elements aren't used in the loop, we don't eliminate them, -// because the swapping side-effect is visible to users of the loop. -TEST_F(WhileLoopSimplifierTest, SwapTupleIndices) { - HloComputation::Builder builder(TestName()); - auto loop_init = builder.AddInstruction(HloInstruction::CreateTuple({ - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0))), - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(1))), - })); - - HloComputation* condition = - MakeAlwaysTrueComputation(loop_init->shape(), &module()); - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - auto param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); - body_builder.AddInstruction(HloInstruction::CreateTuple({ - body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)), - body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)), - })); - body = module().AddEmbeddedComputation(body_builder.Build()); +// A while loop that does nothing else besides swapping tuple elements +// can't be simplified as the result of the swapping is visible to users of the +// loop. +TEST_F(WhileLoopSimplifierTest, LoopSwappingTupleElementsNotSimplified) { + const string hlo_string = R"( + HloModule SwappingTupleElements + SwappingTupleElements.body { + loop_var = (s32[], s32[]) parameter(0) + get-tuple-element = s32[] get-tuple-element((s32[], s32[]) loop_var),index=1 + get-tuple-element.1 = s32[] get-tuple-element((s32[], s32[]) loop_var), + index=0 + ROOT tuple = (s32[], s32[]) tuple(s32[] get-tuple-element, + s32[] get-tuple-element.1) } + SwappingTupleElements.always_true { + param = (s32[], s32[]) parameter(0) + ROOT constant = pred[] constant(true) + } + ENTRY SwappingTupleElements { + x = s32[] parameter(0) + y = s32[] parameter(1) + tuple.1 = (s32[], s32[]) tuple(s32[] x, s32[] y) + ROOT while = (s32[], s32[]) while((s32[], s32[]) tuple.1), + condition=SwappingTupleElements.always_true, + body=SwappingTupleElements.body + } + )"; - builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - - module().AddEntryComputation(builder.Build()); + ParseAndVerifyModule(hlo_string); EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } // Construct a loop where we assign a constant to tuple element 0 in each // iteration. We can't eliminate tuple element 0, even though we never use its // value. -TEST_F(WhileLoopSimplifierTest, UnusedButModifiedTupleElement) { - HloComputation::Builder builder(TestName()); - auto loop_init = builder.AddInstruction( - HloInstruction::CreateTuple({builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0)))})); - - HloComputation* condition = - MakeAlwaysTrueComputation(loop_init->shape(), &module()); - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - body_builder.AddInstruction(HloInstruction::CreateTuple({ - body_builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(1))), - })); - body = module().AddEmbeddedComputation(body_builder.Build()); +TEST_F(WhileLoopSimplifierTest, + LoopWithUnusedButModifiedTupleElementNotSimplified) { + const string hlo_string = R"( + HloModule UnusedButModifiedTupleElement + UnusedButModifiedTupleElement.body { + loop_var = (s32[]) parameter(0) + constant.1 = s32[] constant(1) + ROOT tuple = (s32[]) tuple(s32[] constant.1) } + UnusedButModifiedTupleElement.always_true { + param = (s32[]) parameter(0) + ROOT constant = pred[] constant(true) + } + ENTRY UnusedButModifiedTupleElement { + constant.2 = s32[] constant(0) + tuple.1 = (s32[]) tuple(s32[] constant.2) + ROOT while = (s32[]) while((s32[]) tuple.1), + condition=UnusedButModifiedTupleElement.always_true, + body=UnusedButModifiedTupleElement.body + } + )"; - builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - - module().AddEntryComputation(builder.Build()); + ParseAndVerifyModule(hlo_string); EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } // Nothing to simplify in a while loop whose tuple has 0 elements. -TEST_F(WhileLoopSimplifierTest, EmptyTuple) { - HloComputation::Builder builder(TestName()); - auto loop_init = builder.AddInstruction(HloInstruction::CreateTuple({})); - - HloComputation* condition = - MakeAlwaysTrueComputation(loop_init->shape(), &module()); - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "loop_var")); - body_builder.AddInstruction(HloInstruction::CreateTuple({})); - body = module().AddEmbeddedComputation(body_builder.Build()); +TEST_F(WhileLoopSimplifierTest, LoopWithEmptyTupleNotSimplified) { + const string hlo_string = R"( + HloModule EmptyTuple + EmptyTuple.body { + loop_var = () parameter(0) + ROOT tuple = () tuple() + } + EmptyTuple.always_true { + param = () parameter(0) + ROOT constant = pred[] constant(true) + } + ENTRY EmptyTuple { + tuple.1 = () tuple() + ROOT while = () while(() tuple.1), condition=EmptyTuple.always_true, + body=EmptyTuple.body } + )"; - builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - module().AddEntryComputation(builder.Build()); + ParseAndVerifyModule(hlo_string); EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } // While loop where one tuple element is used twice in the body, and thus can't // be simplified away. -TEST_F(WhileLoopSimplifierTest, ElemUsedTwice) { - HloComputation::Builder builder(TestName()); - auto loop_init = builder.AddInstruction(HloInstruction::CreateTuple({ - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0))), - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(1))), - })); - - HloComputation* condition = - MakeAlwaysTrueComputation(loop_init->shape(), &module()); - - auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - auto* param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_init->shape(), "param0")); - auto* gte0 = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_s32, param, /*index=*/0)); - // get0 is used twice in the loop body's tuple. - body_builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte0})); - body = module().AddEmbeddedComputation(body_builder.Build()); +TEST_F(WhileLoopSimplifierTest, LoopWithElemUsedTwiceNotSimplified) { + const string hlo_string = R"( + HloModule ElemUsedTwice + ElemUsedTwice.body { + param0 = (s32[], s32[]) parameter(0) + get-tuple-element = s32[] get-tuple-element((s32[], s32[]) param0), index=0 + ROOT tuple = (s32[], s32[]) tuple(s32[] get-tuple-element, + s32[] get-tuple-element) + } + ElemUsedTwice.always_true { + param = (s32[], s32[]) parameter(0) + ROOT constant = pred[] constant(true) } + ENTRY ElemUsedTwice { + x = s32[] parameter(0) + y = s32[] parameter(1) + tuple.1 = (s32[], s32[]) tuple(s32[] x, s32[] y) + ROOT while = (s32[], s32[]) while((s32[], s32[]) tuple.1), + condition=ElemUsedTwice.always_true, body=ElemUsedTwice.body + } + )"; - builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - module().AddEntryComputation(builder.Build()); + ParseAndVerifyModule(hlo_string); EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } // This while loop has three tuple elements. Element 0 is unused and should be // removed. Element 1 is used by the loop body, and element 2 is used by the // loop condition; these two should stay. -TEST_F(WhileLoopSimplifierTest, RemoveUnusedOperand) { - HloComputation::Builder builder(TestName()); - auto loop_init = builder.AddInstruction(HloInstruction::CreateTuple({ - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0))), - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0))), - builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0))), - })); - auto loop_shape = loop_init->shape(); - auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); - - HloComputation* condition; - { - HloComputation::Builder cond_builder(TestName() + ".loop_condition"); - auto param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_shape, "param0")); - cond_builder.AddInstruction(HloInstruction::CreateBinary( - ShapeUtil::MakeShape(PRED, {}), HloOpcode::kEq, - cond_builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(0))), - cond_builder.AddInstruction(HloInstruction::CreateGetTupleElement( - scalar_s32, param, /*index=*/2)))); - condition = module().AddEmbeddedComputation(cond_builder.Build()); +TEST_F(WhileLoopSimplifierTest, RemoveUnusedLoopOperands) { + const string hlo_string = R"( + HloModule RemoveUnusedOperands + RemoveUnusedOperands.body { + loop_var = (s32[], s32[], s32[]) parameter(0) + get-tuple-element.1 = s32[] get-tuple-element((s32[], s32[], + s32[]) loop_var), index=0 + get-tuple-element.2 = s32[] get-tuple-element((s32[], s32[], + s32[]) loop_var), index=1 + constant.1 = s32[] constant(1) + add = s32[] add(s32[] get-tuple-element.2, s32[] constant.1) + get-tuple-element.3 = s32[] get-tuple-element((s32[], s32[], s32[]) + loop_var), index=2 + ROOT tuple = (s32[], s32[], s32[]) tuple(s32[] get-tuple-element.1, + s32[] add, s32[] get-tuple-element.3) } - - HloComputation* body; - { - HloComputation::Builder body_builder(TestName() + ".body"); - auto* param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_shape, "loop_var")); - - auto* tuple0 = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_s32, param, /*index=*/0)); - auto* tuple1 = body_builder.AddInstruction(HloInstruction::CreateBinary( - scalar_s32, HloOpcode::kAdd, - body_builder.AddInstruction(HloInstruction::CreateGetTupleElement( - scalar_s32, param, /*index=*/1)), - body_builder.AddInstruction( - HloInstruction::CreateConstant(Literal::CreateR0(1))))); - auto* tuple2 = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_s32, param, /*index=*/2)); - body_builder.AddInstruction( - HloInstruction::CreateTuple({tuple0, tuple1, tuple2})); - - body = module().AddEmbeddedComputation(body_builder.Build()); + RemoveUnusedOperands.loop_condition { + constant.2 = s32[] constant(0) + param0 = (s32[], s32[], s32[]) parameter(0) + get-tuple-element = s32[] get-tuple-element((s32[], s32[], s32[]) param0), + index=2 + ROOT equal-to = pred[] equal-to(s32[] constant.2, s32[] get-tuple-element) } + ENTRY RemoveUnusedOperands { + x = s32[] parameter(0) + constant.3 = s32[] constant(0) + y = s32[] parameter(1) + tuple.1 = (s32[], s32[], s32[]) tuple(s32[] x, s32[] constant.3, + s32[] y) + ROOT while = (s32[], s32[], s32[]) while((s32[], s32[], s32[]) tuple.1), + condition=RemoveUnusedOperands.loop_condition, + body=RemoveUnusedOperands.body + } + )"; + + ParseAndVerifyModule(hlo_string); + HloModule* the_module = &module(); + EXPECT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + + // The original while instruction is still left in the module as a dead + // instruction, find a while instruction with a different name as the new + // while instruction. + HloInstruction* new_while_op = + *std::find_if(the_module->entry_computation()->instructions().begin(), + the_module->entry_computation()->instructions().end(), + [&](const HloInstruction* instr) { + return (instr->opcode() == HloOpcode::kWhile && + instr->name() != "while"); + }); - auto* while_op = builder.AddInstruction(HloInstruction::CreateWhile( - loop_init->shape(), condition, body, loop_init)); - - module().AddEntryComputation(builder.Build()); - EXPECT_TRUE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); - - // We leave most of the checking to HloVerifiedTestBase, which runs the - // verifier on module() at the end of this test. - HloInstruction* new_while_op = *std::find_if( - module().entry_computation()->instructions().begin(), - module().entry_computation()->instructions().end(), - [&](const HloInstruction* instr) { - return instr != while_op && instr->opcode() == HloOpcode::kWhile; - }); + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); EXPECT_TRUE( ShapeUtil::Equal(new_while_op->shape(), ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}))) @@ -418,31 +421,91 @@ 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); +TEST_F(WhileLoopSimplifierTest, LoopWithNonTupleBodyShapeNotSimplified) { + const string hlo_string = R"( + HloModule BodyHasNonTupleRoot + BodyHasNonTupleRoot.passthrough { + ROOT param = (s32[], s32[]) parameter(0) + } + BodyHasNonTupleRoot.always_true { + param.1 = (s32[], s32[]) parameter(0) + ROOT constant = pred[] constant(true) + } + ENTRY BodyHasNonTupleRoot { + init_value = (s32[], s32[]) parameter(0) + ROOT while = (s32[], s32[]) while((s32[], s32[]) init_value), + condition=BodyHasNonTupleRoot.always_true, + body=BodyHasNonTupleRoot.passthrough + } + )"; + + ParseAndVerifyModule(hlo_string); + EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); +} + +TEST_F(WhileLoopSimplifierTest, + LoopWithNonTupleBodyRootInstructionNotSimplified) { + const string hlo_string = R"( + HloModule SimpleLoop + SimpleLoop.body { + loop_var.1 = (s32[], s32[3]{0}) parameter(0) + get-tuple-element.1 = s32[] get-tuple-element(loop_var.1), index=0 + constant.1 = s32[] constant(1) + add = s32[] add(get-tuple-element.1, constant.1) + get-tuple-element.2 = s32[3]{0} get-tuple-element(loop_var.1), index=1 + multiply = s32[3]{0} multiply(get-tuple-element.2, get-tuple-element.2) + ROOT custom-call = (s32[], s32[3]{0}) custom-call(add, multiply), + custom_call_target="x" + } + SimpleLoop.condition { + loop_var.2 = (s32[], s32[3]{0}) parameter(0) + get-tuple-element.3 = s32[] get-tuple-element(loop_var.2), index=0 + constant.2 = s32[] constant(44) + ROOT less-than = pred[] less-than(get-tuple-element.3, constant.2) + } + ENTRY SimpleLoop { + constant.3 = s32[] constant(42) + constant.4 = s32[3]{0} constant({0, 1, 2}) + tuple.1 = (s32[], s32[3]{0}) tuple(constant.3, constant.4) + ROOT while = (s32[], s32[3]{0}) while(tuple.1), condition= + SimpleLoop.condition, body=SimpleLoop.body + } + )"; + + ParseAndVerifyModule(hlo_string); + EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); +} + +TEST_F(WhileLoopSimplifierTest, LoopWithArrayConstantNotSimplified) { + const string hlo_string = R"( + HloModule SimpleLoop + SimpleLoop.body { + loop_var.1 = (s32[], s32[3]{0}, s32[3]{0}) parameter(0) + get-tuple-element.1 = s32[] get-tuple-element(loop_var.1), index=0 + constant.1 = s32[] constant(1) + add = s32[] add(get-tuple-element.1, constant.1) + get-tuple-element.2 = s32[3]{0} get-tuple-element(loop_var.1), index=1 + get-tuple-element.3 = s32[3]{0} get-tuple-element(loop_var.1), index=2 + add.2 = s32[3]{0} add(get-tuple-element.2, get-tuple-element.3) + ROOT tuple = (s32[], s32[3]{0}) tuple(add, add.2, get-tuple-element.3) + } + SimpleLoop.condition { + loop_var.2 = (s32[], s32[3]{0}, s32[3]{0}) parameter(0) + get-tuple-element.4 = s32[] get-tuple-element(loop_var.2), index=0 + constant.2 = s32[] constant(47) + ROOT less-than = pred[] less-than(get-tuple-element.4, constant.2) + } + ENTRY SimpleLoop { + constant.3 = s32[] constant(42) + constant.4 = s32[3]{0} constant({0, 1, 2}) + tuple.1 = (s32[], s32[3]{0}) tuple(constant.3, constant.4, constant.4) + ROOT while = (s32[], s32[3]{0}, s32[3]{0}) while(tuple.1), condition= + SimpleLoop.condition, body=SimpleLoop.body + } + )"; + + ParseAndVerifyModule(hlo_string); + EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); } } // namespace diff --git a/tensorflow/compiler/xla/service/while_util.cc b/tensorflow/compiler/xla/service/while_util.cc index e20b25e4a08a946f6b58575a4d4e557744f8035c..8cd5882f326acab156cbf42d50d83e9cf2d1687f 100644 --- a/tensorflow/compiler/xla/service/while_util.cc +++ b/tensorflow/compiler/xla/service/while_util.cc @@ -15,18 +15,21 @@ 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/hlo_creation_utils.h" #include "tensorflow/compiler/xla/service/tuple_util.h" +#include "tensorflow/core/lib/strings/strcat.h" namespace xla { +using tensorflow::strings::StrCat; + 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())); + HloComputation::Builder builder(StrCat("wide.", narrow_condition->name())); builder.AddInstruction( HloInstruction::CreateParameter(0, wide_shape, "wide_param")); @@ -57,8 +60,7 @@ 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())); + HloComputation::Builder builder(StrCat("wide.", narrow_body->name())); builder.AddInstruction( HloInstruction::CreateParameter(0, wide_shape, "wide_param")); return narrow_body->parent()->AddEmbeddedComputation(builder.Build()); @@ -137,4 +139,108 @@ WhileUtil::MakeInstructionsLiveIn( return std::move(result); } + +static StatusOr> +MakeCountedLoopConditionComputation(const Shape& loop_state_shape, + int64 trip_count) { + Shape scalar_pred = ShapeUtil::MakeShape(PRED, {}); + Shape scalar_s64 = ShapeUtil::MakeShape(S64, {}); + + TF_ASSIGN_OR_RETURN(std::unique_ptr cond_computation, + CreateComputationWithSignature( + {&loop_state_shape}, scalar_pred, "while_cond")); + + HloInstruction* trip_count_constant = cond_computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(trip_count))); + + HloInstruction* param = cond_computation->parameter_instruction(0); + TF_ASSIGN_OR_RETURN(HloInstruction * counter, + MakeGetTupleElementHlo(param, 0)); + TF_ASSIGN_OR_RETURN( + HloInstruction * compare, + MakeBinaryHlo(HloOpcode::kLt, counter, trip_count_constant)); + cond_computation->set_root_instruction(compare); + return std::move(cond_computation); +} + +static StatusOr> MakeCountedLoopBodyComputation( + const Shape& loop_state_shape, + const std::function( + HloInstruction*, const WhileUtil::LoopStateTy&)>& loop_body_generator) { + TF_ASSIGN_OR_RETURN(std::unique_ptr body_computation, + CreateComputationWithSignature( + {&loop_state_shape}, loop_state_shape, "while_body")); + HloInstruction* one = body_computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1))); + + HloInstruction* param = body_computation->parameter_instruction(0); + TF_ASSIGN_OR_RETURN(HloInstruction * indvar, + MakeGetTupleElementHlo(param, 0)); + TF_ASSIGN_OR_RETURN(HloInstruction * next_indvar, + MakeBinaryHlo(HloOpcode::kAdd, indvar, one)); + + std::vector loop_body_generator_args; + for (int64 i = 1, e = loop_state_shape.tuple_shapes_size(); i < e; i++) { + TF_ASSIGN_OR_RETURN(HloInstruction * tuple_element, + MakeGetTupleElementHlo(param, i)); + loop_body_generator_args.push_back(tuple_element); + } + TF_ASSIGN_OR_RETURN(std::vector next_state, + loop_body_generator(indvar, loop_body_generator_args)); + next_state.insert(next_state.begin(), next_indvar); + HloInstruction* next_state_tuple = + body_computation->AddInstruction(HloInstruction::CreateTuple(next_state)); + body_computation->set_root_instruction(next_state_tuple); + + return std::move(body_computation); +} + +static StatusOr MakeInitTupleFromInitValues( + HloComputation* computation, const WhileUtil::LoopStateTy& init_values) { + std::vector init_values_with_indvar; + init_values_with_indvar.reserve(init_values.size() + 1); + HloInstruction* zero = computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0))); + init_values_with_indvar.push_back(zero); + c_copy(init_values, std::back_inserter(init_values_with_indvar)); + return computation->AddInstruction( + HloInstruction::CreateTuple(init_values_with_indvar)); +} + +static Shape MakeLoopStateShape(const WhileUtil::LoopStateTy& init_values) { + std::vector loop_state_shape_components; + loop_state_shape_components.reserve(init_values.size() + 1); + loop_state_shape_components.push_back(ShapeUtil::MakeShape(S64, {})); + c_transform(init_values, std::back_inserter(loop_state_shape_components), + [](HloInstruction* instr) { return instr->shape(); }); + return ShapeUtil::MakeTupleShape(loop_state_shape_components); +} + +/*static*/ StatusOr WhileUtil::MakeCountedLoop( + HloComputation* computation, int64 trip_count, + const WhileUtil::LoopStateTy& init_values, + const WhileUtil::LoopBodyGeneratorTy& loop_body_generator) { + Shape loop_state_shape = MakeLoopStateShape(init_values); + TF_ASSIGN_OR_RETURN( + std::unique_ptr cond, + MakeCountedLoopConditionComputation(loop_state_shape, trip_count)); + TF_ASSIGN_OR_RETURN( + std::unique_ptr body, + MakeCountedLoopBodyComputation(loop_state_shape, loop_body_generator)); + TF_ASSIGN_OR_RETURN(HloInstruction * init_tuple, + MakeInitTupleFromInitValues(computation, init_values)); + HloModule* module = computation->parent(); + HloInstruction* while_instr = + computation->AddInstruction(HloInstruction::CreateWhile( + loop_state_shape, module->AddEmbeddedComputation(std::move(cond)), + module->AddEmbeddedComputation(std::move(body)), init_tuple)); + + std::vector result; + for (int64 i = 0, e = init_values.size(); i < e; i++) { + TF_ASSIGN_OR_RETURN(HloInstruction * user_state, + MakeGetTupleElementHlo(while_instr, i + 1)); + result.push_back(user_state); + } + return result; +} } // namespace xla diff --git a/tensorflow/compiler/xla/service/while_util.h b/tensorflow/compiler/xla/service/while_util.h index 3600b5a80d26e37fdb7d5173c3b8743734306390..80f7e16e64f4d1b1faa73f4fb9b4dd6443bf488b 100644 --- a/tensorflow/compiler/xla/service/while_util.h +++ b/tensorflow/compiler/xla/service/while_util.h @@ -52,6 +52,28 @@ class WhileUtil { static StatusOr MakeInstructionsLiveIn( HloInstruction* while_instr, tensorflow::gtl::ArraySlice instructions); + + using LoopStateTy = std::vector; + using LoopBodyGeneratorTy = std::function( + HloInstruction* /*induction_var*/, + const LoopStateTy& /*current_values*/)>; + + // Creates a while loop in `computation` that runs for `trip_count` + // iterations. The structure of the while loop is as follows, in pseudocode: + // + // loop_state while_loop() { + // indvar = 0; + // loop_state = init_values + // while (indvar < trip_count) { + // loop_state = loop_body_generator(loop_state) + // indvar++; + // } + // return loop_state; + // } + static StatusOr MakeCountedLoop( + HloComputation* computation, int64 trip_count, + const LoopStateTy& init_values, + const LoopBodyGeneratorTy& loop_body_generator); }; } // namespace xla diff --git a/tensorflow/compiler/xla/service_interface.h b/tensorflow/compiler/xla/service_interface.h index 809941d8fe1f63d66bf104e66eea66167a0f509d..d8235113dd800f7bab5ceb70272a598b9dcb1fbe 100644 --- a/tensorflow/compiler/xla/service_interface.h +++ b/tensorflow/compiler/xla/service_interface.h @@ -54,6 +54,9 @@ class ServiceInterface { virtual tensorflow::Status Execute(const ExecuteRequest* arg, ExecuteResponse* result) = 0; + virtual tensorflow::Status ExecuteGraph(const ExecuteGraphRequest* arg, + ExecuteResponse* result) = 0; + virtual tensorflow::Status ExecuteParallel( const ExecuteParallelRequest* arg, ExecuteParallelResponse* result) = 0; diff --git a/tensorflow/compiler/xla/shape_tree.h b/tensorflow/compiler/xla/shape_tree.h index d752619bd65751779c24f061e44e206d66b01465..ffaa40c2d673a2365342371ed8dab59565d1d08f 100644 --- a/tensorflow/compiler/xla/shape_tree.h +++ b/tensorflow/compiler/xla/shape_tree.h @@ -53,7 +53,7 @@ struct ShapeTreeNode { ShapeTreeNode(const ShapeTreeNode& other) : data(other.data), children(other.children.size()) { for (size_t i = 0; i < children.size(); ++i) { - children[i] = MakeUnique(*other.children[i]); + children[i] = ::xla::MakeUnique(*other.children[i]); } } @@ -62,7 +62,7 @@ struct ShapeTreeNode { data = other.data; children.resize(other.children.size()); for (size_t i = 0; i < children.size(); ++i) { - children[i] = MakeUnique(*other.children[i]); + children[i] = ::xla::MakeUnique(*other.children[i]); } } return *this; @@ -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 { @@ -433,7 +445,7 @@ class ShapeTreeIterator : public std::iterator(index, node_->data); + current_ = ::xla::MakeUnique(index, node_->data); return *current_; } @@ -480,7 +492,7 @@ void ShapeTree::InitChildren(const Shape& shape, Node* node) { template ShapeTree::ShapeTree(Shape shape) : root_(), - shape_storage_(MakeUnique(std::move(shape))), + shape_storage_(::xla::MakeUnique(std::move(shape))), shape_(shape_storage_.get()) { // The shape_ field is just used to hold the structure of the shape. // It should not be relied upon to store layout information. @@ -496,7 +508,7 @@ ShapeTree::ShapeTree(const Shape* shape) : root_(), shape_(shape) { template ShapeTree::ShapeTree(Shape shape, const T& init_value) : root_(init_value), - shape_storage_(MakeUnique(std::move(shape))), + shape_storage_(::xla::MakeUnique(std::move(shape))), shape_(shape_storage_.get()) { // The shape_ field is just used to hold the structure of the shape. // It should not be relied upon to store layout information. diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index d63e16ce2bf51cff0d113640d31ec6e70bfaf421..4f604e6f7cb18c1aaf844967d54e3b0e07e54b34 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -609,6 +609,8 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { /* static */ bool ShapeUtil::SameDimensions(const Shape& lhs, const Shape& rhs) { + CHECK(ShapeUtil::IsArray(lhs)); + CHECK(ShapeUtil::IsArray(rhs)); return ContainersEqual(lhs.dimensions(), rhs.dimensions()); } @@ -617,7 +619,10 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { return rhs.element_type() == TUPLE && ContainersEqual(lhs.tuple_shapes(), rhs.tuple_shapes(), Compatible); } - return SameDimensions(lhs, rhs) && SameElementType(lhs, rhs); + if (lhs.element_type() == OPAQUE) { + return rhs.element_type() == OPAQUE; + } + return SameElementType(lhs, rhs) && SameDimensions(lhs, rhs); } /* static */ bool ShapeUtil::CompatibleIgnoringElementType(const Shape& lhs, @@ -627,7 +632,26 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { ContainersEqual(lhs.tuple_shapes(), rhs.tuple_shapes(), CompatibleIgnoringElementType); } - return SameDimensions(lhs, rhs); + if (lhs.element_type() == OPAQUE) { + return rhs.element_type() == OPAQUE; + } + return ShapeUtil::IsArray(rhs) && 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 (lhs.element_type() == OPAQUE) { + return rhs.element_type() == OPAQUE; + } + if (SameElementTypeIgnoringFpPrecision(lhs, rhs)) { + return CompatibleIgnoringElementType(lhs, rhs); + } + return false; } /* static */ int64 ShapeUtil::GetDimension(const Shape& shape, @@ -1060,9 +1084,10 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, /* static */ bool ShapeUtil::TransposeIsBitcast( const Shape& input_shape, const Shape& output_shape, tensorflow::gtl::ArraySlice dimension_mapping) { - // Can't insert bitcasts without layout information. - if (!LayoutUtil::HasLayout(input_shape) && - !LayoutUtil::HasLayout(output_shape)) { + CHECK(LayoutUtil::HasLayout(input_shape) && + LayoutUtil::HasLayout(output_shape)); + + if (!SameElementType(input_shape, output_shape)) { return false; } @@ -1093,9 +1118,10 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, /* static */ bool ShapeUtil::ReshapeIsBitcast(const Shape& input_shape, const Shape& output_shape) { - // Can't convert reshapes into bitcasts without layout information. - if (!LayoutUtil::HasLayout(input_shape) || - !LayoutUtil::HasLayout(output_shape)) { + CHECK(LayoutUtil::HasLayout(input_shape) && + LayoutUtil::HasLayout(output_shape)); + + if (!SameElementType(input_shape, output_shape)) { return false; } diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index 453d4ec04726a4dd3851b8becb439bb7506e4ca9..3e130a02e2ce853ee157e46afb9760f5ff5a5026 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -23,6 +23,8 @@ limitations under the License. #include #include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -62,6 +64,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(); } @@ -204,6 +209,7 @@ class ShapeUtil { // Returns whether the LHS and RHS shapes have the same dimensions; note: does // not check element type. + // Precondition: IsArray(lhs) && IsArray(rhs) static bool SameDimensions(const Shape& lhs, const Shape& rhs); // Returns whether the lhs and rhs shapes have the same element type. @@ -211,6 +217,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. @@ -221,6 +252,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); @@ -287,6 +322,15 @@ class ShapeUtil { static Shape MakeShape(PrimitiveType element_type, tensorflow::gtl::ArraySlice dimensions); + // Creates a Shape with element type corresponding to T and the given + // dimensions + template + static Shape MakeShapeWithType( + tensorflow::gtl::ArraySlice dimensions) { + return ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), + dimensions); + } + // Constructs a new shape with the given minor_to_major order in its Layout. // Returns a value shape such that shape.has_layout(). static Shape MakeShapeWithLayout( @@ -489,12 +533,16 @@ class ShapeUtil { // Returns whether a transpose from input_shape to output_shape with dimension // mapping "dimension_mapping" produces a result which is bit-wise identical // to its input and thus may be replaced with a bitcast. + // + // Precondition: Both input_shape and output_shape have explicit layouts. static bool TransposeIsBitcast( const Shape& input_shape, const Shape& output_shape, tensorflow::gtl::ArraySlice dimension_mapping); // Returns whether a reshape from "input_shape" to "output_shape" is a // bitcast. + // + // Precondition: Both input_shape and output_shape have explicit layouts. static bool ReshapeIsBitcast(const Shape& input_shape, const Shape& output_shape); @@ -527,16 +575,16 @@ class ShapeUtil { // The visitor_function visitor function should return true if it wants to // continue, or false otherwise. // - // visitor_function must be a callable of type bool(const std::vector&) - // or compatible. + // visitor_function must be a callable of type + // StatusOr(ArraySlice) or compatible. template - static void ForEachIndex(const Shape& shape, - tensorflow::gtl::ArraySlice base, - tensorflow::gtl::ArraySlice count, - tensorflow::gtl::ArraySlice incr, - const FnType& visitor_function) { + static Status ForEachIndexWithStatus(const Shape& shape, + tensorflow::gtl::ArraySlice base, + tensorflow::gtl::ArraySlice count, + tensorflow::gtl::ArraySlice incr, + const FnType& visitor_function) { if (ShapeUtil::HasZeroElements(shape)) { - return; + return Status::OK(); } CHECK_EQ(Rank(shape), base.size()); CHECK_EQ(incr.size(), base.size()); @@ -546,7 +594,11 @@ class ShapeUtil { // once with the proper empty indexes. int64 n = -1; std::vector indexes(base.begin(), base.end()); - while (n < rank && visitor_function(indexes)) { + while (n < rank) { + TF_ASSIGN_OR_RETURN(bool should_continue, visitor_function(indexes)); + if (!should_continue) { + break; + } // Increments dimensions in minor to major order. for (n = 0; n < rank; ++n) { int64 dim = LayoutUtil::Minor(shape.layout(), n); @@ -557,6 +609,37 @@ class ShapeUtil { indexes[dim] = base[dim]; } } + + return Status::OK(); + } + + // Simple ergonomic wrapper around ShapeUtil::ForEachIndexWithStatus. + struct IndexIterationSpace { + std::vector index_base; + std::vector index_count; + std::vector index_incr; + }; + + template + static Status ForEachIndexWithStatus( + const Shape& shape, const IndexIterationSpace& iteration_space, + FnTy&& function) { + return ShapeUtil::ForEachIndexWithStatus( + shape, iteration_space.index_base, iteration_space.index_count, + iteration_space.index_incr, std::forward(function)); + } + + template + static void ForEachIndex(const Shape& shape, + tensorflow::gtl::ArraySlice base, + tensorflow::gtl::ArraySlice count, + tensorflow::gtl::ArraySlice incr, + const FnType& visitor_function) { + ForEachIndexWithStatus(shape, base, count, incr, + [&](tensorflow::gtl::ArraySlice indices) { + return StatusOr(visitor_function(indices)); + }) + .IgnoreError(); } private: diff --git a/tensorflow/compiler/xla/shape_util_test.cc b/tensorflow/compiler/xla/shape_util_test.cc index 81ba7afb95265398e830e26122cd0056a32daee3..424cfe37ea44d64884e08695fd1f49ca1970ca62 100644 --- a/tensorflow/compiler/xla/shape_util_test.cc +++ b/tensorflow/compiler/xla/shape_util_test.cc @@ -170,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}); @@ -184,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})}); @@ -193,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})}); @@ -210,6 +238,18 @@ TEST(ShapeUtilTest, IncompatibleTuplesWithDifferentDimensions) { EXPECT_FALSE(ShapeUtil::Compatible(tuple1, tuple2)); } +TEST(ShapeUtilTest, IncompatibleScalarVsTuple) { + Shape shape1 = ShapeUtil::MakeShape(F32, {}); + Shape shape2 = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {3, 2}), ShapeUtil::MakeShape(U32, {})}); + EXPECT_FALSE(ShapeUtil::Compatible(shape1, shape2)); + EXPECT_FALSE(ShapeUtil::Compatible(shape2, shape1)); + EXPECT_FALSE(ShapeUtil::CompatibleIgnoringElementType(shape1, shape2)); + EXPECT_FALSE(ShapeUtil::CompatibleIgnoringElementType(shape2, shape1)); + EXPECT_FALSE(ShapeUtil::CompatibleIgnoringFpPrecision(shape1, shape2)); + EXPECT_FALSE(ShapeUtil::CompatibleIgnoringFpPrecision(shape2, shape1)); +} + TEST(ShapeUtilTest, CompareShapesWithPaddedDimensionsMismatch) { Shape shape1 = ShapeUtil::MakeShape(F32, {20, 30}); shape1.mutable_layout()->add_padded_dimensions(10); @@ -545,10 +585,11 @@ TEST(ShapeUtilTest, ForEachIndex) { Shape shape = ShapeUtil::MakeShape(F32, data.dimensions); // Increments at every invocation. int invocations = 0; - auto increment_func = [&invocations](const std::vector& indexes) { - invocations++; - return true; - }; + auto increment_func = + [&invocations](tensorflow::gtl::ArraySlice indexes) { + invocations++; + return true; + }; std::vector zero_base(data.dimensions.size(), 0); std::vector step(data.dimensions.size(), 1); @@ -560,6 +601,29 @@ TEST(ShapeUtilTest, ForEachIndex) { } } +TEST(ShapeUtilTest, ForEachIndexWithStatus) { + Shape shape = ShapeUtil::MakeShape(F32, {10, 10}); + // Increments at every invocation. + int invocations = 0; + auto increment_func = + [&invocations]( + tensorflow::gtl::ArraySlice indexes) -> StatusOr { + if (++invocations == 5) { + return Unimplemented("Cannot increment beyond 5."); + } + return true; + }; + + Status error_status = ShapeUtil::ForEachIndexWithStatus( + shape, /*base=*/{0, 0}, /*count=*/{10, 10}, /*incr=*/{0, 1}, + increment_func); + + EXPECT_FALSE(error_status.ok()); + EXPECT_THAT(error_status.error_message(), + ::testing::HasSubstr("Cannot increment beyond 5.")); + EXPECT_EQ(invocations, 5); +} + TEST(ShapeUtilTest, DimensionsUnmodifiedByReshape_1x1x1x1_to_1x1x1) { // All output dimensions should be unmodified. One of the input dimensions is // modified because the input rank is larger by one. diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index d4820d1b6d6d8c848794b120a9064609ba72e5a9..025ac129d7040a007493cbb222d07c6cf323567f 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -44,6 +44,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:test", ], + alwayslink = True, ) cc_library( @@ -138,6 +139,7 @@ cc_library( "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_verifier", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -271,6 +273,9 @@ cc_library( xla_test( name = "bad_rng_shape_validation_test", srcs = ["bad_rng_shape_validation_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:test", @@ -290,6 +295,9 @@ xla_test( xla_test( name = "check_execution_arity_test", srcs = ["check_execution_arity_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -309,6 +317,9 @@ xla_test( xla_test( name = "query_inferred_shape_test", srcs = ["query_inferred_shape_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", @@ -325,6 +336,9 @@ xla_test( xla_test( name = "while_test", srcs = ["while_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -366,6 +380,9 @@ xla_test( xla_test( name = "axpy_simple_test", srcs = ["axpy_simple_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", @@ -430,6 +447,9 @@ xla_test( xla_test( name = "pred_test", srcs = ["pred_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla/client:computation_builder", @@ -444,6 +464,9 @@ xla_test( xla_test( name = "select_test", srcs = ["select_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", @@ -460,6 +483,7 @@ xla_test( xla_test( name = "conditional_test", srcs = ["conditional_test.cc"], + tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", @@ -476,6 +500,7 @@ xla_test( xla_test( name = "unary_op_test", srcs = ["unary_op_test.cc"], + tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", @@ -532,6 +557,9 @@ xla_test( xla_test( name = "deconstruct_tuple_test", srcs = ["deconstruct_tuple_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -575,6 +603,27 @@ 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"], @@ -601,6 +650,11 @@ xla_test( xla_test( name = "dot_operation_test", srcs = ["dot_operation_test.cc"], + shard_count = 20, + tags = [ + "enable_for_xla_interpreter", + "optonly", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -618,29 +672,19 @@ xla_test( ], ) -# Tests the dot operation in some cases that can be performed via a -# runtime call on some backends - e.g. a runtime call to Eigen. xla_test( - name = "dot_operation_runtime_test", - srcs = ["dot_operation_test.cc"], + name = "gather_operation_test", + srcs = ["gather_operation_test.cc"], deps = [ - "//tensorflow/compiler/xla:array2d", - "//tensorflow/compiler/xla:array3d", - "//tensorflow/compiler/xla:reference_util", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla/client:computation_builder", - "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/tests:client_library_test_base", - "//tensorflow/compiler/xla/tests:literal_test_util", - "//tensorflow/compiler/xla/tests:test_utils", + ":hlo_test_base", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/tests:xla_internal_test_main", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:test", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", ], ) -# Repeat dot_operation_runtime_test with single-threded eigen. +# Repeat dot_operation_runtime_test with single-threaded eigen. xla_test( name = "dot_operation_single_threaded_runtime_test", srcs = ["dot_operation_test.cc"], @@ -652,6 +696,8 @@ xla_test( "--xla_cpu_multi_thread_eigen=false", ], }, + shard_count = 20, + tags = ["optonly"], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -672,6 +718,9 @@ xla_test( xla_test( name = "transpose_test", srcs = ["transpose_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", @@ -690,6 +739,9 @@ xla_test( xla_test( name = "constants_test", srcs = ["constants_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -850,8 +902,7 @@ xla_test( name = "half_test", srcs = ["half_test.cc"], backends = [ - # TODO(b/72509305): Flaky (fails with SEGV) as of 2018-01-25 - # "cpu", + "cpu", "gpu", ], deps = [ @@ -875,6 +926,9 @@ xla_test( name = "slice_test", srcs = ["slice_test.cc"], shard_count = 40, + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", @@ -891,6 +945,9 @@ xla_test( xla_test( name = "multidimensional_slice_test", srcs = ["multidimensional_slice_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -907,6 +964,9 @@ xla_test( name = "dynamic_ops_test", timeout = "moderate", srcs = ["dynamic_ops_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", @@ -933,6 +993,9 @@ xla_test( xla_test( name = "tuple_test", srcs = ["tuple_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal_util", @@ -953,6 +1016,9 @@ xla_test( xla_test( name = "vector_ops_reduce_test", srcs = ["vector_ops_reduce_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -971,6 +1037,10 @@ xla_test( name = "reduce_test", srcs = ["reduce_test.cc"], shard_count = 40, + tags = [ + "enable_for_xla_interpreter", + "optonly", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1064,6 +1134,9 @@ xla_test( xla_test( name = "copy_test", srcs = ["copy_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ ":client_library_test_base", "//tensorflow/compiler/xla:array2d", @@ -1082,6 +1155,9 @@ xla_test( xla_test( name = "reduce_hlo_test", srcs = ["reduce_hlo_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ ":client_library_test_base", "//tensorflow/compiler/xla/tests:hlo_test_base", @@ -1095,6 +1171,9 @@ xla_test( xla_test( name = "call_test", srcs = ["call_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -1132,6 +1211,9 @@ xla_test( xla_test( name = "binop_scaling_test", srcs = ["binop_scaling_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1148,6 +1230,9 @@ xla_test( xla_test( name = "broadcast_simple_test", srcs = ["broadcast_simple_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1165,6 +1250,9 @@ xla_test( xla_test( name = "pad_test", srcs = ["pad_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1185,6 +1273,9 @@ xla_test( xla_test( name = "fmax_test", srcs = ["fmax_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", @@ -1198,6 +1289,9 @@ xla_test( xla_test( name = "log_test", srcs = ["log_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", @@ -1211,6 +1305,9 @@ xla_test( xla_test( name = "matrix_ops_simple_test", srcs = ["matrix_ops_simple_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal_util", @@ -1225,6 +1322,7 @@ xla_test( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -1253,6 +1351,9 @@ xla_test( name = "reshape_test", srcs = ["reshape_test.cc"], shard_count = 30, + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1279,6 +1380,9 @@ xla_test( xla_test( name = "reverse_test", srcs = ["reverse_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1295,6 +1399,9 @@ xla_test( xla_test( name = "vector_ops_simple_test", srcs = ["vector_ops_simple_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:shape_util", @@ -1318,6 +1425,9 @@ xla_test( xla_test( name = "concat_test", srcs = ["concat_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -1338,8 +1448,12 @@ xla_test( xla_test( name = "convert_test", srcs = ["convert_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", @@ -1394,6 +1508,9 @@ xla_test( xla_test( name = "floor_ceil_test", srcs = ["floor_ceil_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", @@ -1477,6 +1594,9 @@ xla_test( xla_test( name = "replay_test", srcs = ["replay_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:protobuf_util", @@ -1499,6 +1619,9 @@ xla_test( xla_test( name = "broadcast_test", srcs = ["broadcast_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -1555,6 +1678,9 @@ xla_test( xla_test( name = "fusion_test", srcs = ["fusion_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal_util", @@ -1566,6 +1692,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", @@ -1592,6 +1719,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 56fc21d019bb823f8f4631420a15fd607ef46a9a..6e21dda25d8e5151b31b8c2328253260595a94c4 100644 --- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc +++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc @@ -101,6 +101,33 @@ XLA_TEST_F(ArrayElementwiseOpTest, NegConstantC64) { {}, error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, NegConstantS64) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({ + -1, + 1, + 0, + 0x12345678, + static_cast(0xffffffff12345678l), + static_cast(0x8000000000000000LL), + static_cast(0x8000000000000001LL), + }); + auto result = builder.Neg(a); + LOG(INFO) << -static_cast(0x7FFFFFFFFFFFFFFFLL); + + ComputeAndCompareR1(&builder, + { + 1, + -1, + 0, + -0x12345678, + 0xedcba988, + static_cast(0x8000000000000000LL), + -static_cast(0x8000000000000001LL), + }, + {}); +} + XLA_TEST_F(ArrayElementwiseOpTest, IsFiniteZeroElementF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({}); @@ -186,6 +213,86 @@ XLA_TEST_F(ArrayElementwiseOpTest, AddTwoConstantZeroElementC64s) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, AddTwoConstantU64s) { + ComputationBuilder b(client_, TestName()); + + std::vector lhs{0xFFFFFFFF, + static_cast(-1), + 0, + 0, + 0x7FFFFFFFFFFFFFFFLL, + 0x7FFFFFFFFFFFFFFLL, + 0x8000000000000000LL, + 0x8000000000000000LL, + 1}; + std::unique_ptr lhs_literal = Literal::CreateR1({lhs}); + auto lhs_param = b.Parameter(0, lhs_literal->shape(), "lhs_param"); + std::unique_ptr lhs_data = + client_->TransferToServer(*lhs_literal).ConsumeValueOrDie(); + + std::vector rhs{1, + 0x7FFFFFFFFFFFFFFLL, + 0x7FFFFFFFFFFFFFFFLL, + 0x8000000000000000LL, + 0, + static_cast(-1), + 0, + 1, + 0x8000000000000000LL}; + std::unique_ptr rhs_literal = Literal::CreateR1({rhs}); + auto rhs_param = b.Parameter(1, rhs_literal->shape(), "rhs_param"); + std::unique_ptr rhs_data = + client_->TransferToServer(*rhs_literal).ConsumeValueOrDie(); + + auto add = b.Add(lhs_param, rhs_param); + + std::vector expected(lhs.size()); + for (int64 i = 0; i < lhs.size(); ++i) { + expected[i] = lhs[i] + rhs[i]; + } + + ComputeAndCompareR1(&b, expected, {lhs_data.get(), rhs_data.get()}); +} + +XLA_TEST_F(ArrayElementwiseOpTest, SubTwoConstantS64s) { + ComputationBuilder b(client_, TestName()); + + std::vector lhs{static_cast(0x8000000000000000LL), + static_cast(0x8000000000000000LL), + -1, + 0x7FFFFFFFFFFFFFFLL, + 0x7FFFFFFFFFFFFFFFLL, + 1, + 0, + -1}; + std::unique_ptr lhs_literal = Literal::CreateR1({lhs}); + auto lhs_param = b.Parameter(0, lhs_literal->shape(), "lhs_param"); + std::unique_ptr lhs_data = + client_->TransferToServer(*lhs_literal).ConsumeValueOrDie(); + + std::vector rhs{-1, + 0, + static_cast(0x8000000000000000LL), + 1, + 0, + 0x7FFFFFFFFFFFFFFLL, + 0x7FFFFFFFFFFFFFFFLL, + 0x7FFFFFFFFFFFFFFFLL}; + std::unique_ptr rhs_literal = Literal::CreateR1({rhs}); + auto rhs_param = b.Parameter(1, rhs_literal->shape(), "rhs_param"); + std::unique_ptr rhs_data = + client_->TransferToServer(*rhs_literal).ConsumeValueOrDie(); + + auto sub = b.Sub(lhs_param, rhs_param); + + std::vector expected(lhs.size()); + for (int64 i = 0; i < lhs.size(); ++i) { + expected[i] = lhs[i] - rhs[i]; + } + + ComputeAndCompareR1(&b, expected, {lhs_data.get(), rhs_data.get()}); +} + TEST_P(ArrayElementwiseOpTestParamCount, AddManyValues) { const int count = GetParam(); ComputationBuilder builder(client_, TestName()); @@ -847,68 +954,76 @@ XLA_TEST_F(ArrayElementwiseOpTest, NotZeroElementU32R1) { XLA_TEST_F(ArrayElementwiseOpTest, ShiftLeftS32) { ComputationBuilder builder(client_, TestName()); - auto a = - builder.ConstantR1({static_cast(0x12345678), - static_cast(0xF0001000), 1, 3, 77}); - auto b = builder.ConstantR1({4, 8, 2, 7, 15}); + auto a = builder.ConstantR1({static_cast(0x12345678), + static_cast(0xF0001000), 1, 3, 77, + 1, -3, 77}); + auto b = builder.ConstantR1({4, 8, 2, 7, 15, 32, 100, -1}); auto out = builder.ShiftLeft(a, b); - ComputeAndCompareR1( - &builder, - {static_cast(0x23456780), 0x00100000, 0x4, 0x180, 2523136}, {}); + ComputeAndCompareR1(&builder, + {static_cast(0x23456780), 0x00100000, 0x4, + 0x180, 2523136, 0, 0, 0}, + {}); } XLA_TEST_F(ArrayElementwiseOpTest, ShiftRightArithmeticS32) { ComputationBuilder builder(client_, TestName()); - auto a = - builder.ConstantR1({static_cast(0x92345678), - static_cast(0x10001000), 1, 3, 77}); - auto b = builder.ConstantR1({4, 8, 2, 7, 2}); + auto a = builder.ConstantR1({static_cast(0x92345678), + static_cast(0x10001000), 1, 3, 77, + 1, -3, 77}); + auto b = builder.ConstantR1({4, 8, 2, 7, 2, 32, 100, -1}); auto out = builder.ShiftRightArithmetic(a, b); - ComputeAndCompareR1(&builder, - {static_cast(0xF9234567), - static_cast(0x00100010), 0, 0, 19}, - {}); + ComputeAndCompareR1( + &builder, + {static_cast(0xF9234567), static_cast(0x00100010), 0, 0, 19, + 0, -1, 0}, + {}); } XLA_TEST_F(ArrayElementwiseOpTest, ShiftRightLogicalS32) { ComputationBuilder builder(client_, TestName()); - auto a = - builder.ConstantR1({static_cast(0x92345678), - static_cast(0x10001000), 1, 3, 77}); - auto b = builder.ConstantR1({4, 8, 2, 7, 5}); + auto a = builder.ConstantR1({static_cast(0x92345678), + static_cast(0x10001000), 1, 3, 77, + 1, -3, 77}); + auto b = builder.ConstantR1({4, 8, 2, 7, 5, 32, 100, -1}); auto out = builder.ShiftRightLogical(a, b); - ComputeAndCompareR1(&builder, {0x09234567, 0x00100010, 0, 0, 2}, {}); + ComputeAndCompareR1(&builder, + {0x09234567, 0x00100010, 0, 0, 2, 0, 0, 0}, {}); } XLA_TEST_F(ArrayElementwiseOpTest, ShiftLeftU32) { ComputationBuilder builder(client_, TestName()); - auto a = builder.ConstantR1({0x12345678, 0xF0001000, 1, 3, 77}); - auto b = builder.ConstantR1({4, 8, 2, 7, 15}); + auto a = builder.ConstantR1( + {0x12345678, 0xF0001000, 1, 3, 77, 1, ~3u, 77}); + auto b = builder.ConstantR1({4, 8, 2, 7, 15, 32, 100, ~0u}); auto out = builder.ShiftLeft(a, b); ComputeAndCompareR1( - &builder, {0x23456780, 0x00100000, 0x4, 0x180, 2523136}, {}); + &builder, {0x23456780, 0x00100000, 0x4, 0x180, 2523136, 0, 0, 0}, {}); } XLA_TEST_F(ArrayElementwiseOpTest, ShiftRightArithmeticU32) { ComputationBuilder builder(client_, TestName()); - auto a = builder.ConstantR1({0x92345678, 0x10001000, 1, 3, 77}); - auto b = builder.ConstantR1({4, 8, 2, 7, 2}); + auto a = builder.ConstantR1( + {0x92345678, 0x10001000, 1, 3, 77, 1, ~3u, 77}); + auto b = builder.ConstantR1({4, 8, 2, 7, 2, 32, 100, ~0u}); auto out = builder.ShiftRightArithmetic(a, b); - ComputeAndCompareR1(&builder, {0xF9234567, 0x00100010, 0, 0, 19}, {}); + ComputeAndCompareR1( + &builder, {0xF9234567, 0x00100010, 0, 0, 19, 0, ~0u, 0}, {}); } XLA_TEST_F(ArrayElementwiseOpTest, ShiftRightLogicalU32) { ComputationBuilder builder(client_, TestName()); - auto a = builder.ConstantR1({0x92345678, 0x10001000, 1, 3, 77}); - auto b = builder.ConstantR1({4, 8, 2, 7, 5}); + auto a = builder.ConstantR1( + {0x92345678, 0x10001000, 1, 3, 77, 1, ~3u, 77}); + auto b = builder.ConstantR1({4, 8, 2, 7, 5, 32, 100, ~0u}); auto out = builder.ShiftRightLogical(a, b); - ComputeAndCompareR1(&builder, {0x09234567, 0x00100010, 0, 0, 2}, {}); + ComputeAndCompareR1(&builder, + {0x09234567, 0x00100010, 0, 0, 2, 0, 0, 0}, {}); } XLA_TEST_F(ArrayElementwiseOpTest, CompareEqF32s) { @@ -1533,33 +1648,15 @@ XLA_TEST_F(ArrayElementwiseOpTest, SquareIn4DZeroElements) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -// GPU backend emits nvvm intrinsic for fmin and fmax, whose semantics is NOT -// such -// * fmin(NaN, x) = x -// * fmax(NaN, x) = x -// so we only test NAN on CPU. -// -// TODO(b/28180546): Make this compile in a way that is consistent -// among backends. XLA_TEST_F(ArrayElementwiseOpTest, MinF32s) { ComputationBuilder builder(client_, TestName()); -#if !defined(XLA_TEST_BACKEND_CPU) - auto lhs = builder.ConstantR1({1.0f, 1.0f, 2.25f}); - auto rhs = builder.ConstantR1({2.0f, -5.0f, 1.0f}); -#else SetFastMathDisabled(true); auto lhs = builder.ConstantR1({1.0f, 1.0f, 2.25f, NAN, 6.0f}); auto rhs = builder.ConstantR1({2.0f, -5.0f, 1.0f, 10.0f, NAN}); -#endif auto minimum = builder.Min(lhs, rhs); - ComputeAndCompareR1(&builder, -#if !defined(XLA_TEST_BACKEND_CPU) - {1.0f, -5.0f, 1.0f}, -#else - {1.0f, -5.0f, 1.0f, 10.0f, 6.0f}, -#endif - {}, error_spec_); + ComputeAndCompareR1(&builder, {1.0f, -5.0f, 1.0f, NAN, NAN}, {}, + error_spec_); } XLA_TEST_F(ArrayElementwiseOpTest, MinZeroElementF32s) { @@ -1570,50 +1667,26 @@ XLA_TEST_F(ArrayElementwiseOpTest, MinZeroElementF32s) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } -// TODO(b/28180546): Make this compile in a way that is consistent -// among backends. See comment on MinF32s test above. XLA_TEST_F(ArrayElementwiseOpTest, MinF64s) { ComputationBuilder builder(client_, TestName()); -#if !defined(XLA_TEST_BACKEND_CPU) - auto lhs = builder.ConstantR1({1.0, 1.0, 2.25}); - auto rhs = builder.ConstantR1({2.0, -5.0, 1.0}); -#else SetFastMathDisabled(true); auto lhs = builder.ConstantR1({1.0, 1.0, 2.25, NAN, 6.0}); auto rhs = builder.ConstantR1({2.0, -5.0, 1.0, 10.0, NAN}); -#endif auto minimum = builder.Min(lhs, rhs); - ComputeAndCompareR1(&builder, -#if !defined(XLA_TEST_BACKEND_CPU) - {1.0, -5.0, 1.0}, -#else - {1.0, -5.0, 1.0, 10.0, 6.0}, -#endif - {}, error_spec_); + ComputeAndCompareR1(&builder, {1.0, -5.0, 1.0, NAN, NAN}, {}, + error_spec_); } -// TODO(b/28180546): Make this compile in a way that is consistent -// among backends. See comment on MinF32s test above. XLA_TEST_F(ArrayElementwiseOpTest, MaxF32s) { ComputationBuilder builder(client_, TestName()); -#if !defined(XLA_TEST_BACKEND_CPU) - auto lhs = builder.ConstantR1({1.0f, 1.0f, 2.25f}); - auto rhs = builder.ConstantR1({2.0f, -5.0f, 1.0f}); -#else SetFastMathDisabled(true); auto lhs = builder.ConstantR1({1.0f, 1.0f, 2.25f, NAN, 6.0f}); auto rhs = builder.ConstantR1({2.0f, -5.0f, 1.0f, 10.0f, NAN}); -#endif auto maximum = builder.Max(lhs, rhs); - ComputeAndCompareR1(&builder, -#if !defined(XLA_TEST_BACKEND_CPU) - {2.0f, 1.0f, 2.25f}, -#else - {2.0f, 1.0f, 2.25f, 10.0f, 6.0f}, -#endif - {}, error_spec_); + ComputeAndCompareR1(&builder, {2.0f, 1.0f, 2.25f, NAN, NAN}, {}, + error_spec_); } XLA_TEST_F(ArrayElementwiseOpTest, MaxZeroElementF32s) { @@ -1624,27 +1697,15 @@ XLA_TEST_F(ArrayElementwiseOpTest, MaxZeroElementF32s) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } -// TODO(b/28180546): Make this compile in a way that is consistent -// among backends. See comment on MinF32s test above. XLA_TEST_F(ArrayElementwiseOpTest, MaxF64s) { ComputationBuilder builder(client_, TestName()); -#if !defined(XLA_TEST_BACKEND_CPU) - auto lhs = builder.ConstantR1({1.0, 1.0, 2.25}); - auto rhs = builder.ConstantR1({2.0, -5.0, 1.0}); -#else SetFastMathDisabled(true); auto lhs = builder.ConstantR1({1.0, 1.0, 2.25, NAN, 6.0}); auto rhs = builder.ConstantR1({2.0, -5.0, 1.0, 10.0, NAN}); -#endif auto maximum = builder.Max(lhs, rhs); - ComputeAndCompareR1(&builder, -#if !defined(XLA_TEST_BACKEND_CPU) - {2.0, 1.0, 2.25}, -#else - {2.0, 1.0, 2.25, 10.0, 6.0}, -#endif - {}, error_spec_); + ComputeAndCompareR1(&builder, {2.0, 1.0, 2.25, NAN, NAN}, {}, + error_spec_); } XLA_TEST_F(ArrayElementwiseOpTest, MaxS32s) { @@ -1879,20 +1940,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()); @@ -1995,47 +2109,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) // / / 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/broadcast_simple_test.cc b/tensorflow/compiler/xla/tests/broadcast_simple_test.cc index 03f5e08315bfed2bcb43ebb7098aaa0b97228605..97095f1cc427789845051a8fea24c95475286fe2 100644 --- a/tensorflow/compiler/xla/tests/broadcast_simple_test.cc +++ b/tensorflow/compiler/xla/tests/broadcast_simple_test.cc @@ -662,7 +662,7 @@ XLA_TEST_F(BroadcastSimpleTest, InvalidBinaryAndDegenerateBroadcasting) { auto result_status = Execute(&b, {}); EXPECT_FALSE(result_status.ok()); EXPECT_THAT(result_status.status().error_message(), - HasSubstr("broadcast dimension 0 mismatch")); + HasSubstr("dimension 0 mismatch")); } XLA_TEST_F(BroadcastSimpleTest, InvalidInDimensionBroadcasting) { @@ -675,7 +675,7 @@ XLA_TEST_F(BroadcastSimpleTest, InvalidInDimensionBroadcasting) { auto result_status = Execute(&b, {}); EXPECT_FALSE(result_status.ok()); EXPECT_THAT(result_status.status().error_message(), - HasSubstr("binary op BINOP_ADD with incompatible shapes")); + HasSubstr("op BINOP_ADD with incompatible shapes")); } XLA_TEST_F(BroadcastSimpleTest, InvalidDegenerateBroadcasting) { @@ -688,7 +688,7 @@ XLA_TEST_F(BroadcastSimpleTest, InvalidDegenerateBroadcasting) { auto result_status = Execute(&b, {}); EXPECT_FALSE(result_status.ok()); EXPECT_THAT(result_status.status().error_message(), - HasSubstr("binary op BINOP_ADD with incompatible shapes")); + HasSubstr("op BINOP_ADD with incompatible shapes")); } } // namespace diff --git a/tensorflow/compiler/xla/tests/concat_test.cc b/tensorflow/compiler/xla/tests/concat_test.cc index 1bcad5a3f37a37c9d482f3a5a899ac527666cca3..fb0e9c724a69b61801e6e0c2d07ef75b63a00465 100644 --- a/tensorflow/compiler/xla/tests/concat_test.cc +++ b/tensorflow/compiler/xla/tests/concat_test.cc @@ -75,7 +75,7 @@ XLA_TEST_F(ConcatTest, CannotConcatR0WithR0) { StatusOr computation_status = builder.Build(); ASSERT_FALSE(computation_status.ok()); EXPECT_THAT(computation_status.status().ToString(), - HasSubstr("dimension to concatenate along out of bounds: 0")); + HasSubstr("out of bounds: 0")); } XLA_TEST_F(ConcatTest, Concat_R1_L0_With_R1_L0) { diff --git a/tensorflow/compiler/xla/tests/conditional_test.cc b/tensorflow/compiler/xla/tests/conditional_test.cc index bc821674820fb128823786d7149037fc59b22ab6..b917dee77b5400db8f2c0a6a86258fee64723d71 100644 --- a/tensorflow/compiler/xla/tests/conditional_test.cc +++ b/tensorflow/compiler/xla/tests/conditional_test.cc @@ -571,5 +571,56 @@ XLA_TEST_F(ConditionalOpTest, ShapeMismatch) { "only parameter of true_computation")); } +XLA_TEST_F(ConditionalOpTest, SwappedInputsInSequentialConditionals) { + Shape tuple_shape = ShapeUtil::MakeTupleShape({r0f32_, r0f32_}); + Computation swapper; + { + ComputationBuilder builder(client_, TestName() + ".swapper"); + auto param0 = builder.Parameter(0, tuple_shape, "sp0"); + auto x = builder.GetTupleElement(param0, 0); + auto y = builder.GetTupleElement(param0, 1); + builder.Tuple({y, x}); + swapper = builder.Build().ConsumeValueOrDie(); + } + Computation forwarder; + { + ComputationBuilder builder(client_, TestName() + ".forwarder"); + auto param0 = builder.Parameter(0, tuple_shape, "fp0"); + auto x = builder.GetTupleElement(param0, 0); + auto y = builder.GetTupleElement(param0, 1); + builder.Tuple({x, y}); + forwarder = builder.Build().ConsumeValueOrDie(); + } + Computation main; + { + ComputationBuilder builder(client_, TestName() + ".main"); + auto param0 = builder.Parameter(0, tuple_shape, "mp0"); + auto x = builder.GetTupleElement(param0, 0); + auto y = builder.GetTupleElement(param0, 1); + auto lt_pred = builder.Lt(x, y); + auto res = builder.Conditional(lt_pred, param0, forwarder, param0, swapper); + auto ge_pred = builder.Ge(x, y); + builder.Conditional(ge_pred, res, swapper, res, forwarder); + main = builder.Build().ConsumeValueOrDie(); + } + + auto test_swap = [&](float a, float b) { + ComputationBuilder builder(client_, TestName()); + auto x = builder.ConstantR0(a); + auto y = builder.ConstantR0(b); + auto tuple_operand = builder.Tuple({x, y}); + builder.Call(main, {tuple_operand}); + + ComputeAndCompareTuple( + &builder, + *Literal::MakeTuple({Literal::CreateR0(a).get(), + Literal::CreateR0(b).get()}), + {}, error_spec_); + }; + + test_swap(3.11f, 9.4f); + test_swap(11.24f, 5.55f); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/convert_test.cc b/tensorflow/compiler/xla/tests/convert_test.cc index f66e3b57bf45fbc9f8ea786146d6fffe5d55a262..7926767a4f3af37cbfd0d80f669eabcb2bf52402 100644 --- a/tensorflow/compiler/xla/tests/convert_test.cc +++ b/tensorflow/compiler/xla/tests/convert_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" @@ -106,11 +107,126 @@ TEST_F(ConvertTest, ConvertR1F32ToR1S32) { XLA_TEST_F(ConvertTest, ConvertR1S64ToR1F32) { ComputationBuilder builder(client_, TestName()); - auto a = builder.ConstantR1({32, 64}); - builder.ConvertElementType(a, F32); + std::vector arg{ + -9223371216516022272, + -2, + -1, + -0x7FFFFFFF, + -0x80000000, + 0, + 1, + 2, + 1073742145, + 1073742656, + 0x7FFFFFFF, + 0x80000000, + 826720496944058148, + 4296062029846194332, + 0x0007FB72E4000000LL, + 0x0007FB72E4000001LL, + 0x0007FB72E6000000LL, + 0x0007FB72E7000000LL, + 0x0007FB72E7FFFFFFLL, + 0x0007FB72E8000000LL, + 0x0007FB72E8000001LL, + 0x0007FB72EA000000LL, + 0x0007FB72EB000000LL, + 0x0007FB72EBFFFFFFLL, + 0x0007FB72EC000000LL, + 0x7FFFFF0000000000LL, + 0x7FFFFF8000000000LL, + 0x7FFFFFFFFFFFFF00, + static_cast(0xFFFFFFFFFFFFFFFF), + static_cast(0x0000f234e67e0001LL), + static_cast(0x8000000000000000), + static_cast(0x8000000000000000LL), + static_cast(0x8000000000000001LL), + static_cast(0x8000008000000000LL), + static_cast(0x8000010000000000LL), + }; + std::unique_ptr arg_literal = Literal::CreateR1({arg}); + auto arg_param = builder.Parameter(0, arg_literal->shape(), "arg_param"); + std::unique_ptr arg_data = + client_->TransferToServer(*arg_literal).ConsumeValueOrDie(); + + builder.ConvertElementType(arg_param, F32); + + std::vector expected(arg.size()); + for (int64 i = 0; i < arg.size(); ++i) { + expected[i] = static_cast(arg[i]); + } + ComputeAndCompareR1(&builder, expected, {arg_data.get()}); +} - std::vector expected = {32.0, 64.0}; - ComputeAndCompareR1(&builder, expected, {}); +XLA_TEST_F(ConvertTest, ConvertR1U32ToR1F32) { + ComputationBuilder builder(client_, TestName()); + std::vector arg{0, 1, 0x1000, 0x7fffffff, + 0x80000000, 0x80000001, 0x80000002, 0x80000003, + 0x80000080, 0x80000081, 0x80000082, 0xFFFFFFFF}; + std::unique_ptr arg_literal = Literal::CreateR1({arg}); + auto arg_param = builder.Parameter(0, arg_literal->shape(), "arg_param"); + std::unique_ptr arg_data = + client_->TransferToServer(*arg_literal).ConsumeValueOrDie(); + + builder.ConvertElementType(arg_param, F32); + + std::vector expected(arg.size()); + for (int64 i = 0; i < arg.size(); ++i) { + expected[i] = static_cast(arg[i]); + } + ComputeAndCompareR1(&builder, expected, {arg_data.get()}); +} + +XLA_TEST_F(ConvertTest, ConvertR1F32ToR1U32) { + ComputationBuilder builder(client_, TestName()); + std::vector arg{0.0f, 1.0f, 16777216.0f, + 16777218.0f, 2147483647.0f, 4294967040.0f}; + std::unique_ptr arg_literal = Literal::CreateR1({arg}); + auto arg_param = builder.Parameter(0, arg_literal->shape(), "arg_param"); + std::unique_ptr arg_data = + client_->TransferToServer(*arg_literal).ConsumeValueOrDie(); + + builder.ConvertElementType(arg_param, U32); + + std::vector expected(arg.size()); + for (int64 i = 0; i < arg.size(); ++i) { + expected[i] = static_cast(arg[i]); + } + ComputeAndCompareR1(&builder, expected, {arg_data.get()}); +} + +XLA_TEST_F(ConvertTest, ConvertR1U32ToR1S64) { + ComputationBuilder builder(client_, TestName()); + std::vector arg{0, 1, 0x1000, 0x7fffffff, 0x80000082, 0xFFFFFFFF}; + std::unique_ptr arg_literal = Literal::CreateR1({arg}); + auto arg_param = builder.Parameter(0, arg_literal->shape(), "arg_param"); + std::unique_ptr arg_data = + client_->TransferToServer(*arg_literal).ConsumeValueOrDie(); + + builder.ConvertElementType(arg_param, S64); + + std::vector expected(arg.size()); + for (int64 i = 0; i < arg.size(); ++i) { + expected[i] = static_cast(arg[i]); + } + ComputeAndCompareR1(&builder, expected, {arg_data.get()}); +} + +XLA_TEST_F(ConvertTest, ConvertR1S32ToR1S64) { + ComputationBuilder builder(client_, TestName()); + std::vector arg{0, 1, 0x1000, -1, -0x1000}; + std::unique_ptr arg_literal = Literal::CreateR1({arg}); + auto arg_param = builder.Parameter(0, arg_literal->shape(), "arg_param"); + std::unique_ptr arg_data = + client_->TransferToServer(*arg_literal).ConsumeValueOrDie(); + + builder.ConvertElementType(arg_param, S64); + + std::vector expected(arg.size()); + for (int64 i = 0; i < arg.size(); ++i) { + expected[i] = static_cast(arg[i]); + } + ComputeAndCompareR1(&builder, expected, {arg_data.get()}); } XLA_TEST_F(ConvertTest, ConvertR1U8ToR1F32) { @@ -208,5 +324,65 @@ TEST_F(ConvertTest, ConvertReshape) { ComputeAndCompareR0(&builder, 42.0f, {}, ErrorSpec(0.0001)); } +std::vector GetInterestingF16ConversionTestCases() { + float infinity = std::numeric_limits::infinity(); + float half_min_positive_normal = + tensorflow::bit_cast(0x38800000); + float half_max_subnormal = tensorflow::bit_cast(0x387fc000); + float half_min_positive_subnormal = + tensorflow::bit_cast(0x33800000); + float half_max = 65504.0f; + + std::vector test_cases( + {-infinity, -(half_max * 2 + 1), -half_max, -42.0f, -1.0f, + -half_min_positive_subnormal, -half_max_subnormal, + -half_min_positive_normal, -0.0f, 0.0f, half_min_positive_subnormal, + half_max_subnormal, half_min_positive_normal, 1.0f, 42.0f, half_max, + (half_max * 2 + 1), infinity}); + return test_cases; +} + +XLA_TEST_F(ConvertTest, ConvertR1F16ToR1F32) { + std::vector test_cases = GetInterestingF16ConversionTestCases(); + std::vector input; + c_transform(test_cases, std::back_inserter(input), + [](float f) { return Eigen::half(f); }); + std::vector expected_output; + c_transform(input, std::back_inserter(expected_output), + [](Eigen::half h) { return static_cast(h); }); + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr dot_lhs_handle, + client_->TransferToServer(*Literal::CreateR1(input))); + + ComputationBuilder builder(client_, TestName()); + builder.ConvertElementType( + builder.Parameter( + 0, ShapeUtil::MakeShape(F16, {static_cast(input.size())}), + "param"), + F32); + + ComputeAndCompareR1(&builder, expected_output, {dot_lhs_handle.get()}); +} + +XLA_TEST_F(ConvertTest, ConvertR1F32ToR1F16) { + std::vector input = GetInterestingF16ConversionTestCases(); + std::vector expected_output; + c_transform(input, std::back_inserter(expected_output), + [](float f) { return Eigen::half(f); }); + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr dot_lhs_handle, + client_->TransferToServer(*Literal::CreateR1(input))); + + ComputationBuilder builder(client_, TestName()); + builder.ConvertElementType( + builder.Parameter( + 0, ShapeUtil::MakeShape(F32, {static_cast(input.size())}), + "param"), + F16); + + ComputeAndCompareR1(&builder, expected_output, {dot_lhs_handle.get()}); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 0ceb9aff378ae8aa8098be9360310b1d78d31ab2..72715398dea468d0000144759454c5f8d8673516 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -53,157 +53,185 @@ 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()); +#ifdef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +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 +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, {}, error_spec_); +TYPED_TEST_CASE(ForwardPassConvolution_3x3x256_256_OutputZ_Iota, TestTypes); +XLA_TYPED_TEST(ForwardPassConvolution_3x3x256_256_OutputZ_Iota, Types) { + this->RunTest(); } -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); - - 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 Convolve_1x1x1x2_1x1x1x2_Valid : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = ShapeUtil::MakeShapeWithType({1, 1, 1, 2}); + Shape filter_shape = ShapeUtil::MakeShapeWithType({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_); + } +}; - 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 = ShapeUtil::MakeShapeWithType({1, 1, 4, 4}); + Shape filter_shape = ShapeUtil::MakeShapeWithType({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, - {std::move(*Literal::CreateFromArray(input_data)), - std::move(*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, - {std::move(*Literal::CreateFromArray(input_data)), - std::move(*Literal::CreateFromArray(filter_data))}, - error_spec_); -} +template +class Convolve_1x1x4x4_1x1x2x2_Same : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = ShapeUtil::MakeShapeWithType({1, 1, 4, 4}); + Shape filter_shape = ShapeUtil::MakeShapeWithType({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, - {std::move(*Literal::CreateFromArray(input_data)), - std::move(*Literal::CreateFromArray(filter_data))}, - error_spec_); -} +template +class Convolve_1x1x4x4_1x1x3x3_Same : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = ShapeUtil::MakeShapeWithType({1, 1, 4, 4}); + Shape filter_shape = ShapeUtil::MakeShapeWithType({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 +260,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 = ShapeUtil::MakeShapeWithType({1, 2, 5}); + Shape filter_shape = ShapeUtil::MakeShapeWithType({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 +361,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 = ShapeUtil::MakeShapeWithType({1, 2, 5}); + Shape filter_shape = ShapeUtil::MakeShapeWithType({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 +434,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 +457,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 = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(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. @@ -519,67 +581,78 @@ 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 = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(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(); - - std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), 1.0); - - 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, - window_size * input_feature); - auto expected_r1 = Literal::CreateR1(expect_elems); - auto expected_r3 = expected_r1->Reshape({batch, num_windows, output_feature}) - .ConsumeValueOrDie(); +class Convolve1D1WindowTestFloat : public Convolve1D1WindowTestBase {}; - 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_); -} +XLA_TEST_P(Convolve1D1WindowTestFloat, Convolve1D1Window) { TestImpl(); } INSTANTIATE_TEST_CASE_P( - Convolve1D1WindowTest_Instantiation, Convolve1D1WindowTest, + Convolve1D1WindowTest_Instantiation, Convolve1D1WindowTestFloat, ::testing::Values(Convolve1DTestParam{1, 1, 1, 1, 2}, Convolve1DTestParam{160, 1, 1, 5, 1}, Convolve1DTestParam{24, 1, 1, 20, 1}, @@ -608,7 +681,49 @@ INSTANTIATE_TEST_CASE_P( ); -TEST_F(ConvolutionTest, Convolve_bf16_1x1x1x2_1x1x1x2_Valid) { +#if (XLA_TEST_BACKEND_GPU || XLA_TEST_BACKEND_CPU) +class Convolve1D1WindowTestHalf : public Convolve1D1WindowTestBase {}; + +XLA_TEST_P(Convolve1D1WindowTestHalf, Convolve1D1Window) { + TestImpl(); +} + +INSTANTIATE_TEST_CASE_P( + Convolve1D1WindowTest_Instantiation, Convolve1D1WindowTestHalf, + ::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}, +// TODO(b/72566306): The following five tests failed on CPU with unreasonable +// relative errors. Last ran on 2018-02-22. +#if XLA_TEST_BACKEND_GPU + 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, 1}, +#endif + Convolve1DTestParam{1, 10, 130, 1, 2}, + 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}) + +); +#endif + +XLA_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}); diff --git a/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc b/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc index 032c06cd3c9f872f57674d3d7b5adc201c91ea77..3ab0ea4ad48c00724d48e7d285ec024e10d5db31 100644 --- a/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc +++ b/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc @@ -195,7 +195,7 @@ XLA_TEST_F(DeconstructTupleTest, DeconstructNestedTuple) { auto result_status = client_->DeconstructTuple(*global_data); EXPECT_FALSE(result_status.ok()); EXPECT_THAT(result_status.status().error_message(), - HasSubstr("deconstructing nested tuples not yet supported")); + HasSubstr("Deconstructing nested tuples is not implemented")); } } // namespace diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index cc683701e6305510d202721fe645310f1009081c..09b1dd283e4d026a2f0007240d88cd9ac38acb19 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -34,169 +34,194 @@ limitations under the License. namespace xla { namespace { -// TODO(b/34468543): use GUnit typed tests when we can do all tests on all -// backends. class DotOperationTest : public ClientLibraryTestBase { public: ErrorSpec error_spec_{0.0001, 1e-5}; - - protected: - template - void TestOneElementVectorDot(); - template - void TestVectorDot(); - template - void TestSquareMatrixDot(bool lhs_row_major = false, - bool rhs_row_major = false); - template - void TestNonsquareMatrixDot(bool lhs_row_major = false, - bool rhs_row_major = false); }; -XLA_TEST_F(DotOperationTest, ZeroElementVectorDotF32) { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR1({}); - auto rhs = builder.ConstantR1({}); +#if defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16) && \ + defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT64) +using TypesF16F32 = ::testing::Types; +using TypesF16F32F64 = ::testing::Types; +using TypesF16F32F64CF64 = ::testing::Types; +#elif !defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16) && \ + !defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT64) +using TypesF16F32 = ::testing::Types; +using TypesF16F32F64 = ::testing::Types; +using TypesF16F32F64CF64 = + ::testing::Types; +#else +#error "Situation not handled yet" +#endif + +template +class DotOperationTest_F16F32F64CF64 : public DotOperationTest {}; +TYPED_TEST_CASE(DotOperationTest_F16F32F64CF64, TypesF16F32F64CF64); + +XLA_TYPED_TEST(DotOperationTest_F16F32F64CF64, ZeroElementVectorDot) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + + auto lhs = builder.ConstantR1({}); + auto rhs = builder.ConstantR1({}); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR0(&builder, 0.0, {}, error_spec_); + this->template ComputeAndCompareR0(&builder, static_cast(0.0), {}, + this->error_spec_); } -XLA_TEST_F(DotOperationTest, TrivialMatrixVectorDotF32) { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR2({{3.0, 4.0}}); - auto rhs = builder.ConstantR1({3.0, 4.0}); - auto result = builder.Dot(lhs, rhs); +template +class DotOperationTest_F16F32F64 : public DotOperationTest {}; +TYPED_TEST_CASE(DotOperationTest_F16F32F64, TypesF16F32F64); - ComputeAndCompareR1(&builder, {25.0}, {}, error_spec_); -} - -template -void DotOperationTest::TestOneElementVectorDot() { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR1({2.0}); - auto rhs = builder.ConstantR1({3.0}); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, TrivialMatrixVectorDot) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantR2FromArray2D({{3.0f, 4.0f}}); + auto rhs = builder.ConstantFromArray({3.0f, 4.0f}); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR0(&builder, 6.0, {}, error_spec_); + this->template ComputeAndCompareR1(&builder, {static_cast(25.0f)}, {}, + this->error_spec_); } -XLA_TEST_F(DotOperationTest, OneElementVectorDotF32) { - TestOneElementVectorDot(); -} +XLA_TYPED_TEST(DotOperationTest_F16F32F64, OneElementVectorDot) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantR1({static_cast(2.0f)}); + auto rhs = builder.ConstantR1({static_cast(3.0f)}); + auto result = builder.Dot(lhs, rhs); -XLA_TEST_F(DotOperationTest, OneElementVectorDotF64) { - TestOneElementVectorDot(); + this->template ComputeAndCompareR0(&builder, static_cast(6.0f), {}, + this->error_spec_); } -template -void DotOperationTest::TestVectorDot() { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR1({1.0, 2.5, 42.0}); - auto rhs = builder.ConstantR1({11.0, -1.0, 0.5}); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, VectorDot) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantFromArray({1.0f, 2.5f, 42.0f}); + auto rhs = builder.ConstantFromArray({11.0f, -1.0f, 0.5f}); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR0(&builder, 29.5, {}, error_spec_); + this->template ComputeAndCompareR0(&builder, static_cast(29.5f), {}, + this->error_spec_); } -XLA_TEST_F(DotOperationTest, VectorDotF32) { TestVectorDot(); } - -XLA_TEST_F(DotOperationTest, VectorDotF64) { TestVectorDot(); } - -namespace { - std::vector MinorToMajorForIsRowMajor(bool row_major) { return {row_major ? 1 : 0, row_major ? 0 : 1}; } -} // namespace - -XLA_TEST_F(DotOperationTest, Dot_0x2_2x0) { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR2FromArray2D(Array2D(0, 2)); - auto rhs = builder.ConstantR2FromArray2D(Array2D(2, 0)); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, Dot_0x2_2x0) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantR2FromArray2D(Array2D(0, 2)); + auto rhs = builder.ConstantR2FromArray2D(Array2D(2, 0)); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR2(&builder, Array2D(0, 0), {}, error_spec_); + this->template ComputeAndCompareR2(&builder, Array2D(0, 0), {}, + this->error_spec_); } -XLA_TEST_F(DotOperationTest, Dot_0x2_2x3) { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR2FromArray2D(Array2D(0, 2)); - auto rhs = builder.ConstantR2({{7.0, 8.0, 9.0}, {42.0, 77.0, 101.0}}); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, Dot_0x2_2x3) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantR2FromArray2D(Array2D(0, 2)); + auto rhs = builder.ConstantR2FromArray2D( + {{7.0f, 8.0f, 9.0f}, {42.0f, 77.0f, 101.0f}}); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR2(&builder, Array2D(0, 3), {}, error_spec_); + this->template ComputeAndCompareR2(&builder, Array2D(0, 3), {}, + this->error_spec_); } -XLA_TEST_F(DotOperationTest, Dot_3x2_2x0) { - ComputationBuilder builder(client_, TestName()); - auto lhs = - builder.ConstantR2({{7.0, 8.0}, {9.0, 42.0}, {77.0, 101.0}}); - auto rhs = builder.ConstantR2FromArray2D(Array2D(2, 0)); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, Dot_3x2_2x0) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantR2FromArray2D( + {{7.0f, 8.0f}, {9.0f, 42.0f}, {77.0f, 101.0f}}); + auto rhs = builder.ConstantR2FromArray2D(Array2D(2, 0)); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR2(&builder, Array2D(3, 0), {}, error_spec_); + this->template ComputeAndCompareR2(&builder, Array2D(3, 0), {}, + this->error_spec_); } -XLA_TEST_F(DotOperationTest, Dot_2x0_0x2) { - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR2FromArray2D(Array2D(2, 0)); - auto rhs = builder.ConstantR2FromArray2D(Array2D(0, 2)); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, Dot_2x0_0x2) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto lhs = builder.ConstantR2FromArray2D(Array2D(2, 0)); + auto rhs = builder.ConstantR2FromArray2D(Array2D(0, 2)); auto result = builder.Dot(lhs, rhs); - ComputeAndCompareR2(&builder, Array2D(2, 2, 0.0f), {}, - error_spec_); + this->template ComputeAndCompareR2( + &builder, Array2D(2, 2, static_cast(0.0f)), {}, this->error_spec_); } -XLA_TEST_F(DotOperationTest, FusedDot) { - ComputationBuilder builder(client_, TestName()); - auto param0 = builder.Parameter(0, ShapeUtil::MakeShape(F32, {2, 4}), "arg0"); - auto param1 = builder.Parameter(1, ShapeUtil::MakeShape(F32, {4, 1}), "arg1"); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, FusedDot) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto param0 = + builder.Parameter(0, ShapeUtil::MakeShapeWithType({2, 4}), "arg0"); + auto param1 = + builder.Parameter(1, ShapeUtil::MakeShapeWithType({4, 1}), "arg1"); auto exp0 = builder.Exp(param0); auto result = builder.Dot(exp0, param1); - auto lhs_handle = client_ - ->TransferToServer(*Literal::CreateR2( - {{1.0, 2.0, 3.0, 4.0}, {-1.0, -2.0, -3.0, -4.0}})) - .ConsumeValueOrDie(); - auto rhs_handle = client_ - ->TransferToServer(*Literal::CreateR2( - {{1.0}, {2.0}, {3.0}, {4.0}})) - .ConsumeValueOrDie(); - - ComputeAndCompareR2( - &builder, Array2D({{296.14560492846033}, {0.8611737683031964}}), - {lhs_handle.get(), rhs_handle.get()}, error_spec_); -} - -template -void DotOperationTest::TestSquareMatrixDot(bool lhs_row_major, - bool rhs_row_major) { auto lhs_handle = - client_ - ->TransferToServer(*Literal::CreateR2WithLayout( - {{1.0, 2.0}, {3.0, -4.0}}, - LayoutUtil::MakeLayout(MinorToMajorForIsRowMajor(lhs_row_major)))) - .ConsumeValueOrDie(); - auto rhs_handle = - client_ - ->TransferToServer(*Literal::CreateR2WithLayout( - {{1.0, 6.0}, {7.0, -4.0}}, - LayoutUtil::MakeLayout(MinorToMajorForIsRowMajor(rhs_row_major)))) + this->client_ + ->TransferToServer(*Literal::CreateR2FromArray2D( + {{1.0f, 2.0f, 3.0f, 4.0f}, {-1.0f, -2.0f, -3.0f, -4.0f}})) .ConsumeValueOrDie(); + auto rhs_handle = this->client_ + ->TransferToServer(*Literal::CreateR2FromArray2D( + {{1.0f}, {2.0f}, {3.0f}, {4.0f}})) + .ConsumeValueOrDie(); - ComputationBuilder builder(client_, TestName()); - auto prim_type = primitive_util::NativeToPrimitiveType(); - auto result = builder.Dot( - builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {2, 2}), "lhs"), - builder.Parameter(1, ShapeUtil::MakeShape(prim_type, {2, 2}), "rhs")); + if (std::is_same::value) { + this->error_spec_ = ErrorSpec{0.0001, 1e-3}; + } - Array2D expected({{15.0, -2.0}, {-25.0, 34.0}}); - ComputeAndCompareR2( - &builder, expected, {lhs_handle.get(), rhs_handle.get()}, error_spec_); + this->template ComputeAndCompareR2( + &builder, Array2D({{296.14560492846033f}, {0.8611737683031964f}}), + {lhs_handle.get(), rhs_handle.get()}, this->error_spec_); } +template +class SquareMatrixDot : public DotOperationTest { + public: + void TestImpl(bool lhs_row_major, bool rhs_row_major) { + auto lhs_handle = + client_ + ->TransferToServer(*Literal::CreateFromArrayWithLayout( + {{1.0f, 2.0f}, {3.0f, -4.0f}}, + LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(lhs_row_major)))) + .ConsumeValueOrDie(); + auto rhs_handle = + client_ + ->TransferToServer(*Literal::CreateFromArrayWithLayout( + {{1.0f, 6.0f}, {7.0f, -4.0f}}, + LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(rhs_row_major)))) + .ConsumeValueOrDie(); + ComputationBuilder builder(client_, TestName()); + auto prim_type = primitive_util::NativeToPrimitiveType(); + auto result = builder.Dot( + builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {2, 2}), "lhs"), + builder.Parameter(1, ShapeUtil::MakeShape(prim_type, {2, 2}), "rhs")); + + Array2D expected({{15.0f, -2.0f}, {-25.0f, 34.0f}}); + ComputeAndCompareR2(&builder, expected, + {lhs_handle.get(), rhs_handle.get()}, error_spec_); + } +}; + +TYPED_TEST_CASE(SquareMatrixDot, TypesF16F32F64CF64); +XLA_TYPED_TEST(SquareMatrixDot, TypesFF) { this->TestImpl(false, false); } +XLA_TYPED_TEST(SquareMatrixDot, TypesFT) { this->TestImpl(false, true); } +XLA_TYPED_TEST(SquareMatrixDot, TypesTF) { this->TestImpl(true, false); } +XLA_TYPED_TEST(SquareMatrixDot, TypesTT) { this->TestImpl(true, true); } + struct DotTestParam { int m; int k; @@ -225,33 +250,39 @@ string PrintDotTestParam( } class ParametricDotTest : public DotOperationTest, - public ::testing::WithParamInterface {}; + public ::testing::WithParamInterface { + protected: + template + void TestImpl(); +}; -XLA_TEST_P(ParametricDotTest, TestF32) { +template +void ParametricDotTest::TestImpl() { DotTestParam param = GetParam(); - std::unique_ptr> dot_lhs_data = - MakeLinspaceArray2D(0.0, 1.0, param.m, param.k); + std::unique_ptr> dot_lhs_data = + MakeLinspaceArray2D(0.0, 1.0, param.m, param.k); std::unique_ptr dot_lhs_lit = Literal::CreateR2FromArray2DWithLayout( *dot_lhs_data, LayoutUtil::MakeLayout( MinorToMajorForIsRowMajor(param.dot_lhs_row_major))); std::unique_ptr dot_lhs_handle = client_->TransferToServer(*dot_lhs_lit).ConsumeValueOrDie(); - std::unique_ptr> dot_rhs_data = - MakeLinspaceArray2D(0.0, 1.0, param.k, param.n); - std::unique_ptr dot_rhs_lit = Literal::CreateR2FromArray2DWithLayout( - *dot_rhs_data, LayoutUtil::MakeLayout( - MinorToMajorForIsRowMajor(param.dot_rhs_row_major))); + std::unique_ptr> dot_rhs_data = + MakeLinspaceArray2D(0.0, 1.0, param.k, param.n); + Layout rhs_layout = LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(param.dot_rhs_row_major)); + std::unique_ptr dot_rhs_lit = + Literal::CreateR2FromArray2DWithLayout(*dot_rhs_data, rhs_layout); std::unique_ptr dot_rhs_handle = client_->TransferToServer(*dot_rhs_lit).ConsumeValueOrDie(); - std::unique_ptr> addend_data; + std::unique_ptr> addend_data; std::unique_ptr addend_lit; std::unique_ptr addend_handle; if (param.has_addend) { - addend_data = MakeLinspaceArray2D(0.0, 1.0, param.m, param.n); + addend_data = MakeLinspaceArray2D(0.0, 1.0, param.m, param.n); addend_lit = Literal::CreateR2FromArray2DWithLayout( *addend_data, LayoutUtil::MakeLayout( MinorToMajorForIsRowMajor(param.addend_row_major))); @@ -259,24 +290,33 @@ XLA_TEST_P(ParametricDotTest, TestF32) { } ComputationBuilder builder(client_, TestName()); - auto prim_type = primitive_util::NativeToPrimitiveType(); + auto prim_type = primitive_util::NativeToPrimitiveType(); auto result = builder.Dot( - builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {param.m, param.k}), + builder.Parameter(0, + ShapeUtil::MakeShapeWithLayout( + prim_type, {param.m, param.k}, + MinorToMajorForIsRowMajor(param.dot_lhs_row_major)), "dot_lhs"), - builder.Parameter(1, ShapeUtil::MakeShape(prim_type, {param.k, param.n}), + builder.Parameter(1, + ShapeUtil::MakeShapeWithLayout( + prim_type, {param.k, param.n}, + MinorToMajorForIsRowMajor(param.dot_rhs_row_major)), "dot_rhs")); if (param.has_addend) { result = builder.Add( - result, - builder.Parameter( - 2, ShapeUtil::MakeShape(prim_type, {param.m, param.n}), "addend")); + result, builder.Parameter( + 2, + ShapeUtil::MakeShapeWithLayout( + prim_type, {param.m, param.n}, + MinorToMajorForIsRowMajor(param.addend_row_major)), + "addend")); } - std::unique_ptr> expected; + std::unique_ptr> expected; if (param.has_addend) { expected = ReferenceUtil::ApplyElementwise2D( - std::plus(), + std::plus(), *ReferenceUtil::MatmulArray2D(*dot_lhs_data, *dot_rhs_data), *addend_data); } else { @@ -287,8 +327,11 @@ XLA_TEST_P(ParametricDotTest, TestF32) { if (param.has_addend) { args.push_back(addend_handle.get()); } - - ComputeAndCompareR2(&builder, *expected, args, ErrorSpec(0.3, 3e-3)); + ErrorSpec error_spec(0.3, 3e-3); + if (std::is_same::value) { + error_spec = ErrorSpec(0.3, 5e-3); + } + ComputeAndCompareR2(&builder, *expected, args, error_spec); } std::vector CreateDotTestParameters() { @@ -305,30 +348,77 @@ std::vector CreateDotTestParameters() { } }; + add_matrix_matrix_dot_test(/*m=*/12, /*k=*/117, /*n=*/7); + add_matrix_matrix_dot_test(/*m=*/270, /*k=*/270, /*n=*/520); + add_matrix_matrix_dot_test(/*m=*/260, /*k=*/3, /*n=*/520); + + return params; +} + +#ifndef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +XLA_TEST_P(ParametricDotTest, TestF16) { TestImpl(); } +#endif +XLA_TEST_P(ParametricDotTest, TestF32) { TestImpl(); } +XLA_TEST_P(ParametricDotTest, TestF64) { TestImpl(); } + +INSTANTIATE_TEST_CASE_P(DotTests, ParametricDotTest, + ::testing::ValuesIn(CreateDotTestParameters()), + PrintDotTestParam); + +class ParametricDotTestWithoutLayoutAssignment : public ParametricDotTest { + public: + ParametricDotTestWithoutLayoutAssignment() { + execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( + "layout-assignment"); + } +}; + +std::vector CreateNoLayoutAssignmentDotTestParameters() { + std::vector params; + auto add_matrix_vector_dot_test = [&](int k, int n) { - for (bool has_addend : {false, true}) { - params.push_back({/*m=*/1, /*k=*/k, /*n=*/n, - /*dot_lhs_row_major=*/true, /*dot_rhs_row_major=*/true, - /*has_addend=*/has_addend, /*addend_row_major=*/true}); - if (n != 1) { - params.push_back( - {/*m=*/n, /*k=*/k, /*n=*/1, - /*dot_lhs_row_major=*/true, /*dot_rhs_row_major=*/true, - /*has_addend=*/has_addend, /*addend_row_major=*/true}); + for (bool lhs_row_major : {true, false}) { + for (bool rhs_row_major : {true, false}) { + for (bool has_addend : {true, false}) { + params.push_back({/*m=*/1, /*k=*/k, /*n=*/n, + /*dot_lhs_row_major=*/lhs_row_major, + /*dot_rhs_row_major=*/rhs_row_major, + /*has_addend=*/has_addend, + /*addend_row_major=*/true}); + if (has_addend) { + params.push_back({/*m=*/1, /*k=*/k, /*n=*/n, + /*dot_lhs_row_major=*/lhs_row_major, + /*dot_rhs_row_major=*/rhs_row_major, + /*has_addend=*/has_addend, + /*addend_row_major=*/false}); + } + if (n != 1) { + params.push_back({/*m=*/n, /*k=*/k, /*n=*/1, + /*dot_lhs_row_major=*/lhs_row_major, + /*dot_rhs_row_major=*/rhs_row_major, + /*has_addend=*/has_addend, + /*addend_row_major=*/true}); + if (has_addend) { + params.push_back({/*m=*/n, /*k=*/k, /*n=*/1, + /*dot_lhs_row_major=*/lhs_row_major, + /*dot_rhs_row_major=*/rhs_row_major, + /*has_addend=*/has_addend, + /*addend_row_major=*/false}); + } + } + } } } }; - add_matrix_matrix_dot_test(/*m=*/12, /*k=*/117, /*n=*/7); - add_matrix_matrix_dot_test(/*m=*/270, /*k=*/270, /*n=*/520); - add_matrix_matrix_dot_test(/*m=*/260, /*k=*/3, /*n=*/520); - add_matrix_vector_dot_test(/*k=*/8, /*n=*/8); add_matrix_vector_dot_test(/*k=*/130, /*n=*/8); add_matrix_vector_dot_test(/*k=*/8, /*n=*/130); add_matrix_vector_dot_test(/*k=*/290, /*n=*/130); add_matrix_vector_dot_test(/*k=*/1, /*n=*/1); add_matrix_vector_dot_test(/*k=*/1, /*n=*/16); + add_matrix_vector_dot_test(/*k=*/1, /*n=*/4); + add_matrix_vector_dot_test(/*k=*/1, /*n=*/3); add_matrix_vector_dot_test(/*k=*/3, /*n=*/16); add_matrix_vector_dot_test(/*k=*/3, /*n=*/3); add_matrix_vector_dot_test(/*k=*/29, /*n=*/29); @@ -339,109 +429,60 @@ std::vector CreateDotTestParameters() { return params; } -INSTANTIATE_TEST_CASE_P(DotTests, ParametricDotTest, - ::testing::ValuesIn(CreateDotTestParameters()), - PrintDotTestParam); - -XLA_TEST_F(DotOperationTest, SquareMatrixDotF32MinorToMajorFF) { - TestSquareMatrixDot(false, false); -} - -XLA_TEST_F(DotOperationTest, SquareMatrixDotF32MinorToMajorFT) { - TestSquareMatrixDot(false, true); -} - -XLA_TEST_F(DotOperationTest, SquareMatrixDotF32MinorToMajorTF) { - TestSquareMatrixDot(true, false); -} - -XLA_TEST_F(DotOperationTest, SquareMatrixDotF32MinorToMajorTT) { - TestSquareMatrixDot(true, true); +#ifndef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +XLA_TEST_P(ParametricDotTestWithoutLayoutAssignment, TestF16) { + TestImpl(); } - -XLA_TEST_F(DotOperationTest, SquareMatrixDotC64MinorToMajorFF) { - TestSquareMatrixDot(false, false); +#endif +XLA_TEST_P(ParametricDotTestWithoutLayoutAssignment, TestF32) { + TestImpl(); } - -XLA_TEST_F(DotOperationTest, SquareMatrixDotC64MinorToMajorFT) { - TestSquareMatrixDot(false, true); -} - -XLA_TEST_F(DotOperationTest, SquareMatrixDotC64MinorToMajorTF) { - TestSquareMatrixDot(true, false); -} - -XLA_TEST_F(DotOperationTest, SquareMatrixDotC64MinorToMajorTT) { - TestSquareMatrixDot(true, true); -} - -XLA_TEST_F(DotOperationTest, SquareMatrixDotF64) { - TestSquareMatrixDot(); +XLA_TEST_P(ParametricDotTestWithoutLayoutAssignment, TestF64) { + TestImpl(); } -template -void DotOperationTest::TestNonsquareMatrixDot(bool lhs_row_major, - bool rhs_row_major) { - auto lhs_handle = - client_ - ->TransferToServer(*Literal::CreateR2WithLayout( - {{1.0, 2.0, 3.0}, {3.0, -4.0, -1.0}}, - LayoutUtil::MakeLayout(MinorToMajorForIsRowMajor(lhs_row_major)))) - .ConsumeValueOrDie(); - auto rhs_handle = - client_ - ->TransferToServer(*Literal::CreateR2WithLayout( - {{1.0, 6.0}, {2.0, 3.0}, {7.0, -4.0}}, - LayoutUtil::MakeLayout(MinorToMajorForIsRowMajor(rhs_row_major)))) - .ConsumeValueOrDie(); - - ComputationBuilder builder(client_, TestName()); - auto prim_type = primitive_util::NativeToPrimitiveType(); - auto result = builder.Dot( - builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {2, 3}), "lhs"), - builder.Parameter(1, ShapeUtil::MakeShape(prim_type, {3, 2}), "rhs")); - - Array2D expected({{26.0, 0.0}, {-12.0, 10.0}}); - - ComputeAndCompareR2( - &builder, expected, {lhs_handle.get(), rhs_handle.get()}, error_spec_); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotF32MajorToMinorFF) { - TestNonsquareMatrixDot(false, false); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotF32MajorToMinorFT) { - TestNonsquareMatrixDot(false, true); -} +INSTANTIATE_TEST_CASE_P( + DotTests, ParametricDotTestWithoutLayoutAssignment, + ::testing::ValuesIn(CreateNoLayoutAssignmentDotTestParameters()), + PrintDotTestParam); -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotF32MajorToMinorTF) { - TestNonsquareMatrixDot(true, false); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotF32MajorToMinorTT) { - TestNonsquareMatrixDot(true, true); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotF64) { - TestNonsquareMatrixDot(); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotC64MajorToMinorFF) { - TestNonsquareMatrixDot(false, false); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotC64MajorToMinorFT) { - TestNonsquareMatrixDot(false, true); -} - -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotC64MajorToMinorTF) { - TestNonsquareMatrixDot(true, false); -} +template +class NonsquareMatrixDot : public DotOperationTest { + public: + void TestImpl(bool lhs_row_major, bool rhs_row_major) { + auto lhs_handle = + client_ + ->TransferToServer(*Literal::CreateFromArrayWithLayout( + {{1.0f, 2.0f, 3.0f}, {3.0f, -4.0f, -1.0f}}, + LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(lhs_row_major)))) + .ConsumeValueOrDie(); + auto rhs_handle = + client_ + ->TransferToServer(*Literal::CreateFromArrayWithLayout( + {{1.0f, 6.0f}, {2.0f, 3.0f}, {7.0f, -4.0f}}, + LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(rhs_row_major)))) + .ConsumeValueOrDie(); + + ComputationBuilder builder(client_, TestName()); + auto prim_type = primitive_util::NativeToPrimitiveType(); + auto result = builder.Dot( + builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {2, 3}), "lhs"), + builder.Parameter(1, ShapeUtil::MakeShape(prim_type, {3, 2}), "rhs")); + + Array2D expected({{26.0f, 0.0f}, {-12.0f, 10.0f}}); + + ComputeAndCompareR2(&builder, expected, + {lhs_handle.get(), rhs_handle.get()}, error_spec_); + } +}; -XLA_TEST_F(DotOperationTest, NonsquareMatrixDotC64MajorToMinorTT) { - TestNonsquareMatrixDot(true, true); -} +TYPED_TEST_CASE(NonsquareMatrixDot, TypesF16F32F64CF64); +XLA_TYPED_TEST(NonsquareMatrixDot, TestFF) { this->TestImpl(false, false); } +XLA_TYPED_TEST(NonsquareMatrixDot, TestFT) { this->TestImpl(false, true); } +XLA_TYPED_TEST(NonsquareMatrixDot, TestTF) { this->TestImpl(true, false); } +XLA_TYPED_TEST(NonsquareMatrixDot, TestTT) { this->TestImpl(true, true); } XLA_TEST_F(DotOperationTest, MatrixVectorC64) { auto lhs_handle = @@ -468,25 +509,35 @@ XLA_TEST_F(DotOperationTest, MatrixVectorC64) { &builder, expected, {lhs_handle.get(), rhs_handle.get()}, error_spec_); } -XLA_TEST_F(DotOperationTest, ConcurrentMatMul) { - ComputationBuilder builder(client_, TestName()); - auto matrix1 = builder.ConstantR2({{1.0, 2.0}, {3.0, 4.0}}); - auto matrix2 = builder.ConstantR2({{5.0, 6.0}, {7.0, 8.0}}); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, ConcurrentMatMult) { + using T = TypeParam; + + ComputationBuilder builder(this->client_, this->TestName()); + auto matrix1 = builder.ConstantR2FromArray2D({{1.0f, 2.0f}, {3.0f, 4.0f}}); + auto matrix2 = builder.ConstantR2FromArray2D({{5.0f, 6.0f}, {7.0f, 8.0f}}); auto matrix12 = builder.Dot(matrix1, matrix2); auto matrix21 = builder.Dot(matrix2, matrix1); builder.Add(matrix12, matrix21); - Array2D expected({{42.0, 56.0}, {74.0, 96.0}}); - ComputeAndCompareR2(&builder, expected, {}, error_spec_); + Array2D expected({{42.0f, 56.0f}, {74.0f, 96.0f}}); + this->template ComputeAndCompareR2(&builder, expected, {}, + this->error_spec_); } +template +class DotOperationTestForBatchMatMul : public DotOperationTest {}; +TYPED_TEST_CASE(DotOperationTestForBatchMatMul, TypesF16F32F64); + // Regression test for b/32055648. The root of the graph is a kFusion of 4 // bitcasts. Although bitcasts don't map to thunks, the root should still be // sync-dependent on bitcasts' operands. -XLA_TEST_F(DotOperationTest, BatchMatMul) { - ComputationBuilder builder(client_, TestName()); - auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {2, 2, 2, 2}), "x"); - auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {2, 2, 2, 2}), "y"); +XLA_TYPED_TEST(DotOperationTestForBatchMatMul, Types) { + using T = TypeParam; + ComputationBuilder builder(this->client_, this->TestName()); + auto x = + builder.Parameter(0, ShapeUtil::MakeShapeWithType({2, 2, 2, 2}), "x"); + auto y = + builder.Parameter(1, ShapeUtil::MakeShapeWithType({2, 2, 2, 2}), "y"); auto x_flat = builder.Reshape(x, {0, 1, 2, 3}, {4, 2, 2}); auto y_flat = builder.Reshape(y, {0, 1, 2, 3}, {4, 2, 2}); @@ -507,33 +558,79 @@ XLA_TEST_F(DotOperationTest, BatchMatMul) { auto out_flat = builder.ConcatInDim(out_slices, 0); builder.Reshape(out_flat, {0, 1, 2}, {2, 2, 2, 2}); - auto x_data = client_ - ->TransferToServer(*Literal::CreateR4( - {{{{1000, 100}, {10, 1}}, {{2000, 200}, {20, 2}}}, - {{{3000, 300}, {30, 3}}, {{4000, 400}, {40, 4}}}})) - .ConsumeValueOrDie(); - auto y_data = client_ - ->TransferToServer(*Literal::CreateR4( - {{{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}, - {{{11, 22}, {33, 44}}, {{55, 66}, {77, 88}}}})) + auto x_data = this->client_ + ->TransferToServer(*Literal::CreateR4FromArray4D( + {{{{1000.0f, 100.0f}, {10.0f, 1.0f}}, + {{2000.0f, 200.0f}, {20.0f, 2.0f}}}, + {{{3000.0f, 300.0f}, {30.0f, 3.0f}}, + {{4000.0f, 400.0f}, {40.0f, 4.0f}}}})) .ConsumeValueOrDie(); + auto y_data = + this->client_ + ->TransferToServer(*Literal::CreateR4FromArray4D( + {{{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}}, + {{{11.0f, 22.0f}, {33.0f, 44.0f}}, + {{55.0f, 66.0f}, {77.0f, 88.0f}}}})) + .ConsumeValueOrDie(); - ComputeAndCompareR4( + if (std::is_same::value) { + this->error_spec_ = ErrorSpec{0.0001, 1e-3}; + } + this->template ComputeAndCompareR4( &builder, - /*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_); + /*expected=*/ + {{{{1300.0f, 2400.0f}, {13.0f, 24.0f}}, + {{11400.0f, 13600.0f}, {114.0f, 136.0f}}}, + {{{42900.0f, 79200.0f}, {429.0f, 792.0f}}, + {{250800.0f, 299200.0f}, {2508.0f, 2992.0f}}}}, + {x_data.get(), y_data.get()}, this->error_spec_); +} + +XLA_TYPED_TEST(DotOperationTest_F16F32F64, GeneralMatMul) { + using T = TypeParam; + + ComputationBuilder builder(this->client_, this->TestName()); + auto x = + builder.Parameter(0, ShapeUtil::MakeShapeWithType({2, 2, 2}), "x"); + auto y = + builder.Parameter(1, ShapeUtil::MakeShapeWithType({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 = + this->client_ + ->TransferToServer(*Literal::CreateR3FromArray3D( + {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}})) + .ConsumeValueOrDie(); + + auto y_data = + this->client_ + ->TransferToServer(*Literal::CreateR3FromArray3D( + {{{1.0f, 0.0f}, {0.0f, 1.0f}}, {{1.0f, 0.0f}, {0.0f, 1.0f}}})) + .ConsumeValueOrDie(); + + this->template ComputeAndCompareR3( + &builder, + /*expected=*/ + {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}}, + {x_data.get(), y_data.get()}, this->error_spec_); } -TEST_F(DotOperationTest, TransposeFolding) { +XLA_TYPED_TEST(DotOperationTest_F16F32F64, TransposeFolding) { + using T = TypeParam; for (bool transpose_lhs : {false, true}) { for (bool transpose_rhs : {false, true}) { for (bool row_major : {false, true}) { - std::unique_ptr> lhs( - new Array2D({{1.0, 2.0, 3.0}, {3.0, -4.0, -1.0}})); - std::unique_ptr> rhs( - new Array2D({{1.0, 6.0}, {2.0, 3.0}, {7.0, -4.0}})); + std::unique_ptr> lhs( + new Array2D({{1.0f, 2.0f, 3.0f}, {3.0f, -4.0f, -1.0f}})); + std::unique_ptr> rhs( + new Array2D({{1.0f, 6.0f}, {2.0f, 3.0f}, {7.0f, -4.0f}})); if (transpose_lhs) { lhs = ReferenceUtil::TransposeArray2D(*lhs); @@ -542,22 +639,20 @@ TEST_F(DotOperationTest, TransposeFolding) { rhs = ReferenceUtil::TransposeArray2D(*rhs); } auto lhs_handle = - client_ - ->TransferToServer( - *Literal::CreateR2FromArray2DWithLayout( - *lhs, LayoutUtil::MakeLayout( - MinorToMajorForIsRowMajor(row_major)))) + this->client_ + ->TransferToServer(*Literal::CreateR2FromArray2DWithLayout( + *lhs, LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(row_major)))) .ConsumeValueOrDie(); auto rhs_handle = - client_ - ->TransferToServer( - *Literal::CreateR2FromArray2DWithLayout( - *rhs, LayoutUtil::MakeLayout( - MinorToMajorForIsRowMajor(row_major)))) + this->client_ + ->TransferToServer(*Literal::CreateR2FromArray2DWithLayout( + *rhs, LayoutUtil::MakeLayout( + MinorToMajorForIsRowMajor(row_major)))) .ConsumeValueOrDie(); - ComputationBuilder builder(client_, TestName()); - auto prim_type = primitive_util::NativeToPrimitiveType(); + ComputationBuilder builder(this->client_, this->TestName()); + auto prim_type = primitive_util::NativeToPrimitiveType(); auto lhs_arg = builder.Parameter( 0, ShapeUtil::MakeShape(prim_type, {lhs->height(), lhs->width()}), "lhs"); @@ -572,24 +667,27 @@ TEST_F(DotOperationTest, TransposeFolding) { } auto result = builder.Dot(lhs_arg, rhs_arg); - Array2D expected({{26.0, 0.0}, {-12.0, 10.0}}); + Array2D expected({{26.0f, 0.0f}, {-12.0f, 10.0f}}); VLOG(1) << "TestTransposeFolding " << transpose_lhs << " " << transpose_rhs << " " << row_major; - ComputeAndCompareR2(&builder, expected, - {lhs_handle.get(), rhs_handle.get()}, - error_spec_); + this->template ComputeAndCompareR2( + &builder, expected, {lhs_handle.get(), rhs_handle.get()}, + this->error_spec_); } } } } -TEST_F(DotOperationTest, DotOfConcatOptimizationWithConstLHS) { - auto prim_type = primitive_util::NativeToPrimitiveType(); +XLA_TYPED_TEST(DotOperationTest_F16F32F64, + DotOfConcatOptimizationWithConstLHS) { + using T = TypeParam; + auto prim_type = primitive_util::NativeToPrimitiveType(); - std::unique_ptr> constant_lhs_array(new Array2D( - {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}})); + std::unique_ptr> constant_lhs_array( + new Array2D({{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}, + {6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f}})); - ComputationBuilder builder(client_, TestName()); + ComputationBuilder builder(this->client_, this->TestName()); auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array); auto rhs_arg_0 = builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {2, 2}), "rhs_arg_0"); @@ -600,78 +698,80 @@ TEST_F(DotOperationTest, DotOfConcatOptimizationWithConstLHS) { auto result = builder.Dot( lhs_constant, builder.ConcatInDim({rhs_arg_0, rhs_arg_1, rhs_arg_2}, 0)); - std::unique_ptr> arg_0_value_array( - new Array2D({{1.0, 2.0}, {3.0, 4.0}})); - std::unique_ptr> arg_1_value_array( - new Array2D({{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}})); - std::unique_ptr> arg_2_value_array( - new Array2D({{1.0, 2.0}})); + std::unique_ptr> arg_0_value_array( + new Array2D({{1.0f, 2.0f}, {3.0f, 4.0f}})); + std::unique_ptr> arg_1_value_array( + new Array2D({{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}})); + std::unique_ptr> arg_2_value_array(new Array2D({{1.0f, 2.0f}})); TF_ASSERT_OK_AND_ASSIGN( auto arg_0_value, - client_->TransferToServer( - *Literal::CreateR2FromArray2D(*arg_0_value_array))); + this->client_->TransferToServer( + *Literal::CreateR2FromArray2D(*arg_0_value_array))); TF_ASSERT_OK_AND_ASSIGN( auto arg_1_value, - client_->TransferToServer( - *Literal::CreateR2FromArray2D(*arg_1_value_array))); + this->client_->TransferToServer( + *Literal::CreateR2FromArray2D(*arg_1_value_array))); TF_ASSERT_OK_AND_ASSIGN( auto arg_2_value, - client_->TransferToServer( - *Literal::CreateR2FromArray2D(*arg_2_value_array))); + this->client_->TransferToServer( + *Literal::CreateR2FromArray2D(*arg_2_value_array))); - Array2D expected({{53.0, 74.0}, {45.0, 66.0}}); - ComputeAndCompareR2( + Array2D expected({{53.0f, 74.0f}, {45.0f, 66.0f}}); + this->template ComputeAndCompareR2( &builder, expected, - {arg_0_value.get(), arg_1_value.get(), arg_2_value.get()}, error_spec_); -} - -TEST_F(DotOperationTest, DotOfConcatOptimizationWithConstRHS) { - auto prim_type = primitive_util::NativeToPrimitiveType(); - - std::unique_ptr> constant_rhs_array( - new Array2D({{1.0, 2.0}, - {3.0, 4.0}, - {5.0, 6.0}, - {6.0, 5.0}, - {4.0, 3.0}, - {2.0, 1.0}})); - - ComputationBuilder builder(client_, TestName()); + {arg_0_value.get(), arg_1_value.get(), arg_2_value.get()}, + this->error_spec_); +} + +XLA_TYPED_TEST(DotOperationTest_F16F32F64, + DotOfConcatOptimizationWithConstRHS) { + using T = TypeParam; + std::unique_ptr> constant_rhs_array( + new Array2D({{1.0f, 2.0f}, + {3.0f, 4.0f}, + {5.0f, 6.0f}, + {6.0f, 5.0f}, + {4.0f, 3.0f}, + {2.0f, 1.0f}})); + + ComputationBuilder builder(this->client_, this->TestName()); auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array); - auto lhs_arg_0 = builder.Parameter(0, ShapeUtil::MakeShape(prim_type, {2, 2}), + auto lhs_arg_0 = builder.Parameter(0, ShapeUtil::MakeShapeWithType({2, 2}), "lhs_arg_0"); - auto lhs_arg_1 = builder.Parameter(1, ShapeUtil::MakeShape(prim_type, {2, 3}), + auto lhs_arg_1 = builder.Parameter(1, ShapeUtil::MakeShapeWithType({2, 3}), "lhs_arg_1"); - auto lhs_arg_2 = builder.Parameter(2, ShapeUtil::MakeShape(prim_type, {2, 1}), + auto lhs_arg_2 = builder.Parameter(2, ShapeUtil::MakeShapeWithType({2, 1}), "lhs_arg_2"); auto result = builder.Dot( builder.ConcatInDim({lhs_arg_0, lhs_arg_1, lhs_arg_2}, 1), rhs_constant); - std::unique_ptr> arg_0_value_array( - new Array2D({{1.0, 2.0}, {3.0, 4.0}})); - std::unique_ptr> arg_1_value_array( - new Array2D({{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}})); - std::unique_ptr> arg_2_value_array( - new Array2D({{1.0}, {2.0}})); + std::unique_ptr> arg_0_value_array( + new Array2D({{1.0f, 2.0f}, {3.0f, 4.0f}})); + std::unique_ptr> arg_1_value_array( + new Array2D({{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}})); + std::unique_ptr> arg_2_value_array( + new Array2D({{1.0f}, {2.0f}})); TF_ASSERT_OK_AND_ASSIGN( auto arg_0_value, - client_->TransferToServer( - *Literal::CreateR2FromArray2D(*arg_0_value_array))); + this->client_->TransferToServer( + *Literal::CreateR2FromArray2D(*arg_0_value_array))); TF_ASSERT_OK_AND_ASSIGN( auto arg_1_value, - client_->TransferToServer( - *Literal::CreateR2FromArray2D(*arg_1_value_array))); + this->client_->TransferToServer( + *Literal::CreateR2FromArray2D(*arg_1_value_array))); TF_ASSERT_OK_AND_ASSIGN( auto arg_2_value, - client_->TransferToServer( - *Literal::CreateR2FromArray2D(*arg_2_value_array))); + this->client_->TransferToServer( + *Literal::CreateR2FromArray2D(*arg_2_value_array))); - Array2D expected({{38.0, 36.0}, {93.0, 91.0}}); - ComputeAndCompareR2( + Array2D expected({{38.0f, 36.0f}, {93.0f, 91.0f}}); + this->template ComputeAndCompareR2( &builder, expected, - {arg_0_value.get(), arg_1_value.get(), arg_2_value.get()}, error_spec_); + {arg_0_value.get(), arg_1_value.get(), arg_2_value.get()}, + this->error_spec_); } + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc index 877dc7db0eec229a7119b3627f177a33ed0d971b..4f354e6aefe70a51c09be1c0ca151af2bb9f0a2c 100644 --- a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc +++ b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc @@ -206,19 +206,19 @@ 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, UInt64R1) { TestR1(); } 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, Int64R2) { TestR2(); } XLA_TEST_F(DynamicSliceTest, UInt64R2) { TestR2(); } 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, UInt64R3) { TestR3(); } XLA_TEST_F(DynamicSliceTest, Int32R1Pred) { // Slice at dimension start. @@ -506,7 +506,7 @@ XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_CPU_PARALLEL(Int32R1BF16)) { } XLA_TEST_F(DynamicUpdateSliceTest, Int32R1) { TestR1(); } XLA_TEST_F(DynamicUpdateSliceTest, Int64R1) { TestR1(); } -XLA_TEST_F(DynamicUpdateSliceTest, UInt64R1) { 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)) { 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/gather_operation_test.cc b/tensorflow/compiler/xla/tests/gather_operation_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0830e9c8f0ccc38ecbc8dda413602b0fad1c2e0f --- /dev/null +++ b/tensorflow/compiler/xla/tests/gather_operation_test.cc @@ -0,0 +1,339 @@ +/* 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/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#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" + +// NB! TODO(b/74360564): These tests do not test out of bounds behavior since +// that hasn't been specced yet. + +namespace xla { +namespace { + +using tensorflow::gtl::nullopt; + +class GatherOperationTest : public HloTestBase { + protected: + void RunTest(const string& hlo_text, Literal* operand, + Literal* gather_indices) { + RunTest(hlo_text, {operand, gather_indices}); + } + + void RunTest(const string& hlo_text, + tensorflow::gtl::ArraySlice args) { + HloModuleConfig config; + config.set_debug_options(GetDebugOptionsForTest()); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + tools::Parse(hlo_text, config)); + EXPECT_TRUE(RunAndCompare(std::move(module), args, nullopt)); + } +}; + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherV1) { + const string hlo_text = R"( +HloModule TensorFlowGatherV1 + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[2,3] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0}, + gather_dims_to_operand_dims={0}, + index_vector_dim=1, + window_bounds={1, 3} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR1({0, 2}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherV2) { + const string hlo_text = R"( +HloModule TensorFlowGatherV2 + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[3,2] gather(operand, indices), + output_window_dims={0}, + elided_window_dims={1}, + gather_dims_to_operand_dims={1}, + index_vector_dim=1, + window_bounds={3, 1} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR1({0, 2}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherMultipleBatchDims) { + const string hlo_text = R"( +HloModule TensorFlowGatherMultipleBatchDims + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,3,2] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={1}, + gather_dims_to_operand_dims={1}, + index_vector_dim=2, + window_bounds={3, 1} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{0, 2}, {2, 1}}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherNdMultipleBatchDims_0) { + const string hlo_text = R"( +HloModule TensorFlowGatherNdMultipleBatchDims + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2,2,2] parameter(1) + ROOT gather = s32[2,2] gather(operand, indices), + output_window_dims={}, + elided_window_dims={0,1}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=2, + window_bounds={1, 1} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = + Literal::CreateR3({{{0, 2}, {2, 1}}, {{1, 2}, {2, 0}}}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherNdMultipleBatchDims_1) { + const string hlo_text = R"( +HloModule TensorFlowGatherNdMultipleBatchDims + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2,2,2] parameter(1) + ROOT gather = s32[2,1,1,2] gather(operand, indices), + output_window_dims={1,2}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=2, + window_bounds={1, 1} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = + Literal::CreateR3({{{0, 2}, {2, 1}}, {{1, 2}, {2, 0}}}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherNd) { + const string hlo_text = R"( +HloModule TensorFlowGatherNd + +ENTRY main { + operand = s32[3,3,2] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,2] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0,1}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=1, + window_bounds={1,1,2} +} +)"; + std::unique_ptr operand = + Literal::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // + {{-4, 4}, {-5, 5}, {-6, 6}}, // + {{-7, 7}, {-8, 8}, {-9, 9}}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{0, 0}, {1, 0}}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, TensorFlowGatherNdNonDefaultIndexVectorDim) { + const string hlo_text = R"( +HloModule TensorFlowGatherNd + +ENTRY main { + operand = s32[3,3,2] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,2] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0,1}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=0, + window_bounds={1,1,2} +} +)"; + std::unique_ptr operand = + Literal::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // + {{-4, 4}, {-5, 5}, {-6, 6}}, // + {{-7, 7}, {-8, 8}, {-9, 9}}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{0, 0}, {1, 0}}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, DynamicSlice) { + const char* hlo_text = R"( +HloModule DynamicSlice + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[1,1] gather(operand, indices), + output_window_dims={0,1}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=0, + window_bounds={1,1} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR1({1, 1}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, BatchDynamicSlice) { + const string hlo_text = R"( +HloModule BatchDynamicSlice + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2,2] parameter(1) + ROOT gather = s32[2,1,1] gather(operand, indices), + output_window_dims={1,2}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=0, + window_bounds={1,1} +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = + Literal::CreateR2({{2, 1}, {1, 1}}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, ZeroDimBounds) { + const char* hlo_text = R"( +HloModule TensorFlowGatherV1 + +ENTRY main { + operand = s32[3,0] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[2,0] gather(operand, indices), + output_window_dims={1}, + elided_window_dims={0}, + gather_dims_to_operand_dims={0}, + index_vector_dim=1, + window_bounds={1, 0} +} +)"; + std::unique_ptr operand = Literal::CreateR2({{}, {}, {}}); + std::unique_ptr gather_indices = Literal::CreateR1({0, 2}); + RunTest(hlo_text, operand.get(), gather_indices.get()); +} + +XLA_TEST_F(GatherOperationTest, OutOfBoundsIndex) { + // Out of bounds indices must not crash, and the indices in range should + // produce the same values across all backends. + // + // TODO(b/74360564): Once we have a well defined semantics for OOB accesses, + // we should get rid of the mask and check that backends produce the same + // value for OOB indices too. + + const string hlo_text = R"( +HloModule BatchDynamicSlice + +ENTRY main { + operand = s32[3,3]{1,0} parameter(0) + indices = s32[6,2]{1,0} parameter(1) + gather = s32[6,1,1]{2,1,0} gather(operand, indices), + output_window_dims={1,2}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=1, + window_bounds={1,1} + gather_reshaped = s32[6]{0} reshape(gather) + in_bounds_mask = s32[6]{0} parameter(2) + ROOT result = s32[6]{0} multiply(gather_reshaped, in_bounds_mask) +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR2( + {{2, 7}, {2, 1}, {1, 1}, {5, 1}, {2147483647, 1}, {1, 2}}); + std::unique_ptr in_bounds_mask = + Literal::CreateR1({0, 1, 1, 0, 0, 1}); + + RunTest(hlo_text, + {operand.get(), gather_indices.get(), in_bounds_mask.get()}); +} + +XLA_TEST_F(GatherOperationTest, NegativeIndex) { + // Negative indices must not crash, and the indices in range should produce + // the same values across all backends. + // + // TODO(b/74360564): Once we have a well defined semantics for negative + // accesses, we should get rid of the mask and check that backends produce the + // same value for negative indices too. + + const string hlo_text = R"( +HloModule BatchDynamicSlice + +ENTRY main { + operand = s32[3,3]{1,0} parameter(0) + indices = s32[6,2]{1,0} parameter(1) + gather = s32[6,1,1]{2,1,0} gather(operand, indices), + output_window_dims={1,2}, + elided_window_dims={}, + gather_dims_to_operand_dims={0,1}, + index_vector_dim=1, + window_bounds={1,1} + gather_reshaped = s32[6]{0} reshape(gather) + in_bounds_mask = s32[6]{0} parameter(2) + ROOT result = s32[6]{0} multiply(gather_reshaped, in_bounds_mask) +} +)"; + std::unique_ptr operand = + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); + std::unique_ptr gather_indices = Literal::CreateR2( + {{2, -1}, {2, 1}, {1, 1}, {-500, 1}, {-2147483648, 1}, {1, 2}}); + std::unique_ptr in_bounds_mask = + Literal::CreateR1({0, 1, 1, 0, 0, 1}); + + RunTest(hlo_text, + {operand.get(), gather_indices.get(), in_bounds_mask.get()}); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index 9f5806c5e16c30cf198027cffab5f78c315cb957..5f62c44f25dd62b563bd8ce02477bd741f264182 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -140,15 +140,10 @@ StatusOr> HloTestBase::MakeReferenceModule( return std::move(reference_module); } -template StatusOr<::testing::AssertionResult> HloTestBase::RunAndCompareInternal( - std::unique_ptr module, const ArraySlice arguments, + std::unique_ptr module, const ArraySlice arguments, const optional& error, bool run_hlo_passes, const std::function& reference_preprocessor) { - static_assert( - std::is_same::value || - std::is_same, LiteralPtr>::value, - "The LiteralPtr type only accepts Literal* or std::unique_ptr."); TF_RETURN_IF_ERROR( VerifyHloModule(*test_runner_.backend().platform(), module.get())); TF_ASSIGN_OR_RETURN(auto reference_module, @@ -165,9 +160,8 @@ StatusOr<::testing::AssertionResult> HloTestBase::RunAndCompareInternal( error); } -template ::testing::AssertionResult HloTestBase::RunAndCompare( - std::unique_ptr module, const ArraySlice arguments, + std::unique_ptr module, const ArraySlice arguments, const optional& error, const std::function& reference_preprocessor) { auto result = @@ -179,9 +173,8 @@ template return result.ValueOrDie(); } -template ::testing::AssertionResult HloTestBase::RunAndCompareNoHloPasses( - std::unique_ptr module, const ArraySlice arguments, + std::unique_ptr module, const ArraySlice arguments, const optional& error, const std::function& reference_preprocessor) { auto result = @@ -198,8 +191,14 @@ template const std::function& reference_preprocessor) { const auto& fake_arguments = MakeFakeArguments(module.get()).ConsumeValueOrDie(); - return RunAndCompare>( - std::move(module), fake_arguments, error, reference_preprocessor); + + std::vector fake_argument_ptrs; + c_transform( + fake_arguments, std::back_inserter(fake_argument_ptrs), + [](const std::unique_ptr& literal) { return literal.get(); }); + + return RunAndCompare(std::move(module), fake_argument_ptrs, error, + reference_preprocessor); } ::testing::AssertionResult HloTestBase::RunAndCompareNoHloPasses( @@ -207,8 +206,13 @@ template const std::function& reference_preprocessor) { const auto& fake_arguments = MakeFakeArguments(module.get()).ConsumeValueOrDie(); - return RunAndCompareNoHloPasses>( - std::move(module), fake_arguments, error, reference_preprocessor); + std::vector fake_argument_ptrs; + c_transform( + fake_arguments, std::back_inserter(fake_argument_ptrs), + [](const std::unique_ptr& literal) { return literal.get(); }); + + return RunAndCompareNoHloPasses(std::move(module), fake_argument_ptrs, error, + reference_preprocessor); } ::testing::AssertionResult HloTestBase::RunAndCompare( @@ -267,6 +271,28 @@ template reference_preprocessor); } +HloComputation* HloTestBase::FindComputation(HloModule* module, + tensorflow::StringPiece name) { + auto it = c_find_if(module->computations(), + [&](HloComputation* c) { return c->name() == name; }); + if (it == module->computations().end()) { + return nullptr; + } + return *it; +} + +HloInstruction* HloTestBase::FindInstruction(HloModule* module, + tensorflow::StringPiece name) { + for (const HloComputation* c : module->computations()) { + auto it = c_find_if(c->instructions(), + [&](HloInstruction* i) { return i->name() == name; }); + if (it != c->instructions().end()) { + return *it; + } + } + return nullptr; +} + Backend& HloTestBase::backend() { return test_runner_.backend(); } /* static */ diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h index 4aea9fc9fd027231106e529eb16bcd43f23fbe1c..e375f13a44e0618f2a498325859d928c1adb830e 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_test_base.h @@ -44,7 +44,7 @@ namespace xla { // enables, for one, explicitly building a graph of HLO instructions to run. // // This can also be used to write text/file-based test cases. Note that the test -// target is responsible for linking the needed backends. A covenient way to do +// target is responsible for linking the needed backends. A convenient way to do // this is to make it an xla_test: it will generate test targets linking with // the respective backends, which will be used as the test backend; the // interpreter backend is already linked with hlo_test_base so it will be the @@ -104,8 +104,7 @@ class HloTestBase : public ::testing::Test { // Executes the given hlo module on two backends and compares results. // - // 'arguments': the input of the hlo module. The LiteralPtr type accepts - // Literal* or std::unique_ptr. + // 'arguments': the input of the hlo module. // // 'error': if has value, expects the results to be near (within the error // bound). Otherwise, expects the results to be equal. @@ -114,20 +113,18 @@ class HloTestBase : public ::testing::Test { // backend, but it might need to be tailored so that it is able to run on the // reference backend. Note that the program shape of the module must not be // modified. - template ::testing::AssertionResult RunAndCompare( std::unique_ptr module, - const tensorflow::gtl::ArraySlice arguments, + const tensorflow::gtl::ArraySlice arguments, const tensorflow::gtl::optional& error, const std::function& reference_preprocessor = nullptr) TF_MUST_USE_RESULT; // Same as above, except that the module will be executed without Hlo // optimization. - template ::testing::AssertionResult RunAndCompareNoHloPasses( std::unique_ptr module, - const tensorflow::gtl::ArraySlice arguments, + const tensorflow::gtl::ArraySlice arguments, const tensorflow::gtl::optional& error, const std::function& reference_preprocessor = nullptr) TF_MUST_USE_RESULT; @@ -197,6 +194,15 @@ class HloTestBase : public ::testing::Test { ->Clear(); } + // Gets the computation/instruction from the given module with the given name. + // + // This is useful for tests which create HLOs from a string and then want to + // inspect a particular computation or instruction. + HloComputation* FindComputation(HloModule* module, + tensorflow::StringPiece name); + HloInstruction* FindInstruction(HloModule* module, + tensorflow::StringPiece name); + // Return an HLO verifier constructed for the test backend. HloVerifier& verifier() const { return *hlo_verifier_; } @@ -223,10 +229,9 @@ class HloTestBase : public ::testing::Test { // Runs the module on two platforms with or without running hlo passes and // compares the results. Returns whether the results are near or equal. If any // error happens before the results are computed, returns the error status. - template StatusOr<::testing::AssertionResult> RunAndCompareInternal( std::unique_ptr module, - const tensorflow::gtl::ArraySlice arguments, + const tensorflow::gtl::ArraySlice arguments, const tensorflow::gtl::optional& error, bool run_hlo_passes, const std::function& reference_preprocessor); }; diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc b/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc index 506091ddd8d1d8e6519525bb7031f4e8b296b5fb..da4cf4ae0c31bc194cd2ec9b845df36afbde69b0 100644 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_verifier.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/test.h" @@ -40,18 +41,22 @@ void HloVerifiedTestBase::TearDown() { << "TearDown called more than once; it should be called exactly once."; tear_down_called_ = true; if (module_) { - HloVerifier verifier; - xla::StatusOr mutated = verifier.Run(module_.get()); - if (!mutated.ok()) { - ADD_FAILURE() << "HloVerifier failed: " << mutated.status(); - } else { - EXPECT_FALSE(mutated.ValueOrDie()) - << "HloVerifier should never mutate the HloModule"; - } + VerifyModule(); } HloTestBase::TearDown(); } +void HloVerifiedTestBase::VerifyModule() { + HloVerifier verifier; + xla::StatusOr mutated = verifier.Run(module_.get()); + if (!mutated.ok()) { + ADD_FAILURE() << "HloVerifier failed: " << mutated.status(); + } else { + EXPECT_FALSE(mutated.ValueOrDie()) + << "HloVerifier should never mutate the HloModule"; + } +} + HloModule& HloVerifiedTestBase::module() { if (!module_) { module_ = CreateNewModule(); @@ -59,4 +64,10 @@ HloModule& HloVerifiedTestBase::module() { return *module_; } +void HloVerifiedTestBase::ParseAndVerifyModule( + tensorflow::StringPiece hlo_text) { + CHECK(!module_) << "Called ParseModule when test already has a module."; + TF_ASSERT_OK_AND_ASSIGN(module_, tools::Parse(hlo_text)); + VerifyModule(); +} } // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base.h b/tensorflow/compiler/xla/tests/hlo_verified_test_base.h index 492688bf7d682cf991cb8c09399492a0437f651b..e5bb14a8839acbdef8fd2b79bb0f574c46ea3d40 100644 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_verified_test_base.h @@ -44,6 +44,7 @@ class HloVerifiedTestBase : public HloTestBase { // Returns the default HloModule, lazily creating it if necessary via // HloTestBase::CreateNewModule(). HloModule& module(); + void ParseAndVerifyModule(tensorflow::StringPiece hlo_text); // Sets the shape-size function used during hlo verification. If this isn't // called, a default ShapeVerifier is used instead. @@ -55,6 +56,7 @@ class HloVerifiedTestBase : public HloTestBase { std::unique_ptr module_; // Lazily populated. Access via module(). std::unique_ptr shape_verifier_; bool tear_down_called_ = false; + void VerifyModule(); }; } // namespace xla diff --git a/tensorflow/compiler/xla/tests/literal_test_util.cc b/tensorflow/compiler/xla/tests/literal_test_util.cc index 39c07297d69b6683f0fbe75fdd2186effda5043c..81630df34c58526b6d41492b2b4b3892a02a21c2 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util.cc @@ -209,6 +209,11 @@ template <> return CompareFloatsBitwiseEqual(lhs, rhs); } template <> +::testing::AssertionResult CompareEqual(Eigen::half lhs, + Eigen::half rhs) { + return CompareFloatsBitwiseEqual(lhs, rhs); +} +template <> ::testing::AssertionResult CompareEqual(float lhs, float rhs) { return CompareFloatsBitwiseEqual(lhs, rhs); } @@ -376,6 +381,10 @@ class NearComparator { abs_expected_miscompare_sum_ = 0.0; max_rel_err_ = 0.0; max_abs_err_ = 0.0; + 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); @@ -445,8 +454,12 @@ class NearComparator { 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 @@ -468,7 +481,8 @@ class NearComparator { const float abs_diff = std::abs(actual - expected); const float rel_err = abs_diff / std::abs(expected); - const bool nan_mismatch = NanMismatch(expected, actual); + 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; @@ -482,11 +496,11 @@ class NearComparator { 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_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_linear_index_ = linear_index; } @@ -626,9 +640,11 @@ class NearComparator { }; 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 <> diff --git a/tensorflow/compiler/xla/tests/literal_test_util.h b/tensorflow/compiler/xla/tests/literal_test_util.h index 9b0724262d51ec7964a918bb8eb8716308662b96..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. diff --git a/tensorflow/compiler/xla/tests/literal_test_util_test.cc b/tensorflow/compiler/xla/tests/literal_test_util_test.cc index e477784557a3b9340cff644a3695485389d8cc22..3a421f8458268a14dcdd84889bcae4990c095ea4 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util_test.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util_test.cc @@ -97,5 +97,29 @@ TEST(LiteralTestUtilTest, ExpectNearFailurePlacesResultsInTemporaryDirectory) { } } +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_irgen_test_base.cc b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc index 99514baf23cafe61adc28a30dfdfe2691ab82d32..3023df47cda33f5d11abc921fd0355d48f761107 100644 --- a/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc +++ b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/tests/filecheck.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" namespace xla { @@ -49,11 +50,11 @@ void LLVMIRGenTestBase::CompileAndVerifyIr( std::unique_ptr hlo_module, const string& pattern, bool match_optimized_ir) { SetIrHook(match_optimized_ir); - ASSERT_TRUE(CompileToExecutable(std::move(hlo_module)).ok()); + TF_ASSERT_OK(CompileToExecutable(std::move(hlo_module)).status()); ResetIrHook(); StatusOr filecheck_result = RunFileCheck(ir_, pattern); - ASSERT_TRUE(filecheck_result.ok()); + TF_ASSERT_OK(filecheck_result.status()); EXPECT_TRUE(filecheck_result.ValueOrDie()); } diff --git a/tensorflow/compiler/xla/tests/map_test.cc b/tensorflow/compiler/xla/tests/map_test.cc index 2b0f7e6e80c48435ca55432a2afa3b6d69162625..0cd812fd1b4bc69c34b70d3ca0fd0aa6cf57fa4c 100644 --- a/tensorflow/compiler/xla/tests/map_test.cc +++ b/tensorflow/compiler/xla/tests/map_test.cc @@ -531,7 +531,7 @@ TEST_F(MapTest, MapOperantionWithBuildError) { ASSERT_TRUE(!computation_status.ok()); EXPECT_THAT( computation_status.status().ToString(), - ::testing::HasSubstr("error from: ErrorAdd: binary op BINOP_ADD with " + ::testing::HasSubstr("error from: ErrorAdd: Binary op BINOP_ADD with " "different element types: f32[] and u16[]")); } diff --git a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc index 6c86dd5b9ef673c9facffafa37e00a859ce82010..c42f71388baba73e08a361d817e41b03e03bf133 100644 --- a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc +++ b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc @@ -29,6 +29,8 @@ limitations under the License. #include "tensorflow/compiler/xla/test_helpers.h" #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/tests/test_utils.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/logging.h" @@ -38,258 +40,223 @@ limitations under the License. namespace xla { namespace { -class MatOpsSimpleTest : public ClientLibraryTestBase { - protected: - Computation BuildSum() { - // sum(x, y) = x + y - ComputationBuilder builder(client_, "sum"); - auto x_value = - builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x_value"); - auto y_value = - builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y_value"); - builder.Add(x_value, y_value); - auto computation_status = builder.Build(); - TF_CHECK_OK(computation_status.status()); - return computation_status.ConsumeValueOrDie(); - } - - void TestLinspaceMax(int64 rows, int64 cols) { - float from = -128.0, to = 256.0; - std::unique_ptr> alhs = - MakeLinspaceArray2D(from, to, rows, cols); - auto arhs = MakeUnique>(rows, cols, 1.0); - - ComputationBuilder builder( - client_, - tensorflow::strings::Printf("max_%lldx%lld_linspace", rows, cols)); - auto lhs = builder.ConstantR2FromArray2D(*alhs); - auto rhs = builder.ConstantR2FromArray2D(*arhs); - auto max = builder.Max(lhs, rhs); - - Array2D aexpected(rows, cols); - for (int row = 0; row < rows; ++row) { - for (int col = 0; col < cols; ++col) { - aexpected(row, col) = std::max((*alhs)(row, col), (*arhs)(row, col)); - } - } - - ComputeAndCompareR2(&builder, aexpected, {}, ErrorSpec(1e-6)); - } -}; - -TEST_F(MatOpsSimpleTest, ExpTwoByTwoValues) { - ComputationBuilder builder(client_, "exp_2x2"); - auto data = builder.ConstantR2({ - {1.0, 0.0}, // row 0 - {-1.0, 0.5}, // row 1 +#ifdef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +using TypesF16F32 = ::testing::Types; +#else +using TypesF16F32 = ::testing::Types; +#endif + +class MatOpsSimpleTest : public ClientLibraryTestBase {}; + +template +class MatOpsSimpleTest_F16F32 : public MatOpsSimpleTest {}; + +// TODO(bixia): This test for F16 failed on GPU 02-25-2018. +#ifdef XLA_TEST_BACKEND_GPU +TYPED_TEST_CASE(MatOpsSimpleTest_F16F32, ::testing::Types); +#else +TYPED_TEST_CASE(MatOpsSimpleTest_F16F32, TypesF16F32); +#endif + +XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, ExpTwoByTwoValues) { + using T = TypeParam; + ComputationBuilder builder(this->client_, "exp_2x2"); + auto data = builder.ConstantR2FromArray2D({ + {1.0f, 0.0f}, // row 0 + {-1.0f, 0.5f}, // row 1 }); builder.Exp(data); std::unique_ptr expected = - Literal::CreateR2({{2.71828, 1.00000}, // row 0 - {0.36788, 1.64872}}); // row 1 + Literal::CreateR2FromArray2D({{2.71828f, 1.00000f}, // row 0 + {0.36788f, 1.64872f}}); // row 1 - ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-5)); + this->template ComputeAndCompareLiteral(&builder, *expected, {}, + ErrorSpec(1e-5)); } -TEST_F(MatOpsSimpleTest, MapTwoByTwo) { +XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, MapTwoByTwo) { + using T = TypeParam; Computation add_half; { // add_half(x) = x + 0.5 - ComputationBuilder builder(client_, "add_half"); + ComputationBuilder builder(this->client_, "add_half"); auto x_value = - builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x_value"); - auto half = builder.ConstantR0(0.5); + builder.Parameter(0, ShapeUtil::MakeShapeWithType({}), "x_value"); + auto half = builder.ConstantR0(static_cast(0.5)); builder.Add(x_value, half); auto computation_status = builder.Build(); ASSERT_IS_OK(computation_status.status()); add_half = computation_status.ConsumeValueOrDie(); } - ComputationBuilder builder(client_, "map_2x2"); - auto data = builder.ConstantR2({ - {1.0, 0.0}, // row 0 - {-1.0, 0.5}, // row 1 + ComputationBuilder builder(this->client_, "map_2x2"); + auto data = builder.ConstantR2FromArray2D({ + {1.0f, 0.0f}, // row 0 + {-1.0f, 0.5f}, // row 1 }); auto map = builder.Map({data}, add_half, {0, 1}); std::unique_ptr expected = - Literal::CreateR2({{1.5, 0.5}, // row 0 - {-0.5, 1.0}}); // row 1 - ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-5)); + Literal::CreateR2FromArray2D({{1.5f, 0.5f}, // row 0 + {-0.5f, 1.0f}}); // row 1 + this->template ComputeAndCompareLiteral(&builder, *expected, {}, + ErrorSpec(1e-5)); } -TEST_F(MatOpsSimpleTest, MaxTwoByTwoValues) { - ComputationBuilder builder(client_, "max_2x2"); - auto lhs = builder.ConstantR2({ - {7.0, 2.0}, // row 0 - {3.0, -4.0}, // row 1 +XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, MaxTwoByTwoValues) { + using T = TypeParam; + ComputationBuilder builder(this->client_, "max_2x2"); + auto lhs = builder.ConstantR2FromArray2D({ + {7.0f, 2.0f}, // row 0 + {3.0f, -4.0f}, // row 1 }); - auto rhs = builder.ConstantR2({ - {5.0, 6.0}, // row 0 - {1.0, -8.0}, // row 1 + auto rhs = builder.ConstantR2FromArray2D({ + {5.0f, 6.0f}, // row 0 + {1.0f, -8.0f}, // row 1 }); auto max = builder.Max(lhs, rhs); std::unique_ptr expected = - Literal::CreateR2({{7.0, 6.0}, // row 0 - {3.0, -4.0}}); // row 1 - ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-6)); + Literal::CreateR2FromArray2D({{7.0f, 6.0f}, // row 0 + {3.0f, -4.0f}}); // row 1 + this->template ComputeAndCompareLiteral(&builder, *expected, {}, + ErrorSpec(1e-6)); } -TEST_F(MatOpsSimpleTest, Max1x1Linspace) { TestLinspaceMax(1, 1); } - -TEST_F(MatOpsSimpleTest, Max2x2Linspace) { TestLinspaceMax(2, 2); } - -TEST_F(MatOpsSimpleTest, Max3x3Linspace) { TestLinspaceMax(3, 3); } - -TEST_F(MatOpsSimpleTest, Max4x4Linspace) { TestLinspaceMax(4, 4); } - -TEST_F(MatOpsSimpleTest, Max6x6Linspace) { TestLinspaceMax(6, 6); } - -TEST_F(MatOpsSimpleTest, Max8x8Linspace) { TestLinspaceMax(8, 8); } - -TEST_F(MatOpsSimpleTest, Max12x12Linspace) { TestLinspaceMax(12, 12); } - -TEST_F(MatOpsSimpleTest, Max16x16Linspace) { TestLinspaceMax(16, 16); } +struct TestLinspaceMaxParam { + int64 rows; + int64 cols; +}; -TEST_F(MatOpsSimpleTest, Max32x8Linspace) { TestLinspaceMax(32, 8); } +class TestLinspaceMaxParametric + : public MatOpsSimpleTest, + public ::testing::WithParamInterface { + public: + template + void TestImpl() { + TestLinspaceMaxParam param = GetParam(); + int64 rows = param.rows; + int64 cols = param.cols; + float from = -128.0, to = 256.0; + std::unique_ptr> alhs = + MakeLinspaceArray2D(from, to, rows, cols); + auto arhs = MakeUnique>(rows, cols, static_cast(1.0f)); -TEST_F(MatOpsSimpleTest, Max64x8Linspace) { TestLinspaceMax(64, 8); } + ComputationBuilder builder( + client_, + tensorflow::strings::Printf("max_%lldx%lld_linspace", rows, cols)); + auto lhs = builder.ConstantR2FromArray2D(*alhs); + auto rhs = builder.ConstantR2FromArray2D(*arhs); + auto max = builder.Max(lhs, rhs); -class MatOpsDotAddTest - : public ClientLibraryTestBase, - public ::testing::WithParamInterface> {}; - -TEST_P(MatOpsDotAddTest, 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({{1.0, 2.0}, {3.0, 4.0}}); - Array2D rhs({{10.0, 11.0}, {12.0, 13.0}}); - - 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({{47, 52}, {71, 78}}); - } else { - expected = Array2D({{35, 39}, {81, 89}}); + Array2D expected(rows, cols); + for (int row = 0; row < rows; ++row) { + for (int col = 0; col < cols; ++col) { + expected(row, col) = std::max((*alhs)(row, col), (*arhs)(row, col)); + } } - } else { - result = builder.Add(result, rhs_arg); - if (transpose) { - expected = Array2D({{56, 61}, {80, 87}}); - } else { - expected = Array2D({{44, 48}, {90, 98}}); + ErrorSpec error_spec(1e-6); + if (std::is_same::value) { + error_spec = ErrorSpec(1e-6, 2e-4); } + ComputeAndCompareR2(&builder, expected, {}, error_spec); } +}; - ComputeAndCompareR2(&builder, expected, - {lhs_handle.get(), rhs_handle.get()}, - ErrorSpec(1e-6)); +string PrintTestLinspaceMaxParam( + const ::testing::TestParamInfo& test_param) { + const TestLinspaceMaxParam& param = test_param.param; + return tensorflow::strings::StrCat(param.rows, "r", param.cols, "c"); } -INSTANTIATE_TEST_CASE_P(MatOpsDotAddTestInstances, MatOpsDotAddTest, - ::testing::Combine(::testing::Bool(), ::testing::Bool(), - ::testing::Bool())); +#ifndef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +// TODO(bixia): This test failed on GPU 02-25-2018 +#ifdef XLA_TEST_BACKEND_CPU +XLA_TEST_P(TestLinspaceMaxParametric, TestF16) { TestImpl(); } +#endif +#endif +XLA_TEST_P(TestLinspaceMaxParametric, TestF32) { TestImpl(); } + +INSTANTIATE_TEST_CASE_P( + TestLinspaceMax, TestLinspaceMaxParametric, + ::testing::Values(TestLinspaceMaxParam{1, 1}, TestLinspaceMaxParam{2, 2}, + TestLinspaceMaxParam{3, 3}, TestLinspaceMaxParam{4, 4}, + TestLinspaceMaxParam{6, 6}, TestLinspaceMaxParam{8, 8}, + TestLinspaceMaxParam{12, 12}, + TestLinspaceMaxParam{16, 16}, TestLinspaceMaxParam{32, 8}, + TestLinspaceMaxParam{64, 8}), + PrintTestLinspaceMaxParam); -class MatOpsDotAddTest_bf16 +class MatOpsDotAddTest : 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); + public ::testing::WithParamInterface> { + public: + template + void TestImpl() { + bool row_major = std::get<0>(GetParam()); + bool add_lhs = std::get<1>(GetParam()); + bool transpose = std::get<2>(GetParam()); + Array2D lhs({{1.0f, 2.0f}, {3.0f, 4.0f}}); + Array2D rhs({{10.0f, 11.0f}, {12.0f, 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) { - expected = Array2D( - {{bfloat16(47), bfloat16(52)}, {bfloat16(71), bfloat16(78)}}); - } else { - expected = Array2D( - {{bfloat16(35), bfloat16(39)}, {bfloat16(81), bfloat16(89)}}); + lhs_mat_arg = builder.Transpose(lhs_mat_arg, {1, 0}); } - } else { - result = builder.Add(result, rhs_arg); - if (transpose) { - expected = Array2D( - {{bfloat16(56), bfloat16(61)}, {bfloat16(80), bfloat16(87)}}); + 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({{47.0f, 52.0f}, {71.0f, 78.0f}}); + } else { + expected = Array2D({{35.0f, 39.0f}, {81.0f, 89.0f}}); + } } else { - expected = Array2D( - {{bfloat16(44), bfloat16(48)}, {bfloat16(90), bfloat16(98)}}); + result = builder.Add(result, rhs_arg); + if (transpose) { + expected = Array2D({{56.0f, 61.0f}, {80.0f, 87.0f}}); + } else { + expected = Array2D({{44.0f, 48.0f}, {90.0f, 98.0f}}); + } } + + ComputeAndCompareR2(&builder, expected, + {lhs_handle.get(), rhs_handle.get()}, + ErrorSpec(1e-6)); } +}; - ComputeAndCompareR2(&builder, expected, - {lhs_handle.get(), rhs_handle.get()}, - ErrorSpec(1e-6)); -} +XLA_TEST_P(MatOpsDotAddTest, Dot_Add_2x2_2x2BF16) { TestImpl(); } +#ifndef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +XLA_TEST_P(MatOpsDotAddTest, Dot_Add_2x2_2x2F16) { TestImpl(); } +#endif +XLA_TEST_P(MatOpsDotAddTest, Dot_Add_2x2_2x2F32) { TestImpl(); } -INSTANTIATE_TEST_CASE_P(MatOpsDotAddTestInstances, MatOpsDotAddTest_bf16, +INSTANTIATE_TEST_CASE_P(MatOpsDotAddTestInstances, MatOpsDotAddTest, ::testing::Combine(::testing::Bool(), ::testing::Bool(), ::testing::Bool())); diff --git a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc index 6e6cb7ff1e2ac74dc54f14d8811c9a5d3662bbd2..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" @@ -176,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/reduce_test.cc b/tensorflow/compiler/xla/tests/reduce_test.cc index a766fa2db0e193c52171490981855843ab3ee158..50d7b5074d201d2292cf90224ef4cd37efdbb8d3 100644 --- a/tensorflow/compiler/xla/tests/reduce_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_test.cc @@ -494,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; diff --git a/tensorflow/compiler/xla/tests/reduce_window_test.cc b/tensorflow/compiler/xla/tests/reduce_window_test.cc index 7f3c72671d51bcdfba89b050371626d672c6d945..8b736f62f045bb913ac09add2f00d5edb0692d83 100644 --- a/tensorflow/compiler/xla/tests/reduce_window_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_window_test.cc @@ -272,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); @@ -282,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); @@ -800,6 +800,14 @@ const R4ReduceWindowTestData kR4ReduceWindowLargeTestValues[] = { /*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( @@ -952,45 +960,67 @@ struct R2ReduceWindowTestData { int64 base_bounds[2]; int64 window_bounds[2]; int64 strides[2]; + int64 pad_low[2]; + int64 pad_high[2]; int64 layout[2]; - Padding padding; Reducer reducer; } kR2TestCases[] = { {/*base_bounds=*/{4, 18}, /*window_bounds=*/{2, 4}, - /*strides=*/{1, 2}, /*layout=*/{0, 1}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{1, 2}, /*pad_low=*/{0, 1}, /*pad_high=*/{1, 1}, + /*layout=*/{0, 1}, + /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{2, 5}, /*window_bounds=*/{2, 4}, - /*strides=*/{1, 1}, /*layout=*/{0, 1}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{1, 1}, /*pad_low=*/{0, 1}, /*pad_high=*/{1, 2}, + /*layout=*/{0, 1}, + /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{1, 3}, /*window_bounds=*/{2, 3}, - /*strides=*/{1, 1}, /*layout=*/{0, 1}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{1, 1}, /*pad_low=*/{0, 1}, /*pad_high=*/{1, 1}, + /*layout=*/{0, 1}, + /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{3, 129}, /*window_bounds=*/{1, 100}, - /*strides=*/{2, 99}, /*layout=*/{0, 1}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{2, 99}, /*pad_low=*/{0, 0}, /*pad_high=*/{35, 35}, + /*layout=*/{0, 1}, + /*reducer=*/Reducer::kAdd}, +// TODO(b/74260408): This test last failed on GPU on 2018-03-08, likely due to a +// ptxas bug. +#ifndef XLA_TEST_BACKEND_GPU {/*base_bounds=*/{6, 152}, /*window_bounds=*/{2, 25}, - /*strides=*/{5, 4}, /*layout=*/{0, 1}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{5, 4}, /*pad_low=*/{0, 1}, /*pad_high=*/{10, 11}, + /*layout=*/{0, 1}, + /*reducer=*/Reducer::kAdd}, +#endif {/*base_bounds=*/{6, 4}, /*window_bounds=*/{4, 2}, - /*strides=*/{3, 3}, /*layout=*/{0, 1}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{3, 3}, /*pad_low=*/{0, 1}, /*pad_high=*/{0, 1}, + /*layout=*/{0, 1}, + /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{5, 147}, /*window_bounds=*/{1, 36}, - /*strides=*/{4, 5}, /*layout=*/{1, 0}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{4, 5}, /*pad_low=*/{0, 0}, /*pad_high=*/{17, 17}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{4, 153}, /*window_bounds=*/{2, 93}, - /*strides=*/{1, 1}, /*layout=*/{1, 0}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{1, 1}, /*pad_low=*/{0, 1}, /*pad_high=*/{46, 46}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, // Regression test for a bug that appeared in Inception (b/34784899). {/*base_bounds=*/{28, 28}, /*window_bounds=*/{3, 3}, - /*strides=*/{1, 1}, /*layout=*/{1, 0}, - /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + /*strides=*/{1, 1}, /*pad_low=*/{1, 1}, /*pad_high=*/{1, 1}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{4, 4}, /*window_bounds=*/{2, 2}, + /*strides=*/{1, 1}, /*pad_low=*/{0, 0}, /*pad_high=*/{0, 0}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, // Regression test for a bug that appeared in Inception (b/34784899). {/*base_bounds=*/{4, 32}, /*window_bounds=*/{2, 2}, - /*strides=*/{2, 2}, /*layout=*/{1, 0}, - /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, - {/*base_bounds=*/{4, 4}, /*window_bounds=*/{2, 2}, - /*strides=*/{1, 1}, /*layout=*/{1, 0}, - /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + /*strides=*/{2, 2}, /*pad_low=*/{0, 0}, /*pad_high=*/{0, 0}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, + // Regression test for b/73903312: bf16 lacks precision to store result of + // very large windows. Testing with a reasonable window larger than 128. + {/*base_bounds=*/{8, 130}, /*window_bounds=*/{1, 130}, + /*strides=*/{1, 1}, /*pad_low=*/{0, 130}, /*pad_high=*/{0, 0}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, }; string R2ReduceWindowTestDataToString( @@ -1000,10 +1030,11 @@ string R2ReduceWindowTestDataToString( 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], // + tensorflow::str_util::Join(param.window_bounds, "x"), // + "__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"), + "__layout_", param.layout[0], "_", param.layout[1], // "__reducer_", param.reducer == kAdd ? "add" : "max"); if (::testing::get<1>(data.param)) { str = tensorflow::strings::StrCat(str, "_bfloat16"); @@ -1031,17 +1062,29 @@ class R2ReduceWindowTest : public ReduceWindowTestBase, ComputationDataHandle parameter; auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", &b, ¶meter); + std::vector> padding(2); + for (int i = 0; i < 2; ++i) { + padding[i] = {param.pad_low[i], param.pad_high[i]}; + } + auto computation = param.reducer == kAdd + ? CreateScalarAddComputation(FloatType(), &b) + : CreateScalarMaxComputation(FloatType(), &b); 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); + b.ReduceWindowWithGeneralPadding( + /*operand=*/parameter, + /*init_value=*/init_value, + /*computation=*/computation, + /*window_dimensions=*/param.window_bounds, + /*window_strides=*/param.strides, /*padding=*/padding); - auto expected = ReferenceUtil::ReduceWindow2DAdd( - /*operand=*/input, /*init=*/kInitValue, /*window=*/param.window_bounds, - /*stride=*/param.strides, /*padding=*/param.padding); + auto reduce_func = param.reducer == kAdd + ? +[](float a, float b) { return a + b; } + : +[](float a, float b) { return std::max(a, b); }; + auto expected = ReferenceUtil::ReduceWindow2DGeneric( + /*operand=*/input, /*init=*/kInitValue, /*reduce_func=*/reduce_func, + /*window=*/param.window_bounds, + /*stride=*/param.strides, /*padding=*/padding); ComputeAndCompareLiteral(&b, *Literal::CreateFromArray(*expected), {input_arg.get()}, DefaultErrorSpec()); @@ -1066,8 +1109,9 @@ XLA_TEST_P(R2ReduceWindowFailingCpuGpuBf16Test, const R2ReduceWindowTestData kR2FailingValuesCpuGpuBf16Test[] = { {/*base_bounds=*/{8, 128}, /*window_bounds=*/{8, 128}, - /*strides=*/{1, 1}, /*layout=*/{1, 0}, - /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + /*strides=*/{1, 1}, /*pad_low=*/{0, 0}, /*pad_high=*/{0, 0}, + /*layout=*/{1, 0}, + /*reducer=*/Reducer::kAdd}, }; INSTANTIATE_TEST_CASE_P( diff --git a/tensorflow/compiler/xla/tests/scalar_computations_test.cc b/tensorflow/compiler/xla/tests/scalar_computations_test.cc index debf2d2d317fe64ca1ef86cb1f2978e76af1b55d..0c88bef69dfc522fef52422b0bd3a825fa173d44 100644 --- a/tensorflow/compiler/xla/tests/scalar_computations_test.cc +++ b/tensorflow/compiler/xla/tests/scalar_computations_test.cc @@ -163,7 +163,7 @@ XLA_TEST_F(ScalarComputationsTest, CastS64ToF32) { auto a = builder.Parameter(0, ShapeUtil::MakeShape(S64, {}), "a"); builder.ConvertElementType(a, F32); - int64 value = 3LL << 32; + int64 value = 3LL << 35; std::unique_ptr a_literal = Literal::CreateR0(value); std::unique_ptr a_data = client_->TransferToServer(*a_literal).ConsumeValueOrDie(); @@ -737,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. @@ -746,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. @@ -755,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. @@ -806,6 +860,12 @@ XLA_TEST_F(ScalarComputationsTest, MinF32Below) { TestMinMax(-100.1f, 3.1f, -100.1f, &ComputationBuilder::Min); } +XLA_TEST_F(ScalarComputationsTest, MinPropagatesNan) { + SetFastMathDisabled(true); + TestMinMax(NAN, 3.1f, NAN, &ComputationBuilder::Min); + TestMinMax(-3.1f, NAN, NAN, &ComputationBuilder::Min); +} + XLA_TEST_F(ScalarComputationsTest, MaxF32Above) { TestMinMax(10.1f, 3.1f, 10.1f, &ComputationBuilder::Max); } @@ -814,6 +874,12 @@ XLA_TEST_F(ScalarComputationsTest, MaxF32Below) { TestMinMax(-100.1f, 3.1f, 3.1f, &ComputationBuilder::Max); } +XLA_TEST_F(ScalarComputationsTest, MaxPropagatesNan) { + SetFastMathDisabled(true); + TestMinMax(NAN, 3.1f, NAN, &ComputationBuilder::Max); + TestMinMax(-3.1f, NAN, NAN, &ComputationBuilder::Max); +} + XLA_TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionF32) { // Compute the expression (1 * (3 - 1) * (7 + 0) - 4) / 20. ComputationBuilder b(client_, TestName()); @@ -852,5 +918,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 9ee94b8571e5fc8789b60501462986967ce909a0..d268fdcacebcb162bf61bc7dd4b208f4db6c4a5f 100644 --- a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc +++ b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc @@ -252,6 +252,21 @@ XLA_TEST_F(SelectAndScatterTest, R2S32) { ComputeAndCompareR2(&builder_, expected, {}); } +// Test for tie breaking rule in ge_f32_. When a tie is present, the operand +// that has the lower lexicographical order (smaller index) should be chosen. +XLA_TEST_F(SelectAndScatterTest, R2F32Tie) { + const auto operand = builder_.ConstantR2( + {{0.f, 0.f, 0.f}, {0.f, 0.f, 0.f}, {0.f, 0.f, 0.f}}); + const auto source = builder_.ConstantR2( + {{1.0f, 2.0f, 3.0f}, {4.f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}); + Array2D expected( + {{12.f, 9.f, 0.f}, {15.f, 9.f, 0.f}, {0.f, 0.f, 0.f}}); + builder_.SelectAndScatter(operand, ge_f32_, /*window_dimensions=*/{3, 3}, + /*window_strides=*/{1, 1}, Padding::kSame, source, + builder_.ConstantR0(0.0f), add_f32_); + ComputeAndCompareR2(&builder_, expected, {}, ErrorSpec(1e-7)); +} + // Similar to SelectAndScatterTest.R2S32 but the input is transposed. XLA_TEST_F(SelectAndScatterTest, ReshapeR2S32) { const auto operand = builder_.ConstantR2( diff --git a/tensorflow/compiler/xla/tests/slice_test.cc b/tensorflow/compiler/xla/tests/slice_test.cc index ac163df127e0087c02777fa3d5ce7970c51b97b9..fe36df160daacc4fdfbdb0b75f8304f91e1a4245 100644 --- a/tensorflow/compiler/xla/tests/slice_test.cc +++ b/tensorflow/compiler/xla/tests/slice_test.cc @@ -237,6 +237,12 @@ INSTANTIATE_TEST_CASE_P( SliceR1TestInstantiation, SliceR1Test, ::testing::Values( +// TODO(b/69425338): This uses too much memory on GPU. +#ifndef XLA_TEST_BACKEND_GPU + R1Spec{16 * 1024 * 1024, 4 * 1024 * 1024, 12 * 1024 * 1024, 1}, + R1Spec{16 * 1024 * 1024, 4 * 1024 * 1024 + 1, 12 * 1024 * 1024 - 1, 1}, + R1Spec{16 * 1024 * 1024, 4 * 1024 * 1024 - 1, 12 * 1024 * 1024 + 1, 1}, +#endif R1Spec{10, 0, 0, 1}, R1Spec{10, 7, 7, 1}, R1Spec{10, 0, 5, 1}, @@ -267,13 +273,15 @@ INSTANTIATE_TEST_CASE_P( R1Spec{64 * 1024, 1024 + 1, 63 * 1024 - 1, 1}, R1Spec{64 * 1024, 32 * 1024, 33 * 1024, 1}, R1Spec{64 * 1024, 32 * 1024 + 1, 33 * 1024 - 1, 1}, - R1Spec{64 * 1024, 32 * 1024 - 17, 36 * 1024 - 18, 1}, -// TODO(b/69425338): This uses too much memory on GPU. -#ifndef XLA_TEST_BACKEND_GPU - R1Spec{16 * 1024 * 1024, 4 * 1024 * 1024, 12 * 1024 * 1024, 1}, - R1Spec{16 * 1024 * 1024, 4 * 1024 * 1024 + 1, 12 * 1024 * 1024 - 1, 1}, - R1Spec{16 * 1024 * 1024, 4 * 1024 * 1024 - 1, 12 * 1024 * 1024 + 1, 1}, -#endif + R1Spec{64 * 1024, 32 * 1024 - 17, 36 * 1024 - 18, 1} + ), + SliceR1TestDataToString +); + +INSTANTIATE_TEST_CASE_P( + SliceStridedR1TestInstantiation, + SliceR1Test, + ::testing::Values( R1Spec{10, 2, 4, 2}, R1Spec{10, 0, 10, 2}, R1Spec{10, 0, 10, 3}, @@ -285,8 +293,24 @@ INSTANTIATE_TEST_CASE_P( R1Spec{2047, 1024 - 24, 1024 + 160, 31}, R1Spec{2047, 1, 2046, 3 * 128}, R1Spec{4096, 1024 + 3, 4095, 500}, - R1Spec{8192, 0, 8192, 1024 * 3 + 400} - ), + R1Spec{8192, 0, 8192, 1024 * 3 + 400}, + R1Spec{1024 * 1024, 0, 1024 * 1024, 2}, + R1Spec{1024 * 1024, 0, 1024 * 1024, 8}, + R1Spec{1024 * 1024, 0, 1024 * 1024, 7}, + R1Spec{1024 * 1024, 0, 1024 * 1024, 125}, + R1Spec{1024 * 1024, 3, 1024 - 9, 2}, + R1Spec{1024 * 1024, 3, 1024 - 9, 8}, + R1Spec{1024 * 1024, 3, 1024 - 9, 7}, + R1Spec{1024 * 1024, 3, 1024 - 9, 125}, + R1Spec{1024 * 1024, 3, 1024 * 512 - 9, 2}, + R1Spec{1024 * 1024, 3, 1024 * 512 - 9, 8}, + R1Spec{1024 * 1024, 3, 1024 * 512 - 9, 7}, + R1Spec{1024 * 1024, 3, 1024 * 512 - 9, 125}, + R1Spec{1024 * 1024 + 71, 3, 1024 * 512 - 9, 2}, + R1Spec{1024 * 1024 + 71, 3, 1024 * 512 - 9, 8}, + R1Spec{1024 * 1024 + 71, 3, 1024 * 512 - 9, 7}, + R1Spec{1024 * 1024 + 71, 3, 1024 * 512 - 9, 125} + ), SliceR1TestDataToString ); // clang-format on diff --git a/tensorflow/compiler/xla/tests/test_macros.h b/tensorflow/compiler/xla/tests/test_macros.h index cc4eaf62f50d1fa622c705fab810fe1e1b0fbf08..e2d406f66d94f8ec76faa5b7d2d2e84dcaf6db57 100644 --- a/tensorflow/compiler/xla/tests/test_macros.h +++ b/tensorflow/compiler/xla/tests/test_macros.h @@ -161,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 b060fb13b1451aab30cfca73bea0a4a598a9fa3a..0bc7df2a65b44a76f877b6513e6bf93b99fbc1a3 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -287,7 +287,7 @@ StatusOr> MakeFakeLiteral(const Shape& shape) { StatusOr>> MakeFakeArguments( HloModule* const module) { - TF_ASSIGN_OR_RETURN(auto dataflow, HloDataflowAnalysis::Run(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()); diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc index a8bca70d85ddf168bc441231d6f43bead019b10a..2029312f94a14bc81706368b9ecfc2727fd9fe4c 100644 --- a/tensorflow/compiler/xla/tests/tuple_test.cc +++ b/tensorflow/compiler/xla/tests/tuple_test.cc @@ -194,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/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc index 52157b837c383205f77a030ef98b2fd03a41aff5..33d457c70bac84c2da10e3cf9302c2c952cf1bc2 100644 --- a/tensorflow/compiler/xla/tests/while_test.cc +++ b/tensorflow/compiler/xla/tests/while_test.cc @@ -910,7 +910,7 @@ XLA_TEST_F(WhileTest, WhileWithDynamicUpdateSlice) { // Per backend the values generated can be different as the different backends // use different random number generators. // TODO(b/32240857): Extend test to verify outputs. -TEST_F(WhileTest, WhileWithPrngScalarResult) { +TEST_F(WhileTest, DISABLED_ON_INTERPRETER(WhileWithPrngScalarResult)) { auto v6s32 = ShapeUtil::MakeShape(S32, {6}); // Create a computation for the condition: repeat for count iterations. @@ -1166,7 +1166,7 @@ XLA_TEST_F(WhileTest, NestedWhileWithScalarResult) { // while (f(result).get<0>()) { // result = result + 1; // } -TEST_F(WhileTest, WhileWithCallInsideCondition) { +TEST_F(WhileTest, DISABLED_ON_INTERPRETER(WhileWithCallInsideCondition)) { auto result_shape = ShapeUtil::MakeShape(S32, {}); // Create a computation for the condition: repeat for 5 iterations. diff --git a/tensorflow/compiler/xla/tools/BUILD b/tensorflow/compiler/xla/tools/BUILD index 091fa0c3ec807a66449eca0bfbb141285b8eb532..2e55f609d17bf42e410f97c51c7b9c6c0e85576d 100644 --- a/tensorflow/compiler/xla/tools/BUILD +++ b/tensorflow/compiler/xla/tools/BUILD @@ -75,6 +75,7 @@ cc_library( name = "replay_computation_library", srcs = ["replay_computation.cc"], deps = [ + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc index d9c4d094b8ea3b4de9164f35d670fc686c921a27..e60a5a4919f2207939821e787c3c59a08ff3ba4e 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc @@ -220,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_; @@ -340,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. @@ -990,6 +994,20 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, shape, operands, *custom_call_target)); break; } + case HloOpcode::kHostCompute: { + optional channel_name; + optional cost_estimate_ns; + attrs["channel_name"] = {/*required=*/true, AttrTy::kString, + &channel_name}; + attrs["cost_estimate_ns"] = {/*required=*/true, AttrTy::kInt64, + &cost_estimate_ns}; + if (!ParseOperands(&operands) || !ParseAttributes(attrs)) { + return false; + } + instruction = builder->AddInstruction(HloInstruction::CreateHostCompute( + shape, operands, *channel_name, *cost_estimate_ns)); + break; + } case HloOpcode::kDot: { optional> lhs_contracting_dims; attrs["lhs_contracting_dims"] = { @@ -1031,6 +1049,40 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, HloInstruction::CreateDot(shape, operands[0], operands[1], dnum)); break; } + case HloOpcode::kGather: { + optional> output_window_dims; + attrs["output_window_dims"] = { + /*required=*/true, AttrTy::kBracedInt64List, &output_window_dims}; + optional> elided_window_dims; + attrs["elided_window_dims"] = { + /*required=*/true, AttrTy::kBracedInt64List, &elided_window_dims}; + optional> gather_dims_to_operand_dims; + attrs["gather_dims_to_operand_dims"] = {/*required=*/true, + AttrTy::kBracedInt64List, + &gather_dims_to_operand_dims}; + optional index_vector_dim; + attrs["index_vector_dim"] = {/*required=*/true, AttrTy::kInt64, + &index_vector_dim}; + optional> window_bounds; + attrs["window_bounds"] = {/*required=*/true, AttrTy::kBracedInt64List, + &window_bounds}; + + if (!ParseOperands(&operands, /*expected_size=*/2) || + !ParseAttributes(attrs)) { + return false; + } + + GatherDimensionNumbers dim_numbers = HloInstruction::MakeGatherDimNumbers( + /*output_window_dims=*/*output_window_dims, + /*elided_window_dims=*/*elided_window_dims, + /*gather_dims_to_operand_dims=*/*gather_dims_to_operand_dims, + /*index_vector_dim=*/*index_vector_dim); + + instruction = builder->AddInstruction(HloInstruction::CreateGather( + shape, /*operand=*/operands[0], /*gather_indices=*/operands[1], + dim_numbers, *window_bounds)); + break; + } case HloOpcode::kTrace: return TokenError(StrCat("parsing not yet implemented for op: ", HloOpcodeString(opcode))); @@ -1229,13 +1281,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, @@ -1705,12 +1757,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"); @@ -1957,10 +2009,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; } @@ -2576,18 +2630,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 dd76d8d0fee7cdfa22829fe92ff889e44157216e..863081d654390440aa6506bab4576b3cc5c1cbd1 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc @@ -716,6 +716,18 @@ ENTRY %sparse_f32_r1 () -> f32[9] { ROOT %foo = f32[9]sparse{10} constant(f32[9]{1: 2, 3: 4, 5: 6}) } +)" +}, +{ +"gather", +R"(HloModule StringifyGather + +ENTRY %Gather (input_tensor: f32[50,49,48,47,46], gather_indices: s64[10,9,8,7,5]) -> f32[10,9,8,7,30,29,28,27,26] { + %input_tensor = f32[50,49,48,47,46]{4,3,2,1,0} parameter(0) + %gather_indices = s64[10,9,8,7,5]{4,3,2,1,0} parameter(1) + ROOT %gather = f32[10,9,8,7,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} gather(f32[50,49,48,47,46]{4,3,2,1,0} %input_tensor, s64[10,9,8,7,5]{4,3,2,1,0} %gather_indices), output_window_dims={4,5,6,7,8}, elided_window_dims={}, gather_dims_to_operand_dims={0,1,2,3,4}, index_vector_dim=4, window_bounds={30,29,28,27,26} +} + )" }, }); @@ -860,6 +872,18 @@ ENTRY dot { ROOT dot = f32[2,3]{1,0} dot(a, b), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_contracting_dims={0} } +)" +}, +{ +"gather", +R"(HloModule gather + +ENTRY Gather { + input_tensor = f32[50,49,48,47,46]{4,3,2,1,0} parameter(0) + gather_indices = s64[10,9,8,7,5]{4,3,2,1,0} parameter(1) + ROOT gather = f32[10,9,8,7,30,29,28,27,26]{8,7,6,5,4,3,2,1,0} gather(input_tensor, gather_indices), output_window_dims={4,5,6,7,8}, elided_window_dims={}, gather_dims_to_operand_dims={0,1,2,3,4}, index_vector_dim=4, window_bounds={30,29,28,27,26} +} + )" }, }); @@ -1275,6 +1299,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 eda5effbb92db92c9317a956497a00c0ec15c27c..62a353ad09af009e4abf47664a5c5f7bd70a049e 100644 --- a/tensorflow/compiler/xla/tools/replay_computation.cc +++ b/tensorflow/compiler/xla/tools/replay_computation.cc @@ -40,6 +40,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/lib/testing.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/session.pb.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -66,6 +67,7 @@ struct Options { bool use_fake_data = false; bool print_result = true; int num_runs = 1; + bool xla_hlo_profile_last_run = false; }; // Invokes the given computation passing arbitrary data for every (unbound) @@ -122,16 +124,21 @@ StatusOr> ReplayComputation( std::unique_ptr result; for (int i = 0; i < opts.num_runs; ++i) { ExecutionProfile profile; + ExecutionOptions execution_options = CreateDefaultExecutionOptions(); + if (opts.xla_hlo_profile_last_run && i == opts.num_runs - 1) { + execution_options.mutable_debug_options()->set_xla_hlo_profile(true); + } + if (opts.print_result) { - TF_ASSIGN_OR_RETURN(result, client->ExecuteAndTransfer( - computation, execute_arguments, - /*execution_options=*/nullptr, &profile)); + TF_ASSIGN_OR_RETURN( + result, client->ExecuteAndTransfer(computation, execute_arguments, + &execution_options, &profile)); } else { // If we're not printing the result, execute the computation but don't // bother retrieving the result. This can be a significant speedup. TF_RETURN_IF_ERROR(client ->Execute(computation, execute_arguments, - /*execution_options=*/nullptr, &profile) + &execution_options, &profile) .status()); } LOG(INFO) << "Execution took " @@ -191,6 +198,9 @@ int main(int argc, char** argv) { "Number of times to run each computation"), tensorflow::Flag("fake_infeed_shape", &opts.fake_infeed_shape, "Shape of fake data to construct for (infinite) infeed"), + tensorflow::Flag( + "xla_hlo_profile_last_run", &opts.xla_hlo_profile_last_run, + "Pass --xla_hlo_profile the last time we run the computation."), }; xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); bool parse_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); diff --git a/tensorflow/compiler/xla/util.cc b/tensorflow/compiler/xla/util.cc index 1f0c626bbb2d64ef4e67c9ec51485ae96ae73d04..dc4f7a1cb436183f5acfa360fb092795258b6a75 100644 --- a/tensorflow/compiler/xla/util.cc +++ b/tensorflow/compiler/xla/util.cc @@ -15,7 +15,6 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" -#include #include #include @@ -292,7 +291,8 @@ void LogLines(int sev, tensorflow::StringPiece text, const char* fname, } int64 Product(tensorflow::gtl::ArraySlice xs) { - return std::accumulate(xs.begin(), xs.end(), 1, std::multiplies()); + return std::accumulate(xs.begin(), xs.end(), static_cast(1), + std::multiplies()); } std::vector> CommonFactors( diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index 08df5b12b3a53a138f56705531baa3333b23c5d8..ff99d3728d1c3b58fc94d3eb3de78be23407edc9 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include "tensorflow/compiler/xla/status.h" @@ -427,32 +428,97 @@ std::vector> CommonFactors( 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); +bool c_all_of(const Container& container, Predicate&& predicate) { + return std::all_of(std::begin(container), std::end(container), + std::forward(predicate)); +} + +template +bool c_any_of(const Container& container, Predicate&& predicate) { + return std::any_of(std::begin(container), std::end(container), + std::forward(predicate)); } template -OutputIterator c_transform(InputContainer input_container, +OutputIterator c_transform(const InputContainer& input_container, OutputIterator output_iterator, - UnaryOperation unary_op) { + UnaryOperation&& unary_op) { return std::transform(std::begin(input_container), std::end(input_container), - output_iterator, unary_op); + output_iterator, + std::forward(unary_op)); } template -OutputIterator c_copy_if(InputContainer input_container, +OutputIterator c_copy_if(const InputContainer& input_container, OutputIterator output_iterator, - UnaryPredicate predicate) { + UnaryPredicate&& predicate) { return std::copy_if(std::begin(input_container), std::end(input_container), - output_iterator, predicate); + output_iterator, std::forward(predicate)); +} + +template +OutputIterator c_copy(const InputContainer& input_container, + OutputIterator output_iterator) { + return std::copy(std::begin(input_container), std::end(input_container), + output_iterator); +} + +template +void c_sort(InputContainer& input_container) { + std::sort(std::begin(input_container), std::end(input_container)); } template -void c_sort(InputContainer& input_container, Comparator comparator) { - std::sort(input_container.begin(), input_container.end(), comparator); +void c_sort(InputContainer& input_container, Comparator&& comparator) { + std::sort(std::begin(input_container), std::end(input_container), + std::forward(comparator)); +} + +template +bool c_binary_search(const Sequence& sequence, T&& value) { + return std::binary_search(std::begin(sequence), std::end(sequence), + std::forward(value)); +} + +template +bool c_is_sorted(const C& c) { + return std::is_sorted(std::begin(c), std::end(c)); } +template +auto c_adjacent_find(const C& c) -> decltype(std::begin(c)) { + return std::adjacent_find(std::begin(c), std::end(c)); +} + +template +auto c_find_if(const C& c, Pred&& pred) -> decltype(std::begin(c)) { + return std::find_if(std::begin(c), std::end(c), std::forward(pred)); +} + +template +auto c_find(const C& c, Value&& value) -> decltype(std::begin(c)) { + return std::find(std::begin(c), std::end(c), std::forward(value)); +} + +template +void c_reverse(Sequence& sequence) { + std::reverse(std::begin(sequence), std::end(sequence)); +} + +template +typename std::decay::type c_accumulate(const Sequence& sequence, T&& init, + BinaryOp&& binary_op) { + return std::accumulate(std::begin(sequence), std::end(sequence), + std::forward(init), + std::forward(binary_op)); +} + +template +int64 FindIndex(const C& c, Value&& value) { + auto it = c_find(c, std::forward(value)); + return std::distance(c.begin(), it); +} } // 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 55f42ed3a454baa3f8b6adf60a78582488733e9b..93284b80f9e1f82c4b18dc7388754d5c01a7740c 100644 --- a/tensorflow/compiler/xla/window_util.cc +++ b/tensorflow/compiler/xla/window_util.cc @@ -32,6 +32,8 @@ Window MakeWindow(tensorflow::gtl::ArraySlice 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; } diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto index 56162ab44e2e0e3e4478fe631888f243332dc1d8..edf1b07af82b5d43fe67c6efdabdb0a9b4b1edea 100644 --- a/tensorflow/compiler/xla/xla.proto +++ b/tensorflow/compiler/xla/xla.proto @@ -16,6 +16,7 @@ limitations under the License. syntax = "proto3"; import "tensorflow/compiler/xla/xla_data.proto"; +import "tensorflow/compiler/xla/service/hlo.proto"; import "tensorflow/compiler/xla/service/session.proto"; package xla; @@ -342,6 +343,14 @@ message ExecuteRequest { ExecutionOptions execution_options = 5; } +message ExecuteGraphRequest { + HloModuleProto computation = 1; + repeated GlobalDataHandle arguments = 2; + + // Options that affect how XLA compiles and runs code to service this request. + ExecutionOptions execution_options = 3; +} + message ExecuteParallelRequest { repeated ExecuteRequest requests = 1; } diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto index 3aea0217539b89b5d60ecfaf2605eee4b69af728..1f16e6d25178fd9c10a30b0c500e090ee2e08117 100644 --- a/tensorflow/compiler/xla/xla_data.proto +++ b/tensorflow/compiler/xla/xla_data.proto @@ -393,6 +393,37 @@ message Window { repeated WindowDimension dimensions = 1; } +// Describes the dimension numbers for a gather operation. +// +// See https://www.tensorflow.org/performance/xla/operation_semantics#gather for +// more details. +message GatherDimensionNumbers { + // "Window indices" is a term for a set of indices that index into the + // interior of a dynamic-slice from the input tensor, the starting indices for + // which were computed from output_gather_dims (see the operation semantic for + // how this is defined) and the gather_indices tensor. + // + // The window indices for a specific output index Out is computed as: + // + // i = 0 + // for (k : [0, input_tensor_shape.rank)) + // window_indices[k] = + // if k in elided_window_dims + // then 0 + // else Out[output_window_dims[i++]] + repeated int64 output_window_dims = 1; + repeated int64 elided_window_dims = 2; + + // This is interpreted as a map from i to gather_dims_to_operand_dims[i]. It + // transforms the gather index looked up from the gather_indices tensor into + // the starting index in the input space. + repeated int64 gather_dims_to_operand_dims = 3; + + // The dimension in the gather_indices input that contains the starting + // indices. + int64 index_vector_dim = 4; +} + // Operation requests that are all collected as a tagged union with a oneof // field in OpRequest. @@ -519,6 +550,20 @@ message CustomCallRequest { Shape shape = 4; } +message HostComputeRequest { + // Operand to the HostCompute. Supports tuple. + repeated ComputationDataHandle operands = 1; + + // Name used to identify HostSend/Recv channels. + string channel_name = 2; + + // Cost estimate in nanoseconds. + int64 cost_estimate_ns = 3; + + // The shape of any data returned by host. + Shape shape = 4; +} + message DotDimensionNumbers { // The dimension numbers that represent the 'lhs' contracting dimensions. repeated int64 lhs_contracting_dimensions = 1; @@ -880,6 +925,13 @@ message RecvRequest { ChannelHandle channel_handle = 2; } +message GatherRequest { + ComputationDataHandle input = 1; + ComputationDataHandle gather_indices = 2; + GatherDimensionNumbers dimension_numbers = 3; + repeated int64 window_bounds = 4; +} + message OpSharding { enum Type { // This sharding is replicated across all devices (implies maximal, @@ -957,7 +1009,9 @@ message OpRequest { FftRequest fft_request = 41; ConvertRequest bitcast_convert_request = 42; ConditionalRequest conditional_request = 44; - // Next: 45 + HostComputeRequest host_compute_request = 45; + GatherRequest gather_request = 46; + // Next: 47 } } diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD index 8ec4be6dadbc755bac02652416eca48f339f01aa..c2663c5e83352a1088166dc7581a0346c7b104a4 100644 --- a/tensorflow/contrib/BUILD +++ b/tensorflow/contrib/BUILD @@ -7,6 +7,8 @@ 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") +load("//tensorflow:tensorflow.bzl", "if_not_windows") py_library( name = "contrib_py", @@ -24,6 +26,7 @@ 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", @@ -36,7 +39,7 @@ py_library( "//tensorflow/contrib/eager/python:tfe", "//tensorflow/contrib/estimator:estimator_py", "//tensorflow/contrib/factorization:factorization_py", - "//tensorflow/contrib/ffmpeg:ffmpeg_ops_py", + "//tensorflow/contrib/feature_column:feature_column_py", "//tensorflow/contrib/framework:framework_py", "//tensorflow/contrib/fused_conv:fused_conv_py", "//tensorflow/contrib/gan", @@ -48,7 +51,6 @@ 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", @@ -60,7 +62,6 @@ py_library( "//tensorflow/contrib/linalg:linalg_py", "//tensorflow/contrib/linear_optimizer:sdca_estimator_py", "//tensorflow/contrib/linear_optimizer:sdca_ops_py", - "//tensorflow/contrib/lite/python:lite", "//tensorflow/contrib/lookup:lookup_py", "//tensorflow/contrib/losses:losses_py", "//tensorflow/contrib/losses:metric_learning_py", @@ -69,7 +70,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", @@ -106,7 +106,13 @@ py_library( "//tensorflow/contrib/training:training_py", "//tensorflow/contrib/util:util_py", "//tensorflow/python:util", - ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]), + ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + if_tensorrt([ + "//tensorflow/contrib/tensorrt:init_py", + ]) + if_not_windows([ + "//tensorflow/contrib/ffmpeg:ffmpeg_ops_py", # unix dependency, need to fix code + "//tensorflow/contrib/lite/python:lite", # unix dependency, need to fix code + "//tensorflow/contrib/kafka", # has some linking issue on opensssl. + ]), ) cc_library( @@ -116,6 +122,8 @@ cc_library( "//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/kafka:dataset_kernels", "//tensorflow/contrib/factorization/kernels:all_kernels", "//tensorflow/contrib/input_pipeline:input_pipeline_ops_kernels", "//tensorflow/contrib/layers:sparse_feature_cross_op_kernel", @@ -138,10 +146,11 @@ cc_library( "//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/kafka:dataset_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 8f6a3cb1ca4544cae6f42fd1727d509af9fc0233..669d611b01b585d91ab48921b7ba17703dd6bc98 100644 --- a/tensorflow/contrib/__init__.py +++ b/tensorflow/contrib/__init__.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os + # Add projects here, they will show up under tf.contrib. from tensorflow.contrib import batching from tensorflow.contrib import bayesflow @@ -33,6 +35,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 @@ -82,8 +85,8 @@ from tensorflow.contrib import tpu 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 +if os.name != 'nt': + from tensorflow.contrib.lite.python import lite 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 @@ -92,6 +95,7 @@ from tensorflow.contrib.summary import summary from tensorflow.python.util.lazy_loader import LazyLoader ffmpeg = LazyLoader("ffmpeg", globals(), "tensorflow.contrib.ffmpeg") +del os 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 28f60b34996945d573facc665c01d0bc10cf5cd1..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 @@ -752,13 +752,13 @@ def _build_nccl_hybrid(input_tensors, red_op, upper_level_f): 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/cmake/CMakeLists.txt b/tensorflow/contrib/android/cmake/CMakeLists.txt index a115d1610e2334a6626f29674f3dd195e3a3c648..ecf1a103d2981f409a4598d762fb26100217f779 100644 --- a/tensorflow/contrib/android/cmake/CMakeLists.txt +++ b/tensorflow/contrib/android/cmake/CMakeLists.txt @@ -75,7 +75,6 @@ target_link_libraries(tensorflow_inference include_directories( ${PREBUILT_DIR}/proto ${PREBUILT_DIR}/protobuf/include - ${PREBUILT_DIR}/nsync/public ${TENSORFLOW_ROOT_DIR}/tensorflow/contrib/makefile/downloads/eigen ${TENSORFLOW_ROOT_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/..) 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 e51e3f747b57cf1c9fd79ee5cc9fdb8acee349c9..abddadac5bcace9b1f992b69bdcc69c24b29cd13 100644 --- a/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java +++ b/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java @@ -197,9 +197,7 @@ public class TensorFlowInferenceInterface { run(outputNames, enableStats, new String[] {}); } - /** - * An overloaded version of runInference that allows supplying targetNodeNames as well - */ + /** 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(); @@ -211,7 +209,7 @@ public class TensorFlowInferenceInterface { runner.fetch(tid.name, tid.outputIndex); } - // Add targets. + // Add targets. for (String t : targetNodeNames) { runner.addTarget(t); } diff --git a/tensorflow/contrib/batching/python/ops/batch_ops.py b/tensorflow/contrib/batching/python/ops/batch_ops.py index 4e0b3f9af989c414ad88c510c1bfd180dbadd5ea..921d6917a4e478c3e60771fdc3ae99febc33d2e3 100644 --- a/tensorflow/contrib/batching/python/ops/batch_ops.py +++ b/tensorflow/contrib/batching/python/ops/batch_ops.py @@ -53,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 @@ -94,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. @@ -111,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/bayesflow/BUILD b/tensorflow/contrib/bayesflow/BUILD index 11c3c037c4e8b4ba41eae60d28d6aac49f1488f2..c6feec68e0104ff33451bbb6fa7de51d13e0a43c 100644 --- a/tensorflow/contrib/bayesflow/BUILD +++ b/tensorflow/contrib/bayesflow/BUILD @@ -39,7 +39,7 @@ py_library( cuda_py_test( name = "metropolis_hastings_test", - size = "medium", + size = "large", srcs = ["python/kernel_tests/metropolis_hastings_test.py"], additional_deps = [ ":bayesflow_py", @@ -56,87 +56,6 @@ cuda_py_test( ], ) -cuda_py_test( - name = "csiszar_divergence_test", - size = "medium", - srcs = ["python/kernel_tests/csiszar_divergence_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/contrib/distributions:distributions_py", - "//tensorflow/python/ops/distributions", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:gradients", - "//tensorflow/python:linalg_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:nn_ops", - ], - tags = [ - "manual", # b/64490288 - "notap", - ], -) - -cuda_py_test( - name = "custom_grad_test", - size = "small", - srcs = ["python/kernel_tests/custom_grad_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/contrib/layers:layers_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:gradients", - "//tensorflow/python:init_ops", - "//tensorflow/python:platform_test", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - ], -) - -cuda_py_test( - name = "layers_conv_variational_test", - size = "small", - srcs = ["python/kernel_tests/layers_conv_variational_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/contrib/distributions:distributions_py", - "//tensorflow/python/ops/distributions", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:gradients", - "//tensorflow/python:linalg_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:nn_ops", - ], -) - -cuda_py_test( - name = "layers_dense_variational_test", - size = "small", - srcs = ["python/kernel_tests/layers_dense_variational_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/contrib/distributions:distributions_py", - "//tensorflow/python/ops/distributions", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:gradients", - "//tensorflow/python:linalg_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:nn_ops", - ], -) - cuda_py_test( name = "monte_carlo_test", size = "small", @@ -158,28 +77,9 @@ cuda_py_test( ], ) -cuda_py_test( - name = "halton_sequence_test", - size = "small", - srcs = ["python/kernel_tests/halton_sequence_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/python:array_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:platform_test", - "//tensorflow/python:random_ops", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - ], -) - cuda_py_test( name = "hmc_test", - size = "medium", + size = "large", srcs = ["python/kernel_tests/hmc_test.py"], additional_deps = [ ":bayesflow_py", @@ -196,49 +96,7 @@ cuda_py_test( "//tensorflow/python:platform_test", "//tensorflow/python:random_seed", ], -) - -cuda_py_test( - name = "sgld_optimizer_test", - size = "small", - srcs = ["python/kernel_tests/sgld_optimizer_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/contrib/distributions:distributions_py", - "//tensorflow/contrib/layers:layers_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:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform_test", - "//tensorflow/python:random_seed", - ], -) - -cuda_py_test( - name = "variational_sgd_optimizer_test", - size = "small", - srcs = ["python/kernel_tests/variational_sgd_optimizer_test.py"], - additional_deps = [ - ":bayesflow_py", - "//third_party/py/numpy", - "//tensorflow/contrib/distributions:distributions_py", - "//tensorflow/contrib/layers:layers_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:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform_test", - "//tensorflow/python:random_seed", - ], - tags = ["notsan"], + tags = ["nomsan"], ) filegroup( diff --git a/tensorflow/contrib/bayesflow/__init__.py b/tensorflow/contrib/bayesflow/__init__.py index 95b9452b1ada60c44672f37800ced2133d2bd8b2..f86820382682f79e85e6a92c7f63fa15bb8be1a3 100644 --- a/tensorflow/contrib/bayesflow/__init__.py +++ b/tensorflow/contrib/bayesflow/__init__.py @@ -21,29 +21,19 @@ from __future__ import division from __future__ import print_function # pylint: disable=unused-import,line-too-long -from tensorflow.contrib.bayesflow.python.ops import csiszar_divergence -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 metropolis_hastings from tensorflow.contrib.bayesflow.python.ops import monte_carlo -from tensorflow.contrib.bayesflow.python.ops import optimizers # pylint: enable=unused-import,line-too-long from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ - 'csiszar_divergence', - 'custom_grad', 'entropy', - 'halton_sequence', 'hmc', - 'layers', 'metropolis_hastings', 'monte_carlo', - 'optimizers', 'special_math', 'stochastic_variables', 'variational_inference', diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/csiszar_divergence_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/csiszar_divergence_test.py deleted file mode 100644 index 2e94b7206de4f7c40c89f083f3bfa2a22bb7b917..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/csiszar_divergence_test.py +++ /dev/null @@ -1,1004 +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 Csiszar Divergence Ops.""" - -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 csiszar_divergence_impl -from tensorflow.contrib.distributions.python.ops import mvn_diag as mvn_diag_lib -from tensorflow.contrib.distributions.python.ops import mvn_full_covariance as mvn_full_lib -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import linalg_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn_ops -from tensorflow.python.ops.distributions import kullback_leibler -from tensorflow.python.ops.distributions import normal as normal_lib -from tensorflow.python.platform import test - - -cd = csiszar_divergence_impl - - -def tridiag(d, diag_value, offdiag_value): - """d x d matrix with given value on diag, and one super/sub diag.""" - diag_mat = linalg_ops.eye(d) * (diag_value - offdiag_value) - three_bands = array_ops.matrix_band_part( - array_ops.fill([d, d], offdiag_value), 1, 1) - return diag_mat + three_bands - - -class AmariAlphaTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - for alpha in [-1., 0., 1., 2.]: - for normalized in [True, False]: - with self.test_session(graph=ops.Graph()): - self.assertAllClose( - cd.amari_alpha(0., alpha=alpha, - self_normalized=normalized).eval(), - 0.) - - def test_correct_when_alpha0(self): - with self.test_session(): - self.assertAllClose( - cd.amari_alpha(self._logu, alpha=0.).eval(), - -self._logu) - - self.assertAllClose( - cd.amari_alpha(self._logu, alpha=0., self_normalized=True).eval(), - -self._logu + (self._u - 1.)) - - def test_correct_when_alpha1(self): - with self.test_session(): - self.assertAllClose( - cd.amari_alpha(self._logu, alpha=1.).eval(), - self._u * self._logu) - - self.assertAllClose( - cd.amari_alpha(self._logu, alpha=1., self_normalized=True).eval(), - self._u * self._logu - (self._u - 1.)) - - def test_correct_when_alpha_not_01(self): - for alpha in [-2, -1., -0.5, 0.5, 2.]: - with self.test_session(graph=ops.Graph()): - self.assertAllClose( - cd.amari_alpha(self._logu, - alpha=alpha, - self_normalized=False).eval(), - ((self._u**alpha - 1)) / (alpha * (alpha - 1.))) - - self.assertAllClose( - cd.amari_alpha(self._logu, - alpha=alpha, - self_normalized=True).eval(), - ((self._u**alpha - 1.) - - alpha * (self._u - 1)) / (alpha * (alpha - 1.))) - - -class KLReverseTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - for normalized in [True, False]: - with self.test_session(graph=ops.Graph()): - self.assertAllClose( - cd.kl_reverse(0., self_normalized=normalized).eval(), - 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.kl_reverse(self._logu).eval(), - -self._logu) - - self.assertAllClose( - cd.kl_reverse(self._logu, self_normalized=True).eval(), - -self._logu + (self._u - 1.)) - - -class KLForwardTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - for normalized in [True, False]: - with self.test_session(graph=ops.Graph()): - self.assertAllClose( - cd.kl_forward(0., self_normalized=normalized).eval(), - 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.kl_forward(self._logu).eval(), - self._u * self._logu) - - self.assertAllClose( - cd.kl_forward(self._logu, self_normalized=True).eval(), - self._u * self._logu - (self._u - 1.)) - - -class JensenShannonTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.jensen_shannon(0.).eval(), np.log(0.25)) - - def test_symmetric(self): - with self.test_session(): - self.assertAllClose( - cd.jensen_shannon(self._logu).eval(), - cd.symmetrized_csiszar_function( - self._logu, cd.jensen_shannon).eval()) - - self.assertAllClose( - cd.jensen_shannon(self._logu, self_normalized=True).eval(), - cd.symmetrized_csiszar_function( - self._logu, - lambda x: cd.jensen_shannon(x, self_normalized=True)).eval()) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.jensen_shannon(self._logu).eval(), - (self._u * self._logu - - (1 + self._u) * np.log1p(self._u))) - - self.assertAllClose( - cd.jensen_shannon(self._logu, self_normalized=True).eval(), - (self._u * self._logu - - (1 + self._u) * np.log((1 + self._u) / 2))) - - -class ArithmeticGeometricMeanTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.arithmetic_geometric(0.).eval(), np.log(4)) - self.assertAllClose( - cd.arithmetic_geometric(0., self_normalized=True).eval(), 0.) - - def test_symmetric(self): - with self.test_session(): - self.assertAllClose( - cd.arithmetic_geometric(self._logu).eval(), - cd.symmetrized_csiszar_function( - self._logu, cd.arithmetic_geometric).eval()) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.arithmetic_geometric(self._logu).eval(), - (1. + self._u) * np.log((1. + self._u) / np.sqrt(self._u))) - - self.assertAllClose( - cd.arithmetic_geometric(self._logu, self_normalized=True).eval(), - (1. + self._u) * np.log(0.5 * (1. + self._u) / np.sqrt(self._u))) - - -class TotalVariationTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.total_variation(0.).eval(), 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.total_variation(self._logu).eval(), - 0.5 * np.abs(self._u - 1)) - - -class PearsonTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.pearson(0.).eval(), 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.pearson(self._logu).eval(), - np.square(self._u - 1)) - - -class SquaredHellingerTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.squared_hellinger(0.).eval(), 0.) - - def test_symmetric(self): - with self.test_session(): - self.assertAllClose( - cd.squared_hellinger(self._logu).eval(), - cd.symmetrized_csiszar_function( - self._logu, cd.squared_hellinger).eval()) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.squared_hellinger(self._logu).eval(), - np.square(np.sqrt(self._u) - 1)) - - -class TriangularTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.triangular(0.).eval(), 0.) - - def test_symmetric(self): - with self.test_session(): - self.assertAllClose( - cd.triangular(self._logu).eval(), - cd.symmetrized_csiszar_function( - self._logu, cd.triangular).eval()) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.triangular(self._logu).eval(), - np.square(self._u - 1) / (1 + self._u)) - - -class TPowerTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.t_power(0., t=-0.1).eval(), 0.) - self.assertAllClose(cd.t_power(0., t=0.5).eval(), 0.) - self.assertAllClose(cd.t_power(0., t=1.1).eval(), 0.) - self.assertAllClose( - cd.t_power(0., t=-0.1, self_normalized=True).eval(), 0.) - self.assertAllClose( - cd.t_power(0., t=0.5, self_normalized=True).eval(), 0.) - self.assertAllClose( - cd.t_power(0., t=1.1, self_normalized=True).eval(), 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.t_power(self._logu, t=np.float64(-0.1)).eval(), - self._u ** -0.1 - 1.) - self.assertAllClose( - cd.t_power(self._logu, t=np.float64(0.5)).eval(), - -self._u ** 0.5 + 1.) - self.assertAllClose( - cd.t_power(self._logu, t=np.float64(1.1)).eval(), - self._u ** 1.1 - 1.) - - def test_correct_self_normalized(self): - with self.test_session(): - self.assertAllClose( - cd.t_power(self._logu, t=np.float64(-0.1), - self_normalized=True).eval(), - self._u ** -0.1 - 1. + 0.1 * (self._u - 1.)) - self.assertAllClose( - cd.t_power(self._logu, t=np.float64(0.5), - self_normalized=True).eval(), - -self._u ** 0.5 + 1. + 0.5 * (self._u - 1.)) - self.assertAllClose( - cd.t_power(self._logu, t=np.float64(1.1), - self_normalized=True).eval(), - self._u ** 1.1 - 1. - 1.1 * (self._u - 1.)) - - -class Log1pAbsTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.log1p_abs(0.).eval(), 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.log1p_abs(self._logu).eval(), - self._u**(np.sign(self._u - 1)) - 1) - - -class JeffreysTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.jeffreys(0.).eval(), 0.) - - def test_symmetric(self): - with self.test_session(): - self.assertAllClose( - cd.jeffreys(self._logu).eval(), - cd.symmetrized_csiszar_function( - self._logu, cd.jeffreys).eval()) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.jeffreys(self._logu).eval(), - 0.5 * (self._u * self._logu - self._logu)) - - -class ChiSquareTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose(cd.chi_square(0.).eval(), 0.) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.chi_square(self._logu).eval(), - self._u**2 - 1) - - -class ModifiedGanTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10, 100) - self._u = np.exp(self._logu) - - def test_at_zero(self): - with self.test_session(): - self.assertAllClose( - cd.modified_gan(0.).eval(), np.log(2)) - self.assertAllClose( - cd.modified_gan(0., self_normalized=True).eval(), np.log(2)) - - def test_correct(self): - with self.test_session(): - self.assertAllClose( - cd.modified_gan(self._logu).eval(), - np.log1p(self._u) - self._logu) - - self.assertAllClose( - cd.modified_gan(self._logu, self_normalized=True).eval(), - np.log1p(self._u) - self._logu + 0.5 * (self._u - 1)) - - -class SymmetrizedCsiszarFunctionTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10., 100) - self._u = np.exp(self._logu) - - def test_jensen_shannon(self): - with self.test_session(): - - # The following functions come from the claim made in the - # symmetrized_csiszar_function docstring. - def js1(logu): - return (-logu - - (1. + math_ops.exp(logu)) * ( - nn_ops.softplus(logu))) - - def js2(logu): - return 2. * (math_ops.exp(logu) * ( - logu - nn_ops.softplus(logu))) - - self.assertAllClose( - cd.symmetrized_csiszar_function(self._logu, js1).eval(), - cd.jensen_shannon(self._logu).eval()) - - self.assertAllClose( - cd.symmetrized_csiszar_function(self._logu, js2).eval(), - cd.jensen_shannon(self._logu).eval()) - - def test_jeffreys(self): - with self.test_session(): - self.assertAllClose( - cd.symmetrized_csiszar_function(self._logu, cd.kl_reverse).eval(), - cd.jeffreys(self._logu).eval()) - - self.assertAllClose( - cd.symmetrized_csiszar_function(self._logu, cd.kl_forward).eval(), - cd.jeffreys(self._logu).eval()) - - -class DualCsiszarFunctionTest(test.TestCase): - - def setUp(self): - self._logu = np.linspace(-10., 10., 100) - self._u = np.exp(self._logu) - - def test_kl_forward(self): - with self.test_session(): - self.assertAllClose( - cd.dual_csiszar_function(self._logu, cd.kl_forward).eval(), - cd.kl_reverse(self._logu).eval()) - - def test_kl_reverse(self): - with self.test_session(): - self.assertAllClose( - cd.dual_csiszar_function(self._logu, cd.kl_reverse).eval(), - cd.kl_forward(self._logu).eval()) - - -class MonteCarloCsiszarFDivergenceTest(test.TestCase): - - def test_kl_forward(self): - with self.test_session() as sess: - q = normal_lib.Normal( - loc=np.ones(6), - scale=np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0])) - - p = normal_lib.Normal(loc=q.loc + 0.1, scale=q.scale - 0.2) - - approx_kl = cd.monte_carlo_csiszar_f_divergence( - f=cd.kl_forward, - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence( - f=lambda logu: cd.kl_forward(logu, self_normalized=True), - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - exact_kl = kullback_leibler.kl_divergence(p, q) - - [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([ - approx_kl, approx_kl_self_normalized, exact_kl]) - - self.assertAllClose(approx_kl_, exact_kl_, - rtol=0.08, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_, exact_kl_, - rtol=0.02, atol=0.) - - def test_kl_reverse(self): - with self.test_session() as sess: - - q = normal_lib.Normal( - loc=np.ones(6), - scale=np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0])) - - p = normal_lib.Normal(loc=q.loc + 0.1, scale=q.scale - 0.2) - - approx_kl = cd.monte_carlo_csiszar_f_divergence( - f=cd.kl_reverse, - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence( - f=lambda logu: cd.kl_reverse(logu, self_normalized=True), - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - exact_kl = kullback_leibler.kl_divergence(q, p) - - [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([ - approx_kl, approx_kl_self_normalized, exact_kl]) - - self.assertAllClose(approx_kl_, exact_kl_, - rtol=0.07, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_, exact_kl_, - rtol=0.02, atol=0.) - - def test_kl_reverse_multidim(self): - - with self.test_session() as sess: - d = 5 # Dimension - - p = mvn_full_lib.MultivariateNormalFullCovariance( - covariance_matrix=tridiag(d, diag_value=1, offdiag_value=0.5)) - - q = mvn_diag_lib.MultivariateNormalDiag(scale_diag=[0.5]*d) - - approx_kl = cd.monte_carlo_csiszar_f_divergence( - f=cd.kl_reverse, - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence( - f=lambda logu: cd.kl_reverse(logu, self_normalized=True), - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - exact_kl = kullback_leibler.kl_divergence(q, p) - - [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([ - approx_kl, approx_kl_self_normalized, exact_kl]) - - self.assertAllClose(approx_kl_, exact_kl_, - rtol=0.02, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_, exact_kl_, - rtol=0.08, atol=0.) - - def test_kl_forward_multidim(self): - - with self.test_session() as sess: - d = 5 # Dimension - - p = mvn_full_lib.MultivariateNormalFullCovariance( - covariance_matrix=tridiag(d, diag_value=1, offdiag_value=0.5)) - - # Variance is very high when approximating Forward KL, so we make - # scale_diag larger than in test_kl_reverse_multidim. This ensures q - # "covers" p and thus Var_q[p/q] is smaller. - q = mvn_diag_lib.MultivariateNormalDiag(scale_diag=[1.]*d) - - approx_kl = cd.monte_carlo_csiszar_f_divergence( - f=cd.kl_forward, - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence( - f=lambda logu: cd.kl_forward(logu, self_normalized=True), - p_log_prob=p.log_prob, - q=q, - num_draws=int(1e5), - seed=1) - - exact_kl = kullback_leibler.kl_divergence(p, q) - - [approx_kl_, approx_kl_self_normalized_, exact_kl_] = sess.run([ - approx_kl, approx_kl_self_normalized, exact_kl]) - - self.assertAllClose(approx_kl_, exact_kl_, - rtol=0.06, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_, exact_kl_, - rtol=0.05, atol=0.) - - def test_score_trick(self): - - with self.test_session() as sess: - d = 5 # Dimension - num_draws = int(1e5) - seed = 1 - - p = mvn_full_lib.MultivariateNormalFullCovariance( - covariance_matrix=tridiag(d, diag_value=1, offdiag_value=0.5)) - - # Variance is very high when approximating Forward KL, so we make - # scale_diag larger than in test_kl_reverse_multidim. This ensures q - # "covers" p and thus Var_q[p/q] is smaller. - s = array_ops.constant(1.) - q = mvn_diag_lib.MultivariateNormalDiag( - scale_diag=array_ops.tile([s], [d])) - - approx_kl = cd.monte_carlo_csiszar_f_divergence( - f=cd.kl_reverse, - p_log_prob=p.log_prob, - q=q, - num_draws=num_draws, - seed=seed) - - approx_kl_self_normalized = cd.monte_carlo_csiszar_f_divergence( - f=lambda logu: cd.kl_reverse(logu, self_normalized=True), - p_log_prob=p.log_prob, - q=q, - num_draws=num_draws, - seed=seed) - - approx_kl_score_trick = cd.monte_carlo_csiszar_f_divergence( - f=cd.kl_reverse, - p_log_prob=p.log_prob, - q=q, - num_draws=num_draws, - use_reparametrization=False, - seed=seed) - - approx_kl_self_normalized_score_trick = ( - cd.monte_carlo_csiszar_f_divergence( - f=lambda logu: cd.kl_reverse(logu, self_normalized=True), - p_log_prob=p.log_prob, - q=q, - num_draws=num_draws, - use_reparametrization=False, - seed=seed)) - - exact_kl = kullback_leibler.kl_divergence(q, p) - - grad_sum = lambda fs: gradients_impl.gradients(fs, s)[0] - - [ - approx_kl_grad_, - approx_kl_self_normalized_grad_, - approx_kl_score_trick_grad_, - approx_kl_self_normalized_score_trick_grad_, - exact_kl_grad_, - approx_kl_, - approx_kl_self_normalized_, - approx_kl_score_trick_, - approx_kl_self_normalized_score_trick_, - exact_kl_, - ] = sess.run([ - grad_sum(approx_kl), - grad_sum(approx_kl_self_normalized), - grad_sum(approx_kl_score_trick), - grad_sum(approx_kl_self_normalized_score_trick), - grad_sum(exact_kl), - approx_kl, - approx_kl_self_normalized, - approx_kl_score_trick, - approx_kl_self_normalized_score_trick, - exact_kl, - ]) - - # Test average divergence. - self.assertAllClose(approx_kl_, exact_kl_, - rtol=0.02, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_, exact_kl_, - rtol=0.08, atol=0.) - - self.assertAllClose(approx_kl_score_trick_, exact_kl_, - rtol=0.02, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_score_trick_, exact_kl_, - rtol=0.08, atol=0.) - - # Test average gradient-divergence. - self.assertAllClose(approx_kl_grad_, exact_kl_grad_, - rtol=0.007, atol=0.) - - self.assertAllClose(approx_kl_self_normalized_grad_, exact_kl_grad_, - rtol=0.011, atol=0.) - - self.assertAllClose(approx_kl_score_trick_grad_, exact_kl_grad_, - rtol=0.018, atol=0.) - - self.assertAllClose( - approx_kl_self_normalized_score_trick_grad_, exact_kl_grad_, - rtol=0.017, atol=0.) - - -class CsiszarVIMCOTest(test.TestCase): - - def _csiszar_vimco_helper(self, logu): - """Numpy implementation of `csiszar_vimco_helper`.""" - - # Since this is a naive/intuitive implementation, we compensate by using the - # highest precision we can. - logu = np.float128(logu) - n = logu.shape[0] - u = np.exp(logu) - loogeoavg_u = [] # Leave-one-out geometric-average of exp(logu). - for j in range(n): - loogeoavg_u.append(np.exp(np.mean( - [logu[i, ...] for i in range(n) if i != j], - axis=0))) - loogeoavg_u = np.array(loogeoavg_u) - - loosum_u = [] # Leave-one-out sum of exp(logu). - for j in range(n): - loosum_u.append(np.sum( - [u[i, ...] for i in range(n) if i != j], - axis=0)) - loosum_u = np.array(loosum_u) - - # Natural log of the average u except each is swapped-out for its - # leave-`i`-th-out Geometric average. - log_sooavg_u = np.log(loosum_u + loogeoavg_u) - np.log(n) - - log_avg_u = np.log(np.mean(u, axis=0)) - return log_avg_u, log_sooavg_u - - def _csiszar_vimco_helper_grad(self, logu, delta): - """Finite difference approximation of `grad(csiszar_vimco_helper, logu)`.""" - - # This code actually estimates the sum of the Jacobiab because that's what - # TF's `gradients` does. - np_log_avg_u1, np_log_sooavg_u1 = self._csiszar_vimco_helper( - logu[..., None] + np.diag([delta]*len(logu))) - np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper( - logu[..., None]) - return [ - (np_log_avg_u1 - np_log_avg_u) / delta, - np.sum(np_log_sooavg_u1 - np_log_sooavg_u, axis=0) / delta, - ] - - def test_vimco_helper_1(self): - """Tests that function calculation correctly handles batches.""" - - logu = np.linspace(-100., 100., 100).reshape([10, 2, 5]) - with self.test_session() as sess: - np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu) - [log_avg_u, log_sooavg_u] = sess.run(cd.csiszar_vimco_helper(logu)) - self.assertAllClose(np_log_avg_u, log_avg_u, - rtol=1e-8, atol=0.) - self.assertAllClose(np_log_sooavg_u, log_sooavg_u, - rtol=1e-8, atol=0.) - - def test_vimco_helper_2(self): - """Tests that function calculation correctly handles overflow.""" - - # Using 700 (rather than 1e3) since naive numpy version can't handle higher. - logu = np.float32([0., 700, -1, 1]) - with self.test_session() as sess: - np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu) - [log_avg_u, log_sooavg_u] = sess.run(cd.csiszar_vimco_helper(logu)) - self.assertAllClose(np_log_avg_u, log_avg_u, - rtol=1e-6, atol=0.) - self.assertAllClose(np_log_sooavg_u, log_sooavg_u, - rtol=1e-5, atol=0.) - - def test_vimco_helper_3(self): - """Tests that function calculation correctly handles underlow.""" - - logu = np.float32([0., -1000, -1, 1]) - with self.test_session() as sess: - np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu) - [log_avg_u, log_sooavg_u] = sess.run(cd.csiszar_vimco_helper(logu)) - self.assertAllClose(np_log_avg_u, log_avg_u, - rtol=1e-5, atol=0.) - self.assertAllClose(np_log_sooavg_u, log_sooavg_u, - rtol=1e-4, atol=1e-15) - - def test_vimco_helper_gradient_using_finite_difference_1(self): - """Tests that gradient calculation correctly handles batches.""" - - logu_ = np.linspace(-100., 100., 100).reshape([10, 2, 5]) - with self.test_session() as sess: - logu = array_ops.constant(logu_) - - grad = lambda flogu: gradients_impl.gradients(flogu, logu)[0] - log_avg_u, log_sooavg_u = cd.csiszar_vimco_helper(logu) - - [ - grad_log_avg_u, - grad_log_sooavg_u, - ] = sess.run([grad(log_avg_u), grad(log_sooavg_u)]) - - # We skip checking against finite-difference approximation since it - # doesn't support batches. - - # Verify claim in docstring. - self.assertAllClose( - np.ones_like(grad_log_avg_u.sum(axis=0)), - grad_log_avg_u.sum(axis=0)) - self.assertAllClose( - np.ones_like(grad_log_sooavg_u.mean(axis=0)), - grad_log_sooavg_u.mean(axis=0)) - - def test_vimco_helper_gradient_using_finite_difference_2(self): - """Tests that gradient calculation correctly handles overflow.""" - - delta = 1e-3 - logu_ = np.float32([0., 1000, -1, 1]) - with self.test_session() as sess: - logu = array_ops.constant(logu_) - - [ - np_grad_log_avg_u, - np_grad_log_sooavg_u, - ] = self._csiszar_vimco_helper_grad(logu_, delta) - - grad = lambda flogu: gradients_impl.gradients(flogu, logu)[0] - log_avg_u, log_sooavg_u = cd.csiszar_vimco_helper(logu) - - [ - grad_log_avg_u, - grad_log_sooavg_u, - ] = sess.run([grad(log_avg_u), grad(log_sooavg_u)]) - - self.assertAllClose(np_grad_log_avg_u, grad_log_avg_u, - rtol=delta, atol=0.) - self.assertAllClose(np_grad_log_sooavg_u, grad_log_sooavg_u, - rtol=delta, atol=0.) - # Verify claim in docstring. - self.assertAllClose( - np.ones_like(grad_log_avg_u.sum(axis=0)), - grad_log_avg_u.sum(axis=0)) - self.assertAllClose( - np.ones_like(grad_log_sooavg_u.mean(axis=0)), - grad_log_sooavg_u.mean(axis=0)) - - def test_vimco_helper_gradient_using_finite_difference_3(self): - """Tests that gradient calculation correctly handles underlow.""" - - delta = 1e-3 - logu_ = np.float32([0., -1000, -1, 1]) - with self.test_session() as sess: - logu = array_ops.constant(logu_) - - [ - np_grad_log_avg_u, - np_grad_log_sooavg_u, - ] = self._csiszar_vimco_helper_grad(logu_, delta) - - grad = lambda flogu: gradients_impl.gradients(flogu, logu)[0] - log_avg_u, log_sooavg_u = cd.csiszar_vimco_helper(logu) - - [ - grad_log_avg_u, - grad_log_sooavg_u, - ] = sess.run([grad(log_avg_u), grad(log_sooavg_u)]) - - self.assertAllClose(np_grad_log_avg_u, grad_log_avg_u, - rtol=delta, atol=0.) - self.assertAllClose(np_grad_log_sooavg_u, grad_log_sooavg_u, - rtol=delta, atol=0.) - # Verify claim in docstring. - self.assertAllClose( - np.ones_like(grad_log_avg_u.sum(axis=0)), - grad_log_avg_u.sum(axis=0)) - self.assertAllClose( - np.ones_like(grad_log_sooavg_u.mean(axis=0)), - grad_log_sooavg_u.mean(axis=0)) - - def test_vimco_and_gradient(self): - - with self.test_session() as sess: - dims = 5 # Dimension - num_draws = int(20) - num_batch_draws = int(3) - seed = 1 - - f = lambda logu: cd.kl_reverse(logu, self_normalized=False) - np_f = lambda logu: -logu - - p = mvn_full_lib.MultivariateNormalFullCovariance( - covariance_matrix=tridiag(dims, diag_value=1, offdiag_value=0.5)) - - # Variance is very high when approximating Forward KL, so we make - # scale_diag larger than in test_kl_reverse_multidim. This ensures q - # "covers" p and thus Var_q[p/q] is smaller. - s = array_ops.constant(1.) - q = mvn_diag_lib.MultivariateNormalDiag( - scale_diag=array_ops.tile([s], [dims])) - - vimco = cd.csiszar_vimco( - f=f, - p_log_prob=p.log_prob, - q=q, - num_draws=num_draws, - num_batch_draws=num_batch_draws, - seed=seed) - - x = q.sample(sample_shape=[num_draws, num_batch_draws], - seed=seed) - x = array_ops.stop_gradient(x) - logu = p.log_prob(x) - q.log_prob(x) - f_log_sum_u = f(cd.csiszar_vimco_helper(logu)[0]) - - grad_sum = lambda fs: gradients_impl.gradients(fs, s)[0] - - def jacobian(x): - # Warning: this function is slow and may not even finish if prod(shape) - # is larger than, say, 100. - shape = x.shape.as_list() - assert all(s is not None for s in shape) - x = array_ops.reshape(x, shape=[-1]) - r = [grad_sum(x[i]) for i in range(np.prod(shape))] - return array_ops.reshape(array_ops.stack(r), shape=shape) - - [ - logu_, - jacobian_logqx_, - vimco_, - grad_vimco_, - f_log_sum_u_, - grad_mean_f_log_sum_u_, - ] = sess.run([ - logu, - jacobian(q.log_prob(x)), - vimco, - grad_sum(vimco), - f_log_sum_u, - grad_sum(f_log_sum_u) / num_batch_draws, - ]) - - np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(logu_) - - # Test VIMCO loss is correct. - self.assertAllClose(np_f(np_log_avg_u).mean(axis=0), vimco_, - rtol=1e-5, atol=0.) - - # Test gradient of VIMCO loss is correct. - # - # To make this computation we'll inject two gradients from TF: - # - grad[mean(f(log(sum(p(x)/q(x)))))] - # - jacobian[log(q(x))]. - # - # We now justify why using these (and only these) TF values for - # ground-truth does not undermine the completeness of this test. - # - # Regarding `grad_mean_f_log_sum_u_`, note that we validate the - # correctness of the zero-th order derivative (for each batch member). - # Since `cd.csiszar_vimco_helper` itself does not manipulate any gradient - # information, we can safely rely on TF. - self.assertAllClose(np_f(np_log_avg_u), f_log_sum_u_, rtol=1e-4, atol=0.) - # - # Regarding `jacobian_logqx_`, note that testing the gradient of - # `q.log_prob` is outside the scope of this unit-test thus we may safely - # use TF to find it. - - # The `mean` is across batches and the `sum` is across iid samples. - np_grad_vimco = ( - grad_mean_f_log_sum_u_ - + np.mean( - np.sum( - jacobian_logqx_ * (np_f(np_log_avg_u) - - np_f(np_log_sooavg_u)), - axis=0), - axis=0)) - - self.assertAllClose(np_grad_vimco, grad_vimco_, - rtol=1e-5, atol=0.) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/custom_grad_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/custom_grad_test.py deleted file mode 100644 index a95df31ac1fd9f5038abe779391ccba5f7fe408d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/custom_grad_test.py +++ /dev/null @@ -1,157 +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 Custom Gradient Ops.""" - -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 custom_grad_impl -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 gradients_impl -from tensorflow.python.ops import init_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 - - -cg = custom_grad_impl - - -class CustomGradientTest(test.TestCase): - - def test_works_correctly(self): - with self.test_session() as sess: - f = lambda x: x**2 / 2 - g = lambda x: (x - 1)**3 / 3 - x_ = np.linspace(-100, 100, int(1e4)) + [0.] - - x = constant_op.constant(x_) - fx = cg.custom_gradient(f(x), g(x), x) - gx = gradients_impl.gradients(fx, x)[0] - [fx_, gx_] = sess.run([fx, gx]) - - self.assertAllClose(f(x_), fx_) - self.assertAllClose(g(x_), gx_) - - def test_works_correctly_both_f_g_zero(self): - with self.test_session() as sess: - f = lambda x: x**2 / 2 - g = lambda x: x**3 / 3 - x_ = np.linspace(-100, 100, int(1e4)) + [0.] - - x = constant_op.constant(x_) - fx = cg.custom_gradient(f(x), g(x), x) - gx = gradients_impl.gradients(fx, x)[0] - [fx_, gx_] = sess.run([fx, gx]) - - self.assertAllClose(f(x_), fx_) - self.assertAllClose(g(x_), gx_) - - def test_works_correctly_vector_of_vars(self): - with self.test_session() as sess: - x = variable_scope.get_variable( - name="x", - shape=[], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer(2)) - y = variable_scope.get_variable( - name="y", - shape=[], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer(3)) - sess.run([variables.global_variables_initializer()]) - - f = lambda z: z[0] * z[1] - g = lambda z: z[0]**2 * z[1]**2 / 2 - - z = array_ops.stack([x, y]) - fz = cg.custom_gradient(f(z), g(z), z, axis=0) - gz = gradients_impl.gradients(fz, variables.trainable_variables()) - [z_, fz_, gx_, gy_] = sess.run([z, fz, gz[0], gz[1]]) - - self.assertEqual(f(z_), fz_) - self.assertEqual(g(z_), gx_) - self.assertEqual(g(z_), gy_) - - def test_works_correctly_side_vars(self): - with self.test_session() as sess: - x_ = np.float32(2.1) # Adding extra tenth to force imprecision. - y_ = np.float32(3.1) - x = variable_scope.get_variable( - name="x", - shape=[], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer(x_)) - y = variable_scope.get_variable( - name="y", - shape=[], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer(y_)) - sess.run([variables.global_variables_initializer()]) - - f = lambda x: x * y - g = lambda z: math_ops.square(x) * y - - fx = cg.custom_gradient(f(x), g(x), x) - gx = gradients_impl.gradients(fx, variables.trainable_variables()) - [x_, fx_, gx_] = sess.run([x, fx, gx[0]]) - gy_ = gx[1] - - self.assertEqual(x_ * y_, fx_) - self.assertEqual(np.square(x_) * y_, gx_) - self.assertEqual(None, gy_) - - def test_works_correctly_fx_gx_manually_stopped(self): - with self.test_session() as sess: - x_ = np.float32(2.1) # Adding extra tenth to force imprecision. - y_ = np.float32(3.1) - x = variable_scope.get_variable( - name="x", - shape=[], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer(x_)) - y = variable_scope.get_variable( - name="y", - shape=[], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer(y_)) - sess.run([variables.global_variables_initializer()]) - - stop = array_ops.stop_gradient # For readability. - - # Basically we need to stop the `x` portion of `f`. And when we supply the - # arg to `custom_gradient` we need to stop the complement, i.e., the `y` - # part. - f = lambda x: stop(x) * y - g = lambda x: stop(math_ops.square(x)) * y - fx = cg.custom_gradient(f(x), g(x), x + stop(y), - fx_gx_manually_stopped=True) - - gx = gradients_impl.gradients(fx, variables.trainable_variables()) - [x_, fx_, gx_, gy_] = sess.run([x, fx, gx[0], gx[1]]) - - self.assertEqual(x_ * y_, fx_) - self.assertEqual(np.square(x_) * y_, gx_) - self.assertEqual(x_, gy_) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/halton_sequence_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/halton_sequence_test.py deleted file mode 100644 index 0a85862abfd744a86b9a38e10dbb5b985d0a0e94..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/halton_sequence_test.py +++ /dev/null @@ -1,131 +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 halton_sequence.py.""" - -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 halton_sequence as halton -from tensorflow.contrib.bayesflow.python.ops import monte_carlo_impl as monte_carlo_lib -from tensorflow.python.framework import dtypes -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.platform import test - - -mc = monte_carlo_lib - - -class HaltonSequenceTest(test.TestCase): - - def test_known_values_small_bases(self): - with self.test_session(): - # The first five elements of the Halton sequence with base 2 and 3 - expected = np.array(((1. / 2, 1. / 3), - (1. / 4, 2. / 3), - (3. / 4, 1. / 9), - (1. / 8, 4. / 9), - (5. / 8, 7. / 9)), dtype=np.float32) - sample = halton.sample(2, num_samples=5) - self.assertAllClose(expected, sample.eval(), rtol=1e-6) - - def test_sample_indices(self): - with self.test_session(): - dim = 5 - indices = math_ops.range(10, dtype=dtypes.int32) - sample_direct = halton.sample(dim, num_samples=10) - sample_from_indices = halton.sample(dim, sample_indices=indices) - self.assertAllClose(sample_direct.eval(), sample_from_indices.eval(), - rtol=1e-6) - - def test_dtypes_works_correctly(self): - with self.test_session(): - dim = 3 - sample_float32 = halton.sample(dim, num_samples=10, dtype=dtypes.float32) - sample_float64 = halton.sample(dim, num_samples=10, dtype=dtypes.float64) - self.assertEqual(sample_float32.eval().dtype, np.float32) - self.assertEqual(sample_float64.eval().dtype, np.float64) - - def test_normal_integral_mean_and_var_correctly_estimated(self): - n = int(1000) - # This test is almost identical to the similarly named test in - # monte_carlo_test.py. The only difference is that we use the Halton - # samples instead of the random samples to evaluate the expectations. - # MC with pseudo random numbers converges at the rate of 1/ Sqrt(N) - # (N=number of samples). For QMC in low dimensions, the expected convergence - # rate is ~ 1/N. Hence we should only need 1e3 samples as compared to the - # 1e6 samples used in the pseudo-random monte carlo. - with self.test_session(): - mu_p = array_ops.constant([-1.0, 1.0], dtype=dtypes.float64) - mu_q = array_ops.constant([0.0, 0.0], dtype=dtypes.float64) - sigma_p = array_ops.constant([0.5, 0.5], dtype=dtypes.float64) - sigma_q = array_ops.constant([1.0, 1.0], dtype=dtypes.float64) - p = normal_lib.Normal(loc=mu_p, scale=sigma_p) - q = normal_lib.Normal(loc=mu_q, scale=sigma_q) - - cdf_sample = halton.sample(2, num_samples=n, dtype=dtypes.float64) - q_sample = q.quantile(cdf_sample) - - # Compute E_p[X]. - e_x = mc.expectation_importance_sampler( - f=lambda x: x, log_p=p.log_prob, sampling_dist_q=q, z=q_sample, - seed=42) - - # Compute E_p[X^2]. - e_x2 = mc.expectation_importance_sampler( - f=math_ops.square, log_p=p.log_prob, sampling_dist_q=q, z=q_sample, - seed=42) - - stddev = math_ops.sqrt(e_x2 - math_ops.square(e_x)) - # Keep the tolerance levels the same as in monte_carlo_test.py. - self.assertEqual(p.batch_shape, e_x.get_shape()) - self.assertAllClose(p.mean().eval(), e_x.eval(), rtol=0.01) - self.assertAllClose(p.stddev().eval(), stddev.eval(), rtol=0.02) - - def test_docstring_example(self): - # Produce the first 1000 members of the Halton sequence in 3 dimensions. - num_samples = 1000 - dim = 3 - with self.test_session(): - sample = halton.sample(dim, num_samples=num_samples) - - # Evaluate the integral of x_1 * x_2^2 * x_3^3 over the three dimensional - # hypercube. - powers = math_ops.range(1.0, limit=dim + 1) - integral = math_ops.reduce_mean( - math_ops.reduce_prod(sample ** powers, axis=-1)) - true_value = 1.0 / math_ops.reduce_prod(powers + 1.0) - - # Produces a relative absolute error of 1.7%. - self.assertAllClose(integral.eval(), true_value.eval(), rtol=0.02) - - # Now skip the first 1000 samples and recompute the integral with the next - # thousand samples. The sample_indices argument can be used to do this. - - sample_indices = math_ops.range(start=1000, limit=1000 + num_samples, - dtype=dtypes.int32) - sample_leaped = halton.sample(dim, sample_indices=sample_indices) - - integral_leaped = math_ops.reduce_mean( - math_ops.reduce_prod(sample_leaped ** powers, axis=-1)) - self.assertAllClose(integral_leaped.eval(), true_value.eval(), rtol=0.001) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/hmc_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/hmc_test.py index cbc66b6dc13db62c25952de6b6c13b2fdfe27f12..dabadfc7b6a3da8786e88d559fe2d05b44599ca0 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) + event_dims = math_ops.range(independent_chain_ndims, + array_ops.rank(x)) - _, new_m, potential_1, _ = ( - hmc.leapfrog_integrator(step_size, hmc_lf_steps, x, - m, potential_and_grad, grad_0)) - - 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,314 @@ 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 + + log_accept_ratio_, samples_, expected_x_ = sess.run( + [kernel_results.log_accept_ratio, samples, expected_x], + feed_dict) + + actual_x = samples_.mean() + actual_exp_x = np.exp(samples_).mean() + acceptance_probs = np.exp(np.minimum(log_accept_ratio_, 0.)) + + 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 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 accepted a decent amount, even near end of the chain. + self.assertLess( + 0.1, kernel_results_.is_accepted[int(0.9 * num_results):].mean()) + # We better not have any NaNs in states 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. + acceptance_probs = np.exp(np.minimum( + kernel_results_.log_accept_ratio, 0.)) + self.assertAllEqual( + np.zeros_like(pstates_[neg_inf_mask]), + 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) + + [ + log_accept_ratio_, + bad_log_accept_ratio_, + initial_draws_, + updated_draws_, + fake_draws_, + ] = sess.run([ + kernel_results.log_accept_ratio, + bad_kernel_results.log_accept_ratio, + 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) + acceptance_probs = np.exp(np.minimum(log_accept_ratio_, 0.)) + bad_acceptance_probs = np.exp(np.minimum(bad_log_accept_ratio_, 0.)) + 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 +439,28 @@ 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]) - - def testKernelLeavesTargetInvariant12(self): - self._kernel_leaves_target_invariant_wrapper([1, 2]) - - 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) - - def target_log_prob(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) - - 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]) - - true_normalizer = -self._shape_param * np.log(self._rate_param) - true_normalizer += special.gammaln(self._shape_param) - true_normalizer *= normalizer_multiplier - - 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) + self._kernel_leaves_target_invariant_wrapper(2) - 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') - - initial_draws = np.random.normal(size=[30, 2, 1]) - feed_dict = {x_ph: initial_draws} - - self._ais_gets_correct_log_normalizer(x_ph, event_dims, sess, - feed_dict) - - def testAISNullShape(self): - self._ais_gets_correct_log_normalizer_wrapper([]) - - def testAIS1(self): - self._ais_gets_correct_log_normalizer_wrapper([1]) - - def testAIS2(self): - self._ais_gets_correct_log_normalizer_wrapper([2]) - - def testAIS12(self): - self._ais_gets_correct_log_normalizer_wrapper([1, 2]) + def testKernelLeavesTargetInvariant3(self): + self._kernel_leaves_target_invariant_wrapper(3) def testNanRejection(self): """Tests that an update that yields NaN potentials gets rejected. @@ -359,86 +473,265 @@ 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_, log_accept_ratio_ = sess.run( + [initial_x, updated_x, kernel_results.log_accept_ratio]) + acceptance_probs = np.exp(np.minimum(log_accept_ratio_, 0.)) + + logging_ops.vlog(1, "initial_x = {}".format(initial_x_)) + logging_ops.vlog(1, "updated_x = {}".format(updated_x_)) + logging_ops.vlog(1, "log_accept_ratio = {}".format(log_accept_ratio_)) + + 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_, log_accept_ratio_ = sess.run( + [initial_x, updated_x, kernel_results.log_accept_ratio]) + acceptance_probs = np.exp(np.minimum(log_accept_ratio_, 0.)) + + logging_ops.vlog(1, "initial_x = {}".format(initial_x_)) + logging_ops.vlog(1, "updated_x = {}".format(updated_x_)) + logging_ops.vlog(1, "log_accept_ratio = {}".format(log_accept_ratio_)) + + 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_, log_accept_ratio_ = sess.run( + [states, kernel_results.log_accept_ratio]) + self.assertEqual(dtype, states_.dtype) + self.assertEqual(dtype, log_accept_ratio_.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 deleted file mode 100644 index 750afb6654311fea30a1dc6b31b20aa3b4160ae2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/layers_conv_variational_test.py +++ /dev/null @@ -1,521 +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 convolutional Bayesian layers.""" - -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 - - -class Counter(object): - """Helper class to manage incrementing a counting `int`.""" - - def __init__(self): - self._value = -1 - - @property - def value(self): - return self._value - - def __call__(self): - self._value += 1 - return self._value - - -class MockDistribution(independent_lib.Independent): - """Monitors layer calls to the underlying distribution.""" - - def __init__(self, result_sample, result_log_prob, loc=None, scale=None): - self.result_sample = result_sample - self.result_log_prob = result_log_prob - self.result_loc = loc - self.result_scale = scale - self.result_distribution = normal_lib.Normal(loc=0.0, scale=1.0) - if loc is not None and scale is not None: - self.result_distribution = normal_lib.Normal(loc=self.result_loc, - scale=self.result_scale) - self.called_log_prob = Counter() - self.called_sample = Counter() - self.called_loc = Counter() - self.called_scale = Counter() - - def log_prob(self, *args, **kwargs): - self.called_log_prob() - return self.result_log_prob - - def sample(self, *args, **kwargs): - self.called_sample() - return self.result_sample - - @property - def distribution(self): # for dummy check on Independent(Normal) - return self.result_distribution - - @property - def loc(self): - self.called_loc() - return self.result_loc - - @property - def scale(self): - self.called_scale() - return self.result_scale - - -class MockKLDivergence(object): - """Monitors layer calls to the divergence implementation.""" - - def __init__(self, result): - self.result = result - self.args = [] - self.called = Counter() - - def __call__(self, *args, **kwargs): - self.called() - self.args.append(args) - return self.result - - -class ConvVariational(test.TestCase): - - def _testKLPenaltyKernel(self, layer_class): - with self.test_session(): - layer = layer_class(filters=2, kernel_size=3) - 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 in (prob_layers_lib.Conv2DReparameterization, - prob_layers_lib.Conv2DFlipout): - inputs = random_ops.random_uniform([2, 3, 3, 1], seed=1) - 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. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 0) - self.assertListEqual(layer.losses, losses) - - _ = layer(inputs) - - # Yes keys. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 1) - self.assertListEqual(layer.losses, losses) - - def _testKLPenaltyBoth(self, layer_class): - def _make_normal(dtype, *args): # pylint: disable=unused-argument - return normal_lib.Normal( - loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)) - with self.test_session(): - layer = layer_class( - filters=2, - kernel_size=3, - bias_posterior_fn=prob_layers_util.default_mean_field_normal_fn(), - bias_prior_fn=_make_normal) - 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 in (prob_layers_lib.Conv2DReparameterization, - prob_layers_lib.Conv2DFlipout): - inputs = random_ops.random_uniform([2, 3, 3, 1], seed=1) - 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. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 0) - self.assertListEqual(layer.losses, losses) - - _ = layer(inputs) - - # Yes keys. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 2) - self.assertListEqual(layer.losses, losses) - - 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: - (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), - filter_shape=tensor_shape.TensorShape(kernel_shape), - padding="SAME") - expected_outputs = convolution_op(inputs, kernel_posterior.result_sample) - expected_outputs = nn.bias_add(expected_outputs, - bias_posterior.result_sample, - data_format="NHWC") - - [ - expected_outputs_, actual_outputs_, - expected_kernel_, actual_kernel_, - expected_kernel_divergence_, actual_kernel_divergence_, - expected_bias_, actual_bias_, - expected_bias_divergence_, actual_bias_divergence_, - ] = sess.run([ - expected_outputs, outputs, - kernel_posterior.result_sample, layer.kernel_posterior_tensor, - kernel_divergence.result, kl_penalty[0], - bias_posterior.result_sample, layer.bias_posterior_tensor, - bias_divergence.result, kl_penalty[1], - ]) - - self.assertAllClose( - expected_kernel_, actual_kernel_, - rtol=1e-6, atol=0.) - 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, - kernel_posterior.result_sample]], - kernel_divergence.args) - - self.assertAllEqual( - [[bias_posterior.distribution, - bias_prior.distribution, - bias_posterior.result_sample]], - bias_divergence.args) - - 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 testKLPenaltyBothConv3DFlipout(self): - self._testKLPenaltyBoth(prob_layers_lib.Conv3DFlipout) - - def testConv1DReparameterization(self): - self._testConvReparameterization(prob_layers_lib.Conv1DReparameterization) - - def testConv2DReparameterization(self): - self._testConvReparameterization(prob_layers_lib.Conv2DReparameterization) - - def testConv3DReparameterization(self): - self._testConvReparameterization(prob_layers_lib.Conv3DReparameterization) - - def testConv1DFlipout(self): - self._testConvFlipout(prob_layers_lib.Conv1DFlipout) - - def testConv2DFlipout(self): - self._testConvFlipout(prob_layers_lib.Conv2DFlipout) - - def testConv3DFlipout(self): - self._testConvFlipout(prob_layers_lib.Conv3DFlipout) - - def testRandomConv1DFlipout(self): - self._testRandomConvFlipout(prob_layers_lib.Conv1DFlipout) - - -if __name__ == "__main__": - test.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 deleted file mode 100644 index 342f38ccec7ec74db1b393d6cdc22300205cc547..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/layers_dense_variational_test.py +++ /dev/null @@ -1,443 +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 dense Bayesian layers.""" - -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_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 -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 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 - - -class Counter(object): - """Helper class to manage incrementing a counting `int`.""" - - def __init__(self): - self._value = -1 - - @property - def value(self): - return self._value - - def __call__(self): - self._value += 1 - return self._value - - -class MockDistribution(independent_lib.Independent): - """Monitors layer calls to the underlying distribution.""" - - def __init__(self, result_sample, result_log_prob, loc=None, scale=None): - self.result_sample = result_sample - self.result_log_prob = result_log_prob - self.result_loc = loc - self.result_scale = scale - self.result_distribution = normal_lib.Normal(loc=0.0, scale=1.0) - if loc is not None and scale is not None: - self.result_distribution = normal_lib.Normal(loc=self.result_loc, - scale=self.result_scale) - self.called_log_prob = Counter() - self.called_sample = Counter() - self.called_loc = Counter() - self.called_scale = Counter() - - def log_prob(self, *args, **kwargs): - self.called_log_prob() - return self.result_log_prob - - def sample(self, *args, **kwargs): - self.called_sample() - return self.result_sample - - @property - def distribution(self): # for dummy check on Independent(Normal) - return self.result_distribution - - @property - def loc(self): - self.called_loc() - return self.result_loc - - @property - def scale(self): - self.called_scale() - return self.result_scale - - -class MockKLDivergence(object): - """Monitors layer calls to the divergence implementation.""" - - def __init__(self, result): - self.result = result - self.args = [] - self.called = Counter() - - def __call__(self, *args, **kwargs): - self.called() - self.args.append(args) - return self.result - - -class DenseVariational(test.TestCase): - - def _testKLPenaltyKernel(self, layer_class): - with self.test_session(): - layer = layer_class(units=2) - inputs = random_ops.random_uniform([2, 3], seed=1) - - # No keys. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 0) - self.assertListEqual(layer.losses, losses) - - _ = layer(inputs) - - # Yes keys. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 1) - self.assertListEqual(layer.losses, losses) - - def _testKLPenaltyBoth(self, layer_class): - def _make_normal(dtype, *args): # pylint: disable=unused-argument - return normal_lib.Normal( - loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)) - with self.test_session(): - layer = layer_class( - units=2, - bias_posterior_fn=prob_layers_util.default_mean_field_normal_fn(), - bias_prior_fn=_make_normal) - inputs = random_ops.random_uniform([2, 3], seed=1) - - # No keys. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 0) - self.assertListEqual(layer.losses, losses) - - _ = layer(inputs) - - # Yes keys. - losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - self.assertEqual(len(losses), 2) - self.assertListEqual(layer.losses, losses) - - def _testDenseSetUp(self, layer_class, batch_size, in_size, out_size, - **kwargs): - seed = Counter() - inputs = random_ops.random_uniform([batch_size, in_size], seed=seed()) - - kernel_size = [in_size, out_size] - kernel_posterior = MockDistribution( - loc=random_ops.random_uniform(kernel_size, seed=seed()), - scale=random_ops.random_uniform(kernel_size, seed=seed()), - result_log_prob=random_ops.random_uniform(kernel_size, seed=seed()), - result_sample=random_ops.random_uniform(kernel_size, seed=seed())) - kernel_prior = MockDistribution( - result_log_prob=random_ops.random_uniform(kernel_size, seed=seed()), - result_sample=random_ops.random_uniform(kernel_size, seed=seed())) - kernel_divergence = MockKLDivergence( - result=random_ops.random_uniform(kernel_size, seed=seed())) - - bias_size = [out_size] - 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( - units=out_size, - 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) - - def testKLPenaltyKernelReparameterization(self): - self._testKLPenaltyKernel(prob_layers_lib.DenseReparameterization) - - def testKLPenaltyKernelLocalReparameterization(self): - self._testKLPenaltyKernel(prob_layers_lib.DenseLocalReparameterization) - - def testKLPenaltyKernelFlipout(self): - self._testKLPenaltyKernel(prob_layers_lib.DenseFlipout) - - def testKLPenaltyBothReparameterization(self): - self._testKLPenaltyBoth(prob_layers_lib.DenseReparameterization) - - def testKLPenaltyBothLocalReparameterization(self): - self._testKLPenaltyBoth(prob_layers_lib.DenseLocalReparameterization) - - def testKLPenaltyBothFlipout(self): - self._testKLPenaltyBoth(prob_layers_lib.DenseFlipout) - - def testDenseReparameterization(self): - batch_size, in_size, out_size = 2, 3, 4 - with self.test_session() as sess: - (kernel_posterior, kernel_prior, kernel_divergence, - bias_posterior, bias_prior, bias_divergence, layer, inputs, - outputs, kl_penalty) = self._testDenseSetUp( - prob_layers_lib.DenseReparameterization, - batch_size, in_size, out_size) - - expected_outputs = ( - math_ops.matmul(inputs, kernel_posterior.result_sample) + - bias_posterior.result_sample) - - [ - expected_outputs_, actual_outputs_, - expected_kernel_, actual_kernel_, - expected_kernel_divergence_, actual_kernel_divergence_, - expected_bias_, actual_bias_, - expected_bias_divergence_, actual_bias_divergence_, - ] = sess.run([ - expected_outputs, outputs, - kernel_posterior.result_sample, layer.kernel_posterior_tensor, - kernel_divergence.result, kl_penalty[0], - bias_posterior.result_sample, layer.bias_posterior_tensor, - bias_divergence.result, kl_penalty[1], - ]) - - self.assertAllClose( - expected_kernel_, actual_kernel_, - rtol=1e-6, atol=0.) - 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, - kernel_posterior.result_sample]], - kernel_divergence.args) - - self.assertAllEqual( - [[bias_posterior.distribution, - bias_prior.distribution, - bias_posterior.result_sample]], - bias_divergence.args) - - def testDenseLocalReparameterization(self): - batch_size, in_size, out_size = 2, 3, 4 - with self.test_session() as sess: - (kernel_posterior, kernel_prior, kernel_divergence, - bias_posterior, bias_prior, bias_divergence, layer, inputs, - outputs, kl_penalty) = self._testDenseSetUp( - prob_layers_lib.DenseLocalReparameterization, - batch_size, in_size, out_size) - - expected_kernel_posterior_affine = normal_lib.Normal( - loc=math_ops.matmul(inputs, kernel_posterior.result_loc), - scale=math_ops.matmul( - inputs**2., kernel_posterior.result_scale**2)**0.5) - expected_kernel_posterior_affine_tensor = ( - expected_kernel_posterior_affine.sample(seed=42)) - expected_outputs = (expected_kernel_posterior_affine_tensor + - bias_posterior.result_sample) - - [ - 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 testDenseFlipout(self): - batch_size, in_size, out_size = 2, 3, 4 - with self.test_session() as sess: - (kernel_posterior, kernel_prior, kernel_divergence, - bias_posterior, bias_prior, bias_divergence, layer, inputs, - outputs, kl_penalty) = self._testDenseSetUp( - prob_layers_lib.DenseFlipout, - batch_size, in_size, out_size, seed=44) - - 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)) - - sign_input = random_ops.random_uniform( - [batch_size, in_size], - 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( - [batch_size, out_size], - minval=0, - maxval=2, - dtype=dtypes.int32, - seed=distribution_util.gen_new_seed( - layer.seed, salt="dense_flipout")) - sign_output = math_ops.cast(2 * sign_output - 1, inputs.dtype) - perturbed_inputs = math_ops.matmul( - inputs * sign_input, expected_kernel_posterior_affine_tensor) - perturbed_inputs *= sign_output - - expected_outputs = math_ops.matmul(inputs, kernel_posterior.result_loc) - expected_outputs += perturbed_inputs - expected_outputs += bias_posterior.result_sample - - [ - 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 testRandomDenseFlipout(self): - batch_size, in_size, out_size = 2, 3, 4 - with self.test_session() as sess: - seed = Counter() - inputs = random_ops.random_uniform([batch_size, in_size], seed=seed()) - - kernel_posterior = MockDistribution( - loc=random_ops.random_uniform( - [in_size, out_size], seed=seed()), - scale=random_ops.random_uniform( - [in_size, out_size], seed=seed()), - result_log_prob=random_ops.random_uniform( - [in_size, out_size], seed=seed()), - result_sample=random_ops.random_uniform( - [in_size, out_size], seed=seed())) - bias_posterior = MockDistribution( - loc=random_ops.random_uniform( - [out_size], seed=seed()), - scale=random_ops.random_uniform( - [out_size], seed=seed()), - result_log_prob=random_ops.random_uniform( - [out_size], seed=seed()), - result_sample=random_ops.random_uniform( - [out_size], seed=seed())) - layer_one = prob_layers_lib.DenseFlipout( - units=out_size, - 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 = prob_layers_lib.DenseFlipout( - units=out_size, - 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_)), out_size) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/metropolis_hastings_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/metropolis_hastings_test.py index 63d93fad64d077aa385b72428665e841b6784b90..f508e5b114a55fc1aeb07212595fda45fc308c7b 100644 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/metropolis_hastings_test.py +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/metropolis_hastings_test.py @@ -12,34 +12,195 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for metropolis_hastings.py.""" +"""Tests for Metropolis-Hastings.""" 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 metropolis_hastings_impl as mh +from tensorflow.contrib.distributions.python.ops import mvn_tril as mvn_tril_lib +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 init_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.ops.distributions import normal as normal_lib from tensorflow.python.platform import test -class McmcStepTest(test.TestCase): +class MetropolisHastingsTest(test.TestCase): + + def testKernelStateTensor(self): + """Test that transition kernel works with tensor input to `state`.""" + loc = variable_scope.get_variable("loc", initializer=0.) + + def target_log_prob_fn(loc): + return normal_lib.Normal(loc=0.0, scale=0.1).log_prob(loc) + + new_state, _ = mh.kernel( + target_log_prob_fn=target_log_prob_fn, + proposal_fn=mh.proposal_normal(scale=0.05), + current_state=loc, + seed=231251) + loc_update = loc.assign(new_state) + + init = variables.initialize_all_variables() + with self.test_session() as sess: + sess.run(init) + loc_samples = [] + for _ in range(2500): + loc_sample = sess.run(loc_update) + loc_samples.append(loc_sample) + loc_samples = loc_samples[500:] # drop samples for burn-in + + self.assertAllClose(np.mean(loc_samples), 0.0, rtol=1e-5, atol=1e-1) + self.assertAllClose(np.std(loc_samples), 0.1, rtol=1e-5, atol=1e-1) + + def testKernelStateList(self): + """Test that transition kernel works with list input to `state`.""" + num_chains = 2 + loc_one = variable_scope.get_variable( + "loc_one", [num_chains], + initializer=init_ops.zeros_initializer()) + loc_two = variable_scope.get_variable( + "loc_two", [num_chains], initializer=init_ops.zeros_initializer()) + + def target_log_prob_fn(loc_one, loc_two): + loc = array_ops.stack([loc_one, loc_two]) + log_prob = mvn_tril_lib.MultivariateNormalTriL( + loc=constant_op.constant([0., 0.]), + scale_tril=constant_op.constant([[0.1, 0.1], [0.0, 0.1]])).log_prob( + loc) + return math_ops.reduce_sum(log_prob, 0) + + def proposal_fn(loc_one, loc_two): + loc_one_proposal = mh.proposal_normal(scale=0.05) + loc_two_proposal = mh.proposal_normal(scale=0.05) + loc_one_sample, _ = loc_one_proposal(loc_one) + loc_two_sample, _ = loc_two_proposal(loc_two) + return [loc_one_sample, loc_two_sample], None + + new_state, _ = mh.kernel( + target_log_prob_fn=target_log_prob_fn, + proposal_fn=proposal_fn, + current_state=[loc_one, loc_two], + seed=12415) + loc_one_update = loc_one.assign(new_state[0]) + loc_two_update = loc_two.assign(new_state[1]) + + init = variables.initialize_all_variables() + with self.test_session() as sess: + sess.run(init) + loc_one_samples = [] + loc_two_samples = [] + for _ in range(10000): + loc_one_sample, loc_two_sample = sess.run( + [loc_one_update, loc_two_update]) + loc_one_samples.append(loc_one_sample) + loc_two_samples.append(loc_two_sample) + + loc_one_samples = np.array(loc_one_samples) + loc_two_samples = np.array(loc_two_samples) + loc_one_samples = loc_one_samples[1000:] # drop samples for burn-in + loc_two_samples = loc_two_samples[1000:] # drop samples for burn-in + + self.assertAllClose(np.mean(loc_one_samples, 0), + np.array([0.] * num_chains), + rtol=1e-5, atol=1e-1) + self.assertAllClose(np.mean(loc_two_samples, 0), + np.array([0.] * num_chains), + rtol=1e-5, atol=1e-1) + self.assertAllClose(np.std(loc_one_samples, 0), + np.array([0.1] * num_chains), + rtol=1e-5, atol=1e-1) + self.assertAllClose(np.std(loc_two_samples, 0), + np.array([0.1] * num_chains), + rtol=1e-5, atol=1e-1) + + def testKernelResultsUsingTruncatedDistribution(self): + def log_prob(x): + return array_ops.where( + x >= 0., + -x - x**2, + array_ops.fill(x.shape, math_ops.cast(-np.inf, x.dtype))) + # The truncated distribution has the property that it is likely to attract + # the 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. + + num_results = 1000 + # Large step size, will give rejections due to going into a region of + # log_prob = -inf. + step_size = 0.3 + 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) - def test_density_increasing_step_accepted(self): + states = [] + is_accepted = [] + proposed_states = [] + current_state = initial_state + for _ in range(num_results): + current_state, kernel_results = mh.kernel( + target_log_prob_fn=log_prob, + proposal_fn=mh.proposal_uniform(step_size=step_size), + current_state=current_state, + seed=42) + states.append(current_state) + proposed_states.append(kernel_results.proposed_state) + is_accepted.append(kernel_results.is_accepted) + + states = array_ops.stack(states) + proposed_states = array_ops.stack(proposed_states) + is_accepted = array_ops.stack(is_accepted) + states_, pstates_, is_accepted_ = sess.run( + [states, proposed_states, is_accepted]) + + # We better have accepted a decent amount, even near end of the chain. + self.assertLess( + 0.1, is_accepted_[int(0.9 * num_results):].mean()) + # We better not have any NaNs in states. + self.assertAllEqual(np.zeros_like(states_), + np.isnan(states_)) + # We better not have any +inf in states. + self.assertAllEqual(np.zeros_like(states_), + np.isposinf(states_)) + + # The move is accepted ==> state = proposed state. + self.assertAllEqual( + states_[is_accepted_], + pstates_[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 is_accepted_[t, i]: + self.assertNotEqual(states_[t, i], states_[t - 1, i]) + else: + self.assertEqual(states_[t, i], states_[t - 1, i]) + + def testDensityIncreasingStepAccepted(self): """Tests that if a transition increases density, it is always accepted.""" target_log_density = lambda x: - x * x - state = variable_scope.get_variable('state', initializer=10.) + state = variable_scope.get_variable("state", initializer=10.) state_log_density = variable_scope.get_variable( - 'state_log_density', + "state_log_density", initializer=target_log_density(state.initialized_value())) log_accept_ratio = variable_scope.get_variable( - 'log_accept_ratio', initializer=0.) + "log_accept_ratio", initializer=0.) get_next_proposal = lambda x: (x - 1., None) step = mh.evolve(state, state_log_density, log_accept_ratio, @@ -54,7 +215,7 @@ class McmcStepTest(test.TestCase): self.assertAlmostEqual(sample, 9 - j) self.assertAlmostEqual(sample_log_density, - (9 - j) * (9 - j)) - def test_sample_properties(self): + def testSampleProperties(self): """Tests that the samples converge to the target distribution.""" def target_log_density(x): @@ -62,16 +223,16 @@ class McmcStepTest(test.TestCase): return - (x - 2.0) * (x - 2.0) * 0.5 # Use the uniform random walker to generate proposals. - proposal_fn = mh.uniform_random_proposal( + proposal_fn = mh.proposal_uniform( step_size=1.0, seed=1234) - state = variable_scope.get_variable('state', initializer=0.0) + state = variable_scope.get_variable("state", initializer=0.0) state_log_density = variable_scope.get_variable( - 'state_log_density', + "state_log_density", initializer=target_log_density(state.initialized_value())) - log_accept_ratio = variable_scope.get_variable( - 'log_accept_ratio', initializer=0.) + "log_accept_ratio", initializer=0.) + # Random walk MCMC converges slowly so need to put in enough iterations. num_iterations = 5000 step = mh.evolve(state, state_log_density, log_accept_ratio, @@ -98,11 +259,11 @@ class McmcStepTest(test.TestCase): self.assertAlmostEqual(sample_mean, 2.0, delta=0.1) self.assertAlmostEqual(sample_variance, 1.0, delta=0.1) - def test_normal_proposals(self): + def testProposalNormal(self): """Tests that the normal proposals are correctly distributed.""" initial_points = array_ops.ones([10000], dtype=dtypes.float32) - proposal_fn = mh.normal_random_proposal( + proposal_fn = mh.proposal_normal( scale=2.0, seed=1234) proposal_points, _ = proposal_fn(initial_points) @@ -115,7 +276,7 @@ class McmcStepTest(test.TestCase): self.assertAlmostEqual(np.mean(sample), 1.0, delta=0.1) self.assertAlmostEqual(np.std(sample), 2.0, delta=0.1) - def test_docstring_example(self): + def testDocstringExample(self): """Tests the simplified docstring example with multiple chains.""" n = 2 # dimension of the problem @@ -123,7 +284,7 @@ class McmcStepTest(test.TestCase): # Generate 300 initial values randomly. Each of these would be an # independent starting point for a Markov chain. state = variable_scope.get_variable( - 'state', initializer=random_ops.random_normal( + "state", initializer=random_ops.random_normal( [300, n], mean=3.0, dtype=dtypes.float32, seed=42)) # Computes the log(p(x)) for the unit normal density and ignores the @@ -133,12 +294,12 @@ class McmcStepTest(test.TestCase): # Initial log-density value state_log_density = variable_scope.get_variable( - 'state_log_density', + "state_log_density", initializer=log_density(state.initialized_value())) # A variable to store the log_acceptance_ratio: log_acceptance_ratio = variable_scope.get_variable( - 'log_acceptance_ratio', + "log_acceptance_ratio", initializer=array_ops.zeros([300], dtype=dtypes.float32)) # Generates random proposals by moving each coordinate uniformly and @@ -175,5 +336,5 @@ class McmcStepTest(test.TestCase): - np.reshape(covariance, [n**2]))), 0, delta=0.2) -if __name__ == '__main__': +if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/sgld_optimizer_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/sgld_optimizer_test.py deleted file mode 100644 index 756c25683bd4b0c8c77e9e28485ca2a85582999c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/sgld_optimizer_test.py +++ /dev/null @@ -1,212 +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. -# ============================================================================== -"""Functional test for GradientDescent.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -import math -from tensorflow.contrib.bayesflow.python.ops.optimizers import SGLDOptimizer -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.ops import variables -from tensorflow.python.platform import test - - -class SGLDOptimizerTest(test.TestCase): - - def testBasic(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - decay_rate = 0.53 - sgd_optimizer = SGLDOptimizer(3.0, preconditioner_decay_rate=decay_rate) - sgd_op = sgd_optimizer.apply_gradients( - zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - # Validate updated params - grads_scaled = (0.5 * 0.1 / math.sqrt(decay_rate + - (1 - decay_rate) * 0.1**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [1.1 - 3.0 * grads_scaled, 2.1 - 3.0 * grads_scaled], var0.eval()) - grads_scaled = (0.5 * 0.01 / math.sqrt( - decay_rate + (1 - decay_rate) * 0.01**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [3.0 - 3.0 * grads_scaled, 4.0 - 3.0 * grads_scaled], var1.eval()) - self.assertAllCloseAccordingToType(1, sgd_optimizer._counter.eval()) - - def testBasicMultiInstance(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - vara = variables.Variable([1.1, 2.1], dtype=dtype) - varb = variables.Variable([3.0, 4.0], dtype=dtype) - gradsa = constant_op.constant([0.1, 0.1], dtype=dtype) - gradsb = constant_op.constant([0.01, 0.01], dtype=dtype) - decay_rate = 0.5 - sgd_optimizer = SGLDOptimizer(3.0, preconditioner_decay_rate=decay_rate) - sgd_op = sgd_optimizer.apply_gradients( - zip([grads0, grads1], [var0, var1])) - sgd_optimizer2 = SGLDOptimizer( - 3.0, preconditioner_decay_rate=decay_rate) - sgd_op2 = sgd_optimizer2.apply_gradients( - zip([gradsa, gradsb], [vara, varb])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - self.assertAllCloseAccordingToType([1.1, 2.1], vara.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], varb.eval()) - - # Run 1 step of sgd - sgd_op.run() - sgd_op2.run() - # Validate updated params - grads_scaled = (0.5 * 0.1 / math.sqrt(decay_rate + - (1 - decay_rate) * 0.1**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [1.1 - 3.0 * grads_scaled, 2.1 - 3.0 * grads_scaled], var0.eval()) - self.assertAllCloseAccordingToType( - [1.1 - 3.0 * grads_scaled, 2.1 - 3.0 * grads_scaled], vara.eval()) - - grads_scaled = (0.5 * 0.01 / math.sqrt( - decay_rate + (1 - decay_rate) * 0.01**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [3.0 - 3.0 * grads_scaled, 4.0 - 3.0 * grads_scaled], var1.eval()) - self.assertAllCloseAccordingToType( - [3.0 - 3.0 * grads_scaled, 4.0 - 3.0 * grads_scaled], varb.eval()) - self.assertNotEqual(sgd_optimizer.variable_scope, - sgd_optimizer2.variable_scope) - self.assertNotEqual(sgd_optimizer.variable_scope.name, - sgd_optimizer2.variable_scope.name) - self.assertAllCloseAccordingToType(1, sgd_optimizer._counter.eval()) - self.assertAllCloseAccordingToType(1, sgd_optimizer2._counter.eval()) - - def testTensorLearningRate(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - lrate = constant_op.constant(3.0) - decay_rate = 0.5 - sgd_op = SGLDOptimizer( - lrate, preconditioner_decay_rate=constant_op.constant( - decay_rate)).apply_gradients( - zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - # Validate updated params - grads_scaled = (0.5 * 0.1 / math.sqrt(decay_rate + - (1 - decay_rate) * 0.1**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [1.1 - 3.0 * grads_scaled, 2.1 - 3.0 * grads_scaled], var0.eval()) - grads_scaled = (0.5 * 0.01 / math.sqrt( - decay_rate + (1 - decay_rate) * 0.01**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [3.0 - 3.0 * grads_scaled, 4.0 - 3.0 * grads_scaled], var1.eval()) - - def testGradWrtRef(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - opt = SGLDOptimizer(3.0) - values = [1.0, 3.0] - vars_ = [variables.Variable([v], dtype=dtype) for v in values] - grads_and_vars = opt.compute_gradients(vars_[0] + vars_[1], vars_) - variables.global_variables_initializer().run() - for grad, _ in grads_and_vars: - self.assertAllCloseAccordingToType([1.0], grad.eval()) - - def testWithGlobalStep(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - global_step = variables.Variable(0, trainable=False) - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - decay_rate = 0.1 - sgd_op = SGLDOptimizer( - 3.0, preconditioner_decay_rate=decay_rate).apply_gradients( - zip([grads0, grads1], [var0, var1]), global_step=global_step) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - - # Validate updated params and global_step - grads_scaled = (0.5 * 0.1 / math.sqrt(decay_rate + - (1 - decay_rate) * 0.1**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [1.1 - 3.0 * grads_scaled, 2.1 - 3.0 * grads_scaled], var0.eval()) - grads_scaled = (0.5 * 0.01 / math.sqrt( - decay_rate + (1 - decay_rate) * 0.01**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [3.0 - 3.0 * grads_scaled, 4.0 - 3.0 * grads_scaled], var1.eval()) - self.assertAllCloseAccordingToType(1, global_step.eval()) - - def testSparseBasic(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([[1.1], [2.1]], dtype=dtype) - var1 = variables.Variable([[3.0], [4.0]], dtype=dtype) - grads0 = ops.IndexedSlices( - constant_op.constant([0.1], shape=[1, 1], dtype=dtype), - constant_op.constant([0]), constant_op.constant([2, 1])) - grads1 = ops.IndexedSlices( - constant_op.constant([0.01], shape=[1, 1], dtype=dtype), - constant_op.constant([1]), constant_op.constant([2, 1])) - decay_rate = 0.9 - sgd_op = SGLDOptimizer( - 3.0, preconditioner_decay_rate=decay_rate).apply_gradients( - zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([[1.1], [2.1]], var0.eval()) - self.assertAllCloseAccordingToType([[3.0], [4.0]], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - # Validate updated params - grads_scaled = (0.5 * 0.1 / math.sqrt(decay_rate + - (1 - decay_rate) * 0.1**2 + 1e-8)) - self.assertAllCloseAccordingToType([[1.1 - 3.0 * grads_scaled], [2.1]], - var0.eval()) - grads_scaled = (0.5 * 0.01 / math.sqrt( - decay_rate + (1 - decay_rate) * 0.01**2 + 1e-8)) - self.assertAllCloseAccordingToType( - [[3.0 - 3.0 * 0], [4.0 - 3.0 * grads_scaled]], var1.eval()) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/variational_sgd_optimizer_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/variational_sgd_optimizer_test.py deleted file mode 100644 index 83c64dbe0fd586edcb784a5c09a4c133aaa99cff..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/variational_sgd_optimizer_test.py +++ /dev/null @@ -1,268 +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. -# ============================================================================== -"""Functional test for GradientDescent.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from tensorflow.contrib.bayesflow.python.ops.optimizers import VariationalSGDOptimizer -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 variables -from tensorflow.python.platform import test - - -class VariationalSGDOptimizerTest(test.TestCase): - - def testBasic(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - decay_rate = 0.53 - sgd_op = VariationalSGDOptimizer( - 1, - 1, - preconditioner_decay_rate=decay_rate, - max_learning_rate=3.0, - burnin_max_learning_rate=3.0, - use_single_learning_rate=True).apply_gradients( - zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - self.assertAllCloseAccordingToType([1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1], - var0.eval()) - self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01], - var1.eval()) - - def testBasicMultiInstance(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - vara = variables.Variable([1.1, 2.1], dtype=dtype) - varb = variables.Variable([3.0, 4.0], dtype=dtype) - gradsa = constant_op.constant([0.1, 0.1], dtype=dtype) - gradsb = constant_op.constant([0.01, 0.01], dtype=dtype) - decay_rate = 0.5 - batch_size = 2 - total_num_examples = 10 - optimizer = VariationalSGDOptimizer( - batch_size, - total_num_examples, - max_learning_rate=1.0, - burnin_max_learning_rate=3.0, - preconditioner_decay_rate=decay_rate) - sgd_op = optimizer.apply_gradients( - zip([grads0, grads1], [var0, var1])) - optimizer2 = VariationalSGDOptimizer( - batch_size, - total_num_examples, - max_learning_rate=1.0, - burnin_max_learning_rate=10.0, - burnin=0, - preconditioner_decay_rate=decay_rate) - sgd_op2 = optimizer2.apply_gradients( - zip([gradsa, gradsb], [vara, varb])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - self.assertAllCloseAccordingToType([1.1, 2.1], vara.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], varb.eval()) - - # Run 1 step of sgd - sgd_op.run() - sgd_op2.run() - # Validate updated params - self.assertAllCloseAccordingToType([1.1 - 3. * 0.1, 2.1 - 3. * 0.1], - var0.eval()) - self.assertAllCloseAccordingToType([1.1 - 0.1, 2.1 - 0.1], vara.eval()) - - self.assertAllCloseAccordingToType([3.0 - 3. * 0.01, 4.0 - 3. * 0.01], - var1.eval()) - self.assertAllCloseAccordingToType([3.0 - 0.01, 4.0 - 0.01], - varb.eval()) - self.assertNotEqual(optimizer.variable_scope, - optimizer2.variable_scope) - self.assertNotEqual(optimizer.variable_scope.name, - optimizer2.variable_scope.name) - self.assertAllCloseAccordingToType(1, optimizer._counter.eval()) - self.assertAllCloseAccordingToType(1, optimizer2._counter.eval()) - - def testTensorLearningRate(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - lrate = constant_op.constant(3.0) - decay_rate = 0.5 - batch_size = 2 - total_num_examples = 10 - sgd_op = VariationalSGDOptimizer( - batch_size, - total_num_examples, - max_learning_rate=lrate, - burnin=0, - preconditioner_decay_rate=decay_rate).apply_gradients( - zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - # Validate updated params - self.assertAllCloseAccordingToType([1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1], - var0.eval()) - self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01], - var1.eval()) - - def testTensorDecayLearningRate(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - lrate = variables.Variable(3.0) - lrate_decay_op = lrate.assign_add(-3.) - decay_rate = 0.5 - batch_size = 2 - total_num_examples = 10 - optimizer = VariationalSGDOptimizer( - batch_size, - total_num_examples, - max_learning_rate=lrate, - burnin=0, - preconditioner_decay_rate=decay_rate) - sgd_op = optimizer.apply_gradients(zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - # Validate updated params - self.assertAllCloseAccordingToType([1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1], - var0.eval()) - self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01], - var1.eval()) - # Update learning rate to 0 - lrate_decay_op.eval() - sgd_op.run() - # Validate params haven't changed - self.assertAllCloseAccordingToType([1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1], - var0.eval()) - self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01], - var1.eval()) - lrate_decay_op.eval() - - with self.assertRaises(errors.InvalidArgumentError): - sgd_op.run() - - def testGradWrtRef(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - opt = VariationalSGDOptimizer(1, 1, max_learning_rate=1.0) - values = [1.0, 3.0] - vars_ = [variables.Variable([v], dtype=dtype) for v in values] - grads_and_vars = opt.compute_gradients(vars_[0] + vars_[1], vars_) - variables.global_variables_initializer().run() - for grad, _ in grads_and_vars: - self.assertAllCloseAccordingToType([1.0], grad.eval()) - - def testWithGlobalStep(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - global_step = variables.Variable(0, trainable=False) - var0 = variables.Variable([1.1, 2.1], dtype=dtype) - var1 = variables.Variable([3.0, 4.0], dtype=dtype) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - decay_rate = 0.1 - batch_size = 2 - total_num_examples = 10 - sgd_optimizer = VariationalSGDOptimizer( - batch_size, - total_num_examples, - max_learning_rate=3.0, - burnin=0, - preconditioner_decay_rate=decay_rate) - sgd_op = sgd_optimizer.apply_gradients( - zip([grads0, grads1], [var0, var1]), global_step=global_step) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.1, 2.1], var0.eval()) - self.assertAllCloseAccordingToType([3.0, 4.0], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - - # Validate updated params and global_step - self.assertAllCloseAccordingToType([1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1], - var0.eval()) - self.assertAllCloseAccordingToType([3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01], - var1.eval()) - self.assertAllCloseAccordingToType(1, global_step.eval()) - self.assertAllCloseAccordingToType(1, sgd_optimizer._counter.eval()) - - def testSparseBasic(self): - for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session(): - var0 = variables.Variable([[1.1], [2.1]], dtype=dtype) - var1 = variables.Variable([[3.0], [4.0]], dtype=dtype) - grads0 = ops.IndexedSlices( - constant_op.constant([0.1], shape=[1, 1], dtype=dtype), - constant_op.constant([0]), constant_op.constant([2, 1])) - grads1 = ops.IndexedSlices( - constant_op.constant([0.01], shape=[1, 1], dtype=dtype), - constant_op.constant([1]), constant_op.constant([2, 1])) - decay_rate = 0.1 - batch_size = 2 - total_num_examples = 10 - sgd_op = VariationalSGDOptimizer( - batch_size, - total_num_examples, - max_learning_rate=3.0, - burnin=0, - preconditioner_decay_rate=decay_rate).apply_gradients( - zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() - # Fetch params to validate initial values - self.assertAllCloseAccordingToType([[1.1], [2.1]], var0.eval()) - self.assertAllCloseAccordingToType([[3.0], [4.0]], var1.eval()) - # Run 1 step of sgd - sgd_op.run() - # Validate updated params - self.assertAllCloseAccordingToType([[1.1 - 3.0 * 0.1], [2.1]], - var0.eval()) - self.assertAllCloseAccordingToType( - [[3.0 - 3.0 * 0], [4.0 - 3.0 * 0.01]], var1.eval()) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/bayesflow/python/ops/csiszar_divergence_impl.py b/tensorflow/contrib/bayesflow/python/ops/csiszar_divergence_impl.py deleted file mode 100644 index 8efd59d6516924bea538717d45bb4ae303583421..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/csiszar_divergence_impl.py +++ /dev/null @@ -1,1105 +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. -# ============================================================================== -"""Csiszar f-Divergence and helpers. - -@@amari_alpha -@@arithmetic_geometric -@@chi_square -@@csiszar_vimco -@@dual_csiszar_function -@@jeffreys -@@jensen_shannon -@@kl_forward -@@kl_reverse -@@log1p_abs -@@modified_gan -@@monte_carlo_csiszar_f_divergence -@@pearson -@@squared_hellinger -@@symmetrized_csiszar_function -@@total_variation -@@triangular - -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib import framework as contrib_framework -from tensorflow.contrib.bayesflow.python.ops import monte_carlo_impl as monte_carlo -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_ops -from tensorflow.python.ops.distributions import distribution -from tensorflow.python.ops.distributions import util as distribution_util - - -def amari_alpha(logu, alpha=1., self_normalized=False, name=None): - """The Amari-alpha Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True`, the Amari-alpha Csiszar-function is: - - ```none - f(u) = { -log(u) + (u - 1), alpha = 0 - { u log(u) - (u - 1), alpha = 1 - { [(u**alpha - 1) - alpha (u - 1)] / (alpha (alpha - 1)), otherwise - ``` - - When `self_normalized = False` the `(u - 1)` terms are omitted. - - Warning: when `alpha != 0` and/or `self_normalized = True` this function makes - non-log-space calculations and may therefore be numerically unstable for - `|logu| >> 0`. - - For more information, see: - A. Cichocki and S. Amari. "Families of Alpha-Beta-and GammaDivergences: - Flexible and Robust Measures of Similarities." Entropy, vol. 12, no. 6, pp. - 1532-1568, 2010. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - alpha: `float`-like Python scalar. (See Mathematical Details for meaning.) - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When - `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even - when `p, q` are unnormalized measures. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - amari_alpha_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - - Raises: - TypeError: if `alpha` is `None` or a `Tensor`. - TypeError: if `self_normalized` is `None` or a `Tensor`. - """ - with ops.name_scope(name, "amari_alpha", [logu]): - if alpha is None or contrib_framework.is_tensor(alpha): - raise TypeError("`alpha` cannot be `None` or `Tensor` type.") - if self_normalized is None or contrib_framework.is_tensor(self_normalized): - raise TypeError("`self_normalized` cannot be `None` or `Tensor` type.") - - logu = ops.convert_to_tensor(logu, name="logu") - - if alpha == 0.: - f = -logu - elif alpha == 1.: - f = math_ops.exp(logu) * logu - else: - f = math_ops.expm1(alpha * logu) / (alpha * (alpha - 1.)) - - if not self_normalized: - return f - - if alpha == 0.: - return f + math_ops.expm1(logu) - elif alpha == 1.: - return f - math_ops.expm1(logu) - else: - return f - math_ops.expm1(logu) / (alpha - 1.) - - -def kl_reverse(logu, self_normalized=False, name=None): - """The reverse Kullback-Leibler Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True`, the KL-reverse Csiszar-function is: - - ```none - f(u) = -log(u) + (u - 1) - ``` - - When `self_normalized = False` the `(u - 1)` term is omitted. - - Observe that as an f-Divergence, this Csiszar-function implies: - - ```none - D_f[p, q] = KL[q, p] - ``` - - The KL is "reverse" because in maximum likelihood we think of minimizing `q` - as in `KL[p, q]`. - - Warning: when self_normalized = True` this function makes non-log-space - calculations and may therefore be numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When - `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even - when `p, q` are unnormalized measures. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - kl_reverse_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at - `u = exp(logu)`. - - Raises: - TypeError: if `self_normalized` is `None` or a `Tensor`. - """ - - with ops.name_scope(name, "kl_reverse", [logu]): - return amari_alpha(logu, alpha=0., self_normalized=self_normalized) - - -def kl_forward(logu, self_normalized=False, name=None): - """The forward Kullback-Leibler Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True`, the KL-forward Csiszar-function is: - - ```none - f(u) = u log(u) - (u - 1) - ``` - - When `self_normalized = False` the `(u - 1)` term is omitted. - - Observe that as an f-Divergence, this Csiszar-function implies: - - ```none - D_f[p, q] = KL[p, q] - ``` - - The KL is "forward" because in maximum likelihood we think of minimizing `q` - as in `KL[p, q]`. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When - `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even - when `p, q` are unnormalized measures. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - kl_forward_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at - `u = exp(logu)`. - - Raises: - TypeError: if `self_normalized` is `None` or a `Tensor`. - """ - - with ops.name_scope(name, "kl_forward", [logu]): - return amari_alpha(logu, alpha=1., self_normalized=self_normalized) - - -def jensen_shannon(logu, self_normalized=False, name=None): - """The Jensen-Shannon Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True`, the Jensen-Shannon Csiszar-function is: - - ```none - f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2) - ``` - - When `self_normalized = False` the `(u + 1) log(2)` term is omitted. - - Observe that as an f-Divergence, this Csiszar-function implies: - - ```none - D_f[p, q] = KL[p, m] + KL[q, m] - m(x) = 0.5 p(x) + 0.5 q(x) - ``` - - In a sense, this divergence is the "reverse" of the Arithmetic-Geometric - f-Divergence. - - This Csiszar-function induces a symmetric f-Divergence, i.e., - `D_f[p, q] = D_f[q, p]`. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - For more information, see: - Lin, J. "Divergence measures based on the Shannon entropy." IEEE Trans. - Inf. Th., 37, 145-151, 1991. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When - `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even - when `p, q` are unnormalized measures. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - jensen_shannon_of_u: `float`-like `Tensor` of the Csiszar-function - evaluated at `u = exp(logu)`. - """ - - with ops.name_scope(name, "jensen_shannon", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - npdt = logu.dtype.as_numpy_dtype - y = nn_ops.softplus(logu) - if self_normalized: - y -= np.log(2).astype(npdt) - return math_ops.exp(logu) * logu - (1. + math_ops.exp(logu)) * y - - -def arithmetic_geometric(logu, self_normalized=False, name=None): - """The Arithmetic-Geometric Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True` the Arithmetic-Geometric Csiszar-function is: - - ```none - f(u) = (1 + u) log( (1 + u) / sqrt(u) ) - (1 + u) log(2) - ``` - - When `self_normalized = False` the `(1 + u) log(2)` term is omitted. - - Observe that as an f-Divergence, this Csiszar-function implies: - - ```none - D_f[p, q] = KL[m, p] + KL[m, q] - m(x) = 0.5 p(x) + 0.5 q(x) - ``` - - In a sense, this divergence is the "reverse" of the Jensen-Shannon - f-Divergence. - - This Csiszar-function induces a symmetric f-Divergence, i.e., - `D_f[p, q] = D_f[q, p]`. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When - `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even - when `p, q` are unnormalized measures. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - arithmetic_geometric_of_u: `float`-like `Tensor` of the - Csiszar-function evaluated at `u = exp(logu)`. - """ - - with ops.name_scope(name, "arithmetic_geometric", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - y = nn_ops.softplus(logu) - 0.5 * logu - if self_normalized: - y -= np.log(2.).astype(logu.dtype.as_numpy_dtype) - return (1. + math_ops.exp(logu)) * y - - -def total_variation(logu, name=None): - """The Total Variation Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Total-Variation Csiszar-function is: - - ```none - f(u) = 0.5 |u - 1| - ``` - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - total_variation_of_u: `float`-like `Tensor` of the Csiszar-function - evaluated at `u = exp(logu)`. - """ - - with ops.name_scope(name, "total_variation", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return 0.5 * math_ops.abs(math_ops.expm1(logu)) - - -def pearson(logu, name=None): - """The Pearson Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Pearson Csiszar-function is: - - ```none - f(u) = (u - 1)**2 - ``` - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - pearson_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at - `u = exp(logu)`. - """ - - with ops.name_scope(name, "pearson", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return math_ops.square(math_ops.expm1(logu)) - - -def squared_hellinger(logu, name=None): - """The Squared-Hellinger Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Squared-Hellinger Csiszar-function is: - - ```none - f(u) = (sqrt(u) - 1)**2 - ``` - - This Csiszar-function induces a symmetric f-Divergence, i.e., - `D_f[p, q] = D_f[q, p]`. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - squared_hellinger_of_u: `float`-like `Tensor` of the Csiszar-function - evaluated at `u = exp(logu)`. - """ - - with ops.name_scope(name, "squared_hellinger", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return pearson(0.5 * logu) - - -def triangular(logu, name=None): - """The Triangular Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Triangular Csiszar-function is: - - ```none - f(u) = (u - 1)**2 / (1 + u) - ``` - - This Csiszar-function induces a symmetric f-Divergence, i.e., - `D_f[p, q] = D_f[q, p]`. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - triangular_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - """ - - with ops.name_scope(name, "triangular", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return pearson(logu) / (1. + math_ops.exp(logu)) - - -def t_power(logu, t, self_normalized=False, name=None): - """The T-Power Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True` the T-Power Csiszar-function is: - - ```none - f(u) = s [ u**t - 1 - t(u - 1) ] - s = { -1 0 < t < 1 - { +1 otherwise - ``` - - When `self_normalized = False` the `- t(u - 1)` term is omitted. - - This is similar to the `amari_alpha` Csiszar-function, with the associated - divergence being the same up to factors depending only on `t`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - t: `Tensor` of same `dtype` as `logu` and broadcastable shape. - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - t_power_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - """ - with ops.name_scope(name, "t_power", [logu, t]): - logu = ops.convert_to_tensor(logu, name="logu") - t = ops.convert_to_tensor(t, dtype=logu.dtype.base_dtype, name="t") - fu = math_ops.expm1(t * logu) - if self_normalized: - fu -= t * math_ops.expm1(logu) - fu *= array_ops.where(math_ops.logical_and(0. < t, t < 1.), - -array_ops.ones_like(t), - array_ops.ones_like(t)) - return fu - - -def log1p_abs(logu, name=None): - """The log1p-abs Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Log1p-Abs Csiszar-function is: - - ```none - f(u) = u**(sign(u-1)) - 1 - ``` - - This function is so-named because it was invented from the following recipe. - Choose a convex function g such that g(0)=0 and solve for f: - - ```none - log(1 + f(u)) = g(log(u)). - <=> - f(u) = exp(g(log(u))) - 1 - ``` - - That is, the graph is identically `g` when y-axis is `log1p`-domain and x-axis - is `log`-domain. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - log1p_abs_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - """ - - with ops.name_scope(name, "log1p_abs", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return math_ops.expm1(math_ops.abs(logu)) - - -def jeffreys(logu, name=None): - """The Jeffreys Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Jeffreys Csiszar-function is: - - ```none - f(u) = 0.5 ( u log(u) - log(u) ) - = 0.5 kl_forward + 0.5 kl_reverse - = symmetrized_csiszar_function(kl_reverse) - = symmetrized_csiszar_function(kl_forward) - ``` - - This Csiszar-function induces a symmetric f-Divergence, i.e., - `D_f[p, q] = D_f[q, p]`. - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - jeffreys_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - """ - - with ops.name_scope(name, "jeffreys", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return 0.5 * math_ops.expm1(logu) * logu - - -def chi_square(logu, name=None): - """The chi-Square Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Chi-square Csiszar-function is: - - ```none - f(u) = u**2 - 1 - ``` - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - """ - - with ops.name_scope(name, "chi_square", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return math_ops.expm1(2. * logu) - - -def modified_gan(logu, self_normalized=False, name=None): - """The Modified-GAN Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - When `self_normalized = True` the modified-GAN (Generative/Adversarial - Network) Csiszar-function is: - - ```none - f(u) = log(1 + u) - log(u) + 0.5 (u - 1) - ``` - - When `self_normalized = False` the `0.5 (u - 1)` is omitted. - - The unmodified GAN Csiszar-function is identical to Jensen-Shannon (with - `self_normalized = False`). - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When - `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even - when `p, q` are unnormalized measures. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated - at `u = exp(logu)`. - """ - - with ops.name_scope(name, "chi_square", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - y = nn_ops.softplus(logu) - logu - if self_normalized: - y += 0.5 * math_ops.expm1(logu) - return y - - -def dual_csiszar_function(logu, csiszar_function, name=None): - """Calculates the dual Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Csiszar-dual is defined as: - - ```none - f^*(u) = u f(1 / u) - ``` - - where `f` is some other Csiszar-function. - - For example, the dual of `kl_reverse` is `kl_forward`, i.e., - - ```none - f(u) = -log(u) - f^*(u) = u f(1 / u) = -u log(1 / u) = u log(u) - ``` - - The dual of the dual is the original function: - - ```none - f^**(u) = {u f(1/u)}^*(u) = u (1/u) f(1/(1/u)) = f(u) - ``` - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - csiszar_function: Python `callable` representing a Csiszar-function over - log-domain. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - dual_f_of_u: `float`-like `Tensor` of the result of calculating the dual of - `f` at `u = exp(logu)`. - """ - - with ops.name_scope(name, "dual_csiszar_function", [logu]): - return math_ops.exp(logu) * csiszar_function(-logu) - - -def symmetrized_csiszar_function(logu, csiszar_function, name=None): - """Symmetrizes a Csiszar-function in log-space. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The symmetrized Csiszar-function is defined as: - - ```none - f_g(u) = 0.5 g(u) + 0.5 u g (1 / u) - ``` - - where `g` is some other Csiszar-function. - - We say the function is "symmetrized" because: - - ```none - D_{f_g}[p, q] = D_{f_g}[q, p] - ``` - - for all `p << >> q` (i.e., `support(p) = support(q)`). - - There exists alternatives for symmetrizing a Csiszar-function. For example, - - ```none - f_g(u) = max(f(u), f^*(u)), - ``` - - where `f^*` is the dual Csiszar-function, also implies a symmetric - f-Divergence. - - Example: - - When either of the following functions are symmetrized, we obtain the - Jensen-Shannon Csiszar-function, i.e., - - ```none - g(u) = -log(u) - (1 + u) log((1 + u) / 2) + u - 1 - h(u) = log(4) + 2 u log(u / (1 + u)) - ``` - - implies, - - ```none - f_g(u) = f_h(u) = u log(u) - (1 + u) log((1 + u) / 2) - = jensen_shannon(log(u)). - ``` - - Warning: this function makes non-log-space calculations and may therefore be - numerically unstable for `|logu| >> 0`. - - Args: - logu: `float`-like `Tensor` representing `log(u)` from above. - csiszar_function: Python `callable` representing a Csiszar-function over - log-domain. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - symmetrized_g_of_u: `float`-like `Tensor` of the result of applying the - symmetrization of `g` evaluated at `u = exp(logu)`. - """ - - with ops.name_scope(name, "symmetrized_csiszar_function", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - return 0.5 * (csiszar_function(logu) - + dual_csiszar_function(logu, csiszar_function)) - - -def monte_carlo_csiszar_f_divergence( - f, - p_log_prob, - q, - num_draws, - use_reparametrization=None, - seed=None, - name=None): - """Monte-Carlo approximation of the Csiszar f-Divergence. - - A Csiszar-function is a member of, - - ```none - F = { f:R_+ to R : f convex }. - ``` - - The Csiszar f-Divergence for Csiszar-function f is given by: - - ```none - D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ] - ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ), - where x_j ~iid q(X) - ``` - - Tricks: Reparameterization and Score-Gradient - - When q is "reparameterized", i.e., a diffeomorphic transformation of a - parameterless distribution (e.g., - `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and - expectation, i.e., - `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}` - and `s_i = f(x_i), x_i ~iid q(X)`. - - However, if q is not reparameterized, TensorFlow's gradient will be incorrect - since the chain-rule stops at samples of unreparameterized distributions. In - this circumstance using the Score-Gradient trick results in an unbiased - gradient, i.e., - - ```none - grad[ E_q[f(X)] ] - = grad[ int dx q(x) f(x) ] - = int dx grad[ q(x) f(x) ] - = int dx [ q'(x) f(x) + q(x) f'(x) ] - = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ] - = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ] - = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ] - ``` - - Unless `q.reparameterization_type != distribution.FULLY_REPARAMETERIZED` it is - usually preferable to set `use_reparametrization = True`. - - Example Application: - - The Csiszar f-Divergence is a useful framework for variational inference. - I.e., observe that, - - ```none - f(p(x)) = f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] ) - <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ] - := D_f[p(x, Z), q(Z | x)] - ``` - - The inequality follows from the fact that the "perspective" of `f`, i.e., - `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and - `t` is a real. Since the above framework includes the popular Evidence Lower - BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework - "Evidence Divergence Bound Optimization" (EDBO). - - Args: - f: Python `callable` representing a Csiszar-function in log-space, i.e., - takes `p_log_prob(q_samples) - q.log_prob(q_samples)`. - p_log_prob: Python `callable` taking (a batch of) samples from `q` and - returning the natural-log of the probability under distribution `p`. - (In variational inference `p` is the joint distribution.) - q: `tf.Distribution`-like instance; must implement: - `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`. - (In variational inference `q` is the approximate posterior distribution.) - num_draws: Integer scalar number of draws used to approximate the - f-Divergence expectation. - use_reparametrization: Python `bool`. When `None` (the default), - automatically set to: - `q.reparameterization_type == distribution.FULLY_REPARAMETERIZED`. - When `True` uses the standard Monte-Carlo average. When `False` uses the - score-gradient trick. (See above for details.) When `False`, consider - using `csiszar_vimco`. - seed: Python `int` seed for `q.sample`. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo - approximation of the Csiszar f-Divergence. - - Raises: - ValueError: if `q` is not a reparameterized distribution and - `use_reparametrization = True`. A distribution `q` is said to be - "reparameterized" when its samples are generated by transforming the - samples of another distribution which does not depend on the - parameterization of `q`. This property ensures the gradient (with respect - to parameters) is valid. - TypeError: if `p_log_prob` is not a Python `callable`. - """ - with ops.name_scope(name, "monte_carlo_csiszar_f_divergence", [num_draws]): - if use_reparametrization is None: - use_reparametrization = (q.reparameterization_type - == distribution.FULLY_REPARAMETERIZED) - elif (use_reparametrization and - q.reparameterization_type != distribution.FULLY_REPARAMETERIZED): - # TODO(jvdillon): Consider only raising an exception if the gradient is - # requested. - raise ValueError( - "Distribution `q` must be reparameterized, i.e., a diffeomorphic " - "transformation of a parameterless distribution. (Otherwise this " - "function has a biased gradient.)") - if not callable(p_log_prob): - raise TypeError("`p_log_prob` must be a Python `callable` function.") - return monte_carlo.expectation( - f=lambda q_samples: f(p_log_prob(q_samples) - q.log_prob(q_samples)), - samples=q.sample(num_draws, seed=seed), - log_prob=q.log_prob, # Only used if use_reparametrization=False. - use_reparametrization=use_reparametrization) - - -def csiszar_vimco(f, - p_log_prob, - q, - num_draws, - num_batch_draws=1, - seed=None, - name=None): - """Use VIMCO to lower the variance of gradient[csiszar_function(Avg(logu))]. - - This function generalizes "Variational Inference for Monte Carlo Objectives" - (VIMCO), i.e., https://arxiv.org/abs/1602.06725, to Csiszar f-Divergences. - - Note: if `q.reparameterization_type = distribution.FULLY_REPARAMETERIZED`, - consider using `monte_carlo_csiszar_f_divergence`. - - The VIMCO loss is: - - ```none - vimco = f(Avg{logu[i] : i=0,...,m-1}) - where, - logu[i] = log( p(x, h[i]) / q(h[i] | x) ) - h[i] iid~ q(H | x) - ``` - - Interestingly, the VIMCO gradient is not the naive gradient of `vimco`. - Rather, it is characterized by: - - ```none - grad[vimco] - variance_reducing_term - where, - variance_reducing_term = Sum{ grad[log q(h[i] | x)] * - (vimco - f(log Avg{h[j;i] : j=0,...,m-1})) - : i=0, ..., m-1 } - h[j;i] = { u[j] j!=i - { GeometricAverage{ u[k] : k!=i} j==i - ``` - - (We omitted `stop_gradient` for brevity. See implementation for more details.) - - The `Avg{h[j;i] : j}` term is a kind of "swap-out average" where the `i`-th - element has been replaced by the leave-`i`-out Geometric-average. - - This implementation prefers numerical precision over efficiency, i.e., - `O(num_draws * num_batch_draws * prod(batch_shape) * prod(event_shape))`. - (The constant may be fairly large, perhaps around 12.) - - Args: - f: Python `callable` representing a Csiszar-function in log-space. - p_log_prob: Python `callable` representing the natural-log of the - probability under distribution `p`. (In variational inference `p` is the - joint distribution.) - q: `tf.Distribution`-like instance; must implement: `sample(n, seed)`, and - `log_prob(x)`. (In variational inference `q` is the approximate posterior - distribution.) - num_draws: Integer scalar number of draws used to approximate the - f-Divergence expectation. - num_batch_draws: Integer scalar number of draws used to approximate the - f-Divergence expectation. - seed: Python `int` seed for `q.sample`. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - vimco: The Csiszar f-Divergence generalized VIMCO objective. - - Raises: - ValueError: if `num_draws < 2`. - """ - with ops.name_scope(name, "csiszar_vimco", [num_draws, num_batch_draws]): - if num_draws < 2: - raise ValueError("Must specify num_draws > 1.") - stop = array_ops.stop_gradient # For readability. - x = stop(q.sample(sample_shape=[num_draws, num_batch_draws], - seed=seed)) - logqx = q.log_prob(x) - logu = p_log_prob(x) - logqx - f_log_avg_u, f_log_sooavg_u = [f(r) for r in csiszar_vimco_helper(logu)] - dotprod = math_ops.reduce_sum( - logqx * stop(f_log_avg_u - f_log_sooavg_u), - axis=0) # Sum over iid samples. - # We now rewrite f_log_avg_u so that: - # `grad[f_log_avg_u] := grad[f_log_avg_u + dotprod]`. - # To achieve this, we use a trick that - # `f(x) - stop(f(x)) == zeros_like(f(x))` - # but its gradient is grad[f(x)]. - # Note that IEEE754 specifies that `x - x == 0.` and `x + 0. == x`, hence - # this trick loses no precision. For more discussion regarding the relevant - # portions of the IEEE754 standard, see the StackOverflow question, - # "Is there a floating point value of x, for which x-x == 0 is false?" - # http://stackoverflow.com/q/2686644 - f_log_avg_u += dotprod - stop(dotprod) # Add zeros_like(dot_prod). - return math_ops.reduce_mean(f_log_avg_u, axis=0) # Avg over batches. - - -def csiszar_vimco_helper(logu, name=None): - """Helper to `csiszar_vimco`; computes `log_avg_u`, `log_sooavg_u`. - - `axis = 0` of `logu` is presumed to correspond to iid samples from `q`, i.e., - - ```none - logu[j] = log(u[j]) - u[j] = p(x, h[j]) / q(h[j] | x) - h[j] iid~ q(H | x) - ``` - - Args: - logu: Floating-type `Tensor` representing `log(p(x, h) / q(h | x))`. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - log_avg_u: `logu.dtype` `Tensor` corresponding to the natural-log of the - average of `u`. The sum of the gradient of `log_avg_u` is `1`. - log_sooavg_u: `logu.dtype` `Tensor` characterized by the natural-log of the - average of `u`` except that the average swaps-out `u[i]` for the - leave-`i`-out Geometric-average. The mean of the gradient of - `log_sooavg_u` is `1`. Mathematically `log_sooavg_u` is, - ```none - log_sooavg_u[i] = log(Avg{h[j ; i] : j=0, ..., m-1}) - h[j ; i] = { u[j] j!=i - { GeometricAverage{u[k] : k != i} j==i - ``` - - """ - with ops.name_scope(name, "csiszar_vimco_helper", [logu]): - logu = ops.convert_to_tensor(logu, name="logu") - - n = logu.shape.with_rank_at_least(1)[0].value - if n is None: - n = array_ops.shape(logu)[0] - log_n = math_ops.log(math_ops.cast(n, dtype=logu.dtype)) - nm1 = math_ops.cast(n - 1, dtype=logu.dtype) - else: - log_n = np.log(n).astype(logu.dtype.as_numpy_dtype) - nm1 = np.asarray(n - 1, dtype=logu.dtype.as_numpy_dtype) - - # Throughout we reduce across axis=0 since this is presumed to be iid - # samples. - - log_max_u = math_ops.reduce_max(logu, axis=0) - log_sum_u_minus_log_max_u = math_ops.reduce_logsumexp( - logu - log_max_u, axis=0) - - # log_loosum_u[i] = - # = logsumexp(logu[j] : j != i) - # = log( exp(logsumexp(logu)) - exp(logu[i]) ) - # = log( exp(logsumexp(logu - logu[i])) exp(logu[i]) - exp(logu[i])) - # = logu[i] + log(exp(logsumexp(logu - logu[i])) - 1) - # = logu[i] + log(exp(logsumexp(logu) - logu[i]) - 1) - # = logu[i] + softplus_inverse(logsumexp(logu) - logu[i]) - d = log_sum_u_minus_log_max_u + (log_max_u - logu) - # We use `d != 0` rather than `d > 0.` because `d < 0.` should never - # happens; if it does we want to complain loudly (which `softplus_inverse` - # will). - d_ok = math_ops.not_equal(d, 0.) - safe_d = array_ops.where(d_ok, d, array_ops.ones_like(d)) - d_ok_result = logu + distribution_util.softplus_inverse(safe_d) - - inf = np.array(np.inf, dtype=logu.dtype.as_numpy_dtype) - - # When not(d_ok) and is_positive_and_largest then we manually compute the - # log_loosum_u. (We can efficiently do this for any one point but not all, - # hence we still need the above calculation.) This is good because when - # this condition is met, we cannot use the above calculation; its -inf. - # We now compute the log-leave-out-max-sum, replicate it to every - # point and make sure to select it only when we need to. - is_positive_and_largest = math_ops.logical_and( - logu > 0., - math_ops.equal(logu, log_max_u[array_ops.newaxis, ...])) - log_lomsum_u = math_ops.reduce_logsumexp( - array_ops.where(is_positive_and_largest, - array_ops.fill(array_ops.shape(logu), -inf), - logu), - axis=0, keep_dims=True) - log_lomsum_u = array_ops.tile( - log_lomsum_u, - multiples=1 + array_ops.pad([n-1], [[0, array_ops.rank(logu)-1]])) - - d_not_ok_result = array_ops.where( - is_positive_and_largest, - log_lomsum_u, - array_ops.fill(array_ops.shape(d), -inf)) - - log_loosum_u = array_ops.where(d_ok, d_ok_result, d_not_ok_result) - - # The swap-one-out-sum ("soosum") is n different sums, each of which - # replaces the i-th item with the i-th-left-out average, i.e., - # soo_sum_u[i] = [exp(logu) - exp(logu[i])] + exp(mean(logu[!=i])) - # = exp(log_loosum_u[i]) + exp(looavg_logu[i]) - looavg_logu = (math_ops.reduce_sum(logu, axis=0) - logu) / nm1 - log_soosum_u = math_ops.reduce_logsumexp( - array_ops.stack([log_loosum_u, looavg_logu]), - axis=0) - - log_avg_u = log_sum_u_minus_log_max_u + log_max_u - log_n - log_sooavg_u = log_soosum_u - log_n - - log_avg_u.set_shape(logu.shape.with_rank_at_least(1)[1:]) - log_sooavg_u.set_shape(logu.shape) - - return log_avg_u, log_sooavg_u diff --git a/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py b/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py deleted file mode 100644 index d44fe6529a7ff0da0c6747e193fdb98a272a8da3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py +++ /dev/null @@ -1,110 +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. -# ============================================================================== -"""Functions for specifying custom gradients. - -@@custom_gradient - -""" - -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 math_ops - -__all__ = [ - "custom_gradient", -] - - -def custom_gradient(fx, gx, x, axis=(), fx_gx_manually_stopped=False, - name=None): - """Enables specifying a custom gradient. - - This function works by clever application of `stop_gradient`. I.e., observe - that: - - ```none - h(x) = x * stop_gradient(g(x)) + stop_gradient(f(x) - x * 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 - is necessary to reduce `x` to a scalar. This can be done by indicating the - `axis` over which `f` operates or by appropriately `reduce_sum`-ing `x`, prior - to calling this function. - - Partial Custom Gradient: - - Suppose `h(x) = htilde(x, y)`. Note that `dh/dx = stop(g(x))` but `dh/dy = - None`. This is because a `Tensor` cannot have only a portion of its gradient - stopped. To circumvent this issue, one must manually `stop_gradient` the - relevant portions of `f`, `g`. For example see the unit-test, - `test_works_correctly_fx_gx_manually_stopped`. - - Args: - fx: `Tensor`. Output of function evaluated at `x`. - gx: `Tensor`. Gradient of function evaluated at `x`. - x: `Tensor`. Point of evaluation for `f, g`. - axis: 1D `int` `Tensor` representing dimensions of `x` which are the domain - of `f`. If `()` (the default), `f` is assumed scalar-domain/scalar-range. - If `None` `f` is assumed to render one scalar given all of `x`. Otherwise - `f` is assumed to output one scalar for each of `axis` dimensions of `x`. - fx_gx_manually_stopped: Python `bool` indicating that `fx`, `gx` manually - have `stop_gradient` applied. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - fx: Floating-type `Tensor` equal to `f(x)` but which has gradient - `stop_gradient(g(x))`. - """ - with ops.name_scope(name, "custom_gradient", [fx, gx, x]): - fx = ops.convert_to_tensor(fx, name="fx") - # We don't want to bother eagerly computing `gx` since we may not even need - # it. - with ops.control_dependencies([fx]): - gx = ops.convert_to_tensor(gx, dtype=fx.dtype, name="gx") - gx = array_ops.identity(gx, name="gx") - # Proof of correctness: - # - # f(x) = x * stop[gx] + stop[fx - x * gx] - # = stop[fx] - # - # g(x) = grad[fx] - # = stop[gx] + grad[stop[fx - x * gx]] - # = stop[gx] + 0 - # - # Notice that when x is zero it still works: - # grad[x * stop(gx) + stop(fx - x * gx)] = 1 * stop[gx] + 0 = stop[gx] - # - # The proof is similar for the tensor-domain case, except that `x` is - # replaced by `reduce_sum(x)`. - sum_x = math_ops.reduce_sum(x, axis=axis, name="sum_x") - if not fx_gx_manually_stopped: - fx = array_ops.stop_gradient(fx) - gx = array_ops.stop_gradient(gx) - # IEEE754 ensures `(x-x)==0.` and that `0.*x==0.` so we make sure to write - # the code this way, rather than, e.g., - # `sum_x * stop(gx) + stop(fx - sum_x * gx)`. - # For more discussion regarding the relevant portions of the IEEE754 - # standard, see the StackOverflow question, - # "Is there a floating point value of x, for which x-x == 0 is false?" - # http://stackoverflow.com/q/2686644 - return (sum_x - array_ops.stop_gradient(sum_x)) * gx + fx diff --git a/tensorflow/contrib/bayesflow/python/ops/halton_sequence_impl.py b/tensorflow/contrib/bayesflow/python/ops/halton_sequence_impl.py deleted file mode 100644 index 8cabf18903b5f15002470acdfb8fdd3ec31a7413..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/halton_sequence_impl.py +++ /dev/null @@ -1,264 +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. -# ============================================================================== -"""Quasi Monte Carlo support: Halton sequence. - -@@sample -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -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 - - -__all__ = [ - 'sample', -] - - -# The maximum dimension we support. This is limited by the number of primes -# in the _PRIMES array. -_MAX_DIMENSION = 1000 - - -def sample(dim, num_samples=None, sample_indices=None, dtype=None, name=None): - r"""Returns a sample from the `m` dimensional Halton sequence. - - Warning: The sequence elements take values only between 0 and 1. Care must be - taken to appropriately transform the domain of a function if it differs from - the unit cube before evaluating integrals using Halton samples. It is also - important to remember that quasi-random numbers are not a replacement for - pseudo-random numbers in every context. Quasi random numbers are completely - deterministic and typically have significant negative autocorrelation (unless - randomized). - - Computes the members of the low discrepancy Halton sequence in dimension - `dim`. The d-dimensional sequence takes values in the unit hypercube in d - dimensions. Currently, only dimensions up to 1000 are supported. The prime - base for the `k`-th axes is the k-th prime starting from 2. For example, - if dim = 3, then the bases will be [2, 3, 5] respectively and the first - element of the sequence will be: [0.5, 0.333, 0.2]. For a more complete - description of the Halton sequences see: - https://en.wikipedia.org/wiki/Halton_sequence. For low discrepancy sequences - and their applications see: - https://en.wikipedia.org/wiki/Low-discrepancy_sequence. - - The user must supply either `num_samples` or `sample_indices` but not both. - The former is the number of samples to produce starting from the first - element. If `sample_indices` is given instead, the specified elements of - the sequence are generated. For example, sample_indices=tf.range(10) is - equivalent to specifying n=10. - - Example Use: - - ```python - bf = tf.contrib.bayesflow - - # Produce the first 1000 members of the Halton sequence in 3 dimensions. - num_samples = 1000 - dim = 3 - sample = bf.halton_sequence.sample(dim, num_samples=num_samples) - - # Evaluate the integral of x_1 * x_2^2 * x_3^3 over the three dimensional - # hypercube. - powers = tf.range(1.0, limit=dim + 1) - integral = tf.reduce_mean(tf.reduce_prod(sample ** powers, axis=-1)) - true_value = 1.0 / tf.reduce_prod(powers + 1.0) - with tf.Session() as session: - values = session.run((integral, true_value)) - - # Produces a relative absolute error of 1.7%. - print ("Estimated: %f, True Value: %f" % values) - - # Now skip the first 1000 samples and recompute the integral with the next - # thousand samples. The sample_indices argument can be used to do this. - - - sample_indices = tf.range(start=1000, limit=1000 + num_samples, - dtype=tf.int32) - sample_leaped = halton.sample(dim, sample_indices=sample_indices) - - integral_leaped = tf.reduce_mean(tf.reduce_prod(sample_leaped ** powers, - axis=-1)) - with tf.Session() as session: - values = session.run((integral_leaped, true_value)) - # Now produces a relative absolute error of 0.05%. - print ("Leaped Estimated: %f, True Value: %f" % values) - ``` - - Args: - dim: Positive Python `int` representing each sample's `event_size.` Must - not be greater than 1000. - num_samples: (Optional) positive Python `int`. The number of samples to - generate. Either this parameter or sample_indices must be specified but - not both. If this parameter is None, then the behaviour is determined by - the `sample_indices`. - sample_indices: (Optional) `Tensor` of dtype int32 and rank 1. The elements - of the sequence to compute specified by their position in the sequence. - The entries index into the Halton sequence starting with 0 and hence, - must be whole numbers. For example, sample_indices=[0, 5, 6] will produce - the first, sixth and seventh elements of the sequence. If this parameter - is None, then the `num_samples` parameter must be specified which gives - the number of desired samples starting from the first sample. - dtype: (Optional) The dtype of the sample. One of `float32` or `float64`. - Default is `float32`. - name: (Optional) Python `str` describing ops managed by this function. If - not supplied the name of this function is used. - - Returns: - halton_elements: Elements of the Halton sequence. `Tensor` of supplied dtype - and `shape` `[num_samples, dim]` if `num_samples` was specified or shape - `[s, dim]` where s is the size of `sample_indices` if `sample_indices` - were specified. - - Raises: - ValueError: if both `sample_indices` and `num_samples` were specified or - if dimension `dim` is less than 1 or greater than 1000. - """ - if dim < 1 or dim > _MAX_DIMENSION: - raise ValueError( - 'Dimension must be between 1 and {}. Supplied {}'.format(_MAX_DIMENSION, - dim)) - if (num_samples is None) == (sample_indices is None): - raise ValueError('Either `num_samples` or `sample_indices` must be' - ' specified but not both.') - - dtype = dtype or dtypes.float32 - if not dtype.is_floating: - raise ValueError('dtype must be of `float`-type') - - with ops.name_scope(name, 'sample', values=[sample_indices]): - # Here and in the following, the shape layout is as follows: - # [sample dimension, event dimension, coefficient dimension]. - # The coefficient dimension is an intermediate axes which will hold the - # weights of the starting integer when expressed in the (prime) base for - # an event dimension. - indices = _get_indices(num_samples, sample_indices, dtype) - radixes = array_ops.constant(_PRIMES[0:dim], dtype=dtype, shape=[dim, 1]) - - max_sizes_by_axes = _base_expansion_size(math_ops.reduce_max(indices), - radixes) - - max_size = math_ops.reduce_max(max_sizes_by_axes) - - # The powers of the radixes that we will need. Note that there is a bit - # of an excess here. Suppose we need the place value coefficients of 7 - # in base 2 and 3. For 2, we will have 3 digits but we only need 2 digits - # for base 3. However, we can only create rectangular tensors so we - # store both expansions in a [2, 3] tensor. This leads to the problem that - # we might end up attempting to raise large numbers to large powers. For - # example, base 2 expansion of 1024 has 10 digits. If we were in 10 - # dimensions, then the 10th prime (29) we will end up computing 29^10 even - # though we don't need it. We avoid this by setting the exponents for each - # axes to 0 beyond the maximum value needed for that dimension. - exponents_by_axes = array_ops.tile([math_ops.range(max_size)], [dim, 1]) - weight_mask = exponents_by_axes > max_sizes_by_axes - capped_exponents = array_ops.where( - weight_mask, array_ops.zeros_like(exponents_by_axes), exponents_by_axes) - weights = radixes ** capped_exponents - coeffs = math_ops.floor_div(indices, weights) - coeffs *= 1 - math_ops.cast(weight_mask, dtype) - coeffs = (coeffs % radixes) / radixes - return math_ops.reduce_sum(coeffs / weights, axis=-1) - - -def _get_indices(n, sample_indices, dtype, name=None): - """Generates starting points for the Halton sequence procedure. - - The k'th element of the sequence is generated starting from a positive integer - which must be distinct for each `k`. It is conventional to choose the starting - point as `k` itself (or `k+1` if k is zero based). This function generates - the starting integers for the required elements and reshapes the result for - later use. - - Args: - n: Positive `int`. The number of samples to generate. If this - parameter is supplied, then `sample_indices` should be None. - sample_indices: `Tensor` of dtype int32 and rank 1. The entries - index into the Halton sequence starting with 0 and hence, must be whole - numbers. For example, sample_indices=[0, 5, 6] will produce the first, - sixth and seventh elements of the sequence. If this parameter is not None - then `n` must be None. - dtype: The dtype of the sample. One of `float32` or `float64`. - Default is `float32`. - name: Python `str` name which describes ops created by this function. - - Returns: - indices: `Tensor` of dtype `dtype` and shape = `[n, 1, 1]`. - """ - with ops.name_scope(name, 'get_indices', [n, sample_indices]): - if sample_indices is None: - sample_indices = math_ops.range(n, dtype=dtype) - else: - sample_indices = math_ops.cast(sample_indices, dtype) - - # Shift the indices so they are 1 based. - indices = sample_indices + 1 - - # Reshape to make space for the event dimension and the place value - # coefficients. - return array_ops.reshape(indices, [-1, 1, 1]) - - -def _base_expansion_size(num, bases): - """Computes the number of terms in the place value expansion. - - Let num = a0 + a1 b + a2 b^2 + ... ak b^k be the place value expansion of - `num` in base b (ak <> 0). This function computes and returns `k` for each - base `b` specified in `bases`. - - This can be inferred from the base `b` logarithm of `num` as follows: - $$k = Floor(log_b (num)) + 1 = Floor( log(num) / log(b)) + 1$$ - - Args: - num: Scalar `Tensor` of dtype either `float32` or `float64`. The number to - compute the base expansion size of. - bases: `Tensor` of the same dtype as num. The bases to compute the size - against. - - Returns: - Tensor of same dtype and shape as `bases` containing the size of num when - written in that base. - """ - return math_ops.floor(math_ops.log(num) / math_ops.log(bases)) + 1 - - -def _primes_less_than(n): - # Based on - # https://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n-in-python/3035188#3035188 - """Returns sorted array of primes such that `2 <= prime < n`.""" - small_primes = np.array((2, 3, 5)) - if n <= 6: - return small_primes[small_primes < n] - sieve = np.ones(n // 3 + (n % 6 == 2), dtype=np.bool) - sieve[0] = False - m = int(n ** 0.5) // 3 + 1 - for i in range(m): - if not sieve[i]: - continue - k = 3 * i + 1 | 1 - sieve[k ** 2 // 3::2 * k] = False - sieve[(k ** 2 + 4 * k - 2 * k * (i & 1)) // 3::2 * k] = False - return np.r_[2, 3, 3 * np.nonzero(sieve)[0] + 1 | 1] - -_PRIMES = _primes_less_than(7919+1) - -assert len(_PRIMES) == _MAX_DIMENSION diff --git a/tensorflow/contrib/bayesflow/python/ops/hmc.py b/tensorflow/contrib/bayesflow/python/ops/hmc.py index 977d42fc16bb91777a76c45ac24f3c5dc587f5fe..c8a5a195d3d709ded7afd09287255deab2ac2f3c 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,8 @@ 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", + "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..66afcc749746ab5c04114e585c5f93a3f3354d86 100644 --- a/tensorflow/contrib/bayesflow/python/ops/hmc_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/hmc_impl.py @@ -14,17 +14,15 @@ # ============================================================================== """Hamiltonian Monte Carlo, a gradient-based MCMC algorithm. -@@chain -@@update -@@leapfrog_integrator -@@leapfrog_step -@@ais_chain +@@sample_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,283 +30,312 @@ 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", + "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", + [ + "log_accept_ratio", + "current_grads_target_log_prob", # "Current result" means "accepted". + "current_target_log_prob", # "Current result" means "accepted". + "is_accepted", + "proposed_grads_target_log_prob", + "proposed_state", + "proposed_target_log_prob", + ]) + + +def _make_dummy_kernel_results( + dummy_state, + dummy_target_log_prob, + dummy_grads_target_log_prob): + return KernelResults( + log_accept_ratio=dummy_target_log_prob, + current_grads_target_log_prob=dummy_grads_target_log_prob, + current_target_log_prob=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, + ) + + +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. - #### Examples: + Since HMC states are correlated, it is sometimes desirable to produce + additional intermediate states, and then discard them, ending up with a set of + states with decreased autocorrelation. See [1]. Such "thinning" is made + possible by setting `num_steps_between_results > 0`. The chain then takes + `num_steps_between_results` extra steps between the steps that make it into + the results. The extra steps are never materialized (in calls to `sess.run`), + and thus do not increase memory requirements. - ```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) - ``` + [1]: "Statistically efficient thinning of a Markov chain sampler." + Art B. Owen. April 2017. + http://statweb.stanford.edu/~owen/reports/bestthinning.pdf - ```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) - ``` + #### Examples: - ```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) - ``` + ##### Sample from a diagonal-variance Gaussian. ```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_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) ``` - """ - 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): - """Runs annealed importance sampling (AIS) to estimate normalizing constants. - This routine uses Hamiltonian Monte Carlo to sample from a series of - distributions that slowly interpolates between an initial "proposal" - distribution + ##### Sampling from factor-analysis posteriors with known factors. - `exp(proposal_log_prob_fn(x) - proposal_log_normalizer)` + I.e., - and the target distribution - - `exp(target_log_prob_fn(x) - target_log_normalizer)`, + ```none + for i=1..n: + w[i] ~ Normal(0, eye(d)) # prior + x[i] ~ Normal(loc=matmul(w[i], F)) # likelihood + ``` - accumulating importance weights along the way. The product of these - importance weights gives an unbiased estimate of the ratio of the - normalizing constants of the initial distribution and the target - distribution: + where `F` denotes factors. - E[exp(w)] = exp(target_log_normalizer - proposal_log_normalizer). + ```python + 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: - 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. + 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. The number of returned chain states is + still equal to `num_results`. Default value: 0 (i.e., no thinning). + 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: - 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)`. - - #### Examples: - - ```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) - # 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) - ``` - - ```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) - # 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) - ``` + next_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_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): + 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 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 +343,619 @@ 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)) + + next_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(next_x) + + step_size_update = step_size.assign_add( + step_size * tf.where( + tf.exp(tf.minimum(other_results.log_accept_ratio), 0.) > + 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`. + next_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) + log_accept_ratio = -energy_change + + # u < exp(log_accept_ratio), where u~Uniform[0,1) + # ==> log(u) < log_accept_ratio + random_value = random_ops.random_uniform( + shape=array_ops.shape(energy_change), + dtype=energy_change.dtype, + seed=seed) + random_negative = math_ops.log(random_value) + is_accepted = random_negative < log_accept_ratio + + accepted_target_log_prob = array_ops.where(is_accepted, + proposed_target_log_prob, + current_target_log_prob) + + next_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(next_state_parts), + KernelResults( + log_accept_ratio=log_accept_ratio, + current_grads_target_log_prob=accepted_grads_target_log_prob, + current_target_log_prob=accepted_target_log_prob, + 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, + ), + ] + + +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`. - Example: Simple quadratic potential. + #### Examples: + + ##### 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]) + + [ + next_momentums, + next_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( + [next_momentums[0], next_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 `next_state`. + proposed_grads_target_log_prob: Gradient of `proposed_target_log_prob` wrt + `next_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 # potential + lp1 = -proposed_target_log_prob # proposed_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 deleted file mode 100644 index a742b7c1aa593d6c08bf9d8d597c99c9fc4e7aed..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/layers.py +++ /dev/null @@ -1,67 +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. -# ============================================================================== -"""Probabilistic neural layers. - -See ${python/contrib.bayesflow.layers}. -""" - -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.layers_conv_variational 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 = [ - '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', - 'dense_reparameterization', - 'dense_local_reparameterization', - 'dense_flipout', - 'default_loc_scale_fn', - 'default_mean_field_normal_fn', -] - -remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py b/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py deleted file mode 100644 index 7723cfb442712626ff415f1412e3362f2392ce9f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py +++ /dev/null @@ -1,2943 +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. -# ============================================================================== -"""Convolutional variational layer classes and their functional aliases. -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.bayesflow.python.ops import 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.layers import base as layers_lib -from tensorflow.python.layers import utils -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import nn -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): - """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. - - 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. - """ - - 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(_ConvVariational, self).__init__( - trainable=trainable, - name=name, - activity_regularizer=activity_regularizer, - **kwargs) - self.rank = rank - self.filters = filters - self.kernel_size = utils.normalize_tuple(kernel_size, rank, "kernel_size") - self.strides = utils.normalize_tuple(strides, rank, "strides") - self.padding = utils.normalize_padding(padding) - self.data_format = utils.normalize_data_format(data_format) - self.dilation_rate = utils.normalize_tuple( - dilation_rate, rank, "dilation_rate") - self.activation = activation - self.input_spec = layers_lib.InputSpec(ndim=self.rank + 2) - self.kernel_posterior_fn = kernel_posterior_fn - self.kernel_posterior_tensor_fn = kernel_posterior_tensor_fn - self.kernel_prior_fn = kernel_prior_fn - self.kernel_divergence_fn = kernel_divergence_fn - self.bias_posterior_fn = bias_posterior_fn - self.bias_posterior_tensor_fn = bias_posterior_tensor_fn - self.bias_prior_fn = bias_prior_fn - self.bias_divergence_fn = bias_divergence_fn - - def build(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - if self.data_format == "channels_first": - channel_axis = 1 - else: - channel_axis = -1 - if input_shape[channel_axis].value is None: - raise ValueError("The channel dimension of the inputs " - "should be defined. Found `None`.") - input_dim = input_shape[channel_axis].value - kernel_shape = self.kernel_size + (input_dim, self.filters) - dtype = dtypes.as_dtype(self.dtype) - - # Must have a posterior kernel. - self.kernel_posterior = self.kernel_posterior_fn( - dtype, kernel_shape, "kernel_posterior", - self.trainable, self.add_variable) - - if self.kernel_prior_fn is None: - self.kernel_prior = None - else: - self.kernel_prior = self.kernel_prior_fn( - dtype, kernel_shape, "kernel_prior", - self.trainable, self.add_variable) - self._built_kernel_divergence = False - - if self.bias_posterior_fn is None: - self.bias_posterior = None - else: - self.bias_posterior = self.bias_posterior_fn( - dtype, (self.filters,), "bias_posterior", - self.trainable, self.add_variable) - - if self.bias_prior_fn is None: - self.bias_prior = None - else: - self.bias_prior = self.bias_prior_fn( - dtype, (self.filters,), "bias_prior", - self.trainable, self.add_variable) - self._built_bias_divergence = False - - self.input_spec = layers_lib.InputSpec(ndim=self.rank + 2, - axes={channel_axis: input_dim}) - self._convolution_op = nn_ops.Convolution( - input_shape, - filter_shape=tensor_shape.TensorShape(kernel_shape), - dilation_rate=self.dilation_rate, - strides=self.strides, - padding=self.padding.upper(), - data_format=utils.convert_data_format(self.data_format, - self.rank + 2)) - - self.built = True - - def call(self, inputs): - inputs = ops.convert_to_tensor(inputs, dtype=self.dtype) - - outputs = self._apply_variational_kernel(inputs) - outputs = self._apply_variational_bias(outputs) - if self.activation is not None: - outputs = self.activation(outputs) - if not self._built_kernel_divergence: - kernel_posterior = self.kernel_posterior - kernel_prior = self.kernel_prior - if isinstance(self.kernel_posterior, independent_lib.Independent): - kernel_posterior = kernel_posterior.distribution - if isinstance(self.kernel_prior, independent_lib.Independent): - kernel_prior = kernel_prior.distribution - self._apply_divergence(self.kernel_divergence_fn, - kernel_posterior, - kernel_prior, - self.kernel_posterior_tensor, - name="divergence_kernel") - self._built_kernel_divergence = True - if not self._built_bias_divergence: - bias_posterior = self.bias_posterior - bias_prior = self.bias_prior - if isinstance(self.bias_posterior, independent_lib.Independent): - bias_posterior = bias_posterior.distribution - if isinstance(self.bias_prior, independent_lib.Independent): - bias_prior = bias_prior.distribution - self._apply_divergence(self.bias_divergence_fn, - bias_posterior, - bias_prior, - self.bias_posterior_tensor, - name="divergence_bias") - self._built_bias_divergence = True - return outputs - - def _apply_variational_bias(self, inputs): - if self.bias_posterior is None: - self.bias_posterior_tensor = None - return inputs - self.bias_posterior_tensor = self.bias_posterior_tensor_fn( - self.bias_posterior) - outputs = inputs - if self.data_format == "channels_first": - if self.rank == 1: - # nn.bias_add does not accept a 1D input tensor. - bias = array_ops.reshape(self.bias_posterior_tensor, - (1, self.filters, 1)) - outputs += bias - if self.rank == 2: - outputs = nn.bias_add(outputs, - self.bias_posterior_tensor, - data_format="NCHW") - if self.rank == 3: - # As of Mar 2017, direct addition is significantly slower than - # bias_add when computing gradients. To use bias_add, we collapse Z - # and Y into a single dimension to obtain a 4D input tensor. - outputs_shape = outputs.shape.as_list() - outputs_4d = array_ops.reshape(outputs, - [outputs_shape[0], outputs_shape[1], - outputs_shape[2] * outputs_shape[3], - outputs_shape[4]]) - outputs_4d = nn.bias_add(outputs_4d, - self.bias_posterior_tensor, - data_format="NCHW") - outputs = array_ops.reshape(outputs_4d, outputs_shape) - else: - outputs = nn.bias_add(outputs, - self.bias_posterior_tensor, - data_format="NHWC") - return outputs - - def _apply_divergence(self, divergence_fn, posterior, prior, - posterior_tensor, name): - if (divergence_fn is None or - posterior is None or - prior is None): - divergence = None - return - divergence = standard_ops.identity( - divergence_fn( - posterior, prior, posterior_tensor), - name=name) - self.add_loss(divergence) - - 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] - new_space = [] - for i in range(len(space)): - new_dim = utils.conv_output_length( - space[i], - self.kernel_size[i], - padding=self.padding, - stride=self.strides[i], - dilation=self.dilation_rate[i]) - new_space.append(new_dim) - return tensor_shape.TensorShape([input_shape[0]] + new_space + - [self.filters]) - else: - space = input_shape[2:] - new_space = [] - for i in range(len(space)): - new_dim = utils.conv_output_length( - space[i], - self.kernel_size[i], - padding=self.padding, - stride=self.strides[i], - dilation=self.dilation_rate[i]) - new_space.append(new_dim) - return tensor_shape.TensorShape([input_shape[0], self.filters] + - new_space) - - -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 - (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: - 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 - 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. - - #### 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.Conv1DReparameterization(64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu)(net) - net = tf.reshape(net, [-1, 128 * 64]) - 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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - - 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), - name=None, - **kwargs): - super(Conv1DReparameterization, 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, - name=name, **kwargs) - - -def conv1d_reparameterization( - 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), - 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 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: - 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 - 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_reparameterization(net, - filters=64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu) - net = tf.reshape(net, [-1, 128 * 64]) - 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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - layer = Conv1DReparameterization( - 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 Conv2DReparameterization(_ConvReparameterization): - """2D convolution layer (e.g. spatial convolution over images). - - 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: - 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 - 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. - - #### 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.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.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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - - 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), - name=None, - **kwargs): - super(Conv2DReparameterization, 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, - name=name, **kwargs) - - -def conv2d_reparameterization( - 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), - 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 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: - 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 - 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_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_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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - layer = Conv2DReparameterization( - 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 Conv3DReparameterization(_ConvReparameterization): - """3D convolution layer (e.g. spatial convolution over volumes). - - 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: - 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 - 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. - - #### 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.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.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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - - 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), - name=None, - **kwargs): - super(Conv3DReparameterization, 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, - name=name, **kwargs) - - -def conv3d_reparameterization( - 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), - 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 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: - 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 - 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_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_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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "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 = Conv3DFlipout( - 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) - - -# Aliases - -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.py b/tensorflow/contrib/bayesflow/python/ops/layers_dense_variational.py deleted file mode 100644 index 591a8e553de0c194786c7ee8693665f762711b2d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/layers_dense_variational.py +++ /dev/null @@ -1,1176 +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. -# ============================================================================== -"""Dense Bayesian layer using KL-divergence based variational inference. -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.bayesflow.python.ops import 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.layers import base as layers_lib -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import nn -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 _DenseVariational(layers_lib.Layer): - """Abstract densely-connected class (private, used as implementation base). - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - outputs = activation(matmul(inputs, kernel) + 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. - - Args: - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Properties: - units: Python integer, dimensionality of the output space. - 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. - """ - - def __init__( - self, - units, - 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(_DenseVariational, self).__init__( - trainable=trainable, - name=name, - activity_regularizer=activity_regularizer, - **kwargs) - self.units = units - self.activation = activation - self.input_spec = layers_lib.InputSpec(min_ndim=2) - self.kernel_posterior_fn = kernel_posterior_fn - self.kernel_posterior_tensor_fn = kernel_posterior_tensor_fn - self.kernel_prior_fn = kernel_prior_fn - self.kernel_divergence_fn = kernel_divergence_fn - self.bias_posterior_fn = bias_posterior_fn - self.bias_posterior_tensor_fn = bias_posterior_tensor_fn - self.bias_prior_fn = bias_prior_fn - self.bias_divergence_fn = bias_divergence_fn - - def build(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - in_size = input_shape.with_rank_at_least(2)[-1].value - if in_size is None: - raise ValueError("The last dimension of the inputs to `Dense` " - "should be defined. Found `None`.") - self._input_spec = layers_lib.InputSpec(min_ndim=2, axes={-1: in_size}) - dtype = dtypes.as_dtype(self.dtype) - - # Must have a posterior kernel. - self.kernel_posterior = self.kernel_posterior_fn( - dtype, [in_size, self.units], "kernel_posterior", - self.trainable, self.add_variable) - - if self.kernel_prior_fn is None: - self.kernel_prior = None - else: - self.kernel_prior = self.kernel_prior_fn( - dtype, [in_size, self.units], "kernel_prior", - self.trainable, self.add_variable) - self._built_kernel_divergence = False - - if self.bias_posterior_fn is None: - self.bias_posterior = None - else: - self.bias_posterior = self.bias_posterior_fn( - dtype, [self.units], "bias_posterior", - self.trainable, self.add_variable) - - if self.bias_prior_fn is None: - self.bias_prior = None - else: - self.bias_prior = self.bias_prior_fn( - dtype, [self.units], "bias_prior", - self.trainable, self.add_variable) - self._built_bias_divergence = False - - self.built = True - - def call(self, inputs): - inputs = ops.convert_to_tensor(inputs, dtype=self.dtype) - - outputs = self._apply_variational_kernel(inputs) - outputs = self._apply_variational_bias(outputs) - if self.activation is not None: - outputs = self.activation(outputs) # pylint: disable=not-callable - if not self._built_kernel_divergence: - kernel_posterior = self.kernel_posterior - kernel_prior = self.kernel_prior - if isinstance(self.kernel_posterior, independent_lib.Independent): - kernel_posterior = kernel_posterior.distribution - if isinstance(self.kernel_prior, independent_lib.Independent): - kernel_prior = kernel_prior.distribution - self._apply_divergence(self.kernel_divergence_fn, - kernel_posterior, - kernel_prior, - self.kernel_posterior_tensor, - name="divergence_kernel") - self._built_kernel_divergence = True - if not self._built_bias_divergence: - bias_posterior = self.bias_posterior - bias_prior = self.bias_prior - if isinstance(self.bias_posterior, independent_lib.Independent): - bias_posterior = bias_posterior.distribution - if isinstance(self.bias_prior, independent_lib.Independent): - bias_prior = bias_prior.distribution - self._apply_divergence(self.bias_divergence_fn, - bias_posterior, - bias_prior, - self.bias_posterior_tensor, - name="divergence_bias") - self._built_bias_divergence = True - return outputs - - def _apply_variational_bias(self, inputs): - if self.bias_posterior is None: - self.bias_posterior_tensor = None - return inputs - self.bias_posterior_tensor = self.bias_posterior_tensor_fn( - self.bias_posterior) - return nn.bias_add(inputs, self.bias_posterior_tensor) - - def _apply_divergence(self, divergence_fn, posterior, prior, - posterior_tensor, name): - if (divergence_fn is None or - posterior is None or - prior is None): - divergence = None - return - divergence = standard_ops.identity( - divergence_fn( - posterior, prior, posterior_tensor), - name=name) - self.add_loss(divergence) - - def _matmul(self, inputs, kernel): - if inputs.shape.ndims <= 2: - return standard_ops.matmul(inputs, kernel) - # To handle broadcasting, we must use `tensordot`. - return standard_ops.tensordot(inputs, kernel, axes=[[-1], [0]]) - - def _compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).with_rank_at_least(2) - if input_shape[-1].value is None: - raise ValueError( - "The innermost dimension of input_shape must be defined, " - "but saw: {}".format(input_shape)) - return input_shape[:-1].concatenate(self.units) - - -class DenseReparameterization(_DenseVariational): - """Densely-connected layer class with reparameterization estimator. - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - 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. - - Args: - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Properties: - units: Python integer, dimensionality of the output space. - 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. - - #### 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 = tfp.layers.DenseReparameterization( - 512, activation=tf.nn.relu)(features) - 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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - - def __init__( - self, - units, - 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), - 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(DenseReparameterization, self).__init__( - units=units, - 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 - return self._matmul(inputs, self.kernel_posterior_tensor) - - -def dense_reparameterization( - inputs, - units, - 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, - reuse=None): - """Densely-connected layer with reparameterization estimator. - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - 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. - - Args: - inputs: Tensor input. - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Returns: - output: `Tensor` representing a the affine transformed input under a random - draw from the surrogate posterior distribution. - - #### 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 = tfp.layers.dense_reparameterization( - features, 512, activation=tf.nn.relu) - 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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses reparameterization gradients 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]: "Auto-Encoding Variational Bayes." - Diederik P. Kingma, Max Welling. - International Conference on Learning Representations, 2014. - """ - layer = DenseReparameterization( - units, - 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 DenseLocalReparameterization(_DenseVariational): - """Densely-connected layer class with local reparameterization estimator. - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - 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. - - Args: - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Properties: - units: Python integer, dimensionality of the output space. - 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. - - #### 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 = tfp.layers.DenseLocalReparameterization( - 512, activation=tf.nn.relu)(features) - logits = tfp.layers.DenseLocalReparameterization(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 local reparameterization gradients 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]: "Variational Dropout and the Local Reparameterization Trick." - Diederik P. Kingma, Tim Salimans, Max Welling. - Neural Information Processing Systems, 2015. - """ - - def __init__( - self, - units, - 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), - 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(DenseLocalReparameterization, self).__init__( - units=units, - 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): - if (not isinstance(self.kernel_posterior, independent_lib.Independent) or - not isinstance(self.kernel_posterior.distribution, normal_lib.Normal)): - raise TypeError( - "`DenseLocalReparameterization` requires " - "`kernel_posterior_fn` produce an instance of " - "`tf.distributions.Independent(tf.distributions.Normal)` " - "(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( - standard_ops.square(inputs), - standard_ops.square(self.kernel_posterior.distribution.scale)))) - self.kernel_posterior_affine_tensor = ( - self.kernel_posterior_tensor_fn(self.kernel_posterior_affine)) - self.kernel_posterior_tensor = None - return self.kernel_posterior_affine_tensor - - -def dense_local_reparameterization( - inputs, - units, - 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), - 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, - reuse=None): - """Densely-connected layer with local reparameterization estimator. - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - 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. - - Args: - inputs: Tensor input. - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Returns: - output: `Tensor` representing a the affine transformed input under a random - draw from the surrogate posterior distribution. - - #### 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 = tfp.layers.dense_local_reparameterization( - features, 512, activation=tf.nn.relu) - logits = tfp.layers.dense_local_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)) - loss = neg_log_likelihood + kl - train_op = tf.train.AdamOptimizer().minimize(loss) - ``` - - It uses local reparameterization gradients 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]: "Variational Dropout and the Local Reparameterization Trick." - Diederik P. Kingma, Tim Salimans, Max Welling. - Neural Information Processing Systems, 2015. - """ - layer = DenseLocalReparameterization( - units, - 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 DenseFlipout(_DenseVariational): - """Densely-connected layer class with Flipout estimator. - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - 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. - - Args: - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Properties: - units: Python integer, dimensionality of the output space. - 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 = tfp.layers.DenseFlipout( - 512, activation=tf.nn.relu)(features) - 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, - units, - 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), - 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(DenseFlipout, self).__init__( - units=units, - 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( - "`DenseFlipout` requires " - "`kernel_posterior_fn` produce an instance of " - "`tf.distributions.Independent(tf.distributions.Normal)` " - "(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) - self.kernel_posterior_affine_tensor = ( - self.kernel_posterior_tensor_fn(self.kernel_posterior_affine)) - self.kernel_posterior_tensor = None - - input_shape = array_ops.shape(inputs) - batch_shape = input_shape[:-1] - - 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, - seed=distribution_util.gen_new_seed( - self.seed, salt="dense_flipout")) - perturbed_inputs = self._matmul( - inputs * sign_input, self.kernel_posterior_affine_tensor) * sign_output - - outputs = self._matmul(inputs, self.kernel_posterior.distribution.loc) - outputs += perturbed_inputs - return outputs - - -def dense_flipout( - inputs, - units, - 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), - 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): - """Densely-connected layer with Flipout estimator. - - This layer implements the Bayesian variational inference analogue to - a dense layer by assuming the `kernel` and/or the `bias` are drawn - from distributions. By default, the layer implements a stochastic - forward pass via sampling from the kernel and bias posteriors, - - ```none - kernel, bias ~ posterior - 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. - - Args: - inputs: Tensor input. - units: Integer or Long, dimensionality of the output space. - activation: Activation function (`callable`). Set it to None to maintain a - linear activation. - activity_regularizer: 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: Python `str`, the name of the layer. Layers with the same name will - share `tf.Variable`s, but to avoid mistakes we require `reuse=True` in - such cases. - reuse: Python `bool`, whether to reuse the `tf.Variable`s of a previous - layer by the same name. - - Returns: - output: `Tensor` representing a the affine transformed input under a random - draw from the surrogate posterior distribution. - - #### 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 = tfp.layers.dense_flipout( - features, 512, activation=tf.nn.relu) - 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 = DenseFlipout( - units, - 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) diff --git a/tensorflow/contrib/bayesflow/python/ops/layers_util.py b/tensorflow/contrib/bayesflow/python/ops/layers_util.py deleted file mode 100644 index 8c1fb203f7328e8260e49b4326d813fbe133613e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/layers_util.py +++ /dev/null @@ -1,191 +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. -# ============================================================================== -"""Utilities for probabilistic layers. -""" - -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 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 - - -def default_loc_scale_fn( - is_singular=False, - loc_initializer=init_ops.random_normal_initializer(stddev=0.1), - untransformed_scale_initializer=init_ops.random_normal_initializer( - mean=-3., stddev=0.1), - loc_regularizer=None, - untransformed_scale_regularizer=None, - loc_constraint=None, - untransformed_scale_constraint=None): - """Makes closure which creates `loc`, `scale` params from `tf.get_variable`. - - This function produces a closure which produces `loc`, `scale` using - `tf.get_variable`. The closure accepts the following arguments: - - dtype: Type of parameter's event. - shape: Python `list`-like representing the parameter's event shape. - name: Python `str` name prepended to any created (or existing) - `tf.Variable`s. - trainable: Python `bool` indicating all created `tf.Variable`s should be - added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`. - add_variable_fn: `tf.get_variable`-like `callable` used to create (or - access existing) `tf.Variable`s. - - Args: - is_singular: Python `bool` indicating if `scale is None`. Default: `False`. - loc_initializer: Initializer function for the `loc` parameters. - The default is `tf.random_normal_initializer(mean=0., stddev=0.1)`. - untransformed_scale_initializer: Initializer function for the `scale` - parameters. Default value: `tf.random_normal_initializer(mean=-3., - stddev=0.1)`. This implies the softplus transformed result has mean - approximately `0.05` and std. deviation approximately `0.005`. - loc_regularizer: Regularizer function for the `loc` parameters. - The default (`None`) is to use the `tf.get_variable` default. - untransformed_scale_regularizer: Regularizer function for the `scale` - parameters. The default (`None`) is to use the `tf.get_variable` default. - loc_constraint: An optional projection function to be applied to the - loc after being updated by an `Optimizer`. The function must take as input - the unprojected variable and must return the projected variable (which - must have the same shape). Constraints are not safe to use when doing - asynchronous distributed training. - The default (`None`) is to use the `tf.get_variable` default. - untransformed_scale_constraint: An optional projection function to be - applied to the `scale` parameters after being updated by an `Optimizer` - (e.g. used to implement norm constraints or value constraints). The - function must take as input the unprojected variable and must return the - projected variable (which must have the same shape). Constraints are not - safe to use when doing asynchronous distributed training. The default - (`None`) is to use the `tf.get_variable` default. - - Returns: - default_loc_scale_fn: Python `callable` which instantiates `loc`, `scale` - parameters from args: `dtype, shape, name, trainable, add_variable_fn`. - """ - def _fn(dtype, shape, name, trainable, add_variable_fn): - """Creates `loc`, `scale` parameters.""" - loc = add_variable_fn( - name=name + "_loc", - shape=shape, - initializer=loc_initializer, - regularizer=loc_regularizer, - constraint=loc_constraint, - dtype=dtype, - trainable=trainable) - if is_singular: - return loc, None - untransformed_scale = add_variable_fn( - name=name + "_untransformed_scale", - shape=shape, - initializer=untransformed_scale_initializer, - regularizer=untransformed_scale_regularizer, - constraint=untransformed_scale_constraint, - dtype=dtype, - trainable=trainable) - scale = (np.finfo(dtype.as_numpy_dtype).eps + - nn_ops.softplus(untransformed_scale)) - return loc, scale - return _fn - - -def default_mean_field_normal_fn( - is_singular=False, - loc_initializer=None, - untransformed_scale_initializer=None, - loc_regularizer=None, - untransformed_scale_regularizer=None, - loc_constraint=None, - untransformed_scale_constraint=None): - """Creates a function to build Normal distributions with trainable params. - - This function produces a closure which produces `tf.distributions.Normal` - parameterized by a loc` and `scale` each created using `tf.get_variable`. The - produced closure accepts the following arguments: - - name: Python `str` name prepended to any created (or existing) - `tf.Variable`s. - shape: Python `list`-like representing the parameter's event shape. - dtype: Type of parameter's event. - trainable: Python `bool` indicating all created `tf.Variable`s should be - added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`. - add_variable_fn: `tf.get_variable`-like `callable` used to create (or - access existing) `tf.Variable`s. - - Args: - is_singular: Python `bool` if `True`, forces the special case limit of - `scale->0`, i.e., a `Deterministic` distribution. - loc_initializer: Initializer function for the `loc` parameters. - If `None` (default), values are initialized using the default - initializer used by `tf.get_variable`. - untransformed_scale_initializer: Initializer function for the `scale` - parameters. If `None` (default), values are initialized using the default - initializer used by `tf.get_variable`. - loc_regularizer: Regularizer function for the `loc` parameters. - untransformed_scale_regularizer: Regularizer function for the `scale` - parameters. - loc_constraint: An optional projection function to be applied to the - loc after being updated by an `Optimizer`. The function must take as input - the unprojected variable and must return the projected variable (which - must have the same shape). Constraints are not safe to use when doing - asynchronous distributed training. - untransformed_scale_constraint: An optional projection function to be - applied to the `scale` parameters after being updated by an `Optimizer` - (e.g. used to implement norm constraints or value constraints). The - function must take as input the unprojected variable and must return the - projected variable (which must have the same shape). Constraints are not - safe to use when doing asynchronous distributed training. - - Returns: - make_normal_fn: Python `callable` which creates a `tf.distributions.Normal` - using from args: `dtype, shape, name, trainable, add_variable_fn`. - """ - loc_scale_fn_ = default_loc_scale_fn( - is_singular, - loc_initializer, - untransformed_scale_initializer, - loc_regularizer, - untransformed_scale_regularizer, - loc_constraint, - untransformed_scale_constraint) - def _fn(dtype, shape, name, trainable, add_variable_fn): - """Creates multivariate `Deterministic` or `Normal` distribution.""" - loc, scale = loc_scale_fn_(dtype, shape, name, trainable, add_variable_fn) - if scale is None: - dist = deterministic_lib.Deterministic(loc=loc) - else: - dist = normal_lib.Normal(loc=loc, scale=scale) - reinterpreted_batch_ndims = array_ops.shape(dist.batch_shape_tensor())[0] - 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/metropolis_hastings.py b/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings.py index 7bdeaa862d5bb64fa8940df453c7aa2b66023eda..e7fcbc65ef379e84a140a06e020549f74f905a99 100644 --- a/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings.py +++ b/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings.py @@ -25,9 +25,10 @@ from tensorflow.contrib.bayesflow.python.ops.metropolis_hastings_impl import * from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ + 'kernel', 'evolve', - 'uniform_random_proposal', - 'normal_random_proposal', + 'proposal_uniform', + 'proposal_normal', ] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings_impl.py b/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings_impl.py index dc1ac68ce009fa46d6c05a3200a29d9fdf245707..05aa134ed5c11092316af5f3e45ba07fdb491e90 100644 --- a/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/metropolis_hastings_impl.py @@ -12,17 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Functions to create a Markov Chain Monte Carlo Metropolis step. +"""Metropolis-Hastings and proposal distributions. +@@kernel @@evolve -@@uniform_random_proposal -@@normal_random_proposal +@@proposal_uniform +@@proposal_normal """ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections + from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -31,123 +34,198 @@ from tensorflow.python.ops import random_ops from tensorflow.python.ops import state_ops __all__ = [ - 'evolve', - 'uniform_random_proposal', - 'normal_random_proposal', + "kernel", + "evolve", + "proposal_uniform", + "proposal_normal", ] -def _single_iteration(current_state, current_log_density, - log_unnormalized_prob_fn, proposal_fn, seed=None, - name='None'): - """Performs a single Metropolis-Hastings step. +KernelResults = collections.namedtuple( + "KernelResults", + [ + "log_accept_ratio", + "current_target_log_prob", # "Current result" means "accepted". + "is_accepted", + "proposed_state", + ]) + + +def kernel(target_log_prob_fn, + proposal_fn, + current_state, + seed=None, + current_target_log_prob=None, + name=None): + """Runs the Metropolis-Hastings transition kernel. + + 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))`.) Args: - current_state: Float-like `Tensor` (i.e., `dtype` is either - `tf.float16`, `tf.float32` or `tf.float64`) of any shape that can - be consumed by the `log_unnormalized_prob_fn` and `proposal_fn` - callables. - current_log_density: Float-like `Tensor` with `dtype` and shape equivalent - to `log_unnormalized_prob_fn(current_state)`, i.e., matching the result of - `log_unnormalized_prob_fn` invoked at `current_state`. - log_unnormalized_prob_fn: A Python callable evaluated at - `current_state` and returning a float-like `Tensor` of log target-density - up to a normalizing constant. In other words, - `log_unnormalized_prob_fn(x) = log(g(x))`, where - `target_density = g(x)/Z` for some constant `A`. The shape of the input - tensor is the same as the shape of the `current_state`. The shape of the - output tensor is either - (a). Same as the input shape if the density being sampled is one - dimensional, or - (b). If the density is defined for `events` of shape - `event_shape = [E1, E2, ... Ee]`, then the input tensor should be of - shape `batch_shape + event_shape`, where `batch_shape = [B1, ..., Bb]` - and the result must be of shape [B1, ..., Bb]. For example, if the - distribution that is being sampled is a 10 dimensional normal, - then the input tensor may be of shape [100, 10] or [30, 20, 10]. The - last dimension will then be 'consumed' by `log_unnormalized_prob_fn` - and it should return tensors of shape [100] and [30, 20] respectively. - proposal_fn: A callable accepting a real valued `Tensor` of current sample - points and returning a tuple of two `Tensors`. The first element of the - pair is a `Tensor` containing the proposal state and should have - the same shape as the input `Tensor`. The second element of the pair gives - the log of the ratio of the probability of transitioning from the - proposal points to the input points and the probability of transitioning - from the input points to the proposal points. If the proposal is - symmetric (e.g., random walk, where the proposal is either - normal or uniform centered at `current_state`), i.e., - Probability(Proposal -> Current) = Probability(Current -> Proposal) - the second value should be set to `None` instead of explicitly supplying a - tensor of zeros. In addition to being convenient, this also leads to a - more efficient graph. - seed: `int` or None. The random seed for this `Op`. If `None`, no seed is - applied. - name: Python `str` name prefix for ops managed by this function. + 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. + proposal_fn: Python callable which takes an argument like `current_state` + (or `*current_state` if it's a list) and returns a tuple of proposed + states of same shape as `state`, and a log ratio `Tensor` of same shape + as `current_target_log_prob`. The log ratio is the log-probability of + `state` given proposed states minus the log-probability of proposed + states given `state`. If the proposal is symmetric, set the second value + to `None`: this enables more efficient computation than explicitly + supplying a tensor of zeros. + 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))`. + 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). + name: A name of the operation (optional). Returns: - next_state: `Tensor` with `dtype` and shape matching `current_state`. - Created by propagating the chain by one step, starting from + next_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`. - next_log_density: `Tensor` with `dtype` and shape matching - `current_log_density`, which is equal to the value of the unnormalized - `log_unnormalized_prob_fn` computed at `next_state`. - log_accept_ratio: `Tensor` with `dtype` and shape matching - `current_log_density`. Stands for the log of Metropolis-Hastings - acceptance ratio used in generating the `next_state`. - """ + kernel_results: `collections.namedtuple` of internal calculations used to + advance the chain. - with ops.name_scope(name, 'single_iteration', [current_state]): - # The proposed state and the log of the corresponding Hastings ratio. - proposal_state, log_transit_ratio = proposal_fn(current_state) - - # If the log ratio is None, assume that the transitions are symmetric, - # i.e., Prob(Current -> Proposed) = Prob(Proposed -> Current). - if log_transit_ratio is None: - log_transit_ratio = 0. - - # Log-density of the proposal state. - proposal_log_density = log_unnormalized_prob_fn(proposal_state) - - # Ops to compute the log of the acceptance ratio. Recall that the - # acceptance ratio is: [Prob(Proposed) / Prob(Current)] * - # [Prob(Proposed -> Current) / Prob(Current -> Proposed)]. The log of the - # second term is the log_transit_ratio. - with ops.name_scope('accept_reject'): - # The log of the acceptance ratio. - log_accept_ratio = (proposal_log_density - current_log_density - + log_transit_ratio) - - # A proposal is accepted or rejected depending on the acceptance ratio. - # If the acceptance ratio is greater than 1 then it is always accepted. - # If the acceptance ratio is less than 1 then the proposal is accepted - # with probability = acceptance ratio. As we are working in log space to - # prevent over/underflows, this logic is expressed in log terms below. - # If a proposal is accepted we place a True in the acceptance state - # tensor and if it is to be rejected we place a False. - # The log_draws below have to be compared to the log_accept_ratio so we - # make sure that they have the same data type. - log_draws = math_ops.log(random_ops.random_uniform( - array_ops.shape(current_log_density), seed=seed, - dtype=log_accept_ratio.dtype)) - is_proposal_accepted = log_draws < log_accept_ratio - - # The acceptance state decides which elements of the current state are to - # be replaced with the corresponding elements in the proposal state. - with ops.name_scope(name, 'metropolis_single_step', - [current_state, current_log_density]): - next_log_density = array_ops.where(is_proposal_accepted, - proposal_log_density, - current_log_density) - next_state = array_ops.where(is_proposal_accepted, proposal_state, - current_state) - - return next_state, next_log_density, log_accept_ratio + #### Examples + + We illustrate Metropolis-Hastings on a Normal likelihood with + unknown mean. + + ```python + tfd = tf.contrib.distributions + tfp = tf.contrib.bayesflow + + loc = tf.get_variable("loc", initializer=1.) + x = tf.constant([0.0] * 50) + + def make_target_log_prob_fn(x): + def target_log_prob_fn(loc): + prior = tfd.Normal(loc=0., scale=1.) + likelihood = tfd.Independent( + tfd.Normal(loc=loc, scale=0.1), + reinterpreted_batch_ndims=1) + return prior.log_prob(loc) + likelihood.log_prob(x) + return target_log_prob_fn + + next_state, kernel_results = tfp.metropolis_hastings.kernel( + target_log_prob_fn=make_target_log_prob_fn(x), + proposal_fn=tfp.metropolis_hastings.proposal_normal(), + current_state=loc) + loc_update = loc.assign(next_state) + ``` + + We illustrate Metropolis-Hastings on a Normal likelihood with + unknown mean and variance. We apply 4 chains. + + ```python + tfd = tf.contrib.distributions + tfp = tf.contrib.bayesflow + + num_chains = 4 + loc = tf.get_variable("loc", shape=[num_chains], + initializer=tf.random_normal_initializer()) + scale = tf.get_variable("scale", shape=[num_chains], + initializer=tf.ones_initializer()) + x = tf.constant([0.0] * 50) + + def make_target_log_prob_fn(x): + data = tf.reshape(x, shape=[-1, 1]) + def target_log_prob_fn(loc, scale): + prior_loc = tfd.Normal(loc=0., scale=1.) + prior_scale = tfd.InverseGamma(concentration=1., rate=1.) + likelihood = tfd.Independent( + tfd.Normal(loc=loc, scale=scale), + reinterpreted_batch_ndims=1) + return (prior_loc.log_prob(loc) + + prior_scale.log_prob(scale) + + likelihood.log_prob(data)) + return target_log_prob_fn + + def proposal_fn(loc, scale): + loc_proposal = tfp.metropolis_hastings.proposal_normal() + scale_proposal = tfp.metropolis_hastings.proposal_uniform(minval=-1.) + proposed_loc, _ = loc_proposal(loc) + proposed_scale, _ = scale_proposal(scale) + proposed_scale = tf.maximum(proposed_scale, 0.01) + return [proposed_loc, proposed_scale], None + + next_state, kernel_results = tfp.metropolis_hastings.kernel( + target_log_prob_fn=make_target_log_prob_fn(x), + proposal_fn=proposal_fn, + current_state=[loc, scale]) + train_op = tf.group(loc.assign(next_state[0]), + scale.assign(next_state[1])) + ``` + + """ + with ops.name_scope( + name, "metropolis_hastings_kernel", + [current_state, seed, current_target_log_prob]): + with ops.name_scope("initialize"): + maybe_expand = lambda x: list(x) if _is_list_like(x) else [x] + current_state_parts = maybe_expand(current_state) + if current_target_log_prob is None: + current_target_log_prob = target_log_prob_fn(*current_state_parts) + + proposed_state, log_transit_ratio = proposal_fn(*current_state_parts) + proposed_state_parts = maybe_expand(proposed_state) + + proposed_target_log_prob = target_log_prob_fn(*proposed_state_parts) + + with ops.name_scope( + "accept_reject", + [current_state_parts, proposed_state_parts, + current_target_log_prob, proposed_target_log_prob]): + log_accept_ratio = proposed_target_log_prob - current_target_log_prob + if log_transit_ratio is not None: + # If the log_transit_ratio is None, then assume the proposal is + # symmetric, i.e., + # log p(old | new) - log p(new | old) = 0. + log_accept_ratio += log_transit_ratio + + # u < exp(log_accept_ratio), where u~Uniform[0,1) + # ==> log(u) < log_accept_ratio + random_value = random_ops.random_uniform( + array_ops.shape(log_accept_ratio), + dtype=log_accept_ratio.dtype, + seed=seed) + random_negative = math_ops.log(random_value) + is_accepted = random_negative < log_accept_ratio + next_state_parts = [array_ops.where(is_accepted, + proposed_state_part, + current_state_part) + for proposed_state_part, current_state_part in + zip(proposed_state_parts, current_state_parts)] + accepted_log_prob = array_ops.where(is_accepted, + proposed_target_log_prob, + current_target_log_prob) + maybe_flatten = lambda x: x if _is_list_like(current_state) else x[0] + return [ + maybe_flatten(next_state_parts), + KernelResults( + log_accept_ratio=log_accept_ratio, + current_target_log_prob=accepted_log_prob, + is_accepted=is_accepted, + proposed_state=maybe_flatten(proposed_state_parts), + ), + ] def evolve(initial_sample, initial_log_density, initial_log_accept_ratio, - log_unnormalized_prob_fn, + target_log_prob_fn, proposal_fn, n_steps=1, seed=None, @@ -162,9 +240,11 @@ def evolve(initial_sample, The probability distribution may have an unknown normalization constan. We parameterize the probability density as follows: - ``` - f(x) = exp(L(x) + constant) - ``` + + ```none + f(x) = exp(L(x) + constant) + ``` + Here `L(x)` is any continuous function with an (possibly unknown but finite) upper bound, i.e. there exists a number beta such that `L(x)< beta < infinity` for all x. The constant is the normalization needed @@ -188,72 +268,77 @@ def evolve(initial_sample, The following example, demonstrates the use to generate a 1000 uniform random walk Metropolis samplers run in parallel for the normal target distribution. + ```python - n = 3 # dimension of the problem - - # Generate 1000 initial values randomly. Each of these would be an - # independent starting point for a Markov chain. - state = tf.get_variable( - 'state',initializer=tf.random_normal([1000, n], mean=3.0, - dtype=tf.float64, seed=42)) - - # Computes the log(p(x)) for the unit normal density and ignores the - # normalization constant. - def log_density(x): - return - tf.reduce_sum(x * x, reduction_indices=-1) / 2.0 - - # Initial log-density value - state_log_density = tf.get_variable( - 'state_log_density', initializer=log_density(state.initialized_value())) - - # A variable to store the log_acceptance_ratio: - log_acceptance_ratio = tf.get_variable( - 'log_acceptance_ratio', initializer=tf.zeros([1000], dtype=tf.float64)) - - # Generates random proposals by moving each coordinate uniformly and - # independently in a box of size 2 centered around the current value. - # Returns the new point and also the log of the Hastings ratio (the - # ratio of the probability of going from the proposal to origin and the - # probability of the reverse transition). When this ratio is 1, the value - # may be omitted and replaced by None. - def random_proposal(x): - return (x + tf.random_uniform(tf.shape(x), minval=-1, maxval=1, - dtype=x.dtype, seed=12)), None - - # Create the op to propagate the chain for 100 steps. - stepper = mh.evolve( - state, state_log_density, log_acceptance_ratio, - log_density, random_proposal, n_steps=100, seed=123) - init = tf.initialize_all_variables() - with tf.Session() as sess: - sess.run(init) - # Run the chains for a total of 1000 steps and print out the mean across - # the chains every 100 iterations. - for n_iter in range(10): - # Executing the stepper advances the chain to the next state. - sess.run(stepper) - # Print out the current value of the mean(sample) for every dimension. - print(np.mean(sess.run(state), 0)) - # Estimated covariance matrix - samples = sess.run(state) - print('') - print(np.cov(samples, rowvar=False)) + n = 3 # dimension of the problem + + # Generate 1000 initial values randomly. Each of these would be an + # independent starting point for a Markov chain. + state = tf.get_variable( + "state", + initializer=tf.random_normal([1000, n], + mean=3.0, + dtype=tf.float64, + seed=42)) + + # Computes the log(p(x)) for the unit normal density and ignores the + # normalization constant. + def log_density(x): + return -tf.reduce_sum(x * x, reduction_indices=-1) / 2.0 + + # Initial log-density value + state_log_density = tf.get_variable( + "state_log_density", + initializer=log_density(state.initialized_value())) + + # A variable to store the log_acceptance_ratio: + log_acceptance_ratio = tf.get_variable( + "log_acceptance_ratio", + initializer=tf.zeros([1000], dtype=tf.float64)) + + # Generates random proposals by moving each coordinate uniformly and + # independently in a box of size 2 centered around the current value. + # Returns the new point and also the log of the Hastings ratio (the + # ratio of the probability of going from the proposal to origin and the + # probability of the reverse transition). When this ratio is 1, the value + # may be omitted and replaced by None. + def random_proposal(x): + return (x + tf.random_uniform(tf.shape(x), minval=-1, maxval=1, + dtype=x.dtype, seed=12)), None + + # Create the op to propagate the chain for 100 steps. + stepper = mh.evolve( + state, state_log_density, log_acceptance_ratio, + log_density, random_proposal, n_steps=100, seed=123) + init = tf.initialize_all_variables() + with tf.Session() as sess: + sess.run(init) + # Run the chains for a total of 1000 steps and print out the mean across + # the chains every 100 iterations. + for n_iter in range(10): + # Executing the stepper advances the chain to the next state. + sess.run(stepper) + # Print out the current value of the mean(sample) for every dimension. + print(np.mean(sess.run(state), 0)) + # Estimated covariance matrix + samples = sess.run(state) + print(np.cov(samples, rowvar=False)) ``` Args: initial_sample: A float-like `tf.Variable` of any shape that can - be consumed by the `log_unnormalized_prob_fn` and `proposal_fn` + be consumed by the `target_log_prob_fn` and `proposal_fn` callables. initial_log_density: Float-like `tf.Variable` with `dtype` and shape - equivalent to `log_unnormalized_prob_fn(initial_sample)`, i.e., matching - the result of `log_unnormalized_prob_fn` invoked at `current_state`. + equivalent to `target_log_prob_fn(initial_sample)`, i.e., matching + the result of `target_log_prob_fn` invoked at `current_state`. initial_log_accept_ratio: A `tf.Variable` with `dtype` and shape matching `initial_log_density`. Stands for the log of Metropolis-Hastings acceptance ratio after propagating the chain for `n_steps`. - log_unnormalized_prob_fn: A Python callable evaluated at + target_log_prob_fn: A Python callable evaluated at `current_state` and returning a float-like `Tensor` of log target-density up to a normalizing constant. In other words, - `log_unnormalized_prob_fn(x) = log(g(x))`, where + `target_log_prob_fn(x) = log(g(x))`, where `target_density = g(x)/Z` for some constant `A`. The shape of the input tensor is the same as the shape of the `current_state`. The shape of the output tensor is either @@ -265,7 +350,7 @@ def evolve(initial_sample, and the result must be of shape [B1, ..., Bb]. For example, if the distribution that is being sampled is a 10 dimensional normal, then the input tensor may be of shape [100, 10] or [30, 20, 10]. The - last dimension will then be 'consumed' by `log_unnormalized_prob_fn` + last dimension will then be 'consumed' by `target_log_prob_fn` and it should return tensors of shape [100] and [30, 20] respectively. proposal_fn: A callable accepting a real valued `Tensor` of current sample points and returning a tuple of two `Tensors`. The first element of the @@ -289,42 +374,48 @@ def evolve(initial_sample, forward_step: an `Op` to step the Markov chain forward for `n_steps`. """ - with ops.name_scope(name, 'metropolis_hastings', [initial_sample]): + with ops.name_scope(name, "metropolis_hastings", [initial_sample]): current_state = initial_sample - current_log_density = initial_log_density + current_target_log_prob = initial_log_density log_accept_ratio = initial_log_accept_ratio - # Stop condition for the while_loop - def stop_condition(i, _): - return i < n_steps - - def step(i, loop_vars): - """Wrap `_single_iteration` for `while_loop`.""" - state = loop_vars[0] - state_log_density = loop_vars[1] - return i + 1, list(_single_iteration(state, state_log_density, - log_unnormalized_prob_fn, - proposal_fn, seed=seed)) - - loop_vars = [current_state, current_log_density, log_accept_ratio] - # Build an `Op` to evolve the Markov chain for `n_steps` - (_, [end_state, end_log_density, end_log_acceptance]) = ( + def step(i, current_state, current_target_log_prob, log_accept_ratio): + """Wrap single Markov chain iteration in `while_loop`.""" + next_state, kernel_results = kernel( + target_log_prob_fn=target_log_prob_fn, + proposal_fn=proposal_fn, + current_state=current_state, + current_target_log_prob=current_target_log_prob, + seed=seed) + accepted_log_prob = kernel_results.current_target_log_prob + log_accept_ratio = kernel_results.log_accept_ratio + return i + 1, next_state, accepted_log_prob, log_accept_ratio + + (_, accepted_state, accepted_target_log_prob, accepted_log_accept_ratio) = ( control_flow_ops.while_loop( - stop_condition, step, - (0, loop_vars), - parallel_iterations=1, swap_memory=1)) + cond=lambda i, *ignored_args: i < n_steps, + body=step, + loop_vars=[ + 0, # i + current_state, + current_target_log_prob, + log_accept_ratio, + ], + parallel_iterations=1 if seed is not None else 10, + # TODO(b/73775595): Confirm optimal setting of swap_memory. + swap_memory=1)) forward_step = control_flow_ops.group( - state_ops.assign(current_log_density, end_log_density), - state_ops.assign(current_state, end_state), - state_ops.assign(log_accept_ratio, end_log_acceptance)) + state_ops.assign(current_target_log_prob, accepted_target_log_prob), + state_ops.assign(current_state, accepted_state), + state_ops.assign(log_accept_ratio, accepted_log_accept_ratio)) return forward_step -def uniform_random_proposal(step_size=1., - seed=None, - name=None): +def proposal_uniform(step_size=1., + seed=None, + name=None): """Returns a callable that adds a random uniform tensor to the input. This function returns a callable that accepts one `Tensor` argument of any @@ -346,11 +437,13 @@ def uniform_random_proposal(step_size=1., Returns: proposal_fn: A callable accepting one float-like `Tensor` and returning a - 2-tuple. The first value in the tuple is a `Tensor` of the same shape and - dtype as the input argument and the second element of the tuple is None. + 2-tuple. The first value in the tuple is a `Tensor` of the same shape and + dtype as the input argument and the second element of the tuple is None. """ - with ops.name_scope(name, 'uniform_random_proposal', [step_size]): + with ops.name_scope(name, "proposal_uniform", [step_size]): + step_size = ops.convert_to_tensor(step_size, name="step_size") + def proposal_fn(input_state, name=None): """Adds a uniform perturbation to the input state. @@ -359,12 +452,12 @@ def uniform_random_proposal(step_size=1., name: A string that sets the name for this `Op`. Returns: - proposal_state: A float-like `Tensot` with `dtype` and shape matching + proposal_state: A float-like `Tensor` with `dtype` and shape matching `input_state`. log_transit_ratio: `None`. Proposal is symmetric. """ - with ops.name_scope(name, 'proposer', [input_state]): - input_state = ops.convert_to_tensor(input_state, name='input_state') + with ops.name_scope(name, "proposer", [input_state]): + input_state = ops.convert_to_tensor(input_state, name="input_state") return input_state + random_ops.random_uniform( array_ops.shape(input_state), minval=-step_size, @@ -373,9 +466,9 @@ def uniform_random_proposal(step_size=1., return proposal_fn -def normal_random_proposal(scale=1., - seed=None, - name=None): +def proposal_normal(scale=1., + seed=None, + name=None): """Returns a callable that adds a random normal tensor to the input. This function returns a callable that accepts one `Tensor` argument of any @@ -398,11 +491,13 @@ def normal_random_proposal(scale=1., Returns: proposal_fn: A callable accepting one float-like `Tensor` and returning a - 2-tuple. The first value in the tuple is a `Tensor` of the same shape and - dtype as the input argument and the second element of the tuple is None. + 2-tuple. The first value in the tuple is a `Tensor` of the same shape and + dtype as the input argument and the second element of the tuple is None. """ - with ops.name_scope(name, 'normal_random_proposal', [scale]): + with ops.name_scope(name, "proposal_normal", [scale]): + scale = ops.convert_to_tensor(scale, name="scale") + def proposal_fn(input_state, name=None): """Adds a normal perturbation to the input state. @@ -411,16 +506,22 @@ def normal_random_proposal(scale=1., name: A string that sets the name for this `Op`. Returns: - proposal_state: A float-like `Tensot` with `dtype` and shape matching + proposal_state: A float-like `Tensor` with `dtype` and shape matching `input_state`. log_transit_ratio: `None`. Proposal is symmetric. """ - with ops.name_scope(name, 'proposer', [input_state]): - input_state = ops.convert_to_tensor(input_state, name='input_state') + with ops.name_scope(name, "proposer", [input_state]): + input_state = ops.convert_to_tensor(input_state, name="input_state") return input_state + random_ops.random_normal( array_ops.shape(input_state), mean=0., stddev=scale, + dtype=scale.dtype, seed=seed), None return proposal_fn + + +def _is_list_like(x): + """Helper which returns `True` if input is `list`-like.""" + return isinstance(x, (tuple, list)) diff --git a/tensorflow/contrib/bayesflow/python/ops/sgld_optimizer.py b/tensorflow/contrib/bayesflow/python/ops/sgld_optimizer.py deleted file mode 100644 index 7786656398e3c87704227be95b3cd23a38785249..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/sgld_optimizer.py +++ /dev/null @@ -1,220 +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. -# ============================================================================== -"""An optimizer module for stochastic gradient Langevin dynamics.""" -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 check_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import random_ops -from tensorflow.python.ops import variable_scope as varscope_ops -from tensorflow.python.training import optimizer -from tensorflow.python.training import training_ops - - -class SGLDOptimizer(optimizer.Optimizer): - """An optimizer module for stochastic gradient Langevin dynamics. - - This implements the preconditioned Stochastic Gradient Langevin Dynamics - optimizer [1]. The optimization variable is regarded as a sample from the - posterior under Stochastic Gradient Langevin Dynamics with noise rescaled in - each dimension according to RMSProp [2]. - - Note: If a prior is included in the loss, it should be scaled by - `1/num_pseudo_batches`, where num_pseudo_batches is the number of minibatches - in the data. I.e., it should be divided by the `num_pseudo_batches` term - described below. - - [1]: "Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural - Networks." Chunyuan Li, Changyou Chen, David Carlson, Lawrence Carin. - ArXiv:1512.07666, 2015. https://arxiv.org/abs/1512.07666 - [2]: http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf - - Args: - learning_rate: Scalar `float`-like `Tensor`. The base learning rate for the - optimizer. Must be tuned to the specific function being minimized. - preconditioner_decay_rate: Scalar `float`-like `Tensor`. The exponential - decay rate of the rescaling of the preconditioner (RMSprop). (This is - "alpha" in [1]). Should be smaller than but nearly `1` to approximate - sampling from the posterior. (Default: `0.95`) - num_pseudo_batches: Scalar `int`-like `Tensor`. The effective number of - minibatches in the data set. Trades off noise and prior with the SGD - likelihood term. Note: Assumes the loss is taken as the mean over a - minibatch. Otherwise if the sum was taken, divide this number by the - batch size. (Default: `1`) - burnin: Scalar `int`-like `Tensor`. The number of iterations to collect - gradient statistics to update the preconditioner before starting to draw - noisy samples. (Default: `25`) - diagonal_bias: Scalar `float`-like `Tensor`. Term added to the diagonal of - the preconditioner to prevent the preconditioner from degenerating. - (Default: `1e-8`) - name: Python `str` describing ops managed by this function. - (Default: `"SGLDOptimizer"`) - variable_scope: Variable scope used for calls to `tf.get_variable`. - If `None`, a new variable scope is created using name - `ops.get_default_graph().unique_name(name or default_name)`. - - Raises: - InvalidArgumentError: If preconditioner_decay_rate is a `Tensor` not in - `(0,1]`. - """ - - def __init__(self, - learning_rate, - preconditioner_decay_rate=0.95, - num_pseudo_batches=1, - burnin=25, - diagonal_bias=1e-8, - name=None, - variable_scope=None): - default_name = 'SGLDOptimizer' - with ops.name_scope(name, default_name, [ - learning_rate, preconditioner_decay_rate, num_pseudo_batches, burnin, - diagonal_bias - ]): - if variable_scope is None: - var_scope_name = ops.get_default_graph().unique_name( - name or default_name) - with varscope_ops.variable_scope(var_scope_name) as scope: - self._variable_scope = scope - else: - self._variable_scope = variable_scope - - self._preconditioner_decay_rate = ops.convert_to_tensor( - preconditioner_decay_rate, name='preconditioner_decay_rate') - self._num_pseudo_batches = ops.convert_to_tensor( - num_pseudo_batches, name='num_pseudo_batches') - self._burnin = ops.convert_to_tensor(burnin, name='burnin') - self._diagonal_bias = ops.convert_to_tensor( - diagonal_bias, name='diagonal_bias') - self._learning_rate = ops.convert_to_tensor( - learning_rate, name='learning_rate') - - with varscope_ops.variable_scope(self._variable_scope): - self._counter = varscope_ops.get_variable( - 'counter', initializer=0, trainable=False) - - self._preconditioner_decay_rate = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._preconditioner_decay_rate, - message='`preconditioner_decay_rate` must be non-negative'), - check_ops.assert_less_equal( - self._preconditioner_decay_rate, - 1., - message='`preconditioner_decay_rate` must be at most 1.'), - ], self._preconditioner_decay_rate) - - self._num_pseudo_batches = control_flow_ops.with_dependencies([ - check_ops.assert_greater( - self._num_pseudo_batches, - 0, - message='`num_pseudo_batches` must be greater than zero') - ], self._num_pseudo_batches) - - self._burnin = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._burnin, message='`burnin` must be non-negative'), - check_ops.assert_integer( - self._burnin, message='`burnin` must be an integer') - ], self._burnin) - - self._diagonal_bias = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._diagonal_bias, - message='`diagonal_bias` must be non-negative') - ], self._diagonal_bias) - - super(SGLDOptimizer, self).__init__(use_locking=False, - name=name or default_name) - - def _create_slots(self, var_list): - for v in var_list: - init_rms = init_ops.ones_initializer(dtype=v.dtype) - self._get_or_make_slot_with_initializer(v, init_rms, v.get_shape(), - v.dtype, 'rms', self._name) - - def _prepare(self): - # We need to put the conversion and check here because a user will likely - # want to decay the learning rate dynamically. - self._learning_rate_tensor = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._learning_rate, message='`learning_rate` must be non-negative') - ], ops.convert_to_tensor(self._learning_rate, name='learning_rate_tensor')) - self._decay_tensor = ops.convert_to_tensor( - self._preconditioner_decay_rate, name='preconditioner_decay_rate') - - super(SGLDOptimizer, self)._prepare() - - def _apply_dense(self, grad, var): - rms = self.get_slot(var, 'rms') - - with ops.control_dependencies([ - self._update_momentum(rms, grad, math_ops.cast(self._decay_tensor, - var.dtype.base_dtype))]): - new_grad = self._apply_noisy_update(rms, grad) - - return training_ops.apply_gradient_descent( - var, - math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype), - new_grad, - use_locking=self._use_locking).op - - def _apply_sparse(self, grad, var): - rms = self.get_slot(var, 'rms') - - with ops.control_dependencies([ - self._update_momentum(rms, grad, math_ops.cast(self._decay_tensor, - var.dtype.base_dtype))]): - new_grad = self._apply_noisy_update(rms, grad) - - return training_ops.apply_gradient_descent( - var, - math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype), - new_grad, - use_locking=self._use_locking).op - - def _finish(self, update_ops, name_scope): - update_ops.append([self._counter.assign_add(1)]) - return control_flow_ops.group(*update_ops, name=name_scope) - - @property - def variable_scope(self): - """Variable scope of all calls to `tf.get_variable`.""" - return self._variable_scope - - def _apply_noisy_update(self, mom, grad): - # Compute and apply the gradient update following - # preconditioned Langevin dynamics - stddev = array_ops.where( - array_ops.squeeze(self._counter > self._burnin), - math_ops.cast(math_ops.rsqrt(self._learning_rate), grad.dtype), - array_ops.zeros([], grad.dtype)) - - preconditioner = math_ops.rsqrt( - mom + math_ops.cast(self._diagonal_bias, grad.dtype)) - return ( - 0.5 * preconditioner * grad * math_ops.cast(self._num_pseudo_batches, - grad.dtype) + - random_ops.random_normal(array_ops.shape(grad), 1.0, dtype=grad.dtype) * - stddev * math_ops.sqrt(preconditioner)) - - def _update_momentum(self, mom, grad, decay): - # Keep an exponentially weighted moving average of squared gradients. - # Not thread safe - return mom.assign_add((1.0 - decay) * (math_ops.square(grad) - mom)) diff --git a/tensorflow/contrib/bayesflow/python/ops/variational_sgd_optimizer.py b/tensorflow/contrib/bayesflow/python/ops/variational_sgd_optimizer.py deleted file mode 100644 index 4d5f0cfe9713a011b32c5aba8d429847d81f33e2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/bayesflow/python/ops/variational_sgd_optimizer.py +++ /dev/null @@ -1,279 +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. -# ============================================================================== -"""An optimizer module for constant stochastic gradient descent.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from tensorflow.python.framework import errors -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 clip_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import state_ops -from tensorflow.python.ops import variable_scope as varscope_ops -from tensorflow.python.training import optimizer -from tensorflow.python.training import training_ops - - -class VariationalSGDOptimizer(optimizer.Optimizer): - """An optimizer module for constant stochastic gradient descent. - - This implements an optimizer module for the constant stochastic gradient - descent algorithm [1]. The optimization variable is regarded as an - approximate sample from the posterior . - - Note: If a prior is included in the loss, it should be scaled by - `1/num_pseudo_batches`, where num_pseudo_batches is the number of minibatches - in the data. I.e., it should be divided by the `num_pseudo_batches` term - described below. - - [1]: "Stochastic Gradient Descent as Approximate Bayesian Inference - Stephan Mandt, Matthew D. Hoffman, David M. Blei. - ArXiv:1704.04289, 2017. https://arxiv.org/abs/1704.04289 - - Args: - batch_size: Scalar `int`-like `Tensor`. The number of examples in a - minibatch in the data set. Note: Assumes the loss is taken as the mean - over a minibatch. Otherwise if the sum was taken set this to 1. - total_num_examples: Scalar `int`-like `Tensor`. The total number of examples - in the data set. - max_learning_rate: Scalar `float`-like `Tensor`. A maximum allowable - effective coordinate-wise learning rate. The algorithm scales down any - effective learning rate (i.e. after preconditioning) that is larger than - this. (Default: `1`) - preconditioner_decay_rate: Scalar `float`-like `Tensor`. The exponential - decay rate of the rescaling of the preconditioner (RMSprop). (This is - "alpha" in [1]). Should be smaller than but nearly `1` to approximate - sampling from the posterior. (Default: `0.95`) - burnin: Scalar `int`-like `Tensor`. The number of iterations to collect - gradient statistics to update the preconditioner before starting to draw - noisy samples. (Default: `25`) - burnin_max_learning_rate: Scalar `float`-like `Tensor`. Maximum learning - rate to use during the burnin period. - (Default: `1e-8`) - use_single_learning_rate: Boolean Indicates whether one single learning - rate is used or coordinate_wise learning rates are used. - (Default: `False`) - name: Python `str` describing ops managed by this function. - (Default: `"VariationalSGDOptimizer"`) - variable_scope: Variable scope used for calls to `tf.get_variable`. - If `None`, a new variable scope is created using name - `ops.get_default_graph().unique_name(name or default_name)`. - - Raises: - InvalidArgumentError: If preconditioner_decay_rate is a `Tensor` not in - `(0,1]`. - """ - - def __init__(self, - batch_size, - total_num_examples, - max_learning_rate=1.0, - preconditioner_decay_rate=0.95, - burnin=25, - burnin_max_learning_rate=1e-6, - use_single_learning_rate=False, - name=None, - variable_scope=None): - default_name = 'VariationalSGDOptimizer' - with ops.name_scope(name, default_name, [ - max_learning_rate, preconditioner_decay_rate, batch_size, burnin, - burnin_max_learning_rate - ]): - if variable_scope is None: - var_scope_name = ops.get_default_graph().unique_name( - name or default_name) - with varscope_ops.variable_scope(var_scope_name) as scope: - self._variable_scope = scope - else: - self._variable_scope = variable_scope - - self._preconditioner_decay_rate = ops.convert_to_tensor( - preconditioner_decay_rate, name='preconditioner_decay_rate') - self._batch_size = ops.convert_to_tensor(batch_size, name='batch_size') - self._total_num_examples = ops.convert_to_tensor( - total_num_examples, name='total_num_examples') - self._burnin = ops.convert_to_tensor(burnin, name='burnin') - self._burnin_max_learning_rate = ops.convert_to_tensor( - burnin_max_learning_rate, name='burnin_max_learning_rate') - self._max_learning_rate = ops.convert_to_tensor( - max_learning_rate, name='max_learning_rate') - self._use_single_learning_rate = use_single_learning_rate - - with varscope_ops.variable_scope(self._variable_scope): - self._counter = varscope_ops.get_variable( - 'counter', initializer=0, trainable=False) - - self._preconditioner_decay_rate = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._preconditioner_decay_rate, - message='`preconditioner_decay_rate` must be non-negative'), - check_ops.assert_less_equal( - self._preconditioner_decay_rate, - 1., - message='`preconditioner_decay_rate` must be at most 1.'), - ], self._preconditioner_decay_rate) - - self._batch_size = control_flow_ops.with_dependencies([ - check_ops.assert_greater( - self._batch_size, - 0, - message='`batch_size` must be greater than zero') - ], self._batch_size) - - self._total_num_examples = control_flow_ops.with_dependencies([ - check_ops.assert_greater( - self._total_num_examples, - 0, - message='`total_num_examples` must be greater than zero') - ], self._total_num_examples) - - self._burnin = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._burnin, message='`burnin` must be non-negative'), - check_ops.assert_integer( - self._burnin, message='`burnin` must be an integer') - ], self._burnin) - - self._burnin_max_learning_rate = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._burnin_max_learning_rate, - message='`burnin_max_learning_rate` must be non-negative') - ], self._burnin_max_learning_rate) - - self._max_learning_rate = control_flow_ops.with_dependencies([ - check_ops.assert_non_negative( - self._max_learning_rate, - message='`max_learning_rate` must be non-negative') - ], self._max_learning_rate) - - super(VariationalSGDOptimizer, self).__init__( - use_locking=False, name=name or default_name) - - def _create_slots(self, var_list): - for v in var_list: - init_moment = init_ops.zeros_initializer(dtype=v.dtype) - self._get_or_make_slot_with_initializer( - v, init_moment, v.get_shape(), v.dtype, 'first_moment', self._name) - self._get_or_make_slot_with_initializer( - v, init_moment, v.get_shape(), v.dtype, 'second_moment', self._name) - - def _prepare(self): - self._decay_tensor = ops.convert_to_tensor( - self._preconditioner_decay_rate, name='preconditioner_decay_rate') - self._batch_size_tensor = ops.convert_to_tensor( - self._batch_size, name='batch_size_tensor') - - super(VariationalSGDOptimizer, self)._prepare() - - def _get_coordinatewise_learning_rate(self, grad, var): - # Compute the learning rate using a moving average for the diagonal of BB^T - avg_first = self.get_slot(var, 'first_moment') - avg_second = self.get_slot(var, 'second_moment') - decay_tensor = math_ops.cast(self._decay_tensor, var.dtype) - batch_size = math_ops.cast(self._batch_size_tensor, var.dtype) - - # Create an estimator for the moving average of gradient mean and variance - # via Welford's algorithm - if isinstance(grad, ops.Tensor): - delta = grad - avg_first - first_moment_update = avg_first.assign_add( - array_ops.where(self._counter < 1, math_ops.cast(1, var.dtype), - 1. - decay_tensor) * delta) - - with ops.control_dependencies([first_moment_update]): - second_moment_update = avg_second.assign_add( - math_ops.cast(self._counter < 1, var.dtype) * - -(1. - decay_tensor) * ( - avg_second - decay_tensor * math_ops.square(delta))) - diag_preconditioner = control_flow_ops.with_dependencies( - [second_moment_update], - clip_ops.clip_by_value(avg_second, 1e-12, 1e12)) - elif isinstance(grad, ops.IndexedSlices): - delta = grad.values - array_ops.gather_nd(avg_first, grad.indices) - first_moment_update = state_ops.scatter_add( - avg_first, - grad.indices, - array_ops.where(self._counter < 1, - math_ops.cast(1., var.dtype), - 1. - decay_tensor) * delta) - - with ops.control_dependencies([first_moment_update]): - avg_second = state_ops.scatter_add( - avg_second, - grad.indices, - math_ops.cast(self._counter < 1, var.dtype) * - -(1. - decay_tensor) * ( - array_ops.gather_nd(avg_second, grad.indices) - decay_tensor * - math_ops.square(delta))) - avg_second = array_ops.gather_nd(avg_second, grad.indices) - # TODO(b/70783772) - diag_preconditioner = clip_ops.clip_by_value(avg_second, 1e-12, 1e12) - else: - raise errors.InvalidArgumentError( - None, None, 'grad must of type Tensor or IndexedSlice') - - diag_preconditioner *= batch_size - - if self._use_single_learning_rate: - diag_preconditioner = math_ops.reduce_mean(diag_preconditioner) - - # From Theorem 2 Corollary 1 of Mandt et al. 2017 - return 2. * batch_size / ( - math_ops.cast(self._total_num_examples, var.dtype.base_dtype) * - diag_preconditioner) - - def _apply_dense(self, grad, var): - - max_learning_rate = array_ops.where(self._counter < self._burnin, - self._burnin_max_learning_rate, - self._max_learning_rate) - - learn_rates = clip_ops.clip_by_value( - self._get_coordinatewise_learning_rate(grad, var), 0.0, - math_ops.cast(max_learning_rate, var.dtype.base_dtype)) - - newgrad = grad * learn_rates - return training_ops.apply_gradient_descent( - var, - math_ops.cast(1.0, var.dtype), - newgrad, - use_locking=self._use_locking).op - - def _apply_sparse(self, grad, var): - - max_learning_rate = array_ops.where(self._counter < self._burnin, - self._burnin_max_learning_rate, - self._max_learning_rate) - - learn_rate = clip_ops.clip_by_value( - self._get_coordinatewise_learning_rate(grad, var), 0.0, - math_ops.cast(max_learning_rate, var.dtype)) - delta = grad.values * learn_rate - - return state_ops.scatter_sub(var, grad.indices, delta, - use_locking=self._use_locking) - - def _finish(self, update_ops, name_scope): - update_ops.append([self._counter.assign_add(1)]) - return control_flow_ops.group(*update_ops, name=name_scope) - - @property - def variable_scope(self): - """Variable scope of all calls to `tf.get_variable`.""" - return self._variable_scope 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 31f5c444817b9b82723c86bea3504d4934e57eb8..d9b0d89a03dce40d34f76bb1262d26bb587a2dc7 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py @@ -54,7 +54,7 @@ def make_custom_export_strategy(name, An `ExportStrategy`. """ base_strategy = saved_model_export_utils.make_export_strategy( - serving_input_fn=export_input_fn) + serving_input_fn=export_input_fn, strip_default_attrs=True) input_fn = export_input_fn() (sorted_feature_names, dense_floats, sparse_float_indices, _, _, sparse_int_indices, _, _) = gbdt_batch.extract_features( @@ -93,7 +93,9 @@ def make_custom_export_strategy(name, "w") as f: f.write("\n".join("%s, %f" % (k, v) for k, v in sorted_by_importance)) return result_dir - return export_strategy.ExportStrategy(name, export_fn) + + return export_strategy.ExportStrategy( + name, export_fn, strip_default_attrs=True) def convert_to_universal_format(dtec, sorted_feature_names, diff --git a/tensorflow/contrib/boosted_trees/kernels/model_ops.cc b/tensorflow/contrib/boosted_trees/kernels/model_ops.cc index 754b7bc3270d647fc381033b769eadd7b791771e..3bf33186ec13f5ff991db938d59849c0124a30a0 100644 --- a/tensorflow/contrib/boosted_trees/kernels/model_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/model_ops.cc @@ -137,6 +137,61 @@ class TreeEnsembleDeserializeOp : public OpKernel { } }; +class TreeEnsembleUsedHandlersOp : public OpKernel { + public: + explicit TreeEnsembleUsedHandlersOp(OpKernelConstruction* context) + : OpKernel(context) { + OP_REQUIRES_OK(context, + context->GetAttr("num_all_handlers", &num_handlers_)); + } + + void Compute(OpKernelContext* context) override { + boosted_trees::models::DecisionTreeEnsembleResource* ensemble_resource; + + OP_REQUIRES_OK(context, LookupResource(context, HandleFromInput(context, 0), + &ensemble_resource)); + tf_shared_lock l(*ensemble_resource->get_mutex()); + core::ScopedUnref unref_me(ensemble_resource); + + // Get the stamp token. + const Tensor* stamp_token_t; + OP_REQUIRES_OK(context, context->input("stamp_token", &stamp_token_t)); + int64 stamp_token = stamp_token_t->scalar()(); + + // Only the Chief should run this Op and it is guaranteed to be in + // a consistent state so the stamps must always match. + CHECK(ensemble_resource->is_stamp_valid(stamp_token)); + + Tensor* output_used_handlers_t = nullptr; + OP_REQUIRES_OK( + context, context->allocate_output("used_handlers_mask", {num_handlers_}, + &output_used_handlers_t)); + auto output_used_handlers = output_used_handlers_t->vec(); + + Tensor* output_num_used_handlers_t = nullptr; + OP_REQUIRES_OK(context, + context->allocate_output("num_used_handlers", {}, + &output_num_used_handlers_t)); + int handler_idx = 0; + std::vector used_handlers = ensemble_resource->GetUsedHandlers(); + output_num_used_handlers_t->scalar()() = used_handlers.size(); + for (int64 i = 0; i < num_handlers_; ++i) { + if (handler_idx >= used_handlers.size() || + used_handlers[handler_idx] > i) { + output_used_handlers(i) = false; + } else { + OP_REQUIRES(context, used_handlers[handler_idx] == i, + errors::InvalidArgument("Handler IDs should be sorted.")); + ++handler_idx; + output_used_handlers(i) = true; + } + } + } + + private: + int64 num_handlers_; +}; + REGISTER_RESOURCE_HANDLE_KERNEL(DecisionTreeEnsembleResource); REGISTER_KERNEL_BUILDER( @@ -155,5 +210,7 @@ REGISTER_KERNEL_BUILDER(Name("TreeEnsembleSerialize").Device(DEVICE_CPU), REGISTER_KERNEL_BUILDER(Name("TreeEnsembleDeserialize").Device(DEVICE_CPU), TreeEnsembleDeserializeOp); +REGISTER_KERNEL_BUILDER(Name("TreeEnsembleUsedHandlers").Device(DEVICE_CPU), + TreeEnsembleUsedHandlersOp); } // namespace boosted_trees } // namespace tensorflow diff --git a/tensorflow/contrib/boosted_trees/kernels/training_ops.cc b/tensorflow/contrib/boosted_trees/kernels/training_ops.cc index 7f8dea1d3c2a04b725843f6e2932a0cdfbc7733c..1bfeed306641111718984b2097512e5ec3fa8630 100644 --- a/tensorflow/contrib/boosted_trees/kernels/training_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/training_ops.cc @@ -361,27 +361,10 @@ 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, allowed_handlers, partition_ids_list, - gains_list, splits_list, &best_splits); + FindBestSplitsPerPartition(context, partition_ids_list, gains_list, + splits_list, &best_splits); // No-op if no new splits can be considered. if (best_splits.empty()) { @@ -422,19 +405,12 @@ class GrowTreeEnsembleOp : public OpKernel { // and finds the best split for each partition. void FindBestSplitsPerPartition( 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(); 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 cd925f6b65e569538212e9c26aef0abc8482960b..794ba2bcb0aafa26c5e1c90fcd66caf9dd5bf7d5 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 @@ -137,7 +137,7 @@ struct NodeStats { Eigen::MatrixXf hessian = TensorToEigenMatrix(grad_stats.second.t, grad_dim, grad_dim); // I is an identity matrix. - // The gain in general form is -g^T (H+l2 I)^-1 g. + // The gain in general form is g^T (H+l2 I)^-1 g. // The node weights are -(H+l2 I)^-1 g. Eigen::MatrixXf identity; identity.setIdentity(grad_dim, grad_dim); @@ -240,7 +240,7 @@ struct NodeStats { // given regularized Hessian and gradient vector g. void CalculateWeightAndGain(const Eigen::MatrixXf& hessian_and_reg, const Eigen::VectorXf& g) { - // The gain in general form is -g^T (Hessian_and_regularization)^-1 g. + // The gain in general form is g^T (Hessian_and_regularization)^-1 g. // The node weights are -(Hessian_and_regularization)^-1 g. Eigen::VectorXf weight; // If we want to calculate x = K^-1 v, instead of explicitly calculating diff --git a/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h b/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h index da5e7448519cb7f4092f7bbbe1b526271008ec22..7815fa049aa165a944c45872c762b7a5bf91b316 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h @@ -48,9 +48,9 @@ class BatchFeatures { Status GetFeatureColumnSizes(int64* const num_dense_float_features, int64* const num_sparse_float_features, int64* const num_sparse_int_features) const { - QCHECK_NE(num_dense_float_features, nullptr); - QCHECK_NE(num_sparse_float_features, nullptr); - QCHECK_NE(num_sparse_int_features, nullptr); + QCHECK_NE(num_dense_float_features, (int64*) nullptr); + QCHECK_NE(num_sparse_float_features, (int64*) nullptr); + QCHECK_NE(num_sparse_int_features, (int64*) nullptr); *num_dense_float_features = dense_float_feature_columns_.size(); *num_sparse_float_features = sparse_float_feature_columns_.size(); *num_sparse_int_features = sparse_int_feature_columns_.size(); diff --git a/tensorflow/contrib/boosted_trees/ops/model_ops.cc b/tensorflow/contrib/boosted_trees/ops/model_ops.cc index 0786c4166410720e8d4d70960e5747ff111076d8..9d6343c7e80f369bf6a5465821c5f4bacb984cd0 100644 --- a/tensorflow/contrib/boosted_trees/ops/model_ops.cc +++ b/tensorflow/contrib/boosted_trees/ops/model_ops.cc @@ -110,5 +110,32 @@ stamp_token: Token to use as the new value of the resource stamp. tree_ensemble_config: Serialized proto of the ensemble. )doc"); +REGISTER_OP("TreeEnsembleUsedHandlers") + .Attr("num_all_handlers: int >= 0") + .Input("tree_ensemble_handle: resource") + .Input("stamp_token: int64") + .Output("num_used_handlers: int64") + .Output("used_handlers_mask: bool") + .SetShapeFn([](shape_inference::InferenceContext* c) { + shape_inference::ShapeHandle unused_input; + TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused_input)); + c->set_output(0, c->Scalar()); + int num_all_handlers; + c->GetAttr("num_all_handlers", &num_all_handlers).IgnoreError(); + c->set_output(1, {c->Vector(num_all_handlers)}); + + return Status::OK(); + }) + .Doc(R"doc( +Returns the mask of used handlers along with the number of non-zero elements in +this mask. Used in feature selection. + +tree_ensemble_handle: Handle to the tree ensemble. +stamp_token: Token to use as the new value of the resource stamp. +num_used_handlers: number of feature column handlers used in the model. +used_handlers_mask: A boolean vector of showing which handlers are used in the + model. +)doc"); + } // namespace boosted_trees } // namespace tensorflow diff --git a/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc b/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc index ae99d53a2cf805d70d60746cd44f73f7fd9dc6e2..6aa52463987b55a54b7308765920cbe94c15b8d1 100644 --- a/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc +++ b/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc @@ -272,6 +272,20 @@ REGISTER_OP("Quantiles") .Input("sparse_indices: num_sparse_features * int64") .Output("dense_quantiles: num_dense_features * int32") .Output("sparse_quantiles: num_sparse_features * int32") + .SetShapeFn([](InferenceContext* c) { + int num_dense_features; + TF_RETURN_IF_ERROR(c->GetAttr("num_dense_features", &num_dense_features)); + int num_sparse_features; + TF_RETURN_IF_ERROR( + c->GetAttr("num_sparse_features", &num_sparse_features)); + // Set output shapes (dense_quantiles and sparse_quantiles) by the + // relevant inputs (dense_values and sparse_values). Note that the output + // has an additional dimension for dimension_ids. + for (int i = 0; i < num_dense_features + num_sparse_features; ++i) { + c->set_output(i, c->MakeShape({c->Dim(c->input(i), 0), 2})); + } + return Status::OK(); + }) .Doc(R"doc( Computes quantile for each a given list of dense and sparse feature values using the given buckets. diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py index 27c288bbf78b3b593d0807e92ac7fd9afc4d2725..63b9c5fddf0d9967d53077608664b59d9ae00481 100644 --- a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py +++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py @@ -310,6 +310,22 @@ class ModelOpsTest(test_util.TensorFlowTestCase): # The third tree was added after the save. self.assertAllClose(result.eval(), [[-1.1], [-1.1]]) + def testUsedHandlers(self): + with self.test_session(): + tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig() + tree_ensemble_config.growing_metadata.used_handler_ids.append(1) + tree_ensemble_config.growing_metadata.used_handler_ids.append(5) + stamp_token = 3 + tree_ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=stamp_token, + tree_ensemble_config=tree_ensemble_config.SerializeToString(), + name="create_tree") + resources.initialize_resources(resources.shared_resources()).run() + result = model_ops.tree_ensemble_used_handlers( + tree_ensemble_handle, stamp_token, num_all_handlers=6) + self.assertAllEqual([0, 1, 0, 0, 0, 1], result.used_handlers_mask.eval()) + self.assertEqual(2, result.num_used_handlers.eval()) + if __name__ == "__main__": googletest.main() 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 8ca1aabacaf53b66aaba184962922294427d6803..3e524efbeac74ff754d63cae92b3e194411cb2de 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 @@ -1588,7 +1588,7 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): self.assertEqual( 2, tree_ensemble_config.tree_metadata[2].num_tree_weight_updates) - def testGrowExistingEnsembleTreeWithFeatureSelectionCanStillGrow(self): + def testGrowExistingEnsembleTreeWithFeatureSelectionUsedHandlers(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. @@ -1649,7 +1649,6 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): 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( @@ -1668,183 +1667,8 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): 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.constraints.max_number_of_unique_feature_columns = 3 learner_config = learner_config.SerializeToString() # Prepare handler inputs. handler1_partitions = np.array([0], dtype=np.int32) @@ -1876,12 +1700,10 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): _, 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)) + self.assertEqual(3, len(tree_ensemble_config.trees)) + # 2 was already used. handler 0 is being added in this tree. + self.assertAllEqual( + [0, 2], tree_ensemble_config.growing_metadata.used_handler_ids) if __name__ == "__main__": diff --git a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py index 7a5f509047d46549ba81039a23d29ec987ca7920..25b2c9e2fd72bd018717e8a87fce726f26bad968 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py @@ -25,6 +25,7 @@ from tensorflow.contrib.boosted_trees.python.ops.gen_model_ops import tree_ensem from tensorflow.contrib.boosted_trees.python.ops.gen_model_ops import tree_ensemble_serialize # pylint: disable=unused-import from tensorflow.contrib.boosted_trees.python.ops.gen_model_ops import tree_ensemble_stamp_token +from tensorflow.contrib.boosted_trees.python.ops.gen_model_ops import tree_ensemble_used_handlers # pylint: enable=unused-import from tensorflow.python.framework import ops 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 f0b66dcbbe1c5167b9993e66b30b1dc8a839c380..85b909e4f2556c520a5bffe46d5954683d9dda5a 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py @@ -57,6 +57,8 @@ PREDICTIONS = "predictions" PARTITION_IDS = "partition_ids" NUM_LAYERS_ATTEMPTED = "num_layers" NUM_TREES_ATTEMPTED = "num_trees" +NUM_USED_HANDLERS = "num_used_handlers" +USED_HANDLERS_MASK = "used_handlers_mask" _FEATURE_NAME_TEMPLATE = "%s_%d" @@ -70,7 +72,8 @@ def _get_column_by_index(tensor, indices): return array_ops.reshape(array_ops.gather(p_flat, i_flat), [shape[0], -1]) -def _make_predictions_dict(stamp, logits, partition_ids, ensemble_stats): +def _make_predictions_dict(stamp, logits, partition_ids, ensemble_stats, + used_handlers): """Returns predictions for the given logits and n_classes. Args: @@ -79,6 +82,8 @@ def _make_predictions_dict(stamp, logits, partition_ids, ensemble_stats): that contains predictions when no dropout was applied. partition_ids: A rank 1 `Tensor` with shape [batch_size]. ensemble_stats: A TreeEnsembleStatsOp result tuple. + used_handlers: A TreeEnsembleUsedHandlerOp result tuple of an int and a + boolean mask.. Returns: A dict of predictions. @@ -89,6 +94,8 @@ def _make_predictions_dict(stamp, logits, partition_ids, ensemble_stats): result[PARTITION_IDS] = partition_ids result[NUM_LAYERS_ATTEMPTED] = ensemble_stats.attempted_layers result[NUM_TREES_ATTEMPTED] = ensemble_stats.attempted_trees + result[NUM_USED_HANDLERS] = used_handlers.num_used_handlers + result[USED_HANDLERS_MASK] = used_handlers.used_handlers_mask return result @@ -361,6 +368,13 @@ class GradientBoostedDecisionTreeModel(object): """ ensemble_stats = training_ops.tree_ensemble_stats(ensemble_handle, ensemble_stamp) + num_handlers = ( + len(self._dense_floats) + len(self._sparse_float_shapes) + + len(self._sparse_int_shapes)) + # Used during feature selection. + used_handlers = model_ops.tree_ensemble_used_handlers( + ensemble_handle, ensemble_stamp, num_all_handlers=num_handlers) + # We don't need dropout info - we can always restore it based on the # seed. apply_dropout, seed = _dropout_params(mode, ensemble_stats) @@ -395,7 +409,7 @@ class GradientBoostedDecisionTreeModel(object): use_locking=True) return _make_predictions_dict(ensemble_stamp, predictions, partition_ids, - ensemble_stats) + ensemble_stats, used_handlers) def predict(self, mode): """Returns predictions given the features and mode. @@ -710,12 +724,28 @@ class GradientBoostedDecisionTreeModel(object): active_handlers_current_layer = ( active_handlers_current_layer < self._learner_config.feature_fraction_per_tree) - active_handlers = array_ops.stack(active_handlers_current_layer, - array_ops.ones( - [len(handlers)], dtype=dtypes.bool)) + active_handlers = array_ops.stack([ + active_handlers_current_layer, + array_ops.ones([len(handlers)], dtype=dtypes.bool)], axis=1) else: active_handlers = array_ops.ones([len(handlers), 2], dtype=dtypes.bool) + if self._learner_config.constraints.max_number_of_unique_feature_columns: + target = ( + self._learner_config.constraints.max_number_of_unique_feature_columns) + + def _feature_selection_active_handlers(): + # The active list for current and the next iteration. + used_handlers = array_ops.reshape(predictions_dict[USED_HANDLERS_MASK], + [-1, 1]) + used_handlers = array_ops.concat([used_handlers, used_handlers], axis=1) + return math_ops.logical_and(used_handlers, active_handlers) + + active_handlers = ( + control_flow_ops.cond(predictions_dict[NUM_USED_HANDLERS] >= target, + _feature_selection_active_handlers, + lambda: active_handlers)) + # Prepare empty gradients and hessians when handlers are not ready. empty_hess_shape = [1] + hessian_shape.as_list() empty_grad_shape = [1] + gradient_shape.as_list() diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py index dba51d4f527792d2a8dedc693f74c07119fd231d..6411f57a5419123e799af9231a04fce8ae7724d4 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py @@ -47,6 +47,38 @@ def _squared_loss(label, unused_weights, predictions): return loss +def _append_to_leaf(leaf, c_id, w): + """Helper method for building tree leaves. + + Appends weight contributions for the given class index to a leaf node. + + Args: + leaf: leaf node to append to. + c_id: class Id for the weight update. + w: weight contribution value. + """ + leaf.sparse_vector.index.append(c_id) + leaf.sparse_vector.value.append(w) + + +def _set_float_split(split, feat_col, thresh, l_id, r_id): + """Helper method for building tree float splits. + + Sets split feature column, threshold and children. + + Args: + split: split node to update. + feat_col: feature column for the split. + thresh: threshold to split on forming rule x <= thresh. + l_id: left child Id. + r_id: right child Id. + """ + split.feature_column = feat_col + split.threshold = thresh + split.left_id = l_id + split.right_id = r_id + + class GbdtTest(test_util.TensorFlowTestCase): def setUp(self): @@ -917,6 +949,350 @@ class GbdtTest(test_util.TensorFlowTestCase): output.trees[0].nodes[2].leaf.sparse_vector.value[0], atol=1e-4, rtol=1e-4) + def testTrainFnChiefFeatureSelectionReachedLimitNoGoodSplit(self): + """Tests the train function running on chief with feature selection.""" + with self.test_session() as sess: + ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, tree_ensemble_config="", name="tree_ensemble") + learner_config = learner_pb2.LearnerConfig() + learner_config.learning_rate_tuner.fixed.learning_rate = 0.1 + learner_config.num_classes = 2 + learner_config.regularization.l1 = 0 + learner_config.regularization.l2 = 0 + learner_config.constraints.max_tree_depth = 1 + learner_config.constraints.max_number_of_unique_feature_columns = 1 + learner_config.constraints.min_node_weight = 0 + features = {} + features["dense_float_0"] = array_ops.ones([4, 1], dtypes.float32) + # Feature 1 is predictive but it won't be used because we have reached the + # limit of num_used_handlers >= max_number_of_unique_feature_columns + features["dense_float_1"] = array_ops.constant([0, 0, 1, 1], + dtypes.float32) + + gbdt_model = gbdt_batch.GradientBoostedDecisionTreeModel( + is_chief=True, + num_ps_replicas=0, + center_bias=False, + ensemble_handle=ensemble_handle, + examples_per_layer=1, + learner_config=learner_config, + logits_dimension=1, + features=features) + + predictions = array_ops.constant( + [[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32) + partition_ids = array_ops.zeros([4], dtypes.int32) + ensemble_stamp = variables.Variable( + initial_value=0, + name="ensemble_stamp", + trainable=False, + dtype=dtypes.int64) + + predictions_dict = { + "predictions": + predictions, + "predictions_no_dropout": + predictions, + "partition_ids": + partition_ids, + "ensemble_stamp": + ensemble_stamp, + "num_trees": + 12, + "num_used_handlers": + array_ops.constant(1, dtype=dtypes.int64), + "used_handlers_mask": + array_ops.constant([True, False], dtype=dtypes.bool), + } + + labels = array_ops.constant([0, 0, 1, 1], dtypes.float32) + weights = array_ops.ones([4, 1], dtypes.float32) + # Create train op. + train_op = gbdt_model.train( + loss=math_ops.reduce_mean( + _squared_loss(labels, weights, predictions)), + predictions_dict=predictions_dict, + labels=labels) + variables.global_variables_initializer().run() + resources.initialize_resources(resources.shared_resources()).run() + + # On first run, expect no splits to be chosen because the quantile + # buckets will not be ready. + train_op.run() + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + self.assertEquals(len(output.trees), 0) + self.assertEquals(len(output.tree_weights), 0) + self.assertEquals(stamp_token.eval(), 1) + + # Update the stamp to be able to run a second time. + sess.run([ensemble_stamp.assign_add(1)]) + + # On second run, expect a trivial split to be chosen to basically + # predict the average. + train_op.run() + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + self.assertEquals(len(output.trees), 1) + self.assertAllClose(output.tree_weights, [0.1]) + self.assertEquals(stamp_token.eval(), 2) + expected_tree = """ + nodes { + dense_float_binary_split { + feature_column: 0 + threshold: 1.0 + left_id: 1 + right_id: 2 + } + node_metadata { + gain: 0 + } + } + nodes { + leaf { + vector { + value: -0.25 + } + } + } + nodes { + leaf { + vector { + value: 0.0 + } + } + }""" + self.assertProtoEquals(expected_tree, output.trees[0]) + + def testTrainFnChiefFeatureSelectionWithGoodSplits(self): + """Tests the train function running on chief with feature selection.""" + with self.test_session() as sess: + ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, tree_ensemble_config="", name="tree_ensemble") + learner_config = learner_pb2.LearnerConfig() + learner_config.learning_rate_tuner.fixed.learning_rate = 0.1 + learner_config.num_classes = 2 + learner_config.regularization.l1 = 0 + learner_config.regularization.l2 = 0 + learner_config.constraints.max_tree_depth = 1 + learner_config.constraints.max_number_of_unique_feature_columns = 1 + learner_config.constraints.min_node_weight = 0 + features = {} + features["dense_float_0"] = array_ops.ones([4, 1], dtypes.float32) + # Feature 1 is predictive and is in our selected features so it will be + # used even when we're at the limit. + features["dense_float_1"] = array_ops.constant([0, 0, 1, 1], + dtypes.float32) + + gbdt_model = gbdt_batch.GradientBoostedDecisionTreeModel( + is_chief=True, + num_ps_replicas=0, + center_bias=False, + ensemble_handle=ensemble_handle, + examples_per_layer=1, + learner_config=learner_config, + logits_dimension=1, + features=features) + + predictions = array_ops.constant( + [[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32) + partition_ids = array_ops.zeros([4], dtypes.int32) + ensemble_stamp = variables.Variable( + initial_value=0, + name="ensemble_stamp", + trainable=False, + dtype=dtypes.int64) + + predictions_dict = { + "predictions": + predictions, + "predictions_no_dropout": + predictions, + "partition_ids": + partition_ids, + "ensemble_stamp": + ensemble_stamp, + "num_trees": + 12, + "num_used_handlers": + array_ops.constant(1, dtype=dtypes.int64), + "used_handlers_mask": + array_ops.constant([False, True], dtype=dtypes.bool), + } + + labels = array_ops.constant([0, 0, 1, 1], dtypes.float32) + weights = array_ops.ones([4, 1], dtypes.float32) + # Create train op. + train_op = gbdt_model.train( + loss=math_ops.reduce_mean( + _squared_loss(labels, weights, predictions)), + predictions_dict=predictions_dict, + labels=labels) + variables.global_variables_initializer().run() + resources.initialize_resources(resources.shared_resources()).run() + + # On first run, expect no splits to be chosen because the quantile + # buckets will not be ready. + train_op.run() + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + self.assertEquals(len(output.trees), 0) + self.assertEquals(len(output.tree_weights), 0) + self.assertEquals(stamp_token.eval(), 1) + + # Update the stamp to be able to run a second time. + sess.run([ensemble_stamp.assign_add(1)]) + + train_op.run() + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + + self.assertEquals(len(output.trees), 1) + self.assertAllClose(output.tree_weights, [0.1]) + self.assertEquals(stamp_token.eval(), 2) + expected_tree = """ + nodes { + dense_float_binary_split { + feature_column: 1 + left_id: 1 + right_id: 2 + } + node_metadata { + gain: 0.5 + } + } + nodes { + leaf { + vector { + value: 0.0 + } + } + } + nodes { + leaf { + vector { + value: -0.5 + } + } + }""" + self.assertProtoEquals(expected_tree, output.trees[0]) + + def testTrainFnChiefFeatureSelectionReachedLimitIncrementAttemptedLayer(self): + """Tests the train function running on chief with feature selection.""" + with self.test_session() as sess: + tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig() + tree = tree_ensemble_config.trees.add() + + _set_float_split(tree.nodes.add() + .sparse_float_binary_split_default_right.split, 2, 4.0, + 1, 2) + _append_to_leaf(tree.nodes.add().leaf, 0, 0.5) + _append_to_leaf(tree.nodes.add().leaf, 1, 1.2) + tree_ensemble_config.tree_weights.append(1.0) + metadata = tree_ensemble_config.tree_metadata.add() + metadata.is_finalized = False + metadata.num_layers_grown = 1 + tree_ensemble_config = tree_ensemble_config.SerializeToString() + ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, tree_ensemble_config=tree_ensemble_config, + name="tree_ensemble") + learner_config = learner_pb2.LearnerConfig() + learner_config.learning_rate_tuner.fixed.learning_rate = 0.1 + learner_config.num_classes = 2 + learner_config.regularization.l1 = 0 + learner_config.regularization.l2 = 0 + learner_config.constraints.max_tree_depth = 1 + learner_config.constraints.max_number_of_unique_feature_columns = 1 + learner_config.constraints.min_node_weight = 0 + features = {} + # Both features will be disabled since the feature selection limit is + # already reached. + features["dense_float_0"] = array_ops.ones([4, 1], dtypes.float32) + features["dense_float_1"] = array_ops.constant([0, 0, 1, 1], + dtypes.float32) + + gbdt_model = gbdt_batch.GradientBoostedDecisionTreeModel( + is_chief=True, + num_ps_replicas=0, + center_bias=False, + ensemble_handle=ensemble_handle, + examples_per_layer=1, + learner_config=learner_config, + logits_dimension=1, + features=features) + + predictions = array_ops.constant( + [[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32) + partition_ids = array_ops.zeros([4], dtypes.int32) + ensemble_stamp = variables.Variable( + initial_value=0, + name="ensemble_stamp", + trainable=False, + dtype=dtypes.int64) + + predictions_dict = { + "predictions": + predictions, + "predictions_no_dropout": + predictions, + "partition_ids": + partition_ids, + "ensemble_stamp": + ensemble_stamp, + "num_trees": + 12, + # We have somehow reached our limit 1. Both of the handlers will be + # disabled. + "num_used_handlers": + array_ops.constant(1, dtype=dtypes.int64), + "used_handlers_mask": + array_ops.constant([False, False], dtype=dtypes.bool), + } + + labels = array_ops.constant([0, 0, 1, 1], dtypes.float32) + weights = array_ops.ones([4, 1], dtypes.float32) + # Create train op. + train_op = gbdt_model.train( + loss=math_ops.reduce_mean( + _squared_loss(labels, weights, predictions)), + predictions_dict=predictions_dict, + labels=labels) + variables.global_variables_initializer().run() + resources.initialize_resources(resources.shared_resources()).run() + + # On first run, expect no splits to be chosen because the quantile + # buckets will not be ready. + train_op.run() + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + self.assertEquals(len(output.trees), 1) + self.assertEquals(output.growing_metadata.num_layers_attempted, 1) + self.assertEquals(stamp_token.eval(), 1) + + # Update the stamp to be able to run a second time. + sess.run([ensemble_stamp.assign_add(1)]) + + train_op.run() + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + # Make sure the trees are not modified, but the num_layers_attempted is + # incremented so that eventually the training stops. + self.assertEquals(len(output.trees), 1) + self.assertEquals(len(output.trees[0].nodes), 3) + + self.assertEquals(output.growing_metadata.num_layers_attempted, 2) if __name__ == "__main__": googletest.main() 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 3ebf28ea442edf87815c39971ae9e01a2a8aae9a..94aeb2c7bb48c6eddb6c7894f8bf6f1567470113 100644 --- a/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h +++ b/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h @@ -126,7 +126,8 @@ class DecisionTreeEnsembleResource : public StampedResource { return; } used_ids->Add(handler_id); - std::rotate(first, used_ids->end() - 1, used_ids->end()); + // Keep the list of used handlers sorted. + std::sort(used_ids->begin(), used_ids->end()); } std::vector GetUsedHandlers() const { diff --git a/tensorflow/contrib/cluster_resolver/BUILD b/tensorflow/contrib/cluster_resolver/BUILD index 15abd2be0385eb776ff4f76484133efb6e34f076..1a124eca364424b651de86bfaac6f33ad131804b 100644 --- a/tensorflow/contrib/cluster_resolver/BUILD +++ b/tensorflow/contrib/cluster_resolver/BUILD @@ -30,10 +30,12 @@ py_library( "python/training/__init__.py", ], srcs_version = "PY2AND3", + visibility = ["//visibility:public"], deps = [ ":cluster_resolver_py", ":gce_cluster_resolver_py", ":tpu_cluster_resolver_py", + "//tensorflow/python:util", ], ) @@ -108,5 +110,6 @@ tf_py_test( "//tensorflow/python:platform_test", "//tensorflow/python:training", ], + grpc_enabled = True, main = "python/training/tpu_cluster_resolver_test.py", ) 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/cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py index b04822fa9d66465e34a545d3b00c399bbb196514..1c480b25134b1e54200e0ddb780bd7bb0f122341 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py @@ -53,11 +53,16 @@ class ClusterResolver(object): raise NotImplementedError( 'cluster_spec is not implemented for {}.'.format(self)) + @abc.abstractmethod + def master(self): + """...""" + raise NotImplementedError('master is not implemented for {}.'.format(self)) + class SimpleClusterResolver(ClusterResolver): """Simple implementation of ClusterResolver that accepts a ClusterSpec.""" - def __init__(self, cluster_spec): + def __init__(self, cluster_spec, master=''): """Creates a SimpleClusterResolver from a ClusterSpec.""" super(SimpleClusterResolver, self).__init__() @@ -65,10 +70,18 @@ class SimpleClusterResolver(ClusterResolver): raise TypeError('cluster_spec must be a ClusterSpec.') self._cluster_spec = cluster_spec + if not isinstance(master, str): + raise TypeError('master must be a string.') + self._master = master + def cluster_spec(self): """Returns the ClusterSpec passed into the constructor.""" return self._cluster_spec + def master(self): + """Returns the master address to use when creating a session.""" + return self._master + class UnionClusterResolver(ClusterResolver): """Performs a union on underlying ClusterResolvers. @@ -87,9 +100,13 @@ class UnionClusterResolver(ClusterResolver): Raises: TypeError: If any argument is not a subclass of `ClusterResolvers`. + ValueError: If there are no arguments passed. """ super(UnionClusterResolver, self).__init__() + if not args: + raise ValueError('At least one ClusterResolver is required.') + for cluster_resolver in args: if not isinstance(cluster_resolver, ClusterResolver): raise TypeError('All arguments must be a sub-class of ' @@ -169,3 +186,7 @@ class UnionClusterResolver(ClusterResolver): merged_cluster[job_name].update(task_dict) return ClusterSpec(merged_cluster) + + def master(self): + """master returns the master address from the first cluster resolver.""" + return self._cluster_resolvers[0].master() diff --git a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py index dbfb77723cdaab66e29bb41b764593bb5fd61b35..d9c97d53eb3663f6ab2f7b40395592dc7638b896 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py +++ b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py @@ -234,5 +234,7 @@ class UnionClusterResolverTest(test.TestCase): self._verifyClusterSpecEquality(cluster_spec, expected_proto) +# TODO(saeta): Include tests for master resolution + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py index d6f2eced93ba4fda5ac27f9412b6f729981f4f40..3f5824128948453634bc5e5a7d6fdeedae60f5bd 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py @@ -134,3 +134,6 @@ class GceClusterResolver(ClusterResolver): worker_list.sort() return ClusterSpec({self._job_name: worker_list}) + + def master(self): + return '' 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 2e75ac226ea74e879edda5e03dff3d53c8a76569..300b19733e2b4d1b912f966e94ae0286ed9c694d 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py @@ -18,12 +18,14 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os 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 +from tensorflow.python.training import server_lib +from tensorflow.python.util import compat _GOOGLE_API_CLIENT_INSTALLED = True try: @@ -33,6 +35,9 @@ except ImportError: _GOOGLE_API_CLIENT_INSTALLED = False +_GKE_ENV_VARIABLE = 'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS' + + class TPUClusterResolver(ClusterResolver): """Cluster Resolver for Google Cloud TPUs. @@ -46,13 +51,30 @@ class TPUClusterResolver(ClusterResolver): req = Request('http://metadata/computeMetadata/v1/%s' % path, headers={'Metadata-Flavor': 'Google'}) resp = urlopen(req) - return resp.read() + return compat.as_bytes(resp.read()) + + def _shouldResolve(self): + if (self._tpu == compat.as_bytes('') or + self._tpu == compat.as_bytes('local') or + self._tpu.startswith(compat.as_bytes('/bns')) or + self._tpu.startswith(compat.as_bytes('grpc://'))): + return False + return True + + def _inGke(self): + """When running in GKE, the environment variable will be set.""" + return _GKE_ENV_VARIABLE in os.environ + + def _gkeMaster(self): + return os.environ[_GKE_ENV_VARIABLE].split(',')[0] def __init__(self, - tpu_names, + tpu=None, zone=None, project=None, - job_name='tpu_worker', + job_name='worker', + coordinator_name='coordinator', + coordinator_address=None, credentials='default', service=None): """Creates a new TPUClusterResolver object. @@ -61,7 +83,11 @@ class TPUClusterResolver(ClusterResolver): for the IP addresses and ports of each Cloud TPU listed. Args: - tpu_names: A list of names of the target Cloud TPUs. + tpu: Either a string, or a list of strings corresponding to the TPUs to + use. If the single string is the empty string, the string 'local', or a + string that begins with 'grpc://' or '/bns', then it is assumed to not + correspond with a Cloud TPU and will instead be passed as the session + master and no ClusterSpec propagation will be done. 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. @@ -69,6 +95,12 @@ class TPUClusterResolver(ClusterResolver): 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. + coordinator_name: The name to use for the coordinator. Set to None if the + coordinator should not be included in the computed ClusterSpec. + coordinator_address: The address of the coordinator (typically an ip:port + pair). If set to None, a TF server will be started. If coordinator_name + is None, a TF server will not be started even if coordinator_address is + None. credentials: GCE Credentials. If None, then we use default credentials from the oauth2client service: The GCE API object returned by the googleapiclient.discovery @@ -77,29 +109,47 @@ class TPUClusterResolver(ClusterResolver): Raises: ImportError: If the googleapiclient is not installed. + ValueError: If no TPUs are specified. """ + if isinstance(tpu, list): + if not tpu: + raise ValueError('At least one TPU must be specified.') + if len(tpu) != 1: + raise NotImplementedError( + 'Using multiple TPUs in a single session is not yet implemented') + tpu = tpu[0] + + # When using GKE with Cloud TPUs, the env variable will be set. + if tpu is None and self._inGke(): + tpu = self._gkeMaster() + + self._tpu = compat.as_bytes(tpu) # self._tpu is always bytes + self._job_name = job_name + self._credentials = credentials + + should_resolve = self._shouldResolve() - if not project: - project = self._requestComputeMetadata('/project/project-id') + if not project and should_resolve: + project = compat.as_str( + self._requestComputeMetadata('project/project-id')) - if not zone: - zone_path = self._requestComputeMetadata('/instance/zone') + if not zone and should_resolve: + zone_path = compat.as_str(self._requestComputeMetadata('instance/zone')) zone = zone_path.split('/')[-1] self._project = project self._zone = zone - self._tpu_names = tpu_names - self._job_name = job_name - self._credentials = credentials - if credentials == 'default': + if credentials == 'default' and should_resolve: if _GOOGLE_API_CLIENT_INSTALLED: self._credentials = GoogleCredentials.get_application_default() - if service is None: + if service is None and should_resolve: if not _GOOGLE_API_CLIENT_INSTALLED: raise ImportError('googleapiclient must be installed before using the ' - 'TPU cluster resolver') + 'TPU cluster resolver. Execute: `pip install ' + '--upgrade google-api-python-client` to install with ' + 'pip.') self._service = discovery.build( 'tpu', 'v1alpha1', @@ -107,25 +157,41 @@ class TPUClusterResolver(ClusterResolver): else: self._service = service - def get_master(self): - """Get the ClusterSpec grpc master path. + self._coordinator_name = coordinator_name + if coordinator_name and not coordinator_address and should_resolve: + self._start_local_server() + else: + self._coordinator_address = coordinator_address - This returns the grpc path (grpc://1.2.3.4:8470) of first instance in the - ClusterSpec returned by the cluster_spec function. This is suitable for use - for the `master` argument in tf.Session() when you are using one TPU. + def master(self): + """Get the Master string to be used for the session. + + In the normal case, this returns the grpc path (grpc://1.2.3.4:8470) of + first instance in the ClusterSpec returned by the cluster_spec function. + + If a non-TPU name is used when constructing a TPUClusterResolver, that will + be returned instead (e.g. If the tpus argument's value when constructing + this TPUClusterResolver was 'grpc://10.240.1.2:8470', + 'grpc://10.240.1.2:8470' will be returned). Returns: - string, the grpc path of the first instance in the ClusterSpec. + string, the connection string to use when creating a session. Raises: ValueError: If none of the TPUs specified exists. """ + if not self._shouldResolve(): + return self._tpu + job_tasks = self.cluster_spec().job_tasks(self._job_name) if not job_tasks: raise ValueError('No TPUs exists with the specified names exist.') return 'grpc://' + job_tasks[0] + def get_master(self): + return self.master() + def cluster_spec(self): """Returns a ClusterSpec object based on the latest TPU information. @@ -134,16 +200,54 @@ class TPUClusterResolver(ClusterResolver): Returns: A ClusterSpec containing host information returned from Cloud TPUs. - """ - worker_list = [] - - for tpu_name in self._tpu_names: - full_name = 'projects/%s/locations/%s/nodes/%s' % ( - self._project, self._zone, tpu_name) - request = self._service.projects().locations().nodes().get(name=full_name) - response = request.execute() + Raises: + RuntimeError: If the provided TPU is not healthy. + """ + if not self._shouldResolve(): + return server_lib.ClusterSpec({}) + + full_name = 'projects/%s/locations/%s/nodes/%s' % ( + self._project, self._zone, compat.as_text(self._tpu)) + request = self._service.projects().locations().nodes().get(name=full_name) + response = request.execute() + + if 'health' in response and response['health'] != 'HEALTHY': + raise RuntimeError('TPU "%s" is unhealthy: "%s"' % (self._tpu, + response['health'])) + + if 'networkEndpoints' in response: + worker_list = [ + '%s:%s' % (endpoint['ipAddress'], endpoint['port']) + for endpoint in response['networkEndpoints'] + ] + else: + # Fall back to the deprecated response format instance_url = '%s:%s' % (response['ipAddress'], response['port']) - worker_list.append(instance_url) - - return ClusterSpec({self._job_name: worker_list}) + worker_list = [instance_url] + + cluster_spec = {self._job_name: worker_list} + + if self._coordinator_address: + cluster_spec[self._coordinator_name] = [self._coordinator_address] + + return server_lib.ClusterSpec(cluster_spec) + + def _start_local_server(self): + address = self._requestComputeMetadata('instance/network-interfaces/0/ip') + self._server = server_lib.Server( + { + 'local': ['0.0.0.0:0'] + }, protocol='grpc', config=None, start=True) + # self._server.target is of the form: grpc://ipaddress:port + target = compat.as_bytes(self._server.target) + splits = target.split(compat.as_bytes(':')) + assert len(splits) == 3, self._server.target + assert splits[0] == compat.as_bytes('grpc'), self._server.target + self._coordinator_port = compat.as_text(splits[2]) + self._coordinator_address = '%s:%s' % ( + address, compat.as_text(self._coordinator_port)) + + def __deepcopy__(self, memo): + # TODO(b/73668574): Remove this once RunConfig avoids performing deepcopy. + return self 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 0c4730613af4ad9ca87deb6200ab4bb93d3f6a53..48c3f6bb4f2d1643982e03d9ed68db14c10c184a 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 @@ -18,10 +18,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os + from tensorflow.contrib.cluster_resolver.python.training.tpu_cluster_resolver import TPUClusterResolver from tensorflow.python.platform import test from tensorflow.python.training import server_lib - +from tensorflow.python.util import compat mock = test.mock @@ -50,10 +52,12 @@ class MockNodeClass(object): def mock_request_compute_metadata(cls, *args, **kwargs): del cls, kwargs # Unused. - if args[0] == '/project/project-id': + if args[0] == 'project/project-id': return 'test-project' - elif args[0] == '/instance/zone': + elif args[0] == 'instance/zone': return 'projects/test-project/locations/us-central1-c' + elif args[0] == 'instance/network-interfaces/0/ip': + return '10.128.1.2' return '' @@ -71,18 +75,17 @@ class TPUClusterResolverTest(test.TestCase): expected_proto: Expected protobuf """ self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def()) - self.assertProtoEquals( - expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()) self.assertProtoEquals( expected_proto, - server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def()) + server_lib.ClusterSpec(cluster_spec).as_cluster_def()) + self.assertProtoEquals(expected_proto, + server_lib.ClusterSpec( + cluster_spec.as_cluster_def()).as_cluster_def()) + self.assertProtoEquals(expected_proto, + server_lib.ClusterSpec( + cluster_spec.as_dict()).as_cluster_def()) - def mock_service_client( - self, - tpu_map=None): + def mock_service_client(self, tpu_map=None): if tpu_map is None: tpu_map = {} @@ -98,27 +101,59 @@ class TPUClusterResolverTest(test.TestCase): return mock_client - @mock.patch.object(TPUClusterResolver, - '_requestComputeMetadata', + @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' + 'port': '8470', + 'health': 'HEALTHY' } } tpu_cluster_resolver = TPUClusterResolver( project=None, zone=None, - tpu_names=['test-tpu-1'], + tpu=['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' } } + job { + name: 'coordinator' + tasks { key: 0 value: '10.128.1.2:%s' } + } + job { + name: 'worker' + tasks { key: 0 value: '10.1.2.3:8470' } + } + """ % tpu_cluster_resolver._coordinator_port + self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto)) + + @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', + mock_request_compute_metadata) + def testRetrieveProjectAndZoneFromMetadataNoCoordinator(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=['test-tpu-1'], + coordinator_name=None, + credentials=None, + service=self.mock_service_client(tpu_map=tpu_map)) + + actual_cluster_spec = tpu_cluster_resolver.cluster_spec() + expected_proto = """ + job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } } """ self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) @@ -126,80 +161,213 @@ 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' } } tpu_cluster_resolver = TPUClusterResolver( project='test-project', zone='us-central1-c', - tpu_names=['test-tpu-1'], + tpu=['test-tpu-1'], + coordinator_address='10.128.1.5:10203', 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' } } + job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } } + job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } } """ self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - def testMultipleSuccessfulRetrieval(self): + def testNewNetworkEndpointFormat(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470' - }, - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': { - 'ipAddress': '10.4.5.6', - 'port': '8470' + 'health': 'HEALTHY', + 'networkEndpoints': [{ + 'ipAddress': '10.2.3.4', + 'port': 8470, + }] } } tpu_cluster_resolver = TPUClusterResolver( project='test-project', zone='us-central1-c', - tpu_names=['test-tpu-2', 'test-tpu-1'], + tpu='test-tpu-1', + coordinator_address='10.128.1.5:10203', 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.4.5.6:8470' } - tasks { key: 1 value: '10.1.2.3:8470' } } + job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } } + job { name: 'worker' tasks { key: 0 value: '10.2.3.4:8470' } } """ self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) + self.assertEqual('grpc://10.2.3.4:8470', tpu_cluster_resolver.master()) + + @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', + mock_request_compute_metadata) + def testPodResolution(self): + tpu_map = { + 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { + 'health': + 'HEALTHY', + 'networkEndpoints': [ + { + 'ipAddress': '10.2.3.4', + 'port': 8470, + }, + { + 'ipAddress': '10.2.3.5', + 'port': 8470, + }, + { + 'ipAddress': '10.2.3.6', + 'port': 8470, + }, + { + 'ipAddress': '10.2.3.7', + 'port': 8470, + }, + ] + } + } - def testGetMasterMultipleEntries(self): + tpu_cluster_resolver = TPUClusterResolver( + tpu='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: 'coordinator', + tasks { key: 0 value: '10.128.1.2:%s'} + } + job { + name: 'worker' + tasks { key: 0 value: '10.2.3.4:8470' } + tasks { key: 1 value: '10.2.3.5:8470' } + tasks { key: 2 value: '10.2.3.6:8470' } + tasks { key: 3 value: '10.2.3.7:8470' } + } + """ % tpu_cluster_resolver._coordinator_port + self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto)) + + def testPodResolutionNoCoordinator(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470' - }, - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': { - 'ipAddress': '10.4.5.6', - 'port': '8470' + 'health': + 'HEALTHY', + 'networkEndpoints': [ + { + 'ipAddress': '10.2.3.4', + 'port': 8470, + }, + { + 'ipAddress': '10.2.3.5', + 'port': 8470, + }, + { + 'ipAddress': '10.2.3.6', + 'port': 8470, + }, + { + 'ipAddress': '10.2.3.7', + 'port': 8470, + }, + ] } } tpu_cluster_resolver = TPUClusterResolver( project='test-project', zone='us-central1-c', - tpu_names=['test-tpu-2', 'test-tpu-1'], + tpu='test-tpu-1', + coordinator_name=None, credentials=None, service=self.mock_service_client(tpu_map=tpu_map)) - self.assertEqual('grpc://10.4.5.6:8470', tpu_cluster_resolver.get_master()) + + actual_cluster_spec = tpu_cluster_resolver.cluster_spec() + expected_proto = """ + job { + name: 'worker' + tasks { key: 0 value: '10.2.3.4:8470' } + tasks { key: 1 value: '10.2.3.5:8470' } + tasks { key: 2 value: '10.2.3.6:8470' } + tasks { key: 3 value: '10.2.3.7:8470' } + } + """ + self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) def testGetMasterNoEntries(self): tpu_map = {} + with self.assertRaises(ValueError): + TPUClusterResolver( + project='test-project', + zone='us-central1-c', + tpu=[], + coordinator_name=None, + credentials=None, + service=self.mock_service_client(tpu_map=tpu_map)) + + # TODO(saeta): Convert to parameterized test when included in OSS TF. + def verifyShouldResolve(self, tpu, should_resolve): tpu_cluster_resolver = TPUClusterResolver( project='test-project', zone='us-central1-c', - tpu_names=[], + tpu=tpu, + coordinator_name=None, credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - with self.assertRaises(ValueError): - tpu_cluster_resolver.get_master() + service=self.mock_service_client(tpu_map={})) + self.assertEqual(should_resolve, tpu_cluster_resolver._shouldResolve(), + "TPU: '%s'" % tpu) + + def testShouldResolveNoName(self): + self.verifyShouldResolve('', False) + + def testShouldResolveLocal(self): + self.verifyShouldResolve('local', False) + + def testShouldResolveGrpc(self): + self.verifyShouldResolve('grpc://10.1.2.3:8470', False) + + def testShouldResolveBns(self): + self.verifyShouldResolve('/bns/foo/bar', False) + + def testShouldResolveName(self): + self.verifyShouldResolve('mytpu', True) + + def testShouldResolveList(self): + self.verifyShouldResolve(['myothertpu'], True) + + def testShouldResolveGrpcPrefix(self): + self.verifyShouldResolve('grpctpu', True) + + def testNoCallComputeMetadata(self): + tpu_cluster_resolver = TPUClusterResolver(tpu='/bns/foo/bar') + self.assertEqual( + compat.as_bytes('/bns/foo/bar'), tpu_cluster_resolver.master()) + self.assertEqual( + server_lib.ClusterSpec({}), tpu_cluster_resolver.cluster_spec()) + + def testGkeEnvironment(self): + os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470' + self.assertTrue('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS' in os.environ) + tpu_cluster_resolver = TPUClusterResolver() + self.assertTrue(tpu_cluster_resolver._inGke()) + self.assertEqual( + compat.as_bytes('grpc://10.120.27.5:8470'), + compat.as_bytes(tpu_cluster_resolver._gkeMaster())) + self.assertEqual( + compat.as_bytes('grpc://10.120.27.5:8470'), + compat.as_bytes(tpu_cluster_resolver.get_master())) + del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt index 12bfd3c62b49c7f9ef40c84abf25268ceafcdbdc..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) @@ -188,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} @@ -203,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 @@ -286,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 @@ -307,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) @@ -335,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}) @@ -354,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 diff --git a/tensorflow/contrib/cmake/README.md b/tensorflow/contrib/cmake/README.md index 8f85a75ee466dbac524a1266dc2522109ca77cd5..fe83bb32046cd75328c92a74cdb4fdb6ce44560e 100644 --- a/tensorflow/contrib/cmake/README.md +++ b/tensorflow/contrib/cmake/README.md @@ -26,7 +26,7 @@ The CMake files in this directory can build the core TensorFlow runtime, an example C++ binary, and a PIP package containing the runtime and Python bindings. -### Pre-requisites +### Prerequisites * CMake version 3.5 or later. @@ -34,14 +34,16 @@ bindings. * [SWIG](http://www.swig.org/download.html) -* Additional pre-requisites for Microsoft Windows: +* Additional prerequisites for Microsoft Windows: - Visual Studio 2015 - Python 3.5 - - NumPy 1.11.0 or later -* Additional pre-requisites for Linux: +* Additional prerequisites for Linux: - Python 2.7 or later - [Docker](https://www.docker.com/) (for automated testing) + +* Python dependencies: + - wheel - NumPy 1.11.0 or later ### Known-good configurations @@ -102,7 +104,7 @@ ops or APIs. Step-by-step Windows build ========================== -1. Install the pre-requisites detailed above, and set up your environment. +1. Install the prerequisites detailed above, and set up your environment. * The following commands assume that you are using the Windows Command Prompt (`cmd.exe`). You will need to set up your environment to use the diff --git a/tensorflow/contrib/cmake/external/boringssl.cmake b/tensorflow/contrib/cmake/external/boringssl.cmake index 5ad477fdff68feab4adf0c0072c68c8e55390ab8..3c4bb01e24fd121c9d0fc3594cc25de37af0e8a1 100644 --- a/tensorflow/contrib/cmake/external/boringssl.cmake +++ b/tensorflow/contrib/cmake/external/boringssl.cmake @@ -37,6 +37,7 @@ 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 -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} diff --git a/tensorflow/contrib/cmake/external/cub.cmake b/tensorflow/contrib/cmake/external/cub.cmake index 836889895567f679d9960e29ece1600d1a7a58eb..98a8c7e736e5c8c407b90e8eac440cdc7ab21579 100644 --- a/tensorflow/contrib/cmake/external/cub.cmake +++ b/tensorflow/contrib/cmake/external/cub.cmake @@ -14,8 +14,8 @@ # ============================================================================== include (ExternalProject) -set(cub_URL https://mirror.bazel.build/github.com/NVlabs/cub/archive/1.7.4.zip) -set(cub_HASH SHA256=20a1a39fd97e5da7f40f5f2e7fd73fd2ea59f9dc4bb8a6c5f228aa543e727e31) +set(cub_URL https://mirror.bazel.build/github.com/NVlabs/cub/archive/1.8.0.zip) +set(cub_HASH SHA256=6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3) set(cub_BUILD ${CMAKE_CURRENT_BINARY_DIR}/cub/src/cub) set(cub_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/cub/src/cub) set(cub_ARCHIVE_DIR ${CMAKE_CURRENT_BINARY_DIR}/external/cub_archive) 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 28adb4fe84423bb5a21c78dac4e757505ce87d1d..17f65999faaf5c0ca39bfbc968a9140dbff49c2e 100644 --- a/tensorflow/contrib/cmake/external/grpc.cmake +++ b/tensorflow/contrib/cmake/external/grpc.cmake @@ -20,14 +20,22 @@ set(GRPC_BUILD ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc) 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 ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgrpc_unsecure.a + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/third_party/cares/cares/lib/libcares.a ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgpr.a) endif() @@ -40,6 +48,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 861201f97edbce2d9d70a833ce5a8cad46f2470a..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,6 +44,7 @@ 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 -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} diff --git a/tensorflow/contrib/cmake/external/lmdb.cmake b/tensorflow/contrib/cmake/external/lmdb.cmake index 41b314e2857577581eb27eb6c6480b757d0b436c..ed5ab788acc5625b9c8020fce15f027d98433096 100644 --- a/tensorflow/contrib/cmake/external/lmdb.cmake +++ b/tensorflow/contrib/cmake/external/lmdb.cmake @@ -20,10 +20,17 @@ 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} @@ -35,12 +42,6 @@ ExternalProject_Add(lmdb -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 b277be5690387b06876ca89eb88becbf885486a4..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,6 +48,7 @@ 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 diff --git a/tensorflow/contrib/cmake/external/protobuf.cmake b/tensorflow/contrib/cmake/external/protobuf.cmake index aedb793d2aef4bf6950cd074cd065909667eaf75..ab464bc99a43138130bb2758ae28ecef29805c31 100644 --- a/tensorflow/contrib/cmake/external/protobuf.cmake +++ b/tensorflow/contrib/cmake/external/protobuf.cmake @@ -19,11 +19,34 @@ set(PROTOBUF_URL https://github.com/google/protobuf.git) set(PROTOBUF_TAG b04e5cba356212e4e8c66c61bbe0c3a20537c5b9) 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,10 +59,15 @@ 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 "" @@ -47,5 +75,7 @@ ExternalProject_Add(protobuf -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 371d8447f93735e7af2a5a2b16f128a47b5a082a..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,6 +40,7 @@ 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 -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} diff --git a/tensorflow/contrib/cmake/external/snappy.cmake b/tensorflow/contrib/cmake/external/snappy.cmake index fd57734298affda13fa90f4cff560eeeb08e59ab..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,6 +39,7 @@ 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 diff --git a/tensorflow/contrib/cmake/external/sqlite.cmake b/tensorflow/contrib/cmake/external/sqlite.cmake index 8297c60712c49ed6f47a9750691eee1325a5b55e..57c4ae76517e4d7247093edd5e5bd95a83258d87 100644 --- a/tensorflow/contrib/cmake/external/sqlite.cmake +++ b/tensorflow/contrib/cmake/external/sqlite.cmake @@ -36,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}" diff --git a/tensorflow/contrib/cmake/external/zlib.cmake b/tensorflow/contrib/cmake/external/zlib.cmake index 5bec14fb00a50f6e6e8c7d8b703bde681e9d02ae..116d42309394b92407cef79c9d3a975f494bc3ff 100644 --- a/tensorflow/contrib/cmake/external/zlib.cmake +++ b/tensorflow/contrib/cmake/external/zlib.cmake @@ -12,50 +12,75 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -include (ExternalProject) - -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) - -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() - -set(ZLIB_HEADERS - "${ZLIB_INSTALL}/include/zconf.h" - "${ZLIB_INSTALL}/include/zlib.h" -) - -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 - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} - -DCMAKE_BUILD_TYPE:STRING=Release - -DCMAKE_INSTALL_PREFIX:STRING=${ZLIB_INSTALL} -) - -# 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() +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}) + + # 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) + +else (systemlib_ZLIB) + include (ExternalProject) + + 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) + + 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() + + set(ZLIB_HEADERS + "${ZLIB_INSTALL}/include/zconf.h" + "${ZLIB_INSTALL}/include/zlib.h" + ) + + 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} + ) + + # 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/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index 9ce8b3cc9cd4783c4b940ea9c7bf0b57fa2a3f28..0d2a6a23db26af2fb9498849aa93e74379915fe3 100644 --- a/tensorflow/contrib/cmake/python_modules.txt +++ b/tensorflow/contrib/cmake/python_modules.txt @@ -6,6 +6,7 @@ tensorflow/core/example tensorflow/core/framework tensorflow/core/lib tensorflow/core/lib/core +tensorflow/core/profiler tensorflow/core/protobuf tensorflow/core/util tensorflow/examples @@ -164,6 +165,7 @@ tensorflow/contrib/distributions/python tensorflow/contrib/distributions/python/ops tensorflow/contrib/distributions/python/ops/bijectors tensorflow/contrib/eager +tensorflow/contrib/eager/proto tensorflow/contrib/eager/python tensorflow/contrib/estimator tensorflow/contrib/estimator/python @@ -174,6 +176,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 @@ -216,6 +221,8 @@ tensorflow/contrib/input_pipeline/python/ops tensorflow/contrib/integrate tensorflow/contrib/integrate/python tensorflow/contrib/integrate/python/ops +tensorflow/contrib/kafka/python +tensorflow/contrib/kafka/python/ops tensorflow/contrib/keras tensorflow/contrib/keras/api tensorflow/contrib/keras/api/keras @@ -293,7 +300,9 @@ 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 @@ -323,8 +332,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 @@ -356,6 +363,7 @@ 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 @@ -405,6 +413,10 @@ 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/hybrid @@ -415,6 +427,7 @@ 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 @@ -435,6 +448,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_protos.txt b/tensorflow/contrib/cmake/python_protos.txt index 8a9c406d8b118c10ddcaafb0e4fc242aa79cdb57..c03c0c80fe62a4f95d0fcf240ee25725a19d86f0 100644 --- a/tensorflow/contrib/cmake/python_protos.txt +++ b/tensorflow/contrib/cmake/python_protos.txt @@ -4,6 +4,7 @@ tensorflow/python tensorflow/contrib/boosted_trees/proto tensorflow/contrib/cloud/kernels tensorflow/contrib/decision_trees/proto +tensorflow/contrib/eager/proto tensorflow/contrib/gdr tensorflow/contrib/lite/toco tensorflow/contrib/mpi 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..9e355da33a7258119b6086216f5487d7ea94716c --- /dev/null +++ b/tensorflow/contrib/cmake/tests/cuda/compatibility_test.c @@ -0,0 +1,22 @@ +/* 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 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..beb574061bea8d04af8386223749677ae36a5d9b --- /dev/null +++ b/tensorflow/contrib/cmake/tests/cuda/compatibility_test.cc @@ -0,0 +1,22 @@ +/* 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 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 f3cf3e70441de67ef79bc9cedf85549315170c29..f73da0b8ab18af1eca4c2bd577604595f8b8ec6d 100644 --- a/tensorflow/contrib/cmake/tf_cc_ops.cmake +++ b/tensorflow/contrib/cmake/tf_cc_ops.cmake @@ -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..a54cbff33b66d63d7229fa2f50b8a4ca962111ed 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" @@ -57,6 +63,12 @@ file(GLOB_RECURSE tf_core_cpu_exclude_srcs "${tensorflow_source_dir}/tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h" "${tensorflow_source_dir}/tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.cc" ) +file(GLOB_RECURSE tf_core_cpu_whitelisted_srcs + "${tensorflow_source_dir}/tensorflow/core/common_runtime/gpu/gpu_id.h" + "${tensorflow_source_dir}/tensorflow/core/common_runtime/gpu/gpu_id.cc" + "${tensorflow_source_dir}/tensorflow/core/common_runtime/gpu/gpu_id_manager.cc" +) +list(REMOVE_ITEM tf_core_cpu_exclude_srcs ${tf_core_cpu_whitelisted_srcs}) list(REMOVE_ITEM tf_core_cpu_srcs ${tf_core_cpu_exclude_srcs}) if (tensorflow_ENABLE_GPU) @@ -73,6 +85,7 @@ if (tensorflow_ENABLE_GPU) "${tensorflow_source_dir}/tensorflow/core/*test*.cc" ) list(REMOVE_ITEM tf_core_gpu_srcs ${tf_core_gpu_exclude_srcs}) + list(REMOVE_ITEM tf_core_gpu_srcs ${tf_core_cpu_whitelisted_srcs}) list(APPEND tf_core_cpu_srcs ${tf_core_gpu_srcs}) endif() diff --git a/tensorflow/contrib/cmake/tf_core_framework.cmake b/tensorflow/contrib/cmake/tf_core_framework.cmake index 129c208ecd6b574ed63c2fe378e1a6ebb92de558..a1c320347fe60f87806736befc677541a93e7e93 100644 --- a/tensorflow/contrib/cmake/tf_core_framework.cmake +++ b/tensorflow/contrib/cmake/tf_core_framework.cmake @@ -292,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" diff --git a/tensorflow/contrib/cmake/tf_core_kernels.cmake b/tensorflow/contrib/cmake/tf_core_kernels.cmake index 6927bf03f08b68a1f13f6a0978af629af45575e8..998f99ecc19f88921dce14fde892912fb699ad08 100644 --- a/tensorflow/contrib/cmake/tf_core_kernels.cmake +++ b/tensorflow/contrib/cmake/tf_core_kernels.cmake @@ -69,8 +69,11 @@ if(tensorflow_BUILD_CONTRIB_KERNELS) "${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/kernels/threadpool_dataset_op.cc" + "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/unique_dataset_op.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" diff --git a/tensorflow/contrib/cmake/tf_core_ops.cmake b/tensorflow/contrib/cmake/tf_core_ops.cmake index c42bc35ce7e75c72a8719f6730f2f8b63b35febb..59e094812aaf4da2549d96314fc550e5635f9de8 100644 --- a/tensorflow/contrib/cmake/tf_core_ops.cmake +++ b/tensorflow/contrib/cmake/tf_core_ops.cmake @@ -85,7 +85,7 @@ GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_quantiles "${tensorflow_source_dir}/te 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_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake index b7c816c24f82c7747f53b4c127866e6008085ef3..b730ebd3baacafe8ae401e8987104f3062372954 100755 --- a/tensorflow/contrib/cmake/tf_python.cmake +++ b/tensorflow/contrib/cmake/tf_python.cmake @@ -307,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 ) @@ -368,8 +368,8 @@ 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" @@ -541,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 @@ -549,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) @@ -702,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..9738bbeb9aebaeb67495127528e26634887d392c 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 @@ -96,8 +100,7 @@ if(WIN32) endif(WIN32) target_include_directories(tensorflow PUBLIC - $ - $) + $) install(TARGETS tensorflow EXPORT tensorflow_export RUNTIME DESTINATION bin @@ -129,10 +132,6 @@ install(DIRECTORY ${tensorflow_source_dir}/tensorflow/stream_executor/ install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/src/google/ DESTINATION include/google FILES_MATCHING PATTERN "*.h") -# nsync headers -install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/external/nsync/ - DESTINATION include/external/nsync - FILES_MATCHING PATTERN "*.h") # Eigen directory install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/eigen/src/eigen/Eigen/ DESTINATION include/Eigen) diff --git a/tensorflow/contrib/cmake/tf_tests.cmake b/tensorflow/contrib/cmake/tf_tests.cmake index 2e79eadf7f566690a7742757ceb56e147ebd6ea0..b3e5b30826097d6c747245fec975fcbea3785d15 100644 --- a/tensorflow/contrib/cmake/tf_tests.cmake +++ b/tensorflow/contrib/cmake/tf_tests.cmake @@ -156,6 +156,7 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${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" @@ -221,6 +222,7 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${tensorflow_source_dir}/tensorflow/python/debug/cli/curses_ui_test.py" # TFDBG grpc:// mode is not yet available on Windows. "${tensorflow_source_dir}/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py" + "${tensorflow_source_dir}/tensorflow/python/debug/lib/grpc_large_data_test.py" "${tensorflow_source_dir}/tensorflow/python/debug/lib/session_debug_grpc_test.py" "${tensorflow_source_dir}/tensorflow/python/debug/lib/source_remote_test.py" # stl on windows handles overflows different @@ -275,8 +277,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. @@ -310,6 +311,8 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${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}) @@ -473,6 +476,10 @@ if (tensorflow_BUILD_CC_TESTS) "${tensorflow_source_dir}/tensorflow/core/profiler/internal/advisor/*_test.cc" ) + list(REMOVE_ITEM tf_test_src_simple + ${tf_core_profiler_test_srcs} + ) + set(tf_test_lib tf_test_lib) add_library(${tf_test_lib} STATIC ${tf_src_testlib}) 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/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_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..1233c8f251c404c57d9e2b38993e7a386b1e6ceb 100644 --- a/tensorflow/contrib/crf/python/ops/crf.py +++ b/tensorflow/contrib/crf/python/ops/crf.py @@ -105,8 +105,8 @@ def crf_sequence_score(inputs, tag_indices, sequence_lengths, return utils.smart_cond( pred=math_ops.equal(inputs.shape[1].value or array_ops.shape(inputs)[1], 1), - fn1=_single_seq_fn, - fn2=_multi_seq_fn) + true_fn=_single_seq_fn, + false_fn=_multi_seq_fn) def crf_log_norm(inputs, sequence_lengths, transition_params): @@ -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 @@ -511,7 +511,7 @@ def crf_decode(potentials, transition_params, sequence_length): return decode_tags, best_score return utils.smart_cond( - pred=math_ops.equal( - potentials.shape[1].value or array_ops.shape(potentials)[1], 1), - fn1=_single_seq_fn, - fn2=_multi_seq_fn) + pred=math_ops.equal(potentials.shape[1].value or + array_ops.shape(potentials)[1], 1), + true_fn=_single_seq_fn, + false_fn=_multi_seq_fn) 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 56c562a3bad1f9e8a55a25eb3273fb2c9fbdd4b5..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,7 +20,7 @@ from __future__ import print_function import time -from six.moves import xrange +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 49d305cb0dd0387c34b7feb79ef631eac9e935cd..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 @@ -51,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 @@ -316,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 diff --git a/tensorflow/contrib/data/BUILD b/tensorflow/contrib/data/BUILD index 8ecc003348d70379ee48d050e63e93d0dd38efaa..5ba2297e7f86e706e6332b4cad3e8b6d34ead77d 100644 --- a/tensorflow/contrib/data/BUILD +++ b/tensorflow/contrib/data/BUILD @@ -27,13 +27,16 @@ 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", + "//tensorflow/core:lib_proto_parsing", + ], ) 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 daeb6a610533404044d42033709d644deb481024..9212b69700941c190df1d44ed308147105c56fba 100644 --- a/tensorflow/contrib/data/__init__.py +++ b/tensorflow/contrib/data/__init__.py @@ -12,26 +12,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""`tf.contrib.data` API for input pipelines. +"""Experimental API for building input pipelines. -This module contains the experimental (less stable) counterpart to the -`tf.data` API. See @{tf.data.Dataset} and @{tf.data.Iterator} for the -stable classes. +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 +@@SqlDataset @@batch_and_drop_remainder +@@bucket_by_sequence_length @@dense_to_sparse_batch @@enumerate_dataset @@group_by_window @@ignore_errors +@@make_batched_features_dataset @@make_saveable_from_iterator @@map_and_batch @@padded_batch_and_drop_remainder @@ -40,6 +40,7 @@ See the @{$datasets$Importing Data} Programmer's Guide for an overview. @@rejection_resample @@scan @@shuffle_and_repeat +@@sliding_window_batch @@sloppy_interleave @@unbatch @@ -58,23 +59,23 @@ 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 bucket_by_sequence_length 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 make_batched_features_dataset 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.contrib.data.python.ops.sliding import sliding_window_batch 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..c87da7dfaa5943f7918c370f63362673844c7f0e 100644 --- a/tensorflow/contrib/data/kernels/BUILD +++ b/tensorflow/contrib/data/kernels/BUILD @@ -10,6 +10,7 @@ cc_library( name = "prefetching_kernels", srcs = ["prefetching_kernels.cc"], deps = [ + "//tensorflow/core:core_cpu_headers_lib", "//tensorflow/core:framework_headers_lib", "//third_party/eigen3", "@protobuf_archive//:protobuf_headers", @@ -17,6 +18,50 @@ 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 = "threadpool_dataset_op", + srcs = ["threadpool_dataset_op.cc"], + deps = [ + "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", + "@protobuf_archive//:protobuf_headers", + ], +) + +cc_library( + name = "unique_dataset_op", + srcs = ["unique_dataset_op.cc"], + deps = [ + "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", + "@protobuf_archive//:protobuf_headers", + ], +) + +cc_library( + name = "dataset_kernels", + deps = [ + ":ignore_errors_dataset_op", + ":prefetching_kernels", + ":threadpool_dataset_op", + ":unique_dataset_op", + "//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 98% rename from tensorflow/core/kernels/data/ignore_errors_dataset_op.cc rename to tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc index 99df699d719b896df37515fc4147cd48db52a113..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 { diff --git a/tensorflow/contrib/data/kernels/prefetching_kernels.cc b/tensorflow/contrib/data/kernels/prefetching_kernels.cc index d3df14bdd03476e9ee4015b374512e5bb9893a63..1baac3ea5239659e65881e5b2dea4fe1a8c49d1b 100644 --- a/tensorflow/contrib/data/kernels/prefetching_kernels.cc +++ b/tensorflow/contrib/data/kernels/prefetching_kernels.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include +#include "tensorflow/core/common_runtime/process_function_library_runtime.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/resource_op_kernel.h" @@ -35,12 +36,14 @@ using FunctionBufferCallback = std::function; class FunctionBufferingResource : public ResourceBase { public: FunctionBufferingResource(FunctionLibraryRuntime* lib, + std::unique_ptr pflr, const NameAttrList& func, int64 buffer_size, const string& source_device, const string& target_device, const std::vector& func_args, int64 thread_pool_size) : lib_(lib), + pflr_(std::move(pflr)), func_(func), buffer_size_(buffer_size), source_device_(source_device), @@ -222,6 +225,7 @@ class FunctionBufferingResource : public ResourceBase { mutex mu_; FunctionLibraryRuntime* lib_; + std::unique_ptr pflr_; NameAttrList func_; const int64 buffer_size_; const string source_device_; @@ -241,7 +245,7 @@ class FunctionBufferingResource : public ResourceBase { class FunctionBufferResourceHandleOp : public OpKernel { public: explicit FunctionBufferResourceHandleOp(OpKernelConstruction* ctx) - : OpKernel(ctx) { + : OpKernel(ctx), flib_def_(nullptr) { OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &func_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("buffer_size", &buffer_size_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("container", &container_)); @@ -249,6 +253,17 @@ class FunctionBufferResourceHandleOp : public OpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("thread_pool_size", &thread_pool_size_)); } + ~FunctionBufferResourceHandleOp() override { + if (cinfo_.resource_is_private_to_kernel()) { + if (!cinfo_.resource_manager() + ->Delete(cinfo_.container(), + cinfo_.name()) + .ok()) { + // Do nothing; the resource can have been deleted by session resets. + } + } + } + void Compute(OpKernelContext* ctx) override { const Tensor* string_arg; OP_REQUIRES_OK(ctx, ctx->input("string_arg", &string_arg)); @@ -267,28 +282,39 @@ class FunctionBufferResourceHandleOp : public OpKernel { const string& source_device = ctx->device()->name(); - ContainerInfo cinfo; - OP_REQUIRES_OK(ctx, cinfo.Init(ctx->resource_manager(), def())); - // Create the resource. - FunctionBufferingResource* buffer; - OP_REQUIRES_OK( - ctx, ctx->resource_manager()->LookupOrCreate( - cinfo.container(), cinfo.name(), &buffer, - [lib, &source_device, &target_device, func_args, - this](FunctionBufferingResource** ptr) { - *ptr = new FunctionBufferingResource( - lib, func_, buffer_size_, source_device, target_device, - func_args, thread_pool_size_); - return Status::OK(); - })); - OP_REQUIRES_OK(ctx, buffer->Instantiate()); + mutex_lock l(mu_); + if (!initialized_) { + OP_REQUIRES_OK(ctx, cinfo_.Init(ctx->resource_manager(), def())); + FunctionLibraryRuntime* clone_lib; + std::unique_ptr pflr; + OP_REQUIRES_OK(ctx, lib->Clone(&flib_def_, &pflr, &clone_lib)); + // Create the resource. + FunctionBufferingResource* buffer; + OP_REQUIRES_OK( + ctx, + ctx->resource_manager()->LookupOrCreate( + cinfo_.container(), cinfo_.name(), &buffer, + [clone_lib, &pflr, &source_device, &target_device, func_args, + this](FunctionBufferingResource** ptr) { + *ptr = new FunctionBufferingResource( + clone_lib, std::move(pflr), func_, buffer_size_, + source_device, target_device, func_args, thread_pool_size_); + return Status::OK(); + })); + OP_REQUIRES_OK(ctx, buffer->Instantiate()); + initialized_ = true; + } OP_REQUIRES_OK(ctx, MakeResourceHandleToOutput( - ctx, 0, cinfo.container(), cinfo.name(), + ctx, 0, cinfo_.container(), cinfo_.name(), MakeTypeIndex())); } private: + mutex mu_; + ContainerInfo cinfo_ GUARDED_BY(mu_); + bool initialized_ GUARDED_BY(mu_) = false; + std::unique_ptr flib_def_; NameAttrList func_; int64 buffer_size_; string container_; diff --git a/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc b/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..63e19ae3f837c9d3cfb1221df64360ee74117f13 --- /dev/null +++ b/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/core/framework/dataset.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/lib/core/threadpool.h" + +namespace tensorflow { +namespace { + +class ThreadPoolResource : public ResourceBase { + public: + ThreadPoolResource(Env* env, const ThreadOptions& thread_options, + const string& name, int num_threads, bool low_latency_hint) + : thread_pool_(env, thread_options, name, num_threads, low_latency_hint) { + } + + // Schedules fn() for execution in the pool of threads. + void Schedule(std::function fn) { + thread_pool_.Schedule(std::move(fn)); + } + + string DebugString() override { return "ThreadPoolResource"; } + + private: + thread::ThreadPool thread_pool_; +}; + +// Creates a handle to a ThreadPool resource. Note that we don't use +// ResourceOpKernel here because the ThreadPoolResource constructor requires +// access to `OpKernelContext::env()`, which isn't provided by +// `ResourceOpKernel::CreateResource()`. +class ThreadPoolHandleOp : public OpKernel { + public: + explicit ThreadPoolHandleOp(OpKernelConstruction* ctx) : OpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("display_name", &display_name_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("num_threads", &num_threads_)); + OP_REQUIRES( + ctx, num_threads_ > 0, + errors::InvalidArgument("`num_threads` must be greater than zero.")); + } + + // The resource is deleted from the resource manager only when it is private + // to kernel. Ideally the resource should be deleted when it is no longer held + // by anyone, but it would break backward compatibility. + ~ThreadPoolHandleOp() override { + if (cinfo_.resource_is_private_to_kernel()) { + if (!cinfo_.resource_manager() + ->Delete(cinfo_.container(), cinfo_.name()) + .ok()) { + // Do nothing; the resource can have been deleted by session resets. + } + } + } + + void Compute(OpKernelContext* ctx) override LOCKS_EXCLUDED(mu_) { + mutex_lock l(mu_); + if (!initialized_) { + ResourceMgr* mgr = ctx->resource_manager(); + OP_REQUIRES_OK(ctx, cinfo_.Init(mgr, def())); + ThreadPoolResource* resource; + OP_REQUIRES_OK(ctx, mgr->LookupOrCreate( + cinfo_.container(), cinfo_.name(), &resource, + [this, ctx](ThreadPoolResource** ret) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { + *ret = new ThreadPoolResource( + ctx->env(), {}, display_name_, + num_threads_, + false /* low_latency_hint */); + return Status::OK(); + })); + initialized_ = true; + } + OP_REQUIRES_OK(ctx, MakeResourceHandleToOutput( + ctx, 0, cinfo_.container(), cinfo_.name(), + MakeTypeIndex())); + } + + private: + mutex mu_; + ContainerInfo cinfo_ GUARDED_BY(mu_); + bool initialized_ GUARDED_BY(mu_) = false; + string display_name_; + int num_threads_; +}; + +class ThreadPoolDatasetOp : public UnaryDatasetOpKernel { + public: + explicit ThreadPoolDatasetOp(OpKernelConstruction* ctx) + : UnaryDatasetOpKernel(ctx) {} + + void MakeDataset(OpKernelContext* ctx, DatasetBase* input, + DatasetBase** output) override { + ThreadPoolResource* threadpool_resource; + OP_REQUIRES_OK(ctx, LookupResource(ctx, HandleFromInput(ctx, 1), + &threadpool_resource)); + core::ScopedUnref unref_iterator(threadpool_resource); + + *output = new Dataset(ctx, input, threadpool_resource); + } + + private: + class Dataset : public GraphDatasetBase { + public: + Dataset(OpKernelContext* ctx, const DatasetBase* input, + ThreadPoolResource* threadpool) + : GraphDatasetBase(ctx), input_(input), threadpool_(threadpool) { + input_->Ref(); + threadpool_->Ref(); + } + + ~Dataset() override { + input_->Unref(); + threadpool_->Unref(); + } + + std::unique_ptr MakeIterator( + const string& prefix) const override { + return std::unique_ptr( + new Iterator({this, strings::StrCat(prefix, "::ThreadPool")})); + } + + const DataTypeVector& output_dtypes() const override { + return input_->output_dtypes(); + } + const std::vector& output_shapes() const override { + return input_->output_shapes(); + } + + string DebugString() override { return "ThreadPoolDatasetOp::Dataset"; } + + protected: + Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b, + Node** output) const override { + return errors::Unimplemented( + "Cannot currently serialize the thread pool for a " + "ThreadPoolDataset."); + } + + private: + class Iterator : public DatasetIterator { + public: + explicit Iterator(const Params& params) + : DatasetIterator(params), + input_impl_(params.dataset->input_->MakeIterator(params.prefix)) {} + + Status GetNextInternal(IteratorContext* ctx, + std::vector* out_tensors, + bool* end_of_sequence) override { + ThreadPoolResource* pool = dataset()->threadpool_; + IteratorContext::Params params; + params.env = ctx->env(); + params.runner = [pool](std::function c) { + pool->Schedule(std::move(c)); + }; + params.stats_aggregator_getter = ctx->stats_aggregator_getter(); + params.lib = ctx->lib(); + params.function_library = ctx->function_library(); + params.allocator_getter = ctx->allocator_getter(); + IteratorContext threadpool_ctx(params); + return input_impl_->GetNext(&threadpool_ctx, out_tensors, + end_of_sequence); + } + + private: + std::unique_ptr input_impl_; + }; + + const DatasetBase* const input_; + ThreadPoolResource* const threadpool_; + }; +}; + +REGISTER_KERNEL_BUILDER(Name("ThreadPoolHandle").Device(DEVICE_CPU), + ThreadPoolHandleOp); +REGISTER_KERNEL_BUILDER(Name("ThreadPoolDataset").Device(DEVICE_CPU), + ThreadPoolDatasetOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/core/kernels/data/unique_dataset_op.cc b/tensorflow/contrib/data/kernels/unique_dataset_op.cc similarity index 99% rename from tensorflow/core/kernels/data/unique_dataset_op.cc rename to tensorflow/contrib/data/kernels/unique_dataset_op.cc index 7726ee0edf71b34cb65fe5fceb2b60dd30bb58e2..69fbb0fcdcce87951d2c9b84210fda378081b103 100644 --- a/tensorflow/core/kernels/data/unique_dataset_op.cc +++ b/tensorflow/contrib/data/kernels/unique_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/hash/hash.h" namespace tensorflow { diff --git a/tensorflow/contrib/data/ops/prefetching_ops.cc b/tensorflow/contrib/data/ops/dataset_ops.cc similarity index 56% rename from tensorflow/contrib/data/ops/prefetching_ops.cc rename to tensorflow/contrib/data/ops/dataset_ops.cc index 23cb62b6f0dbfed15667dd00ae0039b33aa944d4..a4c1212da11a2410461a120ed5f7116e80e4b903 100644 --- a/tensorflow/contrib/data/ops/prefetching_ops.cc +++ b/tensorflow/contrib/data/ops/dataset_ops.cc @@ -17,6 +17,26 @@ 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("UniqueDataset") + .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 unique elements of `input_dataset`. +)doc"); + REGISTER_OP("FunctionBufferingResource") .Input("string_arg: string") .Input("target_device: string") @@ -55,4 +75,33 @@ output: A list of return values. output_types: The type list for the return values. )doc"); +REGISTER_OP("ThreadPoolDataset") + .Input("input_dataset: variant") + .Input("thread_pool: resource") + .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 uses a custom thread pool to compute `input_dataset`. + +handle: A resource produced by the ThreadPoolHandle op. +)doc"); + +REGISTER_OP("ThreadPoolHandle") + .Output("handle: resource") + .SetShapeFn(shape_inference::ScalarShape) + .Attr("num_threads: int") + .Attr("display_name: string") + .Attr("container: string = ''") + .Attr("shared_name: string = ''") + .Doc(R"doc( +Creates a custom thread pool with the given number of threads. + +handle: A resource that can be consumed by one or more ThreadPoolDataset ops. +num_threads: The number of threads in the thread pool. +display_name: A human-readable name for the threads that may be visible in + some visualizations. +)doc"); + } // namespace tensorflow diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD index 04a21f2b0f5b508fcd2e6bafa6a6ebdb555329b4..2c4d4adfdad6d2b3268896cb91cd0357b2b814d9 100644 --- a/tensorflow/contrib/data/python/kernel_tests/BUILD +++ b/tensorflow/contrib/data/python/kernel_tests/BUILD @@ -52,24 +52,6 @@ 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", @@ -186,8 +168,10 @@ py_test( srcs = ["interleave_dataset_op_test.py"], srcs_version = "PY2AND3", tags = [ + "manual", "no_oss", "no_pip", + "notap", ], deps = [ ":dataset_serialization_test", @@ -208,73 +192,18 @@ py_test( ) tf_py_test( - name = "iterator_ops_cluster_test", + name = "get_single_element_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", - 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", ], ) @@ -365,9 +294,12 @@ py_test( "//tensorflow/python:errors", "//tensorflow/python:framework_ops", "//tensorflow/python:lib", + "//tensorflow/python:math_ops", "//tensorflow/python:parsing_ops", + "//tensorflow/python:string_ops", "//tensorflow/python:util", "//tensorflow/python/data/ops:iterator_ops", + "//third_party/py/numpy", ], ) @@ -426,14 +358,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", ], ) @@ -489,6 +424,20 @@ py_test( ], ) +py_test( + name = "threadpool_dataset_ops_test", + size = "small", + srcs = ["threadpool_dataset_ops_test.py"], + srcs_version = "PY2AND3", + tags = ["no_pip"], + deps = [ + "//tensorflow/contrib/data/python/ops:dataset_ops", + "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:errors", + ], +) + py_test( name = "unique_dataset_op_test", size = "small", @@ -535,7 +484,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", @@ -549,6 +498,23 @@ py_test( ], ) +tf_py_test( + name = "slide_dataset_op_test", + size = "small", + srcs = ["slide_dataset_op_test.py"], + additional_deps = [ + "//tensorflow/contrib/data/python/ops:dataset_ops", + "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:errors", + "//tensorflow/python:math_ops", + "//tensorflow/python:sparse_tensor", + "//third_party/py/numpy", + ], +) + filegroup( name = "all_files", srcs = glob( 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..a2da953c7bfb0566799b9750b5a4dacf7e337dc5 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 = ( @@ -646,6 +387,20 @@ class BatchDatasetTest(test.TestCase): def testBatchAndMapDatasetWithParallelBatching(self): return self._testBatchAndMapDatasetHelper(num_parallel_batches=10) + def testMapAndBatchYieldsPartialBatch(self): + iterator = (dataset_ops.Dataset.range(10) + .apply(batching.map_and_batch( + lambda x: array_ops.reshape(x * x, [1]), 4)) + .make_one_shot_iterator()) + self.assertEqual([None, 1], iterator.output_shapes.as_list()) + next_element = iterator.get_next() + with self.test_session() as sess: + self.assertAllEqual([[0], [1], [4], [9]], sess.run(next_element)) + self.assertAllEqual([[16], [25], [36], [49]], sess.run(next_element)) + self.assertAllEqual([[64], [81]], sess.run(next_element)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + def testMapAndBatchSparse(self): def _sparse(i): @@ -744,6 +499,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 6de93059d8cfbcf43c4e359b6973c33f22fff5fa..94f800e8a58bc34eef3034cd976b931528c01940 100644 --- a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py @@ -17,11 +17,13 @@ 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.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 @@ -379,5 +381,93 @@ class BucketTest(test.TestCase): self.assertEqual(batches, 15) +class BucketBySequenceLength(test.TestCase): + + def testBucket(self): + + boundaries = [10, 20, 30] + batch_sizes = [10, 8, 4, 2] + lengths = [8, 13, 25, 35] + + def element_gen(): + # Produce 1 batch for each bucket + elements = [] + for batch_size, length in zip(batch_sizes, lengths): + for _ in range(batch_size): + elements.append([1] * length) + random.shuffle(elements) + for el in elements: + yield (el,) + + element_len = lambda el: array_ops.shape(el)[0] + dataset = dataset_ops.Dataset.from_generator( + element_gen, (dtypes.int64,), ([None],)).apply( + grouping.bucket_by_sequence_length( + element_len, boundaries, batch_sizes)) + batch, = dataset.make_one_shot_iterator().get_next() + + with self.test_session() as sess: + batches = [] + for _ in range(4): + batches.append(sess.run(batch)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(batch) + batch_sizes_val = [] + lengths_val = [] + for batch in batches: + batch_size = batch.shape[0] + length = batch.shape[1] + batch_sizes_val.append(batch_size) + lengths_val.append(length) + self.assertEqual(sum(batch_sizes_val), sum(batch_sizes)) + self.assertEqual(sorted(batch_sizes), sorted(batch_sizes_val)) + self.assertEqual(sorted(lengths), sorted(lengths_val)) + + def testPadToBoundary(self): + + boundaries = [10, 20, 30] + batch_sizes = [10, 8, 4, 2] + lengths = [8, 13, 25] + + def element_gen(): + # Produce 1 batch for each bucket + elements = [] + for batch_size, length in zip(batch_sizes[:-1], lengths): + for _ in range(batch_size): + elements.append([1] * length) + random.shuffle(elements) + for el in elements: + yield (el,) + for _ in range(batch_sizes[-1]): + el = [1] * (boundaries[-1] + 5) + yield (el,) + + element_len = lambda el: array_ops.shape(el)[0] + dataset = dataset_ops.Dataset.from_generator( + element_gen, (dtypes.int64,), ([None],)).apply( + grouping.bucket_by_sequence_length( + element_len, boundaries, batch_sizes, + pad_to_bucket_boundary=True)) + batch, = dataset.make_one_shot_iterator().get_next() + + with self.test_session() as sess: + batches = [] + for _ in range(3): + batches.append(sess.run(batch)) + with self.assertRaisesOpError("bucket_boundaries"): + sess.run(batch) + batch_sizes_val = [] + lengths_val = [] + for batch in batches: + batch_size = batch.shape[0] + length = batch.shape[1] + batch_sizes_val.append(batch_size) + lengths_val.append(length) + batch_sizes = batch_sizes[:-1] + self.assertEqual(sum(batch_sizes_val), sum(batch_sizes)) + self.assertEqual(sorted(batch_sizes), sorted(batch_sizes_val)) + self.assertEqual(sorted(boundaries), sorted(lengths_val)) + + if __name__ == "__main__": test.main() 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 063c71063601002af8168c4facf4057433061ab7..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 @@ -20,117 +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.data.util import nest -from tensorflow.python.framework import errors -from tensorflow.python.framework import tensor_shape +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.platform import test -class ConcatenateDatasetTest(test.TestCase): - - 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) - - class ConcatenateDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): 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 a90ba30e60cef13156719bba24fb553c0acec391..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,713 +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() - sess.run(iterator.initializer) - - with self.assertRaisesRegexp( - errors.FailedPreconditionError, - "Error while reading resource variable Variable"): - sess.run(iterator.get_next()) - 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 701fc8247e21e4c018c704d82f3a40c7daf1d742..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,6 +24,7 @@ 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 @@ -35,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. @@ -234,6 +250,7 @@ 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) self._initialize(init_op, sess) @@ -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,6 +408,7 @@ 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: self._initialize(init_op, sess) for _ in range(break_point): @@ -484,11 +504,13 @@ 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: @@ -559,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: @@ -578,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) @@ -593,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: @@ -608,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") 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..87b7c6ddb7afcbaaf8fe97cd8be87e6f5af8cd4d --- /dev/null +++ b/tensorflow/contrib/data/python/kernel_tests/get_single_element_test.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. +# ============================================================================== +"""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.framework import sparse_tensor +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=[]) + + def make_sparse(x): + x_1d = array_ops.reshape(x, [1]) + x_2d = array_ops.reshape(x, [1, 1]) + return sparse_tensor.SparseTensor(x_2d, x_1d, x_1d) + + dataset = (dataset_ops.Dataset.range(100) + .skip(skip_value) + .map(lambda x: (x * x, make_sparse(x))) + .take(take_value)) + + element = get_single_element.get_single_element(dataset) + + with self.test_session() as sess: + for x in [0, 5, 10]: + dense_val, sparse_val = sess.run(element, feed_dict={skip_value: x}) + self.assertEqual(x * x, dense_val) + self.assertAllEqual([[x]], sparse_val.indices) + self.assertAllEqual([x], sparse_val.values) + self.assertAllEqual([x], sparse_val.dense_shape) + + 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 b1937c08f347734d0d6871bd30ed209ff520623a..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): 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 dd8247bfd47a9880c7cfe905103702e43b1f2165..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,15 +16,12 @@ 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 as contrib_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 @@ -34,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 @@ -50,231 +41,11 @@ 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 ( - contrib_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 (contrib_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 = ( - contrib_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 = ( - contrib_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 = ( - contrib_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 = ( - contrib_dataset_ops.Dataset.from_tensor_slices(components) + 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() @@ -292,10 +63,9 @@ class MapDatasetTest(test.TestCase): components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) dataset = ( - contrib_dataset_ops.Dataset.from_tensor_slices(components).map( + 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())) + 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() @@ -317,8 +87,8 @@ class MapDatasetTest(test.TestCase): write_string_to_file(filename, filename) dataset = ( - contrib_dataset_ops.Dataset.from_tensor_slices(filenames).map( - io_ops.read_file, num_threads=2, output_buffer_size=2).apply( + 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 @@ -343,350 +113,6 @@ 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 = contrib_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 = ( - contrib_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() - - iterator = ( - contrib_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() - - 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 testCaptureVariable(self): - counter_var = variable_scope.get_variable( - "counter", (), dtypes.int32, use_resource=True) - iterator = ( - contrib_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() - - 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 = ( - contrib_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() - - with self.test_session() as sess: - sess.run(init_op) - with self.assertRaises(errors.NotFoundError): - sess.run(get_next) - - def testSeededStatefulOperatorIsProperlyStateful(self): - iterator = ( - contrib_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)) - - # Randomness is repeatable given same seed - self.assertAllClose(random_values, random_values_2) - - def testMapDict(self): - iterator = (contrib_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() - - 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 testMapNamedtuple(self, count=10): - # construct dataset of tuples - labels = contrib_dataset_ops.Dataset.range(count) - images = labels.map(lambda l: -l) - dataset_tuple = contrib_dataset_ops.Dataset.zip((labels, images)) - - # convert dataset of tuples to dataset of namedtuples - example = namedtuple("Example", ["label", "image"]) - dataset_namedtuple = dataset_tuple.map(example) - - def preprocess_tuple(label, image): - image = 2 * image - return label, image - - def preprocess_namedtuple(example): - return example._replace(image=2 * example.image) - - # preprocess both datasets - dataset_tuple = dataset_tuple.map(preprocess_tuple) - dataset_namedtuple = dataset_namedtuple.map(preprocess_namedtuple) - - next_tuple = dataset_tuple.make_one_shot_iterator().get_next() - next_namedtuple = dataset_namedtuple.make_one_shot_iterator().get_next() - - # 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)) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_namedtuple) - - def testUseStepContainerInMap(self): - row = np.arange(6) - iterator = ( - contrib_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() - - 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 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() - - set_event_during_invocation = 5 - - def _map_py_func(x): - if x == set_event_during_invocation: - ev.set() - return x * x - - def _map_fn(x): - return script_ops.py_func(_map_py_func, [x], x.dtype) - - buffer_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = ( - contrib_dataset_ops.Dataset.range(100).map(_map_fn) - .prefetch(buffer_size_placeholder).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - 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) - - # 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) - - def testReturnList(self): - iterator = ( - contrib_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() - - 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 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 = ( - contrib_dataset_ops.Dataset.range(10).map(_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): - self.assertEqual((i, 37.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 _sparse(i): - return sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0]]), - values=(i * np.array([1])), - dense_shape=np.array([1, 1])) - - iterator = ( - contrib_dataset_ops.Dataset.range(10).map(_sparse) - .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, _sparse(i)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSparseChain(self): - - def _sparse(i): - return sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0]]), - 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 = ( - contrib_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): @@ -695,10 +121,10 @@ class MapDatasetTest(test.TestCase): get_next = iterator.get_next() return x * get_next - return contrib_dataset_ops.Dataset.range(10).map(_map_fn) + return dataset_ops.Dataset.range(10).map(_map_fn) def _build_graph(): - captured_iterator = contrib_dataset_ops.Dataset.range( + captured_iterator = dataset_ops.Dataset.range( 10).make_initializable_iterator() ds = _build_ds(captured_iterator) iterator = ds.make_initializable_iterator() @@ -734,7 +160,7 @@ class MapDatasetSerializationTest( return math_ops.square(x), math_ops.square(y), math_ops.square(z) return ( - contrib_dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) + dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) .repeat(self._num_epochs)) def testSaveRestoreCore(self): @@ -751,7 +177,7 @@ class MapDatasetSerializationTest( return random_ops.random_uniform( (), 0, 10, dtype=dtypes.int32) * math_ops.to_int32(x) - return contrib_dataset_ops.Dataset.range(100).map(_map_fn) + return dataset_ops.Dataset.range(100).map(_map_fn) self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) @@ -760,7 +186,7 @@ class MapDatasetSerializationTest( def _build_ds(): counter_var = variable_scope.get_variable( "counter", (), dtypes.int32, use_resource=True) - return (contrib_dataset_ops.Dataset.from_tensors(0).repeat(10).map( + return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( lambda _: counter_var.assign_add(1))) self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) @@ -769,7 +195,7 @@ class MapDatasetSerializationTest( def _build_ds(): constant_var = constant_op.constant(5) - return (contrib_dataset_ops.Dataset.from_tensors(0).repeat(10).map( + return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( lambda x: x + constant_var)) self.run_core_tests(_build_ds, None, 10) @@ -783,7 +209,7 @@ class MapDatasetSerializationTest( def defun_fn(x): return constant_op.constant(1000) + math_ops.to_int32(x) - return contrib_dataset_ops.Dataset.range(num_outputs).map(defun_fn) + return dataset_ops.Dataset.range(num_outputs).map(defun_fn) self.run_core_tests(_build_ds, None, num_outputs) @@ -801,10 +227,25 @@ class MapDatasetSerializationTest( return constant_op.constant(11000) + defun_fn_deep(math_ops.to_int32(x)) - return contrib_dataset_ops.Dataset.range(num_outputs).map(defun_fn) + return dataset_ops.Dataset.range(num_outputs).map(defun_fn) self.run_core_tests(_build_ds, None, num_outputs) + def testSparseCore(self): + + def _sparse(i): + return sparse_tensor.SparseTensorValue( + indices=np.array([[0, 0]]), + values=(i * np.array([1])), + dense_shape=np.array([1, 1])) + + def _build_ds(num_outputs): + return dataset_ops.Dataset.range(num_outputs).map(_sparse) + + num_outputs = 10 + self.run_core_tests(lambda: _build_ds(num_outputs), + lambda: _build_ds(int(num_outputs / 2)), num_outputs) + class ParallelMapDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): @@ -851,7 +292,8 @@ class ParallelMapDatasetSerializationTest( return random_ops.random_uniform( (), 0, 10, dtype=dtypes.int32) * math_ops.to_int32(x) - return contrib_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) @@ -860,8 +302,9 @@ class ParallelMapDatasetSerializationTest( def _build_ds(): counter_var = variable_scope.get_variable( "counter", (), dtypes.int32, use_resource=True) - return (contrib_dataset_ops.Dataset.from_tensors(0).repeat(10).map( - lambda _: counter_var.assign_add(1))) + return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( + lambda _: counter_var.assign_add(1), + num_parallel_calls=2).prefetch(2)) self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) @@ -869,8 +312,8 @@ class ParallelMapDatasetSerializationTest( def _build_ds(): constant_var = constant_op.constant(5) - return (contrib_dataset_ops.Dataset.from_tensors(0).repeat(10).map( - lambda x: x + constant_var)) + 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) @@ -883,7 +326,8 @@ class ParallelMapDatasetSerializationTest( def defun_fn(x): return constant_op.constant(1000) + math_ops.to_int32(x) - return contrib_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) @@ -901,7 +345,8 @@ class ParallelMapDatasetSerializationTest( return constant_op.constant(11000) + defun_fn_deep(math_ops.to_int32(x)) - return contrib_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) @@ -910,7 +355,7 @@ class IgnoreErrorsSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): def _build_ds(self, components): - return contrib_dataset_ops.Dataset.from_tensor_slices(components).map( + return dataset_ops.Dataset.from_tensor_slices(components).map( lambda x: array_ops.check_numerics(x, "message")).apply( error_ops.ignore_errors()) 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 a431670829ed1d66f1719985af73eafa1fe45982..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 @@ -21,14 +21,13 @@ 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.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 @@ -38,131 +37,6 @@ from tensorflow.python.platform import test class RangeDatasetTest(test.TestCase): - 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) 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..699e8e7865502facd05e0c4d6d4f01b80f7c050c 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 @@ -21,18 +21,23 @@ import gzip import os import zlib +import numpy as np + from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base 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 from tensorflow.python.framework import ops from tensorflow.python.lib.io import python_io from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import parsing_ops +from tensorflow.python.ops import string_ops from tensorflow.python.platform import test from tensorflow.python.util import compat @@ -76,101 +81,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 +133,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 +166,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 +190,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 +221,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) @@ -562,12 +266,19 @@ class ReadBatchFeaturesTest(test.TestCase): self._num_records = 7 self.test_filenames = self._createFiles() - def _read_batch_features(self, filenames, num_epochs, batch_size): + def _read_batch_features(self, + filenames, + num_epochs, + batch_size, + reader_num_threads=1, + parser_num_threads=1, + shuffle=False, + shuffle_seed=None): self.filenames = filenames self.num_epochs = num_epochs self.batch_size = batch_size - return readers.read_batch_features( + return readers.make_batched_features_dataset( file_pattern=self.filenames, batch_size=self.batch_size, features={ @@ -575,9 +286,13 @@ class ReadBatchFeaturesTest(test.TestCase): "record": parsing_ops.FixedLenFeature([], dtypes.int64), "keywords": parsing_ops.VarLenFeature(dtypes.string) }, - reader=readers.TFRecordDataset, - randomize_input=False, - num_epochs=self.num_epochs) + reader=core_readers.TFRecordDataset, + num_epochs=self.num_epochs, + shuffle=shuffle, + shuffle_seed=shuffle_seed, + reader_num_threads=reader_num_threads, + parser_num_threads=parser_num_threads).make_one_shot_iterator( + ).get_next() def _record(self, f, r): example = example_pb2.Example(features=feature_pb2.Features( @@ -612,24 +327,35 @@ class ReadBatchFeaturesTest(test.TestCase): writer.close() return filenames - def _next_actual_batch(self, sess): - file_op = self.outputs["file"] - keywords_indices_op = self.outputs["keywords"].indices - keywords_values_op = self.outputs["keywords"].values - keywords_dense_shape_op = self.outputs["keywords"].dense_shape - record_op = self.outputs["record"] + def _run_actual_batch(self, outputs, sess): + file_op = outputs["file"] + keywords_indices_op = outputs["keywords"].indices + keywords_values_op = outputs["keywords"].values + keywords_dense_shape_op = outputs["keywords"].dense_shape + record_op = outputs["record"] return sess.run([ file_op, keywords_indices_op, keywords_values_op, keywords_dense_shape_op, record_op ]) - def _next_expected_batch(self, file_indices, batch_size, num_epochs): + def _next_actual_batch(self, sess): + return self._run_actual_batch(self.outputs, sess) + + def _next_expected_batch(self, + file_indices, + batch_size, + num_epochs, + cycle_length=1): def _next_record(file_indices): for j in file_indices: for i in range(self._num_records): yield j, i + def _next_record_interleaved(file_indices, cycle_length): + return self._interleave([_next_record([i]) for i in file_indices], + cycle_length) + file_batch = [] keywords_batch_indices = [] keywords_batch_values = [] @@ -637,7 +363,11 @@ class ReadBatchFeaturesTest(test.TestCase): record_batch = [] batch_index = 0 for _ in range(num_epochs): - for record in _next_record(file_indices): + if cycle_length == 1: + next_records = _next_record(file_indices) + else: + next_records = _next_record_interleaved(file_indices, cycle_length) + for record in next_records: f = record[0] r = record[1] file_batch.append(f) @@ -665,14 +395,41 @@ class ReadBatchFeaturesTest(test.TestCase): [len(file_batch), keywords_batch_max_len], record_batch ] - def _verify_records(self, sess, batch_size, file_index=None, num_epochs=1): + def _interleave(self, iterators, cycle_length): + pending_iterators = iterators + open_iterators = [] + num_open = 0 + for i in range(cycle_length): + if pending_iterators: + open_iterators.append(pending_iterators.pop(0)) + num_open += 1 + + while num_open: + for i in range(min(cycle_length, len(open_iterators))): + if open_iterators[i] is None: + continue + try: + yield next(open_iterators[i]) + except StopIteration: + if pending_iterators: + open_iterators[i] = pending_iterators.pop(0) + else: + open_iterators[i] = None + num_open -= 1 + + def _verify_records(self, + sess, + batch_size, + file_index=None, + num_epochs=1, + interleave_cycle_length=1): if file_index is not None: file_indices = [file_index] else: file_indices = range(self._num_files) - for expected_batch in self._next_expected_batch(file_indices, batch_size, - num_epochs): + for expected_batch in self._next_expected_batch( + file_indices, batch_size, num_epochs, interleave_cycle_length): actual_batch = self._next_actual_batch(sess) for i in range(len(expected_batch)): self.assertAllEqual(expected_batch[i], actual_batch[i]) @@ -714,12 +471,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() @@ -736,6 +492,302 @@ class ReadBatchFeaturesTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(next_element) + def testReadWithFusedShuffleRepeatDataset(self): + num_epochs = 5 + total_records = num_epochs * self._num_records + for batch_size in [1, 2]: + # Test that shuffling with same seed produces the same result. + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + outputs1 = self._read_batch_features( + filenames=self.test_filenames[0], + num_epochs=num_epochs, + batch_size=batch_size, + shuffle=True, + shuffle_seed=5) + outputs2 = self._read_batch_features( + filenames=self.test_filenames[0], + num_epochs=num_epochs, + batch_size=batch_size, + shuffle=True, + shuffle_seed=5) + for _ in range(total_records // batch_size): + batch1 = self._run_actual_batch(outputs1, sess) + batch2 = self._run_actual_batch(outputs2, sess) + for i in range(len(batch1)): + self.assertAllEqual(batch1[i], batch2[i]) + + # Test that shuffling with different seeds produces a different order. + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + outputs1 = self._read_batch_features( + filenames=self.test_filenames[0], + num_epochs=num_epochs, + batch_size=batch_size, + shuffle=True, + shuffle_seed=5) + outputs2 = self._read_batch_features( + filenames=self.test_filenames[0], + num_epochs=num_epochs, + batch_size=batch_size, + shuffle=True, + shuffle_seed=15) + all_equal = True + for _ in range(total_records // batch_size): + batch1 = self._run_actual_batch(outputs1, sess) + batch2 = self._run_actual_batch(outputs2, sess) + for i in range(len(batch1)): + all_equal = all_equal and np.array_equal(batch1[i], batch2[i]) + self.assertFalse(all_equal) + + def testParallelReadersAndParsers(self): + num_epochs = 5 + for batch_size in [1, 2]: + for reader_num_threads in [2, 4]: + for parser_num_threads in [2, 4]: + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + self.outputs = self._read_batch_features( + filenames=self.test_filenames, + num_epochs=num_epochs, + batch_size=batch_size, + reader_num_threads=reader_num_threads, + parser_num_threads=parser_num_threads) + self._verify_records( + sess, + batch_size, + num_epochs=num_epochs, + interleave_cycle_length=reader_num_threads) + with self.assertRaises(errors.OutOfRangeError): + self._next_actual_batch(sess) + + +class MakeCsvDatasetTest(test.TestCase): + + COLUMN_TYPES = [ + dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64, dtypes.string + ] + COLUMNS = ["col%d" % i for i in range(len(COLUMN_TYPES))] + LABEL = COLUMNS[0] + + def setUp(self): + super(MakeCsvDatasetTest, self).setUp() + self._num_files = 2 + self._num_records = 7 + self._test_filenames = self._create_files() + + def _csv_values(self, fileno, recordno): + return [ + fileno, + recordno, + fileno * recordno * 0.5, + fileno * recordno + 0.5, + "record %d" % recordno if recordno % 2 == 1 else "", + ] + + def _csv_record(self, fileno, recordno): + return ",".join(str(v) for v in self._csv_values(fileno, recordno)) + + def _create_files(self): + filenames = [] + for i in range(self._num_files): + fn = os.path.join(self.get_temp_dir(), "csv_file%d.csv" % i) + filenames.append(fn) + f = open(fn, "w") + f.write(",".join(self.COLUMNS) + "\n") # header line + for j in range(self._num_records): + f.write(self._csv_record(i, j) + "\n") + f.write("# Some comment goes here. Should be ignored!\n") + f.close() + return filenames + + def _make_csv_dataset(self, + filenames, + defaults, + label_key=LABEL, + batch_size=1, + num_epochs=1, + shuffle=False, + shuffle_seed=None): + return readers.make_csv_dataset( + filenames, + column_keys=self.COLUMNS, + column_defaults=defaults, + label_key=label_key, + batch_size=batch_size, + num_epochs=num_epochs, + shuffle=shuffle, + shuffle_seed=shuffle_seed, + skip=1, + filter_fn= + lambda line: math_ops.not_equal(string_ops.substr(line, 0, 1), "#"), + ) + + def _next_actual_batch(self, file_indices, batch_size, num_epochs): + features = {col: list() for col in self.COLUMNS} + for _ in range(num_epochs): + for i in file_indices: + for j in range(self._num_records): + values = self._csv_values(i, j) + if not values[-1]: + values[-1] = "NULL" # null values in csv are interpreted as default + values[-1] = values[-1].encode("utf-8") + + # Regroup lists by column instead of row + for n, col in enumerate(self.COLUMNS): + features[col].append(values[n]) + if len(list(features.values())[0]) == batch_size: + yield features + features = {col: list() for col in self.COLUMNS} + + def _run_actual_batch(self, outputs, sess): + features, labels = sess.run(outputs) + batch = [features[k] for k in self.COLUMNS if k != self.LABEL] + batch.append(labels) + return batch + + def _verify_records( + self, + sess, + dataset, + file_indices, + label_key=LABEL, + batch_size=1, + num_epochs=1, + ): + iterator = dataset.make_one_shot_iterator() + get_next = iterator.get_next() + + for expected_features in self._next_actual_batch(file_indices, batch_size, + num_epochs): + actual_features = sess.run(get_next) + + if label_key is not None: + expected_labels = expected_features.pop(label_key) + # Compare labels + self.assertAllEqual(expected_labels, actual_features[1]) + actual_features = actual_features[0] # Extract features dict from tuple + + for k in expected_features.keys(): + # Compare features + self.assertAllEqual(expected_features[k], actual_features[k]) + + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + def test_make_csv_dataset(self): + defaults = [ + constant_op.constant([], dtype=d) for d in self.COLUMN_TYPES[:-1] + ] + defaults.append(constant_op.constant(["NULL"], dtype=dtypes.string)) + + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Basic test: read from file 0. + dataset = self._make_csv_dataset(self._test_filenames[0], defaults) + self._verify_records(sess, dataset, [0]) + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Basic test: read from file 1. + dataset = self._make_csv_dataset(self._test_filenames[1], defaults) + self._verify_records(sess, dataset, [1]) + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Read from both files. + dataset = self._make_csv_dataset(self._test_filenames, defaults) + self._verify_records(sess, dataset, range(self._num_files)) + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Read from both files. Exercise the `batch` and `num_epochs` parameters + # of make_csv_dataset and make sure they work. + dataset = self._make_csv_dataset( + self._test_filenames, defaults, batch_size=2, num_epochs=10) + self._verify_records( + sess, dataset, range(self._num_files), batch_size=2, num_epochs=10) + + def test_make_csv_dataset_with_no_label(self): + defaults = [ + constant_op.constant([], dtype=d) for d in self.COLUMN_TYPES[:-1] + ] + defaults.append(constant_op.constant(["NULL"], dtype=dtypes.string)) + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Read from both files. Make sure this works with no label key supplied. + dataset = self._make_csv_dataset( + self._test_filenames, + defaults, + batch_size=2, + num_epochs=10, + label_key=None) + self._verify_records( + sess, + dataset, + range(self._num_files), + batch_size=2, + num_epochs=10, + label_key=None) + + def test_make_csv_dataset_with_types(self): + defaults = [d for d in self.COLUMN_TYPES[:-1]] + defaults.append(constant_op.constant(["NULL"], dtype=dtypes.string)) + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + dataset = self._make_csv_dataset(self._test_filenames, defaults) + self._verify_records(sess, dataset, range(self._num_files)) + + def test_make_csv_dataset_with_shuffle(self): + total_records = self._num_files * self._num_records + defaults = [d for d in self.COLUMN_TYPES[:-1]] + defaults.append(constant_op.constant(["NULL"], dtype=dtypes.string)) + for batch_size in [1, 2]: + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Test that shuffling with the same seed produces the same result + dataset1 = self._make_csv_dataset( + self._test_filenames, + defaults, + batch_size=batch_size, + shuffle=True, + shuffle_seed=5) + dataset2 = self._make_csv_dataset( + self._test_filenames, + defaults, + batch_size=batch_size, + shuffle=True, + shuffle_seed=5) + outputs1 = dataset1.make_one_shot_iterator().get_next() + outputs2 = dataset2.make_one_shot_iterator().get_next() + for _ in range(total_records // batch_size): + batch1 = self._run_actual_batch(outputs1, sess) + batch2 = self._run_actual_batch(outputs2, sess) + for i in range(len(batch1)): + self.assertAllEqual(batch1[i], batch2[i]) + + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + # Test that shuffling with a different seed produces different results + dataset1 = self._make_csv_dataset( + self._test_filenames, + defaults, + batch_size=batch_size, + shuffle=True, + shuffle_seed=5) + dataset2 = self._make_csv_dataset( + self._test_filenames, + defaults, + batch_size=batch_size, + shuffle=True, + shuffle_seed=6) + outputs1 = dataset1.make_one_shot_iterator().get_next() + outputs2 = dataset2.make_one_shot_iterator().get_next() + all_equal = False + for _ in range(total_records // batch_size): + batch1 = self._run_actual_batch(outputs1, sess) + batch2 = self._run_actual_batch(outputs2, sess) + for i in range(len(batch1)): + all_equal = all_equal and np.array_equal(batch1[i], batch2[i]) + self.assertFalse(all_equal) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py index a76c6b1e39c3bae43b275505cc79af2e6476fdf2..3e100a54306053c2eafc548a9f111e78e93914de 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 dtypes from tensorflow.python.framework import errors from tensorflow.python.ops import random_ops 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 45943d56ecb4bc18a6221157d0eeeae4efdf23cc..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,144 +17,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections - 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 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 test -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( dataset_serialization_test_base.DatasetSerializationTestBase): diff --git a/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..33c48e20bea53b88d69a59e715af38b22dd2cbd4 --- /dev/null +++ b/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py @@ -0,0 +1,242 @@ +# 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 the experimental input pipeline ops.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.data.python.ops import sliding +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 +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +class SlideDatasetTest(test.TestCase): + + def testSlideDataset(self): + """Test an dataset that maps a TF function across its input elements.""" + 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=[]) + window_size = array_ops.placeholder(dtypes.int64, shape=[]) + stride = 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) + + # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> + # RepeatDataset(count) -> _SlideDataset(window_size, stride). + iterator = (dataset_ops.Dataset.from_tensor_slices(components) + .map(_map_fn) + .repeat(count) + .apply(sliding.sliding_window_batch(window_size, stride)) + .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: + # Slide over a finite input, where the window_size divides the + # total number of elements. + sess.run(init_op, feed_dict={count: 20, window_size: 14, stride: 7}) + # Same formula with convolution layer. + num_batches = (20 * 7 - 14) // 7 + 1 + 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*7 + j) % 7]**2, + result_component[j]) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Slide over a finite input, where the window_size does not + # divide the total number of elements. + sess.run(init_op, feed_dict={count: 20, window_size: 17, stride: 9}) + + num_batches = (20 * 7 - 17) // 9 + 1 + for i in range(num_batches): + result = sess.run(get_next) + for component, result_component in zip(components, result): + for j in range(17): + self.assertAllEqual(component[(i*9 + j) % 7]**2, + result_component[j]) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Slide over a finite input, which is less than window_size, + # should fail straight away. + sess.run(init_op, feed_dict={count: 1, window_size: 10, stride: 4}) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + sess.run(init_op, feed_dict={count: 1, window_size: 10, stride: 8}) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Slide over an empty input should fail straight away. + sess.run(init_op, feed_dict={count: 0, window_size: 8, stride: 4}) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Empty window_size should be an initialization time error. + with self.assertRaises(errors.InvalidArgumentError): + sess.run(init_op, feed_dict={count: 14, window_size: 0, stride: 0}) + + # Invalid stride should be an initialization time error. + with self.assertRaises(errors.InvalidArgumentError): + sess.run(init_op, feed_dict={count: 14, window_size: 3, stride: 0}) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(init_op, feed_dict={count: 14, window_size: 3, stride: 3}) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(init_op, feed_dict={count: 14, window_size: 3, stride: 5}) + + 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 testSlideSparse(self): + + def _sparse(i): + return sparse_tensor.SparseTensorValue( + indices=[[0]], values=(i * [1]), dense_shape=[1]) + + iterator = dataset_ops.Dataset.range(10).map(_sparse).apply( + sliding.sliding_window_batch(5, 3)).make_initializable_iterator() + init_op = iterator.initializer + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(init_op) + num_batches = (10 - 5) // 3 + 1 + for i in range(num_batches): + actual = sess.run(get_next) + expected = sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 0], [2, 0], [3, 0], [4, 0]], + values=[i * 3, i * 3 + 1, i * 3 + 2, i * 3 + 3, i * 3 + 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 testSlideSparseWithDifferentDenseShapes(self): + + def _sparse(i): + return sparse_tensor.SparseTensorValue( + indices=array_ops.expand_dims( + math_ops.range(i, dtype=dtypes.int64), 1), + values=array_ops.fill([math_ops.to_int32(i)], i), + dense_shape=[i]) + + iterator = dataset_ops.Dataset.range(10).map(_sparse).apply( + sliding.sliding_window_batch(5, 3)).make_initializable_iterator() + init_op = iterator.initializer + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(init_op) + num_batches = (10 - 5) // 3 + 1 + for i in range(num_batches): + actual = sess.run(get_next) + expected_indices = [] + expected_values = [] + for j in range(5): + for k in range(i * 3 + j): + expected_indices.append([j, k]) + expected_values.append(i * 3 + j) + expected = sparse_tensor.SparseTensorValue( + indices=expected_indices, + values=expected_values, + dense_shape=[5, i * 3 + 5 - 1]) + self.assertTrue(sparse_tensor.is_sparse(actual)) + self.assertSparseValuesEqual(actual, expected) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + def testNestedSlideSparse(self): + + def _sparse(i): + return sparse_tensor.SparseTensorValue( + indices=[[0]], values=(i * [1]), dense_shape=[1]) + + iterator = (dataset_ops.Dataset.range(10) + .map(_sparse) + .apply(sliding.sliding_window_batch(4, 2)) + .apply(sliding.sliding_window_batch(3, 1)) + .make_initializable_iterator()) + init_op = iterator.initializer + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(init_op) + # Slide: 1st batch. + actual = sess.run(get_next) + expected = sparse_tensor.SparseTensorValue( + indices=[[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], + [1, 0, 0], [1, 1, 0], [1, 2, 0], [1, 3, 0], + [2, 0, 0], [2, 1, 0], [2, 2, 0], [2, 3, 0]], + values=[0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7], + dense_shape=[3, 4, 1]) + self.assertTrue(sparse_tensor.is_sparse(actual)) + self.assertSparseValuesEqual(actual, expected) + # Slide: 2nd batch. + actual = sess.run(get_next) + expected = sparse_tensor.SparseTensorValue( + indices=[[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], + [1, 0, 0], [1, 1, 0], [1, 2, 0], [1, 3, 0], + [2, 0, 0], [2, 1, 0], [2, 2, 0], [2, 3, 0]], + values=[2, 3, 4, 5, 4, 5, 6, 7, 6, 7, 8, 9], + dense_shape=[3, 4, 1]) + self.assertTrue(sparse_tensor.is_sparse(actual)) + self.assertSparseValuesEqual(actual, expected) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + def testSlideShapeError(self): + + def generator(): + yield [1.0, 2.0, 3.0] + yield [4.0, 5.0, 6.0] + yield [7.0, 8.0, 9.0, 10.0] + + iterator = (dataset_ops.Dataset.from_generator(generator, dtypes.float32, + output_shapes=[None]) + .apply(sliding.sliding_window_batch(3, 1)) + .make_initializable_iterator()) + next_element = iterator.get_next() + + with self.test_session() as sess: + sess.run(iterator.initializer) + with self.assertRaisesRegexp( + errors.InvalidArgumentError, + r"Cannot batch tensors with different shapes in component 0. " + r"First element had shape \[3\] and element 2 had shape \[4\]."): + sess.run(next_element) + + +if __name__ == "__main__": + test.main() 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/threadpool_dataset_ops_test.py b/tensorflow/contrib/data/python/kernel_tests/threadpool_dataset_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9167cb3379bba5cb1ba76a96549395c45dca9e35 --- /dev/null +++ b/tensorflow/contrib/data/python/kernel_tests/threadpool_dataset_ops_test.py @@ -0,0 +1,77 @@ +# 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 the experimental input pipeline statistics gathering ops.""" +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.ops import threadpool +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.ops import script_ops +from tensorflow.python.platform import test + + +class OverrideThreadpoolDatasetTest(test.TestCase): + + def testNumThreads(self): + + def get_thread_id(_): + # Python creates a dummy thread object to represent the current + # thread when called from an "alien" thread (such as a + # `PrivateThreadPool` thread in this case). It does not include + # the TensorFlow-given display name, but it has a unique + # identifier that maps one-to-one with the underlying OS thread. + return np.array(threading.current_thread().ident).astype(np.int64) + + for num_threads in [1, 2, 4, 8, 16]: + + dataset = ( + dataset_ops.Dataset.range(1000).map( + lambda x: script_ops.py_func(get_thread_id, [x], dtypes.int64), + num_parallel_calls=32).apply(unique.unique())) + + dataset = threadpool.override_threadpool( + dataset, + threadpool.PrivateThreadPool( + num_threads, display_name="private_thread_pool_%d" % num_threads)) + + iterator = dataset.make_initializable_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + sess.run(iterator.initializer) + thread_ids = [] + try: + while True: + thread_ids.append(sess.run(next_element)) + except errors.OutOfRangeError: + pass + self.assertEqual(len(thread_ids), len(set(thread_ids))) + self.assertGreater(len(thread_ids), 0) + # NOTE(mrry): We don't control the thread pool scheduling, and + # so cannot guarantee that all of the threads in the pool will + # perform work. + self.assertLessEqual(len(thread_ids), num_threads) + + +if __name__ == "__main__": + test.main() 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..c3331e963602d60fe27dd44b0cc06dfb20ca2b6a 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 = [ @@ -67,6 +67,8 @@ py_library( srcs_version = "PY2AND3", deps = [ ":dataset_ops", + ":shuffle_ops", + "//tensorflow/python:constant_op", "//tensorflow/python:dataset_ops_gen", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", @@ -104,11 +106,15 @@ py_library( "interleave_ops.py", "resampling.py", "scan_ops.py", + "sliding.py", "stats_ops.py", + "threadpool.py", "unique.py", ], srcs_version = "PY2AND3", deps = [ + ":contrib_op_loader", + ":gen_dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", "//tensorflow/python:dataset_ops_gen", @@ -118,10 +124,12 @@ py_library( "//tensorflow/python:logging_ops", "//tensorflow/python:math_ops", "//tensorflow/python:random_ops", + "//tensorflow/python:resource_variable_ops", "//tensorflow/python:tensor_shape", "//tensorflow/python:tensor_util", "//tensorflow/python:util", "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/ops:readers", "//tensorflow/python/data/util:convert", "//tensorflow/python/data/util:nest", "//tensorflow/python/data/util:sparse", @@ -130,36 +138,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 76c07b2c999e1424e8efe4af515fddee73922c9c..6463d75750d21f900e3e6d8fc23ff89a60bf8e90 100644 --- a/tensorflow/contrib/data/python/ops/batching.py +++ b/tensorflow/contrib/data/python/ops/batching.py @@ -374,8 +374,7 @@ class _MapAndBatchDataset(dataset_ops.MapDataset): @property def output_shapes(self): return nest.pack_sequence_as(self._output_shapes, [ - tensor_shape.vector(tensor_util.constant_value( - self._batch_size)).concatenate(s) + tensor_shape.vector(None).concatenate(s) for s in nest.flatten(self._output_shapes) ]) @@ -403,7 +402,7 @@ 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 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 deleted file mode 100644 index fafd231061a9108b2585f4fc9256b6f069b7c37a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/data/python/ops/dataset_ops.py +++ /dev/null @@ -1,691 +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. -# ============================================================================== -"""Python wrappers for Datasets and Iterators.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.data.python.ops import batching -from tensorflow.contrib.data.python.ops import enumerate_ops -from tensorflow.contrib.data.python.ops import error_ops -from tensorflow.contrib.data.python.ops import grouping -from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.util import nest -from tensorflow.python.ops import gen_dataset_ops -from tensorflow.python.ops import gen_io_ops -from tensorflow.python.util import deprecation - - -class Dataset(dataset_ops.Dataset): - """Represents a potentially large set of elements. - - A `Dataset` can be used to represent an input pipeline as a - collection of elements (nested structures of tensors) and a "logical - plan" of transformations that act on those elements. - """ - - def __init__(self, dataset): - super(Dataset, self).__init__() - self._dataset = dataset - - @deprecation.deprecated(None, "Use `ds._as_variant_tensor()`.") - def make_dataset_resource(self): - return self._as_variant_tensor() - - def _as_variant_tensor(self): - return self._dataset._as_variant_tensor() # pylint: disable=protected-access - - @property - def output_classes(self): - return self._dataset.output_classes - - @property - def output_shapes(self): - return self._dataset.output_shapes - - @property - def output_types(self): - return self._dataset.output_types - - @staticmethod - @deprecation.deprecated(None, "Use `tf.data.Dataset.from_tensors()`.") - def from_tensors(tensors): - """Creates a `Dataset` with a single element, comprising the given tensors. - - Args: - tensors: A nested structure of tensors. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.TensorDataset(tensors)) - - @staticmethod - @deprecation.deprecated(None, "Use `tf.data.Dataset.from_tensor_slices()`.") - def from_tensor_slices(tensors): - """Creates a `Dataset` whose elements are slices of the given tensors. - - Args: - tensors: A nested structure of tensors, each having the same size in the - 0th dimension. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.TensorSliceDataset(tensors)) - - @staticmethod - @deprecation.deprecated(None, - "Use `tf.data.Dataset.from_sparse_tensor_slices()`.") - def from_sparse_tensor_slices(sparse_tensor): - """Splits each rank-N `tf.SparseTensor` in this dataset row-wise. - - Args: - sparse_tensor: A `tf.SparseTensor`. - - Returns: - A `Dataset` of rank-(N-1) sparse tensors. - """ - return Dataset(dataset_ops.SparseTensorSliceDataset(sparse_tensor)) - - @staticmethod - @deprecation.deprecated(None, "Use `tf.data.Dataset.from_generator()`.") - def from_generator(generator, output_types, output_shapes=None): - """Creates a `Dataset` whose elements are generated by `generator`. - - The `generator` argument must be a callable object that returns - an object that support the `iter()` protocol (e.g. a generator function). - The elements generated by `generator` must be compatible with the given - `output_types` and (optional) `output_shapes` arguments. - - For example: - - ```python - import itertools - - def gen(): - for i in itertools.count(1): - yield (i, [1] * i) - - ds = Dataset.from_generator( - gen, (tf.int64, tf.int64), (tf.TensorShape([]), tf.TensorShape([None]))) - value = ds.make_one_shot_iterator().get_next() - - sess.run(value) # (1, array([1])) - sess.run(value) # (2, array([1, 1])) - ``` - - Args: - generator: A callable object that takes no arguments and returns an - object that supports the `iter()` protocol. - output_types: A nested structure of `tf.DType` objects corresponding to - each component of an element yielded by `generator`. - output_shapes: (Optional.) A nested structure of `tf.TensorShape` - objects corresponding to each component of an element yielded by - `generator`. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.Dataset.from_generator( - generator, output_types, output_shapes)) - - @staticmethod - @deprecation.deprecated(None, "Use `tf.data.Dataset.range()`.") - def range(*args): - """Creates a `Dataset` of a step-separated range of values. - - For example: - - ```python - Dataset.range(5) == [0, 1, 2, 3, 4] - Dataset.range(2, 5) == [2, 3, 4] - Dataset.range(1, 5, 2) == [1, 3] - Dataset.range(1, 5, -2) == [] - Dataset.range(5, 1) == [] - Dataset.range(5, 1, -2) == [5, 3] - ``` - - Args: - *args: follow same semantics as python's xrange. - len(args) == 1 -> start = 0, stop = args[0], step = 1 - len(args) == 2 -> start = args[0], stop = args[1], step = 1 - len(args) == 3 -> start = args[0], stop = args[1, stop = args[2] - - Returns: - A `RangeDataset`. - - Raises: - ValueError: if len(args) == 0. - """ - return Dataset(dataset_ops.RangeDataset(*args)) - - @staticmethod - @deprecation.deprecated(None, "Use `tf.data.Dataset.zip()`.") - def zip(datasets): - """Creates a `Dataset` by zipping together the given datasets. - - This method has similar semantics to the built-in `zip()` function - in Python, with the main difference being that the `datasets` - argument can be an arbitrary nested structure of `Dataset` objects. - For example: - - ```python - # NOTE: The following examples use `{ ... }` to represent the - # contents of a dataset. - a = { 1, 2, 3 } - b = { 4, 5, 6 } - c = { (7, 8), (9, 10), (11, 12) } - d = { 13, 14 } - - # The nested structure of the `datasets` argument determines the - # structure of elements in the resulting dataset. - Dataset.zip((a, b)) == { (1, 4), (2, 5), (3, 6) } - Dataset.zip((b, a)) == { (4, 1), (5, 2), (6, 3) } - - # The `datasets` argument may contain an arbitrary number of - # datasets. - Dataset.zip((a, b, c)) == { (1, 4, (7, 8)), - (2, 5, (9, 10)), - (3, 6, (11, 12)) } - - # The number of elements in the resulting dataset is the same as - # the size of the smallest dataset in `datasets`. - Dataset.zip((a, d)) == { (1, 13), (2, 14) } - ``` - - Args: - datasets: A nested structure of datasets. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.ZipDataset(datasets)) - - def concatenate(self, dataset): - """Creates a `Dataset` by concatenating given dataset with this dataset. - - ```python - # NOTE: The following examples use `{ ... }` to represent the - # contents of a dataset. - a = { 1, 2, 3 } - b = { 4, 5, 6, 7 } - - # Input dataset and dataset to be concatenated should have same - # nested structures and output types. - # c = { (8, 9), (10, 11), (12, 13) } - # d = { 14.0, 15.0, 16.0 } - # a.concatenate(c) and a.concatenate(d) would result in error. - - a.concatenate(b) == { 1, 2, 3, 4, 5, 6, 7 } - ``` - - Args: - dataset: `Dataset` to be concatenated. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.ConcatenateDataset(self._dataset, dataset)) - - def prefetch(self, buffer_size): - """Creates a `Dataset` that prefetches elements from this dataset. - - Args: - buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the - maximum number elements that will be buffered when prefetching. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.PrefetchDataset(self._dataset, buffer_size)) - - @staticmethod - @deprecation.deprecated(None, "Use `tf.data.Dataset.list_files()`.") - def list_files(file_pattern): - """A dataset of all files matching a pattern. - - Example: - If we had the following files on our filesystem: - - /path/to/dir/a.txt - - /path/to/dir/b.py - - /path/to/dir/c.py - If we pass "/path/to/dir/*.py" as the directory, the dataset would - produce: - - /path/to/dir/b.py - - /path/to/dir/c.py - - Args: - file_pattern: A string or scalar string `tf.Tensor`, representing - the filename pattern that will be matched. - - Returns: - A `Dataset` of strings corresponding to file names. - """ - return Dataset.from_tensor_slices(gen_io_ops.matching_files(file_pattern)) - - def repeat(self, count=None): - """Repeats this dataset `count` times. - - Args: - count: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the - number of times the elements of this dataset should be repeated. The - default behavior (if `count` is `None` or `-1`) is for the elements to - be repeated indefinitely. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.RepeatDataset(self._dataset, count)) - - @deprecation.deprecated( - None, "Use `ds.apply(tf.contrib.data.enumerate_dataset())`.") - def enumerate(self, start=0): - """Deprecated: Use `Dataset.apply(tf.contrib.data.enumerate_dataset(..)`.""" - - return self.apply(enumerate_ops.enumerate_dataset(start)) - - def shuffle(self, buffer_size, seed=None): - """Randomly shuffles the elements of this dataset. - - Args: - buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the - number of elements from this dataset from which the new - dataset will sample. - seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the - random seed that will be used to create the distribution. See - @{tf.set_random_seed} for behavior. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.ShuffleDataset(self._dataset, buffer_size, seed)) - - def cache(self, filename=""): - """Caches the elements in this dataset. - - Args: - filename: A `tf.string` scalar `tf.Tensor`, representing the name of a - directory on the filesystem to use for caching tensors in this Dataset. - If a filename is not provided, the dataset will be cached in memory. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.CacheDataset(self._dataset, filename)) - - def take(self, count): - """Creates a `Dataset` with at most `count` elements from this dataset. - - Args: - count: A `tf.int64` scalar `tf.Tensor`, representing the number of - elements of this dataset that should be taken to form the new dataset. - If `count` is -1, or if `count` is greater than the size of this - dataset, the new dataset will contain all elements of this dataset. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.TakeDataset(self._dataset, count)) - - def skip(self, count): - """Creates a `Dataset` that skips `count` elements from this dataset. - - Args: - count: A `tf.int64` scalar `tf.Tensor`, representing the number - of elements of this dataset that should be skipped to form the - new dataset. If `count` is greater than the size of this - dataset, the new dataset will contain no elements. If `count` - is -1, skips the entire dataset. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.SkipDataset(self._dataset, count)) - - def shard(self, num_shards, index): - """Creates a `Dataset` that includes only 1/`num_shards` of this dataset. - - This dataset operator is very useful when running distributed training, as - it allows each worker to read a unique subset. - - When reading a single input file, you can skip elements as follows: - - ```python - 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) - d = d.map(parser_fn, num_parallel_calls=FLAGS.num_map_threads) - ``` - - Important caveats: - - - Be sure to shard before you use any randomizing operator (such as - shuffle). - - Generally it is best if the shard operator is used early in the dataset - pipeline. For example, when reading from a set of TFRecord files, shard - before converting the dataset to input samples. This avoids reading every - file on every worker. The following is an example of an efficient - sharding strategy within a complete pipeline: - - ```python - 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.interleave(tf.data.TFRecordDataset, - cycle_length=FLAGS.num_readers, block_length=1) - d = d.map(parser_fn, num_parallel_calls=FLAGS.num_map_threads) - ``` - - Args: - num_shards: A `tf.int64` scalar `tf.Tensor`, representing the number of - shards operating in parallel. - index: A `tf.int64` scalar `tf.Tensor`, representing the worker index. - - Returns: - A `Dataset`. - - Raises: - ValueError: if `num_shards` or `index` are illegal values. Note: error - checking is done on a best-effort basis, and aren't guaranteed to be - caught upon dataset creation. (e.g. providing in a placeholder tensor - bypasses the early checking, and will instead result in an error during - a session.run call.) - """ - return Dataset(self._dataset.shard(num_shards, index)) - - @deprecation.deprecated( - None, "Use `ds.apply(tf.contrib.data.ignore_errors())`.") - def ignore_errors(self): - """Deprecated: Use `Dataset.apply(tf.contrib.data.ignore_errors())`.""" - - return self.apply(error_ops.ignore_errors()) - - def batch(self, batch_size): - """Combines consecutive elements of this dataset into batches. - - Args: - batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of - consecutive elements of this dataset to combine in a single batch. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.BatchDataset(self._dataset, batch_size)) - - def padded_batch(self, batch_size, padded_shapes, padding_values=None): - """Combines consecutive elements of this dataset into padded batches. - - Like `Dataset.dense_to_sparse_batch()`, this method combines - multiple consecutive elements of this dataset, which might have - different shapes, into a single element. The tensors in the - resulting element have an additional outer dimension, and are - padded to the respective shape in `padded_shapes`. - - Args: - batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of - consecutive elements of this dataset to combine in a single batch. - padded_shapes: A nested structure of `tf.TensorShape` or - `tf.int64` vector tensor-like objects representing the shape - to which the respective component of each input element should - be padded prior to batching. Any unknown dimensions - (e.g. `tf.Dimension(None)` in a `tf.TensorShape` or `-1` in a - tensor-like object) will be padded to the maximum size of that - dimension in each batch. - padding_values: (Optional.) A nested structure of scalar-shaped - `tf.Tensor`, representing the padding values to use for the - respective components. Defaults are `0` for numeric types and - the empty string for string types. - - Returns: - A `Dataset`. - """ - return Dataset( - dataset_ops.PaddedBatchDataset(self._dataset, batch_size, padded_shapes, - padding_values)) - - @deprecation.deprecated( - None, "Use `ds.apply(tf.contrib.data.dense_to_sparse_batch())`.") - def dense_to_sparse_batch(self, batch_size, row_shape): - """Use: `Dataset.apply(tf.contrib.data.dense_to_sparse_batch(...))`.""" - - return self.apply(batching.dense_to_sparse_batch(batch_size, row_shape)) - - @deprecation.deprecated( - None, "Use `ds.apply(tf.contrib.data.group_by_window())`.") - def group_by_window(self, key_func, reduce_func, window_size): - """Deprecated: Use `Dataset.apply(tf.contrib.data.group_by_window(...))`.""" - - return self.apply( - grouping.group_by_window(key_func, reduce_func, window_size)) - - @deprecation.deprecated_args( - None, - "Replace `num_threads=T` with `num_parallel_calls=T`. Replace " - "`output_buffer_size=N` with `ds.prefetch(N)` on the returned dataset.", - "num_threads", "output_buffer_size") - def map(self, - map_func, - num_threads=None, - output_buffer_size=None, - num_parallel_calls=None): - """Maps `map_func` across this datset. - - Args: - map_func: A function mapping a nested structure of tensors (having - shapes and types defined by `self.output_shapes` and - `self.output_types`) to another nested structure of tensors. - num_threads: (Optional.) Deprecated, use `num_parallel_calls` instead. - output_buffer_size: (Optional.) A `tf.int64` scalar `tf.Tensor`, - representing the maximum number of processed elements that will be - buffered. - num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`, - representing the number elements to process in parallel. If not - specified, elements will be processed sequentially. - - Returns: - A `Dataset`. - """ - if num_threads is None and num_parallel_calls is None: - ret = Dataset(dataset_ops.MapDataset(self._dataset, map_func)) - else: - if num_threads is None: - ret = Dataset( - dataset_ops.ParallelMapDataset(self._dataset, map_func, - num_parallel_calls)) - else: - ret = Dataset( - dataset_ops.ParallelMapDataset(self._dataset, map_func, - num_threads)) - if output_buffer_size is not None: - ret = ret.prefetch(output_buffer_size) - return ret - - def flat_map(self, map_func): - """Maps `map_func` across this dataset and flattens the result. - - Args: - map_func: A function mapping a nested structure of tensors (having shapes - and types defined by `self.output_shapes` and `self.output_types`) to a - `Dataset`. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.FlatMapDataset(self._dataset, map_func)) - - def interleave(self, map_func, cycle_length, block_length=1): - """Maps `map_func` across this dataset, and interleaves the results. - - For example, you can use `Dataset.interleave()` to process many input files - concurrently: - - ```python - # Preprocess 4 files concurrently, and interleave blocks of 16 records from - # each file. - filenames = ["/var/data/file1.txt", "/var/data/file2.txt", ...] - dataset = (Dataset.from_tensor_slices(filenames) - .interleave(lambda x: - TextLineDataset(x).map(parse_fn, num_parallel_calls=1), - cycle_length=4, block_length=16)) - ``` - - The `cycle_length` and `block_length` arguments control the order in which - 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.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 - consuming the next input element each time it reaches the end of an - iterator. - - For example: - - ```python - # NOTE: The following examples use `{ ... }` to represent the - # contents of a dataset. - a = { 1, 2, 3, 4, 5 } - - # NOTE: New lines indicate "block" boundaries. - a.interleave(lambda x: Dataset.from_tensors(x).repeat(6), - cycle_length=2, block_length=4) == { - 1, 1, 1, 1, - 2, 2, 2, 2, - 1, 1, - 2, 2, - 3, 3, 3, 3, - 4, 4, 4, 4, - 3, 3, - 4, 4, - 5, 5, 5, 5, - 5, 5, - } - ``` - - NOTE: The order of elements yielded by this transformation is - deterministic, as long as `map_func` is a pure function. If - `map_func` contains any stateful operations, the order in which - that state is accessed is undefined. - - Args: - map_func: A function mapping a nested structure of tensors (having shapes - and types defined by `self.output_shapes` and `self.output_types`) to a - `Dataset`. - cycle_length: The number of elements from this dataset that will be - processed concurrently. - block_length: The number of consecutive elements to produce from each - input element before cycling to another input element. - - Returns: - A `Dataset`. - """ - return Dataset( - dataset_ops.InterleaveDataset(self._dataset, map_func, cycle_length, - block_length)) - - @deprecation.deprecated(None, "Use `ds.apply(tf.contrib.data.unbatch())`.") - def unbatch(self): - """Deprecated: Use `Dataset.apply(tf.contrib.data.unbatch()`.""" - - return self.apply(batching.unbatch()) - - def filter(self, predicate): - """Filters this dataset according to `predicate`. - - Args: - predicate: 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.bool` tensor. - - Returns: - A `Dataset`. - """ - return Dataset(dataset_ops.FilterDataset(self._dataset, predicate)) - - def apply(self, transformation_func): - """Apply a transformation function to this dataset. - - `apply` enables chaining of custom `Dataset` transformations, which are - represented as functions that take one `Dataset` argument and return a - transformed `Dataset`. - - For example: - - ``` - dataset = (dataset.map(lambda x: x ** 2) - .(group_by_window(key_func, reduce_func, window_size)) - .map(lambda x: x ** 3)) - ``` - - Args: - transformation_func: A function that takes one `Dataset` argument and - returns a `Dataset`. - - Returns: - The `Dataset` returned by applying `transformation_func` to this dataset. - """ - dataset = transformation_func(self) - if not isinstance(dataset, dataset_ops.Dataset): - raise TypeError("`transformation_func` must return a Dataset.") - return Dataset(dataset) - - -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/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..3a07df572748e464284f580d67e3a664e71acdfe --- /dev/null +++ b/tensorflow/contrib/data/python/ops/get_single_element.py @@ -0,0 +1,73 @@ +# 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.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.") + + nested_ret = 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(sparse.as_dense_types( + dataset.output_types, dataset.output_classes)), + output_shapes=nest.flatten(sparse.as_dense_shapes( + dataset.output_shapes, dataset.output_classes)))) + return sparse.deserialize_sparse_tensors( + nested_ret, dataset.output_types, dataset.output_shapes, + dataset.output_classes) diff --git a/tensorflow/contrib/data/python/ops/grouping.py b/tensorflow/contrib/data/python/ops/grouping.py index ef91c56726e969053fdad667dda3e89430045652..a19be222545ef0242502ec07badbdae5c7634a0c 100644 --- a/tensorflow/contrib/data/python/ops/grouping.py +++ b/tensorflow/contrib/data/python/ops/grouping.py @@ -17,13 +17,20 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + 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 constant_op 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.ops import array_ops +from tensorflow.python.ops import check_ops from tensorflow.python.ops import gen_dataset_ops +from tensorflow.python.ops import math_ops def group_by_window(key_func, @@ -45,7 +52,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 @@ -85,6 +92,114 @@ def group_by_window(key_func, return _apply_fn +def bucket_by_sequence_length(element_length_func, + bucket_boundaries, + bucket_batch_sizes, + padded_shapes=None, + padding_values=None, + pad_to_bucket_boundary=False): + """A transformation that buckets elements in a `Dataset` by length. + + Elements of the `Dataset` are grouped together by length and then are padded + and batched. + + This is useful for sequence tasks in which the elements have variable length. + Grouping together elements that have similar lengths reduces the total + fraction of padding in a batch which increases training step efficiency. + + Args: + element_length_func: function from element in `Dataset` to `tf.int64`, + determines the length of the element, which will determine the bucket it + goes into. + bucket_boundaries: `list`, upper length boundaries of the buckets. + bucket_batch_sizes: `list`, batch size per bucket. Length should be + `len(bucket_boundaries) + 1`. + padded_shapes: Nested structure of `tf.TensorShape` to pass to + @{tf.data.Dataset.padded_batch}. If not provided, will use + `dataset.output_shapes`, which will result in variable length dimensions + being padded out to the maximum length in each batch. + padding_values: Values to pad with, passed to + @{tf.data.Dataset.padded_batch}. Defaults to padding with 0. + pad_to_bucket_boundary: bool, if `False`, will pad dimensions with unknown + size to maximum length in batch. If `True`, will pad dimensions with + unknown size to bucket boundary, and caller must ensure that the source + `Dataset` does not contain any elements with length longer than + `max(bucket_boundaries)`. + + Returns: + A `Dataset` transformation function, which can be passed to + @{tf.data.Dataset.apply}. + + Raises: + ValueError: if `len(bucket_batch_sizes) != len(bucket_boundaries) + 1`. + """ + with ops.name_scope("bucket_by_seq_length"): + if len(bucket_batch_sizes) != (len(bucket_boundaries) + 1): + raise ValueError( + "len(bucket_batch_sizes) must equal len(bucket_boundaries) + 1") + + batch_sizes = constant_op.constant(bucket_batch_sizes, dtype=dtypes.int64) + + def element_to_bucket_id(element): + """Return int64 id of the length bucket for this element.""" + seq_length = element_length_func(element) + + boundaries = list(bucket_boundaries) + buckets_min = [np.iinfo(np.int32).min] + boundaries + buckets_max = boundaries + [np.iinfo(np.int32).max] + conditions_c = math_ops.logical_and( + math_ops.less_equal(buckets_min, seq_length), + math_ops.less(seq_length, buckets_max)) + bucket_id = math_ops.reduce_min(array_ops.where(conditions_c)) + + return bucket_id + + def window_size_fn(bucket_id): + # The window size is set to the batch size for this bucket + window_size = batch_sizes[bucket_id] + return window_size + + def make_padded_shapes(shapes, none_filler=None): + padded = [] + for shape in nest.flatten(shapes): + shape = tensor_shape.TensorShape(shape) + shape = [ + none_filler if d.value is None else d + for d in shape + ] + padded.append(shape) + return nest.pack_sequence_as(shapes, padded) + + def batching_fn(bucket_id, grouped_dataset): + """Batch elements in dataset.""" + batch_size = batch_sizes[bucket_id] + none_filler = None + if pad_to_bucket_boundary: + err_msg = ("When pad_to_bucket_boundary=True, elements must have " + "length <= max(bucket_boundaries).") + check = check_ops.assert_less( + bucket_id, + constant_op.constant(len(bucket_batch_sizes) - 1, + dtype=dtypes.int64), + message=err_msg) + with ops.control_dependencies([check]): + boundaries = constant_op.constant(bucket_boundaries, + dtype=dtypes.int64) + bucket_boundary = boundaries[bucket_id] + none_filler = bucket_boundary + shapes = make_padded_shapes( + padded_shapes or grouped_dataset.output_shapes, + none_filler=none_filler) + return grouped_dataset.padded_batch(batch_size, shapes, padding_values) + + def _apply_fn(dataset): + return dataset.apply( + group_by_window(element_to_bucket_id, batching_fn, + window_size_func=window_size_fn)) + + return _apply_fn + + class _VariantDataset(dataset_ops.Dataset): """A Dataset wrapper for a tf.variant-typed function argument.""" diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py index 3124ca1d1540e12d949dded88ce1c66181be3595..91f19da02d4a479820782822475d9121125fc38e 100644 --- a/tensorflow/contrib/data/python/ops/interleave_ops.py +++ b/tensorflow/contrib/data/python/ops/interleave_ops.py @@ -17,101 +17,10 @@ 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 convert -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 function -from tensorflow.python.framework import ops -from tensorflow.python.ops import gen_dataset_ops +from tensorflow.python.data.ops import readers from tensorflow.python.util import deprecation -class ParallelInterleaveDataset(dataset_ops.Dataset): - """A `Dataset` that maps a function over its input and flattens the result.""" - - def __init__(self, input_dataset, map_func, cycle_length, block_length, - sloppy, buffer_output_elements, prefetch_input_elements): - """See `tf.contrib.data.parallel_interleave()` for details.""" - super(ParallelInterleaveDataset, self).__init__() - self._input_dataset = input_dataset - - @function.Defun(*nest.flatten( - sparse.as_dense_types(input_dataset.output_types, - input_dataset.output_classes))) - def tf_map_func(*args): - """A wrapper for Defun that facilitates shape inference.""" - # Pass in shape information from the input_dataset. - dense_shapes = sparse.as_dense_shapes(input_dataset.output_shapes, - input_dataset.output_classes) - for arg, shape in zip(args, nest.flatten(dense_shapes)): - arg.set_shape(shape) - - nested_args = nest.pack_sequence_as(input_dataset.output_types, args) - nested_args = sparse.deserialize_sparse_tensors( - nested_args, input_dataset.output_types, input_dataset.output_shapes, - input_dataset.output_classes) - if dataset_ops._should_unpack_args(nested_args): # pylint: disable=protected-access - dataset = map_func(*nested_args) - else: - dataset = map_func(nested_args) - - if not isinstance(dataset, dataset_ops.Dataset): - raise TypeError("`map_func` must return a `Dataset` object.") - - self._output_classes = dataset.output_classes - self._output_types = dataset.output_types - self._output_shapes = dataset.output_shapes - - return dataset._as_variant_tensor() # pylint: disable=protected-access - - self._map_func = tf_map_func - self._map_func.add_to_graph(ops.get_default_graph()) - - self._cycle_length = ops.convert_to_tensor( - cycle_length, dtype=dtypes.int64, name="cycle_length") - self._block_length = ops.convert_to_tensor( - block_length, dtype=dtypes.int64, name="block_length") - self._sloppy = ops.convert_to_tensor( - sloppy, dtype=dtypes.bool, name="sloppy") - self._buffer_output_elements = convert.optional_param_to_tensor( - "buffer_output_elements", - buffer_output_elements, - argument_default=2 * block_length) - self._prefetch_input_elements = convert.optional_param_to_tensor( - "prefetch_input_elements", - prefetch_input_elements, - argument_default=2 * cycle_length) - - def _as_variant_tensor(self): - return gen_dataset_ops.parallel_interleave_dataset( - self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access - self._map_func.captured_inputs, - self._cycle_length, - self._block_length, - self._sloppy, - self._buffer_output_elements, - self._prefetch_input_elements, - f=self._map_func, - 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_classes(self): - return self._output_classes - - @property - def output_shapes(self): - return self._output_shapes - - @property - def output_types(self): - return self._output_types - - def parallel_interleave(map_func, cycle_length, block_length=1, @@ -162,7 +71,7 @@ def parallel_interleave(map_func, @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): - return ParallelInterleaveDataset( + return readers.ParallelInterleaveDataset( dataset, map_func, cycle_length, block_length, sloppy, buffer_output_elements, prefetch_input_elements) @@ -221,7 +130,7 @@ def sloppy_interleave(map_func, cycle_length, block_length=1): @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): - return ParallelInterleaveDataset( + return readers.ParallelInterleaveDataset( dataset, map_func, cycle_length, diff --git a/tensorflow/contrib/data/python/ops/prefetching_ops.py b/tensorflow/contrib/data/python/ops/prefetching_ops.py index cfe8012b5657995b78d701528ea35cbb3748adb9..7059b358f349e0ec847e85c37652012d48ed910a 100644 --- a/tensorflow/contrib/data/python/ops/prefetching_ops.py +++ b/tensorflow/contrib/data/python/ops/prefetching_ops.py @@ -17,25 +17,23 @@ 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__ # method and provides a get_next() that calls the prefetch op. def function_buffering_resource(string_arg, target_device, - shared_name, f, buffer_size, thread_pool_size=1, container="", + shared_name=None, name=None): - return gen_prefetching_ops.function_buffering_resource( + if shared_name is None: + shared_name = "" + return gen_dataset_ops.function_buffering_resource( string_arg=string_arg, target_device=target_device, shared_name=shared_name, @@ -49,7 +47,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/random_ops.py b/tensorflow/contrib/data/python/ops/random_ops.py index 7d727165feabb101549567f28a2dfa07083de244..28ef5e50f39dd7d1b6f124e58e068fc968ddd6dc 100644 --- a/tensorflow/contrib/data/python/ops/random_ops.py +++ b/tensorflow/contrib/data/python/ops/random_ops.py @@ -19,11 +19,10 @@ 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 random_seed from tensorflow.python.data.util import sparse -from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import random_seed from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import gen_dataset_ops @@ -34,16 +33,7 @@ class RandomDataset(dataset_ops.Dataset): def __init__(self, seed=None): """A `Dataset` of pseudorandom values.""" super(RandomDataset, self).__init__() - seed, seed2 = random_seed.get_seed(seed) - if seed is None: - self._seed = constant_op.constant(0, dtype=dtypes.int64, name="seed") - else: - self._seed = ops.convert_to_tensor(seed, dtype=dtypes.int64, name="seed") - if seed2 is None: - self._seed2 = constant_op.constant(0, dtype=dtypes.int64, name="seed2") - else: - self._seed2 = ops.convert_to_tensor( - seed2, dtype=dtypes.int64, name="seed2") + self._seed, self._seed2 = random_seed.get_seed(seed) def _as_variant_tensor(self): return gen_dataset_ops.random_dataset( diff --git a/tensorflow/contrib/data/python/ops/readers.py b/tensorflow/contrib/data/python/ops/readers.py index 347e5edc7b0d479dfa260e8cec500ffaaba375be..f70f9c881df168564cbf2431bbc2ebdf7e7f7ded 100644 --- a/tensorflow/contrib/data/python/ops/readers.py +++ b/tensorflow/contrib/data/python/ops/readers.py @@ -17,10 +17,12 @@ 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.contrib.data.python.ops import interleave_ops +from tensorflow.contrib.data.python.ops import shuffle_ops from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.ops import readers +from tensorflow.python.data.ops import readers as core_readers from tensorflow.python.data.util import nest +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 @@ -29,78 +31,275 @@ from tensorflow.python.ops import parsing_ops from tensorflow.python.platform import gfile from tensorflow.python.util import deprecation +_ACCEPTABLE_CSV_TYPES = (dtypes.float32, dtypes.float64, dtypes.int32, + dtypes.int64, dtypes.string) + + +def make_csv_dataset( + file_pattern, + batch_size, + column_keys, + column_defaults, + label_key=None, + field_delim=",", + use_quote_delim=True, + skip=0, + filter_fn=None, + num_epochs=None, + shuffle=True, + shuffle_buffer_size=10000, + shuffle_seed=None, + prefetch_buffer_size=1, +): + """Reads CSV files into a dataset. + + Reads CSV files into a dataset, where each element is a (features, labels) + tuple that corresponds to a batch of CSV rows. The features dictionary + maps feature column names to `Tensor`s containing the corresponding + feature data, and labels is a `Tensor` containing the batch's label data. -class TextLineDataset(contrib_dataset_ops.Dataset): - """A `Dataset` comprising lines from one or more text files.""" + Args: + file_pattern: List of files or patterns of file paths containing CSV + records. See @{tf.gfile.Glob} for pattern rules. + batch_size: An int representing the number of consecutive elements of this + dataset to combine in a single batch. + column_keys: A list of strings that corresponds to the CSV columns, in + order. One per column of the input record. + column_defaults: A list of default values for the CSV fields. One item per + column of the input record. Each item in the list is either one of the + following dtypes: float32, float64, int32, int64, or string, or a + `Tensor` with one of the aforementioned types. One item per column of + the input record, with either scalar default value for that column if it + is required, or, if the column is required, an empty `Tensor` or a dtype. + label_key: A optional string corresponding to the label column. If provided, + the data for this column is returned as a separate `Tensor` from the + features dictionary, so that the dataset complies with the format expected + by a `tf.Estimator.train` or `tf.Estimator.evaluate` input function. + field_delim: An optional `string`. Defaults to `","`. Char delimiter to + separate fields in a record. + use_quote_delim: An optional bool. Defaults to `True`. If false, treats + double quotation marks as regular characters inside of the string fields. + skip: An integer that corresponds to the number of lines to skip at the + head of each CSV file. Defaults to 0. + filter_fn: A callable function that takes in a CSV string and returns a + boolean that corresponds to whether the record should be included. If + None, does not filter records. + num_epochs: An int specifying the number of times this dataset is repeated. + If None, cycles through the dataset forever. + shuffle: A bool that indicates whether the input should be shuffled. + shuffle_buffer_size: Buffer size to use for shuffling. A large buffer size + ensures better shuffling, but would increase memory usage and startup + time. + shuffle_seed: Randomization seed to use for shuffling. + prefetch_buffer_size: An int specifying the number of feature batches to + prefetch for performance improvement. Recommended value is the number of + batches consumed per training step. - @deprecation.deprecated(None, "Use `tf.data.TextLineDataset`.") - def __init__(self, filenames, compression_type=None, buffer_size=None): - """Creates a `TextLineDataset`. + Returns: + A dataset, where each element is a (features, labels) tuple that corresponds + to a batch of `batch_size` CSV rows. The features dictionary maps feature + column names to `Tensor`s containing the corresponding column data, and + labels is a `Tensor` containing the column data for the label column + specified by `label_key`. + """ + filenames = _get_file_names(file_pattern, False) + column_defaults = [ + constant_op.constant([], dtype=x) if x in _ACCEPTABLE_CSV_TYPES else x + for x in column_defaults + ] + + dataset = dataset_ops.Dataset.from_tensor_slices(filenames) + if label_key is not None: + assert label_key in column_keys + + def filename_to_dataset(filename): + ds = core_readers.TextLineDataset(filename) + if skip > 0: + ds = ds.skip(skip) + if filter_fn is not None: + ds = ds.filter(filter_fn) + return ds + + def decode_csv(line): + """Decodes csv line into features. 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. + line: String tensor corresponding to one csv record. + Returns: + A dictionary of feature names to values for that particular record. If + label_key is provided, extracts the label feature to be returned as the + second element of the tuple. """ - dataset = readers.TextLineDataset(filenames, compression_type, - buffer_size) - super(TextLineDataset, self).__init__(dataset) + columns = parsing_ops.decode_csv( + line, + column_defaults, + field_delim=field_delim, + use_quote_delim=use_quote_delim) + features = dict(zip(column_keys, columns)) + if label_key is not None: + label = features.pop(label_key) + return features, label + return features + + # TODO(rachelim): interleave records from files for better shuffling + dataset = dataset.flat_map(filename_to_dataset) + # TODO(rachelim): use fused shuffle_and_repeat for perf + if shuffle: + dataset = dataset.shuffle(shuffle_buffer_size, shuffle_seed) + if num_epochs != 1: + dataset = dataset.repeat(num_epochs) + dataset = dataset.batch(batch_size) + dataset = dataset.map(decode_csv) + dataset = dataset.prefetch(prefetch_buffer_size) + return dataset + + +def make_batched_features_dataset(file_pattern, + batch_size, + features, + reader=core_readers.TFRecordDataset, + reader_args=None, + num_epochs=None, + shuffle=True, + shuffle_buffer_size=10000, + shuffle_seed=None, + prefetch_buffer_size=1, + reader_num_threads=1, + parser_num_threads=2, + sloppy_ordering=False): + """Returns a `Dataset` of feature dictionaries from `Example` protos. -class TFRecordDataset(contrib_dataset_ops.Dataset): - """A `Dataset` comprising records from one or more TFRecord files.""" + Example: - @deprecation.deprecated(None, "Use `tf.data.TFRecordDataset`.") - def __init__(self, filenames, compression_type=None, buffer_size=None): - """Creates a `TFRecordDataset`. + ``` + serialized_examples = [ + features { + feature { key: "age" value { int64_list { value: [ 0 ] } } } + feature { key: "gender" value { bytes_list { value: [ "f" ] } } } + feature { key: "kws" value { bytes_list { value: [ "code", "art" ] } } } + }, + features { + feature { key: "age" value { int64_list { value: [] } } } + feature { key: "gender" value { bytes_list { value: [ "f" ] } } } + feature { key: "kws" value { bytes_list { value: [ "sports" ] } } } + } + ] + ``` - 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) + We can use arguments: + + ``` + features: { + "age": FixedLenFeature([], dtype=tf.int64, default_value=-1), + "gender": FixedLenFeature([], dtype=tf.string), + "kws": VarLenFeature(dtype=tf.string), + } + ``` + + And the expected output is: + + ```python + { + "age": [[0], [-1]], + "gender": [["f"], ["f"]], + "kws": SparseTensor( + indices=[[0, 0], [0, 1], [1, 0]], + values=["code", "art", "sports"] + dense_shape=[2, 2]), + } + ``` + Args: + file_pattern: List of files or patterns of file paths containing + `Example` records. See `tf.gfile.Glob` for pattern rules. + batch_size: An int representing the number of consecutive elements of this + dataset to combine in a single batch. + features: A `dict` mapping feature keys to `FixedLenFeature` or + `VarLenFeature` values. See `tf.parse_example`. + reader: A function or class that can be + called with a `filenames` tensor and (optional) `reader_args` and returns + a `Dataset` of `Example` tensors. Defaults to `tf.data.TFRecordDataset`. + reader_args: Additional arguments to pass to the reader class. + num_epochs: Integer specifying the number of times to read through the + dataset. If None, cycles through the dataset forever. Defaults to `None`. + shuffle: A boolean, indicates whether the input should be shuffled. Defaults + to `True`. + shuffle_buffer_size: Buffer size of the ShuffleDataset. A large capacity + ensures better shuffling but would increase memory usage and startup time. + shuffle_seed: Randomization seed to use for shuffling. + prefetch_buffer_size: Number of feature batches to prefetch in order to + improve performance. Recommended value is the number of batches consumed + per training step (default is 1). + reader_num_threads: Number of threads used to read `Example` records. If >1, + the results will be interleaved. + parser_num_threads: Number of threads to use for parsing `Example` tensors + into a dictionary of `Feature` tensors. + sloppy_ordering: If `True`, reading performance will be improved at + the cost of non-deterministic ordering. If `False`, the order of elements + produced is deterministic prior to shuffling (elements are still + randomized if `shuffle=True`. Note that if the seed is set, then order + of elements after shuffling is deterministic). Defaults to `False`. -class FixedLengthRecordDataset(contrib_dataset_ops.Dataset): - """A `Dataset` of fixed-length records from one or more binary files.""" + Returns: + A dataset of `dict` elements. Each `dict` maps feature keys to + `Tensor` or `SparseTensor` objects. + """ + # Create dataset of all matching filenames + if shuffle: + dataset = dataset_ops.Dataset.list_files(file_pattern, shuffle=True) + else: + # TODO(b/73959787): Use Dataset.list_files() once ordering is deterministic. + filenames = _get_file_names(file_pattern, shuffle) + dataset = dataset_ops.Dataset.from_tensor_slices(filenames) + + # Read `Example` records from files as tensor objects. + if reader_args is None: + reader_args = [] + + # Read files sequentially (if reader_num_threads=1) or in parallel + dataset = dataset.apply( + interleave_ops.parallel_interleave( + lambda filename: reader(filename, *reader_args), + cycle_length=reader_num_threads, + sloppy=sloppy_ordering)) + + # Extract values if the `Example` tensors are stored as key-value tuples. + if dataset.output_types == (dtypes.string, dtypes.string): + dataset = dataset.map(lambda _, v: v) - @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`. + # Apply dataset repeat and shuffle transformations. + repeat_dataset = (num_epochs != 1) + if repeat_dataset and shuffle: + # Used fused shuffle_and_repeat operation for better performance + dataset = dataset.apply( + shuffle_ops.shuffle_and_repeat(shuffle_buffer_size, num_epochs, + shuffle_seed)) + elif repeat_dataset: + dataset = dataset.repeat(num_epochs) + elif shuffle: + dataset = dataset.shuffle(shuffle_buffer_size, shuffle_seed) - 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) + dataset = dataset.batch(batch_size) + + # Parse `Example` tensors to a dictionary of `Feature` tensors. + dataset = dataset.map( + lambda x: parsing_ops.parse_example(x, features), + num_parallel_calls=parser_num_threads) + # TODO(rachelim): Add an optional label_key argument for extracting the label + # from the features dictionary, to comply with the type expected by the + # input_fn to a `tf.Estimator.train` or `tf.Estimator.evaluate` function. + dataset = dataset.prefetch(prefetch_buffer_size) + return dataset + +@deprecation.deprecated(None, + "Use `tf.contrib.data.make_batched_features_dataset`") def read_batch_features(file_pattern, batch_size, features, - reader, + reader=core_readers.TFRecordDataset, reader_args=None, randomize_input=True, num_epochs=None, @@ -154,43 +353,38 @@ def read_batch_features(file_pattern, dataset to combine in a single batch. features: A `dict` mapping feature keys to `FixedLenFeature` or `VarLenFeature` values. See `tf.parse_example`. - reader: A function or class that can be called with a `filenames` tensor - and (optional) `reader_args` and returns a `Dataset` of Examples. + reader: A function or class that can be + called with a `filenames` tensor and (optional) `reader_args` and returns + a `Dataset` of `Example` tensors. Defaults to `tf.data.TFRecordDataset`. reader_args: Additional arguments to pass to the reader class. randomize_input: Whether the input should be randomized. num_epochs: Integer specifying the number of times to read through the dataset. If None, cycles through the dataset forever. - capacity: Capacity of the ShuffleDataset. A large capacity ensures better + capacity: Buffer size of the ShuffleDataset. A large capacity ensures better shuffling but would increase memory usage and startup time. - Returns: A dict from keys in features to `Tensor` or `SparseTensor` objects. """ - filenames = _get_file_names(file_pattern, randomize_input) - if reader_args: - dataset = reader(filenames, *reader_args) - else: - dataset = reader(filenames) - if dataset.output_types == (dtypes.string, dtypes.string): - dataset = dataset.map(lambda _, v: v) - if num_epochs != 1: - dataset = dataset.repeat(num_epochs) - if randomize_input: - dataset = dataset.shuffle(capacity) - dataset = dataset.batch(batch_size) - dataset = dataset.map(lambda x: parsing_ops.parse_example(x, features)) - dataset = dataset.prefetch(1) + dataset = make_batched_features_dataset( + file_pattern, + batch_size, + features, + reader=reader, + reader_args=reader_args, + shuffle=randomize_input, + num_epochs=num_epochs, + shuffle_buffer_size=capacity) iterator = dataset.make_one_shot_iterator() outputs = iterator.get_next() return outputs -def _get_file_names(file_pattern, randomize_input): +def _get_file_names(file_pattern, shuffle): """Parse list of file names from pattern, optionally shuffled. Args: file_pattern: File glob pattern, or list of glob patterns. - randomize_input: Whether to shuffle the order of file names. + shuffle: Whether to shuffle the order of file names. Returns: List of file names matching `file_pattern`. @@ -211,19 +405,12 @@ def _get_file_names(file_pattern, randomize_input): raise ValueError("No files match %s." % file_pattern) # Sort files so it will be deterministic for unit tests. - if not randomize_input: + if not shuffle: file_names = sorted(file_names) 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 +442,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/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py index c99cdd3ff23db5340659d9f8e240b66bca4ba181..44de42e94d3d43156fe48cfe028587a7dc15c5e5 100644 --- a/tensorflow/contrib/data/python/ops/resampling.py +++ b/tensorflow/contrib/data/python/ops/resampling.py @@ -54,7 +54,7 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None): def _apply_fn(dataset): """Function from `Dataset` to `Dataset` that applies the transformation.""" dist_estimation_batch_size = 32 - target_dist_t = ops.convert_to_tensor(target_dist, name="initial_dist") + target_dist_t = ops.convert_to_tensor(target_dist, name="target_dist") class_values_ds = dataset.map(class_func) if initial_dist is not None: initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist") @@ -151,7 +151,7 @@ def _calculate_acceptance_probs(initial_probs, target_probs): ``` - A solution for a_i in terms of the other variabes is the following: + A solution for a_i in terms of the other variables is the following: ```a_i = (t_i / p_i) / max_i[t_i / p_i]``` """ # Add tiny to initial_probs to avoid divide by zero. diff --git a/tensorflow/contrib/data/python/ops/shuffle_ops.py b/tensorflow/contrib/data/python/ops/shuffle_ops.py index 99bb79bc06a421f811869ca9169aaa11deaca2f3..f35795abd38000b13cec0f08596e2ff66e86286c 100644 --- a/tensorflow/contrib/data/python/ops/shuffle_ops.py +++ b/tensorflow/contrib/data/python/ops/shuffle_ops.py @@ -19,11 +19,11 @@ 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 random_seed from tensorflow.python.data.util import sparse from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import random_seed from tensorflow.python.ops import gen_dataset_ops @@ -45,17 +45,7 @@ class _ShuffleAndRepeatDataset(dataset_ops.Dataset): else: self._count = ops.convert_to_tensor( count, dtype=dtypes.int64, name="count") - - seed, seed2 = random_seed.get_seed(seed) - if seed is None: - self._seed = constant_op.constant(0, dtype=dtypes.int64, name="seed") - else: - self._seed = ops.convert_to_tensor(seed, dtype=dtypes.int64, name="seed") - if seed2 is None: - self._seed2 = constant_op.constant(0, dtype=dtypes.int64, name="seed2") - else: - self._seed2 = ops.convert_to_tensor( - seed2, dtype=dtypes.int64, name="seed2") + self._seed, self._seed2 = random_seed.get_seed(seed) def _as_variant_tensor(self): # pylint: disable=protected-access diff --git a/tensorflow/contrib/data/python/ops/sliding.py b/tensorflow/contrib/data/python/ops/sliding.py new file mode 100644 index 0000000000000000000000000000000000000000..19cc3cb89fc5c494f79ce1d25ed57c92099c8bd2 --- /dev/null +++ b/tensorflow/contrib/data/python/ops/sliding.py @@ -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. +# ============================================================================== +"""Sliding dataset transformations.""" +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.ops import gen_dataset_ops + + +class _SlideDataset(dataset_ops.Dataset): + """A `Dataset` that passes a sliding window over its input.""" + + def __init__(self, input_dataset, window_size, stride=1): + """See `sliding_window_batch` for details.""" + super(_SlideDataset, self).__init__() + self._input_dataset = input_dataset + self._window_size = ops.convert_to_tensor( + window_size, dtype=dtypes.int64, name="window_size") + self._stride = ops.convert_to_tensor( + stride, dtype=dtypes.int64, name="stride") + + def _as_variant_tensor(self): + return gen_dataset_ops.slide_dataset( + self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access + window_size=self._window_size, + stride=self._stride, + output_shapes=nest.flatten( + sparse.as_dense_shapes(self.output_shapes, self.output_classes)), + output_types=nest.flatten( + sparse.as_dense_types(self.output_types, self.output_classes))) + + @property + def output_classes(self): + return self._input_dataset.output_classes + + @property + def output_shapes(self): + input_shapes = self._input_dataset.output_shapes + return nest.pack_sequence_as(input_shapes, [ + tensor_shape.vector(None).concatenate(s) + for s in nest.flatten(self._input_dataset.output_shapes) + ]) + + @property + def output_types(self): + return self._input_dataset.output_types + + +def sliding_window_batch(window_size, stride=1): + """A sliding window with size of `window_size` and step of `stride`. + + This transformation passes a sliding window over this dataset. The + window size is `window_size` and step size is `stride`. If the left + elements cannot fill up the sliding window, this transformation will + drop the final smaller element. For example: + + ```python + # NOTE: The following examples use `{ ... }` to represent the + # contents of a dataset. + a = { [1], [2], [3], [4], [5], [6] } + + a.apply(tf.contrib.data.sliding_window_batch(window_size=3, stride=2)) == + { + [[1], [2], [3]], + [[3], [4], [5]], + } + ``` + + Args: + window_size: A `tf.int64` scalar `tf.Tensor`, representing the number of + elements in the sliding window. + stride: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the + steps moving the sliding window forward for one iteration. The default + is `1`. It must be in `[1, window_size)`. + + Returns: + A `Dataset` transformation function, which can be passed to + @{tf.data.Dataset.apply}. + """ + def _apply_fn(dataset): + return _SlideDataset(dataset, window_size, stride) + + return _apply_fn diff --git a/tensorflow/contrib/data/python/ops/stats_ops.py b/tensorflow/contrib/data/python/ops/stats_ops.py index 9cd1701c397b5a0bf5cc47c1bcab033704794d80..b5cf0fcfe91ebc22444302fca5d488a278ef2994 100644 --- a/tensorflow/contrib/data/python/ops/stats_ops.py +++ b/tensorflow/contrib/data/python/ops/stats_ops.py @@ -47,7 +47,7 @@ class StatsAggregator(object): dataset = ... iterator = dataset.make_one_shot_iterator() stats_aggregator = stats_ops.StatsAggregator() - set_op = stats_op.set_stats_aggregator_op(iterator, stats_aggregator) + set_op = stats_aggregator.subscribe(iterator) with tf.Session() as sess: # Running `set_op` will associate `iterator` with `stats_aggregator`. diff --git a/tensorflow/contrib/data/python/ops/threadpool.py b/tensorflow/contrib/data/python/ops/threadpool.py new file mode 100644 index 0000000000000000000000000000000000000000..56f67e1766bbaff680bdff6b939df0c3ba68c679 --- /dev/null +++ b/tensorflow/contrib/data/python/ops/threadpool.py @@ -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. +# ============================================================================== +"""Experimental API for controlling threading in `tf.data` pipelines.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import threading + +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.eager import context +from tensorflow.python.ops import resource_variable_ops + +_uid_counter = 0 +_uid_lock = threading.Lock() + + +def _generate_shared_name(prefix): + with _uid_lock: + global _uid_counter + uid = _uid_counter + _uid_counter += 1 + return "{}{}".format(prefix, uid) + + +class PrivateThreadPool(object): + """A stateful resource that represents a private thread pool.""" + + def __init__(self, num_threads, display_name=None): + """Creates a `PrivateThreadPool` with the given number of threads.""" + if context.executing_eagerly(): + shared_name = _generate_shared_name("privatethreadpool") + self._resource = gen_dataset_ops.thread_pool_handle( + num_threads=num_threads, + display_name=display_name, + shared_name=shared_name) + self._resource_deleter = resource_variable_ops.EagerResourceDeleter( + handle=self._resource, handle_device=context.context().device_name) + else: + self._resource = gen_dataset_ops.thread_pool_handle( + num_threads=num_threads, display_name=display_name) + + +class _ThreadPoolDataset(dataset_ops.Dataset): + """A `Dataset` that acts as an identity, and sets a custom threadpool.""" + + def __init__(self, input_dataset, thread_pool): + super(_ThreadPoolDataset, self).__init__() + self._input_dataset = input_dataset + self._thread_pool = thread_pool + + def _as_variant_tensor(self): + return gen_dataset_ops.thread_pool_dataset( + self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access + self._thread_pool._resource, # pylint: disable=protected-access + output_shapes=nest.flatten( + sparse.as_dense_shapes(self.output_shapes, self.output_classes)), + output_types=nest.flatten( + sparse.as_dense_types(self.output_types, self.output_classes))) + + @property + def output_shapes(self): + return self._input_dataset.output_shapes + + @property + def output_types(self): + return self._input_dataset.output_types + + @property + def output_classes(self): + return self._input_dataset.output_classes + + +def override_threadpool(dataset, thread_pool): + """Returns a new dataset that uses the given thread pool for its operations. + + Args: + dataset: A `tf.data.Dataset` object. + thread_pool: A `PrivateThreadPool` object. + + Returns: + A dataset containing the same values as `dataset`, but which uses + `thread_pool` to compute any of its parallel operations (such as + @{tf.data.Dataset.map}). + """ + return _ThreadPoolDataset(dataset, thread_pool) diff --git a/tensorflow/contrib/data/python/ops/unique.py b/tensorflow/contrib/data/python/ops/unique.py index 133e17d20d0fc4c8d52cef3c95c132374e927a0b..765ef3f9b6d42c9d7af3ce4916731d37d65c9260 100644 --- a/tensorflow/contrib/data/python/ops/unique.py +++ b/tensorflow/contrib/data/python/ops/unique.py @@ -17,11 +17,12 @@ 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.framework import dtypes -from tensorflow.python.ops import gen_dataset_ops def unique(): diff --git a/tensorflow/contrib/decision_trees/proto/BUILD b/tensorflow/contrib/decision_trees/proto/BUILD index f6de5998d73a4869d2444cd90c9b64d1a2c889ac..ae3847b8b62452b1afbe472fcb6369181ec60b73 100644 --- a/tensorflow/contrib/decision_trees/proto/BUILD +++ b/tensorflow/contrib/decision_trees/proto/BUILD @@ -25,7 +25,6 @@ tf_proto_library( name = "generic_tree_model", srcs = ["generic_tree_model.proto"], cc_api_version = 2, - go_api_version = 2, java_api_version = 2, visibility = ["//visibility:public"], ) @@ -34,7 +33,6 @@ tf_proto_library( name = "generic_tree_model_extensions", srcs = ["generic_tree_model_extensions.proto"], cc_api_version = 2, - go_api_version = 2, protodeps = [":generic_tree_model"], visibility = ["//visibility:public"], ) diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD index 7f510c42215f48a9e795eb81bd9f66b0a2108335..1bd73ee7044de34988144196f53299db2fb80fcf 100644 --- a/tensorflow/contrib/distributions/BUILD +++ b/tensorflow/contrib/distributions/BUILD @@ -251,6 +251,21 @@ cuda_py_test( ], ) +cuda_py_test( + name = "kumaraswamy_test", + srcs = ["python/kernel_tests/kumaraswamy_test.py"], + additional_deps = [ + ":distributions_py", + "//third_party/py/numpy", + "//tensorflow/python:client", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:nn_ops", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "moving_stats_test", size = "small", @@ -335,6 +350,7 @@ cuda_py_test( "//tensorflow/python:nn_ops", "//tensorflow/python:platform_test", ], + tags = ["nomsan"], ) cuda_py_test( @@ -403,7 +419,7 @@ cuda_py_test( cuda_py_test( name = "poisson_lognormal_test", - size = "small", + size = "medium", srcs = ["python/kernel_tests/poisson_lognormal_test.py"], additional_deps = [ ":distributions_py", @@ -438,6 +454,7 @@ cuda_py_test( "//tensorflow/python:framework_test_lib", "//tensorflow/python:platform_test", ], + tags = ["no_windows"], # TODO: needs investigation on Windows ) cuda_py_test( @@ -459,6 +476,24 @@ cuda_py_test( tags = ["nomsan"], # disable to avoid false positives from scipy. ) +cuda_py_test( + name = "statistical_testing_test", + size = "medium", + srcs = [ + "python/kernel_tests/statistical_testing_test.py", + ], + additional_deps = [ + ":distributions_py", + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + ], + tags = [ + "manual", + "noasan", + "noguitar", + ], +) + cuda_py_test( name = "vector_sinh_arcsinh_diag_test", size = "medium", @@ -782,6 +817,25 @@ cuda_py_test( tags = ["noasan"], # times out b/63678675 ) +cuda_py_test( + name = "affine_scalar_test", + size = "small", + srcs = ["python/kernel_tests/bijectors/affine_scalar_test.py"], + additional_deps = [ + ":bijectors_py", + ":distributions_py", + "//third_party/py/numpy", + "@six_archive//:six", + "//tensorflow/contrib/linalg:linalg_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "affine_linear_operator_test", size = "small", @@ -801,6 +855,22 @@ cuda_py_test( ], ) +cuda_py_test( + name = "batch_normalization_test", + size = "small", + srcs = ["python/kernel_tests/bijectors/batch_normalization_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 = "chain_test", size = "small", @@ -915,6 +985,25 @@ cuda_py_test( ], ) +cuda_py_test( + name = "kumaraswamy_bijector_test", + size = "small", + srcs = ["python/kernel_tests/bijectors/kumaraswamy_bijector_test.py"], + additional_deps = [ + ":bijectors_py", + ":distributions_py", + "//third_party/py/numpy", + "@six_archive//:six", + "//tensorflow/contrib/linalg:linalg_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "masked_autoregressive_test", size = "small", @@ -984,7 +1073,7 @@ cuda_py_test( cuda_py_test( name = "reshape_test", - size = "small", + size = "medium", srcs = ["python/kernel_tests/bijectors/reshape_test.py"], additional_deps = [ ":bijectors_py", @@ -1055,6 +1144,7 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], + tags = ["no_windows"], # TODO: needs investigation on Windows ) cuda_py_test( @@ -1095,6 +1185,25 @@ cuda_py_test( ], ) +cuda_py_test( + name = "square_test", + size = "small", + srcs = ["python/kernel_tests/bijectors/square_test.py"], + additional_deps = [ + ":bijectors_py", + ":distributions_py", + "//third_party/py/numpy", + "@six_archive//:six", + "//tensorflow/contrib/linalg:linalg_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "weibull_test", size = "small", diff --git a/tensorflow/contrib/distributions/__init__.py b/tensorflow/contrib/distributions/__init__.py index 60a187e541df4a794ae3944c30c427944915f7d0..61c411271d0bb8d7b4cc3b14992b82ec1e5674ed 100644 --- a/tensorflow/contrib/distributions/__init__.py +++ b/tensorflow/contrib/distributions/__init__.py @@ -40,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 * @@ -97,7 +98,6 @@ _allowed_symbols = [ 'Autoregressive', 'Binomial', 'Bernoulli', - 'BernoulliWithSigmoidProbs', 'Beta', 'BetaWithSoftplusConcentration', 'Categorical', @@ -115,6 +115,7 @@ _allowed_symbols = [ 'Independent', 'InverseGamma', 'InverseGammaWithSoftplusConcentrationRate', + 'Kumaraswamy', 'Laplace', 'LaplaceWithSoftplusScale', 'Logistic', diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_scalar_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_scalar_test.py new file mode 100644 index 0000000000000000000000000000000000000000..16173a166fd943413345036df12245c2a4ab8343 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_scalar_test.py @@ -0,0 +1,153 @@ +# 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. +# ============================================================================== +"""Affine Scalar Tests.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import AffineScalar +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency +from tensorflow.python.platform import test + + +class AffineScalarBijectorTest(test.TestCase): + """Tests correctness of the Y = scale @ x + shift transformation.""" + + def testProperties(self): + with self.test_session(): + mu = -1. + # scale corresponds to 1. + bijector = AffineScalar(shift=mu) + self.assertEqual("affine_scalar", bijector.name) + + def testNoBatchScalar(self): + with self.test_session() as sess: + + def static_run(fun, x): + return fun(x).eval() + + def dynamic_run(fun, x_value): + x_value = np.array(x_value) + x = array_ops.placeholder(dtypes.float32, name="x") + return sess.run(fun(x), feed_dict={x: x_value}) + + for run in (static_run, dynamic_run): + mu = -1. + # Corresponds to scale = 2 + bijector = AffineScalar(shift=mu, scale=2.) + x = [1., 2, 3] # Three scalar samples (no batches). + self.assertAllClose([1., 3, 5], run(bijector.forward, x)) + self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x)) + self.assertAllClose([-np.log(2.)] * 3, + run(bijector.inverse_log_det_jacobian, x)) + + def testOneBatchScalarViaIdentityIn64BitUserProvidesShiftOnly(self): + with self.test_session() as sess: + + def static_run(fun, x): + return fun(x).eval() + + def dynamic_run(fun, x_value): + x_value = np.array(x_value).astype(np.float64) + x = array_ops.placeholder(dtypes.float64, name="x") + return sess.run(fun(x), feed_dict={x: x_value}) + + for run in (static_run, dynamic_run): + mu = np.float64([1.]) + # One batch, scalar. + # Corresponds to scale = 1. + bijector = AffineScalar(shift=mu) + x = np.float64([1.]) # One sample from one batches. + self.assertAllClose([2.], run(bijector.forward, x)) + self.assertAllClose([0.], run(bijector.inverse, x)) + self.assertAllClose([0.], run(bijector.inverse_log_det_jacobian, x)) + + def testOneBatchScalarViaIdentityIn64BitUserProvidesScaleOnly(self): + with self.test_session() as sess: + + def static_run(fun, x): + return fun(x).eval() + + def dynamic_run(fun, x_value): + x_value = np.array(x_value).astype(np.float64) + x = array_ops.placeholder(dtypes.float64, name="x") + return sess.run(fun(x), feed_dict={x: x_value}) + + for run in (static_run, dynamic_run): + multiplier = np.float64([2.]) + # One batch, scalar. + # Corresponds to scale = 2, shift = 0. + bijector = AffineScalar(scale=multiplier) + x = np.float64([1.]) # One sample from one batches. + self.assertAllClose([2.], run(bijector.forward, x)) + self.assertAllClose([0.5], run(bijector.inverse, x)) + self.assertAllClose([np.log(0.5)], + run(bijector.inverse_log_det_jacobian, x)) + + def testTwoBatchScalarIdentityViaIdentity(self): + with self.test_session() as sess: + + def static_run(fun, x): + return fun(x).eval() + + def dynamic_run(fun, x_value): + x_value = np.array(x_value) + x = array_ops.placeholder(dtypes.float32, name="x") + return sess.run(fun(x), feed_dict={x: x_value}) + + for run in (static_run, dynamic_run): + mu = [1., -1] + # Univariate, two batches. + # Corresponds to scale = 1. + bijector = AffineScalar(shift=mu) + x = [1., 1] # One sample from each of two batches. + self.assertAllClose([2., 0], run(bijector.forward, x)) + self.assertAllClose([0., 2], run(bijector.inverse, x)) + self.assertAllClose([0., 0.], run(bijector.inverse_log_det_jacobian, x)) + + def testTwoBatchScalarIdentityViaScale(self): + with self.test_session() as sess: + + def static_run(fun, x): + return fun(x).eval() + + def dynamic_run(fun, x_value): + x_value = np.array(x_value) + x = array_ops.placeholder(dtypes.float32, name="x") + return sess.run(fun(x), feed_dict={x: x_value}) + + for run in (static_run, dynamic_run): + mu = [1., -1] + # Univariate, two batches. + # Corresponds to scale = 1. + bijector = AffineScalar(shift=mu, scale=[2., 1]) + x = [1., 1] # One sample from each of two batches. + self.assertAllClose([3., 0], run(bijector.forward, x)) + self.assertAllClose([0., 2], run(bijector.inverse, x)) + self.assertAllClose( + [-np.log(2), 0.], run(bijector.inverse_log_det_jacobian, x)) + + def testScalarCongruency(self): + with self.test_session(): + bijector = AffineScalar(shift=3.6, scale=0.42) + assert_scalar_congruency(bijector, lower_x=-2., upper_x=2.) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_test.py index c9158117f7a982e37047e8dd2b534a30040a87d9..077e6176b4e7aecb28369d49edad6d1367cc7259 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_test.py @@ -25,7 +25,6 @@ import numpy as np from tensorflow.contrib.distributions.python.ops.bijectors.affine import Affine from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops -from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency from tensorflow.python.platform import test @@ -36,192 +35,9 @@ class AffineBijectorTest(test.TestCase): with self.test_session(): mu = -1. # scale corresponds to 1. - bijector = Affine(shift=mu, event_ndims=0) + bijector = Affine(shift=mu) self.assertEqual("affine", bijector.name) - def testNoBatchScalarViaIdentity(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = -1. - # Corresponds to scale = 2 - bijector = Affine( - shift=mu, scale_identity_multiplier=2., event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = [1., 2, 3] # Three scalar samples (no batches). - self.assertAllClose([1., 3, 5], run(bijector.forward, x)) - self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x)) - self.assertAllClose(-np.log(2.), - run(bijector.inverse_log_det_jacobian, x)) - - def testNoBatchScalarViaDiag(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = -1. - # Corresponds to scale = 2 - bijector = Affine(shift=mu, scale_identity_multiplier=2., event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = [1., 2, 3] # Three scalar samples (no batches). - self.assertAllClose([1., 3, 5], run(bijector.forward, x)) - self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x)) - self.assertAllClose(-np.log(2.), - run(bijector.inverse_log_det_jacobian, x)) - - def testWeirdSampleNoBatchScalarViaDiagMultiplier(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = -1. - # Corresponds to scale = 2. - bijector = Affine( - shift=mu, scale_identity_multiplier=2., event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = [[1., 2, 3], [4, 5, 6]] # Weird sample shape. - self.assertAllClose([[1., 3, 5], - [7, 9, 11]], - run(bijector.forward, x)) - self.assertAllClose([[1., 1.5, 2.], - [2.5, 3, 3.5]], - run(bijector.inverse, x)) - self.assertAllClose(-np.log(2.), - run(bijector.inverse_log_det_jacobian, x)) - - def testOneBatchScalarViaIdentityIn64BitUserProvidesShiftOnly(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value).astype(np.float64) - x = array_ops.placeholder(dtypes.float64, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = np.float64([1.]) - # One batch, scalar. - # Corresponds to scale = 1. - bijector = Affine(shift=mu, event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = np.float64([1.]) # One sample from one batches. - self.assertAllClose([2.], run(bijector.forward, x)) - self.assertAllClose([0.], run(bijector.inverse, x)) - self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x)) - - def testOneBatchScalarViaIdentityIn64BitUserProvidesMultiplierOnly(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value).astype(np.float64) - x = array_ops.placeholder(dtypes.float64, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - multiplier = np.float64([2.]) - # One batch, scalar. - # Corresponds to scale = 2, shift = 0. - bijector = Affine(scale_identity_multiplier=multiplier, event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = np.float64([1.]) # One sample from one batches. - self.assertAllClose([2.], run(bijector.forward, x)) - self.assertAllClose([0.5], run(bijector.inverse, x)) - self.assertAllClose([np.log(0.5)], - run(bijector.inverse_log_det_jacobian, x)) - - def testOneBatchScalarViaDiagMultiplier(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = [1.] - # One batch, scalar. - # Corresponds to scale = 1. - bijector = Affine(shift=mu, scale_identity_multiplier=1., event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = [1.] # One sample from one batches. - self.assertAllClose([2.], run(bijector.forward, x)) - self.assertAllClose([0.], run(bijector.inverse, x)) - self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x)) - - def testTwoBatchScalarIdentityViaIdentity(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = [1., -1] - # Univariate, two batches. - # Corresponds to scale = 1. - bijector = Affine(shift=mu, event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = [1., 1] # One sample from each of two batches. - self.assertAllClose([2., 0], run(bijector.forward, x)) - self.assertAllClose([0., 2], run(bijector.inverse, x)) - self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x)) - - def testTwoBatchScalarIdentityViaDiagMultiplier(self): - with self.test_session() as sess: - - def static_run(fun, x): - return fun(x).eval() - - def dynamic_run(fun, x_value): - x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) - - for run in (static_run, dynamic_run): - mu = [1., -1] - # Univariate, two batches. - # Corresponds to scale = 1. - bijector = Affine(shift=mu, scale_identity_multiplier=1., event_ndims=0) - self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar" - x = [1., 1] # One sample from each of two batches. - self.assertAllClose([2., 0], run(bijector.forward, x)) - self.assertAllClose([0., 2], run(bijector.inverse, x)) - self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x)) - def testNoBatchMultivariateIdentity(self): with self.test_session() as sess: @@ -238,7 +54,6 @@ class AffineBijectorTest(test.TestCase): # Multivariate # Corresponds to scale = [[1., 0], [0, 1.]] bijector = Affine(shift=mu) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [1., 1] # matmul(sigma, x) + shift # = [-1, -1] + [1, -1] @@ -269,7 +84,6 @@ class AffineBijectorTest(test.TestCase): # Multivariate # Corresponds to scale = [[2., 0], [0, 1.]] bijector = Affine(shift=mu, scale_diag=[2., 1]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [1., 1] # matmul(sigma, x) + shift # = [-1, -1] + [1, -1] @@ -297,22 +111,17 @@ class AffineBijectorTest(test.TestCase): x = array_ops.placeholder(dtypes.float32, name="x") mu = array_ops.placeholder(dtypes.float32, name="mu") scale_diag = array_ops.placeholder(dtypes.float32, name="scale_diag") - event_ndims = array_ops.placeholder(dtypes.int32, name="event_ndims") x_value = np.array([[1., 1]], dtype=np.float32) mu_value = np.array([1., -1], dtype=np.float32) scale_diag_value = np.array([2., 2], dtype=np.float32) - event_ndims_value = np.array(1, dtype=np.int32) feed_dict = { x: x_value, mu: mu_value, scale_diag: scale_diag_value, - event_ndims: event_ndims_value } - bijector = Affine( - shift=mu, scale_diag=scale_diag, event_ndims=event_ndims) - self.assertEqual(1, sess.run(bijector.event_ndims, feed_dict)) + bijector = Affine(shift=mu, scale_diag=scale_diag) self.assertAllClose([[3., 1]], sess.run(bijector.forward(x), feed_dict)) self.assertAllClose([[0., 1]], sess.run(bijector.inverse(x), feed_dict)) self.assertAllClose( @@ -335,7 +144,6 @@ class AffineBijectorTest(test.TestCase): # Corresponds to 1 2x2 matrix, with twos on the diagonal. scale = 2. bijector = Affine(shift=mu, scale_identity_multiplier=scale) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [[[1., 1]]] self.assertAllClose([[[3., 1]]], run(bijector.forward, x)) self.assertAllClose([[[0., 1]]], run(bijector.inverse, x)) @@ -358,7 +166,6 @@ class AffineBijectorTest(test.TestCase): # Corresponds to 1 2x2 matrix, with twos on the diagonal. scale_diag = [[2., 2]] bijector = Affine(shift=mu, scale_diag=scale_diag) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [[[1., 1]]] self.assertAllClose([[[3., 1]]], run(bijector.forward, x)) self.assertAllClose([[[0., 1]]], run(bijector.inverse, x)) @@ -370,23 +177,18 @@ class AffineBijectorTest(test.TestCase): x = array_ops.placeholder(dtypes.float32, name="x") mu = array_ops.placeholder(dtypes.float32, name="mu") scale_diag = array_ops.placeholder(dtypes.float32, name="scale_diag") - event_ndims = array_ops.placeholder(dtypes.int32, name="event_ndims") x_value = np.array([[[1., 1]]], dtype=np.float32) mu_value = np.array([[1., -1]], dtype=np.float32) scale_diag_value = np.array([[2., 2]], dtype=np.float32) - event_ndims_value = 1 feed_dict = { x: x_value, mu: mu_value, scale_diag: scale_diag_value, - event_ndims: event_ndims_value } - bijector = Affine( - shift=mu, scale_diag=scale_diag, event_ndims=event_ndims) - self.assertEqual(1, sess.run(bijector.event_ndims, feed_dict)) + bijector = Affine(shift=mu, scale_diag=scale_diag) self.assertAllClose([[[3., 1]]], sess.run(bijector.forward(x), feed_dict)) self.assertAllClose([[[0., 1]]], sess.run(bijector.inverse(x), feed_dict)) self.assertAllClose([-np.log(4)], @@ -410,9 +212,7 @@ class AffineBijectorTest(test.TestCase): bijector = Affine( shift=mu, scale_identity_multiplier=1., - scale_diag=[1., 1., 1.], - event_ndims=1) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" + scale_diag=[1., 1., 1.]) x = [1., 2, 3] # Three scalar samples (no batches). self.assertAllClose([1., 3, 5], run(bijector.forward, x)) self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x)) @@ -437,7 +237,6 @@ class AffineBijectorTest(test.TestCase): shift=mu, scale_identity_multiplier=1., scale_tril=[[1., 0], [2., 1]]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [[1., 2]] # One multivariate sample. self.assertAllClose([[1., 5]], run(bijector.forward, x)) self.assertAllClose([[1., 0.5]], run(bijector.inverse, x)) @@ -460,7 +259,6 @@ class AffineBijectorTest(test.TestCase): # scale = [[2., 0], [2, 3]] bijector = Affine( shift=mu, scale_diag=[1., 2.], scale_tril=[[1., 0], [2., 1]]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [[1., 2]] # One multivariate sample. self.assertAllClose([[1., 7]], run(bijector.forward, x)) self.assertAllClose([[1., 1 / 3.]], run(bijector.inverse, x)) @@ -486,7 +284,6 @@ class AffineBijectorTest(test.TestCase): scale_identity_multiplier=1.0, scale_diag=[1., 2.], scale_tril=[[1., 0], [2., 1]]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [[1., 2]] # One multivariate sample. self.assertAllClose([[2., 9]], run(bijector.forward, x)) self.assertAllClose([[2 / 3., 5 / 12.]], run(bijector.inverse, x)) @@ -514,7 +311,6 @@ class AffineBijectorTest(test.TestCase): scale_perturb_factor=[[2., 0], [0., 0], [0, 1]]) bijector_ref = Affine(shift=mu, scale_diag=[10., 2, 3]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [1., 2, 3] # Vector. self.assertAllClose([9., 3, 8], run(bijector.forward, x)) self.assertAllClose( @@ -550,7 +346,6 @@ class AffineBijectorTest(test.TestCase): scale_perturb_factor=[[2., 0], [0., 0], [0, 1]]) bijector_ref = Affine(shift=mu, scale_diag=[10., 3, 5]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [1., 2, 3] # Vector. self.assertAllClose([9., 5, 14], run(bijector.forward, x)) self.assertAllClose( @@ -586,7 +381,6 @@ class AffineBijectorTest(test.TestCase): bijector_ref = Affine( shift=mu, scale_tril=[[10., 0, 0], [1, 3, 0], [2, 3, 5]]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [1., 2, 3] # Vector. self.assertAllClose([9., 6, 22], run(bijector.forward, x)) self.assertAllClose( @@ -622,7 +416,6 @@ class AffineBijectorTest(test.TestCase): bijector_ref = Affine( shift=mu, scale_tril=[[6., 0, 0], [1, 3, 0], [2, 3, 5]]) - self.assertEqual(1, bijector.event_ndims.eval()) # "is vector" x = [1., 2, 3] # Vector. self.assertAllClose([5., 6, 22], run(bijector.forward, x)) self.assertAllClose( @@ -647,38 +440,6 @@ class AffineBijectorTest(test.TestCase): with self.assertRaisesOpError("diagonal part must be non-zero"): bijector.forward([1., 1.]).eval() - def testEventNdimsLargerThanOneRaises(self): - with self.test_session(): - mu = [1., -1] - with self.assertRaisesRegexp( - ValueError, (r"event_ndims\(2\) was not 0 or 1")): - # Scale corresponds to 2x2 identity matrix. - bijector = Affine(shift=mu, event_ndims=2, validate_args=True) - bijector.forward([1., 1.]).eval() - - def testScaleZeroScalarRaises(self): - with self.test_session(): - mu = -1. - # Check Identity matrix with zero scaling. - bijector = Affine( - shift=mu, - scale_identity_multiplier=0., - event_ndims=0, - validate_args=True) - with self.assertRaisesOpError("identity_multiplier should be non-zero"): - bijector.forward(1.).eval() - - def testScaleDiagAndEventNdimsZeroRaises(self): - # Check Diag matrix with zero scaling. - with self.assertRaisesRegexp(ValueError, "only scale argument"): - Affine(shift=None, scale_diag=[0.0], event_ndims=0, validate_args=True) - - def testScalarCongruency(self): - with self.test_session(): - bijector = Affine( - shift=3.6, scale_identity_multiplier=0.42, event_ndims=0) - assert_scalar_congruency(bijector, lower_x=-2., upper_x=2.) - def _makeScale(self, x, scale_identity_multiplier=None, @@ -747,14 +508,12 @@ class AffineBijectorTest(test.TestCase): scale_args = dict({"x": x}, **args) scale = self._makeScale(**scale_args) - bijector_args = dict({"event_ndims": 1}, **args) - # We haven't specified enough information for the scale. if scale is None: with self.assertRaisesRegexp(ValueError, ("must be specified.")): - bijector = Affine(shift=shift, **bijector_args) + bijector = Affine(shift=shift, **args) else: - bijector = Affine(shift=shift, **bijector_args) + bijector = Affine(shift=shift, **args) np_x = x # For the case a vector is passed in, we need to make the shape # match the matrix for matmul to work. @@ -829,15 +588,5 @@ class AffineBijectorTest(test.TestCase): x=np.array( [1., 2], dtype=np.float32)) - def testScalarEventIdentityScale(self): - with self.test_session() as sess: - doubler = Affine( - scale_identity_multiplier=2., - event_ndims=0) - doubler2 = doubler.inverse_log_det_jacobian(2.) - doubler2_ildj_ = sess.run([doubler2]) - self.assertAllClose([-np.log(2.)], doubler2_ildj_) - - if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/batch_normalization_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/batch_normalization_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a215a4a2b1ffbea7951bdb9b4352ed567e0b1e41 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/batch_normalization_test.py @@ -0,0 +1,236 @@ +# 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 BatchNorm Bijector.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib import distributions +from tensorflow.contrib.distributions.python.ops import test_util +from tensorflow.contrib.distributions.python.ops.bijectors.batch_normalization import BatchNormalization +from tensorflow.contrib.distributions.python.ops.bijectors.invert import Invert +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.layers import normalization +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_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 +from tensorflow.python.training import adam + + +class BatchNormTest(test_util.VectorDistributionTestHelpers, + test.TestCase): + + def _reduction_axes(self, input_shape, event_dims): + if isinstance(event_dims, int): + event_dims = [event_dims] + ndims = len(input_shape) + # Convert event_dims to non-negative indexing. + event_dims = list(event_dims) + for idx, x in enumerate(event_dims): + if x < 0: + event_dims[idx] = ndims + x + return tuple(i for i in range(ndims) if i not in event_dims) + + def testForwardInverse(self): + """Tests forward and backward passes with different event shapes. + + input_shape: Tuple of shapes for input tensor. + event_dims: Tuple of dimension indices that will be normalized. + training: Boolean of whether bijector runs in training or inference mode. + """ + params = [ + ((5*2, 4), [-1], False), + ((5, 2, 4), [-1], False), + ((5, 2, 4), [1, 2], False), + ((5, 2, 4), [0, 1], False), + ((5*2, 4), [-1], True), + ((5, 2, 4), [-1], True), + ((5, 2, 4), [1, 2], True), + ((5, 2, 4), [0, 1], True) + ] + for input_shape, event_dims, training in params: + x_ = np.arange(5 * 4 * 2).astype(np.float32).reshape(input_shape) + with self.test_session() as sess: + x = constant_op.constant(x_) + # When training, memorize the exact mean of the last + # minibatch that it normalized (instead of moving average assignment). + layer = normalization.BatchNormalization( + axis=event_dims, momentum=0., epsilon=0.) + batch_norm = BatchNormalization( + batchnorm_layer=layer, training=training) + # Minibatch statistics are saved only after norm_x has been computed. + norm_x = batch_norm.inverse(x) + with ops.control_dependencies(batch_norm.batchnorm.updates): + moving_mean = array_ops.identity(batch_norm.batchnorm.moving_mean) + moving_var = array_ops.identity(batch_norm.batchnorm.moving_variance) + denorm_x = batch_norm.forward(array_ops.identity(norm_x)) + fldj = batch_norm.forward_log_det_jacobian(x) + # Use identity to invalidate cache. + ildj = batch_norm.inverse_log_det_jacobian( + array_ops.identity(denorm_x)) + variables.global_variables_initializer().run() + # Update variables. + norm_x_ = sess.run(norm_x) + [ + norm_x_, + moving_mean_, + moving_var_, + denorm_x_, + ildj_, + fldj_, + ] = sess.run([ + norm_x, + moving_mean, + moving_var, + denorm_x, + ildj, + fldj, + ]) + self.assertEqual("batch_normalization", batch_norm.name) + + reduction_axes = self._reduction_axes(input_shape, event_dims) + keepdims = len(event_dims) > 1 + + expected_batch_mean = np.mean( + x_, axis=reduction_axes, keepdims=keepdims) + expected_batch_var = np.var(x_, axis=reduction_axes, keepdims=keepdims) + + if training: + # When training=True, values become normalized across batch dim and + # original values are recovered after de-normalizing. + zeros = np.zeros_like(norm_x_) + self.assertAllClose(np.mean(zeros, axis=reduction_axes), + np.mean(norm_x_, axis=reduction_axes)) + + self.assertAllClose(expected_batch_mean, moving_mean_) + self.assertAllClose(expected_batch_var, moving_var_) + self.assertAllClose(x_, denorm_x_, atol=1e-5) + # Since moving statistics are set to batch statistics after + # normalization, ildj and -fldj should match. + self.assertAllClose(ildj_, -fldj_) + # ildj is computed with minibatch statistics. + expected_ildj = np.sum(np.log(1.) - .5 * np.log( + expected_batch_var + batch_norm.batchnorm.epsilon)) + self.assertAllClose(expected_ildj, ildj_) + else: + # When training=False, moving_mean, moving_var remain at their + # initialized values (0., 1.), resulting in no scale/shift (a small + # shift occurs if epsilon > 0.) + self.assertAllClose(x_, norm_x_) + self.assertAllClose(x_, denorm_x_, atol=1e-5) + # ildj is computed with saved statistics. + expected_ildj = np.sum( + np.log(1.) - .5 * np.log(1. + batch_norm.batchnorm.epsilon)) + self.assertAllClose(expected_ildj, ildj_) + + def testMaximumLikelihoodTraining(self): + # Test Maximum Likelihood training with default bijector. + with self.test_session() as sess: + base_dist = distributions.MultivariateNormalDiag(loc=[0., 0.]) + batch_norm = BatchNormalization(training=True) + dist = transformed_distribution_lib.TransformedDistribution( + distribution=base_dist, + bijector=batch_norm) + target_dist = distributions.MultivariateNormalDiag(loc=[1., 2.]) + target_samples = target_dist.sample(100) + dist_samples = dist.sample(3000) + loss = -math_ops.reduce_mean(dist.log_prob(target_samples)) + with ops.control_dependencies(batch_norm.batchnorm.updates): + train_op = adam.AdamOptimizer(1e-2).minimize(loss) + moving_mean = array_ops.identity(batch_norm.batchnorm.moving_mean) + moving_var = array_ops.identity(batch_norm.batchnorm.moving_variance) + variables.global_variables_initializer().run() + for _ in range(3000): + sess.run(train_op) + [ + dist_samples_, + moving_mean_, + moving_var_ + ] = sess.run([ + dist_samples, + moving_mean, + moving_var + ]) + self.assertAllClose([1., 2.], np.mean(dist_samples_, axis=0), atol=5e-2) + self.assertAllClose([1., 2.], moving_mean_, atol=5e-2) + self.assertAllClose([1., 1.], moving_var_, atol=5e-2) + + def testLogProb(self): + with self.test_session() as sess: + layer = normalization.BatchNormalization(epsilon=0.) + batch_norm = BatchNormalization(batchnorm_layer=layer, training=False) + base_dist = distributions.MultivariateNormalDiag(loc=[0., 0.]) + dist = transformed_distribution_lib.TransformedDistribution( + distribution=base_dist, + bijector=batch_norm, + validate_args=True) + samples = dist.sample(int(1e5)) + # No volume distortion since training=False, bijector is initialized + # to the identity transformation. + base_log_prob = base_dist.log_prob(samples) + dist_log_prob = dist.log_prob(samples) + variables.global_variables_initializer().run() + base_log_prob_, dist_log_prob_ = sess.run([base_log_prob, dist_log_prob]) + self.assertAllClose(base_log_prob_, dist_log_prob_) + + def testMutuallyConsistent(self): + # BatchNorm bijector is only mutually consistent when training=False. + dims = 4 + with self.test_session() as sess: + layer = normalization.BatchNormalization(epsilon=0.) + batch_norm = BatchNormalization(batchnorm_layer=layer, training=False) + dist = transformed_distribution_lib.TransformedDistribution( + distribution=normal_lib.Normal(loc=0., scale=1.), + bijector=batch_norm, + 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=2., + center=0., + rtol=0.02) + + def testInvertMutuallyConsistent(self): + # BatchNorm bijector is only mutually consistent when training=False. + dims = 4 + with self.test_session() as sess: + layer = normalization.BatchNormalization(epsilon=0.) + batch_norm = Invert( + BatchNormalization(batchnorm_layer=layer, training=False)) + dist = transformed_distribution_lib.TransformedDistribution( + distribution=normal_lib.Normal(loc=0., scale=1.), + bijector=batch_norm, + 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=2., + center=0., + rtol=0.02) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/cholesky_outer_product_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/cholesky_outer_product_test.py index 0ff35304283fce9ce3f9e5d31b1258394e384d7b..f392e83d2c3da9dac43c2e87070e952ae2060b34 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/cholesky_outer_product_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/cholesky_outer_product_test.py @@ -18,70 +18,111 @@ 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 bijectors -from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops -from tensorflow.python.ops.distributions import gamma as gamma_lib -from tensorflow.python.ops.distributions import transformed_distribution as transformed_distribution_lib -from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency from tensorflow.python.platform import test -class InvertBijectorTest(test.TestCase): - """Tests the correctness of the Y = Invert(bij) transformation.""" +class CholeskyOuterProductBijectorTest(test.TestCase): + """Tests the correctness of the Y = X @ X.T transformation.""" - def testBijector(self): + def testBijectorMatrix(self): with self.test_session(): - for fwd in [ - bijectors.Identity(), - bijectors.Exp(event_ndims=1), - bijectors.Affine( - shift=[0., 1.], scale_diag=[2., 3.], event_ndims=1), - bijectors.Softplus(event_ndims=1), - bijectors.SoftmaxCentered(event_ndims=1), - bijectors.SigmoidCentered(), - ]: - rev = bijectors.Invert(fwd) - self.assertEqual("_".join(["invert", fwd.name]), rev.name) - x = [[[1., 2.], - [2., 3.]]] - self.assertAllClose(fwd.inverse(x).eval(), rev.forward(x).eval()) - self.assertAllClose(fwd.forward(x).eval(), rev.inverse(x).eval()) - self.assertAllClose( - fwd.forward_log_det_jacobian(x).eval(), - rev.inverse_log_det_jacobian(x).eval()) - self.assertAllClose( - fwd.inverse_log_det_jacobian(x).eval(), - rev.forward_log_det_jacobian(x).eval()) + bijector = bijectors.CholeskyOuterProduct(validate_args=True) + self.assertEqual("cholesky_outer_product", bijector.name) + x = [[[1., 0], [2, 1]], [[np.sqrt(2.), 0], [np.sqrt(8.), 1]]] + y = np.matmul(x, np.transpose(x, axes=(0, 2, 1))) + # Fairly easy to compute differentials since we have 2x2. + dx_dy = [[[2. * 1, 0, 0], + [2, 1, 0], + [0, 2 * 2, 2 * 1]], + [[2 * np.sqrt(2.), 0, 0], + [np.sqrt(8.), np.sqrt(2.), 0], + [0, 2 * np.sqrt(8.), 2 * 1]]] + ildj = -np.sum( + np.log(np.asarray(dx_dy).diagonal( + offset=0, axis1=1, axis2=2)), + axis=1) + self.assertAllEqual((2, 2, 2), bijector.forward(x).get_shape()) + self.assertAllEqual((2, 2, 2), bijector.inverse(y).get_shape()) + self.assertAllClose(y, bijector.forward(x).eval()) + self.assertAllClose(x, bijector.inverse(y).eval()) + self.assertAllClose( + ildj, bijector.inverse_log_det_jacobian(y).eval(), atol=0., rtol=1e-7) + self.assertAllClose( + -bijector.inverse_log_det_jacobian(y).eval(), + bijector.forward_log_det_jacobian(x).eval(), + atol=0., + rtol=1e-7) - def testScalarCongruency(self): - with self.test_session(): - bijector = bijectors.Invert(bijectors.Exp()) - assert_scalar_congruency( - bijector, lower_x=1e-3, upper_x=1.5, rtol=0.05) + def testNoBatchStatic(self): + x = np.array([[1., 0], [2, 1]]) # np.linalg.cholesky(y) + y = np.array([[1., 2], [2, 5]]) # np.matmul(x, x.T) + with self.test_session() as sess: + y_actual = bijectors.CholeskyOuterProduct().forward(x=x) + x_actual = bijectors.CholeskyOuterProduct().inverse(y=y) + [y_actual_, x_actual_] = sess.run([y_actual, x_actual]) + self.assertAllEqual([2, 2], y_actual.get_shape()) + self.assertAllEqual([2, 2], x_actual.get_shape()) + self.assertAllClose(y, y_actual_) + self.assertAllClose(x, x_actual_) - def testShapeGetters(self): - with self.test_session(): - bijector = bijectors.Invert(bijectors.SigmoidCentered(validate_args=True)) - x = tensor_shape.TensorShape([2]) - y = tensor_shape.TensorShape([]) - self.assertAllEqual(y, bijector.forward_event_shape(x)) - self.assertAllEqual( - y.as_list(), - bijector.forward_event_shape_tensor(x.as_list()).eval()) - self.assertAllEqual(x, bijector.inverse_event_shape(y)) - self.assertAllEqual( - x.as_list(), - bijector.inverse_event_shape_tensor(y.as_list()).eval()) + def testNoBatchDeferred(self): + x = np.array([[1., 0], [2, 1]]) # np.linalg.cholesky(y) + y = np.array([[1., 2], [2, 5]]) # np.matmul(x, x.T) + with self.test_session() as sess: + x_pl = array_ops.placeholder(dtypes.float32) + y_pl = array_ops.placeholder(dtypes.float32) + y_actual = bijectors.CholeskyOuterProduct().forward(x=x_pl) + x_actual = bijectors.CholeskyOuterProduct().inverse(y=y_pl) + [y_actual_, x_actual_] = sess.run([y_actual, x_actual], + feed_dict={x_pl: x, y_pl: y}) + self.assertEqual(None, y_actual.get_shape()) + self.assertEqual(None, x_actual.get_shape()) + self.assertAllClose(y, y_actual_) + self.assertAllClose(x, x_actual_) - def testDocstringExample(self): - with self.test_session(): - exp_gamma_distribution = ( - transformed_distribution_lib.TransformedDistribution( - distribution=gamma_lib.Gamma(concentration=1., rate=2.), - bijector=bijectors.Invert(bijectors.Exp()))) - self.assertAllEqual( - [], array_ops.shape(exp_gamma_distribution.sample()).eval()) + def testBatchStatic(self): + x = np.array([[[1., 0], + [2, 1]], + [[3., 0], + [1, 2]]]) # np.linalg.cholesky(y) + y = np.array([[[1., 2], + [2, 5]], + [[9., 3], + [3, 5]]]) # np.matmul(x, x.T) + with self.test_session() as sess: + y_actual = bijectors.CholeskyOuterProduct().forward(x=x) + x_actual = bijectors.CholeskyOuterProduct().inverse(y=y) + [y_actual_, x_actual_] = sess.run([y_actual, x_actual]) + self.assertEqual([2, 2, 2], y_actual.get_shape()) + self.assertEqual([2, 2, 2], x_actual.get_shape()) + self.assertAllClose(y, y_actual_) + self.assertAllClose(x, x_actual_) + + def testBatchDeferred(self): + x = np.array([[[1., 0], + [2, 1]], + [[3., 0], + [1, 2]]]) # np.linalg.cholesky(y) + y = np.array([[[1., 2], + [2, 5]], + [[9., 3], + [3, 5]]]) # np.matmul(x, x.T) + with self.test_session() as sess: + x_pl = array_ops.placeholder(dtypes.float32) + y_pl = array_ops.placeholder(dtypes.float32) + y_actual = bijectors.CholeskyOuterProduct().forward(x=x_pl) + x_actual = bijectors.CholeskyOuterProduct().inverse(y=y_pl) + [y_actual_, x_actual_] = sess.run([y_actual, x_actual], + feed_dict={x_pl: x, y_pl: y}) + self.assertEqual(None, y_actual.get_shape()) + self.assertEqual(None, x_actual.get_shape()) + self.assertAllClose(y, y_actual_) + self.assertAllClose(x, x_actual_) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/invert_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/invert_test.py index 0ff35304283fce9ce3f9e5d31b1258394e384d7b..28e3e3135455348debb002b7d457e785799e1564 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/invert_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/invert_test.py @@ -35,8 +35,7 @@ class InvertBijectorTest(test.TestCase): for fwd in [ bijectors.Identity(), bijectors.Exp(event_ndims=1), - bijectors.Affine( - shift=[0., 1.], scale_diag=[2., 3.], event_ndims=1), + bijectors.Affine(shift=[0., 1.], scale_diag=[2., 3.]), bijectors.Softplus(event_ndims=1), bijectors.SoftmaxCentered(event_ndims=1), bijectors.SigmoidCentered(), diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/kumaraswamy_bijector_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/kumaraswamy_bijector_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ad11d9f2484c4b08c67c5f82aec1320475d1d983 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/kumaraswamy_bijector_test.py @@ -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. +# ============================================================================== +"""Tests for Kumaraswamy Bijector.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distributions.python.ops.bijectors.kumaraswamy import Kumaraswamy +from tensorflow.python.ops.distributions.bijector_test_util import assert_bijective_and_finite +from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency +from tensorflow.python.platform import test + + +class KumaraswamyBijectorTest(test.TestCase): + """Tests correctness of the Kumaraswamy bijector.""" + + def testBijector(self): + with self.test_session(): + a = 2. + b = 0.3 + bijector = Kumaraswamy( + concentration1=a, concentration0=b, + event_ndims=0, validate_args=True) + self.assertEqual("kumaraswamy", bijector.name) + x = np.array([[[0.1], [0.2], [0.3], [0.4], [0.5]]], dtype=np.float32) + # Kumaraswamy cdf. This is the same as inverse(x). + y = 1. - (1. - x ** a) ** b + self.assertAllClose(y, bijector.inverse(x).eval()) + self.assertAllClose(x, bijector.forward(y).eval()) + kumaraswamy_log_pdf = (np.log(a) + np.log(b) + (a - 1) * np.log(x) + + (b - 1) * np.log1p(-x ** a)) + + self.assertAllClose( + # We should lose a dimension from calculating the determinant of the + # jacobian. + kumaraswamy_log_pdf, + bijector.inverse_log_det_jacobian(x).eval()) + self.assertAllClose( + -bijector.inverse_log_det_jacobian(x).eval(), + bijector.forward_log_det_jacobian(y).eval(), + rtol=1e-4, + atol=0.) + + def testScalarCongruency(self): + with self.test_session(): + assert_scalar_congruency( + Kumaraswamy(concentration1=0.5, concentration0=1.1), + lower_x=0., upper_x=1., n=int(10e3), rtol=0.02) + + def testBijectiveAndFinite(self): + with self.test_session(): + concentration1 = 1.2 + concentration0 = 2. + bijector = Kumaraswamy( + concentration1=concentration1, + concentration0=concentration0, validate_args=True) + # Omitting the endpoints 0 and 1, since idlj will be inifinity at these + # endpoints. + y = np.linspace(.01, 0.99, num=10).astype(np.float32) + x = 1 - (1 - y ** concentration1) ** concentration0 + assert_bijective_and_finite(bijector, x, y, rtol=1e-3) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softmax_centered_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softmax_centered_test.py index 62e3869db090e9c9327bc552d10234ff76ba28fd..4a7679daad6f6acc632eb9133078499dda89e43d 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softmax_centered_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softmax_centered_test.py @@ -21,7 +21,9 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.distributions.python.ops.bijectors.softmax_centered import SoftmaxCentered +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.bijector_test_util import assert_bijective_and_finite from tensorflow.python.platform import test @@ -76,6 +78,32 @@ class SoftmaxCenteredBijectorTest(test.TestCase): atol=0., rtol=1e-7) + def testBijectorUnknownShape(self): + with self.test_session(): + softmax = SoftmaxCentered(event_ndims=1) + self.assertEqual("softmax_centered", softmax.name) + x = array_ops.placeholder(shape=[2, None], dtype=dtypes.float32) + real_x = np.log([[2., 3, 4], [4., 8, 12]]) + y = array_ops.placeholder(shape=[2, None], dtype=dtypes.float32) + real_y = [[0.2, 0.3, 0.4, 0.1], [0.16, 0.32, 0.48, 0.04]] + self.assertAllClose(real_y, softmax.forward(x).eval( + feed_dict={x: real_x})) + self.assertAllClose(real_x, softmax.inverse(y).eval( + feed_dict={y: real_y})) + self.assertAllClose( + -np.sum(np.log(real_y), axis=1), + softmax.inverse_log_det_jacobian(y).eval( + feed_dict={y: real_y}), + atol=0., + rtol=1e-7) + self.assertAllClose( + -softmax.inverse_log_det_jacobian(y).eval( + feed_dict={y: real_y}), + softmax.forward_log_det_jacobian(x).eval( + feed_dict={x: real_x}), + atol=0., + rtol=1e-7) + def testShapeGetters(self): with self.test_session(): for x, y, b in ((tensor_shape.TensorShape([]), diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f03d6f1343a11ae4517f9034ceb0c99ca6fe7fa2 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py @@ -0,0 +1,58 @@ +# 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 Bijector.""" + +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 bijectors +from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency +from tensorflow.python.platform import test + + +class SquareBijectorTest(test.TestCase): + """Tests the correctness of the Y = X ** 2 transformation.""" + + def testBijectorScalar(self): + with self.test_session(): + bijector = bijectors.Square(validate_args=True) + self.assertEqual("square", bijector.name) + x = [[[1., 5], + [2, 1]], + [[np.sqrt(2.), 3], + [np.sqrt(8.), 1]]] + y = np.square(x) + ildj = -np.log(2.) - np.log(x) + self.assertAllClose(y, bijector.forward(x).eval()) + self.assertAllClose(x, bijector.inverse(y).eval()) + self.assertAllClose( + ildj, bijector.inverse_log_det_jacobian(y).eval(), atol=0., rtol=1e-7) + self.assertAllClose( + -bijector.inverse_log_det_jacobian(y).eval(), + bijector.forward_log_det_jacobian(x).eval(), + atol=0., + rtol=1e-7) + + def testScalarCongruency(self): + with self.test_session(): + bijector = bijectors.Square(validate_args=True) + assert_scalar_congruency(bijector, lower_x=1e-3, upper_x=1.5, rtol=0.05) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/independent_test.py b/tensorflow/contrib/distributions/python/kernel_tests/independent_test.py index 06318ca09dec851cf025fa35c83732b85824cbee..6a69f9e60b99a17c657f074597a075890265a93b 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/independent_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/independent_test.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops.distributions import bernoulli as bernoulli_lib +from tensorflow.python.ops.distributions import kullback_leibler from tensorflow.python.ops.distributions import normal as normal_lib from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging @@ -126,6 +127,100 @@ class ProductDistributionTest(test.TestCase): self.assertAllClose(sample_entropy_, actual_entropy_, rtol=0.01, atol=0.) self.assertAllClose(loc, actual_mode_, rtol=1e-6, atol=0.) + def testKLRaises(self): + ind1 = independent_lib.Independent( + distribution=normal_lib.Normal( + loc=np.float32([-1., 1]), + scale=np.float32([0.1, 0.5])), + reinterpreted_batch_ndims=1) + ind2 = independent_lib.Independent( + distribution=normal_lib.Normal( + loc=np.float32(-1), + scale=np.float32(0.5)), + reinterpreted_batch_ndims=0) + + with self.assertRaisesRegexp( + ValueError, "Event shapes do not match"): + kullback_leibler.kl_divergence(ind1, ind2) + + ind1 = independent_lib.Independent( + distribution=normal_lib.Normal( + loc=np.float32([-1., 1]), + scale=np.float32([0.1, 0.5])), + reinterpreted_batch_ndims=1) + ind2 = independent_lib.Independent( + distribution=mvn_diag_lib.MultivariateNormalDiag( + loc=np.float32([-1., 1]), + scale_diag=np.float32([0.1, 0.5])), + reinterpreted_batch_ndims=0) + + with self.assertRaisesRegexp( + NotImplementedError, "different event shapes"): + kullback_leibler.kl_divergence(ind1, ind2) + + def testKLScalarToMultivariate(self): + normal1 = normal_lib.Normal( + loc=np.float32([-1., 1]), + scale=np.float32([0.1, 0.5])) + ind1 = independent_lib.Independent( + distribution=normal1, reinterpreted_batch_ndims=1) + + normal2 = normal_lib.Normal( + loc=np.float32([-3., 3]), + scale=np.float32([0.3, 0.3])) + ind2 = independent_lib.Independent( + distribution=normal2, reinterpreted_batch_ndims=1) + + normal_kl = kullback_leibler.kl_divergence(normal1, normal2) + ind_kl = kullback_leibler.kl_divergence(ind1, ind2) + self.assertAllClose( + self.evaluate(math_ops.reduce_sum(normal_kl, axis=-1)), + self.evaluate(ind_kl)) + + def testKLIdentity(self): + normal1 = normal_lib.Normal( + loc=np.float32([-1., 1]), + scale=np.float32([0.1, 0.5])) + # This is functionally just a wrapper around normal1, + # and doesn't change any outputs. + ind1 = independent_lib.Independent( + distribution=normal1, reinterpreted_batch_ndims=0) + + normal2 = normal_lib.Normal( + loc=np.float32([-3., 3]), + scale=np.float32([0.3, 0.3])) + # This is functionally just a wrapper around normal2, + # and doesn't change any outputs. + ind2 = independent_lib.Independent( + distribution=normal2, reinterpreted_batch_ndims=0) + + normal_kl = kullback_leibler.kl_divergence(normal1, normal2) + ind_kl = kullback_leibler.kl_divergence(ind1, ind2) + self.assertAllClose( + self.evaluate(normal_kl), self.evaluate(ind_kl)) + + def testKLMultivariateToMultivariate(self): + # (1, 1, 2) batch of MVNDiag + mvn1 = mvn_diag_lib.MultivariateNormalDiag( + loc=np.float32([[[[-1., 1, 3.], [2., 4., 3.]]]]), + scale_diag=np.float32([[[0.2, 0.1, 5.], [2., 3., 4.]]])) + ind1 = independent_lib.Independent( + distribution=mvn1, reinterpreted_batch_ndims=2) + + # (1, 1, 2) batch of MVNDiag + mvn2 = mvn_diag_lib.MultivariateNormalDiag( + loc=np.float32([[[[-2., 3, 2.], [1., 3., 2.]]]]), + scale_diag=np.float32([[[0.1, 0.5, 3.], [1., 2., 1.]]])) + + ind2 = independent_lib.Independent( + distribution=mvn2, reinterpreted_batch_ndims=2) + + mvn_kl = kullback_leibler.kl_divergence(mvn1, mvn2) + ind_kl = kullback_leibler.kl_divergence(ind1, ind2) + self.assertAllClose( + self.evaluate(math_ops.reduce_sum(mvn_kl, axis=[-1, -2])), + self.evaluate(ind_kl)) + def _testMnistLike(self, static_shape): sample_shape = [4, 5] batch_shape = [10] 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..2980e2bfe93b2e2aa01d38fc9fa4650a015efc06 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/kumaraswamy_test.py @@ -0,0 +1,386 @@ +# 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 non-negative"): + dist.prob([-1., 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("Mode undefined for concentration1 <= 1."): + 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("Mode undefined for concentration0 <= 1."): + 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/poisson_test.py b/tensorflow/contrib/distributions/python/kernel_tests/poisson_test.py index d9c9008417cdb20b62390630cf887d3bd888a0d3..19a7472d91758a2dbd00c4d918853d7bae33685d 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/poisson_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/poisson_test.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function import numpy as np +from scipy import special from scipy import stats from tensorflow.contrib.distributions.python.ops import poisson as poisson_lib from tensorflow.python.framework import constant_op @@ -110,7 +111,7 @@ class PoissonTest(test.TestCase): batch_size = 6 lam = constant_op.constant([3.0] * batch_size) lam_v = 3.0 - x = [2.2, 3.1, 4., 5.5, 6., 7.] + x = [2., 3., 4., 5., 6., 7.] poisson = self._make_poisson(rate=lam) log_cdf = poisson.log_cdf(x) @@ -121,12 +122,31 @@ class PoissonTest(test.TestCase): self.assertEqual(cdf.get_shape(), (6,)) self.assertAllClose(cdf.eval(), stats.poisson.cdf(x, lam_v)) + def testPoissonCDFNonIntegerValues(self): + with self.test_session(): + batch_size = 6 + lam = constant_op.constant([3.0] * batch_size) + lam_v = 3.0 + x = np.array([2.2, 3.1, 4., 5.5, 6., 7.], dtype=np.float32) + + poisson = self._make_poisson(rate=lam) + cdf = poisson.cdf(x) + self.assertEqual(cdf.get_shape(), (6,)) + + # The Poisson CDF should be valid on these non-integer values, and + # equal to igammac(1 + x, rate). + self.assertAllClose(cdf.eval(), special.gammaincc(1. + x, lam_v)) + + with self.assertRaisesOpError("cannot contain fractional components"): + poisson_validate = self._make_poisson(rate=lam, validate_args=True) + poisson_validate.cdf(x).eval() + def testPoissonCdfMultidimensional(self): with self.test_session(): batch_size = 6 lam = constant_op.constant([[2.0, 4.0, 5.0]] * batch_size) lam_v = [2.0, 4.0, 5.0] - x = np.array([[2.2, 3.1, 4., 5.5, 6., 7.]], dtype=np.float32).T + x = np.array([[2., 3., 4., 5., 6., 7.]], dtype=np.float32).T poisson = self._make_poisson(rate=lam) log_cdf = poisson.log_cdf(x) diff --git a/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py b/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3548ac18078a0b40f117c2bf9e2b34d20cee163b --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py @@ -0,0 +1,166 @@ +# 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 the statistical testing library.""" + +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 statistical_testing as st +from tensorflow.python.framework import errors +from tensorflow.python.ops import check_ops +from tensorflow.python.platform import test + + +class StatisticalTestingTest(test.TestCase): + + def test_dkwm_design_mean_one_sample_soundness(self): + numbers = [1e-5, 1e-2, 1.1e-1, 0.9, 1., 1.02, 2., 10., 1e2, 1e5, 1e10] + rates = [1e-6, 1e-3, 1e-2, 1.1e-1, 0.2, 0.5, 0.7, 1.] + with self.test_session() as sess: + for ff in rates: + for fp in rates: + sufficient_n = st.min_num_samples_for_dkwm_mean_test( + numbers, 0., 1., false_fail_rate=ff, false_pass_rate=fp) + detectable_d = st.min_discrepancy_of_true_means_detectable_by_dkwm( + sufficient_n, 0., 1., false_fail_rate=ff, false_pass_rate=fp) + sess.run(check_ops.assert_less_equal(detectable_d, numbers)) + + def test_dkwm_design_mean_two_sample_soundness(self): + numbers = [1e-5, 1e-2, 1.1e-1, 0.9, 1., 1.02, 2., 10., 1e2, 1e5, 1e10] + rates = [1e-6, 1e-3, 1e-2, 1.1e-1, 0.2, 0.5, 0.7, 1.] + with self.test_session() as sess: + for ff in rates: + for fp in rates: + (sufficient_n1, + sufficient_n2) = st.min_num_samples_for_dkwm_mean_two_sample_test( + numbers, 0., 1., 0., 1., + false_fail_rate=ff, false_pass_rate=fp) + d_fn = st.min_discrepancy_of_true_means_detectable_by_dkwm_two_sample + detectable_d = d_fn( + sufficient_n1, 0., 1., sufficient_n2, 0., 1., + false_fail_rate=ff, false_pass_rate=fp) + sess.run(check_ops.assert_less_equal(detectable_d, numbers)) + + def test_true_mean_confidence_interval_by_dkwm_one_sample(self): + rng = np.random.RandomState(seed=0) + + num_samples = 5000 + # 5000 samples is chosen to be enough to find discrepancies of + # size 0.1 or more with assurance 1e-6, as confirmed here: + with self.test_session() as sess: + d = st.min_discrepancy_of_true_means_detectable_by_dkwm( + num_samples, 0., 1., false_fail_rate=1e-6, false_pass_rate=1e-6) + d = sess.run(d) + self.assertLess(d, 0.1) + + # Test that the confidence interval computed for the mean includes + # 0.5 and excludes 0.4 and 0.6. + with self.test_session() as sess: + samples = rng.uniform(size=num_samples).astype(np.float32) + (low, high) = st.true_mean_confidence_interval_by_dkwm( + samples, 0., 1., error_rate=1e-6) + low, high = sess.run([low, high]) + self.assertGreater(low, 0.4) + self.assertLess(low, 0.5) + self.assertGreater(high, 0.5) + self.assertLess(high, 0.6) + + def test_dkwm_mean_one_sample_assertion(self): + rng = np.random.RandomState(seed=0) + num_samples = 5000 + + # Test that the test assertion agrees that the mean of the standard + # uniform distribution is 0.5. + samples = rng.uniform(size=num_samples).astype(np.float32) + with self.test_session() as sess: + sess.run(st.assert_true_mean_equal_by_dkwm( + samples, 0., 1., 0.5, false_fail_rate=1e-6)) + + # Test that the test assertion confirms that the mean of the + # standard uniform distribution is not 0.4. + with self.assertRaises(errors.InvalidArgumentError): + sess.run(st.assert_true_mean_equal_by_dkwm( + samples, 0., 1., 0.4, false_fail_rate=1e-6)) + + # Test that the test assertion confirms that the mean of the + # standard uniform distribution is not 0.6. + with self.assertRaises(errors.InvalidArgumentError): + sess.run(st.assert_true_mean_equal_by_dkwm( + samples, 0., 1., 0.6, false_fail_rate=1e-6)) + + def test_dkwm_mean_two_sample_assertion(self): + rng = np.random.RandomState(seed=0) + num_samples = 15000 + + # 15000 samples is chosen to be enough to find discrepancies of + # size 0.1 or more with assurance 1e-6, as confirmed here: + with self.test_session() as sess: + d = st.min_discrepancy_of_true_means_detectable_by_dkwm_two_sample( + num_samples, 0., 1., num_samples, 0., 1., + false_fail_rate=1e-6, false_pass_rate=1e-6) + d = sess.run(d) + self.assertLess(d, 0.1) + + # Test that the test assertion agrees that the standard + # uniform distribution has the same mean as itself. + samples1 = rng.uniform(size=num_samples).astype(np.float32) + samples2 = rng.uniform(size=num_samples).astype(np.float32) + with self.test_session() as sess: + sess.run(st.assert_true_mean_equal_by_dkwm_two_sample( + samples1, 0., 1., samples2, 0., 1., false_fail_rate=1e-6)) + + # Test that the test assertion confirms that the mean of the + # standard uniform distribution is different from the mean of beta(2, 1). + beta_high_samples = rng.beta(2, 1, size=num_samples).astype(np.float32) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(st.assert_true_mean_equal_by_dkwm_two_sample( + samples1, 0., 1., + beta_high_samples, 0., 1., + false_fail_rate=1e-6)) + + # Test that the test assertion confirms that the mean of the + # standard uniform distribution is different from the mean of beta(1, 2). + beta_low_samples = rng.beta(1, 2, size=num_samples).astype(np.float32) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(st.assert_true_mean_equal_by_dkwm_two_sample( + samples1, 0., 1., + beta_low_samples, 0., 1., + false_fail_rate=1e-6)) + + def test_dkwm_argument_validity_checking(self): + rng = np.random.RandomState(seed=0) + samples = rng.uniform(size=5000).astype(np.float32) + + # Test that the test library complains if the given samples fall + # outside the purported bounds. + with self.test_session() as sess: + with self.assertRaises(errors.InvalidArgumentError): + sess.run(st.true_mean_confidence_interval_by_dkwm( + samples, 0., 0.5, error_rate=0.5)) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(st.true_mean_confidence_interval_by_dkwm( + samples, 0.5, 1., error_rate=0.5)) + + # But doesn't complain if they don't. + op = st.true_mean_confidence_interval_by_dkwm( + samples, 0., 1., error_rate=0.5) + _ = sess.run(op) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/transformed_distribution_test.py b/tensorflow/contrib/distributions/python/kernel_tests/transformed_distribution_test.py index cbaf74d3f66253ae5727e1ba579e2d49235b748e..af13553c32bdb6ef4038daa5e4bbef3251cff2f3 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/transformed_distribution_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/transformed_distribution_test.py @@ -245,9 +245,8 @@ class TransformedDistributionTest(test.TestCase): with self.test_session() as sess: exp2 = self._cls()( ds.Exponential(rate=0.25), - bijector=ds.bijectors.Affine( - scale_identity_multiplier=2., - event_ndims=0)) + bijector=ds.bijectors.AffineScalar(scale=2.) + ) log_prob = exp2.log_prob(1.) log_prob_ = sess.run(log_prob) base_log_prob = -0.5 * 0.25 + np.log(0.25) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py index 93923c3f083c7f5136b55e9021cbd6323684b976..452f1caa30fdbf5442274cbcc7f3549081b80ae9 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py @@ -17,7 +17,9 @@ @@AbsoluteValue @@Affine @@AffineLinearOperator +@@AffineScalar @@Bijector +@@BatchNormalization @@Chain @@CholeskyOuterProduct @@ConditionalBijector @@ -26,6 +28,7 @@ @@Identity @@Inline @@Invert +@@Kumaraswamy @@MaskedAutoregressiveFlow @@Permute @@PowerTransform @@ -36,6 +39,7 @@ @@SinhArcsinh @@SoftmaxCentered @@Softplus +@@Square @@Weibull @@masked_autoregressive_default_template @@ -52,6 +56,8 @@ from __future__ import print_function from tensorflow.contrib.distributions.python.ops.bijectors.absolute_value import * from tensorflow.contrib.distributions.python.ops.bijectors.affine import * from tensorflow.contrib.distributions.python.ops.bijectors.affine_linear_operator import * +from tensorflow.contrib.distributions.python.ops.bijectors.affine_scalar import * +from tensorflow.contrib.distributions.python.ops.bijectors.batch_normalization import * from tensorflow.contrib.distributions.python.ops.bijectors.chain import * from tensorflow.contrib.distributions.python.ops.bijectors.cholesky_outer_product import * from tensorflow.contrib.distributions.python.ops.bijectors.conditional_bijector import * @@ -59,6 +65,7 @@ from tensorflow.contrib.distributions.python.ops.bijectors.exp import * from tensorflow.contrib.distributions.python.ops.bijectors.gumbel import * from tensorflow.contrib.distributions.python.ops.bijectors.inline import * from tensorflow.contrib.distributions.python.ops.bijectors.invert import * +from tensorflow.contrib.distributions.python.ops.bijectors.kumaraswamy 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 * @@ -69,6 +76,7 @@ from tensorflow.contrib.distributions.python.ops.bijectors.sigmoid_centered impo from tensorflow.contrib.distributions.python.ops.bijectors.sinh_arcsinh import * from tensorflow.contrib.distributions.python.ops.bijectors.softmax_centered import * from tensorflow.contrib.distributions.python.ops.bijectors.softplus import * +from tensorflow.contrib.distributions.python.ops.bijectors.square import * from tensorflow.python.ops.distributions.bijector import * from tensorflow.python.ops.distributions.identity_bijector import Identity diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/affine.py b/tensorflow/contrib/distributions/python/ops/bijectors/affine.py index 05bb9c2f9bdf35e222c94db3491157893da64ebd..7fe73ada4466d38a7d352f23a55d6b90ed38c84a 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/affine.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/affine.py @@ -104,7 +104,6 @@ class Affine(bijector.Bijector): scale_tril=None, scale_perturb_factor=None, scale_perturb_diag=None, - event_ndims=1, validate_args=False, name="affine"): """Instantiates the `Affine` bijector. @@ -157,8 +156,6 @@ class Affine(bijector.Bijector): matrix. `scale_perturb_diag` has shape [N1, N2, ... r], which represents an `r x r` diagonal matrix. When `None` low rank updates will take the form `scale_perturb_factor * scale_perturb_factor.T`. - event_ndims: Scalar `int` `Tensor` indicating the number of dimensions - associated with a particular draw from the distribution. Must be 0 or 1. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. @@ -187,23 +184,6 @@ class Affine(bijector.Bijector): with self._name_scope("init", values=[ shift, scale_identity_multiplier, scale_diag, scale_tril, scale_perturb_diag, scale_perturb_factor]): - event_ndims = ops.convert_to_tensor(event_ndims, name="event_ndims") - event_ndims_const = tensor_util.constant_value(event_ndims) - if event_ndims_const is not None and event_ndims_const not in (0, 1): - raise ValueError("event_ndims(%s) was not 0 or 1" % event_ndims_const) - else: - if validate_args: - # Shape tool will catch if event_ndims is negative. - event_ndims = control_flow_ops.with_dependencies( - [check_ops.assert_less( - event_ndims, 2, message="event_ndims must be 0 or 1")], - event_ndims) - - if event_ndims_const == 0 and not self._is_only_identity_multiplier: - raise ValueError( - "If event_ndims == 0, the only scale argument you can pass is " - "scale_identity_multiplier. All others operate on vectors.") - # In the absence of `loc` and `scale`, we'll assume `dtype` is `float32`. dtype = dtypes.float32 @@ -251,12 +231,11 @@ class Affine(bijector.Bijector): self._scale = scale self._shaper = _DistributionShape( batch_ndims=batch_ndims, - event_ndims=event_ndims, + event_ndims=1, validate_args=validate_args) super(Affine, self).__init__( - event_ndims=event_ndims, + event_ndims=1, graph_parents=( - [event_ndims] + [self._scale] if tensor_util.is_tensor(self._scale) else self._scale.graph_parents + [self._shift] if self._shift is not None else []), @@ -388,9 +367,7 @@ class Affine(bijector.Bijector): if self._is_only_identity_multiplier: # We don't pad in this case and instead let the fldj be applied # via broadcast. - event_size = distribution_util.pick_vector( - math_ops.equal(self._shaper.event_ndims, 0), - [1], array_ops.shape(x))[-1] + event_size = array_ops.shape(x)[-1] event_size = math_ops.cast(event_size, dtype=self._scale.dtype) return math_ops.log(math_ops.abs(self._scale)) * event_size return self.scale.log_abs_determinant() diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/affine_scalar.py b/tensorflow/contrib/distributions/python/ops/bijectors/affine_scalar.py new file mode 100644 index 0000000000000000000000000000000000000000..8adaa54c843d1b243a02967402a37b7c63fabbdf --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/bijectors/affine_scalar.py @@ -0,0 +1,138 @@ +# 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. +# ============================================================================== +"""Affine bijector.""" + +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 check_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops.distributions import bijector + + +__all__ = [ + "AffineScalar", +] + + +class AffineScalar(bijector.Bijector): + """Compute `Y = g(X; shift, scale) = scale * X + shift`. + + Examples: + + ```python + # Y = X + b = AffineScalar() + + # Y = X + shift + b = AffineScalar(shift=[1., 2, 3]) + + # Y = 2 * X + shift + b = AffineScalar( + shift=[1., 2, 3], + scale=2.) + ``` + + """ + + def __init__(self, + shift=None, + scale=None, + validate_args=False, + name="affine_scalar"): + """Instantiates the `AffineScalar` bijector. + + This `Bijector` is initialized with `shift` `Tensor` and `scale` arguments, + giving the forward operation: + + ```none + Y = g(X) = scale * X + shift + ``` + + if `scale` is not specified, then the bijector has the semantics of + `scale = 1.`. Similarly, if `shift` is not specified, then the bijector + has the semantics of `shift = 0.`. + + Args: + shift: Floating-point `Tensor`. If this is set to `None`, no shift is + applied. + scale: Floating-point `Tensor`. If this is set to `None`, no scale is + applied. + validate_args: Python `bool` indicating whether arguments should be + checked for correctness. + name: Python `str` name given to ops managed by this object. + """ + self._graph_parents = [] + self._name = name + self._validate_args = validate_args + + with self._name_scope("init", values=[scale, shift]): + self._shift = shift + self._scale = scale + + if self._shift is not None: + self._shift = ops.convert_to_tensor(shift, name="shift") + + if self._scale is not None: + self._scale = ops.convert_to_tensor(self._scale, name="scale") + if validate_args: + self._scale = control_flow_ops.with_dependencies( + [check_ops.assert_none_equal( + self._scale, + array_ops.zeros([], dtype=self._scale.dtype))], + self._scale) + + super(AffineScalar, self).__init__( + event_ndims=0, + is_constant_jacobian=True, + validate_args=validate_args, + name=name) + + @property + def shift(self): + """The `shift` `Tensor` in `Y = scale @ X + shift`.""" + return self._shift + + @property + def scale(self): + """The `scale` `LinearOperator` in `Y = scale @ X + shift`.""" + return self._scale + + def _forward(self, x): + y = array_ops.identity(x) + if self.scale is not None: + y *= self.scale + if self.shift is not None: + y += self.shift + return y + + def _inverse(self, y): + x = array_ops.identity(y) + if self.shift is not None: + x -= self.shift + if self.scale is not None: + x /= self.scale + return x + + def _forward_log_det_jacobian(self, x): + log_det_jacobian = array_ops.zeros_like(x) + if self.scale is None: + return log_det_jacobian + log_det_jacobian += math_ops.log(math_ops.abs(self.scale)) + return log_det_jacobian diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py b/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..be72ff3081225b9f9fdb6541322b7fc3d4aaa41e --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py @@ -0,0 +1,258 @@ +# 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. +# ============================================================================== +"""Batch Norm bijector.""" + +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.layers import normalization +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn +from tensorflow.python.ops.distributions import bijector + + +__all__ = [ + "BatchNormalization", +] + + +def _undo_batch_normalization(x, + mean, + variance, + offset, + scale, + variance_epsilon, + name=None): + r"""Inverse of tf.nn.batch_normalization. + + Args: + x: Input `Tensor` of arbitrary dimensionality. + mean: A mean `Tensor`. + variance: A variance `Tensor`. + offset: An offset `Tensor`, often denoted `beta` in equations, or + None. If present, will be added to the normalized tensor. + scale: A scale `Tensor`, often denoted `gamma` in equations, or + `None`. If present, the scale is applied to the normalized tensor. + variance_epsilon: A small `float` added to the minibatch `variance` to + prevent dividing by zero. + name: A name for this operation (optional). + + Returns: + batch_unnormalized: The de-normalized, de-scaled, de-offset `Tensor`. + """ + with ops.name_scope( + name, "undo_batchnorm", [x, mean, variance, scale, offset]): + # inv = math_ops.rsqrt(variance + variance_epsilon) + # if scale is not None: + # inv *= scale + # return x * inv + ( + # offset - mean * inv if offset is not None else -mean * inv) + rescale = math_ops.sqrt(variance + variance_epsilon) + if scale is not None: + rescale /= scale + batch_unnormalized = x * rescale + ( + mean - offset * rescale if offset is not None else mean) + return batch_unnormalized + + +class BatchNormalization(bijector.Bijector): + """Compute `Y = g(X) s.t. X = g^-1(Y) = (Y - mean(Y)) / std(Y)`. + + Applies Batch Normalization [1] to samples from a data distribution. This can + be used to stabilize training of normalizing flows [2, 3]. + + When training Deep Neural Networks (DNNs), it is common practice to + normalize or whiten features by shifting them to have zero mean and + scaling them to have unit variance. + + The `inverse()` method of the BatchNorm bijector, which is used in the + log-likelihood computation of data samples, implements the normalization + procedure (shift-and-scale) using the mean and standard deviation of the + current minibatch. + + Conversely, the `forward()` method of the bijector de-normalizes samples (e.g. + `X*std(Y) + mean(Y)` with the running-average mean and standard deviation + computed at training-time. De-normalization is useful for sampling. + + + ```python + + dist = tfd.TransformedDistribution( + distribution=tfd.Normal()), + bijector=tfb.BatchNorm()) + + y = tfd.MultivariateNormalDiag(loc=1., scale=2.).sample(100) # ~ N(1, 2) + x = dist.bijector.inverse(y) # ~ N(0, 1) + y = dist.sample() # ~ N(1, 2) + ``` + + During training time, `BatchNorm.inverse` and `BatchNorm.forward` are not + guaranteed to be inverses of each other because `inverse(y)` uses statistics + of the current minibatch, while `forward(x)` uses running-average statistics + accumulated from training. In other words, + `BatchNorm.inverse(BatchNorm.forward(...))` and + `BatchNorm.forward(BatchNorm.inverse(...))` will be identical when + `training=False` but may be different when `training=True`. + + [1]: "Batch Normalization: Accelerating Deep Network Training by Reducing + Internal Covariate Shift." + Sergey Ioffe, Christian Szegedy. Arxiv. 2015. + https://arxiv.org/abs/1502.03167 + + [2]: "Density Estimation using Real NVP." + Laurent Dinh, Jascha Sohl-Dickstein, Samy Bengio. ICLR. 2017. + https://arxiv.org/abs/1605.08803 + + [3]: "Masked Autoregressive Flow for Density Estimation." + George Papamakarios, Theo Pavlakou, Iain Murray. Arxiv. 2017. + https://arxiv.org/abs/1705.07057 + + """ + + def __init__(self, + batchnorm_layer=None, + training=True, + validate_args=False, + name="batch_normalization"): + """Instantiates the `BatchNorm` bijector. + + Args: + batchnorm_layer: `tf.layers.BatchNormalization` layer object. If `None`, + defaults to + `tf.layers.BatchNormalization(gamma_constraint=nn_ops.relu(x) + 1e-6)`. + This ensures positivity of the scale variable. + + training: If True, updates running-average statistics during call to + `inverse()`. + 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 bn_layer is not an instance of + `tf.layers.BatchNormalization`, or if it is specified with `renorm=True` + or a virtual batch size. + """ + # Scale must be positive. + g_constraint = lambda x: nn.relu(x) + 1e-6 + self.batchnorm = batchnorm_layer or normalization.BatchNormalization( + gamma_constraint=g_constraint) + self._validate_bn_layer(self.batchnorm) + self._training = training + super(BatchNormalization, self).__init__( + validate_args=validate_args, name=name) + + def _validate_bn_layer(self, layer): + """Check for valid BatchNormalization layer. + + Args: + layer: Instance of `tf.layers.BatchNormalization`. + Raises: + ValueError: If batchnorm_layer argument is not an instance of + `tf.layers.BatchNormalization`, or if `batchnorm_layer.renorm=True` or + if `batchnorm_layer.virtual_batch_size` is specified. + """ + if not isinstance(layer, normalization.BatchNormalization): + raise ValueError( + "batchnorm_layer must be an instance of BatchNormalization layer.") + if layer.renorm: + raise ValueError("BatchNorm Bijector does not support renormalization.") + if layer.virtual_batch_size: + raise ValueError( + "BatchNorm Bijector does not support virtual batch sizes.") + + def _get_broadcast_fn(self, x): + # Compute shape to broadcast scale/shift parameters to. + if not x.shape.is_fully_defined(): + raise ValueError("Input must have shape known at graph construction.") + input_shape = np.int32(x.shape.as_list()) + + ndims = len(input_shape) + # event_dims = self._compute_event_dims(x) + reduction_axes = [i for i in range(ndims) if i not in self.batchnorm.axis] + # Broadcasting only necessary for single-axis batch norm where the axis is + # not the last dimension + broadcast_shape = [1] * ndims + broadcast_shape[self.batchnorm.axis[0]] = ( + input_shape[self.batchnorm.axis[0]]) + def _broadcast(v): + if (v is not None and + len(v.get_shape()) != ndims and + reduction_axes != list(range(ndims - 1))): + return array_ops.reshape(v, broadcast_shape) + return v + return _broadcast + + def _normalize(self, y): + return self.batchnorm.apply(y, training=self._training) + + def _de_normalize(self, x): + # Uses the saved statistics. + if not self.batchnorm.built: + input_shape = x.get_shape() + self.batchnorm.build(input_shape) + broadcast_fn = self._get_broadcast_fn(x) + mean = broadcast_fn(self.batchnorm.moving_mean) + variance = broadcast_fn(self.batchnorm.moving_variance) + beta = broadcast_fn(self.batchnorm.beta) if self.batchnorm.center else None + gamma = broadcast_fn(self.batchnorm.gamma) if self.batchnorm.scale else None + return _undo_batch_normalization( + x, mean, variance, beta, gamma, self.batchnorm.epsilon) + + def _forward(self, x): + return self._de_normalize(x) + + def _inverse(self, y): + return self._normalize(y) + + def _forward_log_det_jacobian(self, x): + # Uses saved statistics to compute volume distortion. + return -self._inverse_log_det_jacobian(x, use_saved_statistics=True) + + def _inverse_log_det_jacobian(self, y, use_saved_statistics=False): + if not y.shape.is_fully_defined(): + raise ValueError("Input must have shape known at graph construction.") + input_shape = np.int32(y.shape.as_list()) + + if not self.batchnorm.built: + # Create variables. + self.batchnorm.build(input_shape) + + event_dims = self.batchnorm.axis + reduction_axes = [i for i in range(len(input_shape)) if i not in event_dims] + + if use_saved_statistics or not self._training: + log_variance = math_ops.log( + self.batchnorm.moving_variance + self.batchnorm.epsilon) + else: + # At training-time, ildj is computed from the mean and log-variance across + # the current minibatch. + _, v = nn.moments(y, axes=reduction_axes, keep_dims=True) + log_variance = math_ops.log(v + self.batchnorm.epsilon) + + # `gamma` and `log Var(y)` reductions over event_dims. + # Log(total change in area from gamma term). + log_total_gamma = math_ops.reduce_sum(math_ops.log(self.batchnorm.gamma)) + + # Log(total change in area from log-variance term). + log_total_variance = math_ops.reduce_sum(log_variance) + # The ildj is scalar, as it does not depend on the values of x and are + # constant across minibatch elements. + return log_total_gamma - 0.5 * log_total_variance diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py index cbd60f92a60612c6cf791b2c7708a3310c6e2b6b..43208ff088b469b70ebc08757daac277d4432b37 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py @@ -20,8 +20,6 @@ from __future__ import print_function import numpy as np -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 check_ops from tensorflow.python.ops import control_flow_ops @@ -39,8 +37,6 @@ __all__ = [ class CholeskyOuterProduct(bijector.Bijector): """Compute `g(X) = X @ X.T`; X is lower-triangular, positive-diagonal matrix. - `event_ndims` must be 0 or 2, i.e., scalar or matrix. - Note: the upper-triangular part of X is ignored (whether or not its zero). The surjectivity of g as a map from the set of n x n positive-diagonal @@ -64,46 +60,31 @@ class CholeskyOuterProduct(bijector.Bijector): Examples: ```python - bijector.CholeskyOuterProduct(event_ndims=2).forward(x=[[1., 0], [2, 1]]) + bijector.CholeskyOuterProduct().forward(x=[[1., 0], [2, 1]]) # Result: [[1., 2], [2, 5]], i.e., x @ x.T - bijector.CholeskyOuterProduct(event_ndims=2).inverse(y=[[1., 2], [2, 5]]) + bijector.CholeskyOuterProduct().inverse(y=[[1., 2], [2, 5]]) # Result: [[1., 0], [2, 1]], i.e., cholesky(y). ``` """ - def __init__(self, event_ndims=2, validate_args=False, - name="cholesky_outer_product"): + def __init__(self, validate_args=False, name="cholesky_outer_product"): """Instantiates the `CholeskyOuterProduct` bijector. Args: - event_ndims: `constant` `int32` scalar `Tensor` indicating the number of - dimensions associated with a particular draw from the distribution. Must - be 0 or 2. 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 event_ndims is neither 0 or 2. """ self._graph_parents = [] self._name = name - with self._name_scope("init", values=[event_ndims]): - event_ndims = ops.convert_to_tensor(event_ndims, name="event_ndims") - event_ndims = tensor_util.constant_value(event_ndims) - if event_ndims is None or event_ndims not in [0, 2]: - raise ValueError("`event_ndims` must be a TF constant which is 0 or 2") - self._static_event_ndims = event_ndims super(CholeskyOuterProduct, self).__init__( - event_ndims=event_ndims, + event_ndims=2, validate_args=validate_args, name=name) def _forward(self, x): - if self._static_event_ndims == 0: - return math_ops.square(x) if self.validate_args: is_matrix = check_ops.assert_rank_at_least(x, 2) shape = array_ops.shape(x) @@ -114,11 +95,7 @@ class CholeskyOuterProduct(bijector.Bijector): return math_ops.matmul(x, x, adjoint_b=True) def _inverse(self, y): - return (math_ops.sqrt(y) if self._static_event_ndims == 0 - else linalg_ops.cholesky(y)) - - def _inverse_log_det_jacobian(self, y): - return -self._forward_log_det_jacobian(x=self._inverse(y)) + return linalg_ops.cholesky(y) def _forward_log_det_jacobian(self, x): # Let Y be a symmetric, positive definite matrix and write: @@ -161,13 +138,6 @@ class CholeskyOuterProduct(bijector.Bijector): # Since there is a 2 X[j,j] term for every lower-triangular element of X we # conclude: # |Jac(d vec[Y]/d vec[X])| = 2^p prod_{j=0}^{p-1} X[j,j]^{p-j}. - if self._static_event_ndims == 0: - if self.validate_args: - is_positive = check_ops.assert_positive( - x, message="All elements must be positive.") - x = control_flow_ops.with_dependencies([is_positive], x) - return np.log(2.) + math_ops.log(x) - diag = array_ops.matrix_diag_part(x) # We now ensure diag is columnar. Eg, if `diag = [1, 2, 3]` then the output diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py b/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py new file mode 100644 index 0000000000000000000000000000000000000000..f5de052c9ed18b1ebf4c174aeea3a951b1ddcd9d --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py @@ -0,0 +1,153 @@ +# 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. +# ============================================================================== +"""Kumaraswamy bijector.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +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 check_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops.distributions import bijector + +__all__ = [ + "Kumaraswamy", +] + + +class Kumaraswamy(bijector.Bijector): + """Compute `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a), X in [0, 1]`. + + This bijector maps inputs from `[0, 1]` to [0, 1]`. The inverse of the + bijector applied to a uniform random variable `X ~ U(0, 1) gives back a + random variable with the [Kumaraswamy distribution]( + https://en.wikipedia.org/wiki/Kumaraswamy_distribution): + + ```none + Y ~ Kumaraswamy(a, b) + pdf(y; a, b, 0 <= y <= 1) = a * b * y ** (a - 1) * (1 - y**a) ** (b - 1) + ``` + """ + + def __init__(self, + concentration1=None, + concentration0=None, + event_ndims=0, + validate_args=False, + name="kumaraswamy"): + """Instantiates the `Kumaraswamy` bijector. + + Args: + concentration1: Python `float` scalar indicating the transform power, + i.e., `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a)` where `a` is + `concentration1`. + concentration0: Python `float` scalar indicating the transform power, + i.e., `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a)` where `b` is + `concentration0`. + event_ndims: Python scalar indicating the number of dimensions associated + with a particular draw from the distribution. Currently only zero is + supported. + 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 `event_ndims` is not zero. + """ + self._graph_parents = [] + self._name = name + self._validate_args = validate_args + + event_ndims = ops.convert_to_tensor(event_ndims, name="event_ndims") + event_ndims_const = tensor_util.constant_value(event_ndims) + if event_ndims_const is not None and event_ndims_const not in (0,): + raise ValueError("event_ndims(%s) was not 0" % event_ndims_const) + else: + if validate_args: + event_ndims = control_flow_ops.with_dependencies( + [check_ops.assert_equal( + event_ndims, 0, message="event_ndims was not 0")], + event_ndims) + + with self._name_scope("init", values=[concentration1, concentration0]): + concentration1 = self._maybe_assert_valid_concentration( + ops.convert_to_tensor(concentration1, name="concentration1"), + validate_args=validate_args) + concentration0 = self._maybe_assert_valid_concentration( + ops.convert_to_tensor(concentration0, name="concentration0"), + validate_args=validate_args) + + self._concentration1 = concentration1 + self._concentration0 = concentration0 + super(Kumaraswamy, self).__init__( + event_ndims=0, + validate_args=validate_args, + name=name) + + @property + def concentration1(self): + """The `a` in: `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a)`.""" + return self._concentration1 + + @property + def concentration0(self): + """The `b` in: `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a)`.""" + return self._concentration0 + + def _forward(self, x): + x = self._maybe_assert_valid(x) + return math_ops.exp( + math_ops.log1p(-math_ops.exp(math_ops.log1p(-x) / self.concentration0)) + / self.concentration1) + + def _inverse(self, y): + y = self._maybe_assert_valid(y) + return math_ops.exp(math_ops.log1p( + -(1 - y**self.concentration1)**self.concentration0)) + + def _inverse_log_det_jacobian(self, y): + y = self._maybe_assert_valid(y) + event_dims = self._event_dims_tensor(y) + return math_ops.reduce_sum( + math_ops.log(self.concentration1) + math_ops.log(self.concentration0) + + (self.concentration1 - 1) * math_ops.log(y) + + (self.concentration0 - 1) * math_ops.log1p(-y**self.concentration1), + axis=event_dims) + + def _maybe_assert_valid_concentration(self, concentration, validate_args): + """Checks the validity of a concentration parameter.""" + if not validate_args: + return concentration + return control_flow_ops.with_dependencies([ + check_ops.assert_positive( + concentration, + message="Concentration parameter must be positive."), + ], concentration) + + def _maybe_assert_valid(self, x): + if not self.validate_args: + return x + return control_flow_ops.with_dependencies([ + check_ops.assert_non_negative( + x, + message="sample must be non-negative"), + check_ops.assert_less_equal( + x, array_ops.ones([], self.concentration0.dtype), + message="sample must be no larger than `1`."), + ], x) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py b/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py index a9dcce6c526600f3b26c6bceb730417000917ce7..24add40445c60db533aac6d0c8eb537774895c65 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py @@ -18,8 +18,6 @@ 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 distribution_util from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -161,33 +159,16 @@ class SoftmaxCentered(bijector.Bijector): # x[i] = log(exp(x[i])) - log(y[end]) - log(normalization) # = log(exp(x[i])/normalization) - log(y[end]) # = log(y[i]) - log(y[end]) - shape = (np.asarray(y.shape.as_list(), dtype=np.int32) - if y.shape.is_fully_defined() - else array_ops.shape(y, name="shape")) - ndims = distribution_util.prefer_static_rank(y) # Do this first to make sure CSE catches that it'll happen again in # _inverse_log_det_jacobian. x = math_ops.log(y) - # We now extract the last coordinate of the rightmost dimension. - # Our trick is to slice from [0,0,...,shape[-1]-1] to shape[:-1]+[1]. - begin = array_ops.one_hot(indices=ndims-1, - depth=ndims, - on_value=shape[-1]-np.array(1, dtype=shape.dtype), - dtype=shape.dtype) - size = array_ops.concat([shape[:-1], np.asarray([1], dtype=shape.dtype)], 0) - log_normalization = -array_ops.strided_slice(x, begin, begin + size) - - # Here we slice out all but the last coordinate; see above for idea. - begin = array_ops.zeros_like(shape) - size = array_ops.concat([shape[:-1], [shape[-1] - 1]], 0) - x = array_ops.strided_slice(x, begin, begin + size) - - x += log_normalization + log_normalization = (-x[..., -1])[..., array_ops.newaxis] + x = x[..., :-1] + log_normalization if self._static_event_ndims == 0: - x = array_ops.squeeze(x, squeeze_dims=[ndims-1]) + x = array_ops.squeeze(x, squeeze_dims=-1) # Set shape hints. if y.shape.ndims is not None: diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/square.py b/tensorflow/contrib/distributions/python/ops/bijectors/square.py new file mode 100644 index 0000000000000000000000000000000000000000..2831a92df8e0ad2bf681f13533cdb6f5d2089a57 --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/bijectors/square.py @@ -0,0 +1,84 @@ +# 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. +# ============================================================================== +"""Square bijector.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +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.distributions import bijector + + +__all__ = [ + "Square", +] + + +class Square(bijector.Bijector): + """Compute `g(X) = X^2`; X is a positive real number. + + g is a bijection between the non-negative real numbers (R_+) and the + non-negative real numbers. + + Examples: + + ```python + bijector.Square().forward(x=[[1., 0], [2, 1]]) + # Result: [[1., 0], [4, 1]], i.e., x^2 + + bijector.Square().inverse(y=[[1., 4], [9, 1]]) + # Result: [[1., 2], [3, 1]], i.e., sqrt(y). + ``` + + """ + + def __init__(self, validate_args=False, name="square"): + """Instantiates the `Square` bijector. + + Args: + validate_args: Python `bool` indicating whether arguments should be + checked for correctness. + name: Python `str` name given to ops managed by this object. + """ + self._name = name + super(Square, self).__init__( + event_ndims=0, + validate_args=validate_args, + name=name) + + def _forward(self, x): + x = self._maybe_assert_valid(x) + return math_ops.square(x) + + def _inverse(self, y): + y = self._maybe_assert_valid(y) + return math_ops.sqrt(y) + + def _forward_log_det_jacobian(self, x): + x = self._maybe_assert_valid(x) + return np.log(2.) + math_ops.log(x) + + def _maybe_assert_valid(self, t): + if not self.validate_args: + return t + is_valid = check_ops.assert_non_negative( + t, message="All elements must be non-negative.") + return control_flow_ops.with_dependencies([is_valid], t) + diff --git a/tensorflow/contrib/distributions/python/ops/chi2.py b/tensorflow/contrib/distributions/python/ops/chi2.py index bdd5571c966a74e58e4f9f8eed2628f131a1b92e..e610f469e5d5f446b75c734cc39811de30a8cb9a 100644 --- a/tensorflow/contrib/distributions/python/ops/chi2.py +++ b/tensorflow/contrib/distributions/python/ops/chi2.py @@ -21,6 +21,8 @@ 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 check_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops.distributions import gamma @@ -87,7 +89,11 @@ class Chi2(gamma.Gamma): # allow_nan_stats=True # through to the parent class results in unnecessary asserts. with ops.name_scope(name, values=[df]): - self._df = ops.convert_to_tensor(df, name="df") + with ops.control_dependencies([ + check_ops.assert_positive(df), + ] if validate_args else []): + self._df = array_ops.identity(df, name="df") + super(Chi2, self).__init__( concentration=0.5 * self._df, rate=constant_op.constant(0.5, dtype=self._df.dtype), diff --git a/tensorflow/contrib/distributions/python/ops/gumbel.py b/tensorflow/contrib/distributions/python/ops/gumbel.py index d0efaefb8e78ddf4436e9e5a112d2c1cdddaf3b5..8d05ad6b8032fb8bada99389959091fb1c28beda 100644 --- a/tensorflow/contrib/distributions/python/ops/gumbel.py +++ b/tensorflow/contrib/distributions/python/ops/gumbel.py @@ -190,9 +190,6 @@ class _Gumbel(distribution.Distribution): def _log_prob(self, x): return self._log_unnormalized_prob(x) - self._log_normalization() - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - def _log_cdf(self, x): return -math_ops.exp(-self._z(x)) diff --git a/tensorflow/contrib/distributions/python/ops/independent.py b/tensorflow/contrib/distributions/python/ops/independent.py index cbce005013281ff3c58c94d525d5ce7a865d725a..7dcb3e3ac4db1855adacb7ec0fa8554c45d9c859 100644 --- a/tensorflow/contrib/distributions/python/ops/independent.py +++ b/tensorflow/contrib/distributions/python/ops/independent.py @@ -28,6 +28,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops.distributions import distribution as distribution_lib +from tensorflow.python.ops.distributions import kullback_leibler class Independent(distribution_lib.Distribution): @@ -254,3 +255,58 @@ class Independent(distribution_lib.Distribution): else: which_maximum = np.maximum return which_maximum(0, ndims - 1) + + +@kullback_leibler.RegisterKL(Independent, Independent) +def _kl_independent(a, b, name="kl_independent"): + """Batched KL divergence `KL(a || b)` for Independent distributions. + + We can leverage the fact that + ``` + KL(Independent(a) || Independent(b)) = sum(KL(a || b)) + ``` + where the sum is over the `reinterpreted_batch_ndims`. + + Args: + a: Instance of `Independent`. + b: Instance of `Independent`. + name: (optional) name to use for created ops. Default "kl_independent". + + Returns: + Batchwise `KL(a || b)`. + + Raises: + ValueError: If the event space for `a` and `b`, or their underlying + distributions don't match. + """ + p = a.distribution + q = b.distribution + + # The KL between any two (non)-batched distributions is a scalar. + # Given that the KL between two factored distributions is the sum, i.e. + # KL(p1(x)p2(y) || q1(x)q2(y)) = KL(p1 || q1) + KL(q1 || q2), we compute + # KL(p || q) and do a `reduce_sum` on the reinterpreted batch dimensions. + if a.event_shape.is_fully_defined() and b.event_shape.is_fully_defined(): + if a.event_shape == b.event_shape: + if p.event_shape == q.event_shape: + num_reduce_dims = a.event_shape.ndims - p.event_shape.ndims + reduce_dims = [-i - 1 for i in range(0, num_reduce_dims)] + + return math_ops.reduce_sum( + kullback_leibler.kl_divergence(p, q, name=name), axis=reduce_dims) + else: + raise NotImplementedError("KL between Independents with different " + "event shapes not supported.") + else: + raise ValueError("Event shapes do not match.") + else: + with ops.control_dependencies([ + check_ops.assert_equal(a.event_shape_tensor(), b.event_shape_tensor()), + check_ops.assert_equal(p.event_shape_tensor(), q.event_shape_tensor()) + ]): + num_reduce_dims = ( + array_ops.shape(a.event_shape_tensor()[0]) - + array_ops.shape(p.event_shape_tensor()[0])) + reduce_dims = math_ops.range(-num_reduce_dims - 1, -1, 1) + return math_ops.reduce_sum( + kullback_leibler.kl_divergence(p, q, name=name), axis=reduce_dims) diff --git a/tensorflow/contrib/distributions/python/ops/inverse_gamma.py b/tensorflow/contrib/distributions/python/ops/inverse_gamma.py index ee4d86867d48b20e97757bcec57d452085814b80..51ac61dcf640ca89f22c47127bda71316a179ca4 100644 --- a/tensorflow/contrib/distributions/python/ops/inverse_gamma.py +++ b/tensorflow/contrib/distributions/python/ops/inverse_gamma.py @@ -192,12 +192,6 @@ class InverseGamma(distribution.Distribution): def _log_prob(self, x): return self._log_unnormalized_prob(x) - self._log_normalization() - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - - def _log_cdf(self, x): - return math_ops.log(self._cdf(x)) - def _cdf(self, x): x = self._maybe_assert_valid_sample(x) # Note that igammac returns the upper regularized incomplete gamma diff --git a/tensorflow/contrib/distributions/python/ops/kumaraswamy.py b/tensorflow/contrib/distributions/python/ops/kumaraswamy.py new file mode 100644 index 0000000000000000000000000000000000000000..120b38db3cf72e8fce56a7e9293cdf25e75784e2 --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/kumaraswamy.py @@ -0,0 +1,235 @@ +# 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.contrib.distributions.python.ops import bijectors +from tensorflow.contrib.distributions.python.ops import distribution_util +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 control_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import special_math_ops +from tensorflow.python.ops.distributions import distribution +from tensorflow.python.ops.distributions import transformed_distribution +from tensorflow.python.ops.distributions import uniform +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(transformed_distribution.TransformedDistribution): + """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. + """ + concentration1 = ops.convert_to_tensor( + concentration1, name="concentration1") + concentration0 = ops.convert_to_tensor( + concentration0, name="concentration0") + super(Kumaraswamy, self).__init__( + distribution=uniform.Uniform( + low=array_ops.zeros([], dtype=concentration1.dtype), + high=array_ops.ones([], dtype=concentration1.dtype), + allow_nan_stats=allow_nan_stats), + bijector=bijectors.Kumaraswamy( + concentration1=concentration1, concentration0=concentration0, + validate_args=validate_args), + batch_shape=distribution_util.get_broadcast_shape( + concentration1, concentration0), + name=name) + self._reparameterization_type = distribution.FULLY_REPARAMETERIZED + + @property + def concentration1(self): + """Concentration parameter associated with a `1` outcome.""" + return self.bijector.concentration1 + + @property + def concentration0(self): + """Concentration parameter associated with a `0` outcome.""" + return self.bijector.concentration0 + + 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.""" + total_concentration = self.concentration1 + self.concentration0 + expanded_concentration1 = array_ops.ones_like( + total_concentration, dtype=self.dtype) * self.concentration1 + expanded_concentration0 = array_ops.ones_like( + 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.concentration1.dtype), + self.concentration1, + message="Mode undefined for concentration1 <= 1."), + check_ops.assert_less( + array_ops.ones([], dtype=self.concentration0.dtype), + self.concentration0, + message="Mode undefined for concentration0 <= 1.") + ], mode) diff --git a/tensorflow/contrib/distributions/python/ops/logistic.py b/tensorflow/contrib/distributions/python/ops/logistic.py index 473677f8d91b184e029f345bb05f5c5d63df7a40..68e6bca5a554b29a450911073eb5c4fe55f313c6 100644 --- a/tensorflow/contrib/distributions/python/ops/logistic.py +++ b/tensorflow/contrib/distributions/python/ops/logistic.py @@ -185,9 +185,6 @@ class Logistic(distribution.Distribution): def _log_prob(self, x): return self._log_unnormalized_prob(x) - self._log_normalization() - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - def _log_cdf(self, x): return -nn_ops.softplus(-self._z(x)) diff --git a/tensorflow/contrib/distributions/python/ops/onehot_categorical.py b/tensorflow/contrib/distributions/python/ops/onehot_categorical.py index b76cebf79fad09ebec68f2459c6fe80794ea81c0..46c2cc8b7a8c536a90176fbb2b2d52fed61e4705 100644 --- a/tensorflow/contrib/distributions/python/ops/onehot_categorical.py +++ b/tensorflow/contrib/distributions/python/ops/onehot_categorical.py @@ -203,9 +203,6 @@ class OneHotCategorical(distribution.Distribution): ret = array_ops.reshape(ret, logits_shape) return ret - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - def _entropy(self): return -math_ops.reduce_sum( nn_ops.log_softmax(self.logits) * self.probs, axis=-1) diff --git a/tensorflow/contrib/distributions/python/ops/poisson.py b/tensorflow/contrib/distributions/python/ops/poisson.py index e967dcc90d0712ffc346fb61ee67c44a6d9207cb..02e97c0a2fd004c4fa9382d5367af9f5b034a869 100644 --- a/tensorflow/contrib/distributions/python/ops/poisson.py +++ b/tensorflow/contrib/distributions/python/ops/poisson.py @@ -35,9 +35,15 @@ __all__ = [ _poisson_sample_note = """ -Note that the input value must be a non-negative floating point tensor with -dtype `dtype` and whose shape can be broadcast with `self.rate`. `x` is only -legal if it is non-negative and its components are equal to integer values. +The Poisson distribution is technically only defined for non-negative integer +values. When `validate_args=False`, non-integral inputs trigger an assertion. + +When `validate_args=False` calculations are otherwise unchanged despite +integral or non-integral inputs. + +When `validate_args=False`, evaluating the pmf at non-integral values, +corresponds to evaluations of an unnormalized distribution, that does not +correspond to evaluations of the cdf. """ @@ -150,10 +156,6 @@ class Poisson(distribution.Distribution): def _cdf(self, x): if self.validate_args: x = distribution_util.embed_check_nonnegative_integer_form(x) - else: - # Whether or not x is integer-form, the following is well-defined. - # However, scipy takes the floor, so we do too. - x = math_ops.floor(x) return math_ops.igammac(1. + x, self.rate) def _log_normalization(self): @@ -162,9 +164,6 @@ class Poisson(distribution.Distribution): def _log_unnormalized_prob(self, x): if self.validate_args: x = distribution_util.embed_check_nonnegative_integer_form(x) - else: - # For consistency with cdf, we take the floor. - x = math_ops.floor(x) return x * self.log_rate - math_ops.lgamma(1. + x) def _mean(self): diff --git a/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py b/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py index b6becfa9fc93f189a1a7bf7b2a7af8dc1f2e9720..ff33f327c7a77597e516208cacad8c4aed65d1c9 100644 --- a/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py +++ b/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py @@ -278,16 +278,13 @@ 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 ret = array_ops.reshape(log_prob, logits_shape) return ret - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - def _assert_valid_sample(self, x): if not self.validate_args: return x diff --git a/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py b/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py index c4b8f055b7fbc3f0835b503eddd7617610326d8c..0d8a1926913766da374cb65767dccfa28bf75579 100644 --- a/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py +++ b/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py @@ -174,13 +174,12 @@ class SinhArcsinh(transformed_distribution.TransformedDistribution): skewness=skewness.dtype.as_numpy_dtype(0.), tailweight=tailweight, event_ndims=0) - # Make the Affine bijector, Z --> loc + scale * Z (2 / F_0(2)) + # Make the AffineScalar bijector, Z --> loc + scale * Z (2 / F_0(2)) c = 2 * scale / f_noskew.forward(ops.convert_to_tensor(2, dtype=dtype)) - affine = bijectors.Affine( + affine = bijectors.AffineScalar( shift=loc, - scale_identity_multiplier=c, - validate_args=validate_args, - event_ndims=0) + scale=c, + validate_args=validate_args) bijector = bijectors.Chain([affine, f]) diff --git a/tensorflow/contrib/distributions/python/ops/statistical_testing.py b/tensorflow/contrib/distributions/python/ops/statistical_testing.py new file mode 100644 index 0000000000000000000000000000000000000000..d66c34cc1a45cc09da5138a5f72ae3817690db49 --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/statistical_testing.py @@ -0,0 +1,728 @@ +# 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. +# ============================================================================== +"""Statistical test assertions calibrated for their error rates.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import itertools + +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 control_flow_ops +from tensorflow.python.ops import gen_math_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn_ops + +__all__ = [ + "true_mean_confidence_interval_by_dkwm", + "assert_true_mean_equal_by_dkwm", + "min_discrepancy_of_true_means_detectable_by_dkwm", + "min_num_samples_for_dkwm_mean_test", + "assert_true_mean_equal_by_dkwm_two_sample", + "min_discrepancy_of_true_means_detectable_by_dkwm_two_sample", + "min_num_samples_for_dkwm_mean_two_sample_test", +] + + +def _batch_sort_vector(x, ascending=True, name=None): + with ops.name_scope(name, "sort_each_row", [x]): + x = ops.convert_to_tensor(x, name="x") + n = array_ops.shape(x)[-1] + if ascending: + y, _ = nn_ops.top_k(-x, k=n, sorted=True) + y = -y + else: + y, _ = nn_ops.top_k(x, k=n, sorted=True) + y.set_shape(x.shape) + return y + + +def _do_maximum_mean(samples, envelope, high, name=None): + """Common code between maximum_mean and minimum_mean.""" + with ops.name_scope(name, "do_maximum_mean", [samples, envelope, high]): + n = array_ops.rank(samples) + # Move the batch dimension of `samples` to the rightmost position, + # where the _batch_sort_vector function wants it. + perm = array_ops.concat([math_ops.range(1, n), [0]], axis=0) + samples = array_ops.transpose(samples, perm) + + samples = _batch_sort_vector(samples) + batch_shape = array_ops.shape(samples)[:-1] + n = array_ops.shape(samples)[-1] + step = 1. / math_ops.cast(n, dtype=samples.dtype.base_dtype) + + def _loop_body(iter_, total, to_skip): + total = array_ops.where( + step <= to_skip, + total, + array_ops.where( + to_skip > 0., + total + (step - to_skip) * samples[..., iter_], + total + step * samples[..., iter_])) + to_skip = array_ops.where(step <= to_skip, to_skip - step, 0.) + return [iter_ + 1, total, to_skip] + + _, total, _ = control_flow_ops.while_loop( + cond=lambda iter_, *args: iter_ < n, + body=_loop_body, + loop_vars=[ + 0, + array_ops.zeros(batch_shape, dtype=samples.dtype.base_dtype), + envelope, # to_skip + ]) + + return total + envelope * high + + +def _maximum_mean(samples, envelope, high, name=None): + """Returns a stochastic upper bound on the mean of a scalar distribution. + + The idea is that if the true CDF is within an `eps`-envelope of the + empirical CDF of the samples, and the support is bounded above, then + the mean is bounded above as well. In symbols, + + ```none + sup_x(|F_n(x) - F(x)|) < eps + ``` + + The 0th dimension of `samples` is interpreted as independent and + identically distributed samples. The remaining dimensions are + broadcast together with `envelope` and `high`, and operated on + separately. + + Args: + samples: Floating-point tensor of samples from the distribution(s) + of interest. Entries are assumed IID across the 0th dimension. + The other dimensions must broadcast with `envelope` and `high`. + envelope: Floating-point tensor of sizes of admissible CDF + envelopes (i.e., the `eps` above). + high: Floating-point tensor of upper bounds on the distributions' + supports. + name: A name for this operation (optional). + + Returns: + bound: Floating-point tensor of upper bounds on the true means. + + Raises: + InvalidArgumentError: If some `sample` is found to be larger than + the corresponding `high`. + """ + with ops.name_scope(name, "maximum_mean", [samples, envelope, high]): + samples = ops.convert_to_tensor(samples, name="samples") + envelope = ops.convert_to_tensor(envelope, name="envelope") + high = ops.convert_to_tensor(high, name="high") + + xmax = math_ops.reduce_max(samples, axis=[-1]) + msg = "Given sample maximum value exceeds expectations" + check_op = check_ops.assert_less_equal(xmax, high, message=msg) + with ops.control_dependencies([check_op]): + return array_ops.identity(_do_maximum_mean(samples, envelope, high)) + + +def _minimum_mean(samples, envelope, low, name=None): + """Returns a stochastic lower bound on the mean of a scalar distribution. + + The idea is that if the true CDF is within an `eps`-envelope of the + empirical CDF of the samples, and the support is bounded below, then + the mean is bounded below as well. In symbols, + + ```none + sup_x(|F_n(x) - F(x)|) < eps + ``` + + The 0th dimension of `samples` is interpreted as independent and + identically distributed samples. The remaining dimensions are + broadcast together with `envelope` and `low`, and operated on + separately. + + Args: + samples: Floating-point tensor of samples from the distribution(s) + of interest. Entries are assumed IID across the 0th dimension. + The other dimensions must broadcast with `envelope` and `low`. + envelope: Floating-point tensor of sizes of admissible CDF + envelopes (i.e., the `eps` above). + low: Floating-point tensor of lower bounds on the distributions' + supports. + name: A name for this operation (optional). + + Returns: + bound: Floating-point tensor of lower bounds on the true means. + + Raises: + InvalidArgumentError: If some `sample` is found to be smaller than + the corresponding `low`. + """ + with ops.name_scope(name, "minimum_mean", [samples, envelope, low]): + samples = ops.convert_to_tensor(samples, name="samples") + envelope = ops.convert_to_tensor(envelope, name="envelope") + low = ops.convert_to_tensor(low, name="low") + + xmin = math_ops.reduce_min(samples, axis=[-1]) + msg = "Given sample minimum value falls below expectations" + check_op = check_ops.assert_greater_equal(xmin, low, message=msg) + with ops.control_dependencies([check_op]): + return - _do_maximum_mean(-samples, envelope, -low) + + +def _dkwm_cdf_envelope(n, error_rate, name=None): + """Computes the CDF envelope that the DKWM inequality licenses. + + The [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval) + gives a stochastic bound on the distance between the true cumulative + distribution function (CDF) of any distribution and its empirical + CDF. To wit, for `n` iid samples from any distribution with CDF F, + + ```none + P(sup_x |F_n(x) - F(x)| > eps) < 2exp(-2n eps^2) + ``` + + This function computes the envelope size `eps` as a function of the + number of samples `n` and the desired limit on the left-hand + probability above. + + Args: + n: Tensor of numbers of samples drawn. + error_rate: Floating-point tensor of admissible rates of mistakes. + name: A name for this operation (optional). + + Returns: + eps: Tensor of maximum distances the true CDF can be from the + empirical CDF. This scales as `O(sqrt(-log(error_rate)))` and + as `O(1 / sqrt(n))`. The shape is the broadcast of `n` and + `error_rate`. + """ + with ops.name_scope(name, "dkwm_cdf_envelope", [n, error_rate]): + n = math_ops.cast(n, dtype=error_rate.dtype) + return math_ops.sqrt(-gen_math_ops.log(error_rate / 2.) / (2. * n)) + + +def _check_shape_dominates(tensor, tensors): + """Check that broadcasting `tensor` against `tensors` does not expand it. + + Why? Because I want to be very sure that the samples tensor is not + accidentally enlarged by broadcasting against tensors that are + supposed to be describing the distribution(s) sampled from, lest the + sample counts end up inflated. + + Args: + tensor: A Tensor whose shape is to be protected against broadcasting. + tensors: A list of Tensors to check + + Returns: + tensor: `tf.identity(tensor)` with control dependencies attached; + be sure to use that downstream. + """ + def check(t): + target = array_ops.shape(tensor)[1:] + result = array_ops.broadcast_dynamic_shape(target, array_ops.shape(t)) + # This rank check ensures that I don't get a wrong answer from the + # _shapes_ broadcasting against each other. + gt = check_ops.assert_greater(array_ops.rank(target), array_ops.rank(t)) + eq = check_ops.assert_equal(target, result) + return gt, eq + checks = list(itertools.chain(*[check(t) for t in tensors])) + with ops.control_dependencies(checks): + return array_ops.identity(array_ops.identity(tensor)) + + +def true_mean_confidence_interval_by_dkwm( + samples, low, high, error_rate=1e-6, name=None): + """Computes a confidence interval for the mean of a scalar distribution. + + In batch mode, computes confidence intervals for all distributions + in the batch (which need not be identically distributed). + + Relies on the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval). + + The probability (over the randomness of drawing the given samples) + that any true mean is outside the corresponding returned interval is + no more than the given `error_rate`. The size of the intervals + scale as + `O(1 / sqrt(#samples))`, as `O(high - low)`, and as `O(-log(error_rate))`. + + Note that `error_rate` is a total error rate for all the confidence + intervals in the batch. As such, if the batch is nontrivial, the + error rate is not broadcast but divided (evenly) among the batch + members. + + Args: + samples: Floating-point tensor of samples from the distribution(s) + of interest. Entries are assumed IID across the 0th dimension. + The other dimensions must broadcast with `low` and `high`. + low: Floating-point tensor of lower bounds on the distributions' + supports. + high: Floating-point tensor of upper bounds on the distributions' + supports. + error_rate: *Scalar* admissible total rate of mistakes. + name: A name for this operation (optional). + + Returns: + low: A floating-point tensor of stochastic lower bounds on the true means. + high: A floating-point tensor of stochastic upper bounds on the true means. + """ + with ops.name_scope( + name, "true_mean_confidence_interval_by_dkwm", + [samples, low, high, error_rate]): + samples = ops.convert_to_tensor(samples, name="samples") + low = ops.convert_to_tensor(low, name="low") + high = ops.convert_to_tensor(high, name="high") + error_rate = ops.convert_to_tensor(error_rate, name="error_rate") + samples = _check_shape_dominates(samples, [low, high]) + check_ops.assert_scalar(error_rate) # Static shape + error_rate = _itemwise_error_rate(error_rate, [low, high], samples) + n = array_ops.shape(samples)[0] + envelope = _dkwm_cdf_envelope(n, error_rate) + min_mean = _minimum_mean(samples, envelope, low) + max_mean = _maximum_mean(samples, envelope, high) + return min_mean, max_mean + + +def _itemwise_error_rate( + total_error_rate, param_tensors, sample_tensor=None, name=None): + with ops.name_scope( + name, "itemwise_error_rate", + [total_error_rate, param_tensors, sample_tensor]): + result_shape = [1] + for p_tensor in param_tensors: + result_shape = array_ops.broadcast_dynamic_shape( + array_ops.shape(p_tensor), result_shape) + if sample_tensor is not None: + result_shape = array_ops.broadcast_dynamic_shape( + array_ops.shape(sample_tensor)[1:], result_shape) + num_items = math_ops.reduce_prod(result_shape) + return total_error_rate / math_ops.cast( + num_items, dtype=total_error_rate.dtype) + + +def assert_true_mean_equal_by_dkwm( + samples, low, high, expected, false_fail_rate=1e-6, name=None): + """Asserts the mean of the given distribution is as expected. + + More precisely, fails if there is enough evidence (using the + [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval)) + that the true mean of some distribution from which the given samples are + drawn is _not_ the given expected mean with statistical significance + `false_fail_rate` or stronger, otherwise passes. If you also want to + check that you are gathering enough evidence that a pass is not + spurious, see `min_num_samples_for_dkwm_mean_test` and + `min_discrepancy_of_true_means_detectable_by_dkwm`. + + Note that `false_fail_rate` is a total false failure rate for all + the assertions in the batch. As such, if the batch is nontrivial, + the assertion will insist on stronger evidence to fail any one member. + + Args: + samples: Floating-point tensor of samples from the distribution(s) + of interest. Entries are assumed IID across the 0th dimension. + The other dimensions must broadcast with `low` and `high`. + low: Floating-point tensor of lower bounds on the distributions' + supports. + high: Floating-point tensor of upper bounds on the distributions' + supports. + expected: Floating-point tensor of expected true means. + false_fail_rate: *Scalar* admissible total rate of mistakes. + name: A name for this operation (optional). + + Returns: + check: Op that raises `InvalidArgumentError` if any expected mean is + outside the corresponding confidence interval. + """ + with ops.name_scope( + name, "assert_true_mean_equal_by_dkwm", + [samples, low, high, expected, false_fail_rate]): + samples = ops.convert_to_tensor(samples, name="samples") + low = ops.convert_to_tensor(low, name="low") + high = ops.convert_to_tensor(high, name="high") + expected = ops.convert_to_tensor(expected, name="expected") + false_fail_rate = ops.convert_to_tensor( + false_fail_rate, name="false_fail_rate") + samples = _check_shape_dominates(samples, [low, high, expected]) + min_mean, max_mean = true_mean_confidence_interval_by_dkwm( + samples, low, high, error_rate=false_fail_rate) + less_op = check_ops.assert_less( + min_mean, expected, message="Mean confidence interval too high") + with ops.control_dependencies([less_op]): + return check_ops.assert_greater( + max_mean, expected, message="Mean confidence interval too low") + + +def min_discrepancy_of_true_means_detectable_by_dkwm( + n, low, high, false_fail_rate, false_pass_rate, name=None): + """Returns the minimum mean discrepancy that a DKWM-based test can detect. + + DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval). + + Note that `false_fail_rate` is a total false failure rate for all + the tests in the batch. As such, if the batch is nontrivial, each + member will demand more samples. The `false_pass_rate` is also + interpreted as a total, but is treated asymmetrically: If each test + in the batch detects its corresponding discrepancy with probability + at least `1 - false_pass_rate`, then running all those tests and + failing if any one fails will jointly detect all those discrepancies + with the same `false_pass_rate`. + + Args: + n: Tensor of numbers of samples to be drawn from the distributions + of interest. + low: Floating-point tensor of lower bounds on the distributions' + supports. + high: Floating-point tensor of upper bounds on the distributions' + supports. + false_fail_rate: *Scalar* admissible total rate of false failures. + false_pass_rate: *Scalar* admissible rate of false passes. + name: A name for this operation (optional). + + Returns: + discr: Tensor of lower bounds on the distances between true + means detectable by a DKWM-based test. + + For each batch member `i`, of `K` total, drawing `n[i]` samples from + some scalar distribution supported on `[low[i], high[i]]` is enough + to detect a difference in means of size `discr[i]` or more. + Specifically, we guarantee that (a) if the true mean is the expected + mean, `assert_true_mean_equal_by_dkwm` will fail with probability at + most `false_fail_rate / K` (which amounts to `false_fail_rate` if + applied to the whole batch at once), and (b) if the true mean + differs from the expected mean by at least `discr[i]`, + `assert_true_mean_equal_by_dkwm` will pass with probability at most + `false_pass_rate`. + + The detectable discrepancy scales as + + - `O(high[i] - low[i])`, + - `O(1 / sqrt(n[i]))`, + - `O(-log(false_fail_rate/K))`, and + - `O(-log(false_pass_rate))`. + """ + with ops.name_scope( + name, "min_discrepancy_of_true_means_detectable_by_dkwm", + [n, low, high, false_fail_rate, false_pass_rate]): + n = ops.convert_to_tensor(n, name="n") + low = ops.convert_to_tensor(low, name="low") + high = ops.convert_to_tensor(high, name="high") + false_fail_rate = ops.convert_to_tensor( + false_fail_rate, name="false_fail_rate") + false_pass_rate = ops.convert_to_tensor( + false_pass_rate, name="false_pass_rate") + # Algorithm: Assume a true CDF F. The DKWM inequality gives a + # stochastic bound on how far the observed empirical CDF F_n can be. + # Then, using the DKWM inequality again gives a stochastic bound on + # the farthest candidate true CDF F' that + # true_mean_confidence_interval_by_dkwm might consider. At worst, these + # errors may go in the same direction, so the distance between F and + # F' is bounded by the sum. + # On batching: false fail rates sum, so I need to reduce + # the input to account for the batching. False pass rates + # max, so I don't. + sampling_envelope = _dkwm_cdf_envelope(n, false_pass_rate) + false_fail_rate = _itemwise_error_rate(false_fail_rate, [n, low, high]) + analysis_envelope = _dkwm_cdf_envelope(n, false_fail_rate) + return (high - low) * (sampling_envelope + analysis_envelope) + + +def min_num_samples_for_dkwm_mean_test( + discrepancy, low, high, + false_fail_rate=1e-6, false_pass_rate=1e-6, name=None): + """Returns how many samples suffice for a one-sample DKWM mean test. + + To wit, returns an upper bound on the number of samples necessary to + guarantee detecting a mean difference of at least the given + `discrepancy`, with the given `false_fail_rate` and `false_pass_rate`, + using the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval) + on a scalar distribution supported on `[low, high]`. + + Args: + discrepancy: Floating-point tensor of desired upper limits on mean + differences that may go undetected with probability higher than + `1 - false_pass_rate`. + low: Tensor of lower bounds on the distributions' support. + high: Tensor of upper bounds on the distributions' support. + false_fail_rate: *Scalar* admissible total rate of false failures. + false_pass_rate: *Scalar* admissible rate of false passes. + name: A name for this operation (optional). + + Returns: + n: Tensor of numbers of samples to be drawn from the distributions + of interest. + + The `discrepancy`, `low`, and `high` tensors must have + broadcast-compatible shapes. + + For each batch member `i`, of `K` total, drawing `n[i]` samples from + some scalar distribution supported on `[low[i], high[i]]` is enough + to detect a difference in means of size `discrepancy[i]` or more. + Specifically, we guarantee that (a) if the true mean is the expected + mean, `assert_true_mean_equal_by_dkwm` will fail with probability at + most `false_fail_rate / K` (which amounts to `false_fail_rate` if + applied to the whole batch at once), and (b) if the true mean + differs from the expected mean by at least `discrepancy[i]`, + `assert_true_mean_equal_by_dkwm` will pass with probability at most + `false_pass_rate`. + + The required number of samples scales + as `O((high[i] - low[i])**2)`, `O(-log(false_fail_rate/K))`, + `O(-log(false_pass_rate))`, and `O(1 / discrepancy[i]**2)`. + """ + with ops.name_scope( + name, "min_num_samples_for_dkwm_mean_test", + [low, high, false_fail_rate, false_pass_rate, discrepancy]): + discrepancy = ops.convert_to_tensor( + discrepancy, name="discrepancy") + low = ops.convert_to_tensor(low, name="low") + high = ops.convert_to_tensor(high, name="high") + false_fail_rate = ops.convert_to_tensor( + false_fail_rate, name="false_fail_rate") + false_pass_rate = ops.convert_to_tensor( + false_pass_rate, name="false_pass_rate") + # Could choose to cleverly allocate envelopes, but this is sound. + envelope1 = discrepancy / (2. * (high - low)) + envelope2 = envelope1 + false_fail_rate = _itemwise_error_rate( + false_fail_rate, [low, high, discrepancy]) + n1 = -math_ops.log(false_fail_rate / 2.) / (2. * envelope1**2) + n2 = -math_ops.log(false_pass_rate / 2.) / (2. * envelope2**2) + return math_ops.maximum(n1, n2) + + +def assert_true_mean_equal_by_dkwm_two_sample( + samples1, low1, high1, samples2, low2, high2, + false_fail_rate=1e-6, name=None): + """Asserts the means of the given distributions are equal. + + More precisely, fails if there is enough evidence (using the + [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval)) + that the means of the distributions from which the given samples are + drawn are _not_ equal with statistical significance `false_fail_rate` + or stronger, otherwise passes. If you also want to check that you + are gathering enough evidence that a pass is not spurious, see + `min_num_samples_for_dkwm_mean_two_sample_test` and + `min_discrepancy_of_true_means_detectable_by_dkwm_two_sample`. + + Note that `false_fail_rate` is a total false failure rate for all + the assertions in the batch. As such, if the batch is nontrivial, + the assertion will insist on stronger evidence to fail any one member. + + Args: + samples1: Floating-point tensor of samples from the + distribution(s) A. Entries are assumed IID across the 0th + dimension. The other dimensions must broadcast with `low1`, + `high1`, `low2`, and `high2`. + low1: Floating-point tensor of lower bounds on the supports of the + distributions A. + high1: Floating-point tensor of upper bounds on the supports of + the distributions A. + samples2: Floating-point tensor of samples from the + distribution(s) B. Entries are assumed IID across the 0th + dimension. The other dimensions must broadcast with `low1`, + `high1`, `low2`, and `high2`. + low2: Floating-point tensor of lower bounds on the supports of the + distributions B. + high2: Floating-point tensor of upper bounds on the supports of + the distributions B. + false_fail_rate: *Scalar* admissible total rate of mistakes. + name: A name for this operation (optional). + + Returns: + check: Op that raises `InvalidArgumentError` if any pair of confidence + intervals true for corresponding true means do not overlap. + """ + with ops.name_scope( + name, "assert_true_mean_equal_by_dkwm_two_sample", + [samples1, low1, high1, samples2, low2, high2, false_fail_rate]): + samples1 = ops.convert_to_tensor(samples1, name="samples1") + low1 = ops.convert_to_tensor(low1, name="low1") + high1 = ops.convert_to_tensor(high1, name="high1") + samples2 = ops.convert_to_tensor(samples2, name="samples2") + low2 = ops.convert_to_tensor(low2, name="low2") + high2 = ops.convert_to_tensor(high2, name="high2") + false_fail_rate = ops.convert_to_tensor( + false_fail_rate, name="false_fail_rate") + samples1 = _check_shape_dominates(samples1, [low1, high1]) + samples2 = _check_shape_dominates(samples2, [low2, high2]) + compatible_samples = check_ops.assert_equal( + array_ops.shape(samples1)[1:], array_ops.shape(samples2)[1:]) + with ops.control_dependencies([compatible_samples]): + # Could in principle play games with cleverly allocating + # significance instead of the even split below. It may be possible + # to get tighter intervals, in order to obtain a higher power test. + # Any allocation strategy that depends only on the support bounds + # and sample counts should be valid; however, because the intervals + # scale as O(-log(false_fail_rate)), there doesn't seem to be much + # room to win. + min_mean_1, max_mean_1 = true_mean_confidence_interval_by_dkwm( + samples1, low1, high1, false_fail_rate / 2.) + min_mean_2, max_mean_2 = true_mean_confidence_interval_by_dkwm( + samples2, low2, high2, false_fail_rate / 2.) + # I want to assert + # not (max_mean_1 < min_mean_2 or min_mean_1 > max_mean_2), + # but I think I only have and-combination of asserts, so use DeMorgan. + clause1_op = check_ops.assert_greater_equal(max_mean_1, min_mean_2) + with ops.control_dependencies([clause1_op]): + return check_ops.assert_less_equal(min_mean_1, max_mean_2) + + +def min_discrepancy_of_true_means_detectable_by_dkwm_two_sample( + n1, low1, high1, n2, low2, high2, + false_fail_rate, false_pass_rate, name=None): + """Returns the minimum mean discrepancy for a two-sample DKWM-based test. + + DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval). + + Note that `false_fail_rate` is a total false failure rate for all + the tests in the batch. As such, if the batch is nontrivial, each + member will demand more samples. The `false_pass_rate` is also + interpreted as a total, but is treated asymmetrically: If each test + in the batch detects its corresponding discrepancy with probability + at least `1 - false_pass_rate`, then running all those tests and + failing if any one fails will jointly detect all those discrepancies + with the same `false_pass_rate`. + + Args: + n1: Tensor of numbers of samples to be drawn from the distributions A. + low1: Floating-point tensor of lower bounds on the supports of the + distributions A. + high1: Floating-point tensor of upper bounds on the supports of + the distributions A. + n2: Tensor of numbers of samples to be drawn from the distributions B. + low2: Floating-point tensor of lower bounds on the supports of the + distributions B. + high2: Floating-point tensor of upper bounds on the supports of + the distributions B. + false_fail_rate: *Scalar* admissible total rate of false failures. + false_pass_rate: *Scalar* admissible rate of false passes. + name: A name for this operation (optional). + + Returns: + discr: Tensor of lower bounds on the distances between true means + detectable by a two-sample DKWM-based test. + + For each batch member `i`, of `K` total, drawing `n1[i]` samples + from scalar distribution A supported on `[low1[i], high1[i]]` and `n2[i]` + samples from scalar distribution B supported on `[low2[i], high2[i]]` + is enough to detect a difference in their true means of size + `discr[i]` or more. Specifically, we guarantee that (a) if their + true means are equal, `assert_true_mean_equal_by_dkwm_two_sample` + will fail with probability at most `false_fail_rate/K` (which + amounts to `false_fail_rate` if applied to the whole batch at once), + and (b) if their true means differ by at least `discr[i]`, + `assert_true_mean_equal_by_dkwm_two_sample` will pass with + probability at most `false_pass_rate`. + + The detectable distribution scales as + + - `O(high1[i] - low1[i])`, `O(high2[i] - low2[i])`, + - `O(1 / sqrt(n1[i]))`, `O(1 / sqrt(n2[i]))`, + - `O(-log(false_fail_rate/K))`, and + - `O(-log(false_pass_rate))`. + """ + with ops.name_scope( + name, "min_discrepancy_of_true_means_detectable_by_dkwm_two_sample", + [n1, low1, high1, n2, low2, high2, false_fail_rate, false_pass_rate]): + n1 = ops.convert_to_tensor(n1, name="n1") + low1 = ops.convert_to_tensor(low1, name="low1") + high1 = ops.convert_to_tensor(high1, name="high1") + n2 = ops.convert_to_tensor(n2, name="n2") + low2 = ops.convert_to_tensor(low2, name="low2") + high2 = ops.convert_to_tensor(high2, name="high2") + false_fail_rate = ops.convert_to_tensor( + false_fail_rate, name="false_fail_rate") + false_pass_rate = ops.convert_to_tensor( + false_pass_rate, name="false_pass_rate") + det_disc1 = min_discrepancy_of_true_means_detectable_by_dkwm( + n1, low1, high1, false_fail_rate / 2., false_pass_rate / 2.) + det_disc2 = min_discrepancy_of_true_means_detectable_by_dkwm( + n2, low2, high2, false_fail_rate / 2., false_pass_rate / 2.) + return det_disc1 + det_disc2 + + +def min_num_samples_for_dkwm_mean_two_sample_test( + discrepancy, low1, high1, low2, high2, + false_fail_rate=1e-6, false_pass_rate=1e-6, name=None): + """Returns how many samples suffice for a two-sample DKWM mean test. + + DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality] + (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval). + + Args: + discrepancy: Floating-point tensor of desired upper limits on mean + differences that may go undetected with probability higher than + `1 - false_pass_rate`. + low1: Floating-point tensor of lower bounds on the supports of the + distributions A. + high1: Floating-point tensor of upper bounds on the supports of + the distributions A. + low2: Floating-point tensor of lower bounds on the supports of the + distributions B. + high2: Floating-point tensor of upper bounds on the supports of + the distributions B. + false_fail_rate: *Scalar* admissible total rate of false failures. + false_pass_rate: *Scalar* admissible rate of false passes. + name: A name for this operation (optional). + + Returns: + n1: Tensor of numbers of samples to be drawn from the distributions A. + n2: Tensor of numbers of samples to be drawn from the distributions B. + + For each batch member `i`, of `K` total, drawing `n1[i]` samples + from scalar distribution A supported on `[low1[i], high1[i]]` and `n2[i]` + samples from scalar distribution B supported on `[low2[i], high2[i]]` + is enough to detect a difference in their true means of size + `discr[i]` or more. Specifically, we guarantee that (a) if their + true means are equal, `assert_true_mean_equal_by_dkwm_two_sample` + will fail with probability at most `false_fail_rate/K` (which + amounts to `false_fail_rate` if applied to the whole batch at once), + and (b) if their true means differ by at least `discr[i]`, + `assert_true_mean_equal_by_dkwm_two_sample` will pass with + probability at most `false_pass_rate`. + + The required number of samples scales as + + - `O((high1[i] - low1[i])**2)`, `O((high2[i] - low2[i])**2)`, + - `O(-log(false_fail_rate/K))`, + - `O(-log(false_pass_rate))`, and + - `O(1 / discrepancy[i]**2)`. + """ + with ops.name_scope( + name, "min_num_samples_for_dkwm_mean_two_sample_test", + [low1, high1, low2, high2, + false_fail_rate, false_pass_rate, discrepancy]): + discrepancy = ops.convert_to_tensor(discrepancy, name="discrepancy") + low1 = ops.convert_to_tensor(low1, name="low1") + high1 = ops.convert_to_tensor(high1, name="high1") + low2 = ops.convert_to_tensor(low2, name="low2") + high2 = ops.convert_to_tensor(high2, name="high2") + false_fail_rate = ops.convert_to_tensor( + false_fail_rate, name="false_fail_rate") + false_pass_rate = ops.convert_to_tensor( + false_pass_rate, name="false_pass_rate") + # Could choose to cleverly allocate discrepancy tolerances and + # failure probabilities, but this is sound. + n1 = min_num_samples_for_dkwm_mean_test( + discrepancy / 2., low1, high1, + false_fail_rate / 2., false_pass_rate / 2.) + n2 = min_num_samples_for_dkwm_mean_test( + discrepancy / 2., low2, high2, + false_fail_rate / 2., false_pass_rate / 2.) + return n1, n2 diff --git a/tensorflow/contrib/distributions/python/ops/vector_sinh_arcsinh_diag.py b/tensorflow/contrib/distributions/python/ops/vector_sinh_arcsinh_diag.py index e1ccf116457a97261b9ce3965552764771d3bdd2..003c66b9413fdcad20fbcc8b4bf47259692932e7 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_sinh_arcsinh_diag.py +++ b/tensorflow/contrib/distributions/python/ops/vector_sinh_arcsinh_diag.py @@ -227,7 +227,7 @@ class VectorSinhArcsinhDiag(transformed_distribution.TransformedDistribution): c = 2 * scale_diag_part / f_noskew.forward( ops.convert_to_tensor(2, dtype=dtype)) affine = bijectors.Affine( - shift=loc, scale_diag=c, validate_args=validate_args, event_ndims=1) + shift=loc, scale_diag=c, validate_args=validate_args) bijector = bijectors.Chain([affine, f]) 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 e984c63af7ce2b32ab30121bf34bb2de4dfeb218..eb810e06ddba51dbb69eab2918e181f6f7370343 100644 --- a/tensorflow/contrib/eager/python/BUILD +++ b/tensorflow/contrib/eager/python/BUILD @@ -11,12 +11,14 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ + ":checkpointable_utils", ":datasets", ":metrics", ":network", ":saver", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", + "//tensorflow/python:gradients", "//tensorflow/python:numerics", "//tensorflow/python:resource_variable_ops", "//tensorflow/python:script_ops", @@ -26,7 +28,6 @@ py_library( "//tensorflow/python/eager:backprop", "//tensorflow/python/eager:context", "//tensorflow/python/eager:core", - "//tensorflow/python/eager:custom_gradient", "//tensorflow/python/eager:execution_callbacks", "//tensorflow/python/eager:function", ], @@ -52,7 +53,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", @@ -69,6 +70,8 @@ cuda_py_test( srcs = ["datasets_test.py"], additional_deps = [ ":datasets", + ":checkpointable_utils", + "//tensorflow/contrib/data/python/ops:transformation_ops", "//tensorflow/contrib/lookup:lookup_py", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", @@ -115,6 +118,7 @@ py_library( srcs_version = "PY2AND3", visibility = ["//tensorflow:internal"], deps = [ + "//tensorflow/contrib/eager/python:checkpointable_utils", "//tensorflow/contrib/summary:summary_ops", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", @@ -220,35 +224,52 @@ 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"], - srcs_version = "PY2AND3", - deps = [ - ":checkpointable", +cuda_py_test( + name = "checkpointable_utils_test", + srcs = ["checkpointable_utils_test.py"], + additional_deps = [ + ":checkpointable_utils", ":network", + "@six_archive//:six", "//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", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", - "@six_archive//:six", + "//tensorflow/python/keras", + ], + tags = [ + "no_windows", # TODO: needs investigation on Windows + "notsan", ], ) 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 ff419614f580d3bace9d99648478cc2204d7801d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/checkpointable_test.py +++ /dev/null @@ -1,271 +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 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( - name=name, - shape=initial_value.get_shape(), - initializer=initial_value, - getter=_variable_getter) - - self._non_slot_dict[key] = v - - return v - - # 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..389d4a03c823aef6edf520d955eec11f24d71cf2 --- /dev/null +++ b/tensorflow/contrib/eager/python/checkpointable_utils.py @@ -0,0 +1,882 @@ +"""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 abc +import collections +import weakref + +from tensorflow.contrib.eager.proto import checkpointable_object_graph_pb2 +from tensorflow.python import pywrap_tensorflow +from tensorflow.python.client import session as session_lib +from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors_impl +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 checkpointable_utils as core_checkpointable_utils +from tensorflow.python.training import optimizer as optimizer_lib +from tensorflow.python.training import saver as saver_lib +from tensorflow.python.util import deprecation + + +_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_factory in ( + checkpointable._gather_saveables_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)) + if callable(saveable_factory): + saveable = saveable_factory(name=attribute.checkpoint_key) + else: + saveable = saveable_factory + # 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) + + +def gather_initializers(root_checkpointable): + """Traverse the object graph and find initialization ops. + + Looks for `Checkpointable` objects which are dependencies of + `root_checkpointable` and which have an `initializer` property. Includes + initializers for slot variables only if the variable they are slotting for and + the optimizer are dependencies of `root_checkpointable` (i.e. if they would be + saved with a checkpoint). + + Args: + root_checkpointable: A `Checkpointable` object to gather initializers for. + Returns: + A list of initialization ops. + """ + # TODO(allenl): Extract out gathering logic so the naming logic doesn't have + # to run. + 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)} + _serialize_slot_variables( + checkpointable_objects=checkpointable_objects, + node_ids=node_ids, + object_names=object_names) + return [c.initializer for c in checkpointable_objects + if hasattr(c, "initializer") and c.initializer is not None] + + +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() + + +class _LoadStatus(object): + """Abstract base for load status callbacks.""" + + @abc.abstractmethod + def assert_consumed(self): + """Raises an exception unless a non-trivial restoration has completed.""" + pass + + @abc.abstractmethod + def run_restore_ops(self, session=None): + """Runs restore ops from the checkpoint. Requires a valid checkpoint.""" + pass + + @abc.abstractmethod + def initialize_or_restore(self, session=None): + """Runs restore ops from the checkpoint, or initializes variables.""" + pass + + +class CheckpointLoadStatus(_LoadStatus): + """Checks the status of checkpoint loading and manages restore ops. + + Returned from `Saver.restore`. Since `restore` may defer the loading of values + in the checkpoint which don't yet have corresponding Python objects, + `CheckpointLoadStatus` provides a callback to verify that checkpoint loading + is complete (`assert_consumed`). + + When graph building, `restore` does not run restore ops itself since their + creation may be deferred. The `run_restore_ops` method must be called once all + Python objects with values to restore have been created and added to the + dependency graph (this does not necessarily have to be the whole checkpoint; + calling `run_restore_ops` while `assert_consumed` fails is supported and will + partially restore the checkpoint). + + See `Saver.restore` for usage examples. + """ + + def __init__(self, checkpoint, feed_dict): + self._checkpoint = checkpoint + self._feed_dict = feed_dict + + def assert_consumed(self): + """Asserts that all objects in the checkpoint have been created/matched. + + Returns: + `self` for chaining. + Raises: + AssertionError: If there are any Python objects in the dependency graph + which have not been restored from this checkpoint or a later `restore`, + or if there are any checkpointed values which have not been matched to + Python objects. + """ + 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,)) + if self._checkpoint.slot_restorations: + # Sanity check; this collection should be clear if everything has been + # restored. + raise AssertionError("Unresolved slot restorations: %s" % ( + self._checkpoint.slot_restorations,)) + if self._checkpoint.unused_attributes: + raise AssertionError( + ("Unused attributes in these objects (the attributes exist in the " + "checkpoint but not in the objects): %s") % ( + self._checkpoint.unused_attributes.items(),)) + return self + + def run_restore_ops(self, session=None): + """Run operations to restore objects in the dependency graph.""" + if context.executing_eagerly(): + return # Run eagerly + if session is None: + session = ops.get_default_session() + session.run(self._checkpoint.restore_ops, feed_dict=self._feed_dict) + + def initialize_or_restore(self, session=None): + """Alias for `run_restore_ops`. + + This method has a sibling in `InitializationOnlyStatus` which instead + initializes variables. That type is returned if no checkpoint is specified + in `Saver.restore`. + + Args: + session: The session to run restore ops in. If `None`, uses the default + session. + """ + self.run_restore_ops(session=session) + + +class InitializationOnlyStatus(_LoadStatus): + """Returned from `Saver.restore` when no checkpoint has been specified. + + Objects of this type have the same `assert_consumed` method as + `CheckpointLoadStatus`, but it always fails. However, + `initialize_or_restore` works on objects of both types, and will + initialize variables in `InitializationOnlyStatus` objects or restore them + otherwise. + """ + + def __init__(self, root_checkpointable): + self._root_checkpointable = root_checkpointable + + def assert_consumed(self): + """Assertion for consistency with `CheckpointLoadStatus`. Always fails.""" + raise AssertionError( + "No checkpoint specified (save_path=None); nothing is being restored.") + + def run_restore_ops(self, session=None): + """For consistency with `CheckpointLoadStatus`. + + Use `initialize_or_restore` for initializing if no checkpoint was passed + to `Saver.restore` and restoring otherwise. + + Args: + session: Not used. + """ + raise AssertionError( + "No checkpoint specified, so no restore ops are available " + "(save_path=None to Saver.restore).") + + def initialize_or_restore(self, session=None): + """Runs initialization ops for variables. + + Only objects which would be saved by `Saver.save` will be initialized. See + `gather_initializers` for details. + + This method does nothing when executing eagerly (initializers get run + eagerly). + + Args: + session: The session to run initialization ops in. If `None`, uses the + default session. + """ + if context.executing_eagerly(): + return # run eagerly + if session is None: + session = ops.get_default_session() + session.run(gather_initializers(self._root_checkpointable)) + + +_DEPRECATED_RESTORE_INSTRUCTIONS = ( + "Restoring a name-based tf.train.Saver checkpoint using the object-based " + "restore API. This mode uses global names to match variables, and so is " + "somewhat fragile. It also adds new restore ops to the graph each time it " + "is called. Prefer re-encoding training checkpoints in the object-based " + "format: run save() on the object-based saver (the same one this message " + "is coming from) and use that checkpoint in the future.") + + +class NameBasedSaverStatus(_LoadStatus): + """Status for loading a name-based training checkpoint.""" + + def __init__(self, object_saver, save_path): + self._object_saver = object_saver + self._save_path = save_path + + def assert_consumed(self): + """Assertion for consistency with `CheckpointLoadStatus`. Always fails.""" + raise AssertionError( + "Restoring a name-based checkpoint. No load status is available.") + + @deprecation.deprecated( + date=None, instructions=_DEPRECATED_RESTORE_INSTRUCTIONS) + def run_restore_ops(self, session=None): + """Load the name-based training checkpoint using a new `tf.train.Saver`.""" + if session is None and not context.executing_eagerly(): + session = ops.get_default_session() + with ops.device("/cpu:0"): + saver_lib.Saver(self._object_saver._global_variable_names()).restore( # pylint: disable=protected-access + sess=session, save_path=self._save_path) + + def initialize_or_restore(self, session=None): + """Alias for `run_restore_ops`.""" + self.run_restore_ops(session=session) + + +class _SessionWithFeedDictAdditions(session_lib.SessionInterface): + """Pretends to be a session, inserts extra feeds on run().""" + + def __init__(self, session, feed_additions): + self._wrapped_session = session + self._feed_additions = feed_additions + + def run(self, fetches, feed_dict=None, **kwargs): + if feed_dict is None: + feed_dict = {} + else: + feed_dict = feed_dict.copy() + feed_dict.update(self._feed_additions) + return self._wrapped_session.run( + fetches=fetches, feed_dict=feed_dict, **kwargs) + + +class CheckpointableSaver(object): + """Saves and restores a `Checkpointable` object and its dependencies. + + See `Checkpointable` for details of dependency management. `Saver` wraps + `tf.train.Saver` for saving, including extra information about the graph of + dependencies between Python objects. When restoring, it uses this information + about the save-time dependency graph to more robustly match objects with their + checkpointed values. When executing eagerly, it supports restoring variables + on object creation (see `Saver.restore`). + + Values in a checkpoint are mapped to `Checkpointable` Python objects + (`Variable`s, `Optimizer`s, `Layer`s) based on the names provided when the + checkpoint was written. To avoid breaking existing checkpoints when modifying + a class, dependency names (the names of attributes to which `Checkpointable` + objects are assigned) may not change. These names are local to objects, in + contrast to the `Variable.name`-based save/restore from `tf.train.Saver`, and + so allow additional program transformations. + """ + + def __init__(self, root_checkpointable): + """Configure saving. + + Args: + root_checkpointable: The root of the object graph to save/restore. This + object and all of its dependencies are saved in the checkpoint. When + restoring, objects are matched and restored starting from this root. + """ + # Allow passing in a weak reference to avoid reference cycles when + # `Checkpointable` objects save themselves. + self._root_checkpointable_ref = root_checkpointable + if not context.executing_eagerly(): + with ops.device("/cpu:0"): + self._file_prefix_placeholder = constant_op.constant("model") + else: + self._file_prefix_placeholder = None + + # Op caching for save + self._object_graph_feed_tensor = None + self._last_save_object_graph = None + self._last_save_saver = None + + # Op caching for restore + self._object_graph_restore_tensor = None + self._last_restore_object_graph = None + self._last_restore_checkpoint = None + + @property + def _root_checkpointable(self): + if isinstance(self._root_checkpointable_ref, weakref.ref): + derefed = self._root_checkpointable_ref() + assert derefed is not None + return derefed + else: + return self._root_checkpointable_ref + + def save(self, file_prefix, checkpoint_number=None, session=None): + """Save a training checkpoint. + + The saved checkpoint includes variables created by this object and any + Checkpointable objects it depends on at the time `Saver.save()` is called. + + 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 `checkpoint_number`, if provided. + 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, graph_proto = _serialize_object_graph( + self._root_checkpointable) + if not context.executing_eagerly(): + if session is None: + session = ops.get_default_session() + if self._object_graph_feed_tensor is None: + with ops.device("/cpu:0"): + self._object_graph_feed_tensor = constant_op.constant( + "", dtype=dtypes.string) + object_graph_tensor = self._object_graph_feed_tensor + feed_additions = {object_graph_tensor: graph_proto.SerializeToString()} + else: + session = None + with ops.device("/cpu:0"): + object_graph_tensor = constant_op.constant( + graph_proto.SerializeToString(), dtype=dtypes.string) + feed_additions = None + assert _OBJECT_GRAPH_PROTO_KEY not in named_variables + named_variables[_OBJECT_GRAPH_PROTO_KEY] = _NoRestoreSaveable( + tensor=object_graph_tensor, + name=_OBJECT_GRAPH_PROTO_KEY) + if self._last_save_object_graph != graph_proto: + if self._last_save_object_graph is not None: + raise NotImplementedError( + "Using a single Saver to save a mutated object graph is not " + "currently supported when graph building. Use a different Saver " + "when the object graph changes (save ops will be duplicated when " + "graph building), or file a feature request if this limitation " + "bothers you.") + saver = saver_lib.Saver(var_list=named_variables) + self._last_save_saver = saver + self._last_save_object_graph = graph_proto + else: + saver = self._last_save_saver + with ops.device("/cpu:0"): + save_path = saver.save( + sess=_SessionWithFeedDictAdditions( + session=session, feed_additions=feed_additions), + save_path=file_prefix, + write_meta_graph=False, + global_step=checkpoint_number) + return save_path + + def _global_variable_names(self): + """Generate a `tf.train.Saver`-style `var_list` using `variable.name`s.""" + named_saveables, graph_proto = _serialize_object_graph( + self._root_checkpointable) + saver_names = {} + for object_proto in graph_proto.nodes: + for attribute_proto in object_proto.attributes: + saver_names[attribute_proto.full_name] = named_saveables[ + attribute_proto.checkpoint_key] + return saver_names + + def restore(self, save_path, session=None): + """Restore a training checkpoint. + + Restores `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 the `root_checkpointable` passed to the + constructor after this call will be matched if they have a corresponding + object in the checkpoint. + + When building a graph, restorations are added to the graph but not run. A + session is required to retrieve checkpoint metadata. + + To disallow deferred loading, assert immediately that all checkpointed + variables have been matched to variable objects: + + ```python + saver = Saver(root) + saver.restore(path).assert_consumed() + ``` + + An exception will be raised unless every object was matched and its + variables already exist. + + When graph building, `assert_consumed()` indicates that all of the restore + ops which will be created for this checkpoint have been created. They can be + run via the `run_restore_ops()` function of the status object: + + ```python + saver.restore(path).assert_consumed().run_restore_ops() + ``` + + If the checkpoint has not been consumed completely, then the list of restore + ops will grow as more objects are added to the dependency graph. + + Name-based `tf.train.Saver` checkpoints can be loaded using this + method. There is no deferred loading, and names are used to match + variables. No restore ops are created/run until `run_restore_ops()` or + `initialize_or_restore()` are called on the returned status object, even + when executing eagerly. Re-encode name-based checkpoints using this + object-based `Saver.save` as soon as possible. + + 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), returns an + object which may run initializers for objects in the dependency + graph. If the checkpoint was written by the name-based `tf.train.Saver`, + names are used to match variables. + session: The session to retrieve metadata with. Ignored when executing + eagerly. If not provided when graph building, the default session is + used. + + Returns: + A load status object, which can be used to make assertions about the + status of checkpoint restoration and run initialization/restore ops + (of type `CheckpointLoadStatus`, or `InitializationOnlyStatus` if + `save_path` is `None`). + + If `save_path` points to a name-based checkpoint, a `NameBasedSaverStatus` + object is returned which runs restore ops from a name-based saver. + """ + if save_path is None: + return InitializationOnlyStatus(self._root_checkpointable) + in_graph_mode = not context.executing_eagerly() + if in_graph_mode: + if session is None: + session = ops.get_default_session() + file_prefix_tensor = self._file_prefix_placeholder + file_prefix_feed_dict = {self._file_prefix_placeholder: save_path} + else: + session = None + with ops.device("/cpu:0"): + file_prefix_tensor = constant_op.constant(save_path) + file_prefix_feed_dict = None + try: + if not in_graph_mode or self._object_graph_restore_tensor is None: + with ops.device("/cpu:0"): + object_graph_string, = io_ops.restore_v2( + prefix=file_prefix_tensor, + tensor_names=[_OBJECT_GRAPH_PROTO_KEY], + shape_and_slices=[""], + dtypes=[dtypes.string], + name="object_graph_proto_read") + if in_graph_mode: + self._object_graph_restore_tensor = object_graph_string + if in_graph_mode: + object_graph_string = session.run( + self._object_graph_restore_tensor, + feed_dict=file_prefix_feed_dict) + else: + object_graph_string = object_graph_string.numpy() + except errors_impl.NotFoundError: + # The object graph proto does not exist in this checkpoint. Try again with + # name-based saving. + return NameBasedSaverStatus(self, save_path) + + object_graph_proto = ( + checkpointable_object_graph_pb2.CheckpointableObjectGraph()) + object_graph_proto.ParseFromString(object_graph_string) + if in_graph_mode and object_graph_proto == self._last_restore_object_graph: + checkpoint = self._last_restore_checkpoint + else: + if in_graph_mode: + dtype_map = None + else: + reader = pywrap_tensorflow.NewCheckpointReader(save_path) + dtype_map = reader.get_variable_to_dtype_map() + checkpoint = core_checkpointable_utils._Checkpoint( # pylint: disable=protected-access + object_graph_proto=object_graph_proto, + save_path=file_prefix_tensor, + dtype_map=dtype_map) + if in_graph_mode: + if self._last_restore_object_graph is not None: + raise NotImplementedError( + "Using a single Saver to restore different object graphs is not " + "currently supported when graph building. Use a different Saver " + "for each object graph (restore ops will be duplicated), or " + "file a feature request if this limitation bothers you.") + self._last_restore_checkpoint = checkpoint + self._last_restore_object_graph = object_graph_proto + core_checkpointable._CheckpointPosition( # pylint: disable=protected-access + checkpoint=checkpoint, proto_id=0).restore(self._root_checkpointable) + load_status = CheckpointLoadStatus( + checkpoint, feed_dict=file_prefix_feed_dict) + return load_status + + +class Checkpoint(core_checkpointable.Checkpointable): + """A utility class which groups `Checkpointable` objects. + + Accepts arbitrary keyword arguments to its constructor and saves those values + with a checkpoint. Maintains a `save_counter` for numbering checkpoints. + + Example usage: + + ```python + import tensorflow as tf + import tensorflow.contrib.eager as tfe + import os + + checkpoint_directory = "/tmp/training_checkpoints" + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + root = tfe.Checkpoint(optimizer=optimizer, model=model) + root.restore(tf.train.latest_checkpoint(checkpoint_directory)) + for _ in range(num_training_steps): + optimizer.minimize( ... ) + root.save(file_prefix=checkpoint_prefix) + ``` + + For more manual control over saving, use `tfe.CheckpointableSaver` directly. + + Attributes: + save_counter: Incremented when `save()` is called. Used to number + checkpoints. + """ + + def __init__(self, **kwargs): + """Group objects into a training checkpoint. + + Args: + **kwargs: Keyword arguments are set as attributes of this object, and are + saved with the checkpoint. Attribute values must derive from + `CheckpointableBase`. + Raises: + ValueError: If objects in `kwargs` are not Checkpointable. + """ + super(Checkpoint, self).__init__() + for k, v in sorted(kwargs.items(), key=lambda item: item[0]): + if not isinstance(v, core_checkpointable.CheckpointableBase): + raise ValueError( + ("`Checkpoint` was expecting an object derived from " + "`CheckpointableBase`, got %s.") % (v,)) + setattr(self, k, v) + self._save_counter = None # Created lazily for restore-on-create. + self._saver = CheckpointableSaver(weakref.ref(self)) + + def _maybe_create_save_counter(self): + """Create a save counter if it does not yet exist.""" + if self._save_counter is None: + # Initialized to 0 and incremented before saving. + with ops.device("/cpu:0"): + self._save_counter = add_variable( + self, name="save_counter", initializer=0, dtype=dtypes.int64) + + @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. + """ + self._maybe_create_save_counter() + return self._save_counter + + def save(self, file_prefix, session=None): + """Save a checkpoint. Wraps `tfe.CheckpointableSaver.save`.""" + in_graph_mode = not context.executing_eagerly() + if in_graph_mode: + if session is None: + session = ops.get_default_session() + if self._save_counter is None: + # When graph building, if this is a new save counter variable then it + # needs to be initialized before assign_add. This is only an issue if + # restore() has not been called first. + session.run(self.save_counter.initializer) + with ops.colocate_with(self.save_counter): + assign_op = self.save_counter.assign_add(1) + if in_graph_mode: + session.run(assign_op) + return self._saver.save( + file_prefix=file_prefix, + checkpoint_number=self.save_counter, + session=session) + + def restore(self, save_path): + """Restore a checkpoint. Wraps `tfe.CheckpointableSaver.restore`.""" + status = self._saver.restore(save_path=save_path) + # Create the save counter now so it gets initialized with other variables + # when graph building. Creating it earlier would lead to double + # initialization when executing eagerly. + self._maybe_create_save_counter() + return status 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..4e0a9923ff50584a14a7bf5b5c5df18a06c0ed9f --- /dev/null +++ b/tensorflow/contrib/eager/python/checkpointable_utils_test.py @@ -0,0 +1,1170 @@ +# 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 six + +from tensorflow.contrib.eager.python import checkpointable_utils +from tensorflow.python.client import session as session_lib +from tensorflow.python.eager import backprop +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.keras._impl.keras.engine import training +from tensorflow.python.layers import core +from tensorflow.python.ops import control_flow_ops +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 template +from tensorflow.python.ops import variable_scope +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 NonLayerCheckpointable(checkpointable.Checkpointable): + + def __init__(self): + super(NonLayerCheckpointable, self).__init__() + self.a_variable = checkpointable_utils.add_variable( + self, name="a_variable", shape=[]) + + +# pylint: disable=not-callable +class MyModel(training.Model): + """A concrete Model for testing.""" + + def __init__(self): + super(MyModel, self).__init__() + self._named_dense = core.Dense(1, use_bias=True) + self._second = core.Dense(1, use_bias=False) + # We can still track Checkpointables which aren't Layers. + self._non_layer = NonLayerCheckpointable() + + def call(self, values): + ret = self._second(self._named_dense(values)) + return ret + + +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=[]) + + self.evaluate(checkpointable_utils.gather_initializers(obj)) + 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.executing_eagerly(): + # When executing eagerly, there's no uniquification of variable names. The + # checkpoint name will be the same. + self.assertEqual("duplicate:0", duplicate.name) + else: + # The .name attribute may be globally influenced, but the checkpoint name + # won't be (tested below). + self.assertEqual("duplicate_1: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 _MirroringSaveable( + core_saver.BaseSaverBuilder.ResourceVariableSaveable): + + def __init__(self, primary_variable, mirrored_variable, name): + self._primary_variable = primary_variable + self._mirrored_variable = mirrored_variable + super(_MirroringSaveable, self).__init__( + self._primary_variable, "", name) + + def restore(self, restored_tensors, restored_shapes): + """Restore the same value into both variables.""" + tensor, = restored_tensors + return control_flow_ops.group( + self._primary_variable.assign(tensor), + self._mirrored_variable.assign(tensor)) + + +class _OwnsMirroredVariables(checkpointable.CheckpointableBase): + """A Checkpointable object which returns a more complex SaveableObject.""" + + def __init__(self): + self.non_dep_variable = variable_scope.get_variable( + name="non_dep_variable", initializer=6., use_resource=True) + self.mirrored = variable_scope.get_variable( + name="mirrored", initializer=15., use_resource=True) + + def _gather_saveables_for_checkpoint(self): + def _saveable_factory(name=self.non_dep_variable.name): + return _MirroringSaveable( + primary_variable=self.non_dep_variable, + mirrored_variable=self.mirrored, + name=name) + return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} + + # The Saver sorts by name before parsing, so we need a name property. + @property + def name(self): + return self.non_dep_variable.name + + +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.]]) + model = MyModel() + # A nuisance Model using the same optimizer. Its slot variables should not + # go in the checkpoint, since it is never depended on. + other_model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + optimizer_step = training_util.get_or_create_global_step() + root_checkpointable = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, optimizer_step=optimizer_step) + if context.executing_eagerly(): + optimizer.minimize( + lambda: model(input_value), + global_step=optimizer_step) + optimizer.minimize( + lambda: other_model(input_value), + global_step=optimizer_step) + else: + train_op = optimizer.minimize( + model(input_value), global_step=optimizer_step) + optimizer.minimize( + other_model(input_value), + global_step=optimizer_step) + self.evaluate(checkpointable_utils.gather_initializers( + root_checkpointable)) + 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", + "model/_second/kernel", + "model/_named_dense/kernel", + "model/_named_dense/bias", + # non-Layer dependency of the model + "model/_non_layer/a_variable", + # The optimizer creates two non-slot variables + "optimizer/beta1_power", + "optimizer/beta2_power", + # Slot variables + "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m", + "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v", + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m", + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v", + "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m", + "model/_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_model/dense_1/kernel:0", + named_variables["model/_second/kernel" + suffix].name) + self.assertEqual( + "my_model/dense/kernel:0", + named_variables["model/_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_model/dense/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_model/dense/kernel/Adam", + serialized_graph.nodes[optimizer_node.slot_variables[0] + .slot_variable_node_id] + .attributes[0].full_name) + self.assertEqual( + "my_model/dense/kernel/Adam:0", + optimizer.get_slot( + var=named_variables["model/_named_dense/kernel" + suffix], + name="m").name) + self.assertEqual( + "model/_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( + "model/_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(assert_no_eager_garbage=True) + def testMoreComplexSaveableReturned(self): + v = _OwnsMirroredVariables() + checkpoint = checkpointable_utils.Checkpoint(v=v) + test_dir = self.get_temp_dir() + prefix = os.path.join(test_dir, "ckpt") + self.evaluate(v.non_dep_variable.assign(42.)) + save_path = checkpoint.save(prefix) + self.evaluate(v.non_dep_variable.assign(43.)) + self.evaluate(v.mirrored.assign(44.)) + checkpoint.restore(save_path).assert_consumed().initialize_or_restore() + self.assertEqual(42., self.evaluate(v.non_dep_variable)) + self.assertEqual(42., self.evaluate(v.mirrored)) + + @test_util.run_in_graph_and_eager_modes() + def testMoreComplexSaveableReturnedWithGlobalName(self): + # The same object can also be saved using the name-based saver. + v = _OwnsMirroredVariables() + saver = core_saver.Saver(var_list=[v]) + test_dir = self.get_temp_dir() + prefix = os.path.join(test_dir, "ckpt") + self.evaluate(v.non_dep_variable.assign(42.)) + with self.test_session() as sess: + save_path = saver.save(sess, prefix) + self.evaluate(v.non_dep_variable.assign(43.)) + self.evaluate(v.mirrored.assign(44.)) + saver.restore(sess, save_path) + self.assertEqual(42., self.evaluate(v.non_dep_variable)) + self.assertEqual(42., self.evaluate(v.mirrored)) + + @test_util.run_in_graph_and_eager_modes() + def testSaveRestore(self): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root_checkpointable = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model) + input_value = constant_op.constant([[3.]]) + if context.executing_eagerly(): + optimizer.minimize( + lambda: model(input_value)) + else: + train_op = optimizer.minimize(model(input_value)) + # TODO(allenl): Make initialization more pleasant when graph building. + root_checkpointable.save_counter # pylint: disable=pointless-statement + self.evaluate(checkpointable_utils.gather_initializers( + root_checkpointable)) + self.evaluate(train_op) + prefix = os.path.join(self.get_temp_dir(), "ckpt") + self.evaluate(state_ops.assign(model._named_dense.variables[1], [42.])) + m_bias_slot = optimizer.get_slot(model._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(model._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 + status = root_checkpointable.restore(save_path=save_path).assert_consumed() + status.run_restore_ops() + self.assertAllEqual([42.], self.evaluate(model._named_dense.variables[1])) + self.assertAllEqual(1, self.evaluate(root_checkpointable.save_counter)) + self.assertAllEqual([1.5], self.evaluate(m_bias_slot)) + if not context.executing_eagerly(): + return # Restore-on-create is only supported when executing eagerly + on_create_model = MyModel() + on_create_optimizer = adam.AdamOptimizer(0.001) + on_create_root = checkpointable_utils.Checkpoint( + optimizer=on_create_optimizer, model=on_create_model) + # Deferred restoration + status = on_create_root.restore(save_path=save_path) + on_create_model(constant_op.constant([[3.]])) # create variables + self.assertAllEqual(1, self.evaluate(on_create_root.save_counter)) + self.assertAllEqual([42.], + self.evaluate( + on_create_model._named_dense.variables[1])) + on_create_m_bias_slot = on_create_optimizer.get_slot( + on_create_model._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)) + + # TODO(allenl): Debug garbage created by this test in python3. + 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): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, + 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: model(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(): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + input_value = constant_op.constant([[3.]]) + train_op = optimizer.minimize( + model(input_value), + global_step=root.global_step) + checkpoint_path = core_saver.latest_checkpoint(checkpoint_directory) + with self.test_session(graph=ops.get_default_graph()) as session: + status = root.restore(save_path=checkpoint_path) + status.initialize_or_restore(session=session) + if checkpoint_path is None: + self.assertEqual(0, training_continuation) + with self.assertRaises(AssertionError): + status.assert_consumed() + else: + status.assert_consumed() + 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)) + + @test_util.run_in_graph_and_eager_modes() + def testAgnosticUsage(self): + """Graph/eager agnostic usage.""" + # Does create garbage when executing eagerly due to ops.Graph() creation. + 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(), self.test_session( + graph=ops.get_default_graph()), test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + checkpoint_path = core_saver.latest_checkpoint(checkpoint_directory) + status = root.restore(save_path=checkpoint_path) + input_value = constant_op.constant([[3.]]) + train_fn = functools.partial( + optimizer.minimize, + functools.partial(model, input_value), + global_step=root.global_step) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + status.initialize_or_restore() + for _ in range(num_training_steps): + train_fn() + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + self.evaluate(root.global_step)) + self.assertEqual(training_continuation + 1, + self.evaluate(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") + + def testAnonymousVarsInInit(self): + + class Model(training.Model): + + def __init__(self): + super(Model, self).__init__() + self.w = resource_variable_ops.ResourceVariable(0.0) + self.b = resource_variable_ops.ResourceVariable(0.0) + self.vars = [self.w, self.b] + + def call(self, x): + return x * self.w + self.b + + with context.eager_mode(): + model = Model() + optimizer = adam.AdamOptimizer(learning_rate=0.05) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + checkpoint = checkpointable_utils.Checkpoint( + model=model, optimizer=optimizer) + for _ in range(2): + with backprop.GradientTape() as tape: + loss = (constant_op.constant(1.) + - model(constant_op.constant(1.))) ** 2 + grad = tape.gradient(loss, model.vars) + optimizer.apply_gradients( + [(g, v) for g, v in zip(grad, model.vars)]) + checkpoint.save(checkpoint_prefix) + + @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.CheckpointableSaver( + original).save(checkpoint_prefix) + load_into = LateDependencies() + status = checkpointable_utils.CheckpointableSaver( + load_into).restore(save_path) + with self.assertRaises(AssertionError): + status.assert_consumed() + load_into.add_dep() + status.assert_consumed() + status.run_restore_ops() + 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.CheckpointableSaver(dep_after_var).save( + checkpoint_prefix) + + loaded_dep_after_var = DepAfterVar() + status = checkpointable_utils.CheckpointableSaver( + loaded_dep_after_var).restore(save_path) + loaded_dep_after_var.add_dep() + status.assert_consumed() + status.run_restore_ops() + 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 = adam.AdamOptimizer(0.1) + if context.executing_eagerly(): + optimizer.minimize(root.var.read_value) + else: + train_op = optimizer.minimize(root.var) + # Note that `optimizer` has not been added as a dependency of + # `root`. Create a one-off grouping so that slot variables for `root.var` + # get initialized too. + self.evaluate(checkpointable_utils.gather_initializers( + checkpointable_utils.Checkpoint(root=root, optimizer=optimizer))) + self.evaluate(train_op) + self.evaluate(state_ops.assign(root.var, 12.)) + no_slots_path = checkpointable_utils.CheckpointableSaver(root).save( + os.path.join(checkpoint_directory, "no_slots")) + 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.CheckpointableSaver(root).save( + os.path.join(checkpoint_directory, "with_slots")) + new_root = checkpointable.Checkpointable() + # Load the slot-containing checkpoint (deferred), then immediately overwrite + # the non-slot variable (also deferred). + slot_status = checkpointable_utils.CheckpointableSaver( + new_root).restore(slots_path) + no_slot_status = checkpointable_utils.CheckpointableSaver( + new_root).restore(no_slots_path) + with self.assertRaises(AssertionError): + no_slot_status.assert_consumed() + new_root.var = checkpointable_utils.add_variable( + new_root, name="var", shape=[]) + no_slot_status.assert_consumed() + no_slot_status.run_restore_ops() + self.assertEqual(12., self.evaluate(new_root.var)) + new_root.optimizer = adam.AdamOptimizer(0.1) + with self.assertRaisesRegexp(AssertionError, "beta1_power"): + slot_status.assert_consumed() + self.assertEqual(12., self.evaluate(new_root.var)) + if context.executing_eagerly(): + # Slot variables are only created with restoring initializers when + # executing eagerly. + self.assertEqual(14., self.evaluate( + new_root.optimizer.get_slot(name="m", var=new_root.var))) + else: + self.assertIs(new_root.optimizer.get_slot(name="m", var=new_root.var), + None) + if context.executing_eagerly(): + new_root.optimizer.minimize(new_root.var.read_value) + else: + train_op = new_root.optimizer.minimize(new_root.var) + # The slot variable now exists; restore() didn't create it, but we should + # now have a restore op for it. + slot_status.run_restore_ops() + self.assertEqual(14., self.evaluate( + new_root.optimizer.get_slot(name="m", var=new_root.var))) + self.evaluate(train_op) + 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.)) + saver = checkpointable_utils.CheckpointableSaver(save_root) + first_path = saver.save(os.path.join(checkpoint_directory, "first")) + self.evaluate(state_ops.assign(save_root.dep.var, 13.)) + second_path = saver.save(os.path.join(checkpoint_directory, "second")) + + first_root = checkpointable.Checkpointable() + second_root = checkpointable.Checkpointable() + first_status = checkpointable_utils.CheckpointableSaver( + first_root).restore(first_path) + second_status = checkpointable_utils.CheckpointableSaver( + second_root).restore(second_path) + 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() + first_status.run_restore_ops() + self.assertEqual(12., self.evaluate(load_dep.var)) + second_root.dep = load_dep + second_status.assert_consumed() + second_status.run_restore_ops() + 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.CheckpointableSaver( + second_root).restore(second_path) + first_status = checkpointable_utils.CheckpointableSaver( + first_root).restore(first_path) + 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() + first_status.run_restore_ops() + self.assertEqual(12., self.evaluate(load_dep.var)) + second_root.dep = load_dep + second_status.assert_consumed() + second_status.run_restore_ops() + 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(checkpointable_utils.gather_initializers(save_root)) + save_path = checkpointable_utils.CheckpointableSaver(save_root).save( + os.path.join(checkpoint_directory, "ckpt")) + load_root = checkpointable.Checkpointable() + checkpointable_utils.CheckpointableSaver(load_root).restore(save_path) + 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(checkpointable_utils.gather_initializers(save_root)) + save_path = checkpointable_utils.CheckpointableSaver(save_root).save( + os.path.join(checkpoint_directory, "ckpt")) + 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) + status = checkpointable_utils.CheckpointableSaver(load_root).restore( + save_path).assert_consumed() + status.run_restore_ops() + 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(checkpointable_utils.gather_initializers(first)) + checkpoint_directory = self.get_temp_dir() + save_path = checkpointable_utils.CheckpointableSaver(first).save( + os.path.join(checkpoint_directory, "ckpt")) + + # Test deferred loading + first_load = checkpointable.Checkpointable() + status = checkpointable_utils.CheckpointableSaver( + first_load).restore(save_path) + 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() + status.run_restore_ops() + 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)) + status = checkpointable_utils.CheckpointableSaver(first_load).restore( + save_path).assert_consumed() + status.run_restore_ops() + 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.CheckpointableSaver(first).save( + checkpoint_prefix) + 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.) + status = checkpointable_utils.CheckpointableSaver( + second).restore(save_path) + recreated_var1 = variable_scope.get_variable( + name="outside_var", initializer=0.) + status.run_restore_ops() + self.assertEqual(8., self.evaluate(second.var2)) + self.evaluate(recreated_var1.assign(-2.)) + self.assertEqual(-2., self.evaluate(recreated_var1)) + second.var1 = recreated_var1 + status.run_restore_ops() + self.assertEqual(4., self.evaluate(recreated_var1)) + + 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 = adam.AdamOptimizer(0.1) + obj.opt.minimize(obj.var.read_value()) + self.evaluate(checkpointable_utils.gather_initializers(obj)) + saver = checkpointable_utils.CheckpointableSaver(obj) + saver.save(checkpoint_prefix) + before_ops = graph.get_operations() + saver.save(checkpoint_prefix) + self.assertEqual(before_ops, graph.get_operations()) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testCheckpointCleanup(self): + 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.) + self.evaluate(checkpointable_utils.gather_initializers(obj)) + saver = checkpointable_utils.Checkpoint(obj=obj) + for _ in range(10): + saver.save(checkpoint_prefix) + expected_filenames = ["checkpoint"] + for checkpoint_number in range(6, 11): + expected_filenames.append("ckpt-%d.index" % (checkpoint_number,)) + expected_filenames.append( + "ckpt-%d.data-00000-of-00001" % (checkpoint_number,)) + six.assertCountEqual( + self, + expected_filenames, + os.listdir(checkpoint_directory)) + + 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 = adam.AdamOptimizer(0.1) + obj.opt.minimize(obj.var.read_value()) + self.evaluate(checkpointable_utils.gather_initializers(obj)) + saver = checkpointable_utils.CheckpointableSaver(obj) + save_path = saver.save(checkpoint_prefix) + saver.restore(save_path) + before_ops = graph.get_operations() + saver.restore(save_path) + self.assertEqual(before_ops, graph.get_operations()) + + def testMultipleGraphsNonSlotVariables(self): + with context.graph_mode(): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + optimizer = adam.AdamOptimizer(0.001) + # Construct a model in one graph + first_graph = ops.Graph() + first_session = session_lib.Session(graph=first_graph) + with first_graph.as_default(), first_session.as_default(): + first_variable = resource_variable_ops.ResourceVariable([1.]) + first_root_checkpointable = checkpointable_utils.Checkpoint( + optimizer=optimizer, variable=first_variable) + train_op = optimizer.minimize(first_variable.read_value) + self.evaluate(checkpointable_utils.gather_initializers( + first_root_checkpointable)) + self.evaluate(train_op) + self.evaluate(first_variable.assign([1.])) + self.evaluate(optimizer.get_slot( + var=first_variable, name="m").assign([2.])) + beta1_power, _ = optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(3.)) + + # Save and load in a second graph + second_graph = ops.Graph() + with second_graph.as_default(), session_lib.Session(graph=second_graph): + second_variable = resource_variable_ops.ResourceVariable([1.]) + second_root_checkpointable = checkpointable_utils.Checkpoint( + optimizer=optimizer, variable=second_variable) + train_op = optimizer.minimize(second_variable.read_value) + second_root_checkpointable.restore(None).initialize_or_restore() + self.evaluate(train_op) + self.evaluate(second_variable.assign([4.])) + self.evaluate(optimizer.get_slot( + var=second_variable, name="m").assign([5.])) + beta1_power, _ = optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(6.)) + save_path = second_root_checkpointable.save(checkpoint_prefix) + self.evaluate(second_variable.assign([7.])) + self.evaluate(optimizer.get_slot( + var=second_variable, name="m").assign([8.])) + beta1_power, _ = optimizer._get_beta_accumulators() + self.assertAllEqual(6., self.evaluate(beta1_power)) + status = second_root_checkpointable.restore(save_path) + status.assert_consumed().run_restore_ops() + self.assertAllEqual([4.], self.evaluate(second_variable)) + self.assertAllEqual([5.], self.evaluate(optimizer.get_slot( + var=second_variable, name="m"))) + beta1_power, _ = optimizer._get_beta_accumulators() + self.assertAllEqual(6., self.evaluate(beta1_power)) + + # Check that the first graph is unmolested + with first_graph.as_default(), first_session.as_default(): + self.assertAllEqual([1.], self.evaluate(first_variable)) + self.assertAllEqual([2.], self.evaluate(optimizer.get_slot( + var=first_variable, name="m"))) + beta1_power, _ = optimizer._get_beta_accumulators() + self.assertAllEqual(3., self.evaluate(beta1_power)) + + +class TemplateTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def test_checkpointable_save_restore(self): + + def _templated(): + v = variable_scope.get_variable( + "v", shape=[1], initializer=init_ops.zeros_initializer()) + v2 = variable_scope.get_variable( + "v2", shape=[1], initializer=init_ops.zeros_initializer()) + return v, v + 1., v2 + + save_template = template.make_template("s1", _templated) + save_root = checkpointable_utils.Checkpoint(my_template=save_template) + v1_save, _, v2_save = save_template() + self.evaluate(v1_save.assign([12.])) + self.evaluate(v2_save.assign([14.])) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = save_root.save(checkpoint_prefix) + + load_template = template.make_template("s2", _templated) + load_root = checkpointable_utils.Checkpoint(my_template=load_template) + status = load_root.restore(save_path) + var, var_plus_one, var2 = load_template() + self.assertEqual(2, len(load_template._checkpoint_dependencies)) + self.assertEqual("v", load_template._checkpoint_dependencies[0].name) + self.assertEqual("v2", load_template._checkpoint_dependencies[1].name) + status.assert_consumed().run_restore_ops() + self.assertAllEqual([12.], self.evaluate(var)) + self.assertAllEqual([13.], self.evaluate(var_plus_one)) + self.assertAllEqual([14.], self.evaluate(var2)) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def test_checkpointable_save_restore_nested(self): + + def _inner_template(): + v = variable_scope.get_variable( + "v", shape=[1], initializer=init_ops.zeros_initializer()) + return v + + def _outer_template(): + first_inner = template.make_template("i1", _inner_template) + second_inner = template.make_template("i2", _inner_template) + v1 = first_inner() + v2 = second_inner() + v3 = second_inner() + return (first_inner, second_inner), (v1, v2, v3) + + with variable_scope.variable_scope("ignored"): + save_template = template.make_template("s1", _outer_template) + save_root = checkpointable_utils.Checkpoint(my_template=save_template) + (inner_template_one, inner_template_two), _ = save_template() + self.evaluate(inner_template_one.variables[0].assign([20.])) + self.evaluate(inner_template_two.variables[0].assign([25.])) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = save_root.save(checkpoint_prefix) + + load_template = template.make_template("s2", _outer_template) + load_root = checkpointable_utils.Checkpoint(my_template=load_template) + status = load_root.restore(save_path) + (inner_template_one, inner_template_two), (v1, v2, v3) = load_template() + outer_template_dependencies = load_root.my_template._checkpoint_dependencies + self.assertEqual(2, len(outer_template_dependencies)) + self.assertEqual("i1", outer_template_dependencies[0].name) + self.assertIs(inner_template_one, outer_template_dependencies[0].ref) + self.assertEqual("i2", outer_template_dependencies[1].name) + self.assertIs(inner_template_two, outer_template_dependencies[1].ref) + self.assertEqual(1, len(inner_template_one._checkpoint_dependencies)) + self.assertEqual("v", inner_template_one._checkpoint_dependencies[0].name) + self.assertEqual(1, len(inner_template_two._checkpoint_dependencies)) + self.assertEqual("v", inner_template_two._checkpoint_dependencies[0].name) + status.assert_consumed().run_restore_ops() + self.assertAllEqual([20.], self.evaluate(v1)) + self.assertAllEqual([25.], self.evaluate(v2)) + self.assertAllEqual([25.], self.evaluate(v3)) + + +class CheckpointCompatibilityTests(test.TestCase): + + def _initialized_model(self): + input_value = constant_op.constant([[3.]]) + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + optimizer_step = training_util.get_or_create_global_step() + root_checkpointable = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, optimizer_step=optimizer_step) + train_op = optimizer.minimize( + functools.partial(model, input_value), + global_step=optimizer_step) + self.evaluate(checkpointable_utils.gather_initializers( + root_checkpointable)) + self.evaluate(train_op) + # A regular variable, a slot variable, and a non-slot Optimizer variable + # with known values to check when loading. + self.evaluate(model._named_dense.bias.assign([1.])) + self.evaluate(optimizer.get_slot( + var=model._named_dense.bias, name="m").assign([2.])) + beta1_power, _ = optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(3.)) + return root_checkpointable + + def _set_sentinels(self, root_checkpointable): + self.evaluate(root_checkpointable.model._named_dense.bias.assign([101.])) + self.evaluate( + root_checkpointable.optimizer.get_slot( + var=root_checkpointable.model._named_dense.bias, name="m") + .assign([102.])) + beta1_power, _ = root_checkpointable.optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(103.)) + + def _check_sentinels(self, root_checkpointable): + self.assertAllEqual( + [1.], self.evaluate(root_checkpointable.model._named_dense.bias)) + self.assertAllEqual([2.], self.evaluate( + root_checkpointable.optimizer.get_slot( + var=root_checkpointable.model._named_dense.bias, name="m"))) + beta1_power, _ = root_checkpointable.optimizer._get_beta_accumulators() + self.assertAllEqual(3., self.evaluate(beta1_power)) + + def _write_name_based_checkpoint(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.test_session( + graph=save_graph) as session: + root = self._initialized_model() + name_saver = core_saver.Saver() + return name_saver.save( + sess=session, save_path=checkpoint_prefix, + global_step=root.optimizer_step) + + @test_util.run_in_graph_and_eager_modes() + def testLoadFromNameBasedSaver(self): + """Save a name-based checkpoint, load it using the object-based API.""" + with test_util.device(use_gpu=True): + save_path = self._write_name_based_checkpoint() + root = self._initialized_model() + self._set_sentinels(root) + with self.assertRaises(AssertionError): + self._check_sentinels(root) + object_saver = checkpointable_utils.CheckpointableSaver(root) + status = object_saver.restore(save_path) + with self.assertRaises(AssertionError): + status.assert_consumed() + status.run_restore_ops() + self._check_sentinels(root) + self._set_sentinels(root) + status.initialize_or_restore() + self._check_sentinels(root) + + # TODO(allenl): Test for the core name-based saver loading object-based + # checkpoints once object-based checkpointing is in core. + + def testSaveGraphLoadEager(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.test_session( + graph=save_graph) as session: + root = self._initialized_model() + object_saver = checkpointable_utils.CheckpointableSaver(root) + save_path = object_saver.save( + session=session, file_prefix=checkpoint_prefix) + with context.eager_mode(): + root = self._initialized_model() + self._set_sentinels(root) + root.restore(save_path).assert_consumed() + self._check_sentinels(root) + + def testSaveEagerLoadGraph(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.eager_mode(): + root = self._initialized_model() + object_saver = checkpointable_utils.CheckpointableSaver(root) + save_path = object_saver.save(file_prefix=checkpoint_prefix) + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.test_session( + graph=save_graph): + root = self._initialized_model() + self._set_sentinels(root) + root.restore(save_path).assert_consumed().run_restore_ops() + self._check_sentinels(root) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/eager/python/datasets.py b/tensorflow/contrib/eager/python/datasets.py index d177bfeab2d1fdc05d7ced54df8723fae2c77fdb..a4c3283dac9194880a1297371ea7591af6dddb2b 100644 --- a/tensorflow/contrib/eager/python/datasets.py +++ b/tensorflow/contrib/eager/python/datasets.py @@ -27,11 +27,12 @@ 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 -from tensorflow.python.framework import errors from tensorflow.python.framework import function from tensorflow.python.framework import ops from tensorflow.python.ops import gen_dataset_ops from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.training import checkpointable +from tensorflow.python.training.saver import BaseSaverBuilder _uid_counter = 0 _uid_lock = threading.Lock() @@ -45,8 +46,13 @@ def _generate_shared_name(prefix): return "{}{}".format(prefix, uid) -class Iterator(object): - """An iterator producing tf.Tensor objects from a tf.data.Dataset.""" +class Iterator(iterator_ops.EagerIterator, checkpointable.CheckpointableBase): + """An iterator producing tf.Tensor objects from a tf.data.Dataset. + + NOTE: Unlike the iterator created by the + @{tf.data.Dataset.make_one_shot_iterator} method, this class enables + additional experimental functionality, such as prefetching to the GPU. + """ def __init__(self, dataset): """Creates a new iterator over the given dataset. @@ -67,37 +73,12 @@ class Iterator(object): Raises: RuntimeError: When invoked without eager execution enabled. """ - - if not context.in_eager_mode(): - raise RuntimeError( - "{} objects can only be used when eager execution is enabled, use " - "tf.data.Dataset.make_iterator or " - "tf.data.Dataset.make_one_shot_iterator for graph construction". - 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( - 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( - 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) - # Delete the resource when this object is deleted - self._resource_deleter = resource_variable_ops.EagerResourceDeleter( - handle=self._resource, handle_device="/device:CPU:0") - self._device = context.context().device_name - self._buffer_resource_handle = None + super(Iterator, self).__init__(dataset) if not context.context().device_spec.device_type: is_remote_device = False else: is_remote_device = context.context().device_spec.device_type != "CPU" + self._buffer_resource_handle = None if is_remote_device: with ops.device("/device:CPU:0"): iter_string_handle = gen_dataset_ops.iterator_to_string_handle( @@ -106,7 +87,7 @@ class Iterator(object): @function.Defun(dtypes.string) def remote_fn(h): remote_iterator = iterator_ops.Iterator.from_string_handle( - h, self._output_types, self._output_shapes) + h, self.output_types, self.output_shapes, self.output_classes) return remote_iterator.get_next() remote_fn.add_to_graph(None) @@ -124,89 +105,43 @@ class Iterator(object): handle=self._buffer_resource_handle, handle_device=self._device) - def __iter__(self): - return self - - def __next__(self): # For Python 3 compatibility - return self.next() - def _next_internal(self): """Returns a nested structure of `tf.Tensor`s containing the next element. """ - with ops.device(self._device): - if self._buffer_resource_handle is not None: + if self._buffer_resource_handle is not None: + with ops.device(self._device): 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. + return sparse.deserialize_sparse_tensors( + nest.pack_sequence_as(self._output_types, ret), self._output_types, + self._output_shapes, self._output_classes) + else: + return super(Iterator, self)._next_internal() - Returns: - A nested structure of `tf.TensorShape` objects corresponding to each - component of an element of this dataset. - """ - return self._output_shapes + # TODO(shivaniagrawal): Expose checkpointable stateful objects from dataset + # attributes(potential). - @property - def output_types(self): - """Returns the type of each component of an element of this iterator. + class _Saveable(BaseSaverBuilder.SaveableObject): + """SaveableObject for saving/restoring iterator state.""" - Returns: - A nested structure of `tf.DType` objects corresponding to each component - of an element of this dataset. - """ - return self._output_types + def __init__(self, iterator_resource, name): + serialized_iterator = gen_dataset_ops.serialize_iterator( + iterator_resource) + specs = [ + BaseSaverBuilder.SaveSpec(serialized_iterator, "", name + "_STATE") + ] + # pylint: disable=protected-access + super(Iterator._Saveable, self).__init__(iterator_resource, specs, name) - def get_next(self, name=None): - """Returns a nested structure of `tf.Tensor`s containing the next element. + def restore(self, restored_tensors, restored_shapes): + with ops.colocate_with(self.op): + return gen_dataset_ops.deserialize_iterator(self.op, + restored_tensors[0]) - Args: - name: (Optional.) A name for the created operation. Currently unused. + def _gather_saveables_for_checkpoint(self): - Returns: - A nested structure of `tf.Tensor` objects. + def _saveable_factory(name): + return self._Saveable(self._resource, name) - Raises: - `tf.errors.OutOfRangeError`: If the end of the dataset has been reached. - """ - del name - return self._next_internal() + return {"ITERATOR": _saveable_factory} diff --git a/tensorflow/contrib/eager/python/datasets_test.py b/tensorflow/contrib/eager/python/datasets_test.py index a1611e92b113839c2dd2a3b2560b0ba90c0a7ef0..c658505de41bb6a0007440f4850fef720c3e97f1 100644 --- a/tensorflow/contrib/eager/python/datasets_test.py +++ b/tensorflow/contrib/eager/python/datasets_test.py @@ -16,11 +16,17 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os + +import threading import time import numpy as np from tensorflow.contrib import lookup +from tensorflow.contrib.data.python.ops import threadpool +from tensorflow.contrib.data.python.ops import unique +from tensorflow.contrib.eager.python import checkpointable_utils from tensorflow.contrib.eager.python import datasets from tensorflow.python.data import Dataset from tensorflow.python.eager import test @@ -41,6 +47,18 @@ class IteratorTest(test.TestCase): got.append(t.numpy()) self.assertAllEqual([0, 1, 2, 3], got) + def testBasicOneShotIterator(self): + got = [] + for t in Dataset.range(4).make_one_shot_iterator(): + got.append(t.numpy()) + self.assertAllEqual([0, 1, 2, 3], got) + + def testBasicImplicitIterator(self): + got = [] + for t in Dataset.range(4): + 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()) @@ -50,6 +68,15 @@ class IteratorTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): iterator.get_next() + def testGetNextOneShotIterator(self): + iterator = Dataset.range(4).make_one_shot_iterator() + 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) @@ -165,6 +192,93 @@ class IteratorTest(test.TestCase): x = math_ops.add(x, x) self.assertAllEqual([0., 2.], x.numpy()) + def testOverrideThreadPool(self): + + def get_thread_id(_): + # Python creates a dummy thread object to represent the current + # thread when called from an "alien" thread (such as a + # `PrivateThreadPool` thread in this case). It does not include + # the TensorFlow-given display name, but it has a unique + # identifier that maps one-to-one with the underlying OS thread. + return np.array(threading.current_thread().ident).astype(np.int64) + + for num_threads in [1, 2, 4, 8, 16]: + + dataset = ( + Dataset.range(1000).map( + lambda x: script_ops.py_func(get_thread_id, [x], dtypes.int64), + num_parallel_calls=32).apply(unique.unique())) + + dataset = threadpool.override_threadpool( + dataset, + threadpool.PrivateThreadPool( + num_threads, display_name='private_thread_pool_%d' % num_threads)) + + thread_ids = [] + for next_element in datasets.Iterator(dataset): + thread_ids.append(next_element) + self.assertEqual(len(thread_ids), len(set(thread_ids))) + self.assertGreater(len(thread_ids), 0) + # NOTE(mrry): We don't control the thread pool scheduling, and + # so cannot guarantee that all of the threads in the pool will + # perform work. + self.assertLessEqual(len(thread_ids), num_threads) + + def testSaveRestore(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt') + dataset = Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + dataset = dataset.map(math_ops.square).batch(2) + iterator = datasets.Iterator(dataset) + checkpoint = checkpointable_utils.Checkpoint(iterator=iterator) + self.assertAllEqual([1, 4], iterator.get_next().numpy()) + save_path = checkpoint.save(checkpoint_prefix) + self.assertAllEqual([9, 16], iterator.get_next().numpy()) + self.assertAllEqual([25, 36], iterator.get_next().numpy()) + checkpoint.restore(save_path) + self.assertAllEqual([9, 16], iterator.get_next().numpy()) + self.assertAllEqual([25, 36], iterator.get_next().numpy()) + + def testSaveRestoreMultipleIterator(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt') + dataset = Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + dataset = dataset.map(math_ops.square).batch(2) + iterator_1 = datasets.Iterator(dataset) + iterator_2 = datasets.Iterator(dataset) + dataset_2 = Dataset.range(10) + iterator_3 = datasets.Iterator(dataset_2) + + checkpoint = checkpointable_utils.Checkpoint( + iterator_1=iterator_1, iterator_2=iterator_2, iterator_3=iterator_3) + self.assertAllEqual([1, 4], iterator_1.get_next().numpy()) + self.assertEqual(0, iterator_3.get_next().numpy()) + self.assertEqual(1, iterator_3.get_next().numpy()) + self.assertEqual(2, iterator_3.get_next().numpy()) + + save_path = checkpoint.save(checkpoint_prefix) + self.assertAllEqual([1, 4], iterator_2.get_next().numpy()) + self.assertAllEqual([9, 16], iterator_2.get_next().numpy()) + self.assertEqual(3, iterator_3.get_next().numpy()) + checkpoint.restore(save_path) + self.assertAllEqual([9, 16], iterator_1.get_next().numpy()) + self.assertAllEqual([1, 4], iterator_2.get_next().numpy()) + self.assertEqual(3, iterator_3.get_next().numpy()) + + def testRestoreExhaustedIterator(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt') + dataset = Dataset.range(3) + iterator = datasets.Iterator(dataset) + + checkpoint = checkpointable_utils.Checkpoint(iterator=iterator) + self.assertEqual(0, iterator.get_next().numpy()) + self.assertEqual(1, iterator.get_next().numpy()) + save_path = checkpoint.save(checkpoint_prefix) + self.assertEqual(2, iterator.get_next().numpy()) + checkpoint.restore(save_path) + self.assertEqual(2, iterator.get_next().numpy()) + class DatasetConstructorBenchmark(test.Benchmark): diff --git a/tensorflow/contrib/eager/python/evaluator.py b/tensorflow/contrib/eager/python/evaluator.py index 68e7b5421fec7f73f10e381ca45f9d900de299d7..37c8f0d47adbde6932bf409cdcae9a1845d700b5 100644 --- a/tensorflow/contrib/eager/python/evaluator.py +++ b/tensorflow/contrib/eager/python/evaluator.py @@ -57,7 +57,7 @@ class Evaluator(object): self._model = model self._metrics = {} self._evaluators = {} - if context.in_graph_mode(): + if not context.executing_eagerly(): self.call = function.defun(self.call) # ---- API for users ---- @@ -90,7 +90,7 @@ class Evaluator(object): Only for graph execution. @end_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Evaluator.init_variables() not needed when " "eager execution is enabled.") return control_flow_ops.group([m.init_variables() for _, m in self.metrics]) @@ -113,7 +113,8 @@ class Evaluator(object): with summary_ops.create_file_writer( summary_logdir).as_default(), summary_ops.always_record_summaries(): return self._all_metric_results() - if context.in_eager_mode(): + + if context.executing_eagerly(): return f() else: return function.defun(f)() @@ -158,16 +159,16 @@ class Evaluator(object): @end_compatibility """ summary_logdir = kwargs.pop("summary_logdir", None) - if context.in_graph_mode(): - call_op = self.__call__(dataset.make_one_shot_iterator().get_next(), - *args, **kwargs) - init_op = self.init_variables() - results_op = self.all_metric_results(summary_logdir) - return (init_op, call_op, results_op) - # Eager case - for example in datasets.Iterator(dataset): - self.__call__(example, *args, **kwargs) - return self.all_metric_results(summary_logdir) + if context.executing_eagerly(): + for example in datasets.Iterator(dataset): + self.__call__(example, *args, **kwargs) + return self.all_metric_results(summary_logdir) + # Graph construction + call_op = self.__call__(dataset.make_one_shot_iterator().get_next(), *args, + **kwargs) + init_op = self.init_variables() + results_op = self.all_metric_results(summary_logdir) + return (init_op, call_op, results_op) @staticmethod def run_evaluation(init_op, call_op, results_op, sess=None): @@ -192,7 +193,7 @@ class Evaluator(object): Only for graph execution. @end_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Evaluator.run_evaluation() not supported when " "eager execution is enabled.") sess = sess or ops.get_default_session() diff --git a/tensorflow/contrib/eager/python/examples/BUILD b/tensorflow/contrib/eager/python/examples/BUILD index 15a21885f66eface291a39fa0ee1ff28bc297548..c1fd9e0ed020beeb722204edf1adfe1dfcf8ff03 100644 --- a/tensorflow/contrib/eager/python/examples/BUILD +++ b/tensorflow/contrib/eager/python/examples/BUILD @@ -8,7 +8,6 @@ py_library( 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", "//tensorflow/contrib/eager/python/examples/rnn_colorbot", "//tensorflow/contrib/eager/python/examples/rnn_ptb", diff --git a/tensorflow/contrib/eager/python/examples/gan/README.md b/tensorflow/contrib/eager/python/examples/gan/README.md index e8c9db1a1e2eb5881b08a4d3866c82b24d64be12..208a64b05d47eea10b49a1bf967a5453677bfd21 100644 --- a/tensorflow/contrib/eager/python/examples/gan/README.md +++ b/tensorflow/contrib/eager/python/examples/gan/README.md @@ -11,7 +11,7 @@ Other eager execution examples can be found under the parent directory. - `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 trainig and using the models using +- `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. diff --git a/tensorflow/contrib/eager/python/examples/gan/mnist.py b/tensorflow/contrib/eager/python/examples/gan/mnist.py index b9ac79f46c83bb709918e3b72830b90ddcfd71b4..2b7e199fad08c9a5e320b51b3a4de92c2d7dbb1a 100644 --- a/tensorflow/contrib/eager/python/examples/gan/mnist.py +++ b/tensorflow/contrib/eager/python/examples/gan/mnist.py @@ -35,7 +35,7 @@ from tensorflow.examples.tutorials.mnist import input_data FLAGS = None -class Discriminator(tfe.Network): +class Discriminator(tf.keras.Model): """GAN Discriminator. A network to differentiate between generated and real handwritten digits. @@ -56,19 +56,15 @@ class Discriminator(tfe.Network): 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)) + self.conv1 = tf.layers.Conv2D( + 64, 5, padding='SAME', data_format=data_format, activation=tf.tanh) + self.pool1 = tf.layers.AveragePooling2D(2, 2, data_format=data_format) + self.conv2 = tf.layers.Conv2D( + 128, 5, data_format=data_format, activation=tf.tanh) + self.pool2 = tf.layers.AveragePooling2D(2, 2, data_format=data_format) + self.flatten = tf.layers.Flatten() + self.fc1 = tf.layers.Dense(1024, activation=tf.tanh) + self.fc2 = tf.layers.Dense(1, activation=None) def call(self, inputs): """Return two logits per image estimating input authenticity. @@ -95,7 +91,7 @@ class Discriminator(tfe.Network): return x -class Generator(tfe.Network): +class Generator(tf.keras.Model): """Generator of handwritten digits similar to the ones in the MNIST dataset. """ @@ -116,18 +112,17 @@ class Generator(tfe.Network): 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)) + self.fc1 = 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)) + self.conv1 = 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)) + self.conv2 = 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. @@ -168,7 +163,8 @@ def discriminator_loss(discriminator_real_outputs, discriminator_gen_outputs): """ loss_on_real = tf.losses.sigmoid_cross_entropy( - tf.ones_like(discriminator_real_outputs), discriminator_real_outputs, + 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) @@ -198,9 +194,9 @@ def generator_loss(discriminator_gen_outputs): return loss -def train_one_epoch(generator, discriminator, - generator_optimizer, discriminator_optimizer, - dataset, log_interval, noise_dim): +def train_one_epoch(generator, discriminator, generator_optimizer, + discriminator_optimizer, dataset, step_counter, + log_interval, noise_dim): """Trains `generator` and `discriminator` models on `dataset`. Args: @@ -209,7 +205,8 @@ def train_one_epoch(generator, discriminator, 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 + step_counter: An integer variable, used to write summaries regularly. + log_interval: How many steps to wait between logging and collecting summaries. noise_dim: Dimension of noise vector to use. """ @@ -218,18 +215,23 @@ def train_one_epoch(generator, discriminator, 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) + tf.assign_add(step_counter, 1) - with tf.contrib.summary.record_summaries_every_n_global_steps(log_interval): + with tf.contrib.summary.record_summaries_every_n_global_steps( + log_interval, global_step=step_counter): current_batch_size = images.shape[0] - noise = tf.random_uniform(shape=[current_batch_size, noise_dim], - minval=-1., maxval=1., seed=batch_index) + 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) + 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) @@ -244,18 +246,16 @@ def train_one_epoch(generator, discriminator, 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)) + generator_optimizer.apply_gradients( + zip(generator_grad, generator.variables)) + 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)) + 'Average Discriminator Loss: %.6f' % + (batch_index, total_generator_loss / batch_index, + total_discriminator_loss / batch_index)) def main(_): @@ -266,18 +266,18 @@ def main(_): # 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) + dataset = ( + tf.data.Dataset.from_tensor_slices(data.train.images).shuffle(60000) + .batch(FLAGS.batch_size)) + + # Create the models and optimizers. + model_objects = { + 'generator': Generator(data_format), + 'discriminator': Discriminator(data_format), + 'generator_optimizer': tf.train.AdamOptimizer(FLAGS.lr), + 'discriminator_optimizer': tf.train.AdamOptimizer(FLAGS.lr), + 'step_counter': tf.train.get_or_create_global_step(), + } # Prepare summary writer and checkpoint info summary_writer = tf.contrib.summary.create_summary_file_writer( @@ -286,28 +286,22 @@ def main(_): latest_cpkt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir) if latest_cpkt: print('Using latest checkpoint at ' + latest_cpkt) + checkpoint = tfe.Checkpoint(**model_objects) + # Restore variables on creation if a checkpoint exists. + checkpoint.restore(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) + for _ in range(100): + start = time.time() + with summary_writer.as_default(): + train_one_epoch(dataset=dataset, log_interval=FLAGS.log_interval, + noise_dim=FLAGS.noise, **model_objects) + end = time.time() + checkpoint.save(checkpoint_prefix) + print('\nTrain time for epoch #%d (step %d): %f' % + (checkpoint.save_counter.numpy(), + checkpoint.step_counter.numpy(), + end - start)) if __name__ == '__main__': diff --git a/tensorflow/contrib/eager/python/examples/gan/mnist_test.py b/tensorflow/contrib/eager/python/examples/gan/mnist_test.py index 4a3ca8d82bc2619b05a734f6d2e58431c1a45995..bd35e50c1f434d167c5a8c5aa7d224912523ce28 100644 --- a/tensorflow/contrib/eager/python/examples/gan/mnist_test.py +++ b/tensorflow/contrib/eager/python/examples/gan/mnist_test.py @@ -62,7 +62,7 @@ class MnistEagerGanBenchmark(tf.test.Benchmark): for _ in range(measure_batches)] measure_dataset = tf.data.Dataset.from_tensor_slices(measure_images) - tf.train.get_or_create_global_step() + step_counter = tf.train.get_or_create_global_step() with tf.device(device()): # Create the models and optimizers generator = mnist.Generator(data_format()) @@ -78,13 +78,15 @@ class MnistEagerGanBenchmark(tf.test.Benchmark): # warm up mnist.train_one_epoch(generator, discriminator, generator_optimizer, discriminator_optimizer, - burn_dataset, log_interval=SUMMARY_INTERVAL, + burn_dataset, step_counter, + 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, + measure_dataset, step_counter, + log_interval=SUMMARY_INTERVAL, noise_dim=NOISE_DIM) self._report('train', start, measure_batches, batch_size) diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD index f86331af6f7928f0f86c888e22706c6e0a5978b2..2f6cfdf31e852d5d69a7a87980c9a441da504cf2 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD +++ b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD @@ -22,6 +22,7 @@ cuda_py_test( ":linear_regression", "//tensorflow:tensorflow_py", ], + tags = ["no_windows"], # TODO: needs investigation on Windows ) cuda_py_test( 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 6ce4de6ee0bf50400eff339ac04e132252a2b53e..6ab847cb78a09ab0a38beefff56f87d8314c0713 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py +++ b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py @@ -33,23 +33,13 @@ import tensorflow as tf import tensorflow.contrib.eager as tfe -class LinearModel(tfe.Network): - """A TensorFlow linear regression model. - - Uses TensorFlow's eager execution. - - 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 and returns the MSE loss value. - The `fit()` performs gradient-descent training on the model's weights - and bias. - """ +class LinearModel(tf.keras.Model): + """A TensorFlow linear regression model.""" def __init__(self): """Constructs a LinearModel object.""" super(LinearModel, self).__init__() - self._hidden_layer = self.track_layer(tf.layers.Dense(1)) + self._hidden_layer = tf.layers.Dense(1) def call(self, xs): """Invoke the linear model. @@ -64,7 +54,7 @@ class LinearModel(tfe.Network): def mean_square_loss(model, xs, ys): - return tf.reduce_mean(tf.square(model(xs) - ys)) + return tf.reduce_mean(tf.square(tf.subtract(model(xs), ys))) def fit(model, dataset, optimizer, verbose=False, logdir=None): diff --git a/tensorflow/contrib/eager/python/examples/mnist/BUILD b/tensorflow/contrib/eager/python/examples/mnist/BUILD deleted file mode 100644 index c61ec2dbae60a782c0e6589701554b045dcb92ae..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/examples/mnist/BUILD +++ /dev/null @@ -1,36 +0,0 @@ -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/mnist/README.md b/tensorflow/contrib/eager/python/examples/mnist/README.md index e987996b88ccf54a322749aadec4f9840760a90f..d1c079ff6b5cb187bbcfe2742293982b1bedd2d4 100644 --- a/tensorflow/contrib/eager/python/examples/mnist/README.md +++ b/tensorflow/contrib/eager/python/examples/mnist/README.md @@ -1,10 +1 @@ -Classification model for the MNIST dataset using eager execution. - -To run: - -``` -python mnist.py -``` - -`mnist_graph_test.py` demonstrates that the same code that is executed eagerly -in `mnist.py` is used to construct a TensorFlow graph. +See https://github.com/tensorflow/models/tree/master/official/mnist/mnist_eager.py diff --git a/tensorflow/contrib/eager/python/examples/mnist/mnist.py b/tensorflow/contrib/eager/python/examples/mnist/mnist.py deleted file mode 100644 index 2a7be95811f6fff06e2c489890703561ed879c42..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/examples/mnist/mnist.py +++ /dev/null @@ -1,266 +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. -# ============================================================================== -"""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 MNISTModel(tfe.Network): - """MNIST Network. - - Network structure is equivalent to: - https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/examples/tutorials/mnist/mnist_deep.py - and - https://github.com/tensorflow/models/blob/master/tutorials/image/mnist/convolutional.py - - But written using the tf.layers API. - """ - - def __init__(self, data_format): - """Creates a model for classifying a hand-written digit. - - 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(MNISTModel, 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(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): - """Computes labels from inputs. - - 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, 784]. - training: True if invoked in the context of training (causing dropout to - be applied). False otherwise. - - Returns: - A Tensor with shape [batch_size, 10] containing the predicted logits - for each image in the batch, for each of the 10 classes. - """ - - x = tf.reshape(inputs, self._input_shape) - x = self.conv1(x) - x = self.max_pool2d(x) - x = self.conv2(x) - x = self.max_pool2d(x) - x = tf.layers.flatten(x) - x = self.fc1(x) - if training: - x = self.dropout(x) - x = self.fc2(x) - return x - - -def loss(predictions, labels): - return tf.reduce_mean( - tf.nn.softmax_cross_entropy_with_logits( - logits=predictions, labels=labels)) - - -def compute_accuracy(predictions, labels): - return tf.reduce_sum( - tf.cast( - tf.equal( - tf.argmax(predictions, axis=1, - output_type=tf.int64), - tf.argmax(labels, axis=1, - output_type=tf.int64)), - dtype=tf.float32)) / float(predictions.shape[0].value) - - -def train_one_epoch(model, optimizer, dataset, log_interval=None): - """Trains model on `dataset` using `optimizer`.""" - - tf.train.get_or_create_global_step() - - for (batch, (images, labels)) in enumerate(tfe.Iterator(dataset)): - with tf.contrib.summary.record_summaries_every_n_global_steps(10): - 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, loss_value)) - - -def test(model, dataset): - """Perform an evaluation of `model` on the examples from `dataset`.""" - avg_loss = tfe.metrics.Mean('loss') - accuracy = tfe.metrics.Accuracy('accuracy') - - for (images, labels) in tfe.Iterator(dataset): - predictions = model(images, training=False) - avg_loss(loss(predictions, labels)) - accuracy(tf.argmax(predictions, axis=1, output_type=tf.int64), - tf.argmax(labels, axis=1, output_type=tf.int64)) - print('Test set: Average loss: %.4f, Accuracy: %4f%%\n' % - (avg_loss.result(), 100 * accuracy.result())) - with tf.contrib.summary.always_record_summaries(): - tf.contrib.summary.scalar('loss', avg_loss.result()) - tf.contrib.summary.scalar('accuracy', accuracy.result()) - - -def load_data(data_dir): - """Returns training and test tf.data.Dataset objects.""" - data = input_data.read_data_sets(data_dir, one_hot=True) - train_ds = tf.data.Dataset.from_tensor_slices((data.train.images, - data.train.labels)) - test_ds = tf.data.Dataset.from_tensors((data.test.images, data.test.labels)) - return (train_ds, test_ds) - - -def main(_): - tfe.enable_eager_execution() - - (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 - (train_ds, test_ds) = load_data(FLAGS.data_dir) - train_ds = train_ds.shuffle(60000).batch(FLAGS.batch_size) - - # Create the model and optimizer - model = MNISTModel(data_format) - optimizer = tf.train.MomentumOptimizer(FLAGS.lr, FLAGS.momentum) - - if FLAGS.output_dir: - train_dir = os.path.join(FLAGS.output_dir, 'train') - test_dir = os.path.join(FLAGS.output_dir, 'eval') - tf.gfile.MakeDirs(FLAGS.output_dir) - else: - train_dir = None - test_dir = None - summary_writer = tf.contrib.summary.create_file_writer( - train_dir, flush_millis=10000) - test_summary_writer = tf.contrib.summary.create_file_writer( - test_dir, flush_millis=10000, name='test') - checkpoint_prefix = os.path.join(FLAGS.checkpoint_dir, 'ckpt') - - with tf.device(device): - for epoch in range(1, 11): - with tfe.restore_variables_on_create( - tf.train.latest_checkpoint(FLAGS.checkpoint_dir)): - global_step = tf.train.get_or_create_global_step() - start = time.time() - with summary_writer.as_default(): - train_one_epoch(model, optimizer, train_ds, FLAGS.log_interval) - end = time.time() - print('\nTrain time for epoch #%d (global step %d): %f' % ( - epoch, global_step.numpy(), end - start)) - with test_summary_writer.as_default(): - test(model, test_ds) - all_variables = ( - model.variables - + optimizer.variables() - + [global_step]) - tfe.Saver(all_variables).save( - checkpoint_prefix, global_step=global_step) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument( - '--data-dir', - type=str, - default='/tmp/tensorflow/mnist/input_data', - help='Directory for storing input data') - parser.add_argument( - '--batch-size', - type=int, - default=64, - metavar='N', - help='input batch size for training (default: 64)') - parser.add_argument( - '--log-interval', - type=int, - default=10, - metavar='N', - help='how many batches to wait before logging training status') - parser.add_argument( - '--output_dir', - type=str, - default=None, - metavar='N', - help='Directory to write TensorBoard summaries') - parser.add_argument( - '--checkpoint_dir', - type=str, - default='/tmp/tensorflow/mnist/checkpoints/', - metavar='N', - help='Directory to save checkpoints in (once per epoch)') - parser.add_argument( - '--lr', - type=float, - default=0.01, - metavar='LR', - help='learning rate (default: 0.01)') - parser.add_argument( - '--momentum', - type=float, - default=0.5, - metavar='M', - help='SGD momentum (default: 0.5)') - 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/mnist/mnist_graph_test.py b/tensorflow/contrib/eager/python/examples/mnist/mnist_graph_test.py deleted file mode 100644 index 1af26553120b34d4682b17b1c29c81dc65e421d4..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/examples/mnist/mnist_graph_test.py +++ /dev/null @@ -1,65 +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 numpy as np -import tensorflow as tf -from tensorflow.contrib.eager.python.examples.mnist import mnist - - -def data_format(): - return "channels_first" if tf.test.is_gpu_available() else "channels_last" - - -class MNISTGraphTest(tf.test.TestCase): - - def testTrainGraph(self): - # The MNISTModel class can be executed eagerly (as in mnist.py and - # mnist_test.py) and also be used to construct a TensorFlow graph, which is - # then trained in a session. - with tf.Graph().as_default(): - # Generate some random data. - batch_size = 64 - images = np.random.randn(batch_size, 784).astype(np.float32) - digits = np.random.randint(low=0, high=10, size=batch_size) - labels = np.zeros((batch_size, 10)) - labels[np.arange(batch_size), digits] = 1. - - # Create a model, optimizer, and dataset as would be done - # for eager execution as well. - model = mnist.MNISTModel(data_format()) - optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0) - dataset = tf.data.Dataset.from_tensors((images, labels)) - - # Define the loss tensor (as opposed to a loss function when - # using eager execution). - (images, labels) = dataset.make_one_shot_iterator().get_next() - predictions = model(images, training=True) - loss = mnist.loss(predictions, labels) - - train_op = optimizer.minimize(loss) - init = tf.global_variables_initializer() - with tf.Session() as sess: - # Variables have to be initialized in the session. - sess.run(init) - # Train using the optimizer. - sess.run(train_op) - - -if __name__ == "__main__": - tf.test.main() diff --git a/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py b/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py deleted file mode 100644 index 136085eba21284a42282395e54f32c33bf63b5c3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py +++ /dev/null @@ -1,80 +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 tensorflow as tf - -import tensorflow.contrib.eager as tfe -from tensorflow.contrib.eager.python.examples.mnist import mnist - - -def device(): - return "/device:GPU:0" if tfe.num_gpus() else "/device:CPU:0" - - -def data_format(): - return "channels_first" if tfe.num_gpus() else "channels_last" - - -def random_dataset(): - batch_size = 64 - images = tf.random_normal([batch_size, 784]) - digits = tf.random_uniform([batch_size], minval=0, maxval=10, dtype=tf.int32) - labels = tf.one_hot(digits, 10) - 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): - train_one_epoch(defun=False) - - def testTest(self): - evaluate(defun=False) - - def testTrainOneEpochWithDefunCall(self): - train_one_epoch(defun=True) - - def testTestWithDefunCall(self): - evaluate(defun=True) - - -if __name__ == "__main__": - tfe.enable_eager_execution() - tf.test.main() diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py index 9982fdb07eefa665379e7be095f4f8017d92cf97..6b59413141f78fc85474850e109454ecdeb68cd3 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py +++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py @@ -27,10 +27,9 @@ from __future__ import print_function import functools import tensorflow as tf -import tensorflow.contrib.eager as tfe -class _IdentityBlock(tfe.Network): +class _IdentityBlock(tf.keras.Model): """_IdentityBlock is the block that has no conv layer at shortcut. Args: @@ -50,31 +49,24 @@ class _IdentityBlock(tfe.Network): bn_name_base = 'bn' + str(stage) + block + '_branch' bn_axis = 1 if data_format == 'channels_first' else 3 - self.conv2a = self.track_layer( - tf.layers.Conv2D( - filters1, (1, 1), - name=conv_name_base + '2a', - data_format=data_format)) - self.bn2a = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')) - - self.conv2b = self.track_layer( - tf.layers.Conv2D( - filters2, - kernel_size, - padding='same', - data_format=data_format, - name=conv_name_base + '2b')) - self.bn2b = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')) - - self.conv2c = self.track_layer( - tf.layers.Conv2D( - filters3, (1, 1), - name=conv_name_base + '2c', - data_format=data_format)) - self.bn2c = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')) + self.conv2a = tf.layers.Conv2D( + filters1, (1, 1), name=conv_name_base + '2a', data_format=data_format) + self.bn2a = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '2a') + + self.conv2b = tf.layers.Conv2D( + filters2, + kernel_size, + padding='same', + data_format=data_format, + name=conv_name_base + '2b') + self.bn2b = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '2b') + + self.conv2c = tf.layers.Conv2D( + filters3, (1, 1), name=conv_name_base + '2c', data_format=data_format) + self.bn2c = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '2c') def call(self, input_tensor, training=False): x = self.conv2a(input_tensor) @@ -92,7 +84,7 @@ class _IdentityBlock(tfe.Network): return tf.nn.relu(x) -class _ConvBlock(tfe.Network): +class _ConvBlock(tf.keras.Model): """_ConvBlock is the block that has a conv layer at shortcut. Args: @@ -121,41 +113,35 @@ class _ConvBlock(tfe.Network): bn_name_base = 'bn' + str(stage) + block + '_branch' bn_axis = 1 if data_format == 'channels_first' else 3 - self.conv2a = self.track_layer( - tf.layers.Conv2D( - filters1, (1, 1), - strides=strides, - name=conv_name_base + '2a', - data_format=data_format)) - self.bn2a = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')) - - self.conv2b = self.track_layer( - tf.layers.Conv2D( - filters2, - kernel_size, - padding='same', - name=conv_name_base + '2b', - data_format=data_format)) - self.bn2b = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')) - - self.conv2c = self.track_layer( - tf.layers.Conv2D( - filters3, (1, 1), - name=conv_name_base + '2c', - data_format=data_format)) - self.bn2c = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')) - - self.conv_shortcut = self.track_layer( - tf.layers.Conv2D( - filters3, (1, 1), - strides=strides, - name=conv_name_base + '1', - data_format=data_format)) - self.bn_shortcut = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '1')) + self.conv2a = tf.layers.Conv2D( + filters1, (1, 1), + strides=strides, + name=conv_name_base + '2a', + data_format=data_format) + self.bn2a = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '2a') + + self.conv2b = tf.layers.Conv2D( + filters2, + kernel_size, + padding='same', + name=conv_name_base + '2b', + data_format=data_format) + self.bn2b = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '2b') + + self.conv2c = tf.layers.Conv2D( + filters3, (1, 1), name=conv_name_base + '2c', data_format=data_format) + self.bn2c = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '2c') + + self.conv_shortcut = tf.layers.Conv2D( + filters3, (1, 1), + strides=strides, + name=conv_name_base + '1', + data_format=data_format) + self.bn_shortcut = tf.layers.BatchNormalization( + axis=bn_axis, name=bn_name_base + '1') def call(self, input_tensor, training=False): x = self.conv2a(input_tensor) @@ -176,7 +162,8 @@ class _ConvBlock(tfe.Network): return tf.nn.relu(x) -class ResNet50(tfe.Network): +# pylint: disable=not-callable +class ResNet50(tf.keras.Model): """Instantiates the ResNet50 architecture. Args: @@ -220,32 +207,28 @@ class ResNet50(tfe.Network): self.include_top = include_top def conv_block(filters, stage, block, strides=(2, 2)): - l = _ConvBlock( + return _ConvBlock( 3, filters, stage=stage, block=block, data_format=data_format, strides=strides) - return self.track_layer(l) def id_block(filters, stage, block): - l = _IdentityBlock( + return _IdentityBlock( 3, filters, stage=stage, block=block, data_format=data_format) - return self.track_layer(l) - - self.conv1 = self.track_layer( - tf.layers.Conv2D( - 64, (7, 7), - strides=(2, 2), - data_format=data_format, - padding='same', - name='conv1')) + + self.conv1 = tf.layers.Conv2D( + 64, (7, 7), + strides=(2, 2), + data_format=data_format, + padding='same', + name='conv1') bn_axis = 1 if data_format == 'channels_first' else 3 - self.bn_conv1 = self.track_layer( - tf.layers.BatchNormalization(axis=bn_axis, name='bn_conv1')) - self.max_pool = self.track_layer( - tf.layers.MaxPooling2D((3, 3), strides=(2, 2), data_format=data_format)) + self.bn_conv1 = tf.layers.BatchNormalization(axis=bn_axis, name='bn_conv1') + self.max_pool = tf.layers.MaxPooling2D( + (3, 3), strides=(2, 2), data_format=data_format) self.l2a = conv_block([64, 64, 256], stage=2, block='a', strides=(1, 1)) self.l2b = id_block([64, 64, 256], stage=2, block='b') @@ -267,13 +250,11 @@ class ResNet50(tfe.Network): self.l5b = id_block([512, 512, 2048], stage=5, block='b') self.l5c = id_block([512, 512, 2048], stage=5, block='c') - self.avg_pool = self.track_layer( - tf.layers.AveragePooling2D( - (7, 7), strides=(7, 7), data_format=data_format)) + self.avg_pool = tf.layers.AveragePooling2D( + (7, 7), strides=(7, 7), data_format=data_format) if self.include_top: - self.fc1000 = self.track_layer( - tf.layers.Dense(classes, name='fc1000')) + self.fc1000 = tf.layers.Dense(classes, name='fc1000') else: reduction_indices = [1, 2] if data_format == 'channels_last' else [2, 3] reduction_indices = tf.constant(reduction_indices) @@ -288,7 +269,7 @@ class ResNet50(tfe.Network): else: self.global_pooling = None - def call(self, input_tensor, training=False): + def call(self, input_tensor, training): x = self.conv1(input_tensor) x = self.bn_conv1(x, training=training) x = tf.nn.relu(x) diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py index 23317886e712323f4b520000e0fd372734fc53a1..551c76b0df71c88919df9cd6d81b4176b23b0ba3 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py +++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py @@ -55,7 +55,7 @@ class ResNet50GraphTest(tf.test.TestCase): with tf.Graph().as_default(): images = tf.placeholder(tf.float32, image_shape(None)) model = resnet50.ResNet50(data_format()) - predictions = model(images) + predictions = model(images, training=False) init = tf.global_variables_initializer() @@ -114,7 +114,7 @@ class ResNet50Benchmarks(tf.test.Benchmark): with tf.Graph().as_default(): images = tf.placeholder(tf.float32, image_shape(None)) model = resnet50.ResNet50(data_format()) - predictions = model(images) + predictions = model(images, training=False) init = tf.global_variables_initializer() diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py index 1f7beee68525e9cf338847caa0bb2dcc8bd60f62..d6923293a374f29ab77be70fa9fea44efd1ea40b 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py +++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py @@ -22,7 +22,7 @@ import gc import tempfile import time -from six.moves import xrange +from six.moves import xrange # pylint: disable=redefined-builtin import tensorflow as tf import tensorflow.contrib.eager as tfe @@ -64,28 +64,35 @@ def train_one_step(model, images, labels, optimizer): class ResNet50Test(tf.test.TestCase): - def _apply(self, defun=False): + def _apply(self, defun=False, execution_mode=None): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format) if defun: model.call = tfe.defun(model.call) - with tf.device(device): + with tf.device(device), tfe.execution_mode(execution_mode): images, _ = random_batch(2) - output = model(images) + output = model(images, training=False) + tfe.async_wait() self.assertEqual((2, 1000), output.shape) def test_apply(self): self._apply(defun=False) + def test_apply_async(self): + self._apply(defun=False, execution_mode=tfe.ASYNC) + def test_apply_with_defun(self): self._apply(defun=True) + def test_apply_with_defun_async(self): + self._apply(defun=True, execution_mode=tfe.ASYNC) + def test_apply_no_top(self): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format, include_top=False) with tf.device(device): images, _ = random_batch(2) - output = model(images) + output = model(images, training=False) output_shape = ((2, 2048, 1, 1) if data_format == 'channels_first' else (2, 1, 1, 2048)) self.assertEqual(output_shape, output.shape) @@ -95,10 +102,10 @@ class ResNet50Test(tf.test.TestCase): model = resnet50.ResNet50(data_format, include_top=False, pooling='avg') with tf.device(device): images, _ = random_batch(2) - output = model(images) + output = model(images, training=False) self.assertEqual((2, 2048), output.shape) - def test_train(self): + def _test_train(self, execution_mode=None): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format) tf.train.get_or_create_global_step() @@ -106,15 +113,22 @@ class ResNet50Test(tf.test.TestCase): with tf.contrib.summary.create_file_writer( logdir, max_queue=0, name='t0').as_default(), tf.contrib.summary.always_record_summaries(): - with tf.device(device): + with tf.device(device), tfe.execution_mode(execution_mode): optimizer = tf.train.GradientDescentOptimizer(0.1) images, labels = random_batch(2) train_one_step(model, images, labels, optimizer) self.assertEqual(320, len(model.variables)) + tfe.async_wait() events = summary_test_util.events_from_logdir(logdir) self.assertEqual(len(events), 2) self.assertEqual(events[1].summary.value[0].tag, 'loss') + def test_train(self): + self._test_train() + + def test_train_async(self): + self._test_train(execution_mode=tfe.ASYNC) + def test_no_garbage(self): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format) @@ -183,59 +197,84 @@ class ResNet50Benchmarks(tf.test.Benchmark): # a sync. This is a roundabout way, yes. tf.constant(1.).cpu() - 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 - with tf.device(device): - images, _ = random_batch(batch_size) - for _ in xrange(num_burn): - model(images).cpu() - gc.collect() - start = time.time() - for _ in xrange(num_iters): - model(images).cpu() - 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, 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 + def _benchmark_eager_apply(self, label, defun=False, execution_mode=None): + with tfe.execution_mode(execution_mode): + device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format) if defun: model.call = tfe.defun(model.call) - optimizer = tf.train.GradientDescentOptimizer(0.1) - + batch_size = 64 + num_burn = 5 + num_iters = 30 with tf.device(device): - iterator = make_iterator((images, labels)) + images, _ = random_batch(batch_size) for _ in xrange(num_burn): - (images, labels) = iterator.next() - train_one_step(model, images, labels, optimizer) - self._force_gpu_sync() + model(images, training=False).cpu() + if execution_mode: + tfe.async_wait() gc.collect() - start = time.time() for _ in xrange(num_iters): - (images, labels) = iterator.next() - train_one_step(model, images, labels, optimizer) - self._force_gpu_sync() + model(images, training=False).cpu() + if execution_mode: + tfe.async_wait() 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_async(self): + self._benchmark_eager_apply( + 'eager_apply_async', defun=False, execution_mode=tfe.ASYNC) + + 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, + defun=False, + execution_mode=None): + with tfe.execution_mode(execution_mode): + 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): + iterator = make_iterator((images, labels)) + for _ in xrange(num_burn): + (images, labels) = iterator.next() + train_one_step(model, images, labels, optimizer) + if execution_mode: + tfe.async_wait() + self._force_gpu_sync() + gc.collect() + + start = time.time() + for _ in xrange(num_iters): + (images, labels) = iterator.next() + train_one_step(model, images, labels, optimizer) + if execution_mode: + tfe.async_wait() + self._force_gpu_sync() + self._report(label, start, num_iters, device, batch_size, data_format) + def benchmark_eager_train(self): self._benchmark_eager_train('eager_train', MockIterator, defun=False) + def benchmark_eager_train_async(self): + self._benchmark_eager_train( + 'eager_train_async', + MockIterator, + defun=False, + execution_mode=tfe.ASYNC) + def benchmark_eager_train_with_defun(self): self._benchmark_eager_train( 'eager_train_with_defun', MockIterator, defun=True) 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..29f02324544ede172500f799cd84068984d7d87b 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 @@ -110,7 +109,7 @@ def load_dataset(data_dir, url, batch_size): # pylint: disable=not-callable -class RNNColorbot(tfe.Network): +class RNNColorbot(tf.keras.Model): """Multi-layer (LSTM) RNN that regresses on real-valued vector labels. """ @@ -128,23 +127,20 @@ class RNNColorbot(tfe.Network): self.label_dimension = label_dimension self.keep_prob = keep_prob - # Note the calls to `track_layer` below; these calls register the layers as - # network components that house trainable variables. - self.cells = [ - self.track_layer(tf.nn.rnn_cell.BasicLSTMCell(size)) - for size in rnn_cell_sizes - ] - self.relu = self.track_layer( - tf.layers.Dense(label_dimension, activation=tf.nn.relu, name="relu")) + self.cells = self._add_cells( + [tf.nn.rnn_cell.BasicLSTMCell(size) for size in rnn_cell_sizes]) + self.relu = tf.layers.Dense( + label_dimension, activation=tf.nn.relu, name="relu") - def call(self, chars, sequence_length, training=False): + def call(self, inputs, training=False): """Implements the RNN logic and prediction generation. Args: - chars: a Tensor of dimension [batch_size, time_steps, 256] holding a - batch of one-hot encoded color names - sequence_length: a Tensor of dimension [batch_size] holding the length - of each character sequence (i.e., color name) + inputs: A tuple (chars, sequence_length), where chars is a batch of + one-hot encoded color names represented as a Tensor with dimensions + [batch_size, time_steps, 256] and sequence_length holds the length + of each character sequence (color name) as a Tensor with dimension + [batch_size]. training: whether the invocation is happening during training Returns: @@ -152,6 +148,7 @@ class RNNColorbot(tfe.Network): passing chars through a multi-layer RNN and applying a ReLU to the final hidden state. """ + (chars, sequence_length) = inputs # Transpose the first and second dimensions so that chars is of shape # [time_steps, batch_size, dimension]. chars = tf.transpose(chars, [1, 0, 2]) @@ -182,6 +179,14 @@ class RNNColorbot(tfe.Network): hidden_states = tf.gather_nd(chars, indices) return self.relu(hidden_states) + def _add_cells(self, cells): + # "Magic" required for keras.Model classes to track all the variables in + # a list of tf.layers.Layer objects. + # TODO(ashankar): Figure out API so user code doesn't have to do this. + for i, c in enumerate(cells): + setattr(self, "cell-%d" % i, c) + return cells + def loss(labels, predictions): """Computes mean squared loss.""" @@ -192,7 +197,7 @@ def test(model, eval_data): """Computes the average loss on eval_data, which should be a Dataset.""" avg_loss = tfe.metrics.Mean("loss") for (labels, chars, sequence_length) in tfe.Iterator(eval_data): - predictions = model(chars, sequence_length, training=False) + predictions = model((chars, sequence_length), training=False) avg_loss(loss(labels, predictions)) print("eval/loss: %.6f\n" % avg_loss.result()) with tf.contrib.summary.always_record_summaries(): @@ -205,7 +210,7 @@ def train_one_epoch(model, optimizer, train_data, log_interval=10): tf.train.get_or_create_global_step() def model_loss(labels, chars, sequence_length): - predictions = model(chars, sequence_length, training=True) + predictions = model((chars, sequence_length), training=True) loss_value = loss(labels, predictions) tf.contrib.summary.scalar("loss", loss_value) return loss_value @@ -278,7 +283,7 @@ def main(_): (chars, length) = (tf.identity(chars), tf.identity(length)) chars = tf.expand_dims(chars, 0) length = tf.expand_dims(length, 0) - preds = tf.unstack(model(chars, length, training=False)[0]) + preds = tf.unstack(model((chars, length), training=False)[0]) # Predictions cannot be negative, as they are generated by a ReLU layer; # they may, however, be greater than 1. 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 d34e9ea68b76373d4b5a9ee9e3852c60a7c81525..69cd16d12c32c8c7c4744d8f0b4b1feedf946aa1 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py +++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py @@ -39,21 +39,23 @@ from tensorflow.contrib.cudnn_rnn.python.layers import cudnn_rnn from tensorflow.contrib.eager.python import tfe -class RNN(tfe.Network): +class RNN(tf.keras.Model): """A static RNN. - Similar to tf.nn.static_rnn, implemented as a tf.layer.Layer. + Similar to tf.nn.static_rnn, implemented as a class. """ def __init__(self, hidden_dim, num_layers, keep_ratio): super(RNN, self).__init__() self.keep_ratio = keep_ratio - for _ in range(num_layers): - self.track_layer(tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_dim)) + self.cells = self._add_cells([ + tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_dim) + for _ in range(num_layers) + ]) def call(self, input_seq, training): batch_size = int(input_seq.shape[1]) - for c in self.layers: + for c in self.cells: state = c.zero_state(batch_size, tf.float32) outputs = [] input_seq = tf.unstack(input_seq, num=int(input_seq.shape[0]), axis=0) @@ -64,7 +66,19 @@ class RNN(tfe.Network): input_seq = tf.stack(outputs, axis=0) if training: input_seq = tf.nn.dropout(input_seq, self.keep_ratio) - return input_seq, None + # Returning a list instead of a single tensor so that the line: + # y = self.rnn(y, ...)[0] + # in PTBModel.call works for both this RNN and CudnnLSTM (which returns a + # tuple (output, output_states). + return [input_seq] + + def _add_cells(self, cells): + # "Magic" required for keras.Model classes to track all the variables in + # a list of tf.layers.Layer objects. + # TODO(ashankar): Figure out API so user code doesn't have to do this. + for i, c in enumerate(cells): + setattr(self, "cell-%d" % i, c) + return cells class Embedding(tf.layers.Layer): @@ -87,7 +101,8 @@ class Embedding(tf.layers.Layer): return tf.nn.embedding_lookup(self.embedding, x) -class PTBModel(tfe.Network): +# pylint: disable=not-callable +class PTBModel(tf.keras.Model): """LSTM for word language modeling. Model described in: @@ -109,19 +124,16 @@ class PTBModel(tfe.Network): self.keep_ratio = 1 - dropout_ratio self.use_cudnn_rnn = use_cudnn_rnn - self.embedding = self.track_layer(Embedding(vocab_size, embedding_dim)) + self.embedding = Embedding(vocab_size, embedding_dim) if self.use_cudnn_rnn: self.rnn = cudnn_rnn.CudnnLSTM( num_layers, hidden_dim, dropout=dropout_ratio) else: self.rnn = RNN(hidden_dim, num_layers, self.keep_ratio) - self.track_layer(self.rnn) - self.linear = self.track_layer( - tf.layers.Dense( - vocab_size, - kernel_initializer=tf.random_uniform_initializer(-0.1, 0.1))) + self.linear = tf.layers.Dense( + vocab_size, kernel_initializer=tf.random_uniform_initializer(-0.1, 0.1)) self._output_shape = [-1, embedding_dim] def call(self, input_seq, training): @@ -136,7 +148,7 @@ class PTBModel(tfe.Network): y = self.embedding(input_seq) if training: y = tf.nn.dropout(y, self.keep_ratio) - y, _ = self.rnn(y, training=training) + y = self.rnn(y, training=training)[0] return self.linear(tf.reshape(y, self._output_shape)) @@ -148,7 +160,7 @@ def clip_gradients(grads_and_vars, clip_ratio): def loss_fn(model, inputs, targets, training): labels = tf.reshape(targets, [-1]) - outputs = model(inputs, training) + outputs = model(inputs, training=training) return tf.reduce_mean( tf.nn.sparse_softmax_cross_entropy_with_logits( labels=labels, logits=outputs)) @@ -339,8 +351,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 epochs.") + 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/BUILD b/tensorflow/contrib/eager/python/examples/spinn/BUILD index a1f8a759e2a556bc219f0aa13942f293c4f34cfa..98d01ad1d5a70788d2d4cb07031a8d76a6bf628f 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/BUILD +++ b/tensorflow/contrib/eager/python/examples/spinn/BUILD @@ -38,5 +38,8 @@ cuda_py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:framework_test_lib", ], - tags = ["no_pip"], # because spinn.py is under third_party/. + tags = [ + "no_cuda_on_cpu_tap", + "no_pip", # because spinn.py is under third_party/. + ], ) diff --git a/tensorflow/contrib/eager/python/examples/spinn/data.py b/tensorflow/contrib/eager/python/examples/spinn/data.py index fcaae0a4f8c0bad916d74bd9b80fcfa55a63d84a..3bc3bb49bcbbc26f7a3134a8bfc385ec080dde1e 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/data.py +++ b/tensorflow/contrib/eager/python/examples/spinn/data.py @@ -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.""" 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 19b0104c807cb43b16c3cd47dbacee6e890021db..081b0af14fcc983a3f85d2a50e2bb04d2f2493b3 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py +++ b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py @@ -26,7 +26,7 @@ import tempfile import time import numpy as np -from six.moves import xrange +from six.moves import xrange # pylint: disable=redefined-builtin import tensorflow as tf # pylint: disable=g-bad-import-order @@ -36,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 @@ -66,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, @@ -86,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): @@ -288,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) @@ -337,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( @@ -351,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. @@ -363,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 ffc1d0332eae605ce0444a225e53baa68954cae0..b73dc17e5f9cb15a51426f85e966a49604145f1d 100644 --- a/tensorflow/contrib/eager/python/g3doc/guide.md +++ b/tensorflow/contrib/eager/python/g3doc/guide.md @@ -22,11 +22,10 @@ to models defined without using eager execution. Eager execution is included in TensorFlow versions 1.5 and above. Installation instructions at https://www.tensorflow.org/install/ -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 either [building from -source](https://www.tensorflow.org/install/install_sources) -or the try latest nightly builds. The nightly builds are available as: +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` packages](https://github.com/tensorflow/tensorflow/blob/master/README.md#installation) and @@ -570,54 +569,50 @@ for i in range(20001): print("Loss on test set: %f" % loss(model, data.test.images, data.test.labels).numpy()) ``` -For a more complete example, see -[`tensorflow/contrib/eager/python/examples/mnist.py`](https://www.tensorflow.org/code/tensorflow/contrib/eager/python/examples/mnist/mnist.py) +For a more complete example, see [the example in the tensorflow/models +repository](https://github.com/tensorflow/models/tree/master/official/mnist/mnist_eager.py). ### Checkpointing trained variables -TensorFlow Variables (`tfe.Variable`) provides a way to represent shared, -persistent state of your model. The `tfe.Saver` class (which is a thin wrapper -over the -[`tf.train.Saver`](https://www.tensorflow.org/api_docs/python/tf/train/Saver) -class) provides a means to save and restore variables to and from _checkpoints_. +TensorFlow Variables (`tfe.Variable`) provide a way to represent shared, +persistent state of your model. The `tfe.Checkpoint` class provides a means to +save and restore variables to and from _checkpoints_. For example: ```python # Create variables. -x = tfe.Variable(10., name='x') -y = tfe.Variable(5., name='y') +x = tfe.Variable(10.) +y = tfe.Variable(5.) -# Create a Saver. -saver = tfe.Saver([x, y]) +# Indicate that the variables should be saved as "x" and "y". +checkpoint = tfe.Checkpoint(x=x, y=y) # Assign new values to the variables and save. x.assign(2.) -saver.save('/tmp/ckpt') +checkpoint.save('/tmp/ckpt') # Change the variable after saving. x.assign(11.) assert 16. == (x + y).numpy() # 11 + 5 # Restore the values in the checkpoint. -saver.restore('/tmp/ckpt') +checkpoint.restore('/tmp/ckpt-1') assert 7. == (x + y).numpy() # 2 + 5 ``` -### `tfe.Network` +### `tf.keras.Model` You may often want to organize your models using classes, like the `MNISTModel` -class described above. We recommend inheriting from the `tfe.Network` class as -it provides conveniences like keeping track of all model variables and methods -to save and restore from checkpoints. +class described above. We recommend inheriting from the `tf.keras.Model` class +as it provides conveniences like keeping track of all model variables. -Sub-classes of `tfe.Network` may register `Layer`s (like classes in -[`tf.layers`](https://www.tensorflow.org/api_docs/python/tf/layers), -or [Keras -layers](https://www.tensorflow.org/api_docs/python/tf/keras/layers)) -using a call to `self.track_layer()` and define the computation in an -implementation of `call()`. +Sub-classes of `tf.keras.Model` may register `Layer`s (like classes in +[`tf.layers`](https://www.tensorflow.org/api_docs/python/tf/layers), or [Keras +layers](https://www.tensorflow.org/api_docs/python/tf/keras/layers)) by +assigning them to attributes (`self.name = layer_object`) and define the +computation in an implementation of `call()`. Note that `tf.layers.Layer` objects (like `tf.layers.Dense`) create variables lazily, when the first input is encountered. @@ -625,12 +620,11 @@ lazily, when the first input is encountered. For example, consider the following two-layer neural network: ```python -class TwoLayerNet(tfe.Network): +class TwoLayerNet(tf.keras.Model): def __init__(self): super(TwoLayerNet, self).__init__() - self.layer1 = self.track_layer( - tf.layers.Dense(2, activation=tf.nn.relu, use_bias=False)) - self.layer2 = self.track_layer(tf.layers.Dense(3, use_bias=False)) + self.layer1 = tf.layers.Dense(2, activation=tf.nn.relu, use_bias=False) + self.layer2 = tf.layers.Dense(3, use_bias=False) def call(self, x): return self.layer2(self.layer1(x)) @@ -654,15 +648,16 @@ assert [1, 2] == net.variables[0].shape.as_list() # weights of layer1. assert [2, 3] == net.variables[1].shape.as_list() # weights of layer2. ``` -The `tfe.Network` class is itself a sub-class of `tf.layers.Layer`. This allows -instances of `tfe.Network` to be embedded in other networks. For example: +The `tf.keras.Model` class is itself a sub-class of `tf.layers.Layer`. This +allows instances of `tf.keras.Model` to be embedded in other models. For +example: ```python -class ThreeLayerNet(tfe.Network): +class ThreeLayerNet(tf.keras.Model): def __init__(self): super(ThreeLayerNet, self).__init__() - self.a = self.track_layer(TwoLayerNet()) - self.b = self.track_layer(tf.layers.Dense(4, use_bias=False)) + self.a = TwoLayerNet() + self.b = tf.layers.Dense(4, use_bias=False) def call(self, x): return self.b(self.a(x)) @@ -679,9 +674,8 @@ assert [3, 4] == net.variables[2].shape.as_list() See more examples in [`tensorflow/contrib/eager/python/examples`](https://www.tensorflow.org/code/tensorflow/contrib/eager/python/examples). -`tfe.Saver` in combination with `tfe.restore_variables_on_create` provides a -convenient way to save and load checkpoints without changing the program once -the checkpoint has been created. For example, we can set an objective for the +`tfe.Checkpoint` provides a convenient way to save and load training +checkpoints. Let's define something simple to train. We set an objective for the output of our network, choose an optimizer, and a location for the checkpoint: ```python @@ -692,30 +686,27 @@ checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt') net = ThreeLayerNet() ``` -Note that variables have not been created yet. We want them to be restored from -a checkpoint, if one exists, so we create them inside a -`tfe.restore_variables_on_create` context manager. Then our training loop is the -same whether starting training or resuming from a previous checkpoint: +We group them in a `tfe.Checkpoint` and request that it be restored. This +ensures that variables created by these objects are restored before their values +are used. Our training loop is the same whether starting training or resuming +from a previous checkpoint: ```python -with tfe.restore_variables_on_create( - tf.train.latest_checkpoint(checkpoint_directory)): - global_step = tf.train.get_or_create_global_step() - for _ in range(100): - loss_fn = lambda: tf.norm(net(inp) - objective) - optimizer.minimize(loss_fn, global_step=global_step) - if tf.equal(global_step % 20, 0): - print("Step %d, output %s" % (global_step.numpy(), - net(inp).numpy())) - all_variables = ( - net.variables - + optimizer.variables() - + [global_step]) - # Save the checkpoint. - tfe.Saver(all_variables).save(checkpoint_prefix, global_step=global_step) -``` - -The first time it runs, `Network` variables are initialized randomly. Then the +global_step = tf.train.get_or_create_global_step() +checkpoint = tfe.Checkpoint( + global_step=global_step, optimizer=optimizer, network=net) +checkpoint.restore(tf.train.latest_checkpoint(checkpoint_directory)) +for _ in range(100): + loss_fn = lambda: tf.norm(net(inp) - objective) + optimizer.minimize(loss_fn, global_step=global_step) + if tf.equal(global_step % 20, 0): + print("Step %d, output %s" % (global_step.numpy(), + net(inp).numpy())) + # Save the checkpoint. + checkpoint.save(checkpoint_prefix) +``` + +The first time it runs, `Model` variables are initialized randomly. Then the output is trained to match the objective we've set: ``` @@ -860,11 +851,9 @@ eagerly or constructing graphs. This means that you can iteratively develop your model with eager execution enabled and later, if needed, use the same code to reap the benefits of representing models as computational graphs. -For example, -[`mnist.py`](https://www.tensorflow.org/code/tensorflow/contrib/eager/python/examples/mnist/mnist.py) -defines a model that is eagerly executed. That same code is used to construct -and execute a graph in -[`mnist_graph_test.py`](https://www.tensorflow.org/code/tensorflow/contrib/eager/python/examples/mnist/mnist_graph_test.py). +For example, the same model definition used to construct a graph in +[mnist.py`](https://github.com/tensorflow/models/tree/master/official/mnist/mnist.py) +can be trained with eager execution enabled as in [`mnist_eager.py`](https://github.com/tensorflow/models/tree/master/official/mnist/mnist_eager.py). Other models in the [examples directory](https://www.tensorflow.org/code/tensorflow/contrib/eager/python/examples/) diff --git a/tensorflow/contrib/eager/python/metrics_impl.py b/tensorflow/contrib/eager/python/metrics_impl.py index ea8dbf2b46ea4bd0e33645ae3c590c4dd13f7a52..1490c2ccacd55156bcc1cf8c07d9941336e18e1b 100644 --- a/tensorflow/contrib/eager/python/metrics_impl.py +++ b/tensorflow/contrib/eager/python/metrics_impl.py @@ -30,12 +30,12 @@ from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope - +from tensorflow.python.training import checkpointable _to_replace = re.compile("[^A-Za-z0-9.]") -class Metric(object): +class Metric(checkpointable.CheckpointableBase): """A metric holds state for aggregating statistics over an evaluation run. Example use with eager execution: @@ -93,11 +93,12 @@ class Metric(object): `aggregate()`, it is for use by TensorFlow infrastructure. """ - def __init__(self, name=None): + def __init__(self, name=None, use_global_variables=False): self._built = False self._vars = [] self._initial_values = {} self._updates = [] + self._use_global_variables = use_global_variables name = name or self.__class__.__name__ # Replace things like spaces in name to create a valid scope name. scope_name = _to_replace.sub("_", name) @@ -108,13 +109,13 @@ class Metric(object): pos = scope.name.rfind(scope_name) self._name = name + scope.name[pos + len(scope_name):] self._scope = scope - if context.in_graph_mode(): + if context.executing_eagerly(): + self._construction_scope = context.eager_mode + else: # We make self.call() into a graph callable here, so that we can # return a single op that performs all of the variable updates. self._construction_scope = ops.get_default_graph().as_default self.call = function.defun(self.call) - else: - self._construction_scope = context.eager_mode # ---- API for users ---- def __call__(self, *args, **kwargs): @@ -155,10 +156,11 @@ class Metric(object): initialization. Under eager execution, the variables are reset to their initial values as a side effect and this function returns None. """ - if context.in_graph_mode(): + if context.executing_eagerly(): + for v in self._vars: + v.assign(self._initial_values[v]) + else: return control_flow_ops.group([v.initializer for v in self._vars]) - for v in self._vars: - v.assign(self._initial_values[v]) # ---- To be implemented by descendants --- def build(self, *args, **kwargs): @@ -200,10 +202,10 @@ class Metric(object): def value(self): """In graph mode returns the result Tensor while in eager the callable.""" - if context.in_graph_mode(): - return self.result() - else: + if context.executing_eagerly(): return self.result + else: + return self.result() # We can support two different strategies of for doing data-parallel # distributed metric computations: @@ -245,19 +247,31 @@ class Metric(object): """***Only for use by descendants of Metric***.""" if self._built: raise RuntimeError("Can't call add_variable() except in build().") - collections = None if context.in_eager_mode() else [ - ops.GraphKeys.LOCAL_VARIABLES, ops.GraphKeys.METRIC_VARIABLES - ] - v = variable_scope.get_variable( - name, - shape, - dtype, - initializer, + if context.executing_eagerly(): + collections = None + else: + if self._use_global_variables: + collections = [ops.GraphKeys.GLOBAL_VARIABLES] + else: + collections = [ops.GraphKeys.LOCAL_VARIABLES] + collections += [ops.GraphKeys.METRIC_VARIABLES] + # Variables are Checkpointable dependencies of Metrics regardless of the + # global/local distinction. Users can avoid saving variables by not adding a + # dependency on the Metric. + v = self._add_variable_with_custom_getter( + name=name, + shape=shape, + dtype=dtype, + initializer=initializer, trainable=False, collections=collections, - use_resource=True) + use_resource=True, + getter=variable_scope.get_variable, + # Raise duplicate variable exceptions from get_variable rather than + # Checkpointable. + overwrite=True) self._vars.append(v) - if context.in_eager_mode(): + if context.executing_eagerly(): self._initial_values[v] = v.value() return v @@ -267,8 +281,10 @@ class Mean(Metric): # TODO(josh11b): Maybe have a dtype argument that defaults to tf.float64? # Or defaults to type of the input if it is tf.float32, else tf.float64? - def __init__(self, name=None, dtype=dtypes.float64): - super(Mean, self).__init__(name=name) + def __init__(self, name=None, dtype=dtypes.float64, + use_global_variables=False): + super(Mean, self).__init__(name=name, + use_global_variables=use_global_variables) self.dtype = dtype def build(self, *args, **kwargs): diff --git a/tensorflow/contrib/eager/python/metrics_test.py b/tensorflow/contrib/eager/python/metrics_test.py index a9ecaa3f8bced3043ea0eb0ac3aa8bfa65e9e1ff..6b5450ba89bdfa6e0195f488b75f596b58c463d5 100644 --- a/tensorflow/contrib/eager/python/metrics_test.py +++ b/tensorflow/contrib/eager/python/metrics_test.py @@ -18,8 +18,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os import tempfile +from tensorflow.contrib.eager.python import checkpointable_utils from tensorflow.contrib.eager.python import metrics from tensorflow.contrib.summary import summary_ops from tensorflow.contrib.summary import summary_test_util @@ -50,6 +52,19 @@ class MetricsTest(test.TestCase): self.assertEqual( set(m.variables), set(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES))) + self.assertEqual(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES), []) + self.assertEqual( + set(m.variables), + set(ops.get_collection(ops.GraphKeys.METRIC_VARIABLES))) + + def testUseGlobalVariablesCollections(self): + with context.graph_mode(), ops.Graph().as_default(): + m = metrics.Mean(use_global_variables=True) + m(1000) + self.assertEqual( + set(m.variables), + set(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))) + self.assertEqual(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES), []) self.assertEqual( set(m.variables), set(ops.get_collection(ops.GraphKeys.METRIC_VARIABLES))) @@ -193,6 +208,31 @@ class MetricsTest(test.TestCase): self.assertAllEqual(m2.result().eval(), 2.0) self.assertAllEqual(m1.result().eval(), 1.0) + @test_util.run_in_graph_and_eager_modes() + def testSaveRestore(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + mean = metrics.Mean() + checkpoint = checkpointable_utils.Checkpoint(mean=mean) + mean.build() + mean._built = True + self.evaluate(mean.init_variables()) + self.evaluate(mean(100.)) + self.evaluate(mean(200.)) + save_path = checkpoint.save(checkpoint_prefix) + self.evaluate(mean(1000.)) + checkpoint.restore(save_path).assert_consumed().run_restore_ops() + self.evaluate(mean(300.)) + self.assertAllEqual(200., self.evaluate(mean.value())) + + restore_mean = metrics.Mean() + restore_checkpoint = checkpointable_utils.Checkpoint(mean=restore_mean) + status = restore_checkpoint.restore(save_path) + restore_update = restore_mean(300.) + status.assert_consumed().run_restore_ops() + self.evaluate(restore_update) + self.assertAllEqual(200., self.evaluate(restore_mean.value())) + self.assertEqual(3, self.evaluate(restore_mean.denom)) if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/eager/python/network.py b/tensorflow/contrib/eager/python/network.py index e3c13cbd2e8ccd2ab79da74e0e97905c6ed5c02d..4c937716e8df7c8cda26d6431885ce33346b77fb 100644 --- a/tensorflow/contrib/eager/python/network.py +++ b/tensorflow/contrib/eager/python/network.py @@ -639,7 +639,7 @@ def _make_custom_getter_for_deferred_restorations(): # Mark as already restored from this checkpoint. delayed_restoration.checkpointed_variables_to_restore[ checkpoint_name] = None - if context.in_graph_mode(): + if not context.executing_eagerly(): delayed_restoration.session.run(variable.initializer) if found_value: # Error checking should run even if we've already restored a value. @@ -772,7 +772,7 @@ def save_network_checkpoint( variable_map[mapped_name]._shared_name, variable._shared_name, network.scope_name)) - if context.in_eager_mode(): + if context.executing_eagerly(): sess = None else: sess = ops.get_default_session() @@ -853,7 +853,7 @@ def _restore_existing_variables(network, save_path, map_func, user_map_func): network_name=network.name, network_scope_name=network.scope_name)) if existing_variables_by_checkpoint_name: - if context.in_eager_mode(): + if context.executing_eagerly(): sess = None else: sess = ops.get_default_session() @@ -880,7 +880,7 @@ def _set_restore_on_create(network, save_path, map_func, user_map_func, # _DeferredRestoration objects once a Network has been built (so that # restoring in a loop does not take increasing amounts of memory). if checkpointed_variables_to_restore: - if context.in_eager_mode(): + if context.executing_eagerly(): sess = None else: sess = ops.get_default_session() diff --git a/tensorflow/contrib/eager/python/network_test.py b/tensorflow/contrib/eager/python/network_test.py index 81c77e41acf420fa84857ccb366aa2fbd6055f42..3329fc6c513265deff41a368f5688dd605209c14 100644 --- a/tensorflow/contrib/eager/python/network_test.py +++ b/tensorflow/contrib/eager/python/network_test.py @@ -539,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): diff --git a/tensorflow/contrib/eager/python/saver.py b/tensorflow/contrib/eager/python/saver.py index 62421849c766a1124c726812428985c913c653a3..fdaca90fd13576e6ca8a3408aaf528dbc2384b0c 100644 --- a/tensorflow/contrib/eager/python/saver.py +++ b/tensorflow/contrib/eager/python/saver.py @@ -73,7 +73,7 @@ def restore_variables_on_create(save_path, map_func=None): NotFoundError: If the variable is not found in checkpoint. ValueError: If not used in eager mode or map_func is not callable. """ - if context.in_graph_mode(): + if not context.executing_eagerly(): raise ValueError( "Currently, restore_variables_on_create can only be used with " "eager execution enabled.") @@ -131,7 +131,7 @@ class Saver(object): Raises: RuntimeError: if invoked when eager execution has not been enabled. """ - if context.in_graph_mode(): + if not context.executing_eagerly(): raise RuntimeError("tfe.Saver can only be used when eager " "execution is enabled. Use tf.train.Saver when " "building graphs.") 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 712d1cb94d2f565bf6216f6c07a45d3d855efe9c..c6f3f20e781147140f2c4b339ed465ab7e919d37 100644 --- a/tensorflow/contrib/eager/python/tfe.py +++ b/tensorflow/contrib/eager/python/tfe.py @@ -56,15 +56,24 @@ To use, at program startup, call `tfe.enable_eager_execution()`. @@save_network_checkpoint @@restore_network_checkpoint +@@Checkpoint +@@Checkpointable +@@CheckpointableSaver + +@@executing_eagerly @@in_eager_mode -@@in_graph_mode +@@set_execution_mode +@@execution_mode +@@async_wait +@@async_clear_error -@@IsolateTest @@run_test_in_graph_and_eager_modes @@DEVICE_PLACEMENT_EXPLICIT @@DEVICE_PLACEMENT_WARN @@DEVICE_PLACEMENT_SILENT +@@SYNC +@@ASYNC """ from __future__ import absolute_import @@ -75,6 +84,8 @@ from __future__ import print_function # pylint:disable=g-bad-import-order,g-import-not-at-top,unused-import # from tensorflow.contrib.eager.python import metrics +from tensorflow.contrib.eager.python.checkpointable_utils import CheckpointableSaver +from tensorflow.contrib.eager.python.checkpointable_utils import Checkpoint from tensorflow.contrib.eager.python.datasets import Iterator from tensorflow.contrib.eager.python.network import Network from tensorflow.contrib.eager.python.network import Sequential @@ -88,11 +99,15 @@ from tensorflow.python.eager import function from tensorflow.python.eager.context import DEVICE_PLACEMENT_EXPLICIT from tensorflow.python.eager.context import DEVICE_PLACEMENT_WARN from tensorflow.python.eager.context import DEVICE_PLACEMENT_SILENT -from tensorflow.python.eager.context import in_eager_mode -from tensorflow.python.eager.context import in_graph_mode +from tensorflow.python.eager.context import executing_eagerly from tensorflow.python.eager.context import list_devices +from tensorflow.python.eager.context import set_execution_mode +from tensorflow.python.eager.context import execution_mode +from tensorflow.python.eager.context import async_wait +from tensorflow.python.eager.context import async_clear_error +from tensorflow.python.eager.context import SYNC +from tensorflow.python.eager.context import ASYNC from tensorflow.python.eager.context import num_gpus -from tensorflow.python.eager.custom_gradient import custom_gradient from tensorflow.python.eager.execution_callbacks import add_execution_callback from tensorflow.python.eager.execution_callbacks import clear_execution_callbacks from tensorflow.python.eager.execution_callbacks import inf_callback @@ -101,12 +116,13 @@ 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.custom_gradient import custom_gradient 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.training.checkpointable import Checkpointable from tensorflow.python.util.all_util import remove_undocumented py_func = script_ops.eager_py_func @@ -117,5 +133,6 @@ implicit_value_and_gradients = backprop.implicit_val_and_grad gradients_function = backprop.gradients_function value_and_gradients_function = backprop.val_and_grad_function GradientTape = backprop.GradientTape # pylint: disable=invalid-name +in_eager_mode = executing_eagerly remove_undocumented(__name__) diff --git a/tensorflow/contrib/eager/python/tfe_test.py b/tensorflow/contrib/eager/python/tfe_test.py index 0dedb2fd7c0905801cd87c239ff2ee09eecb6080..e80ccbb74d8623e977a98cb7fa5eb41f3c9bf250 100644 --- a/tensorflow/contrib/eager/python/tfe_test.py +++ b/tensorflow/contrib/eager/python/tfe_test.py @@ -47,7 +47,8 @@ class TFETest(test_util.TensorFlowTestCase): def testVariableError(self): with self.assertRaisesRegexp( - RuntimeError, r'Variable not supported in Eager mode'): + RuntimeError, + r'Variable not supported when eager execution is enabled'): variables.Variable(initial_value=1.0) def testGradients(self): @@ -102,10 +103,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 6cdbed5b896577f5622b1bd0123c289c798bc0a5..26d6bc5ae67721286d98b89e32a46bf1a901857e 100644 --- a/tensorflow/contrib/estimator/BUILD +++ b/tensorflow/contrib/estimator/BUILD @@ -138,6 +138,7 @@ py_test( size = "medium", srcs = ["python/estimator/extenders_test.py"], srcs_version = "PY2AND3", + tags = ["notsan"], # b/62863147 deps = [ ":extenders", "//tensorflow/contrib/data/python/ops:dataset_ops", @@ -169,6 +170,7 @@ py_library( "//tensorflow/python:lookup_ops", "//tensorflow/python:math_ops", "//tensorflow/python:metrics", + "//tensorflow/python:nn", "//tensorflow/python:sparse_ops", "//tensorflow/python:sparse_tensor", "//tensorflow/python:summary", @@ -191,6 +193,7 @@ py_test( ":head", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", + "//tensorflow/python:check_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", "//tensorflow/python:control_flow_ops", @@ -351,6 +354,7 @@ cuda_py_test( size = "medium", srcs = ["python/estimator/replicate_model_fn_test.py"], additional_deps = [ + "//third_party/py/absl/testing:parameterized", "//tensorflow/python/estimator", "//tensorflow/python/estimator:dnn", "//tensorflow/python/estimator:export_export", diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py index 0f75b77050b0ba4c752a6a74fdc7024170b6f318..6b9f9575b606f1822d760e8597c55994dd8af04c 100644 --- a/tensorflow/contrib/estimator/__init__.py +++ b/tensorflow/contrib/estimator/__init__.py @@ -39,6 +39,7 @@ _allowed_symbols = [ 'multi_class_head', 'multi_head', 'multi_label_head', + 'poisson_regression_head', 'regression_head', 'DNNEstimator', 'DNNLinearCombinedEstimator', diff --git a/tensorflow/contrib/estimator/python/estimator/extenders.py b/tensorflow/contrib/estimator/python/estimator/extenders.py index c99bf8badb35e6fffb7cae8761db9d402b8b3a8f..2b6881b81487dfdb682d5d6261a0318c59d461f6 100644 --- a/tensorflow/contrib/estimator/python/estimator/extenders.py +++ b/tensorflow/contrib/estimator/python/estimator/extenders.py @@ -33,7 +33,7 @@ _VALID_METRIC_FN_ARGS = set(['features', 'labels', 'predictions', 'config']) def add_metrics(estimator, metric_fn): - """Creates a new ${tf.estimator.Estimator} which has given metrics. + """Creates a new @{tf.estimator.Estimator} which has given metrics. Example: @@ -60,7 +60,7 @@ def add_metrics(estimator, metric_fn): ``` Args: - estimator: A ${tf.estimator.Estimator} object. + estimator: A @{tf.estimator.Estimator} object. metric_fn: A function which should obey the following signature: - Args: can only have following four arguments in any order: * predictions: Predictions `Tensor` or dict of `Tensor` created by given @@ -78,7 +78,7 @@ def add_metrics(estimator, metric_fn): function, namely a `(metric_tensor, update_op)` tuple. Returns: - A new ${tf.estimator.Estimator} which has a union of original metrics with + A new @{tf.estimator.Estimator} which has a union of original metrics with given ones. """ _verify_metric_fn_args(metric_fn) @@ -161,14 +161,14 @@ def forward_features(estimator, keys=None): ``` Args: - estimator: A ${tf.estimator.Estimator} object. + estimator: A @{tf.estimator.Estimator} object. keys: a `string` or a `list` of `string`. If it is `None`, all of the `features` in `dict` is forwarded to the `predictions`. If it is a `string`, only given key is forwarded. If it is a `list` of strings, all the given `keys` are forwarded. Returns: - A new ${tf.estimator.Estimator} which forwards features to predictions. + A new @{tf.estimator.Estimator} which forwards features to predictions. Raises: ValueError: 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 238cf287b768eee28b20202084eb244c085c8b75..f95fcc8039cb54c26543781b31013a7676168b0b 100644 --- a/tensorflow/contrib/estimator/python/estimator/head.py +++ b/tensorflow/contrib/estimator/python/estimator/head.py @@ -31,6 +31,7 @@ from tensorflow.python.ops import check_ops from tensorflow.python.ops import lookup_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics as metrics_lib +from tensorflow.python.ops import nn from tensorflow.python.ops import sparse_ops from tensorflow.python.ops.losses import losses from tensorflow.python.saved_model import signature_constants @@ -177,6 +178,7 @@ def regression_head(weight_column=None, label_dimension=1, loss_reduction=losses.Reduction.SUM, loss_fn=None, + inverse_link_fn=None, name=None): """Creates a `_Head` for regression using the `mean_squared_error` loss. @@ -195,10 +197,16 @@ 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 + 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]`. + Also supports custom `inverse_link_fn`, also known as 'mean function'. + `inverse_link_fn` takes `logits` as argument and returns predicted values. + This function is the inverse of the link function defined in + https://en.wikipedia.org/wiki/Generalized_linear_model#Link_function + Namely, for poisson regression, set `inverse_link_fn=tf.exp`. + Args: weight_column: A string or a `_NumericColumn` created by `tf.feature_column.numeric_column` defining feature column representing @@ -209,7 +217,9 @@ def regression_head(weight_column=None, `[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. + loss_fn: Optional loss function. Defaults to `mean_squared_error`. + inverse_link_fn: Optional inverse link function, also known as 'mean + function'. Defaults to identity. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. @@ -224,6 +234,67 @@ def regression_head(weight_column=None, label_dimension=label_dimension, loss_reduction=loss_reduction, loss_fn=loss_fn, + inverse_link_fn=inverse_link_fn, + name=name) + + +def poisson_regression_head( + weight_column=None, + label_dimension=1, + loss_reduction=losses.Reduction.SUM, + compute_full_loss=True, + name=None): + """Creates a `_Head` for poisson regression using `tf.nn.log_poisson_loss`. + + The loss is the weighted sum over all input dimensions. Namely, if the input + labels have shape `[batch_size, label_dimension]`, the loss is the weighted + sum over both `batch_size` and `label_dimension`. + + The head expects `logits` with shape `[D0, D1, ... DN, label_dimension]`. + In many applications, the shape is `[batch_size, label_dimension]`. + + The `labels` shape must match `logits`, namely + `[D0, D1, ... DN, label_dimension]`. If `label_dimension=1`, shape + `[D0, D1, ... DN]` is also supported. + + If `weight_column` is specified, weights must be of shape + `[D0, D1, ... DN]`, `[D0, D1, ... DN, 1]` or + `[D0, D1, ... DN, label_dimension]`. + + This is implemented as a generalized linear model, see + https://en.wikipedia.org/wiki/Generalized_linear_model. + + Args: + weight_column: A string or a `_NumericColumn` created by + `tf.feature_column.numeric_column` defining feature column representing + weights. It is used to down weight or boost examples during training. It + will be multiplied by the loss of the example. + 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`. + compute_full_loss: Whether to include the constant `log(z!)` term in + computing the poisson loss. See `tf.nn.log_poisson_loss` for the full + documentation. + 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 poisson regression. + + Raises: + ValueError: If `label_dimension` or `loss_reduction` is invalid. + """ + def _poisson_loss(labels, logits): + return nn.log_poisson_loss( + targets=labels, log_input=logits, compute_full_loss=compute_full_loss) + 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=_poisson_loss, + inverse_link_fn=math_ops.exp, name=name) diff --git a/tensorflow/contrib/estimator/python/estimator/head_test.py b/tensorflow/contrib/estimator/python/estimator/head_test.py index 43cdfec9689879201305385499b3b784e1593d60..dc30dde877ab5f912e3f6a724d481b151a3ed044 100644 --- a/tensorflow/contrib/estimator/python/estimator/head_test.py +++ b/tensorflow/contrib/estimator/python/estimator/head_test.py @@ -32,6 +32,7 @@ from tensorflow.python.framework import errors 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 math_ops from tensorflow.python.ops import string_ops @@ -1106,5 +1107,75 @@ class MultiLabelHead(test.TestCase): expected_metrics=expected_metrics) +class PoissonRegressionHead(test.TestCase): + + def setUp(self): + ops.reset_default_graph() + + def test_train(self): + head = head_lib.poisson_regression_head() + + # Create estimator spec. + logits = np.array([[0], [-1], [1]], dtype=np.float32) + labels = np.array([[1], [2], [3]], dtype=np.int32) + # With x = exp(logits), z = labels. + # loss = -ln(exp(-x) * (x^z) / z!) + # = x - z * ln(x) + ln(z!) + # = exp(logits) - labels * logits - ln(labels!) + # But for ln(z!) and z > 1, the Stirling approximation is used + # ln(z!) = z*ln(z) - z + 0.5*ln(2*pi*z) + # loss = [exp(0) - 1 * 0 + ln(1!), + # exp(-1) - 2 * (-1) + 2*ln(2) - 2 + 0.5*ln(2*pi*2), + # exp(1) - 3 * 1 + 3*ln(3) - 3 + 0.5*ln(2*pi*3)] + # = [1.0, 3.020, 1.482] + # sum_loss = 5.502 + expected_loss = 5.502 + atol = 0.001 + expected_train_result = b'my_train_op' + def _train_op_fn(loss): + with ops.control_dependencies((check_ops.assert_near( + math_ops.to_float(expected_loss), math_ops.to_float(loss), + atol=atol, name='assert_loss'),)): + return constant_op.constant(expected_train_result) + + 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) + + with self.test_session() as sess: + _initialize_variables(self, spec.scaffold) + loss, train_result = sess.run([spec.loss, spec.train_op]) + self.assertAlmostEqual(expected_loss, loss, delta=atol) + self.assertEqual(expected_train_result, train_result) + + def test_predict(self): + head = head_lib.poisson_regression_head() + + # Create estimator spec. + logits = np.array([[0], [-1], [1]], dtype=np.float32) + expected_predictions = np.exp(logits) + spec = head.create_estimator_spec( + features={'x': np.array(((42.,),), dtype=np.int32)}, + mode=model_fn.ModeKeys.PREDICT, + logits=logits) + + # Assert spec contains expected tensors. + keys = prediction_keys.PredictionKeys + self.assertItemsEqual( + (keys.PREDICTIONS, keys.LOGITS), spec.predictions.keys()) + self.assertEqual(dtypes.float32, spec.predictions[keys.PREDICTIONS].dtype) + self.assertEqual(dtypes.float32, spec.predictions[keys.LOGITS].dtype) + + # Assert predictions. + with self.test_session(): + _initialize_variables(self, spec.scaffold) + self.assertAllClose( + expected_predictions, spec.predictions[keys.PREDICTIONS].eval()) + self.assertAllClose(logits, spec.predictions[keys.LOGITS].eval()) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py index caa9dd83233b6b850385335fde96431271d85c3a..e0fae2c99292385c6dd32cc6002cee2076a2bb20 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py @@ -110,7 +110,8 @@ def replicate_model_fn(model_fn, Certain algorithms were chosen for aggregating results of computations on multiple towers: - Losses from all towers are reduced according to `loss_reduction`. - - Gradients are reduced using sum for each trainable variable. + - Gradients from all towers are reduced according to `loss_reduction` + for each trainable variable. - `eval_metrics_ops` are reduced per metric using `reduce_mean`. - `EstimatorSpec.predictions` and `EstimatorSpec.export_outputs` are reduced using concatenation. @@ -195,7 +196,7 @@ 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] + 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] @@ -457,6 +458,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)] @@ -467,6 +475,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 @@ -476,6 +485,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: @@ -483,6 +493,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 @@ -780,7 +791,7 @@ def _extract_tensors(tensors_and_vars): tensor, _ = tensor_and_var if isinstance(tensor, ops_lib.IndexedSlices): tensors.append(tensor.values) - else: + elif tensor is not None: tensors.append(tensor) return tensors 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 03d31226af613960a19ce116b19b30153b1fdcee..144b45982c8aec2e2b115c812b24e8843d60ce1e 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import re import shutil import tempfile +from absl.testing import parameterized import numpy as np import six @@ -37,6 +38,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 @@ -56,26 +58,19 @@ from tensorflow.python.training import gradient_descent from tensorflow.python.training import training -# TODO(isaprykin): Parametrize all the tests on -# replicate_model_fn._VariableDistributionMode when it's supported. -class DNNClassifierIntegrationTest(test_util.TensorFlowTestCase): +class DNNClassifierIntegrationTest(test_util.TensorFlowTestCase, + parameterized.TestCase): def setUp(self): self._model_dir = tempfile.mkdtemp() - def test_complete_flow_with_public_version(self): - return self._complete_flow_with_mode(mode=None) - - def test_complete_flow_with_mode_local_ps_server(self): - return self._complete_flow_with_mode( - replicate_model_fn._VariableDistributionMode. - SHARED_LOCAL_PARAMETER_SERVER) - - def test_complete_flow_with_mode_round_robin(self): - return self._complete_flow_with_mode( - replicate_model_fn._VariableDistributionMode.SHARED_ROUND_ROBIN) - - def _complete_flow_with_mode(self, mode): + @parameterized.named_parameters( + ('PublicInterface', None), + ('ParameterServerMode', replicate_model_fn._VariableDistributionMode. + SHARED_LOCAL_PARAMETER_SERVER), + ('RoundRobinMode', + replicate_model_fn._VariableDistributionMode.SHARED_ROUND_ROBIN)) + def test_complete_flow_with_mode(self, mode): n_classes = 3 input_dimension = 2 batch_size = 12 @@ -239,6 +234,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( @@ -433,12 +435,51 @@ 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): @@ -981,8 +1022,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( @@ -995,7 +1041,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( @@ -1008,7 +1054,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( @@ -1021,7 +1067,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] @@ -1035,6 +1081,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]} diff --git a/tensorflow/contrib/factorization/BUILD b/tensorflow/contrib/factorization/BUILD index 180f1b68f3b56113dfbbfc100bd04efc3bb8b31f..ad8568ad44ea84f96b97e98567a276c70520d53d 100644 --- a/tensorflow/contrib/factorization/BUILD +++ b/tensorflow/contrib/factorization/BUILD @@ -66,6 +66,7 @@ tf_custom_op_py_library( "//tensorflow/python:variables", "//tensorflow/python/estimator", "//tensorflow/python/estimator:model_fn", + "//tensorflow/python/feature_column:feature_column_py", "//third_party/py/numpy", ], ) @@ -223,7 +224,10 @@ py_test( srcs = ["python/ops/kmeans_test.py"], shard_count = 4, srcs_version = "PY2AND3", - tags = ["notsan"], # b/67512932 + tags = [ + "nomac", # b/73741358 + "notsan", # b/67512932 + ], deps = [ ":factorization_py", ":factorization_py_CYCLIC_DEPENDENCIES_THAT_NEED_TO_GO", @@ -238,6 +242,7 @@ py_test( "//tensorflow/python:random_ops", "//tensorflow/python:training", "//tensorflow/python/estimator:run_config", + "//tensorflow/python/feature_column:feature_column_py", "//third_party/py/numpy", ], ) 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 4d0f9b24240ccbafe89ef912b4d3252cefb1f7f2..38faca119d0b5ee883de3b215428a0db8a021016 100644 --- a/tensorflow/contrib/factorization/python/ops/kmeans.py +++ b/tensorflow/contrib/factorization/python/ops/kmeans.py @@ -26,6 +26,7 @@ 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.feature_column import feature_column as fc from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -61,8 +62,8 @@ class _LossRelativeChangeHook(session_run_hook.SessionRunHook): loss = run_values.results assert loss is not None if self._prev_loss: - relative_change = (abs(loss - self._prev_loss) / - (1 + abs(self._prev_loss))) + relative_change = ( + abs(loss - self._prev_loss) / (1 + abs(self._prev_loss))) if relative_change < self._tolerance: run_context.request_stop() self._prev_loss = loss @@ -105,24 +106,32 @@ class _InitializeClustersHook(session_run_hook.SessionRunHook): logging.info(e) -def _parse_tensor_or_dict(features): +def _parse_features_if_necessary(features, feature_columns): """Helper function to convert the input points into a usable format. Args: - features: The input points. + features: The input features. + feature_columns: An optionable iterable containing all the feature columns + used by the model. All items in the set should be feature column instances + that can be passed to `tf.feature_column.input_layer`. If this is None, + all features will be used. Returns: - If `features` is a dict of `k` features, each of which is a vector of `n` - scalars, the return value is a Tensor of shape `(n, k)` representing `n` - input points, where the items in the `k` dimension are sorted - lexicographically by `features` key. If `features` is not a dict, it is - returned unmodified. + If `features` is a dict of `k` features (optionally filtered by + `feature_columns`), each of which is a vector of `n` scalars, the return + value is a Tensor of shape `(n, k)` representing `n` input points, where the + items in the `k` dimension are sorted lexicographically by `features` key. + If `features` is not a dict, it is returned unmodified. """ - if isinstance(features, dict): - keys = sorted(features.keys()) - with ops.colocate_with(features[keys[0]]): - features = array_ops.concat([features[k] for k in keys], axis=1) - return features + if not isinstance(features, dict): + return features + + if feature_columns: + return fc.input_layer(features, feature_columns) + + keys = sorted(features.keys()) + with ops.colocate_with(features[keys[0]]): + return array_ops.concat([features[k] for k in keys], axis=1) class _ModelFn(object): @@ -130,7 +139,8 @@ class _ModelFn(object): def __init__(self, num_clusters, initial_clusters, distance_metric, random_seed, use_mini_batch, mini_batch_steps_per_iteration, - kmeans_plus_plus_num_retries, relative_tolerance): + kmeans_plus_plus_num_retries, relative_tolerance, + feature_columns): self._num_clusters = num_clusters self._initial_clusters = initial_clusters self._distance_metric = distance_metric @@ -139,11 +149,12 @@ class _ModelFn(object): self._mini_batch_steps_per_iteration = mini_batch_steps_per_iteration self._kmeans_plus_plus_num_retries = kmeans_plus_plus_num_retries self._relative_tolerance = relative_tolerance + self._feature_columns = feature_columns 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: @@ -166,7 +177,7 @@ class _ModelFn(object): # input_points is a single Tensor. Therefore, the sharding functionality # in clustering_ops is unused, and some of the values below are lists of a # single item. - input_points = _parse_tensor_or_dict(features) + input_points = _parse_features_if_necessary(features, self._feature_columns) # Let N = the number of input_points. # all_distances: A list of one matrix of shape (N, num_clusters). Each value @@ -233,7 +244,57 @@ class _ModelFn(object): # TODO(agarwal,ands): support sharded input. class KMeansClustering(estimator.Estimator): - """An Estimator for K-Means clustering.""" + """An Estimator for K-Means clustering. + + Example: + ``` + import numpy as np + import tensorflow as tf + + num_points = 100 + dimensions = 2 + points = np.random.uniform(0, 1000, [num_points, dimensions]) + + def input_fn(): + return tf.train.limit_epochs( + tf.convert_to_tensor(points, dtype=tf.float32), num_epochs=1) + + num_clusters = 5 + kmeans = tf.contrib.factorization.KMeansClustering( + num_clusters=num_clusters, use_mini_batch=False) + + # train + num_iterations = 10 + previous_centers = None + for _ in xrange(num_iterations): + kmeans.train(input_fn) + cluster_centers = kmeans.cluster_centers() + if previous_centers is not None: + print 'delta:', cluster_centers - previous_centers + previous_centers = cluster_centers + print 'score:', kmeans.score(input_fn) + print 'cluster centers:', cluster_centers + + # map the input points to their clusters + cluster_indices = list(kmeans.predict_cluster_index(input_fn)) + for i, point in enumerate(points): + cluster_index = cluster_indices[i] + center = cluster_centers[cluster_index] + print 'point:', point, 'is in cluster', cluster_index, 'centered at', center + ``` + + The `SavedModel` saved by the `export_savedmodel` method does not include the + cluster centers. However, the cluster centers may be retrieved by the + latest checkpoint saved during training. Specifically, + ``` + kmeans.cluster_centers() + ``` + is equivalent to + ``` + tf.train.load_variable( + kmeans.model_dir, KMeansClustering.CLUSTER_CENTERS_VAR_NAME) + ``` + """ # Valid values for the distance_metric constructor argument. SQUARED_EUCLIDEAN_DISTANCE = clustering_ops.SQUARED_EUCLIDEAN_DISTANCE @@ -253,6 +314,9 @@ class KMeansClustering(estimator.Estimator): CLUSTER_INDEX = 'cluster_index' ALL_DISTANCES = 'all_distances' + # Variable name used by cluster_centers(). + CLUSTER_CENTERS_VAR_NAME = clustering_ops.CLUSTERS_VAR_NAME + def __init__(self, num_clusters, model_dir=None, @@ -263,7 +327,8 @@ class KMeansClustering(estimator.Estimator): mini_batch_steps_per_iteration=1, kmeans_plus_plus_num_retries=2, relative_tolerance=None, - config=None): + config=None, + feature_columns=None): """Creates an Estimator for running KMeans training and inference. This Estimator implements the following variants of the K-means algorithm: @@ -330,6 +395,10 @@ class KMeansClustering(estimator.Estimator): iterations. Stops learning if the loss changes less than this amount. This may not work correctly if `use_mini_batch=True`. config: See @{tf.estimator.Estimator}. + feature_columns: An optionable iterable containing all the feature columns + used by the model. All items in the set should be feature column + instances that can be passed to `tf.feature_column.input_layer`. If this + is None, all features will be used. Raises: ValueError: An invalid argument was passed to `initial_clusters` or @@ -349,7 +418,8 @@ class KMeansClustering(estimator.Estimator): model_fn=_ModelFn( num_clusters, initial_clusters, distance_metric, random_seed, use_mini_batch, mini_batch_steps_per_iteration, - kmeans_plus_plus_num_retries, relative_tolerance).model_fn, + kmeans_plus_plus_num_retries, relative_tolerance, + feature_columns).model_fn, model_dir=model_dir, config=config) @@ -406,4 +476,4 @@ class KMeansClustering(estimator.Estimator): def cluster_centers(self): """Returns the cluster centers.""" - return self.get_variable_value(clustering_ops.CLUSTERS_VAR_NAME) + return self.get_variable_value(KMeansClustering.CLUSTER_CENTERS_VAR_NAME) diff --git a/tensorflow/contrib/factorization/python/ops/kmeans_test.py b/tensorflow/contrib/factorization/python/ops/kmeans_test.py index f9598bfc08c05ea3bba88b3135da0cf2e6bb0c95..0103cc44394b772a1bb1bb8bed78f5445297950a 100644 --- a/tensorflow/contrib/factorization/python/ops/kmeans_test.py +++ b/tensorflow/contrib/factorization/python/ops/kmeans_test.py @@ -27,6 +27,7 @@ from sklearn.cluster import KMeans as SklearnKMeans # pylint: disable=g-import-not-at-top from tensorflow.contrib.factorization.python.ops import kmeans as kmeans_lib from tensorflow.python.estimator import run_config +from tensorflow.python.feature_column import feature_column as fc from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -226,6 +227,44 @@ class KMeansTest(KMeansTestBase): self._infer_helper(kmeans, clusters, 10) self._infer_helper(kmeans, clusters, 1) + def _parse_feature_dict_helper(self, features, parsed_feature_dict): + # Perform a sanity check. + self.assertEqual(features.shape, parsed_feature_dict.shape) + self.assertEqual(features.dtype, parsed_feature_dict.dtype) + # Then check that running the tensor yields the original list of points. + with self.test_session() as sess: + parsed_points = sess.run(parsed_feature_dict) + self.assertAllEqual(self.points, parsed_points) + + def test_parse_features(self): + """Tests the various behaviours of kmeans._parse_features_if_necessary.""" + + # No-op if a tensor is passed in. + features = constant_op.constant(self.points) + parsed_features = kmeans_lib._parse_features_if_necessary(features, None) + self.assertAllEqual(features, parsed_features) + + # All values from a feature dict are transformed into a tensor. + feature_dict = { + 'x': [[point[0]] for point in self.points], + 'y': [[point[1]] for point in self.points] + } + parsed_feature_dict = kmeans_lib._parse_features_if_necessary( + feature_dict, None) + self._parse_feature_dict_helper(features, parsed_feature_dict) + + # Only the feature_columns of a feature dict are transformed into a tensor. + feature_dict_with_extras = { + 'foo': 'bar', + 'x': [[point[0]] for point in self.points], + 'baz': {'fizz': 'buzz'}, + 'y': [[point[1]] for point in self.points] + } + feature_columns = [fc.numeric_column(key='x'), fc.numeric_column(key='y')] + parsed_feature_dict = kmeans_lib._parse_features_if_necessary( + feature_dict_with_extras, feature_columns) + self._parse_feature_dict_helper(features, parsed_feature_dict) + class KMeansTestMultiStageInit(KMeansTestBase): @@ -394,7 +433,6 @@ class KMeansCosineDistanceTest(KMeansTestBase): true_assignments = [0] * 2 + [1] * 2 + [2] * 8 true_score = len(points) - np.tensordot( normalize(points), true_centers[true_assignments]) - kmeans = kmeans_lib.KMeansClustering( 3, initial_clusters=self.initial_clusters, diff --git a/tensorflow/contrib/feature_column/BUILD b/tensorflow/contrib/feature_column/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..3614b2b15a6cbdd73f9f24c7e4e4534228d31499 --- /dev/null +++ b/tensorflow/contrib/feature_column/BUILD @@ -0,0 +1,66 @@ +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 = [ + ":sequence_feature_column", + "//tensorflow/python:util", + ], +) + +py_library( + name = "sequence_feature_column", + srcs = ["python/feature_column/sequence_feature_column.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:array_ops", + "//tensorflow/python:check_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:parsing_ops", + "//tensorflow/python:sparse_ops", + "//tensorflow/python:tensor_shape", + "//tensorflow/python:variable_scope", + "//tensorflow/python/feature_column", + ], +) + +py_test( + name = "sequence_feature_column_test", + srcs = ["python/feature_column/sequence_feature_column_test.py"], + srcs_version = "PY2AND3", + tags = ["no_pip"], + deps = [ + ":sequence_feature_column", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:errors", + "//tensorflow/python:framework_ops", + "//tensorflow/python:sparse_tensor", + "//tensorflow/python:training", + "//tensorflow/python/feature_column", + "//third_party/py/numpy", + ], +) diff --git a/tensorflow/contrib/bayesflow/python/ops/optimizers.py b/tensorflow/contrib/feature_column/__init__.py similarity index 55% rename from tensorflow/contrib/bayesflow/python/ops/optimizers.py rename to tensorflow/contrib/feature_column/__init__.py index fb70628d1083836281e9327e83e109493276c64f..baa8c1567a5aeb39976ab04c54ae2728ba050a7c 100644 --- a/tensorflow/contrib/bayesflow/python/ops/optimizers.py +++ b/tensorflow/contrib/feature_column/__init__.py @@ -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. @@ -12,25 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Probabilistic optimizer modules. - -See ${python/contrib.bayesflow.optimizers}. -""" +"""Experimental utilities for tf.feature_column.""" 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.sgld_optimizer import * -from tensorflow.contrib.bayesflow.python.ops.variational_sgd_optimizer import * -# pylint: enable=wildcard-import +# pylint: disable=unused-import,line-too-long,wildcard-import +from tensorflow.contrib.feature_column.python.feature_column.sequence_feature_column import * + from tensorflow.python.util.all_util import remove_undocumented +# pylint: enable=unused-import,line-too-long,wildcard-import _allowed_symbols = [ - 'SGLDOptimizer', - 'VariationalSGDOptimizer', + 'sequence_categorical_column_with_hash_bucket', + 'sequence_categorical_column_with_identity', + 'sequence_categorical_column_with_vocabulary_list', + 'sequence_categorical_column_with_vocabulary_file', + 'sequence_input_layer', + 'sequence_numeric_column', ] -remove_undocumented(__name__, _allowed_symbols) +remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py new file mode 100644 index 0000000000000000000000000000000000000000..e60116966fc8d8bb0745f50a0238f10f02af4167 --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py @@ -0,0 +1,417 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Experimental methods for tf.feature_column sequence input.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +import collections + + +from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import check_ops +from tensorflow.python.ops import parsing_ops +from tensorflow.python.ops import sparse_ops +from tensorflow.python.ops import variable_scope + +# pylint: disable=protected-access +# TODO(b/73827486): Support SequenceExample. + + +def sequence_input_layer( + features, + feature_columns, + weight_collections=None, + trainable=True): + """"Builds input layer for sequence input. + + All `feature_columns` must be sequence dense columns with the same + `sequence_length`. The output of this method can be fed into sequence + networks, such as RNN. + + The output of this method is a 3D `Tensor` of shape `[batch_size, T, D]`. + `T` is the maximum sequence length for this batch, which could differ from + batch to batch. + + If multiple `feature_columns` are given with `Di` `num_elements` each, their + outputs are concatenated. So, the final `Tensor` has shape + `[batch_size, T, D0 + D1 + ... + Dn]`. + + Example: + + ```python + rating = sequence_numeric_column('rating') + watches = sequence_categorical_column_with_identity( + 'watches', num_buckets=1000) + watches_embedding = embedding_column(watches, dimension=10) + columns = [rating, watches] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Args: + features: A dict mapping keys to tensors. + feature_columns: An iterable of dense sequence columns. Valid columns are + - `embedding_column` that wraps a `sequence_categorical_column_with_*` + - `sequence_numeric_column`. + weight_collections: A list of collection names to which the Variable will be + added. Note that variables will also be added to collections + `tf.GraphKeys.GLOBAL_VARIABLES` and `ops.GraphKeys.MODEL_VARIABLES`. + trainable: If `True` also add the variable to the graph collection + `GraphKeys.TRAINABLE_VARIABLES`. + + Returns: + An `(input_layer, sequence_length)` tuple where: + - input_layer: A float `Tensor` of shape `[batch_size, T, D]`. + `T` is the maximum sequence length for this batch, which could differ + from batch to batch. `D` is the sum of `num_elements` for all + `feature_columns`. + - sequence_length: An int `Tensor` of shape `[batch_size]`. The sequence + length for each example. + + Raises: + ValueError: If any of the `feature_columns` is the wrong type. + """ + feature_columns = fc._clean_feature_columns(feature_columns) + for c in feature_columns: + if not isinstance(c, fc._SequenceDenseColumn): + raise ValueError( + 'All feature_columns must be of type _SequenceDenseColumn. ' + 'You can wrap a sequence_categorical_column with an embedding_column ' + 'or indicator_column. ' + 'Given (type {}): {}'.format(type(c), c)) + + with variable_scope.variable_scope( + None, default_name='sequence_input_layer', values=features.values()): + builder = fc._LazyBuilder(features) + output_tensors = [] + sequence_lengths = [] + ordered_columns = [] + for column in sorted(feature_columns, key=lambda x: x.name): + ordered_columns.append(column) + with variable_scope.variable_scope( + None, default_name=column._var_scope_name): + dense_tensor, sequence_length = column._get_sequence_dense_tensor( + builder, + weight_collections=weight_collections, + trainable=trainable) + # Flattens the final dimension to produce a 3D Tensor. + num_elements = column._variable_shape.num_elements() + shape = array_ops.shape(dense_tensor) + output_tensors.append( + array_ops.reshape( + dense_tensor, + shape=array_ops.concat([shape[:2], [num_elements]], axis=0))) + sequence_lengths.append(sequence_length) + fc._verify_static_batch_size_equality(output_tensors, ordered_columns) + fc._verify_static_batch_size_equality(sequence_lengths, ordered_columns) + sequence_length = _assert_all_equal_and_return(sequence_lengths) + return array_ops.concat(output_tensors, -1), sequence_length + + +def sequence_categorical_column_with_identity( + key, num_buckets, default_value=None): + """Returns a feature column that represents sequences of integers. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + watches = sequence_categorical_column_with_identity( + 'watches', num_buckets=1000) + watches_embedding = embedding_column(watches, dimension=10) + columns = [watches_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Args: + key: A unique string identifying the input feature. + num_buckets: Range of inputs. Namely, inputs are expected to be in the + range `[0, num_buckets)`. + default_value: If `None`, this column's graph operations will fail for + out-of-range inputs. Otherwise, this value must be in the range + `[0, num_buckets)`, and will replace out-of-range inputs. + + Returns: + A `_SequenceCategoricalColumn`. + """ + return fc._SequenceCategoricalColumn( + fc.categorical_column_with_identity( + key=key, + num_buckets=num_buckets, + default_value=default_value)) + + +def sequence_categorical_column_with_hash_bucket( + key, hash_bucket_size, dtype=dtypes.string): + """A sequence of categorical terms where ids are set by hashing. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + tokens = sequence_categorical_column_with_hash_bucket( + 'tokens', hash_bucket_size=1000) + tokens_embedding = embedding_column(tokens, dimension=10) + columns = [tokens_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Args: + key: A unique string identifying the input feature. + hash_bucket_size: An int > 1. The number of buckets. + dtype: The type of features. Only string and integer types are supported. + + Returns: + A `_SequenceCategoricalColumn`. + """ + return fc._SequenceCategoricalColumn( + fc.categorical_column_with_hash_bucket( + key=key, + hash_bucket_size=hash_bucket_size, + dtype=dtype)) + + +def sequence_categorical_column_with_vocabulary_file( + key, vocabulary_file, vocabulary_size=None, num_oov_buckets=0, + default_value=None, dtype=dtypes.string): + """A sequence of categorical terms where ids use a vocabulary file. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + states = sequence_categorical_column_with_vocabulary_file( + key='states', vocabulary_file='/us/states.txt', vocabulary_size=50, + num_oov_buckets=5) + states_embedding = embedding_column(states, dimension=10) + columns = [states_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Args: + key: A unique string identifying the input feature. + vocabulary_file: The vocabulary file name. + vocabulary_size: Number of the elements in the vocabulary. This must be no + greater than length of `vocabulary_file`, if less than length, later + values are ignored. If None, it is set to the length of `vocabulary_file`. + num_oov_buckets: Non-negative integer, the number of out-of-vocabulary + buckets. All out-of-vocabulary inputs will be assigned IDs in the range + `[vocabulary_size, vocabulary_size+num_oov_buckets)` based on a hash of + the input value. A positive `num_oov_buckets` can not be specified with + `default_value`. + default_value: The integer ID value to return for out-of-vocabulary feature + values, defaults to `-1`. This can not be specified with a positive + `num_oov_buckets`. + dtype: The type of features. Only string and integer types are supported. + + Returns: + A `_SequenceCategoricalColumn`. + """ + return fc._SequenceCategoricalColumn( + fc.categorical_column_with_vocabulary_file( + key=key, + vocabulary_file=vocabulary_file, + vocabulary_size=vocabulary_size, + num_oov_buckets=num_oov_buckets, + default_value=default_value, + dtype=dtype)) + + +def sequence_categorical_column_with_vocabulary_list( + key, vocabulary_list, dtype=None, default_value=-1, num_oov_buckets=0): + """A sequence of categorical terms where ids use an in-memory list. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + colors = sequence_categorical_column_with_vocabulary_list( + key='colors', vocabulary_list=('R', 'G', 'B', 'Y'), + num_oov_buckets=2) + colors_embedding = embedding_column(colors, dimension=3) + columns = [colors_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Args: + key: A unique string identifying the input feature. + vocabulary_list: An ordered iterable defining the vocabulary. Each feature + is mapped to the index of its value (if present) in `vocabulary_list`. + Must be castable to `dtype`. + dtype: The type of features. Only string and integer types are supported. + If `None`, it will be inferred from `vocabulary_list`. + default_value: The integer ID value to return for out-of-vocabulary feature + values, defaults to `-1`. This can not be specified with a positive + `num_oov_buckets`. + num_oov_buckets: Non-negative integer, the number of out-of-vocabulary + buckets. All out-of-vocabulary inputs will be assigned IDs in the range + `[len(vocabulary_list), len(vocabulary_list)+num_oov_buckets)` based on a + hash of the input value. A positive `num_oov_buckets` can not be specified + with `default_value`. + + Returns: + A `_SequenceCategoricalColumn`. + """ + return fc._SequenceCategoricalColumn( + fc.categorical_column_with_vocabulary_list( + key=key, + vocabulary_list=vocabulary_list, + dtype=dtype, + default_value=default_value, + num_oov_buckets=num_oov_buckets)) + + +def sequence_numeric_column( + key, + shape=(1,), + default_value=0., + dtype=dtypes.float32): + """Returns a feature column that represents sequences of numeric data. + + Example: + + ```python + temperature = sequence_numeric_column('temperature') + columns = [temperature] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Args: + key: A unique string identifying the input features. + shape: The shape of the input data per sequence id. E.g. if `shape=(2,)`, + each example must contain `2 * sequence_length` values. + default_value: A single value compatible with `dtype` that is used for + padding the sparse data into a dense `Tensor`. + dtype: The type of values. + + Returns: + A `_SequenceNumericColumn`. + """ + # TODO(b/73160931): Add validations. + return _SequenceNumericColumn( + key, + shape=shape, + default_value=default_value, + dtype=dtype) + + +def _assert_all_equal_and_return(tensors, name=None): + """Asserts that all tensors are equal and returns the first one.""" + with ops.name_scope(name, 'assert_all_equal', values=tensors): + if len(tensors) == 1: + return tensors[0] + assert_equal_ops = [] + for t in tensors[1:]: + assert_equal_ops.append(check_ops.assert_equal(tensors[0], t)) + with ops.control_dependencies(assert_equal_ops): + return array_ops.identity(tensors[0]) + + +class _SequenceNumericColumn( + fc._SequenceDenseColumn, + collections.namedtuple( + '_SequenceNumericColumn', + ['key', 'shape', 'default_value', 'dtype'])): + """Represents sequences of numeric data.""" + + @property + def name(self): + return self.key + + @property + def _parse_example_spec(self): + return {self.key: parsing_ops.VarLenFeature(self.dtype)} + + def _transform_feature(self, inputs): + return inputs.get(self.key) + + @property + def _variable_shape(self): + return tensor_shape.TensorShape(self.shape) + + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + # Do nothing with weight_collections and trainable since no variables are + # created in this function. + del weight_collections + del trainable + sp_tensor = inputs.get(self) + dense_tensor = sparse_ops.sparse_tensor_to_dense( + sp_tensor, default_value=self.default_value) + # Reshape into [batch_size, T, variable_shape]. + dense_shape = array_ops.concat( + [array_ops.shape(dense_tensor)[:1], [-1], self._variable_shape], + axis=0) + dense_tensor = array_ops.reshape(dense_tensor, shape=dense_shape) + sequence_length = fc._sequence_length_from_sparse_tensor( + sp_tensor, num_elements=self._variable_shape.num_elements()) + return fc._SequenceDenseColumn.TensorSequenceLengthPair( + dense_tensor=dense_tensor, sequence_length=sequence_length) + +# pylint: enable=protected-access diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b64f086376dad65c1f32bee4bfce9334a60fd24a --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py @@ -0,0 +1,790 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for sequential_feature_column.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import numpy as np + +from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column as sfc +from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.feature_column.feature_column import _LazyBuilder +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops +from tensorflow.python.framework import sparse_tensor +from tensorflow.python.platform import test +from tensorflow.python.training import monitored_session + + +class SequenceInputLayerTest(test.TestCase): + + def test_embedding_column(self): + vocabulary_size = 3 + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [1] + # example 1, ids [2, 0] + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 0), + dense_shape=(2, 2)) + + embedding_dimension_a = 2 + embedding_values_a = ( + (1., 2.), # id 0 + (3., 4.), # id 1 + (5., 6.) # id 2 + ) + embedding_dimension_b = 3 + embedding_values_b = ( + (11., 12., 13.), # id 0 + (14., 15., 16.), # id 1 + (17., 18., 19.) # id 2 + ) + def _get_initializer(embedding_dimension, embedding_values): + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + return _initializer + + expected_input_layer = [ + # example 0, ids_a [2], ids_b [1] + [[5., 6., 14., 15., 16.], [0., 0., 0., 0., 0.]], + # example 1, ids_a [0, 1], ids_b [2, 0] + [[1., 2., 17., 18., 19.], [3., 4., 11., 12., 13.]], + ] + expected_sequence_length = [1, 2] + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column_a = fc.embedding_column( + categorical_column_a, dimension=embedding_dimension_a, + initializer=_get_initializer(embedding_dimension_a, embedding_values_a)) + categorical_column_b = sfc.sequence_categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + embedding_column_b = fc.embedding_column( + categorical_column_b, dimension=embedding_dimension_b, + initializer=_get_initializer(embedding_dimension_b, embedding_values_b)) + + input_layer, sequence_length = sfc.sequence_input_layer( + features={ + 'aaa': sparse_input_a, + 'bbb': sparse_input_b, + }, + # Test that columns are reordered alphabetically. + feature_columns=[embedding_column_b, embedding_column_a]) + + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual( + ('sequence_input_layer/aaa_embedding/embedding_weights:0', + 'sequence_input_layer/bbb_embedding/embedding_weights:0'), + tuple([v.name for v in global_vars])) + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(embedding_values_a, global_vars[0].eval(session=sess)) + self.assertAllEqual(embedding_values_b, global_vars[1].eval(session=sess)) + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_embedding_column_with_non_sequence_categorical(self): + """Tests that error is raised for non-sequence categorical column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column_a = fc.embedding_column( + categorical_column_a, dimension=2) + + with self.assertRaisesRegexp( + ValueError, + r'In embedding_column: aaa_embedding\. categorical_column must be of ' + r'type _SequenceCategoricalColumn to use sequence_input_layer\.'): + _, _ = sfc.sequence_input_layer( + features={'aaa': sparse_input}, + feature_columns=[embedding_column_a]) + + def test_indicator_column(self): + vocabulary_size_a = 3 + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + vocabulary_size_b = 2 + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [1] + # example 1, ids [1, 0] + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 1, 0), + dense_shape=(2, 2)) + + expected_input_layer = [ + # example 0, ids_a [2], ids_b [1] + [[0., 0., 1., 0., 1.], [0., 0., 0., 0., 0.]], + # example 1, ids_a [0, 1], ids_b [1, 0] + [[1., 0., 0., 0., 1.], [0., 1., 0., 1., 0.]], + ] + expected_sequence_length = [1, 2] + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size_a) + indicator_column_a = fc.indicator_column(categorical_column_a) + categorical_column_b = sfc.sequence_categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size_b) + indicator_column_b = fc.indicator_column(categorical_column_b) + input_layer, sequence_length = sfc.sequence_input_layer( + features={ + 'aaa': sparse_input_a, + 'bbb': sparse_input_b, + }, + # Test that columns are reordered alphabetically. + feature_columns=[indicator_column_b, indicator_column_a]) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_indicator_column_with_non_sequence_categorical(self): + """Tests that error is raised for non-sequence categorical column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column_a = fc.indicator_column(categorical_column_a) + + with self.assertRaisesRegexp( + ValueError, + r'In indicator_column: aaa_indicator\. categorical_column must be of ' + r'type _SequenceCategoricalColumn to use sequence_input_layer\.'): + _, _ = sfc.sequence_input_layer( + features={'aaa': sparse_input}, + feature_columns=[indicator_column_a]) + + def test_numeric_column(self): + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + expected_input_layer = [ + [[0.], [1.]], + [[10.], [0.]], + ] + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa') + + input_layer, sequence_length = sfc.sequence_input_layer( + features={'aaa': sparse_input}, + feature_columns=[numeric_column]) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_numeric_column_multi_dim(self): + """Tests sequence_input_layer for multi-dimensional numeric_column.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]] + # example 1, [[[10., 11.], [12., 13.]]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), + (1, 0), (1, 1), (1, 2), (1, 3)), + values=(0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + dense_shape=(2, 8)) + # The output of numeric_column._get_dense_tensor should be flattened. + expected_input_layer = [ + [[0., 1., 2., 3.], [4., 5., 6., 7.]], + [[10., 11., 12., 13.], [0., 0., 0., 0.]], + ] + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) + + input_layer, sequence_length = sfc.sequence_input_layer( + features={'aaa': sparse_input}, + feature_columns=[numeric_column]) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_not_equal(self): + """Tests that an error is raised when sequence lengths are not equal.""" + # Input a with sequence_length = [2, 1] + sparse_input_a = sparse_tensor.SparseTensorValue( + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + # Input b with sequence_length = [1, 1] + sparse_input_b = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0)), + values=(1., 10.), + dense_shape=(2, 2)) + numeric_column_a = sfc.sequence_numeric_column('aaa') + numeric_column_b = sfc.sequence_numeric_column('bbb') + + _, sequence_length = sfc.sequence_input_layer( + features={ + 'aaa': sparse_input_a, + 'bbb': sparse_input_b, + }, + feature_columns=[numeric_column_a, numeric_column_b]) + + with monitored_session.MonitoredSession() as sess: + with self.assertRaisesRegexp( + errors.InvalidArgumentError, + r'\[Condition x == y did not hold element-wise:\] ' + r'\[x \(sequence_input_layer/aaa/sequence_length:0\) = \] \[2 1\] ' + r'\[y \(sequence_input_layer/bbb/sequence_length:0\) = \] \[1 1\]'): + sess.run(sequence_length) + + +class InputLayerTest(test.TestCase): + """Tests input_layer with sequence feature columns.""" + + def test_embedding_column(self): + """Tests that error is raised for sequence embedding column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column_a = fc.embedding_column( + categorical_column_a, dimension=2) + + with self.assertRaisesRegexp( + ValueError, + r'In embedding_column: aaa_embedding\. categorical_column must not be ' + r'of type _SequenceCategoricalColumn\.'): + _ = fc.input_layer( + features={'aaa': sparse_input}, + feature_columns=[embedding_column_a]) + + def test_indicator_column(self): + """Tests that error is raised for sequence indicator column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column_a = fc.indicator_column(categorical_column_a) + + with self.assertRaisesRegexp( + ValueError, + r'In indicator_column: aaa_indicator\. categorical_column must not be ' + r'of type _SequenceCategoricalColumn\.'): + _ = fc.input_layer( + features={'aaa': sparse_input}, + feature_columns=[indicator_column_a]) + + +def _assert_sparse_tensor_value(test_case, expected, actual): + _assert_sparse_tensor_indices_shape(test_case, expected, actual) + + test_case.assertEqual( + np.array(expected.values).dtype, np.array(actual.values).dtype) + test_case.assertAllEqual(expected.values, actual.values) + + +def _assert_sparse_tensor_indices_shape(test_case, expected, actual): + test_case.assertEqual(np.int64, np.array(actual.indices).dtype) + test_case.assertAllEqual(expected.indices, actual.indices) + + test_case.assertEqual(np.int64, np.array(actual.dense_shape).dtype) + test_case.assertAllEqual(expected.dense_shape, actual.dense_shape) + + +class SequenceCategoricalColumnWithIdentityTest(test.TestCase): + + def test_get_sparse_tensors(self): + column = sfc.sequence_categorical_column_with_identity( + 'aaa', num_buckets=3) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 0), + dense_shape=(2, 2)) + expected_sparse_ids = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + values=np.array((1, 2, 0), dtype=np.int64), + dense_shape=(2, 2, 1)) + + id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs})) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, + expected_sparse_ids, + id_weight_pair.id_tensor.eval(session=sess)) + + def test_get_sparse_tensors_inputs3d(self): + """Tests _get_sparse_tensors when the input is already 3D Tensor.""" + column = sfc.sequence_categorical_column_with_identity( + 'aaa', num_buckets=3) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + values=(1, 2, 0), + dense_shape=(2, 2, 1)) + + with self.assertRaisesRegexp( + errors.InvalidArgumentError, + r'Column aaa expected ID tensor of rank 2\.\s*' + r'id_tensor shape:\s*\[2 2 1\]'): + id_weight_pair = column._get_sparse_tensors( + _LazyBuilder({'aaa': inputs})) + with monitored_session.MonitoredSession() as sess: + id_weight_pair.id_tensor.eval(session=sess) + + +class SequenceCategoricalColumnWithHashBucketTest(test.TestCase): + + def test_get_sparse_tensors(self): + column = sfc.sequence_categorical_column_with_hash_bucket( + 'aaa', hash_bucket_size=10) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=('omar', 'stringer', 'marlo'), + dense_shape=(2, 2)) + + expected_sparse_ids = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + # Ignored to avoid hash dependence in test. + values=np.array((0, 0, 0), dtype=np.int64), + dense_shape=(2, 2, 1)) + + id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs})) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_indices_shape( + self, + expected_sparse_ids, + id_weight_pair.id_tensor.eval(session=sess)) + + +class SequenceCategoricalColumnWithVocabularyFileTest(test.TestCase): + + def _write_vocab(self, vocab_strings, file_name): + vocab_file = os.path.join(self.get_temp_dir(), file_name) + with open(vocab_file, 'w') as f: + f.write('\n'.join(vocab_strings)) + return vocab_file + + def setUp(self): + super(SequenceCategoricalColumnWithVocabularyFileTest, self).setUp() + + vocab_strings = ['omar', 'stringer', 'marlo'] + self._wire_vocabulary_file_name = self._write_vocab(vocab_strings, + 'wire_vocabulary.txt') + self._wire_vocabulary_size = 3 + + def test_get_sparse_tensors(self): + column = sfc.sequence_categorical_column_with_vocabulary_file( + key='aaa', + vocabulary_file=self._wire_vocabulary_file_name, + vocabulary_size=self._wire_vocabulary_size) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=('marlo', 'skywalker', 'omar'), + dense_shape=(2, 2)) + expected_sparse_ids = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + values=np.array((2, -1, 0), dtype=np.int64), + dense_shape=(2, 2, 1)) + + id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs})) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, + expected_sparse_ids, + id_weight_pair.id_tensor.eval(session=sess)) + + +class SequenceCategoricalColumnWithVocabularyListTest(test.TestCase): + + def test_get_sparse_tensors(self): + column = sfc.sequence_categorical_column_with_vocabulary_list( + key='aaa', + vocabulary_list=('omar', 'stringer', 'marlo')) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=('marlo', 'skywalker', 'omar'), + dense_shape=(2, 2)) + expected_sparse_ids = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + values=np.array((2, -1, 0), dtype=np.int64), + dense_shape=(2, 2, 1)) + + id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs})) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, + expected_sparse_ids, + id_weight_pair.id_tensor.eval(session=sess)) + + +class SequenceEmbeddingColumnTest(test.TestCase): + + def test_get_sequence_dense_tensor(self): + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + indices=((0, 0), (1, 0), (1, 1), (3, 0)), + values=(2, 0, 1, 1), + dense_shape=(4, 2)) + + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 5.), # id 1 + (7., 11.) # id 2 + ) + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + expected_lookups = [ + # example 0, ids [2] + [[7., 11.], [0., 0.]], + # example 1, ids [0, 1] + [[1., 2.], [3., 5.]], + # example 2, ids [] + [[0., 0.], [0., 0.]], + # example 3, ids [1] + [[3., 5.], [0., 0.]], + ] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = fc.embedding_column( + categorical_column, dimension=embedding_dimension, + initializer=_initializer) + + embedding_lookup, _ = embedding_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual( + ('embedding_weights:0',), tuple([v.name for v in global_vars])) + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(embedding_values, global_vars[0].eval(session=sess)) + self.assertAllEqual(expected_lookups, embedding_lookup.eval(session=sess)) + + def test_sequence_length(self): + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + expected_sequence_length = [1, 2] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = fc.embedding_column( + categorical_column, dimension=2) + + _, sequence_length = embedding_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + sequence_length = sess.run(sequence_length) + self.assertAllEqual(expected_sequence_length, sequence_length) + self.assertEqual(np.int64, sequence_length.dtype) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [] + # example 1, ids [2] + # example 2, ids [0, 1] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [] + indices=((1, 0), (2, 0), (2, 1), (4, 0)), + values=(2, 0, 1, 1), + dense_shape=(6, 2)) + expected_sequence_length = [0, 1, 2, 0, 1, 0] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = fc.embedding_column( + categorical_column, dimension=2) + + _, sequence_length = embedding_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +class SequenceIndicatorColumnTest(test.TestCase): + + def test_get_sequence_dense_tensor(self): + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + indices=((0, 0), (1, 0), (1, 1), (3, 0)), + values=(2, 0, 1, 1), + dense_shape=(4, 2)) + + expected_lookups = [ + # example 0, ids [2] + [[0., 0., 1.], [0., 0., 0.]], + # example 1, ids [0, 1] + [[1., 0., 0.], [0., 1., 0.]], + # example 2, ids [] + [[0., 0., 0.], [0., 0., 0.]], + # example 3, ids [1] + [[0., 1., 0.], [0., 0., 0.]], + ] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column = fc.indicator_column(categorical_column) + + indicator_tensor, _ = indicator_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected_lookups, indicator_tensor.eval(session=sess)) + + def test_sequence_length(self): + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + expected_sequence_length = [1, 2] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column = fc.indicator_column(categorical_column) + + _, sequence_length = indicator_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + sequence_length = sess.run(sequence_length) + self.assertAllEqual(expected_sequence_length, sequence_length) + self.assertEqual(np.int64, sequence_length.dtype) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [] + # example 1, ids [2] + # example 2, ids [0, 1] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [] + indices=((1, 0), (2, 0), (2, 1), (4, 0)), + values=(2, 0, 1, 1), + dense_shape=(6, 2)) + expected_sequence_length = [0, 1, 2, 0, 1, 0] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column = fc.indicator_column(categorical_column) + + _, sequence_length = indicator_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +class SequenceNumericColumnTest(test.TestCase): + + def test_get_sequence_dense_tensor(self): + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + expected_dense_tensor = [ + [[0.], [1.]], + [[10.], [0.]], + ] + numeric_column = sfc.sequence_numeric_column('aaa') + + dense_tensor, _ = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + def test_get_sequence_dense_tensor_with_shape(self): + """Tests get_sequence_dense_tensor with shape !=(1,).""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0., 1., 2.], [3., 4., 5.]] + # example 1, [[10., 11., 12.]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), + (1, 0), (1, 1), (1, 2)), + values=(0., 1., 2., 3., 4., 5., 10., 11., 12.), + dense_shape=(2, 6)) + expected_dense_tensor = [ + [[0., 1., 2.], [3., 4., 5.]], + [[10., 11., 12.], [0., 0., 0.]], + ] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(3,)) + + dense_tensor, _ = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + def test_get_dense_tensor_multi_dim(self): + """Tests get_sequence_dense_tensor for multi-dim numeric_column.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]] + # example 1, [[[10., 11.], [12., 13.]]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), + (1, 0), (1, 1), (1, 2), (1, 3)), + values=(0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + dense_shape=(2, 8)) + expected_dense_tensor = [ + [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]], + [[[10., 11.], [12., 13.]], [[0., 0.], [0., 0.]]], + ] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) + + dense_tensor, _ = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + def test_sequence_length(self): + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0., 1., 2.], [3., 4., 5.]] + # example 1, [[10., 11., 12.]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), + (1, 0), (1, 1), (1, 2)), + values=(0., 1., 2., 3., 4., 5., 10., 11., 12.), + dense_shape=(2, 6)) + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(3,)) + + _, sequence_length = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + sequence_length = sess.run(sequence_length) + self.assertAllEqual(expected_sequence_length, sequence_length) + self.assertEqual(np.int64, sequence_length.dtype) + + def test_sequence_length_with_shape(self): + """Tests _sequence_length with shape !=(1,).""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa') + + _, sequence_length = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [] + # example 1, values [[0.], [1.]] + # example 2, [[2.]] + # example 3, values [] + # example 4, [[3.]] + # example 5, values [] + indices=((1, 0), (1, 1), (2, 0), (4, 0)), + values=(0., 1., 2., 3.), + dense_shape=(6, 2)) + expected_sequence_length = [0, 2, 1, 0, 1, 0] + numeric_column = sfc.sequence_numeric_column('aaa') + + _, sequence_length = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +if __name__ == '__main__': + test.main() 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..4ed7268e7a921284eed7767d870e56ecac39a3b1 --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column.py @@ -0,0 +1,325 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Experimental methods for tf.feature_column sequence input.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +import abc +import collections + + +from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import check_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import parsing_ops +from tensorflow.python.ops import sparse_ops +from tensorflow.python.ops import variable_scope + +# TODO(b/73160931): Fix pydoc. +# pylint: disable=g-doc-args,missing-docstring,protected-access +# TODO(b/73827486): Support SequenceExample. + + +def sequence_input_layer( + features, + feature_columns, + weight_collections=None, + trainable=True, + scope=None): + """"Builds input layer for sequence input. + + All `feature_columns` must be sequence dense columns with the same + `sequence_length`. The output of this method can be fed into sequence + networks, such as RNN. + + The output of this method is a 3D `Tensor` of shape `[batch_size, T, D]`. + `T` is the maximum sequence length for this batch, which could differ from + batch to batch. + + If multiple `feature_columns` are given with `Di` `num_elements` each, their + outputs are concatenated. So, the final `Tensor` has shape + `[batch_size, T, D0 + D1 + ... + Dn]`. + + Example: + + ```python + rating = sequence_numeric_column('rating') + watches = sequence_categorical_column_with_identity( + 'watches', num_buckets=1000) + watches_embedding = embedding_column(watches, dimension=10) + columns = [rating, watches] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + input_layer, sequence_length = sequence_input_layer(features, columns) + + rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) + outputs, state = tf.nn.dynamic_rnn( + rnn_cell, inputs=input_layer, sequence_length=sequence_length) + ``` + + Returns: + An `(input_layer, sequence_length)` tuple where: + - input_layer: A float `Tensor` of shape `[batch_size, T, D]`. + `T` is the maximum sequence length for this batch, which could differ + from batch to batch. `D` is the sum of `num_elements` for all + `feature_columns`. + - sequence_length: An int `Tensor` of shape `[batch_size]`. The sequence + length for each example. + Raises: + ValueError: If any of the `feature_columns` is the wrong type. + """ + feature_columns = fc._clean_feature_columns(feature_columns) + for c in feature_columns: + if not isinstance(c, _SequenceDenseColumn): + raise ValueError( + 'All feature_columns must be of type _SequenceDenseColumn. ' + 'Given (type {}): {}'.format(type(c), c)) + + with variable_scope.variable_scope( + scope, default_name='sequence_input_layer', values=features.values()): + builder = fc._LazyBuilder(features) + output_tensors = [] + sequence_lengths = [] + ordered_columns = [] + for column in sorted(feature_columns, key=lambda x: x.name): + ordered_columns.append(column) + with variable_scope.variable_scope( + None, default_name=column._var_scope_name): + dense_tensor, sequence_length = column._get_sequence_dense_tensor( + builder, + weight_collections=weight_collections, + trainable=trainable) + # Flattens the final dimension to produce a 3D Tensor. + num_elements = column._variable_shape.num_elements() + shape = array_ops.shape(dense_tensor) + output_tensors.append( + array_ops.reshape( + dense_tensor, + shape=array_ops.concat([shape[:2], [num_elements]], axis=0))) + sequence_lengths.append(sequence_length) + fc._verify_static_batch_size_equality(output_tensors, ordered_columns) + # TODO(b/73160931): Verify sequence_length equality. + return array_ops.concat(output_tensors, -1), sequence_lengths[0] + + +# TODO(b/73160931): Add remaining categorical columns. +def sequence_categorical_column_with_identity( + key, num_buckets, default_value=None): + return _SequenceCategoricalColumn( + fc.categorical_column_with_identity( + key=key, + num_buckets=num_buckets, + default_value=default_value)) + + +# TODO(b/73160931): Merge with embedding_column +def _sequence_embedding_column( + categorical_column, dimension, initializer=None, ckpt_to_load_from=None, + tensor_name_in_ckpt=None, max_norm=None, trainable=True): + if not isinstance(categorical_column, _SequenceCategoricalColumn): + raise ValueError( + 'categorical_column must be of type _SequenceCategoricalColumn. ' + 'Given (type {}): {}'.format( + type(categorical_column), categorical_column)) + return _SequenceEmbeddingColumn( + fc.embedding_column( + categorical_column, + dimension=dimension, + initializer=initializer, + ckpt_to_load_from=ckpt_to_load_from, + tensor_name_in_ckpt=tensor_name_in_ckpt, + max_norm=max_norm, + trainable=trainable)) + + +def sequence_numeric_column( + key, + shape=(1,), + default_value=0., + dtype=dtypes.float32): + # TODO(b/73160931): Add validations. + return _SequenceNumericColumn( + key, + shape=shape, + default_value=default_value, + dtype=dtype) + + +class _SequenceDenseColumn(fc._FeatureColumn): + """Represents dense sequence data.""" + + __metaclass__ = abc.ABCMeta + + TensorSequenceLengthPair = collections.namedtuple( # pylint: disable=invalid-name + 'TensorSequenceLengthPair', ['dense_tensor', 'sequence_length']) + + @abc.abstractproperty + def _variable_shape(self): + """`TensorShape` without batch and sequence dimensions.""" + pass + + @abc.abstractmethod + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + """Returns a `TensorSequenceLengthPair`.""" + pass + + +def _sequence_length_from_sparse_tensor(sp_tensor, num_elements=1): + with ops.name_scope(None, 'sequence_length') as name_scope: + row_ids = sp_tensor.indices[:, 0] + column_ids = sp_tensor.indices[:, 1] + column_ids += array_ops.ones_like(column_ids) + seq_length = ( + math_ops.segment_max(column_ids, segment_ids=row_ids) / num_elements) + # If the last n rows do not have ids, seq_length will have shape + # [batch_size - n]. Pad the remaining values with zeros. + n_pad = array_ops.shape(sp_tensor)[:1] - array_ops.shape(seq_length)[:1] + padding = array_ops.zeros(n_pad, dtype=seq_length.dtype) + return array_ops.concat([seq_length, padding], axis=0, name=name_scope) + + +class _SequenceCategoricalColumn( + fc._CategoricalColumn, + collections.namedtuple( + '_SequenceCategoricalColumn', ['categorical_column'])): + + @property + def name(self): + return self.categorical_column.name + + @property + def _parse_example_spec(self): + return self.categorical_column._parse_example_spec + + def _transform_feature(self, inputs): + return self.categorical_column._transform_feature(inputs) + + @property + def _num_buckets(self): + return self.categorical_column._num_buckets + + def _get_sparse_tensors(self, inputs, weight_collections=None, + trainable=None): + sparse_tensors = self.categorical_column._get_sparse_tensors(inputs) + id_tensor = sparse_tensors.id_tensor + weight_tensor = sparse_tensors.weight_tensor + # Expands final dimension, so that embeddings are not combined during + # embedding lookup. + check_id_rank = check_ops.assert_equal( + array_ops.rank(id_tensor), 2, + data=[ + 'Column {} expected ID tensor of rank 2. '.format(self.name), + 'id_tensor shape: ', array_ops.shape(id_tensor)]) + with ops.control_dependencies([check_id_rank]): + id_tensor = sparse_ops.sparse_reshape( + id_tensor, + shape=array_ops.concat([id_tensor.dense_shape, [1]], axis=0)) + if weight_tensor is not None: + check_weight_rank = check_ops.assert_equal( + array_ops.rank(weight_tensor), 2, + data=[ + 'Column {} expected weight tensor of rank 2.'.format(self.name), + 'weight_tensor shape:', array_ops.shape(weight_tensor)]) + with ops.control_dependencies([check_weight_rank]): + weight_tensor = sparse_ops.sparse_reshape( + weight_tensor, + shape=array_ops.concat([weight_tensor.dense_shape, [1]], axis=0)) + return fc._CategoricalColumn.IdWeightPair(id_tensor, weight_tensor) + + def _sequence_length(self, inputs): + sparse_tensors = self.categorical_column._get_sparse_tensors(inputs) + return _sequence_length_from_sparse_tensor(sparse_tensors.id_tensor) + + +class _SequenceEmbeddingColumn( + _SequenceDenseColumn, + collections.namedtuple('_SequenceEmbeddingColumn', ['embedding_column'])): + + @property + def name(self): + return self.embedding_column.name + + @property + def _parse_example_spec(self): + return self.embedding_column._parse_example_spec + + def _transform_feature(self, inputs): + return self.embedding_column._transform_feature(inputs) + + @property + def _variable_shape(self): + return self.embedding_column._variable_shape + + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + dense_tensor = self.embedding_column._get_dense_tensor( + inputs=inputs, + weight_collections=weight_collections, + trainable=trainable) + sequence_length = self.embedding_column.categorical_column._sequence_length( + inputs) + return _SequenceDenseColumn.TensorSequenceLengthPair( + dense_tensor=dense_tensor, sequence_length=sequence_length) + + +class _SequenceNumericColumn( + _SequenceDenseColumn, + collections.namedtuple( + '_SequenceNumericColumn', + ['key', 'shape', 'default_value', 'dtype'])): + + @property + def name(self): + return self.key + + @property + def _parse_example_spec(self): + return {self.key: parsing_ops.VarLenFeature(self.dtype)} + + def _transform_feature(self, inputs): + return inputs.get(self.key) + + @property + def _variable_shape(self): + return tensor_shape.TensorShape(self.shape) + + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + # Do nothing with weight_collections and trainable since no variables are + # created in this function. + del weight_collections + del trainable + sp_tensor = inputs.get(self) + dense_tensor = sparse_ops.sparse_tensor_to_dense( + sp_tensor, default_value=self.default_value) + # Reshape into [batch_size, T, variable_shape]. + dense_shape = array_ops.concat( + [array_ops.shape(dense_tensor)[:1], [-1], self._variable_shape], + axis=0) + dense_tensor = array_ops.reshape(dense_tensor, shape=dense_shape) + sequence_length = _sequence_length_from_sparse_tensor( + sp_tensor, num_elements=self._variable_shape.num_elements()) + return _SequenceDenseColumn.TensorSequenceLengthPair( + dense_tensor=dense_tensor, sequence_length=sequence_length) + +# pylint: enable=g-doc-args,missing-docstring,protected-access diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column_test.py new file mode 100644 index 0000000000000000000000000000000000000000..59674869a27c3a40ab9cb3dcede384d1cda7ce27 --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column_test.py @@ -0,0 +1,471 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for sequential_feature_column.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.feature_column.python.feature_column import sequential_feature_column as sfc +from tensorflow.python.feature_column.feature_column import _LazyBuilder +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops +from tensorflow.python.framework import sparse_tensor +from tensorflow.python.platform import test +from tensorflow.python.training import monitored_session + + +class SequenceInputLayerTest(test.TestCase): + + def test_embedding_column(self): + vocabulary_size = 3 + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [1] + # example 1, ids [2, 0] + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 0), + dense_shape=(2, 2)) + + embedding_dimension_a = 2 + embedding_values_a = ( + (1., 2.), # id 0 + (3., 4.), # id 1 + (5., 6.) # id 2 + ) + embedding_dimension_b = 3 + embedding_values_b = ( + (11., 12., 13.), # id 0 + (14., 15., 16.), # id 1 + (17., 18., 19.) # id 2 + ) + def _get_initializer(embedding_dimension, embedding_values): + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + return _initializer + + expected_input_layer = [ + # example 0, ids_a [2], ids_b [1] + [[5., 6., 14., 15., 16.], [0., 0., 0., 0., 0.]], + # example 1, ids_a [0, 1], ids_b [2, 0] + [[1., 2., 17., 18., 19.], [3., 4., 11., 12., 13.]], + ] + expected_sequence_length = [1, 2] + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column_a = sfc._sequence_embedding_column( + categorical_column_a, dimension=embedding_dimension_a, + initializer=_get_initializer(embedding_dimension_a, embedding_values_a)) + categorical_column_b = sfc.sequence_categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + embedding_column_b = sfc._sequence_embedding_column( + categorical_column_b, dimension=embedding_dimension_b, + initializer=_get_initializer(embedding_dimension_b, embedding_values_b)) + + input_layer, sequence_length = sfc.sequence_input_layer( + features={ + 'aaa': sparse_input_a, + 'bbb': sparse_input_b, + }, + # Test that columns are reordered alphabetically. + feature_columns=[embedding_column_b, embedding_column_a]) + + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual( + ('sequence_input_layer/aaa_embedding/embedding_weights:0', + 'sequence_input_layer/bbb_embedding/embedding_weights:0'), + tuple([v.name for v in global_vars])) + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(embedding_values_a, global_vars[0].eval(session=sess)) + self.assertAllEqual(embedding_values_b, global_vars[1].eval(session=sess)) + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_numeric_column(self): + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + expected_input_layer = [ + [[0.], [1.]], + [[10.], [0.]], + ] + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa') + + input_layer, sequence_length = sfc.sequence_input_layer( + features={'aaa': sparse_input}, + feature_columns=[numeric_column]) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_numeric_column_multi_dim(self): + """Tests sequence_input_layer for multi-dimensional numeric_column.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]] + # example 1, [[[10., 11.], [12., 13.]]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), + (1, 0), (1, 1), (1, 2), (1, 3)), + values=(0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + dense_shape=(2, 8)) + # The output of numeric_column._get_dense_tensor should be flattened. + expected_input_layer = [ + [[0., 1., 2., 3.], [4., 5., 6., 7.]], + [[10., 11., 12., 13.], [0., 0., 0., 0.]], + ] + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) + + input_layer, sequence_length = sfc.sequence_input_layer( + features={'aaa': sparse_input}, + feature_columns=[numeric_column]) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +def _assert_sparse_tensor_value(test_case, expected, actual): + test_case.assertEqual(np.int64, np.array(actual.indices).dtype) + test_case.assertAllEqual(expected.indices, actual.indices) + + test_case.assertEqual( + np.array(expected.values).dtype, np.array(actual.values).dtype) + test_case.assertAllEqual(expected.values, actual.values) + + test_case.assertEqual(np.int64, np.array(actual.dense_shape).dtype) + test_case.assertAllEqual(expected.dense_shape, actual.dense_shape) + + +class SequenceCategoricalColumnWithIdentityTest(test.TestCase): + + def test_get_sparse_tensors(self): + column = sfc.sequence_categorical_column_with_identity( + 'aaa', num_buckets=3) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 0), + dense_shape=(2, 2)) + expected_sparse_ids = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + values=np.array((1, 2, 0), dtype=np.int64), + dense_shape=(2, 2, 1)) + + id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs})) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, + expected_sparse_ids, + id_weight_pair.id_tensor.eval(session=sess)) + + def test_get_sparse_tensors_inputs3d(self): + """Tests _get_sparse_tensors when the input is already 3D Tensor.""" + column = sfc.sequence_categorical_column_with_identity( + 'aaa', num_buckets=3) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)), + values=(1, 2, 0), + dense_shape=(2, 2, 1)) + + with self.assertRaisesRegexp( + errors.InvalidArgumentError, + r'Column aaa expected ID tensor of rank 2\.\s*' + r'id_tensor shape:\s*\[2 2 1\]'): + id_weight_pair = column._get_sparse_tensors( + _LazyBuilder({'aaa': inputs})) + with monitored_session.MonitoredSession() as sess: + id_weight_pair.id_tensor.eval(session=sess) + + def test_sequence_length(self): + column = sfc.sequence_categorical_column_with_identity( + 'aaa', num_buckets=3) + inputs = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 0), + dense_shape=(2, 2)) + expected_sequence_length = [1, 2] + + sequence_length = column._sequence_length(_LazyBuilder({'aaa': inputs})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_with_zeros(self): + column = sfc.sequence_categorical_column_with_identity( + 'aaa', num_buckets=3) + inputs = sparse_tensor.SparseTensorValue( + indices=((1, 0), (3, 0), (3, 1)), + values=(1, 2, 0), + dense_shape=(5, 2)) + expected_sequence_length = [0, 1, 0, 2, 0] + + sequence_length = column._sequence_length(_LazyBuilder({'aaa': inputs})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +class SequenceEmbeddingColumnTest(test.TestCase): + + def test_get_sequence_dense_tensor(self): + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + indices=((0, 0), (1, 0), (1, 1), (3, 0)), + values=(2, 0, 1, 1), + dense_shape=(4, 2)) + + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 5.), # id 1 + (7., 11.) # id 2 + ) + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + expected_lookups = [ + # example 0, ids [2] + [[7., 11.], [0., 0.]], + # example 1, ids [0, 1] + [[1., 2.], [3., 5.]], + # example 2, ids [] + [[0., 0.], [0., 0.]], + # example 3, ids [1] + [[3., 5.], [0., 0.]], + ] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = sfc._sequence_embedding_column( + categorical_column, dimension=embedding_dimension, + initializer=_initializer) + + embedding_lookup, _ = embedding_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual( + ('embedding_weights:0',), tuple([v.name for v in global_vars])) + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(embedding_values, global_vars[0].eval(session=sess)) + self.assertAllEqual(expected_lookups, embedding_lookup.eval(session=sess)) + + def test_sequence_length(self): + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + expected_sequence_length = [1, 2] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = sfc._sequence_embedding_column( + categorical_column, dimension=2) + + _, sequence_length = embedding_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [] + # example 1, ids [2] + # example 2, ids [0, 1] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [] + indices=((1, 0), (2, 0), (2, 1), (4, 0)), + values=(2, 0, 1, 1), + dense_shape=(6, 2)) + expected_sequence_length = [0, 1, 2, 0, 1, 0] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = sfc._sequence_embedding_column( + categorical_column, dimension=2) + + _, sequence_length = embedding_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +class SequenceNumericColumnTest(test.TestCase): + + def test_get_sequence_dense_tensor(self): + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + expected_dense_tensor = [ + [[0.], [1.]], + [[10.], [0.]], + ] + numeric_column = sfc.sequence_numeric_column('aaa') + + dense_tensor, _ = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + def test_get_sequence_dense_tensor_with_shape(self): + """Tests get_sequence_dense_tensor with shape !=(1,).""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0., 1., 2.], [3., 4., 5.]] + # example 1, [[10., 11., 12.]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), + (1, 0), (1, 1), (1, 2)), + values=(0., 1., 2., 3., 4., 5., 10., 11., 12.), + dense_shape=(2, 6)) + expected_dense_tensor = [ + [[0., 1., 2.], [3., 4., 5.]], + [[10., 11., 12.], [0., 0., 0.]], + ] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(3,)) + + dense_tensor, _ = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + def test_get_dense_tensor_multi_dim(self): + """Tests get_sequence_dense_tensor for multi-dim numeric_column.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]] + # example 1, [[[10., 11.], [12., 13.]]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), + (1, 0), (1, 1), (1, 2), (1, 3)), + values=(0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + dense_shape=(2, 8)) + expected_dense_tensor = [ + [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]], + [[[10., 11.], [12., 13.]], [[0., 0.], [0., 0.]]], + ] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) + + dense_tensor, _ = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + def test_sequence_length(self): + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0., 1., 2.], [3., 4., 5.]] + # example 1, [[10., 11., 12.]] + indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), + (1, 0), (1, 1), (1, 2)), + values=(0., 1., 2., 3., 4., 5., 10., 11., 12.), + dense_shape=(2, 6)) + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa', shape=(3,)) + + _, sequence_length = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_with_shape(self): + """Tests _sequence_length with shape !=(1,).""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + expected_sequence_length = [2, 1] + numeric_column = sfc.sequence_numeric_column('aaa') + + _, sequence_length = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [] + # example 1, values [[0.], [1.]] + # example 2, [[2.]] + # example 3, values [] + # example 4, [[3.]] + # example 5, values [] + indices=((1, 0), (1, 1), (2, 0), (4, 0)), + values=(0., 1., 2., 3.), + dense_shape=(6, 2)) + expected_sequence_length = [0, 2, 1, 0, 1, 0] + numeric_column = sfc.sequence_numeric_column('aaa') + + _, sequence_length = numeric_column._get_sequence_dense_tensor( + _LazyBuilder({'aaa': sparse_input})) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc index e61221a6b0d34373279a379f356c99c379488182..35341406a08dc681c861aea30fcff784e3b963ef 100644 --- a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc +++ b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc @@ -256,6 +256,9 @@ Status ReadInfoFile(const string& filename, uint32* width, uint32* height, if (p != std::string::npos) { string rgb24 = line.substr(p + 9, line.find(" ", p + 9)); rgb24 = rgb24.substr(0, rgb24.find(",")); + // Strip anything after " ", in case the format is + // `640x360 [SAR 1:1 DAR 16:9]` + rgb24 = rgb24.substr(0, rgb24.find(" ")); string rgb24_width = rgb24.substr(0, rgb24.find("x")); string rgb24_height = rgb24.substr(rgb24_width.length() + 1); if (strings::safe_strtou32(rgb24_width, &width_value) && @@ -270,8 +273,10 @@ Status ReadInfoFile(const string& filename, uint32* width, uint32* height, // We only look for the first stream mapping to have the number of the // frames. // Once processed we will not further process stream mapping section. - if (line.find("frame= ") == 0) { - string number = line.substr(8, line.find(" ", 8)); + if (line.find("frame=") == 0) { + // The format might be `frame= 166 ` or `frame=12488 ` + string number = line.substr(6); + number = number.substr(number.find_first_not_of(" ")); number = number.substr(0, number.find(" ")); if (strings::safe_strtou32(number, &frames_value)) { in_mapping = false; diff --git a/tensorflow/contrib/framework/BUILD b/tensorflow/contrib/framework/BUILD index 9e5f54f0973eae899ca65e4098358107053cb7d4..ac043fda0638e61f422e769ab3047a53a1b377bd 100644 --- a/tensorflow/contrib/framework/BUILD +++ b/tensorflow/contrib/framework/BUILD @@ -28,7 +28,6 @@ tf_custom_op_py_library( "python/framework/graph_util.py", "python/framework/tensor_util.py", "python/ops/__init__.py", - "python/ops/accumulate_n_v2.py", "python/ops/arg_scope.py", "python/ops/audio_ops.py", "python/ops/checkpoint_ops.py", @@ -63,7 +62,9 @@ tf_custom_op_py_library( "//tensorflow/python:math_ops", "//tensorflow/python:platform", "//tensorflow/python:pywrap_tensorflow", + "//tensorflow/python:resource_variable_ops", "//tensorflow/python:script_ops", + "//tensorflow/python:smart_cond", "//tensorflow/python:sparse_tensor", "//tensorflow/python:state_ops", "//tensorflow/python:state_ops_gen", @@ -161,23 +162,6 @@ py_test( ], ) -py_test( - name = "accumulate_n_v2_test", - size = "small", - srcs = ["python/ops/accumulate_n_v2_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":framework_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:gradients", - "//tensorflow/python:platform_test", - "//tensorflow/python:variables", - "//third_party/py/numpy", - ], -) - cuda_py_test( name = "critical_section_test", size = "medium", @@ -185,31 +169,14 @@ cuda_py_test( additional_deps = [ "//tensorflow/python:client_testlib", ":framework_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:control_flow_ops", "//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", - srcs = ["python/ops/accumulate_n_v2_eager_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":framework_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python/eager:backprop", - "//tensorflow/python/eager:context", - "//tensorflow/python/eager:tape", - "//third_party/py/numpy", + "//tensorflow/python:tensor_array_ops", ], ) diff --git a/tensorflow/contrib/framework/__init__.py b/tensorflow/contrib/framework/__init__.py index 503b868aaa67b26faaad9eb82f6b864d896a44af..3398b3fd1c1036091bfadf548f7d44dbf9eb1046 100644 --- a/tensorflow/contrib/framework/__init__.py +++ b/tensorflow/contrib/framework/__init__.py @@ -71,6 +71,7 @@ See the @{$python/contrib.framework} guide. @@model_variable @@variable @@VariableDeviceChooser +@@convolutional_delta_orthogonal @@zero_initializer @@load_checkpoint @@ -85,7 +86,16 @@ See the @{$python/contrib.framework} guide. @@py_func @@sort +@@get_placeholders + +@@smart_cond +@@smart_constant_value +@@smart_case + @@CriticalSection + +@@BoundedTensorSpec +@@TensorSpec """ from __future__ import absolute_import @@ -99,7 +109,12 @@ 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.smart_cond import smart_case +from tensorflow.python.framework.smart_cond import smart_cond +from tensorflow.python.framework.smart_cond import smart_constant_value +from tensorflow.python.framework.tensor_spec import BoundedTensorSpec +from tensorflow.python.framework.tensor_spec import TensorSpec +from tensorflow.python.ops.init_ops import convolutional_delta_orthogonal from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = ['nest'] diff --git a/tensorflow/contrib/framework/python/framework/experimental_test.py b/tensorflow/contrib/framework/python/framework/experimental_test.py index 8e54e09e04ee3c0ddbd4fa84cc0912cb70c93e62..cfdc7df7d8fd4c1406bf447a79038ac33b11e047 100644 --- a/tensorflow/contrib/framework/python/framework/experimental_test.py +++ b/tensorflow/contrib/framework/python/framework/experimental_test.py @@ -49,7 +49,6 @@ class ExperimentalTest(test.TestCase): "\nTHIS FUNCTION IS EXPERIMENTAL. It may change or " "be removed at any time, and without warning." "\n" - "\n" "\nArgs:" "\n arg0: Arg 0." "\n arg1: Arg 1." diff --git a/tensorflow/contrib/framework/python/framework/graph_util.py b/tensorflow/contrib/framework/python/framework/graph_util.py index a18ff2320d99726bb355ff6179fc97a070c2fec7..2703224b1bf62831b6088558d4f93950fe938c10 100644 --- a/tensorflow/contrib/framework/python/framework/graph_util.py +++ b/tensorflow/contrib/framework/python/framework/graph_util.py @@ -85,14 +85,19 @@ def fuse_op(graph_def, input_nodes, output_nodes, output_dtypes, if n not in reachable_by_input and n not in output_nodes_set: # n is between input and output, i.e., part of the fused op next_to_visit = [n] + visited = set() while next_to_visit: cur_node = next_to_visit[0] + visited.add(cur_node) del next_to_visit[0] if cur_node in reachable_by_input and cur_node not in input_nodes_set: raise TypeError("Node %s uses input %s not in input_nodes." % (n, cur_node)) if cur_node not in input_nodes_set: - next_to_visit += name_to_input_name[cur_node] + next_to_visit += [ + input_node for input_node in name_to_input_name[cur_node] + if input_node not in visited + ] elif n not in reachable_by_input: nodes_post_output.append(n) @@ -133,6 +138,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/graph_util_test.py b/tensorflow/contrib/framework/python/framework/graph_util_test.py index b8a6d109e19211d271c2b15bac66ddacd38fe395..812c5fbd8cb759aef6eb1aad532c03794b2ceaf4 100644 --- a/tensorflow/contrib/framework/python/framework/graph_util_test.py +++ b/tensorflow/contrib/framework/python/framework/graph_util_test.py @@ -42,7 +42,8 @@ class GraphUtilTest(test.TestCase): graph_def = graph_pb2.GraphDef() node_a = GetNewNode('A', 'Placeholder', []) node_b = GetNewNode('B', 'Op1', ['A']) - node_c = GetNewNode('C', 'Op1', ['B']) + # A loop in the part that will be fused. + node_c = GetNewNode('C', 'Op1', ['B', 'C']) node_d = GetNewNode('D', 'Op1', ['C']) node_e = GetNewNode('E', 'Op1', ['D']) graph_def.node.extend([node_a, node_b, node_c, node_d, node_e]) diff --git a/tensorflow/contrib/framework/python/framework/tensor_util_test.py b/tensorflow/contrib/framework/python/framework/tensor_util_test.py index 8cdb340f2ddd9b3a7f55c1937ef045f4627e99be..a2834b648933772cab53002462c3edbe9a553e94 100644 --- a/tensorflow/contrib/framework/python/framework/tensor_util_test.py +++ b/tensorflow/contrib/framework/python/framework/tensor_util_test.py @@ -209,6 +209,7 @@ class WithShapeTest(test.TestCase): self.assertRaisesRegexp(errors_impl.OpError, "Wrong shape", tensor_2x2.eval, {tensor_no_shape: [42.0]}) + @test_util.enable_c_shapes def test_with_shape_partial(self): with self.test_session(): tensor_partial_shape = array_ops.placeholder(dtypes.float32) diff --git a/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py b/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py deleted file mode 100644 index 2375ee4f550616ff60d20b87b5773704d8fbbe1e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py +++ /dev/null @@ -1,111 +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. -# ============================================================================== -"""Ops that will eventually be folded into tensorflow/python/ops/math_ops.py -""" -from __future__ import absolute_import -from __future__ import division -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 gen_math_ops -from tensorflow.python.ops import math_ops - - - -def accumulate_n_v2(inputs, shape=None, tensor_dtype=None, name=None): - """Returns the element-wise sum of a list of tensors. - - Optionally, pass `shape` and `tensor_dtype` for shape and type checking, - otherwise, these are inferred. - - `tf.accumulate_n_v2` performs the same operation as `tf.add_n`, but does not - wait for all of its inputs to be ready before beginning to sum. This can - save memory if inputs are ready at different times, since minimum temporary - storage is proportional to the output size rather than the inputs size. - - Unlike the original `accumulate_n`, `accumulate_n_v2` is differentiable. - - For example: - - ```python - a = tf.constant([[1, 2], [3, 4]]) - b = tf.constant([[5, 0], [0, 6]]) - tf.accumulate_n_v2([a, b, a]) # [[7, 4], [6, 14]] - - # Explicitly pass shape and type - tf.accumulate_n_v2([a, b, a], shape=[2, 2], tensor_dtype=tf.int32) - # [[7, 4], - # [6, 14]] - ``` - - Args: - inputs: A list of `Tensor` objects, each with same shape and type. - shape: Shape of elements of `inputs`. - tensor_dtype: The type of `inputs`. - name: A name for the operation (optional). - - Returns: - A `Tensor` of same shape and type as the elements of `inputs`. - - Raises: - ValueError: If `inputs` don't all have same shape and dtype or the shape - cannot be inferred. - """ - _INPUTS_ERR_MSG = ValueError("inputs must be a list of at least one Tensor" - "with the same dtype and shape") - if not inputs or not isinstance(inputs, (list, tuple)): - raise _INPUTS_ERR_MSG - inputs = ops.convert_n_to_tensor_or_indexed_slices(inputs) - if not all(isinstance(x, ops.Tensor) for x in inputs): - raise _INPUTS_ERR_MSG - if not all(x.dtype == inputs[0].dtype for x in inputs): - raise _INPUTS_ERR_MSG - if shape is not None: - shape = tensor_shape.as_shape(shape) - else: - shape = tensor_shape.unknown_shape() - for input_tensor in inputs: - if isinstance(input_tensor, ops.Tensor): - shape = shape.merge_with(input_tensor.get_shape()) - - # tensor_dtype is for safety only; operator's output type computed in C++ - if tensor_dtype is not None and tensor_dtype != inputs[0].dtype: - raise TypeError("tensor_dtype is {}, but input is of type {}" - .format(tensor_dtype, inputs[0].dtype)) - - if len(inputs) == 1 and name is None: - return inputs[0] - elif len(inputs) == 1 and name is not None: - return array_ops.identity(inputs[0], name=name) - elif context.in_eager_mode(): - # TemporaryVariable not currently supported in eager mode; fall back - # onto AddN for now. - # TODO(frreiss) remove this once the lifetime of eager variables gets - # addressed - return math_ops.add_n(inputs, name=name) - else: - return gen_math_ops._accumulate_nv2(inputs, name=name, shape=shape) - -# The following code should eventually be merged into -# tensorflow/python/ops/math_grad.py -@ops.RegisterGradient("AccumulateNV2") -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/arg_scope.py b/tensorflow/contrib/framework/python/ops/arg_scope.py index 409657fe1da0e5540cd2ad6070d86737c039e91f..3cad1fee1984042e3a9ab91a0af70cbaca25cece 100644 --- a/tensorflow/contrib/framework/python/ops/arg_scope.py +++ b/tensorflow/contrib/framework/python/ops/arg_scope.py @@ -142,7 +142,7 @@ def arg_scope(list_ops_or_scope, **kwargs): else: # Assumes that list_ops_or_scope is a list/tuple of ops with kwargs. if not isinstance(list_ops_or_scope, (list, tuple)): - raise TypeError('list_ops_or_scope must either be a list/tuple or reused' + raise TypeError('list_ops_or_scope must either be a list/tuple or reused ' 'scope (i.e. dict)') try: current_scope = current_arg_scope().copy() diff --git a/tensorflow/contrib/framework/python/ops/critical_section_ops.py b/tensorflow/contrib/framework/python/ops/critical_section_ops.py index 182fec924febb74a23b82b1664d137f033f3b1b4..cc19372acf956371c2d029c7b8eb5534c3789413 100644 --- a/tensorflow/contrib/framework/python/ops/critical_section_ops.py +++ b/tensorflow/contrib/framework/python/ops/critical_section_ops.py @@ -27,7 +27,11 @@ 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.framework import tensor_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_resource_variable_ops +from tensorflow.python.ops import tensor_array_ops from tensorflow.python.util import nest @@ -38,7 +42,8 @@ CRITICAL_SECTION_EXECUTIONS = "critical_section_executions" class _ExecutionSignature( collections.namedtuple("_ExecutionSignature", - ("op", "exclusive_resource_access"))): + ("op", "handle", + "resources", "exclusive_resource_access"))): """A class storing an `ExecuteInCriticalResource` op and associated attrs.""" pass @@ -112,16 +117,18 @@ class CriticalSection(object): ``` """ - def __init__(self, name=None, critical_section_def=None, import_scope=None): + def __init__(self, name=None, shared_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.") + raise ValueError("critical_section_def and shared_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) + self._init_from_args(name, shared_name) - def _init_from_proto(self, critical_section_def, import_scope): + def _init_from_proto(self, critical_section_def, import_scope): # pylint: disable=invalid-name raise NotImplementedError("Not yet implemented") # TODO(ebrevdo): Re-enable once CriticalSection is in core. # assert isinstance( @@ -133,19 +140,21 @@ class CriticalSection(object): # critical_section_def.critical_section_name, # import_scope=import_scope)) - def _init_from_args(self, name): + def _init_from_args(self, name, shared_name): # pylint: disable=invalid-name """Initialize the CriticalSection from constructor arguments.""" with ops.name_scope(name, "CriticalSection", []) as name: - with ops.control_dependencies(None): + with ops.init_scope(): # pylint: disable=protected-access - handle_name = ops._name_from_scope_name(name) container = ops.get_default_graph()._container # pylint: enable=protected-access + if shared_name is None: + shared_name = name 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(): + self._handle = gen_resource_variable_ops.mutex_v2( + shared_name=shared_name, container=container, name=name) + + if not context.executing_eagerly(): ops.add_to_collections(CRITICAL_SECTIONS, self) @property @@ -183,68 +192,98 @@ class CriticalSection(object): 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: + lock = gen_resource_variable_ops.mutex_lock(self._handle) + + with ops.control_dependencies([lock]): + c_known_ops = set() + c_captured_tensors = set() + + def add_op_internal(op): + c_known_ops.add(op) + for i in op.inputs: + if i.op not in c_known_ops: + c_captured_tensors.add(i) + + c = function.HelperContext(add_op_internal) + with c: + r = fn(*args, **kwargs) + + resource_inputs = set([ + x for x in + list(nest.flatten(args)) + nest.flatten(kwargs.values()) + + list(c_captured_tensors) + if tensor_util.is_tensor(x) and x.dtype == dtypes.resource]) + + if self._handle in resource_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 in which it would be running. " + "This is illegal and would cause deadlocks. " "CriticalSection: %s." % self._handle) - if context.in_graph_mode(): + if not context.executing_eagerly(): # 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: + sg_handle_name = ops.convert_to_tensor(sg.handle).name + self_handle_name = ops.convert_to_tensor(self._handle).name + if sg_handle_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]) + resource_intersection = resource_inputs.intersection(sg.resources) + if resource_intersection: + raise ValueError( + "This execution would access resources: %s. Either this " + "lock (CriticalSection: %s) or lock '%s' " + "(CriticalSection: %s) requested exclusive resource access " + "of this resource. Did you mean to call execute with keyword " + "argument exclusive_resource_access=False?" % + (list(resource_intersection), self._handle.name, + sg.op.name, sg.handle.name)) + + def identity(x): # pylint: disable=invalid-name + if isinstance(x, tensor_array_ops.TensorArray): + return x.identity() + elif isinstance(x, ops.Operation): + return control_flow_ops.group(x) + elif context.executing_eagerly() and x is None: + return None + else: + return array_ops.identity(x) + + r_flat = [identity(x) for x in nest.flatten(r)] + + with ops.control_dependencies(r_flat): + # The identity must run on the same machine as self._handle + with ops.colocate_with(self._handle): + # Do not use array_ops.identity as there are special + # optimizations within TensorFlow which seem to elide it + # even when optimizations are disabled(!). + ensure_lock_exists = gen_resource_variable_ops.consume_mutex_lock( + lock) + + # Make sure that if any element of r is accessed, all of + # them are executed together. + r = nest.pack_sequence_as( + r, control_flow_ops.tuple(nest.flatten(r))) + + with ops.control_dependencies([ensure_lock_exists]): + outputs = nest.map_structure(identity, r) + + if not context.executing_eagerly(): signature = _ExecutionSignature( - op=op, + op=lock.op, + handle=self._handle, + resources=list(resource_inputs), 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)) + return outputs # TODO(ebrevdo): Re-enable once CriticalSection is in core. @@ -276,6 +315,7 @@ class CriticalSection(object): # def _execution_to_proto_fn(execution_signature, export_scope=None): # """Converts `_ExecutionSignature` to a `CriticalSectionExecutionDef`. +# # TODO(ebrevdo): Update for _ExecutionSignature storing resource list. # Args: # execution_signature: Instance of `_ExecutionSignature`. @@ -298,6 +338,7 @@ class CriticalSection(object): # def _execution_from_proto_fn(op_def, import_scope=None): # """Converts a `CriticalSectionExecutionDef` to a `_ExecutionSignature`.""" +# # TODO(ebrevdo): Update for _ExecutionSignature storing resource list. # assert isinstance( # op_def, critical_section_pb2.CriticalSectionExecutionDef) diff --git a/tensorflow/contrib/framework/python/ops/critical_section_test.py b/tensorflow/contrib/framework/python/ops/critical_section_test.py index a416724d3ba1719471d70667e140f9cd2daf86c7..c916592ce1979fe3a79cf28ad4bdac44284cce97 100644 --- a/tensorflow/contrib/framework/python/ops/critical_section_test.py +++ b/tensorflow/contrib/framework/python/ops/critical_section_test.py @@ -19,12 +19,10 @@ 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 control_flow_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.platform import test # TODO(ebrevdo): Re-enable once CriticalSection is in core. @@ -35,7 +33,7 @@ class CriticalSectionTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testCreateCriticalSection(self): - cs = critical_section_ops.CriticalSection(name="cs") + cs = critical_section_ops.CriticalSection(shared_name="cs") v = resource_variable_ops.ResourceVariable(0.0, name="v") def fn(a, b): @@ -45,16 +43,72 @@ class CriticalSectionTest(test.TestCase): with ops.control_dependencies([nv]): return array_ops.identity(c) - num_concurrent = 1000 + num_concurrent = 100 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 testCriticalSectionWithControlFlow(self): + for outer_cond in [False, True]: + for inner_cond in [False, True]: + cs = critical_section_ops.CriticalSection(shared_name="cs") + v = resource_variable_ops.ResourceVariable(0.0, name="v") + num_concurrent = 100 + + # pylint: disable=cell-var-from-loop + def fn(a, b): + c = v.read_value() + def true_fn(): + with ops.control_dependencies([c]): + nv = v.assign_add(a * b) + with ops.control_dependencies([nv]): + return array_ops.identity(c) + return control_flow_ops.cond( + array_ops.identity(inner_cond), true_fn, lambda: c) + + def execute(): + return cs.execute(fn, 1.0, 2.0) + + r = [ + control_flow_ops.cond(array_ops.identity(outer_cond), + execute, + v.read_value) + for _ in range(num_concurrent) + ] + # pylint: enable=cell-var-from-loop + + self.evaluate(v.initializer) + r_value = self.evaluate(r) + if inner_cond and outer_cond: + self.assertAllClose([2.0 * i for i in range(num_concurrent)], + sorted(r_value)) + else: + self.assertAllClose([0] * num_concurrent, r_value) + + def testCriticalSectionInParallelDoesntDeadlockOnError(self): + # No eager mode execution of this test because eager does not + # run fn() in parallel, which is where the deadlock could + # potentially occur (in graph mode). + cs = critical_section_ops.CriticalSection(shared_name="cs") + v = resource_variable_ops.ResourceVariable(0.0, name="v") + + def fn(i): + error = control_flow_ops.Assert((i % 2) == 1, ["Error"]) + with ops.control_dependencies([error]): + return v.read_value() + num_concurrent = 2 + r = [cs.execute(fn, i) for i in range(num_concurrent)] + self.evaluate(v.initializer) + for _ in range(100): + with self.assertRaisesOpError("Error"): + self.evaluate(r) + @test_util.run_in_graph_and_eager_modes() def testCreateCriticalSectionFnReturnsOp(self): - cs = critical_section_ops.CriticalSection(name="cs") + cs = critical_section_ops.CriticalSection(shared_name="cs") v = resource_variable_ops.ResourceVariable(0.0, name="v") def fn_return_op(a, b): @@ -62,7 +116,7 @@ class CriticalSectionTest(test.TestCase): with ops.control_dependencies([c]): nv = v.assign_add(a * b) with ops.control_dependencies([nv]): - return () + return control_flow_ops.no_op() num_concurrent = 100 r = [cs.execute(fn_return_op, 1.0, 2.0) for _ in range(num_concurrent)] @@ -71,47 +125,25 @@ class CriticalSectionTest(test.TestCase): 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") + cs = critical_section_ops.CriticalSection(shared_name="cs") self.assertIn( cs, ops.get_collection(critical_section_ops.CRITICAL_SECTIONS)) - execute_op = cs.execute(lambda x: x + 1, 1.0).op + execute = cs.execute(lambda x: x + 1, 1.0, name="my_execute") + execute_op = [ + x for x in execute.graph.get_operations() + if "my_execute" in x.name and "MutexLock" in x.type + ][0] 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") + # This does not work properly in eager mode. Eager users will + # just hit a deadlock if they do this. But at least it'll be easier + # to debug. + cs = critical_section_ops.CriticalSection(shared_name="cs") def fn(x): return cs.execute(lambda x: x+1, x) with self.assertRaisesRegexp( @@ -167,7 +199,7 @@ class CriticalSectionTest(test.TestCase): # self.assertEqual(restored_exec[0].op.name, "imported/%s" % r.op.name) # def testToProto(self): - # cs = critical_section_ops.CriticalSection(name="cs") + # cs = critical_section_ops.CriticalSection(shared_name="cs") # proto = cs.to_proto() # self.assertEqual(proto.critical_section_name, cs._handle.name) # cs_copy = critical_section_ops.CriticalSection.from_proto(proto) 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/gan/BUILD b/tensorflow/contrib/gan/BUILD index 5db34f0f8db93620b8b4a6b71f63b66ac718ee30..ff6f3b744190c9a7c74fb88878e5f13412251e79 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", @@ -353,6 +354,7 @@ py_test( name = "classifier_metrics_test", srcs = ["python/eval/python/classifier_metrics_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":classifier_metrics", "//tensorflow/core:protos_all_py", 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 986a5ff6dcbeb2ff996f49137adc6d34e14c979f..b17975811247aebc759bc4b16e5ad703df4bb1b9 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 @@ -43,11 +44,11 @@ from tensorflow.python.ops import functional_ops from tensorflow.python.ops import image_ops from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn_impl from tensorflow.python.ops import nn_ops from tensorflow.python.platform import gfile from tensorflow.python.platform import resource_loader - __all__ = [ 'get_graph_def_from_disk', 'get_graph_def_from_resource', @@ -61,10 +62,11 @@ __all__ = [ 'frechet_inception_distance', 'frechet_classifier_distance', 'frechet_classifier_distance_from_activations', + 'mean_only_frechet_classifier_distance_from_activations', + 'diagonal_only_frechet_classifier_distance_from_activations', 'INCEPTION_DEFAULT_IMAGE_SIZE', ] - INCEPTION_URL = 'http://download.tensorflow.org/models/frozen_inception_v1_2015_12_05.tar.gz' INCEPTION_FROZEN_GRAPH = 'inceptionv1_for_inception_score.pb' INCEPTION_INPUT = 'Mul:0' @@ -76,8 +78,7 @@ INCEPTION_DEFAULT_IMAGE_SIZE = 299 def _validate_images(images, image_size): images = ops.convert_to_tensor(images) images.shape.with_rank(4) - images.shape.assert_is_compatible_with( - [None, image_size, image_size, None]) + images.shape.assert_is_compatible_with([None, image_size, image_size, None]) return images @@ -108,9 +109,10 @@ def _symmetric_matrix_square_root(mat, eps=1e-10): math_ops.matmul(u, array_ops.diag(si)), v, transpose_b=True) -def preprocess_image( - images, height=INCEPTION_DEFAULT_IMAGE_SIZE, - width=INCEPTION_DEFAULT_IMAGE_SIZE, scope=None): +def preprocess_image(images, + height=INCEPTION_DEFAULT_IMAGE_SIZE, + width=INCEPTION_DEFAULT_IMAGE_SIZE, + scope=None): """Prepare a batch of images for evaluation. This is the preprocessing portion of the graph from @@ -189,20 +191,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, @@ -257,8 +273,11 @@ def run_inception(images, return activations -def run_image_classifier(tensor, graph_def, input_tensor, - output_tensor, scope='RunClassifier'): +def run_image_classifier(tensor, + graph_def, + input_tensor, + output_tensor, + scope='RunClassifier'): """Runs a network from a frozen graph. Args: @@ -302,7 +321,7 @@ def classifier_score(images, classifier_fn, num_batches=1): 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 + large batches, use classifier_score_from_logits(), which this method also uses. Args: @@ -418,8 +437,8 @@ def trace_sqrt_product(sigma, sigma_v): sqrt_sigma = _symmetric_matrix_square_root(sigma) # This is sqrt(A sigma_v A) above - sqrt_a_sigmav_a = math_ops.matmul( - sqrt_sigma, math_ops.matmul(sigma_v, sqrt_sigma)) + sqrt_a_sigmav_a = math_ops.matmul(sqrt_sigma, + math_ops.matmul(sigma_v, sqrt_sigma)) return math_ops.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a)) @@ -435,9 +454,9 @@ def frechet_classifier_distance(real_images, 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 - C and C_w, this function calcuates + C and C_w, this function calculates - |m - m_w|^2 + Tr(C + C_w - 2(C * C_w)^(1/2)) + |m - m_w|^2 + Tr(C + C_w - 2(C * C_w)^(1/2)) which captures how different the distributions of real images and generated images (or more accurately, their visual features) are. Note that unlike the @@ -448,7 +467,7 @@ def frechet_classifier_distance(real_images, 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 + sample size to compute Frechet classifier distance when comparing two generative models. NOTE: This function consumes images, computes their activations, and then @@ -496,10 +515,140 @@ def frechet_classifier_distance(real_images, return frechet_classifier_distance_from_activations(real_a, gen_a) -def frechet_classifier_distance_from_activations( +def mean_only_frechet_classifier_distance_from_activations( real_activations, generated_activations): """Classifier distance for evaluating a generative model from activations. + Given two Gaussian distribution with means m and m_w and covariance matrices + C and C_w, this function calcuates + + |m - m_w|^2 + + which captures how different the distributions of real images and generated + images (or more accurately, their visual features) are. Note that unlike the + 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. + + In this variant, we only compute the difference between the means of the + fitted Gaussians. The computation leads to O(n) vs. O(n^2) memory usage, yet + still retains much of the same information as FID. + + Args: + real_activations: 2D array of activations of real images of size + [num_images, num_dims] to use to compute Frechet Inception distance. + generated_activations: 2D array of activations of generated images of size + [num_images, num_dims] to use to compute Frechet Inception distance. + + Returns: + The mean-only 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) + + activations_dtype = real_activations.dtype + if activations_dtype != dtypes.float64: + real_activations = math_ops.to_double(real_activations) + generated_activations = math_ops.to_double(generated_activations) + + # Compute means of activations. + m = math_ops.reduce_mean(real_activations, 0) + m_w = math_ops.reduce_mean(generated_activations, 0) + + # Next the distance between means. + mean = math_ops.square(linalg_ops.norm(m - m_w)) # This uses the L2 norm. + mofid = mean + if activations_dtype != dtypes.float64: + mofid = math_ops.cast(mofid, activations_dtype) + + return mofid + + +def diagonal_only_frechet_classifier_distance_from_activations( + real_activations, generated_activations): + """Classifier distance for evaluating a generative model. + + This is based on the Frechet Inception distance, but for an arbitrary + classifier. + + 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 + C and C_w, this function calcuates + + |m - m_w|^2 + (sigma + sigma_w - 2(sigma x sigma_w)^(1/2)) + + which captures how different the distributions of real images and generated + images (or more accurately, their visual features) are. Note that unlike the + Inception score, this is a true distance and utilizes information about real + world images. In this variant, we compute diagonal-only covariance matrices. + As a result, instead of computing an expensive matrix square root, we can do + something much simpler, and has O(n) vs O(n^2) space complexity. + + 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. + + Returns: + The diagonal-only Frechet Inception distance. A floating-point scalar of + the same type as the output of the activations. + + Raises: + ValueError: If the shape of the variance and mean vectors are not equal. + """ + real_activations.shape.assert_has_rank(2) + generated_activations.shape.assert_has_rank(2) + + activations_dtype = real_activations.dtype + if activations_dtype != dtypes.float64: + real_activations = math_ops.to_double(real_activations) + generated_activations = math_ops.to_double(generated_activations) + + # Compute mean and covariance matrices of activations. + m, var = nn_impl.moments(real_activations, axes=[0]) + m_w, var_w = nn_impl.moments(generated_activations, axes=[0]) + + actual_shape = var.get_shape() + expected_shape = m.get_shape() + + if actual_shape != expected_shape: + raise ValueError('shape: {} must match expected shape: {}'.format( + actual_shape, expected_shape)) + + # Compute the two components of FID. + + # First the covariance component. + # Here, note that trace(A + B) = trace(A) + trace(B) + trace = math_ops.reduce_sum( + (var + var_w) - 2.0 * math_ops.sqrt(math_ops.multiply(var, var_w))) + + # Next the distance between means. + mean = math_ops.square(linalg_ops.norm(m - m_w)) # This uses the L2 norm. + dofid = trace + mean + if activations_dtype != dtypes.float64: + dofid = math_ops.cast(dofid, activations_dtype) + + return dofid + + +def frechet_classifier_distance_from_activations(real_activations, + generated_activations): + """Classifier distance for evaluating a generative model. + 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 @@ -508,15 +657,22 @@ def frechet_classifier_distance_from_activations( 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 - C and C_w, this function calcuates + C and C_w, this function calculates - |m - m_w|^2 + Tr(C + C_w - 2(C * C_w)^(1/2)) + |m - m_w|^2 + Tr(C + C_w - 2(C * C_w)^(1/2)) which captures how different the distributions of real images and generated images (or more accurately, their visual features) are. Note that unlike the 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: 2D Tensor containing activations of real data. Shape is [batch_size, activation_size]. @@ -538,36 +694,37 @@ def frechet_classifier_distance_from_activations( # Compute mean and covariance matrices of activations. m = math_ops.reduce_mean(real_activations, 0) - m_v = math_ops.reduce_mean(generated_activations, 0) + m_w = math_ops.reduce_mean(generated_activations, 0) num_examples = math_ops.to_double(array_ops.shape(real_activations)[0]) # sigma = (1 / (n - 1)) * (X - mu) (X - mu)^T real_centered = real_activations - m sigma = math_ops.matmul( - real_centered, real_centered, transpose_a=True) / (num_examples - 1) + real_centered, real_centered, transpose_a=True) / ( + num_examples - 1) - gen_centered = generated_activations - m_v - sigma_v = math_ops.matmul( - gen_centered, gen_centered, transpose_a=True) / (num_examples - 1) + gen_centered = generated_activations - m_w + sigma_w = math_ops.matmul( + gen_centered, gen_centered, transpose_a=True) / ( + num_examples - 1) - # Find the Tr(sqrt(sigma sigma_v)) component of FID - sqrt_trace_component = trace_sqrt_product(sigma, sigma_v) + # Find the Tr(sqrt(sigma sigma_w)) component of FID + sqrt_trace_component = trace_sqrt_product(sigma, sigma_w) # Compute the two components of FID. # First the covariance component. # Here, note that trace(A + B) = trace(A) + trace(B) - trace = math_ops.trace(sigma + sigma_v) - 2.0 * sqrt_trace_component + trace = math_ops.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component # Next the distance between means. - mean = math_ops.square(linalg_ops.norm(m - m_v)) # This uses the L2 norm. + mean = math_ops.square(linalg_ops.norm(m - m_w)) # This uses the L2 norm. fid = trace + mean if activations_dtype != dtypes.float64: fid = math_ops.cast(fid, activations_dtype) return fid - frechet_inception_distance = functools.partial( frechet_classifier_distance, classifier_fn=functools.partial( 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..663e49bdca3cb2dd9257da326488c877fcc4256d 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py @@ -50,6 +50,26 @@ def _expected_inception_score(logits): return np.exp(np.mean(per_example_logincscore)) +def _expected_mean_only_fid(real_imgs, gen_imgs): + m = np.mean(real_imgs, axis=0) + m_v = np.mean(gen_imgs, axis=0) + mean = np.square(m - m_v).sum() + mofid = mean + return mofid + + +def _expected_diagonal_only_fid(real_imgs, gen_imgs): + m = np.mean(real_imgs, axis=0) + m_v = np.mean(gen_imgs, axis=0) + var = np.var(real_imgs, axis=0) + var_v = np.var(gen_imgs, axis=0) + sqcc = np.sqrt(var * var_v) + mean = (np.square(m - m_v)).sum() + trace = (var + var_v - 2 * sqcc).sum() + dofid = mean + trace + return dofid + + def _expected_fid(real_imgs, gen_imgs): m = np.mean(real_imgs, axis=0) m_v = np.mean(gen_imgs, axis=0) @@ -181,7 +201,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 +216,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 +233,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 +274,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 +294,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) @@ -276,6 +305,46 @@ class ClassifierMetricsTest(test.TestCase): self.assertAllClose(_expected_inception_score(logits), incscore_np) + def test_mean_only_frechet_classifier_distance_value(self): + """Test that `frechet_classifier_distance` gives the correct value.""" + np.random.seed(0) + + pool_real_a = np.float32(np.random.randn(256, 2048)) + pool_gen_a = np.float32(np.random.randn(256, 2048)) + + tf_pool_real_a = array_ops.constant(pool_real_a) + tf_pool_gen_a = array_ops.constant(pool_gen_a) + + mofid_op = classifier_metrics.mean_only_frechet_classifier_distance_from_activations( # pylint: disable=line-too-long + tf_pool_real_a, tf_pool_gen_a) + + with self.test_session() as sess: + actual_mofid = sess.run(mofid_op) + + expected_mofid = _expected_mean_only_fid(pool_real_a, pool_gen_a) + + self.assertAllClose(expected_mofid, actual_mofid, 0.0001) + + def test_diagonal_only_frechet_classifier_distance_value(self): + """Test that `frechet_classifier_distance` gives the correct value.""" + np.random.seed(0) + + pool_real_a = np.float32(np.random.randn(256, 2048)) + pool_gen_a = np.float32(np.random.randn(256, 2048)) + + tf_pool_real_a = array_ops.constant(pool_real_a) + tf_pool_gen_a = array_ops.constant(pool_gen_a) + + dofid_op = classifier_metrics.diagonal_only_frechet_classifier_distance_from_activations( # pylint: disable=line-too-long + tf_pool_real_a, tf_pool_gen_a) + + with self.test_session() as sess: + actual_dofid = sess.run(dofid_op) + + expected_dofid = _expected_diagonal_only_fid(pool_real_a, pool_gen_a) + + self.assertAllClose(expected_dofid, actual_dofid, 0.0001) + def test_frechet_classifier_distance_value(self): """Test that `frechet_classifier_distance` gives the correct value.""" np.random.seed(0) @@ -285,9 +354,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 +367,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_impl.py b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_impl.py index 9bebcacbe46d85fc4226c4275b71b3ecbde57a97..4b10bc0f8e607c02763d8ea622d6f8f2572c586d 100644 --- a/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_impl.py @@ -212,7 +212,7 @@ def sliced_wasserstein_distance(real_images, Args: real_images: (tensor) Real images (batch, height, width, channels). fake_images: (tensor) Fake images (batch, height, width, channels). - resolution_min: (int) Minimum resolution for the Laplacion pyramid. + resolution_min: (int) Minimum resolution for the Laplacian pyramid. patches_per_image: (int) Number of patches to extract per image per Laplacian level. patch_size: (int) Width of a square patch. @@ -221,7 +221,7 @@ def sliced_wasserstein_distance(real_images, use_svd: experimental method to compute a more accurate distance. Returns: List of tuples (distance_real, distance_fake) for each level of the - Laplacian pyramid from the highest resoluion to the lowest. + Laplacian pyramid from the highest resolution to the lowest. distance_real is the Wasserstein distance between real images distance_fake is the Wasserstein distance between real and fake images. Raises: 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 74811ff4096eb5215148f0565bf094b83408014c..0d1afad72da8a8e087239868e25ddebe23490d1e 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py @@ -39,12 +39,13 @@ 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. @@ -83,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, diff --git a/tensorflow/contrib/gan/python/eval/python/summaries_test.py b/tensorflow/contrib/gan/python/eval/python/summaries_test.py index a02d8772e130a2a927735e56c4272aba1f1a6996..45eb108586bed07434ac29595164745eac6054c1 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_test.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_test.py @@ -72,8 +72,10 @@ def get_cyclegan_model(): class SummariesTest(test.TestCase): def _test_add_gan_model_image_summaries_impl(self, get_model_fn, - expected_num_summary_ops): - summaries.add_gan_model_image_summaries(get_model_fn(), grid_size=2) + expected_num_summary_ops, + model_summaries): + summaries.add_gan_model_image_summaries(get_model_fn(), grid_size=2, + model_summaries=model_summaries) self.assertEquals(expected_num_summary_ops, len(ops.get_collection(ops.GraphKeys.SUMMARIES))) @@ -82,10 +84,13 @@ class SummariesTest(test.TestCase): summary.merge_all().eval() def test_add_gan_model_image_summaries(self): - self._test_add_gan_model_image_summaries_impl(get_gan_model, 5) + 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) def test_add_gan_model_image_summaries_for_cyclegan(self): - self._test_add_gan_model_image_summaries_impl(get_cyclegan_model, 10) + 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): diff --git a/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py b/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py index cd31c62667fc048b1003d334377405b284f32af5..e2594faf85bcf91cbe09f266e4d4211d20bdee17 100644 --- a/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py +++ b/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Miscellanous utilities for TFGAN code and examples. +"""Miscellaneous utilities for TFGAN code and examples. Includes: 1) Conditioning the value of a Tensor, based on techniques from 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 4cfae0de4451880cf8229903b0eb74b1c6e2e04d..9e4ec59e7098443efc53506a4ba159e84b5c1618 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 @@ -17,7 +17,7 @@ We use this to keep a history of values created by a generator, such that a discriminator can randomly be trained on some older samples, not just the current one. This can help to not let the discriminator get too far ahead of the -generator and also to keep the system from oscilating, if the discriminator +generator and also to keep the system from oscillating, if the discriminator forgets too fast what past samples from the generator looked like. See the following papers for more details. @@ -97,7 +97,7 @@ def tensor_pool(input_values, dtypes=[v.dtype for v in input_values], shapes=None) - # In pseudeo code this code does the following: + # In pseudo code this code does the following: # if not pool_full: # enqueue(input_values) # return input_values diff --git a/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_test.py b/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_test.py index 845f89827b6e60eda41a55a80671f43460247b05..2fe06a287284ff994326d5a977a2e4d4634268ae 100644 --- a/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_test.py +++ b/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_test.py @@ -148,7 +148,7 @@ class VirtualBatchnormTest(test.TestCase): self.assertAllClose(bn_np[i, ...], vb_np) def test_minibatch_independent(self): - """Test that virtual batch normalized exampels are independent. + """Test that virtual batch normalized examples are independent. Unlike batch normalization, virtual batch normalization has the property that the virtual batch normalized value of an example is independent of the diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl.py b/tensorflow/contrib/gan/python/losses/python/losses_impl.py index 23a3b60cc0055917bfc5243b0ebdbaea7b61edb9..39588b7219ebac1cc4855532be3fcc38e6381134 100644 --- a/tensorflow/contrib/gan/python/losses/python/losses_impl.py +++ b/tensorflow/contrib/gan/python/losses/python/losses_impl.py @@ -305,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, @@ -324,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 @@ -374,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) 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 56ac45554da3633149a61155a416fa7cb6cff553..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)) diff --git a/tensorflow/contrib/gan/python/train.py b/tensorflow/contrib/gan/python/train.py index 5d0ac93aec7869bb1d9b8a174ba50d4bec2c2826..776eb11ecb1624544d24611d8fe6ca19768b8313 100644 --- a/tensorflow/contrib/gan/python/train.py +++ b/tensorflow/contrib/gan/python/train.py @@ -460,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, @@ -481,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 @@ -539,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( 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/graph_editor/reroute.py b/tensorflow/contrib/graph_editor/reroute.py index 7ffdbb7139281734917fdb715601b317eb58b82f..95c02a64d47c26e731ef2628fb551529e9bc3f4d 100644 --- a/tensorflow/contrib/graph_editor/reroute.py +++ b/tensorflow/contrib/graph_editor/reroute.py @@ -471,9 +471,10 @@ def remove_control_inputs(op, cops): if cop not in op.control_inputs: raise ValueError("{} is not a control_input of {}".format(op.name, cop.name)) + control_inputs = [cop for cop in op.control_inputs if cop not in cops] # pylint: disable=protected-access - op._control_inputs = [cop for cop in op._control_inputs if cop not in cops] - op._recompute_node_def() + op._remove_all_control_inputs() + op._add_control_inputs(control_inputs) # pylint: enable=protected-access @@ -496,9 +497,6 @@ def add_control_inputs(op, cops): if cop in op.control_inputs: raise ValueError("{} is already a control_input of {}".format(cop.name, op.name)) - # pylint: disable=protected-access - op._control_inputs += cops - op._recompute_node_def() - # pylint: enable=protected-access + op._add_control_inputs(cops) # pylint: disable=protected-access remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py b/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py index 252788140f8c1906718c150574b963385b6ecfa1..bcd2a34c4e791a2ab66a439109145d6b78c14e22 100644 --- a/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py +++ b/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py @@ -110,7 +110,7 @@ class GridRNNCell(rnn.RNNCell): logging.warning('%s: Using a concatenated state is slower and will ' 'soon be deprecated. Use state_is_tuple=True.', self) if not output_is_tuple: - logging.warning('%s: Using a concatenated output is slower and will' + logging.warning('%s: Using a concatenated output is slower and will ' 'soon be deprecated. Use output_is_tuple=True.', self) if num_dims < 1: diff --git a/tensorflow/contrib/hvx/README.md b/tensorflow/contrib/hvx/README.md index cb3a1087de5b83e7737e0dac2c04fbe152a9d676..163993a3f6bb1bedcdffb32944a98c7cc846878e 100644 --- a/tensorflow/contrib/hvx/README.md +++ b/tensorflow/contrib/hvx/README.md @@ -141,16 +141,16 @@ Configuring the installer for this system's environment... Launching installer... -An internal LaunchAnywhere application error has occured and this application cannot proceed. (LAX) +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) + 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) ``` It can be solved by temporarily assigning the `PS1` environment variable to something simple, such as '$'. 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/segmentation_ops.cc b/tensorflow/contrib/image/kernels/segmentation_ops.cc index fe8bf6e21c7b7310527668324571774e8bc50893..93722896233f0278c6cbb44af7203345e58c3172 100644 --- a/tensorflow/contrib/image/kernels/segmentation_ops.cc +++ b/tensorflow/contrib/image/kernels/segmentation_ops.cc @@ -101,8 +101,8 @@ struct ImageConnectedComponentsFunctor { 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) { + [&union_find, 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 = 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 index bf0c97245fc5c70469350ec66023f4d1474930e2..3f4029e558d92a2b6539456bf9cf49ec2d21c9f3 100644 --- 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 @@ -18,13 +18,8 @@ from __future__ import absolute_import 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.image.python.ops.single_image_random_dot_stereograms \ import single_image_random_dot_stereograms -from tensorflow.python.client import session from tensorflow.python.framework import constant_op from tensorflow.python.framework import test_util from tensorflow.python.platform import googletest diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py index 63377ae50310db51a3111c5a6e00df7d75dccc0b..c139ae89d8d682d6b87813c3a21703ffa762f28e 100644 --- a/tensorflow/contrib/image/python/ops/image_ops.py +++ b/tensorflow/contrib/image/python/ops/image_ops.py @@ -40,7 +40,7 @@ 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) @@ -290,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") @@ -346,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: diff --git a/tensorflow/contrib/kafka/BUILD b/tensorflow/contrib/kafka/BUILD index f7593aa462c4ca86d0ffc4f065e1aae849364561..1c3974871c62911c0cb47677eb92d28286837142 100644 --- a/tensorflow/contrib/kafka/BUILD +++ b/tensorflow/contrib/kafka/BUILD @@ -1,66 +1,93 @@ -package( - default_visibility = ["//visibility:private"], -) +package(default_visibility = ["//tensorflow:internal"]) 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") +load( + "//tensorflow:tensorflow.bzl", + "tf_gen_op_wrapper_py", + "tf_kernel_library", + "tf_custom_op_library", + "tf_custom_op_py_library", + "tf_gen_op_libs", + "tf_py_test", +) -tf_kernel_library( - name = "kafka_kernels", +py_library( + name = "kafka", + srcs = ["__init__.py"], + srcs_version = "PY2AND3", + deps = [ + ":dataset_ops", + ], +) + +tf_custom_op_library( + name = "_dataset_ops.so", + srcs = ["ops/dataset_ops.cc"], + deps = [":dataset_kernels"], +) + +tf_gen_op_libs( + op_lib_names = ["dataset_ops"], +) + +cc_library( + name = "dataset_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", + "//tensorflow/core:framework_headers_lib", "//third_party/eigen3", - "@kafka//:kafka", + "@kafka", + "@protobuf_archive//:protobuf_headers", ], + alwayslink = 1, ) -tf_gen_op_libs( - op_lib_names = ["kafka_ops"], +py_library( + name = "dataset_ops", + srcs = [ + "python/ops/kafka_dataset_ops.py", + ], + srcs_version = "PY2AND3", deps = [ - "//tensorflow/core:lib", + ":kafka_op_loader", + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:util", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", ], ) 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"], + name = "gen_dataset_ops", + out = "python/ops/gen_dataset_ops.py", + deps = ["//tensorflow/contrib/kafka:dataset_ops_op_lib"], ) -py_library( - name = "kafka", - srcs = [ - "__init__.py", - "python/ops/kafka_dataset_ops.py", +tf_kernel_library( + name = "dataset_ops_kernels", + deps = [ + ":dataset_kernels", + "//tensorflow/core:framework", + ], + alwayslink = 1, +) + +tf_custom_op_py_library( + name = "kafka_op_loader", + srcs = ["python/ops/kafka_op_loader.py"], + dso = ["//tensorflow/contrib/kafka:_dataset_ops.so"], + kernels = [ + ":dataset_ops_kernels", + "//tensorflow/contrib/kafka:dataset_ops_op_lib", ], srcs_version = "PY2AND3", - visibility = ["//visibility:public"], deps = [ - ":gen_kafka_ops", + ":gen_dataset_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", ], ) @@ -88,13 +115,17 @@ tf_py_test( ], tags = [ "manual", + "no_windows", + "notap", ], ) filegroup( name = "all_files", srcs = glob( - ["**/*"], + include = [ + "**/*", + ], exclude = [ "**/METADATA", "**/OWNERS", diff --git a/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc b/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc index 88ef5f357113372b0a2d0cb13382ac980a61252d..a4cd4a2cc4b99b5906185bd2b942ed15c1ddf5e4 100644 --- a/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc +++ b/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc @@ -13,9 +13,7 @@ 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 "tensorflow/core/framework/dataset.h" #include "src-cpp/rdkafkacpp.h" diff --git a/tensorflow/contrib/kafka/ops/kafka_ops.cc b/tensorflow/contrib/kafka/ops/dataset_ops.cc similarity index 100% rename from tensorflow/contrib/kafka/ops/kafka_ops.cc rename to tensorflow/contrib/kafka/ops/dataset_ops.cc diff --git a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py index 94cf6b5ace6a10b7c0471d7d25b5bce789ac322b..621911876fc502ece76b08eb6c28697b3c12c863 100644 --- a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py +++ b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py @@ -18,21 +18,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np -import os - from tensorflow.contrib.kafka.python.ops import kafka_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.framework import tensor_shape -from tensorflow.python.lib.io import python_io from tensorflow.python.ops import array_ops -from tensorflow.python.ops import io_ops from tensorflow.python.platform import test -from tensorflow.python.util import compat + class KafkaDatasetTest(test.TestCase): @@ -64,52 +56,58 @@ class KafkaDatasetTest(test.TestCase): 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}) + 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)) + 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}) + 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)) + 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}) + 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)) + 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}) + 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)) + 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}) + 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)], + 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)], + self.assertAllEqual(["D" + str(i + 5) for i in range(5)], sess.run(get_next)) diff --git a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh index 7997c12731189e56fc491a0f4de8b19c39d262b2..adf027b8e714124cde2b4618546e20c6b7162e1f 100644 --- a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh +++ b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh @@ -1,4 +1,18 @@ #!/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 diff --git a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py index 6590d86ebb7d9da836e5777af7d517919f4e2eff..a1624614d1ab1be31463c5cdc0b4cfb653165a0c 100644 --- a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py +++ b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py @@ -17,20 +17,24 @@ 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.contrib.util import loader -from tensorflow.python.data.ops.readers import Dataset -from tensorflow.python.framework import common_shapes +from tensorflow.contrib.kafka.python.ops import kafka_op_loader # pylint: disable=unused-import +from tensorflow.contrib.kafka.python.ops import gen_dataset_ops +from tensorflow.python.data.ops.dataset_ops import Dataset from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape -from tensorflow.python.platform import resource_loader + class KafkaDataset(Dataset): """A Kafka Dataset that consumes the message. """ - def __init__(self, topics, servers="localhost", group="", eof=False, timeout=1000): + def __init__(self, + topics, + servers="localhost", + group="", + eof=False, + timeout=1000): """Create a KafkaReader. Args: @@ -50,14 +54,13 @@ class KafkaDataset(Dataset): 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._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) + return gen_dataset_ops.kafka_dataset(self._topics, self._servers, + self._group, self._eof, self._timeout) @property def output_classes(self): diff --git a/tensorflow/contrib/kafka/python/ops/kafka_op_loader.py b/tensorflow/contrib/kafka/python/ops/kafka_op_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..ec2fdea962ef946d3f8f32b9e630b92649d612fe --- /dev/null +++ b/tensorflow/contrib/kafka/python/ops/kafka_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 kafka 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/kfac/python/kernel_tests/BUILD b/tensorflow/contrib/kfac/python/kernel_tests/BUILD index f4ed978174a9ddd8b54a88e60bfb48a67a2e76d2..d1c449402a697dd5f8876c82a6682dde2d18b4df 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/BUILD +++ b/tensorflow/contrib/kfac/python/kernel_tests/BUILD @@ -36,6 +36,7 @@ py_test( srcs = ["fisher_factors_test.py"], srcs_version = "PY2AND3", deps = [ + "//tensorflow/contrib/kfac/python/ops:fisher_blocks", "//tensorflow/contrib/kfac/python/ops:fisher_factors", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", @@ -113,6 +114,7 @@ py_test( name = "utils_test", srcs = ["utils_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ "//tensorflow/contrib/kfac/python/ops:utils", "//tensorflow/contrib/tpu", diff --git a/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py b/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py index bfdb69ad02caaa57827e0ae6b3c9fc0d0ed03754..30c5404e03910eedb48132b0d69b2eabb89a9149 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py @@ -90,40 +90,63 @@ class EstimatorTest(test.TestCase): def testEstimatorInitManualRegistration(self): with self._graph.as_default(): # We should be able to build an estimator for only the registered vars. - estimator.FisherEstimator([self.weights], 0.1, 0.2, self.layer_collection) + estimator.FisherEstimator([self.weights], 0.1, 0.2, + self.layer_collection) # Check that we throw an error if we try to build an estimator for vars # that were not manually registered. with self.assertRaises(ValueError): - estimator.FisherEstimator([self.weights, self.bias], 0.1, 0.2, - self.layer_collection) + est = estimator.FisherEstimator([self.weights, self.bias], 0.1, 0.2, + self.layer_collection) + est.make_ops_and_vars() # Check that we throw an error if we don't include registered variables, # i.e. self.weights with self.assertRaises(ValueError): - estimator.FisherEstimator([], 0.1, 0.2, self.layer_collection) + est = estimator.FisherEstimator([], 0.1, 0.2, self.layer_collection) + est.make_ops_and_vars() @test.mock.patch.object(utils.SubGraph, "variable_uses", return_value=42) def testVariableWrongNumberOfUses(self, mock_uses): with self.assertRaises(ValueError): - estimator.FisherEstimator([self.weights], 0.1, 0.2, self.layer_collection) + est = estimator.FisherEstimator([self.weights], 0.1, 0.2, + self.layer_collection) + est.make_ops_and_vars() def testInvalidEstimationMode(self): with self.assertRaises(ValueError): - estimator.FisherEstimator([self.weights], 0.1, 0.2, self.layer_collection, - "not_a_real_mode") + est = estimator.FisherEstimator([self.weights], 0.1, 0.2, + self.layer_collection, + estimation_mode="not_a_real_mode") + est.make_ops_and_vars() - def testModeListCorrect(self): + def testGradientsModeBuild(self): with self._graph.as_default(): est = estimator.FisherEstimator([self.weights], 0.1, 0.2, - self.layer_collection) - self.assertItemsEqual(_ALL_ESTIMATION_MODES, est._gradient_fns.keys()) + self.layer_collection, + estimation_mode="gradients") + est.make_ops_and_vars() + + def testEmpiricalModeBuild(self): + with self._graph.as_default(): + est = estimator.FisherEstimator([self.weights], 0.1, 0.2, + self.layer_collection, + estimation_mode="empirical") + est.make_ops_and_vars() + + def testCurvaturePropModeBuild(self): + with self._graph.as_default(): + est = estimator.FisherEstimator([self.weights], 0.1, 0.2, + self.layer_collection, + estimation_mode="curvature_prop") + est.make_ops_and_vars() - def testAllModesBuild(self): - for mode in _ALL_ESTIMATION_MODES: - with self._graph.as_default(): - estimator.FisherEstimator([self.weights], 0.1, 0.2, - self.layer_collection, mode) + def testExactModeBuild(self): + with self._graph.as_default(): + est = estimator.FisherEstimator([self.weights], 0.1, 0.2, + self.layer_collection, + estimation_mode="exact") + est.make_ops_and_vars() def test_cov_update_thunks(self): """Ensures covariance update ops run once per global_step.""" @@ -131,16 +154,19 @@ class EstimatorTest(test.TestCase): fisher_estimator = estimator.FisherEstimator( variables=[self.weights], layer_collection=self.layer_collection, - cov_ema_decay=0.0, - damping=0.0) + damping=0.2, + cov_ema_decay=0.0) # Construct an op that executes one covariance update per step. global_step = training_util.get_or_create_global_step() + (cov_variable_thunks, cov_update_op_thunks, + _, _) = fisher_estimator.create_ops_and_vars_thunks() + for thunk in cov_variable_thunks: + thunk() 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)]) @@ -178,17 +204,22 @@ class EstimatorTest(test.TestCase): fisher_estimator = estimator.FisherEstimator( variables=[self.weights], layer_collection=self.layer_collection, - cov_ema_decay=0.0, - damping=0.0) + damping=0.2, + cov_ema_decay=0.0) # Construct op that updates one inverse per global step. global_step = training_util.get_or_create_global_step() + (cov_variable_thunks, _, inv_variable_thunks, + inv_update_op_thunks) = fisher_estimator.create_ops_and_vars_thunks() + for thunk in cov_variable_thunks: + thunk() + for thunk in inv_variable_thunks: + thunk() inv_matrices = [ matrix for fisher_factor in self.layer_collection.get_factors() - for matrix in fisher_factor._inverses_by_damping.values() + for matrix in fisher_factor._matpower_by_exp_and_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)]) 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 82accd57f0c37d140238f1884fce956654d14227..b70c700f0936c2d8a2eca6e0836a3ee4ffe4e46d 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 @@ -93,6 +94,9 @@ class FullFBTest(test.TestCase): block.register_additional_minibatch(32) grads = (params[0]**2, math_ops.sqrt(params[1])) block.instantiate_factors((grads,), 0.5) + block._factor.instantiate_cov_variables() + block.register_inverse() + block._factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -111,6 +115,9 @@ class FullFBTest(test.TestCase): block.register_additional_minibatch(32) grads = params**2 block.instantiate_factors((grads,), 0.5) + block._factor.instantiate_cov_variables() + block.register_inverse() + block._factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -130,6 +137,9 @@ class FullFBTest(test.TestCase): grads = (array_ops.constant([2., 3.]), array_ops.constant(4.)) damping = 0.5 block.instantiate_factors((grads,), damping) + block._factor.instantiate_cov_variables() + block.register_inverse() + block._factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(state_ops.assign(block._factor._cov, _make_psd(3))) @@ -184,6 +194,7 @@ class NaiveDiagonalFBTest(test.TestCase): block.register_additional_minibatch(32) grads = (params[0]**2, math_ops.sqrt(params[1])) block.instantiate_factors((grads,), 0.5) + block._factor.instantiate_cov_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -202,6 +213,7 @@ class NaiveDiagonalFBTest(test.TestCase): block.register_additional_minibatch(32) grads = params**2 block.instantiate_factors((grads,), 0.5) + block._factor.instantiate_cov_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -220,6 +232,7 @@ class NaiveDiagonalFBTest(test.TestCase): grads = (params[0]**2, math_ops.sqrt(params[1])) damping = 0.5 block.instantiate_factors((grads,), damping) + block._factor.instantiate_cov_variables() cov = array_ops.reshape(array_ops.constant([2., 3., 4.]), [-1, 1]) sess.run(state_ops.assign(block._factor._cov, cov)) @@ -236,10 +249,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 @@ -366,6 +379,7 @@ class FullyConnectedDiagonalFB(test.TestCase): block.register_additional_minibatch(i, o) block.instantiate_factors((output_grads,), damping=0.0) + block._factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) sess.run(block._factor.make_covariance_update_op(0.0)) @@ -375,6 +389,70 @@ 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) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() + + # 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): @@ -396,7 +474,7 @@ class FullyConnectedKFACBasicFBTest(test.TestCase): block.register_additional_minibatch(inputs, outputs) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) def testInstantiateFactorsNoBias(self): with ops.Graph().as_default(): @@ -407,7 +485,7 @@ class FullyConnectedKFACBasicFBTest(test.TestCase): block.register_additional_minibatch(inputs, outputs) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) def testMultiplyInverseTuple(self): with ops.Graph().as_default(), self.test_session() as sess: @@ -417,7 +495,13 @@ class FullyConnectedKFACBasicFBTest(test.TestCase): block = fb.FullyConnectedKFACBasicFB(lc.LayerCollection(), has_bias=False) block.register_additional_minibatch(inputs, outputs) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) + + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -443,7 +527,12 @@ class FullyConnectedKFACBasicFBTest(test.TestCase): block = fb.FullyConnectedKFACBasicFB(lc.LayerCollection(), has_bias=False) block.register_additional_minibatch(inputs, outputs) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -467,10 +556,17 @@ class FullyConnectedKFACBasicFBTest(test.TestCase): block.register_additional_minibatch(inputs, outputs) grads = outputs**2 damping = 0. # This test is only valid without damping. - block.instantiate_factors(([grads],), damping) + block.instantiate_factors(((grads,),), damping) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() sess.run(state_ops.assign(block._input_factor._cov, _make_psd(3))) sess.run(state_ops.assign(block._output_factor._cov, _make_psd(2))) + + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() + sess.run(block._input_factor.make_inverse_update_ops()) sess.run(block._output_factor.make_inverse_update_ops()) @@ -658,6 +754,7 @@ class ConvDiagonalFBTest(test.TestCase): block.register_additional_minibatch(i, o) block.instantiate_factors((output_grads,), damping=0.0) + block._factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) sess.run(block._factor.make_covariance_update_op(0.0)) @@ -667,6 +764,54 @@ class ConvDiagonalFBTest(test.TestCase): return multiply_result, multiply_inverse_result +class DepthwiseConvKFCBasicFBTest(test.TestCase): + + def testInstantiateFactors(self): + with ops.Graph().as_default(): + random_seed.set_random_seed(200) + params = random_ops.random_normal((3, 3, 8, 2)) + inputs = random_ops.random_normal((32, 5, 5, 8)) + outputs = random_ops.random_normal((32, 5, 5, 16)) + layer_collection = lc.LayerCollection() + block = fb.DepthwiseConvKFCBasicFB( + layer_collection, params=params, strides=[1, 1, 1, 1], padding='SAME') + block.register_additional_minibatch(inputs, outputs) + grads = outputs**2 + block.instantiate_factors(([grads],), 0.5) + + def testMultiplyInverse(self): + with ops.Graph().as_default(), self.test_session() as sess: + random_seed.set_random_seed(200) + params = random_ops.random_normal((3, 3, 8, 2)) + inputs = random_ops.random_normal((32, 5, 5, 8)) + outputs = random_ops.random_normal((32, 5, 5, 16)) + layer_collection = lc.LayerCollection() + block = fb.DepthwiseConvKFCBasicFB( + layer_collection, params=params, strides=[1, 1, 1, 1], padding='SAME') + block.register_additional_minibatch(inputs, outputs) + grads = outputs**2 + block.instantiate_factors(([grads],), 0.5) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() + + # Ensure inverse update op doesn't crash. + sess.run(tf_variables.global_variables_initializer()) + sess.run([ + factor.make_inverse_update_ops() + for factor in layer_collection.get_factors() + ]) + + # Ensure inverse-vector multiply doesn't crash. + output = block.multiply_inverse(params) + sess.run(output) + + # Ensure same shape. + self.assertAllEqual(output.shape, params.shape) + + class ConvKFCBasicFBTest(test.TestCase): def _testConvKFCBasicFBInitParams(self, params): @@ -678,16 +823,17 @@ class ConvKFCBasicFBTest(test.TestCase): params = array_ops.constant(params) inputs = random_ops.random_normal((2, 2, 2)) outputs = random_ops.random_normal((2, 2, 2)) - block = fb.ConvKFCBasicFB(lc.LayerCollection(), params, [1, 1, 1], 'SAME') + block = fb.ConvKFCBasicFB( + lc.LayerCollection(), params=params, padding='SAME') block.register_additional_minibatch(inputs, outputs) self.assertAllEqual([outputs], block.tensors_to_compute_grads()) def testConvKFCBasicFBInitParamsParamsTuple(self): - self._testConvKFCBasicFBInitParams([np.array([1., 2.]), np.array(3.)]) + self._testConvKFCBasicFBInitParams([np.ones([1, 2, 2]), np.ones([2])]) def testConvKFCBasicFBInitParamsParamsSingle(self): - self._testConvKFCBasicFBInitParams([np.array([1., 2.])]) + self._testConvKFCBasicFBInitParams([np.ones([1, 2, 2])]) def testMultiplyInverseTuple(self): with ops.Graph().as_default(), self.test_session() as sess: @@ -695,11 +841,16 @@ class ConvKFCBasicFBTest(test.TestCase): params = random_ops.random_normal((2, 2, 2, 2)) inputs = random_ops.random_normal((2, 2, 2, 2)) outputs = random_ops.random_normal((2, 2, 2, 2)) - block = fb.ConvKFCBasicFB(lc.LayerCollection(), params, (1, 1, 1, 1), - 'SAME') + block = fb.ConvKFCBasicFB( + lc.LayerCollection(), params=params, padding='SAME') block.register_additional_minibatch(inputs, outputs) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -721,12 +872,17 @@ class ConvKFCBasicFBTest(test.TestCase): params = random_ops.random_normal((2, 2, 2, 2)) inputs = random_ops.random_normal((2, 2, 2, 2)) outputs = random_ops.random_normal((2, 2, 2, 2)) - block = fb.ConvKFCBasicFB(lc.LayerCollection(), params, (1, 1, 1, 1), - 'SAME') + block = fb.ConvKFCBasicFB( + lc.LayerCollection(), params=params, padding='SAME') block.register_additional_minibatch(inputs, outputs) self.assertFalse(block._has_bias) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -744,12 +900,17 @@ class ConvKFCBasicFBTest(test.TestCase): params = [random_ops.random_normal((2, 2, 2, 2))] inputs = random_ops.random_normal((2, 2, 2, 2)) outputs = random_ops.random_normal((2, 2, 2, 2)) - block = fb.ConvKFCBasicFB(lc.LayerCollection(), params, (1, 1, 1, 1), - 'SAME') + block = fb.ConvKFCBasicFB( + lc.LayerCollection(), params=params, padding='SAME') block.register_additional_minibatch(inputs, outputs) self.assertTrue(block._has_bias) grads = outputs**2 - block.instantiate_factors(([grads],), 0.5) + block.instantiate_factors(((grads,),), 0.5) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() # Make sure our inverse is something other than the identity. sess.run(tf_variables.global_variables_initializer()) @@ -767,12 +928,17 @@ class ConvKFCBasicFBTest(test.TestCase): params = array_ops.zeros((2, 2, 2, 2)) inputs = array_ops.zeros((2, 2, 2, 2)) outputs = array_ops.zeros((2, 2, 2, 2)) - block = fb.ConvKFCBasicFB(lc.LayerCollection(), params, (1, 1, 1, 1), - 'SAME') + block = fb.ConvKFCBasicFB( + lc.LayerCollection(), params=params, padding='SAME') block.register_additional_minibatch(inputs, outputs) grads = outputs**2 damping = 0. # This test is only valid without damping. - block.instantiate_factors(([grads],), damping) + block.instantiate_factors(((grads,),), damping) + block._input_factor.instantiate_cov_variables() + block._output_factor.instantiate_cov_variables() + block.register_inverse() + block._input_factor.instantiate_inv_variables() + block._output_factor.instantiate_inv_variables() sess.run(state_ops.assign(block._input_factor._cov, _make_psd(8))) sess.run(state_ops.assign(block._output_factor._cov, _make_psd(2))) @@ -797,9 +963,9 @@ class FullyConnectedSeriesFBTest(test.TestCase): random_seed.set_random_seed(200) inputs = array_ops.constant([1., 2.]) outputs = array_ops.constant([3., 4.]) - block = fb.FullyConnectedSeriesFB( - lc.LayerCollection(), inputs=[inputs], outputs=[outputs]) - self.assertAllEqual([outputs], block.tensors_to_compute_grads()) + block = fb.FullyConnectedSeriesFB(lc.LayerCollection()) + block.register_additional_minibatch([inputs], [outputs]) + self.assertAllEqual([[outputs]], block.tensors_to_compute_grads()) def testInstantiateFactorsHasBias(self): with ops.Graph().as_default(): @@ -808,11 +974,10 @@ class FullyConnectedSeriesFBTest(test.TestCase): outputs = array_ops.constant([[3., 4.], [5., 6.]]) block = fb.FullyConnectedSeriesFB( lc.LayerCollection(), - inputs=[inputs], - outputs=[outputs], has_bias=True) + block.register_additional_minibatch([inputs], [outputs]) grads = outputs**2 - block.instantiate_factors(((grads,),), 0.5) + block.instantiate_factors((((grads,),),), 0.5) def testInstantiateFactorsNoBias(self): with ops.Graph().as_default(): @@ -821,11 +986,10 @@ class FullyConnectedSeriesFBTest(test.TestCase): outputs = array_ops.constant([[3., 4.], [5., 6.]]) block = fb.FullyConnectedSeriesFB( lc.LayerCollection(), - inputs=[inputs], - outputs=[outputs], has_bias=False) + block.register_additional_minibatch([inputs], [outputs]) grads = outputs**2 - block.instantiate_factors(((grads,),), 0.5) + block.instantiate_factors((((grads,),),), 0.5) def as_tensors(tensor_or_tuple): 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 753378d9f4a0d8762bafbee2ec27d6c71783dda1..16f02f1199ad8a404b0e6944fc89df32ce08609c 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import numpy as np import numpy.random as npr +from tensorflow.contrib.kfac.python.ops import fisher_blocks as fb from tensorflow.contrib.kfac.python.ops import fisher_factors as ff from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -29,36 +30,13 @@ 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 math_ops +from tensorflow.python.ops import random_ops from tensorflow.python.ops import variables as tf_variables 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): - 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): - 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()) +def make_damping_func(damping): + return fb._package_func(lambda: damping, damping) class FisherFactorTestingDummy(ff.FisherFactor): @@ -89,6 +67,24 @@ 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_matpower(self, x, exp, damping): + return NotImplementedError + + def right_multiply_matpower(self, x, exp, damping): + return NotImplementedError + + def instantiate_inv_variables(self): + return NotImplementedError + class InverseProvidingFactorTestingDummy(ff.InverseProvidingFactor): """Dummy class to test the non-abstract methods on ff.InverseProvidingFactor. @@ -231,21 +227,24 @@ class InverseProvidingFactorTest(test.TestCase): factor = InverseProvidingFactorTestingDummy(shape) factor_var_scope = 'dummy/a_b_c' - dampings = 0.1, 1e-1, 0.00001, 1e-5 + damping_funcs = [make_damping_func(0.1), + make_damping_func(0.1), + make_damping_func(1e-5), + make_damping_func(1e-5)] + for damping_func in damping_funcs: + factor.register_inverse(damping_func) - for damping in dampings: - factor.register_damped_inverse(damping) + factor.instantiate_inv_variables() - self.assertEqual(set(dampings), set(factor._inverses_by_damping.keys())) - inv = factor._inverses_by_damping[dampings[0]] - self.assertEqual(inv, factor._inverses_by_damping[dampings[1]]) - self.assertNotEqual(inv, factor._inverses_by_damping[dampings[2]]) - self.assertEqual(factor._inverses_by_damping[dampings[2]], - factor._inverses_by_damping[dampings[3]]) + inv = factor.get_inverse(damping_funcs[0]) + self.assertEqual(inv, factor.get_inverse(damping_funcs[1])) + self.assertNotEqual(inv, factor.get_inverse(damping_funcs[2])) + self.assertEqual(factor.get_inverse(damping_funcs[2]), + factor.get_inverse(damping_funcs[3])) factor_vars = tf_ops.get_collection(tf_ops.GraphKeys.GLOBAL_VARIABLES, factor_var_scope) - self.assertListEqual([inv, factor._inverses_by_damping[dampings[2]]], - factor_vars) + self.assertEqual(set([inv, factor.get_inverse(damping_funcs[2])]), + set(factor_vars)) self.assertEqual(shape, inv.get_shape()) def testRegisterMatpower(self): @@ -255,17 +254,22 @@ class InverseProvidingFactorTest(test.TestCase): factor = InverseProvidingFactorTestingDummy(shape) factor_var_scope = 'dummy/a_b_c' - factor.register_matpower(1, 0.5) - factor.register_matpower(2, 0.5) + # TODO(b/74201126): Change to using the same func for both once + # Topohash is in place. + damping_func_1 = make_damping_func(0.5) + damping_func_2 = make_damping_func(0.5) + + factor.register_matpower(-0.5, damping_func_1) + factor.register_matpower(2, damping_func_2) + + factor.instantiate_inv_variables() - self.assertEqual( - set([(1, 0.5), (2, 0.5)]), - set(factor._matpower_by_exp_and_damping.keys())) factor_vars = tf_ops.get_collection(tf_ops.GraphKeys.GLOBAL_VARIABLES, factor_var_scope) - matpower1 = factor.get_matpower(1, 0.5) - matpower2 = factor.get_matpower(2, 0.5) - self.assertListEqual([matpower1, matpower2], factor_vars) + matpower1 = factor.get_matpower(-0.5, damping_func_1) + matpower2 = factor.get_matpower(2, damping_func_2) + + self.assertEqual(set([matpower1, matpower2]), set(factor_vars)) self.assertEqual(shape, matpower1.get_shape()) self.assertEqual(shape, matpower2.get_shape()) @@ -284,17 +288,24 @@ class InverseProvidingFactorTest(test.TestCase): factor = InverseProvidingFactorTestingDummy(cov.shape) factor._cov = array_ops.constant(cov, dtype=dtypes.float32) + damping_funcs = [] for i in range(1, ff.EIGENVALUE_DECOMPOSITION_THRESHOLD + 1): - factor.register_damped_inverse(1. / i) + damping_funcs.append(make_damping_func(1./i)) + + for i in range(ff.EIGENVALUE_DECOMPOSITION_THRESHOLD): + factor.register_inverse(damping_funcs[i]) + + factor.instantiate_inv_variables() ops = factor.make_inverse_update_ops() self.assertEqual(1, len(ops)) sess.run(tf_variables.global_variables_initializer()) new_invs = [] sess.run(ops) - for i in range(1, ff.EIGENVALUE_DECOMPOSITION_THRESHOLD + 1): + for i in range(ff.EIGENVALUE_DECOMPOSITION_THRESHOLD): # The inverse op will assign the damped inverse of cov to the inv var. - new_invs.append(sess.run(factor._inverses_by_damping[1. / i])) + new_invs.append(sess.run(factor.get_inverse(damping_funcs[i]))) + # We want to see that the new invs are all different from each other. for i in range(len(new_invs)): for j in range(i + 1, len(new_invs)): @@ -309,14 +320,16 @@ class InverseProvidingFactorTest(test.TestCase): factor._cov = array_ops.constant(cov, dtype=dtypes.float32) exp = 2 # NOTE(mattjj): must be int to test with np.linalg.matrix_power damping = 0.5 + damping_func = make_damping_func(damping) - factor.register_matpower(exp, damping) + factor.register_matpower(exp, damping_func) + factor.instantiate_inv_variables() ops = factor.make_inverse_update_ops() self.assertEqual(1, len(ops)) sess.run(tf_variables.global_variables_initializer()) sess.run(ops[0]) - matpower = sess.run(factor._matpower_by_exp_and_damping[(exp, damping)]) + matpower = sess.run(factor.get_matpower(exp, damping_func)) matpower_np = np.linalg.matrix_power(cov + np.eye(2) * damping, exp) self.assertAllClose(matpower, matpower_np) @@ -327,18 +340,21 @@ class InverseProvidingFactorTest(test.TestCase): factor = InverseProvidingFactorTestingDummy(cov.shape) factor._cov = array_ops.constant(cov, dtype=dtypes.float32) - factor.register_damped_inverse(0) + damping_func = make_damping_func(0) + + factor.register_inverse(damping_func) + factor.instantiate_inv_variables() ops = factor.make_inverse_update_ops() self.assertEqual(1, len(ops)) sess.run(tf_variables.global_variables_initializer()) # The inverse op will assign the damped inverse of cov to the inv var. - old_inv = sess.run(factor._inverses_by_damping[0]) + old_inv = sess.run(factor.get_inverse(damping_func)) self.assertAllClose( sess.run(ff.inverse_initializer(cov.shape, dtypes.float32)), old_inv) sess.run(ops) - new_inv = sess.run(factor._inverses_by_damping[0]) + new_inv = sess.run(factor.get_inverse(damping_func)) self.assertAllClose(new_inv, np.linalg.inv(cov)) @@ -349,6 +365,7 @@ class FullFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3), name='a/b/c') factor = ff.FullFactor((tensor,), 32) + factor.instantiate_cov_variables() self.assertEqual([6, 6], factor.get_cov().get_shape().as_list()) def testFullFactorInitFloat64(self): @@ -357,6 +374,7 @@ class FullFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3), dtype=dtype, name='a/b/c') factor = ff.FullFactor((tensor,), 32) + factor.instantiate_cov_variables() cov = factor.get_cov() self.assertEqual(cov.dtype, dtype) self.assertEqual([6, 6], cov.get_shape().as_list()) @@ -366,6 +384,7 @@ class FullFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.constant([1., 2.], name='a/b/c') factor = ff.FullFactor((tensor,), 2) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) @@ -379,7 +398,8 @@ 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()) + factor.instantiate_cov_variables() + self.assertEqual([6, 1], factor.get_cov_var().get_shape().as_list()) def testNaiveDiagonalFactorInitFloat64(self): with tf_ops.Graph().as_default(): @@ -387,7 +407,8 @@ 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() + factor.instantiate_cov_variables() + cov = factor.get_cov_var() self.assertEqual(cov.dtype, dtype) self.assertEqual([6, 1], cov.get_shape().as_list()) @@ -396,12 +417,149 @@ class NaiveDiagonalFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.constant([1., 2.], name='a/b/c') factor = ff.NaiveDiagonalFactor((tensor,), 2) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) 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) + factor.instantiate_cov_variables() + 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) + factor.instantiate_cov_variables() + 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 ConvDiagonalFactorTest(test.TestCase): + + def setUp(self): + self.batch_size = 10 + self.height = self.width = 32 + self.in_channels = 3 + self.out_channels = 1 + self.kernel_height = self.kernel_width = 3 + self.strides = [1, 2, 2, 1] + self.data_format = 'NHWC' + self.padding = 'SAME' + self.kernel_shape = [ + self.kernel_height, self.kernel_width, self.in_channels, + self.out_channels + ] + + def testInit(self): + with tf_ops.Graph().as_default(): + inputs = random_ops.random_uniform( + [self.batch_size, self.height, self.width, self.in_channels]) + outputs_grads = [ + random_ops.random_uniform([ + self.batch_size, self.height // self.strides[1], + self.width // self.strides[2], self.out_channels + ]) for _ in range(3) + ] + + factor = ff.ConvDiagonalFactor( + inputs, + outputs_grads, + self.kernel_shape, + self.strides, + self.padding, + data_format=self.data_format) + factor.instantiate_cov_variables() + + # Ensure covariance matrix's shape makes sense. + self.assertEqual([ + self.kernel_height * self.kernel_width * self.in_channels, + self.out_channels + ], + factor.get_cov_var().shape.as_list()) + + def testMakeCovarianceUpdateOp(self): + with tf_ops.Graph().as_default(): + # Construct all arguments such that convolution kernel is applied in + # exactly one spatial location. + inputs = np.random.randn( + 1, # batch_size + self.kernel_height, + self.kernel_width, + self.in_channels) # in_channels + outputs_grad = np.random.randn( + 1, # batch_size + 1, # output_height + 1, # output_width + self.out_channels) + + factor = ff.ConvDiagonalFactor( + constant_op.constant(inputs), [constant_op.constant(outputs_grad)], + self.kernel_shape, + strides=[1, 1, 1, 1], + padding='VALID') + factor.instantiate_cov_variables() + + # Completely forget initial value on first update. + cov_update_op = factor.make_covariance_update_op(0.0) + + # Ensure new covariance value is same as outer-product of inputs/outputs + # vectorized, squared. + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + cov = sess.run(cov_update_op) + expected_cov = np.outer(inputs.flatten(), outputs_grad.flatten())**2 + self.assertAllClose(expected_cov, cov) + + def testHasBias(self): + with tf_ops.Graph().as_default(): + inputs = random_ops.random_uniform( + [self.batch_size, self.height, self.width, self.in_channels]) + outputs_grads = [ + random_ops.random_uniform([ + self.batch_size, self.height // self.strides[1], + self.width // self.strides[2], self.out_channels + ]) for _ in range(3) + ] + + factor = ff.ConvDiagonalFactor( + inputs, + outputs_grads, + self.kernel_shape, + self.strides, + self.padding, + data_format=self.data_format, + has_bias=True) + factor.instantiate_cov_variables() + + # Ensure shape accounts for bias. + self.assertEqual([ + self.kernel_height * self.kernel_width * self.in_channels + 1, + self.out_channels + ], + factor.get_cov_var().shape.as_list()) + + # Ensure update op doesn't crash. + cov_update_op = factor.make_covariance_update_op(0.0) + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + sess.run(cov_update_op) + + class FullyConnectedKroneckerFactorTest(test.TestCase): def _testFullyConnectedKroneckerFactorInit(self, @@ -412,6 +570,7 @@ class FullyConnectedKroneckerFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3), dtype=dtype, name='a/b/c') factor = ff.FullyConnectedKroneckerFactor((tensor,), has_bias=has_bias) + factor.instantiate_cov_variables() cov = factor.get_cov() self.assertEqual(cov.dtype, dtype) self.assertEqual(final_shape, cov.get_shape().as_list()) @@ -429,6 +588,7 @@ class FullyConnectedKroneckerFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c') factor = ff.FullyConnectedKroneckerFactor((tensor,), has_bias=True) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) @@ -439,39 +599,170 @@ class FullyConnectedKroneckerFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c') factor = ff.FullyConnectedKroneckerFactor((tensor,)) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) self.assertAllClose([[3, 3.5], [3.5, 5.5]], new_cov) -class ConvInputKroneckerFactorTest(test.TestCase): +class ConvFactorTestCase(test.TestCase): + + def assertMatrixRank(self, rank, matrix, atol=1e-5): + assert rank <= matrix.shape[0], 'Rank cannot be larger than matrix size.' + eigvals = np.linalg.eigvals(matrix) + nnz_eigvals = np.sum(eigvals > atol) + self.assertEqual( + rank, + nnz_eigvals, + msg=('Found %d of %d expected non-zero eigenvalues: %s.' % + (nnz_eigvals, rank, eigvals))) + + +class ConvInputKroneckerFactorTest(ConvFactorTestCase): + + def test3DConvolution(self): + with tf_ops.Graph().as_default(): + batch_size = 1 + width = 3 + in_channels = 3**3 + out_channels = 4 + + factor = ff.ConvInputKroneckerFactor( + inputs=random_ops.random_uniform( + (batch_size, width, width, width, in_channels), seed=0), + filter_shape=(width, width, width, in_channels, out_channels), + padding='SAME', + strides=(2, 2, 2), + extract_patches_fn='extract_convolution_patches', + has_bias=False) + factor.instantiate_cov_variables() + + # Ensure shape of covariance matches input size of filter. + input_size = in_channels * (width**3) + self.assertEqual([input_size, input_size], + factor.get_cov_var().shape.as_list()) + + # Ensure cov_update_op doesn't crash. + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + sess.run(factor.make_covariance_update_op(0.0)) + cov = sess.run(factor.get_cov_var()) + + # Cov should be rank-8, as the filter will be applied at each corner of + # the 4-D cube. + self.assertMatrixRank(8, cov) + + def testPointwiseConv2d(self): + with tf_ops.Graph().as_default(): + batch_size = 1 + width = 3 + in_channels = 3**2 + out_channels = 4 + + factor = ff.ConvInputKroneckerFactor( + inputs=random_ops.random_uniform( + (batch_size, width, width, in_channels), seed=0), + filter_shape=(1, 1, in_channels, out_channels), + padding='SAME', + strides=(1, 1, 1, 1), + extract_patches_fn='extract_pointwise_conv2d_patches', + has_bias=False) + factor.instantiate_cov_variables() + + # Ensure shape of covariance matches input size of filter. + self.assertEqual([in_channels, in_channels], + factor.get_cov_var().shape.as_list()) + + # Ensure cov_update_op doesn't crash. + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + sess.run(factor.make_covariance_update_op(0.0)) + cov = sess.run(factor.get_cov_var()) + + # Cov should be rank-9, as the filter will be applied at each location. + self.assertMatrixRank(9, cov) + + def testStrides(self): + with tf_ops.Graph().as_default(): + batch_size = 1 + width = 3 + in_channels = 3**2 + out_channels = 4 + + factor = ff.ConvInputKroneckerFactor( + inputs=random_ops.random_uniform( + (batch_size, width, width, in_channels), seed=0), + filter_shape=(1, 1, in_channels, out_channels), + padding='SAME', + strides=(1, 2, 1, 1), + extract_patches_fn='extract_image_patches', + has_bias=False) + factor.instantiate_cov_variables() + + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + sess.run(factor.make_covariance_update_op(0.0)) + cov = sess.run(factor.get_cov_var()) + + # Cov should be the sum of 3 * 2 = 6 outer products. + self.assertMatrixRank(6, cov) + + def testDilationRate(self): + with tf_ops.Graph().as_default(): + batch_size = 1 + width = 3 + in_channels = 2 + out_channels = 4 + + factor = ff.ConvInputKroneckerFactor( + inputs=random_ops.random_uniform( + (batch_size, width, width, in_channels), seed=0), + filter_shape=(3, 3, in_channels, out_channels), + padding='SAME', + extract_patches_fn='extract_image_patches', + strides=(1, 1, 1, 1), + dilation_rate=(1, width, width, 1), + has_bias=False) + factor.instantiate_cov_variables() + + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + sess.run(factor.make_covariance_update_op(0.0)) + cov = sess.run(factor.get_cov_var()) + + # Cov should be rank = in_channels, as only the center of the filter + # receives non-zero input for each input channel. + self.assertMatrixRank(in_channels, cov) def testConvInputKroneckerFactorInitNoBias(self): with tf_ops.Graph().as_default(): - random_seed.set_random_seed(200) - tensor = array_ops.ones((2, 3), name='a/b/c') + tensor = array_ops.ones((64, 1, 2, 3), name='a/b/c') factor = ff.ConvInputKroneckerFactor( - tensor, (1, 2, 3, 4), 3, 2, has_bias=False) + inputs=tensor, + filter_shape=(1, 2, 3, 4), + padding='SAME', + has_bias=False) + factor.instantiate_cov_variables() self.assertEqual([1 * 2 * 3, 1 * 2 * 3], factor.get_cov().get_shape().as_list()) def testConvInputKroneckerFactorInit(self): with tf_ops.Graph().as_default(): - random_seed.set_random_seed(200) - tensor = array_ops.ones((2, 3), name='a/b/c') + tensor = array_ops.ones((64, 1, 2, 3), name='a/b/c') factor = ff.ConvInputKroneckerFactor( - tensor, (1, 2, 3, 4), 3, 2, has_bias=True) + tensor, filter_shape=(1, 2, 3, 4), padding='SAME', has_bias=True) + factor.instantiate_cov_variables() self.assertEqual([1 * 2 * 3 + 1, 1 * 2 * 3 + 1], factor.get_cov().get_shape().as_list()) def testConvInputKroneckerFactorInitFloat64(self): with tf_ops.Graph().as_default(): dtype = dtypes.float64_ref - random_seed.set_random_seed(200) - tensor = array_ops.ones((2, 3), dtype=dtype, name='a/b/c') + tensor = array_ops.ones((64, 1, 2, 3), name='a/b/c', dtype=dtypes.float64) factor = ff.ConvInputKroneckerFactor( - tensor, (1, 2, 3, 4), 3, 2, has_bias=True) + tensor, filter_shape=(1, 2, 3, 4), padding='SAME', has_bias=True) + factor.instantiate_cov_variables() cov = factor.get_cov() self.assertEqual(cov.dtype, dtype) self.assertEqual([1 * 2 * 3 + 1, 1 * 2 * 3 + 1], @@ -479,37 +770,67 @@ class ConvInputKroneckerFactorTest(test.TestCase): def testMakeCovarianceUpdateOpWithBias(self): with tf_ops.Graph().as_default(), self.test_session() as sess: - random_seed.set_random_seed(200) + input_shape = (2, 1, 1, 1) tensor = array_ops.constant( - np.arange(1., 17.).reshape(2, 2, 2, 2), dtype=dtypes.float32) + np.arange(1, 1 + np.prod(input_shape)).reshape(input_shape).astype( + np.float32)) factor = ff.ConvInputKroneckerFactor( - tensor, (1, 2, 1, 1), [1, 1, 1, 1], 'SAME', has_bias=True) + tensor, filter_shape=(1, 1, 1, 1), padding='SAME', has_bias=True) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) - new_cov = sess.run(factor.make_covariance_update_op(.5)) - self.assertAllClose([[34.375, 37, 3.125], [37, 41, 3.5], [3.125, 3.5, 1]], - new_cov) + new_cov = sess.run(factor.make_covariance_update_op(0.)) + self.assertAllClose( + [ + [(1. + 4.) / 2., (1. + 2.) / 2.], # + [(1. + 2.) / 2., (1. + 1.) / 2.] + ], # + new_cov) def testMakeCovarianceUpdateOpNoBias(self): with tf_ops.Graph().as_default(), self.test_session() as sess: - random_seed.set_random_seed(200) + input_shape = (2, 1, 1, 1) tensor = array_ops.constant( - np.arange(1., 17.).reshape(2, 2, 2, 2), dtype=dtypes.float32) - factor = ff.ConvInputKroneckerFactor(tensor, (1, 2, 1, 1), [1, 1, 1, 1], - 'SAME') + np.arange(1, 1 + np.prod(input_shape)).reshape(input_shape).astype( + np.float32)) + factor = ff.ConvInputKroneckerFactor( + tensor, filter_shape=(1, 1, 1, 1), padding='SAME') + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) - new_cov = sess.run(factor.make_covariance_update_op(.5)) - self.assertAllClose([[34.375, 37], [37, 41]], new_cov) + new_cov = sess.run(factor.make_covariance_update_op(0.)) + self.assertAllClose([[(1. + 4.) / 2.]], new_cov) -class ConvOutputKroneckerFactorTest(test.TestCase): +class ConvOutputKroneckerFactorTest(ConvFactorTestCase): + + def test3DConvolution(self): + with tf_ops.Graph().as_default(): + batch_size = 1 + width = 3 + out_channels = width**3 + + factor = ff.ConvOutputKroneckerFactor(outputs_grads=[ + random_ops.random_uniform( + (batch_size, width, width, width, out_channels), seed=0) + ]) + factor.instantiate_cov_variables() + + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + sess.run(factor.make_covariance_update_op(0.0)) + cov = sess.run(factor.get_cov()) + + # Cov should be rank 3^3, as each spatial position donates a rank-1 + # update. + self.assertMatrixRank(width**3, cov) def testConvOutputKroneckerFactorInit(self): with tf_ops.Graph().as_default(): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3, 4, 5), name='a/b/c') factor = ff.ConvOutputKroneckerFactor((tensor,)) + factor.instantiate_cov_variables() self.assertEqual([5, 5], factor.get_cov().get_shape().as_list()) def testConvOutputKroneckerFactorInitFloat64(self): @@ -518,22 +839,17 @@ class ConvOutputKroneckerFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3, 4, 5), dtype=dtype, name='a/b/c') factor = ff.ConvOutputKroneckerFactor((tensor,)) + factor.instantiate_cov_variables() cov = factor.get_cov() self.assertEqual(cov.dtype, dtype) self.assertEqual([5, 5], cov.get_shape().as_list()) - def testConvOutputKroneckerFactorInitNotEnoughDims(self): - with tf_ops.Graph().as_default(): - random_seed.set_random_seed(200) - tensor = array_ops.ones((2, 3), name='a/b/c') - with self.assertRaises(IndexError): - ff.ConvOutputKroneckerFactor(tensor) - def testMakeCovarianceUpdateOp(self): with tf_ops.Graph().as_default(), self.test_session() as sess: random_seed.set_random_seed(200) tensor = np.arange(1, 17).reshape(2, 2, 2, 2).astype(np.float32) factor = ff.ConvOutputKroneckerFactor((array_ops.constant(tensor),)) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) @@ -548,6 +864,7 @@ class FullyConnectedMultiKFTest(test.TestCase): tensor = array_ops.ones((2, 3), name='a/b/c') tensor_list = [tensor] factor = ff.FullyConnectedMultiKF((tensor_list,), has_bias=False) + factor.instantiate_cov_variables() self.assertEqual([3, 3], factor.get_cov().get_shape().as_list()) def testFullyConnectedMultiKFInitFloat64(self): @@ -557,6 +874,7 @@ class FullyConnectedMultiKFTest(test.TestCase): tensor = array_ops.ones((2, 3), dtype=dtype, name='a/b/c') tensor_list = [tensor] factor = ff.FullyConnectedMultiKF((tensor_list,), has_bias=False) + factor.instantiate_cov_variables() cov = factor.get_cov() self.assertEqual(cov.dtype, dtype) self.assertEqual([3, 3], cov.get_shape().as_list()) @@ -567,6 +885,7 @@ class FullyConnectedMultiKFTest(test.TestCase): tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c') tensor_list = [tensor] factor = ff.FullyConnectedMultiKF((tensor_list,), has_bias=True) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) @@ -578,6 +897,7 @@ class FullyConnectedMultiKFTest(test.TestCase): tensor = array_ops.constant([[1., 2.], [3., 4.]], name='a/b/c') tensor_list = [tensor] factor = ff.FullyConnectedMultiKF((tensor_list,)) + factor.instantiate_cov_variables() sess.run(tf_variables.global_variables_initializer()) new_cov = sess.run(factor.make_covariance_update_op(.5)) diff --git a/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py b/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py index b8ccbeadd0a9d69edb41fef50e3edb090457adf2..bae6bd7a3bd47bc50378afe95d26d57535377f6f 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/layer_collection_test.py @@ -104,14 +104,31 @@ class LayerCollectionTest(test.TestCase): array_ops.constant(3), approx=layer_collection.APPROX_DIAGONAL_NAME) lc.register_conv2d( - array_ops.constant(4), [1, 1, 1, 1], 'SAME', - array_ops.ones((1, 1, 1, 1)), array_ops.constant(3)) + params=array_ops.ones((2, 3, 4, 5)), + strides=[1, 1, 1, 1], + padding='SAME', + inputs=array_ops.ones((1, 2, 3, 4)), + outputs=array_ops.ones((1, 1, 1, 5))) lc.register_conv2d( - array_ops.constant(4), [1, 1, 1, 1], - 'SAME', - array_ops.ones((1, 1, 1, 1)), - array_ops.constant(3), + params=array_ops.ones((2, 3, 4, 5)), + strides=[1, 1, 1, 1], + padding='SAME', + inputs=array_ops.ones((1, 2, 3, 4)), + outputs=array_ops.ones((1, 1, 1, 5)), approx=layer_collection.APPROX_DIAGONAL_NAME) + lc.register_separable_conv2d( + depthwise_params=array_ops.ones((3, 3, 1, 2)), + pointwise_params=array_ops.ones((1, 1, 2, 4)), + inputs=array_ops.ones((32, 5, 5, 1)), + depthwise_outputs=array_ops.ones((32, 5, 5, 2)), + pointwise_outputs=array_ops.ones((32, 5, 5, 4)), + strides=[1, 1, 1, 1], + padding='SAME') + lc.register_convolution( + params=array_ops.ones((3, 3, 1, 8)), + inputs=array_ops.ones((32, 5, 5, 1)), + outputs=array_ops.ones((32, 5, 5, 8)), + padding='SAME') lc.register_generic( array_ops.constant(5), 16, approx=layer_collection.APPROX_FULL_NAME) lc.register_generic( @@ -119,7 +136,7 @@ class LayerCollectionTest(test.TestCase): 16, approx=layer_collection.APPROX_DIAGONAL_NAME) - self.assertEqual(6, len(lc.get_blocks())) + self.assertEqual(9, len(lc.get_blocks())) def testRegisterBlocksMultipleRegistrations(self): with ops.Graph().as_default(): @@ -237,16 +254,16 @@ class LayerCollectionTest(test.TestCase): # Create a new loss function by name. lc.register_categorical_predictive_distribution(logits, name='loss1') - self.assertEqual(1, len(lc.losses)) + self.assertEqual(1, len(lc.towers_by_loss)) # Add logits to same loss function. lc.register_categorical_predictive_distribution( logits, name='loss1', reuse=True) - self.assertEqual(1, len(lc.losses)) + self.assertEqual(1, len(lc.towers_by_loss)) # Add another new loss function. lc.register_categorical_predictive_distribution(logits, name='loss2') - self.assertEqual(2, len(lc.losses)) + self.assertEqual(2, len(lc.towers_by_loss)) def testLossFunctionWithoutName(self): """Ensure loss functions get unique names if 'name' not specified.""" @@ -298,13 +315,9 @@ class LayerCollectionTest(test.TestCase): name='loss1', reuse=layer_collection.VARIABLE_SCOPE) - self.assertEqual(len(lc.losses), 1) - loss = lc.losses[0] - + self.assertEqual(len(lc.towers_by_loss), 1) # Three successful registrations. - self.assertEqual(loss.params.shape.as_list(), - [3 * batch_size, output_size]) - self.assertEqual(loss.targets.shape.as_list(), [3 * batch_size]) + self.assertEqual(len(lc.towers_by_loss[0]), 3) def testRegisterCategoricalPredictiveDistributionBatchSize1(self): with ops.Graph().as_default(): @@ -479,17 +492,6 @@ class LayerCollectionTest(test.TestCase): variables = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) self.assertTrue(all([var.name.startswith(scope) for var in variables])) - def testGetUseCountMap(self): - """Ensure get_use_count_map() sums 'num_registered_minibatches'.""" - lc = layer_collection.LayerCollection() - lc.fisher_blocks = { - 'a': MockFisherBlock(), - ('a', 'c'): MockFisherBlock(), - ('b', 'c'): MockFisherBlock() - } - use_count_map = lc.get_use_count_map() - self.assertDictEqual({'a': 4, 'b': 2, 'c': 4}, use_count_map) - def testIdentifyLinkedParametersSomeRegisteredInOtherTuples(self): x = variable_scope.get_variable('x', shape=()) y = variable_scope.get_variable('y', shape=()) @@ -550,6 +552,32 @@ class LayerCollectionTest(test.TestCase): self.assertIsInstance(lc.fisher_blocks[b_0], fisher_blocks.FullFB) self.assertIsInstance(lc.fisher_blocks[b_1], fisher_blocks.NaiveDiagonalFB) + def testDefaultLayerCollection(self): + with ops.Graph().as_default(): + # Can't get default if there isn't one set. + with self.assertRaises(ValueError): + layer_collection.get_default_layer_collection() + + # Can't set default twice. + lc = layer_collection.LayerCollection() + layer_collection.set_default_layer_collection(lc) + with self.assertRaises(ValueError): + layer_collection.set_default_layer_collection(lc) + + # Same as one set. + self.assertTrue(lc is layer_collection.get_default_layer_collection()) + + # Can set to None. + layer_collection.set_default_layer_collection(None) + with self.assertRaises(ValueError): + layer_collection.get_default_layer_collection() + + # as_default() is the same as setting/clearing. + with lc.as_default(): + self.assertTrue(lc is layer_collection.get_default_layer_collection()) + with self.assertRaises(ValueError): + layer_collection.get_default_layer_collection() + if __name__ == '__main__': test.main() 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 ae787b6f1ac90218f2ac73d37fb270df0b822de2..c00af5593f085e3b1f3e030a24f4b821115cc869 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py @@ -24,7 +24,6 @@ from tensorflow.contrib.kfac.python.ops import loss_functions 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 random_ops from tensorflow.python.platform import test @@ -97,22 +96,6 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase): # 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.CategoricalLogitsNegativeLogProbLoss(logits) - else: - loss.register_additional_minibatch(logits) - 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.]) @@ -203,23 +186,5 @@ class OnehotCategoricalLogitsNegativeLogProbLossTest(test.TestCase): # 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 97a97adbf5577cd2694d3055acaa59258ad27964..2cee01212a11595669e9df0fc95a5657926c1038 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py @@ -29,6 +29,8 @@ 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 nn_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 @@ -325,6 +327,84 @@ class UtilsTest(test.TestCase): ], values) + def testExtractConvolutionPatches(self): + with ops.Graph().as_default(), self.test_session() as sess: + batch_size = 10 + image_spatial_shape = [9, 10, 11] + in_channels = out_channels = 32 + kernel_spatial_shape = [5, 3, 3] + spatial_strides = [1, 2, 1] + spatial_dilation = [1, 1, 1] + padding = 'SAME' + + images = random_ops.random_uniform( + [batch_size] + image_spatial_shape + [in_channels], seed=0) + kernel_shape = kernel_spatial_shape + [in_channels, out_channels] + kernel = random_ops.random_uniform(kernel_shape, seed=1) + + # Ensure shape matches expectation. + patches = utils.extract_convolution_patches( + images, + kernel_shape, + padding, + strides=spatial_strides, + dilation_rate=spatial_dilation) + result_spatial_shape = ( + patches.shape.as_list()[1:1 + len(image_spatial_shape)]) + self.assertEqual(patches.shape.as_list(), + [batch_size] + result_spatial_shape + + kernel_spatial_shape + [in_channels]) + + # Ensure extract...patches() + matmul() and convolution() implementation + # give the same answer. + outputs = nn_ops.convolution( + images, + kernel, + padding, + strides=spatial_strides, + dilation_rate=spatial_dilation) + + patches_flat = array_ops.reshape( + patches, [-1, np.prod(kernel_spatial_shape) * in_channels]) + kernel_flat = array_ops.reshape(kernel, [-1, out_channels]) + outputs_flat = math_ops.matmul(patches_flat, kernel_flat) + + outputs_, outputs_flat_ = sess.run([outputs, outputs_flat]) + self.assertAllClose(outputs_.flatten(), outputs_flat_.flatten()) + + def testExtractPointwiseConv2dPatches(self): + with ops.Graph().as_default(), self.test_session() as sess: + batch_size = 10 + image_height = image_width = 8 + in_channels = out_channels = 3 + kernel_height = kernel_width = 1 + strides = [1, 1, 1, 1] + padding = 'VALID' + + images = random_ops.random_uniform( + [batch_size, image_height, image_width, in_channels], seed=0) + kernel_shape = [kernel_height, kernel_width, in_channels, out_channels] + kernel = random_ops.random_uniform(kernel_shape, seed=1) + + # Ensure shape matches expectation. + patches = utils.extract_pointwise_conv2d_patches(images, kernel_shape) + self.assertEqual(patches.shape.as_list(), [ + batch_size, image_height, image_width, kernel_height, kernel_width, + in_channels + ]) + + # Ensure extract...patches() + matmul() and conv2d() implementation + # give the same answer. + outputs = nn_ops.conv2d(images, kernel, strides, padding) + + patches_flat = array_ops.reshape( + patches, [-1, kernel_height * kernel_width * in_channels]) + kernel_flat = array_ops.reshape(kernel, [-1, out_channels]) + outputs_flat = math_ops.matmul(patches_flat, kernel_flat) + + outputs_, outputs_flat_ = sess.run([outputs, outputs_flat]) + self.assertAllClose(outputs_.flatten(), outputs_flat_.flatten()) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/kfac/python/ops/BUILD b/tensorflow/contrib/kfac/python/ops/BUILD index ee6549b109399766579b6ea18a987ae2c8275983..c26230c2a82ae9529ab13b523b9ec287d17debaf 100644 --- a/tensorflow/contrib/kfac/python/ops/BUILD +++ b/tensorflow/contrib/kfac/python/ops/BUILD @@ -144,10 +144,13 @@ py_library( ":fisher_estimator", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:linalg_ops", "//tensorflow/python:math_ops", + "//tensorflow/python:state_ops", "//tensorflow/python:training", + "//tensorflow/python:variable_scope", "//tensorflow/python:variables", ], ) diff --git a/tensorflow/contrib/kfac/python/ops/estimator.py b/tensorflow/contrib/kfac/python/ops/estimator.py index a7b1f9d35c931fc44408be804479e758f28f7110..64755be65c4b5686397dbfd798fec1ed70ae61dc 100644 --- a/tensorflow/contrib/kfac/python/ops/estimator.py +++ b/tensorflow/contrib/kfac/python/ops/estimator.py @@ -27,6 +27,7 @@ from tensorflow.contrib.kfac.python.ops import utils from tensorflow.python.framework import ops as tf_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import variable_scope from tensorflow.python.util import nest @@ -65,6 +66,13 @@ class _DeviceContextGenerator(object): yield +def _make_thunk_on_device(func, device): + def thunk(): + with tf_ops.device(device): + return func() + return thunk + + class FisherEstimator(object): """Fisher estimator class supporting various approximations of the Fisher. @@ -87,10 +95,10 @@ class FisherEstimator(object): cov_ema_decay, damping, layer_collection, + exps=(-1,), estimation_mode="gradients", colocate_gradients_with_ops=True, - cov_devices=None, - inv_devices=None): + name="FisherEstimator"): """Create a FisherEstimator object. Args: @@ -99,14 +107,19 @@ class FisherEstimator(object): None). cov_ema_decay: The decay factor used when calculating the covariance estimate moving averages. - damping: The damping factor used to stabilize training due to errors in - the local approximation with the Fisher information matrix, and to - regularize the update direction by making it closer to the gradient. - (Higher damping means the update looks more like a standard gradient - update - see Tikhonov regularization.) + damping: float. The damping factor used to stabilize training due to + errors in the local approximation with the Fisher information matrix, + and to regularize the update direction by making it closer to the + gradient. (Higher damping means the update looks more like a standard + gradient update - see Tikhonov regularization.) layer_collection: The layer collection object, which holds the fisher blocks, kronecker factors, and losses associated with the graph. + exps: List of floats or ints. These represent the different matrix + powers of the approximate Fisher that the FisherEstimator will be able + to multiply vectors by. If the user asks for a matrix power other + one of these (or 1, which is always supported), there will be a + failure. (Default: (-1,)) estimation_mode: The type of estimator to use for the Fishers. Can be 'gradients', 'empirical', 'curvature_prop', or 'exact'. (Default: 'gradients'). 'gradients' is the basic estimation approach @@ -125,24 +138,17 @@ class FisherEstimator(object): equal to the output dimension, roughly speaking. colocate_gradients_with_ops: Whether we should request gradients 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. - inv_devices: Iterable of device strings (e.g. '/gpu:0'). Inversion - computations will be placed on these devices in a round-robin fashion. - Can be None, which means that no devices are specified. - + name: A string. A name given to this estimator, which is added to the + variable scope when constructing variables and ops. + (Default: "FisherEstimator") Raises: ValueError: If no losses have been registered with layer_collection. """ - - self._cov_ema_decay = cov_ema_decay self._variables = variables + self._cov_ema_decay = cov_ema_decay self._damping = damping self._estimation_mode = estimation_mode self._layers = layer_collection - self._layers.create_subgraph() - self._layers.check_registration(variables) self._gradient_fns = { "gradients": self._get_grads_lists_gradients, "empirical": self._get_grads_lists_empirical, @@ -151,30 +157,10 @@ class FisherEstimator(object): } 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) + self._made_vars = False + self._exps = exps - 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") + self._name = name @property def variables(self): @@ -184,6 +170,20 @@ class FisherEstimator(object): def damping(self): return self._damping + @property + def blocks(self): + """All registered FisherBlocks.""" + return self._layers.get_blocks() + + @property + def factors(self): + """All registered FisherFactors.""" + return self._layers.get_factors() + + @property + def name(self): + return self._name + def _apply_transformation(self, vecs_and_vars, transform): """Applies an block-wise transformation to the corresponding vectors. @@ -217,9 +217,7 @@ class FisherEstimator(object): A list of (transformed vector, var) pairs in the same order as vecs_and_vars. """ - - return self._apply_transformation(vecs_and_vars, - lambda fb, vec: fb.multiply_inverse(vec)) + return self.multiply_matpower(-1, vecs_and_vars) def multiply(self, vecs_and_vars): """Multiplies the vectors by the corresponding (damped) blocks. @@ -231,9 +229,22 @@ class FisherEstimator(object): A list of (transformed vector, var) pairs in the same order as vecs_and_vars. """ + return self.multiply_matpower(1, vecs_and_vars) - return self._apply_transformation(vecs_and_vars, - lambda fb, vec: fb.multiply(vec)) + def multiply_matpower(self, exp, vecs_and_vars): + """Multiplies the vecs by the corresponding matrix powers of the blocks. + + Args: + exp: A float representing the power to raise the blocks by before + multiplying it by the vector. + vecs_and_vars: List of (vector, variable) pairs. + + Returns: + A list of (transformed vector, var) pairs in the same order as + vecs_and_vars. + """ + fcn = lambda fb, vec: fb.multiply_matpower(vec, exp) + return self._apply_transformation(vecs_and_vars, fcn) def _instantiate_factors(self): """Instantiates FisherFactors' variables. @@ -241,9 +252,9 @@ class FisherEstimator(object): Raises: ValueError: If estimation_mode was improperly specified at construction. """ - fisher_blocks_list = self._layers.get_blocks() + blocks = self.blocks tensors_to_compute_grads = [ - fb.tensors_to_compute_grads() for fb in fisher_blocks_list + block.tensors_to_compute_grads() for block in blocks ] try: @@ -253,45 +264,283 @@ class FisherEstimator(object): raise ValueError("Unrecognized value {} for estimation_mode.".format( self._estimation_mode)) - # TODO(b/68033310): This loop round-robins the "concat" operations which - # gather the inputs for the cov_updates. In future, we might do these - # computations locally then communicate the results, which would require a - # modification to this code. - for grads_list, fb in zip(grads_lists, fisher_blocks_list): - with self._cov_device_context_generator(): - fb.instantiate_factors(grads_list, self.damping) + for grads_list, block in zip(grads_lists, blocks): + block.instantiate_factors(grads_list, self.damping) + + def _check_vars_unmade_and_set_made_flag(self): + if self._made_vars: + raise Exception("Already made variables.") + self._made_vars = True + + def made_vars(self): + return self._made_vars + + def _register_matrix_functions(self): + for exp in self._exps: + for block in self.blocks: + block.register_matpower(exp) + + def _finalize_layer_collection(self): + self._layers.create_subgraph() + self._layers.check_registration(self.variables) + self._instantiate_factors() + self._register_matrix_functions() + + def make_ops_and_vars(self, scope=None): + """Make ops and vars with no specific device placement. + + See make_ops_and_vars_round_robin for further details. + + Args: + scope: A string or None. If None it will be set to the name of this + estimator (given by the name property). All variables will be created, + and all ops will execute, inside of a variable scope of the given + name. (Default: None) + Returns: + cov_update_ops: List of ops that compute the cov updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + cov_update_op: cov_update_ops grouped into a single op. + inv_update_ops: List of ops that compute the inv updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + inv_update_op: inv_update_ops grouped into a single op. + cov_update_thunks: Thunks that make the ops in cov_update_ops. + inv_update_thunks: Thunks that make the ops in inv_update_ops. + """ + return self.make_ops_and_vars_round_robin(scope=scope) + + # TODO(b/70674513): Factor device placement outside of this class. + def make_ops_and_vars_round_robin(self, scope=None, cov_devices=None, + inv_devices=None): + """Make ops and vars with a round-robin device placement strategy. + + For each factor, all of that factor's cov variables and their associated + update ops will be placed on a particular device. A new device is chosen + for each factor by cycling through list of devices in the cov_devices + argument. If cov_devices is None then no explicit device placement occurs. + + An analogous strategy is followed for inverse update ops, with the list of + devices being given by the inv_devices argument. - def _create_cov_update_thunk(self, factor): + Inverse variables on the other hand are not placed on any specific device + (they will just use the current the device placement context, whatever + that happens to be). The idea is that the inverse variable belong where + they will be accessed most often, which is the device that actually applies + the preconditioner to the gradient. The user will be responsible for setting + the device context for this. + + Args: + scope: A string or None. If None it will be set to the name of this + estimator (given by the name property). All variables will be created, + and all ops will execute, inside of a variable scope of the given + name. (Default: None) + 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. + inv_devices: Iterable of device strings (e.g. '/gpu:0'). Inversion + computations will be placed on these devices in a round-robin fashion. + Can be None, which means that no devices are specified. + + Returns: + cov_update_ops: List of ops that compute the cov updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + cov_update_op: cov_update_ops grouped into a single op. + inv_update_ops: List of ops that compute the inv updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + inv_update_op: inv_update_ops grouped into a single op. + cov_update_thunks: Thunks that make the ops in cov_update_ops. + inv_update_thunks: Thunks that make the ops in inv_update_ops. + """ + (cov_update_thunks, + inv_update_thunks) = self.make_vars_and_create_op_thunks_round_robin( + scope=scope, + cov_devices=cov_devices, + inv_devices=inv_devices) + cov_update_ops = [thunk() for thunk in cov_update_thunks] + inv_update_ops = [thunk() for thunk in inv_update_thunks] + + scope = self.name if scope is None else scope + with variable_scope.variable_scope(scope): + cov_update_op = control_flow_ops.group(cov_update_ops, + name="cov_update_op") + inv_update_op = control_flow_ops.group(inv_update_ops, + name="inv_update_op") + + return (cov_update_ops, cov_update_op, inv_update_ops, inv_update_op, + cov_update_thunks, inv_update_thunks) + + def make_vars_and_create_op_thunks_round_robin(self, + scope=None, + cov_devices=None, + inv_devices=None): + """Make vars and create op thunks w/ a round-robin device placement strat. + + For each factor, all of that factor's cov variables and their associated + update ops will be placed on a particular device. A new device is chosen + for each factor by cycling through list of devices in the cov_devices + argument. If cov_devices is None then no explicit device placement occurs. + + An analogous strategy is followed for inverse update ops, with the list of + devices being given by the inv_devices argument. + + Inverse variables on the other hand are not placed on any specific device + (they will just use the current the device placement context, whatever + that happens to be). The idea is that the inverse variable belong where + they will be accessed most often, which is the device that actually applies + the preconditioner to the gradient. The user will be responsible for setting + the device context for this. + + Args: + scope: A string or None. If None it will be set to the name of this + estimator (given by the name property). All variables will be created, + and all thunks will execute, inside of a variable scope of the given + name. (Default: None) + 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. + inv_devices: Iterable of device strings (e.g. '/gpu:0'). Inversion + computations will be placed on these devices in a round-robin fashion. + Can be None, which means that no devices are specified. + Returns: + cov_update_thunks: List of cov update thunks. Corresponds one-to-one with + the list of factors given by the "factors" property. + inv_update_thunks: List of inv update thunks. Corresponds one-to-one with + the list of factors given by the "factors" property. + """ + + (cov_variable_thunks_raw, cov_update_thunks_raw, inv_variable_thunks_raw, + inv_update_thunks_raw) = self.create_ops_and_vars_thunks(scope=scope) + + if cov_devices: + cov_update_thunks = [] + for cov_variable_thunk, cov_update_thunk, device in zip( + cov_variable_thunks_raw, cov_update_thunks_raw, + itertools.cycle(cov_devices)): + with tf_ops.device(device): + cov_variable_thunk() + cov_update_thunks.append(_make_thunk_on_device(cov_update_thunk, + device)) + else: + for cov_variable_thunk in cov_variable_thunks_raw: + cov_variable_thunk() + cov_update_thunks = cov_update_thunks_raw + + for inv_variable_thunk in inv_variable_thunks_raw: + inv_variable_thunk() + + if inv_devices: + inv_update_thunks = [] + for inv_update_thunk, device in zip(inv_update_thunks_raw, + itertools.cycle(inv_devices)): + inv_update_thunks.append(_make_thunk_on_device(inv_update_thunk, + device)) + else: + inv_update_thunks = inv_update_thunks_raw + + return cov_update_thunks, inv_update_thunks + + def create_ops_and_vars_thunks(self, scope=None): + """Create thunks that make the ops and vars on demand. + + This function returns 4 lists of thunks: cov_variable_thunks, + cov_update_thunks, inv_variable_thunks, and inv_update_thunks. + + The length of each list is the number of factors and the i-th element of + each list corresponds to the i-th factor (given by the "factors" property). + + Note that the execution of these thunks must happen in a certain + partial order. The i-th element of cov_variable_thunks must execute + before the i-th element of cov_update_thunks (and also the i-th element + of inv_update_thunks). Similarly, the i-th element of inv_variable_thunks + must execute before the i-th element of inv_update_thunks. + + TL;DR (oversimplified): Execute the thunks according to the order that + they are returned. + + Args: + scope: A string or None. If None it will be set to the name of this + estimator (given by the name property). All thunks will execute inside + of a variable scope of the given name. (Default: None) + Returns: + cov_variable_thunks: A list of thunks that make the cov variables. + cov_update_thunks: A list of thunks that make the cov update ops. + inv_variable_thunks: A list of thunks that make the inv variables. + inv_update_thunks: A list of thunks that make the inv update ops. + """ + self._check_vars_unmade_and_set_made_flag() + + self._finalize_layer_collection() + + scope = self.name if scope is None else scope + + cov_variable_thunks = [ + self._create_cov_variable_thunk(factor, scope) + for factor in self.factors + ] + cov_update_thunks = [ + self._create_cov_update_thunk(factor, scope) for factor in self.factors + ] + inv_variable_thunks = [ + self._create_inv_variable_thunk(factor, scope) + for factor in self.factors + ] + inv_update_thunks = [ + self._create_inv_update_thunk(factor, scope) for factor in self.factors + ] + + return (cov_variable_thunks, cov_update_thunks, + inv_variable_thunks, inv_update_thunks) + + def _create_cov_variable_thunk(self, factor, scope): + """Constructs a covariance variable thunk for a single FisherFactor.""" + + def thunk(): + with variable_scope.variable_scope(scope): + return factor.instantiate_cov_variables() + + return thunk + + def _create_cov_update_thunk(self, factor, scope): """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]): + with variable_scope.variable_scope(scope): return factor.make_covariance_update_op(self._cov_ema_decay) return thunk - def _create_inv_update_thunk(self, factor): + def _create_inv_variable_thunk(self, factor, scope): + """Constructs a inverse variable thunk for a single FisherFactor.""" + + def thunk(): + with variable_scope.variable_scope(scope): + return factor.instantiate_inv_variables() + + return thunk + + def _create_inv_update_thunk(self, factor, scope): """Constructs an inverse update thunk for a single FisherFactor.""" 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()) + with variable_scope.variable_scope(scope): + return control_flow_ops.group(factor.make_inverse_update_ops()) return thunk def _get_grads_lists_gradients(self, tensors): + # Passing in a list of loss values is better than passing in the sum as + # the latter creates unnessesary ops on the default device grads_flat = gradients_impl.gradients( - self._layers.total_sampled_loss(), + self._layers.eval_losses_on_samples(), nest.flatten(tensors), colocate_gradients_with_ops=self._colocate_gradients_with_ops) grads_all = nest.pack_sequence_as(tensors, grads_flat) return tuple((grad,) for grad in grads_all) def _get_grads_lists_empirical(self, tensors): + # Passing in a list of loss values is better than passing in the sum as + # the latter creates unnessesary ops on the default device grads_flat = gradients_impl.gradients( - self._layers.total_loss(), + self._layers.eval_losses(), nest.flatten(tensors), colocate_gradients_with_ops=self._colocate_gradients_with_ops) grads_all = nest.pack_sequence_as(tensors, grads_flat) @@ -300,9 +549,10 @@ class FisherEstimator(object): def _get_transformed_random_signs(self): transformed_random_signs = [] for loss in self._layers.losses: - transformed_random_signs.append( - loss.multiply_fisher_factor( - utils.generate_random_signs(loss.fisher_factor_inner_shape))) + with tf_ops.colocate_with(self._layers.loss_colocation_ops[loss]): + transformed_random_signs.append( + loss.multiply_fisher_factor( + utils.generate_random_signs(loss.fisher_factor_inner_shape))) return transformed_random_signs def _get_grads_lists_curvature_prop(self, tensors): @@ -321,13 +571,14 @@ class FisherEstimator(object): # Loop over all coordinates of all losses. grads_all = [] for loss in self._layers.losses: - for index in np.ndindex(*loss.fisher_factor_inner_static_shape[1:]): - transformed_one_hot = loss.multiply_fisher_factor_replicated_one_hot( - index) - grads_flat = gradients_impl.gradients( - loss.inputs, - nest.flatten(tensors), - grad_ys=transformed_one_hot, - colocate_gradients_with_ops=self._colocate_gradients_with_ops) - grads_all.append(nest.pack_sequence_as(tensors, grads_flat)) + with tf_ops.colocate_with(self._layers.loss_colocation_ops[loss]): + for index in np.ndindex(*loss.fisher_factor_inner_static_shape[1:]): + transformed_one_hot = loss.multiply_fisher_factor_replicated_one_hot( + index) + grads_flat = gradients_impl.gradients( + loss.inputs, + nest.flatten(tensors), + grad_ys=transformed_one_hot, + colocate_gradients_with_ops=self._colocate_gradients_with_ops) + grads_all.append(nest.pack_sequence_as(tensors, grads_flat)) return zip(*grads_all) diff --git a/tensorflow/contrib/kfac/python/ops/fisher_blocks.py b/tensorflow/contrib/kfac/python/ops/fisher_blocks.py index 0d2fa706f5853570bb8c04a9b9ac3378e2f2386e..31f4689fbfbbf13872c237913a37478f3c2debe0 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_blocks.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_blocks.py @@ -40,10 +40,12 @@ from __future__ import print_function import abc import enum # pylint: disable=g-bad-import-order +import numpy as np import six from tensorflow.contrib.kfac.python.ops import fisher_factors from tensorflow.contrib.kfac.python.ops import utils +from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -92,10 +94,22 @@ 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) @@ -109,12 +123,44 @@ def compute_pi_adjusted_damping(left_cov, right_cov, damping): return (damping, damping) +class PackagedFunc(object): + """A Python thunk with a stable ID. + + Enables stable names for lambdas. + """ + + def __init__(self, func, func_id): + """Initializes PackagedFunc. + + Args: + func: a zero-arg Python function. + func_id: a hashable, function that produces a hashable, or a list/tuple + thereof. + """ + self._func = func + func_id = func_id if isinstance(func_id, (tuple, list)) else (func_id,) + self._func_id = func_id + + def __call__(self): + return self._func() + + @property + def func_id(self): + """A hashable identifier for this function.""" + return tuple(elt() if callable(elt) else elt for elt in self._func_id) + + +def _package_func(func, func_id): + return PackagedFunc(func, func_id) + + @six.add_metaclass(abc.ABCMeta) class FisherBlock(object): """Abstract base class for objects modeling approximate Fisher matrix blocks. - Subclasses must implement multiply_inverse(), instantiate_factors(), and - tensors_to_compute_grads() methods. + Subclasses must implement register_matpower, multiply_matpower, + instantiate_factors, tensors_to_compute_grads, and num_registered_minibatches + methods. """ def __init__(self, layer_collection): @@ -133,6 +179,32 @@ class FisherBlock(object): pass @abc.abstractmethod + def register_matpower(self, exp): + """Registers a matrix power to be computed by the block. + + Args: + exp: A float representing the power to raise the block by. + """ + pass + + def register_inverse(self): + """Registers a matrix inverse to be computed by the block.""" + self.register_matpower(-1) + + @abc.abstractmethod + def multiply_matpower(self, vector, exp): + """Multiplies the vector by the (damped) matrix-power of the block. + + Args: + vector: The vector (a Tensor or tuple of Tensors) to be multiplied. + exp: A float representing the power to raise the block by before + multiplying it by the vector. + + Returns: + The vector left-multiplied by the (damped) matrix-power of the block. + """ + pass + def multiply_inverse(self, vector): """Multiplies the vector by the (damped) inverse of the block. @@ -142,9 +214,8 @@ class FisherBlock(object): Returns: The vector left-multiplied by the (damped) inverse of the block. """ - pass + return self.multiply_matpower(vector, -1) - @abc.abstractmethod def multiply(self, vector): """Multiplies the vector by the (damped) block. @@ -154,7 +225,7 @@ class FisherBlock(object): Returns: The vector left-multiplied by the (damped) block. """ - pass + return self.multiply_matpower(vector, 1) @abc.abstractmethod def tensors_to_compute_grads(self): @@ -195,21 +266,18 @@ class FullFB(FisherBlock): super(FullFB, self).__init__(layer_collection) def instantiate_factors(self, grads_list, damping): - self._damping = damping + self._damping_func = _package_func(lambda: damping, (damping,)) + self._factor = self._layer_collection.make_or_get_factor( fisher_factors.FullFactor, (grads_list, self._batch_size)) - 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)) - return utils.column_to_tensors(vector, out_flat) + def register_matpower(self, exp): + self._factor.register_matpower(exp, self._damping_func) - def multiply(self, vector): + def multiply_matpower(self, vector, exp): 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_matpower( + vector_flat, exp, self._damping_func) return utils.column_to_tensors(vector, out_flat) def full_fisher_block(self): @@ -259,22 +327,24 @@ class NaiveDiagonalFB(FisherBlock): super(NaiveDiagonalFB, self).__init__(layer_collection) def instantiate_factors(self, grads_list, damping): - self._damping = damping + self._damping_func = _package_func(lambda: damping, (damping,)) + self._factor = self._layer_collection.make_or_get_factor( fisher_factors.NaiveDiagonalFactor, (grads_list, self._batch_size)) - def multiply_inverse(self, vector): - vector_flat = utils.tensors_to_column(vector) - out_flat = vector_flat / (self._factor.get_cov() + self._damping) - return utils.column_to_tensors(vector, out_flat) + def register_matpower(self, exp): + # Not needed for this. Matrix powers are computed on demand in the + # diagonal case + pass - def multiply(self, vector): + def multiply_matpower(self, vector, exp): vector_flat = utils.tensors_to_column(vector) - out_flat = vector_flat * (self._factor.get_cov() + self._damping) + out_flat = self._factor.left_multiply_matpower( + vector_flat, exp, self._damping_func) 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 @@ -296,7 +366,89 @@ class NaiveDiagonalFB(FisherBlock): return math_ops.reduce_sum(self._batch_sizes) -class FullyConnectedDiagonalFB(FisherBlock): +class InputOutputMultiMinibatch(object): + """Mix-in class for blocks with inputs & outputs and multiple mini-batches.""" + + def __init__(self, *args, **kwargs): + self.__inputs = [] + self.__outputs = [] + super(InputOutputMultiMinibatch, self).__init__(*args, **kwargs) + + def tensors_to_compute_grads(self): + """Tensors to compute derivative of loss with respect to.""" + return self._outputs + + def register_additional_minibatch(self, inputs, outputs): + self._inputs.append(inputs) + self._outputs.append(outputs) + + @property + def num_registered_minibatches(self): + result = len(self._inputs) + assert result == len(self._outputs) + return result + + @property + def _inputs(self): + return self.__inputs + + @property + def _outputs(self): + return self.__outputs + + def _package_minibatches(self, grads_list): + """Constructs PartitionedTensor for inputs, grads_list. + + The purpose of this method is to package up the towers/minibatch dimension + of these arrays into PartitionedTensor objects. + + Args: + grads_list: 2-D list of Tensors. First index is for source, second + index for tower. + + Returns: + inputs: PartitionedTensor. + grads_list: Tuple of PartitionedTensors, one per source. + """ + inputs = utils.PartitionedTensor(self._inputs) + grads_list = tuple(utils.PartitionedTensor(grads) for grads in grads_list) + + return inputs, grads_list + + def _package_minibatches_multi(self, grads_list): + """Constructs PartitionedTensors for inputs, grads_list. + + The purpose of this method is to package up the towers/minibatch dimension + of these arrays into PartitionedTensor objects. + + This version of this function is for use with FisherBlocks that deal with + multiple uses or time-steps. One PartitionedTensor is created for each + use/time-step. + + Args: + grads_list: 3-D tuple of Tensors. First index is for source, second + index is for tower, third is for use/time-step. + + Returns: + inputs: A tuple of PartitionedTensor's, one per use/time-step. + grads_list: 2-D tuple of PartitionedTensors. First index is for source, + second is for use/time-step. + """ + # self._inputs is a 2-D tuple. First index is tower/mini-batch, second is + # use/time-step. + inputs = self._inputs + num_uses = len(inputs[0]) + assert all(len(input_) == num_uses for input_ in inputs) + assert all(len(grad) == num_uses for grads in grads_list for grad in grads) + + inputs = tuple(utils.PartitionedTensor(input_) for input_ in zip(*inputs)) + grads_list = tuple(tuple(utils.PartitionedTensor(grad) + for grad in zip(*grads)) for grads in grads_list) + + return inputs, grads_list + + +class FullyConnectedDiagonalFB(InputOutputMultiMinibatch, FisherBlock): """FisherBlock for fully-connected (dense) layers using a diagonal approx. Estimates the Fisher Information matrix's diagonal entries for a fully @@ -328,78 +480,46 @@ class FullyConnectedDiagonalFB(FisherBlock): has_bias: Whether the component Kronecker factors have an additive bias. (Default: False) """ - self._inputs = [] - self._outputs = [] self._has_bias = has_bias super(FullyConnectedDiagonalFB, self).__init__(layer_collection) def instantiate_factors(self, grads_list, damping): - inputs = _concat_along_batch_dim(self._inputs) - grads_list = tuple(_concat_along_batch_dim(grads) for grads in grads_list) + inputs, grads_list = self._package_minibatches(grads_list) - self._damping = damping self._factor = self._layer_collection.make_or_get_factor( fisher_factors.FullyConnectedDiagonalFactor, (inputs, grads_list, self._has_bias)) - def multiply_inverse(self, vector): - """Approximate damped inverse Fisher-vector product. + self._damping_func = _package_func(lambda: damping, (damping,)) - Args: - vector: Tensor or 2-tuple of Tensors. if self._has_bias, Tensor of shape - [input_size, output_size] corresponding to layer's weights. If not, a - 2-tuple of the former and a Tensor of shape [output_size] corresponding - to the layer's bias. + def register_matpower(self, exp): + # Not needed for this. Matrix powers are computed on demand in the + # diagonal case + pass - Returns: - 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) - return utils.mat2d_to_layer_params(vector, reshaped_out) - - def multiply(self, vector): - """Approximate damped Fisher-vector product. + def multiply_matpower(self, vector, exp): + """Multiplies the vector by the (damped) matrix-power of the block. Args: vector: Tensor or 2-tuple of Tensors. if self._has_bias, Tensor of shape [input_size, output_size] corresponding to layer's weights. If not, a 2-tuple of the former and a Tensor of shape [output_size] corresponding to the layer's bias. + exp: A scalar representing the power to raise the block before multiplying + it by the vector. Returns: - Tensor of the same shape, corresponding to the Fisher-vector product. + The vector left-multiplied by the (damped) matrix-power of the block. """ - 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_matpower( + reshaped_vec, exp, self._damping_func) return utils.mat2d_to_layer_params(vector, reshaped_out) - def tensors_to_compute_grads(self): - """Tensors to compute derivative of loss with respect to.""" - 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): - result = len(self._inputs) - assert result == len(self._outputs) - return result - - -class ConvDiagonalFB(FisherBlock): - """FisherBlock for convolutional layers using a diagonal approx. +class ConvDiagonalFB(InputOutputMultiMinibatch, FisherBlock): + """FisherBlock for 2-D convolutional layers using a diagonal approx. Estimates the Fisher Information matrix's diagonal entries for a convolutional layer. Unlike NaiveDiagonalFB this uses the low-variance "sum of squares" @@ -423,7 +543,13 @@ class ConvDiagonalFB(FisherBlock): to the layer's parameters 'w'. """ - def __init__(self, layer_collection, params, strides, padding): + def __init__(self, + layer_collection, + params, + strides, + padding, + data_format=None, + dilations=None): """Creates a ConvDiagonalFB block. Args: @@ -435,65 +561,72 @@ class ConvDiagonalFB(FisherBlock): containing the previous and a Tensor of shape [out_channels]. strides: The stride size in this layer (1-D Tensor of length 4). padding: The padding in this layer (e.g. "SAME"). + data_format: str or None. Format of input data. + dilations: List of 4 ints or None. Rate for dilation along all dimensions. + + Raises: + ValueError: if strides is not length-4. + ValueError: if dilations is not length-4. + ValueError: if channel is not last dimension. """ - self._inputs = [] - self._outputs = [] - self._strides = tuple(strides) if isinstance(strides, list) else strides + if len(strides) != 4: + raise ValueError("strides must contain 4 numbers.") + + if dilations is None: + dilations = [1, 1, 1, 1] + + if len(dilations) != 4: + raise ValueError("dilations must contain 4 numbers.") + + if not utils.is_data_format_channel_last(data_format): + raise ValueError("data_format must be channels-last.") + + self._strides = maybe_tuple(strides) self._padding = padding + self._data_format = data_format + self._dilations = maybe_tuple(dilations) self._has_bias = isinstance(params, (tuple, list)) fltr = params[0] if self._has_bias else params self._filter_shape = tuple(fltr.shape.as_list()) + if len(self._filter_shape) != 4: + raise ValueError( + "Convolution filter must be of shape" + " [filter_height, filter_width, in_channels, out_channels].") + super(ConvDiagonalFB, self).__init__(layer_collection) def instantiate_factors(self, grads_list, damping): - # Concatenate inputs, grads_list into single Tensors. - inputs = _concat_along_batch_dim(self._inputs) - grads_list = tuple(_concat_along_batch_dim(grads) for grads in grads_list) + inputs, grads_list = self._package_minibatches(grads_list) # Infer number of locations upon which convolution is applied. - inputs_shape = tuple(inputs.shape.as_list()) - self._num_locations = ( - inputs_shape[1] * inputs_shape[2] // - (self._strides[1] * self._strides[2])) - - self._damping = (self._num_locations - * normalize_damping(damping, self._num_locations)) + self._num_locations = num_conv_locations(inputs.shape.as_list(), + self._strides) self._factor = self._layer_collection.make_or_get_factor( fisher_factors.ConvDiagonalFactor, (inputs, grads_list, self._filter_shape, self._strides, self._padding, - self._has_bias)) + self._data_format, self._dilations, self._has_bias)) - def multiply_inverse(self, vector): - reshaped_vect = utils.layer_params_to_mat2d(vector) - reshaped_out = reshaped_vect / (self._factor.get_cov() + 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) - return utils.mat2d_to_layer_params(vector, reshaped_out) + def damping_func(): + return self._num_locations * normalize_damping(damping, + self._num_locations) - def tensors_to_compute_grads(self): - return self._outputs + damping_id = (self._num_locations, "mult", "normalize_damping", damping, + self._num_locations) + self._damping_func = _package_func(damping_func, damping_id) - def register_additional_minibatch(self, inputs, outputs): - """Registers an additional minibatch to the FisherBlock. - - Args: - inputs: Tensor of shape [batch_size, height, width, input_size]. Inputs to - the convolution. - outputs: Tensor of shape [batch_size, height, width, output_size]. Layer - preactivations. - """ - self._inputs.append(inputs) - self._outputs.append(outputs) + def register_matpower(self, exp): + # Not needed for this. Matrix powers are computed on demand in the + # diagonal case + pass - @property - def num_registered_minibatches(self): - return len(self._inputs) + def multiply_matpower(self, vector, exp): + reshaped_vect = utils.layer_params_to_mat2d(vector) + reshaped_out = self._factor.left_multiply_matpower( + reshaped_vect, exp, self._damping_func) + return utils.mat2d_to_layer_params(vector, reshaped_out) class KroneckerProductFB(FisherBlock): @@ -503,22 +636,40 @@ class KroneckerProductFB(FisherBlock): output factors. """ - def _register_damped_input_and_output_inverses(self, damping): - """Registers damped inverses for both the input and output factors. + def __init__(self, layer_collection): + super(KroneckerProductFB, self).__init__(layer_collection) + + def _setup_damping(self, damping, normalization=None): + """Makes functions that compute the damping values for both factors.""" + def compute_damping(): + if normalization is not None: + maybe_normalized_damping = normalize_damping(damping, normalization) + else: + maybe_normalized_damping = damping + + return compute_pi_adjusted_damping(self._input_factor.get_cov(), + self._output_factor.get_cov(), + maybe_normalized_damping**0.5) + + if normalization is not None: + damping_id = ("compute_pi_adjusted_damping", + "cov", self._input_factor.name, + "cov", self._output_factor.name, + "normalize_damping", damping, normalization, "power", 0.5) + else: + damping_id = ("compute_pi_adjusted_damping", + "cov", self._input_factor.name, + "cov", self._output_factor.name, + damping, "power", 0.5) - Sets the instance members _input_damping and _output_damping. Requires the - instance members _input_factor and _output_factor. + self._input_damping_func = _package_func(lambda: compute_damping()[0], + damping_id + ("ref", 0)) + self._output_damping_func = _package_func(lambda: compute_damping()[1], + damping_id + ("ref", 1)) - 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_factor.get_cov(), - self._output_factor.get_cov(), - damping**0.5) - - self._input_factor.register_damped_inverse(self._input_damping) - self._output_factor.register_damped_inverse(self._output_damping) + def register_matpower(self, exp): + self._input_factor.register_matpower(exp, self._input_damping_func) + self._output_factor.register_matpower(exp, self._output_damping_func) @property def _renorm_coeff(self): @@ -532,32 +683,15 @@ 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) + def multiply_matpower(self, vector, exp): 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_matpower( + reshaped_vector, exp, self._output_damping_func) + reshaped_out = self._input_factor.left_multiply_matpower( + reshaped_out, exp, self._input_damping_func) 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) - if self._renorm_coeff != 1.0: - reshaped_out *= math_ops.cast( - self._renorm_coeff, dtype=reshaped_out.dtype) + renorm_coeff = math_ops.cast(self._renorm_coeff, dtype=reshaped_out.dtype) + reshaped_out *= math_ops.cast(renorm_coeff**exp, dtype=reshaped_out.dtype) return utils.mat2d_to_layer_params(vector, reshaped_out) def full_fisher_block(self): @@ -574,7 +708,50 @@ class KroneckerProductFB(FisherBlock): right_factor) -class FullyConnectedKFACBasicFB(KroneckerProductFB): +class EmbeddingKFACFB(InputOutputMultiMinibatch, 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._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. + """ + inputs, grads_list = self._package_minibatches(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._setup_damping(damping) + + +class FullyConnectedKFACBasicFB(InputOutputMultiMinibatch, KroneckerProductFB): """K-FAC FisherBlock for fully-connected (dense) layers. This uses the Kronecker-factorized approximation from the original @@ -590,8 +767,6 @@ class FullyConnectedKFACBasicFB(KroneckerProductFB): has_bias: Whether the component Kronecker factors have an additive bias. (Default: False) """ - self._inputs = [] - self._outputs = [] self._has_bias = has_bias super(FullyConnectedKFACBasicFB, self).__init__(layer_collection) @@ -606,12 +781,7 @@ class FullyConnectedKFACBasicFB(KroneckerProductFB): 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) + inputs, grads_list = self._package_minibatches(grads_list) self._input_factor = self._layer_collection.make_or_get_factor( # fisher_factors.FullyConnectedKroneckerFactor, # @@ -619,29 +789,11 @@ class FullyConnectedKFACBasicFB(KroneckerProductFB): 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 + self._setup_damping(damping) - 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 ConvKFCBasicFB(KroneckerProductFB): - """FisherBlock for 2D convolutional layers using the basic KFC approx. +class ConvKFCBasicFB(InputOutputMultiMinibatch, KroneckerProductFB): + """FisherBlock for convolutional layers using the basic KFC approx. Estimates the Fisher Information matrix's blog for a convolutional layer. @@ -664,23 +816,40 @@ class ConvKFCBasicFB(KroneckerProductFB): See equation 23 in https://arxiv.org/abs/1602.01407 for details. """ - def __init__(self, layer_collection, params, strides, padding): + def __init__(self, + layer_collection, + params, + padding, + strides=None, + dilation_rate=None, + data_format=None, + extract_patches_fn=None): """Creates a ConvKFCBasicFB block. Args: layer_collection: The collection of all layers in the K-FAC approximate Fisher information matrix to which this FisherBlock belongs. params: The parameters (Tensor or tuple of Tensors) of this layer. If - kernel alone, a Tensor of shape [kernel_height, kernel_width, + kernel alone, a Tensor of shape [..spatial_filter_shape.., in_channels, out_channels]. If kernel and bias, a tuple of 2 elements containing the previous and a Tensor of shape [out_channels]. - strides: The stride size in this layer (1-D Tensor of length 4). - padding: The padding in this layer (1-D of Tensor length 4). + padding: str. Padding method. + strides: List of ints or None. Contains [..spatial_filter_strides..] if + 'extract_patches_fn' is compatible with tf.nn.convolution(), else + [1, ..spatial_filter_strides, 1]. + dilation_rate: List of ints or None. Rate for dilation along each spatial + dimension if 'extract_patches_fn' is compatible with + tf.nn.convolution(), else [1, ..spatial_dilation_rates.., 1]. + data_format: str or None. Format of input data. + extract_patches_fn: str or None. Name of function that extracts image + patches. One of "extract_convolution_patches", "extract_image_patches", + "extract_pointwise_conv2d_patches". """ - self._inputs = [] - self._outputs = [] - self._strides = tuple(strides) if isinstance(strides, list) else strides self._padding = padding + self._strides = maybe_tuple(strides) + self._dilation_rate = maybe_tuple(dilation_rate) + self._data_format = data_format + self._extract_patches_fn = extract_patches_fn self._has_bias = isinstance(params, (tuple, list)) fltr = params[0] if self._has_bias else params @@ -689,135 +858,314 @@ class ConvKFCBasicFB(KroneckerProductFB): super(ConvKFCBasicFB, self).__init__(layer_collection) def instantiate_factors(self, grads_list, damping): - # 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) + inputs, grads_list = self._package_minibatches(grads_list) # Infer number of locations upon which convolution is applied. - self._num_locations = num_conv_locations(inputs.shape.as_list(), + self._num_locations = num_conv_locations(self._inputs[0].shape.as_list(), self._strides) self._input_factor = self._layer_collection.make_or_get_factor( fisher_factors.ConvInputKroneckerFactor, - (inputs, self._filter_shape, self._strides, self._padding, + (inputs, self._filter_shape, self._padding, self._strides, + self._dilation_rate, self._data_format, self._extract_patches_fn, self._has_bias)) self._output_factor = self._layer_collection.make_or_get_factor( fisher_factors.ConvOutputKroneckerFactor, (grads_list,)) - damping = normalize_damping(damping, self._num_locations) - self._register_damped_input_and_output_inverses(damping) - self._damping = damping + self._setup_damping(damping, normalization=self._num_locations) @property def _renorm_coeff(self): return self._num_locations - def tensors_to_compute_grads(self): - return self._outputs - def register_additional_minibatch(self, inputs, outputs): - """Registers an additional minibatch to the FisherBlock. +class DepthwiseConvDiagonalFB(ConvDiagonalFB): + """FisherBlock for depthwise_conv2d(). + + Equivalent to ConvDiagonalFB applied to each input channel in isolation. + """ + + def __init__(self, + layer_collection, + params, + strides, + padding, + rate=None, + data_format=None): + """Creates a DepthwiseConvKFCBasicFB block. Args: - inputs: Tensor of shape [batch_size, height, width, input_size]. Inputs to - the convolution. - outputs: Tensor of shape [batch_size, height, width, output_size]. Layer - preactivations. + layer_collection: The collection of all layers in the K-FAC approximate + Fisher information matrix to which this FisherBlock belongs. + params: Tensor of shape [filter_height, filter_width, in_channels, + channel_multiplier]. + strides: List of 4 ints. Strides along all dimensions. + padding: str. Padding method. + rate: List of 4 ints or None. Rate for dilation along all dimensions. + data_format: str or None. Format of input data. + + Raises: + NotImplementedError: If parameters contains bias. + ValueError: If filter is not 4-D. + ValueError: If strides is not length-4. + ValueError: If rates is not length-2. + ValueError: If channels are not last dimension. """ - self._inputs.append(inputs) - self._outputs.append(outputs) + if isinstance(params, (tuple, list)): + raise NotImplementedError("Bias not yet supported.") - @property - def num_registered_minibatches(self): - return len(self._inputs) + if params.shape.ndims != 4: + raise ValueError("Filter must be 4-D.") + + if len(strides) != 4: + raise ValueError("strides must account for 4 dimensions.") + + if rate is not None: + if len(rate) != 2: + raise ValueError("rate must only account for spatial dimensions.") + rate = [1, rate[0], rate[1], 1] # conv2d expects 4-element rate. + + if not utils.is_data_format_channel_last(data_format): + raise ValueError("data_format must be channels-last.") + + super(DepthwiseConvDiagonalFB, self).__init__( + layer_collection=layer_collection, + params=params, + strides=strides, + padding=padding, + dilations=rate, + data_format=data_format) + + # This is a hack to overwrite the same setting in ConvKFCBasicFB.__init__(). + filter_height, filter_width, in_channels, channel_multiplier = ( + params.shape.as_list()) + self._filter_shape = (filter_height, filter_width, in_channels, + in_channels * channel_multiplier) + + def multiply_matpower(self, vector, exp): + conv2d_vector = depthwise_conv2d_filter_to_conv2d_filter(vector) + conv2d_result = super(DepthwiseConvDiagonalFB, self).multiply_matpower( + conv2d_vector, exp) + return conv2d_filter_to_depthwise_conv2d_filter(conv2d_result) + + +class DepthwiseConvKFCBasicFB(ConvKFCBasicFB): + """FisherBlock for depthwise_conv2d(). + + Equivalent to ConvKFCBasicFB applied to each input channel in isolation. + """ + + def __init__(self, + layer_collection, + params, + strides, + padding, + rate=None, + data_format=None): + """Creates a DepthwiseConvKFCBasicFB block. + + Args: + layer_collection: The collection of all layers in the K-FAC approximate + Fisher information matrix to which this FisherBlock belongs. + params: Tensor of shape [filter_height, filter_width, in_channels, + channel_multiplier]. + strides: List of 4 ints. Strides along all dimensions. + padding: str. Padding method. + rate: List of 4 ints or None. Rate for dilation along all dimensions. + data_format: str or None. Format of input data. + + Raises: + NotImplementedError: If parameters contains bias. + ValueError: If filter is not 4-D. + ValueError: If strides is not length-4. + ValueError: If rates is not length-2. + ValueError: If channels are not last dimension. + """ + if isinstance(params, (tuple, list)): + raise NotImplementedError("Bias not yet supported.") + + if params.shape.ndims != 4: + raise ValueError("Filter must be 4-D.") + + if len(strides) != 4: + raise ValueError("strides must account for 4 dimensions.") + + if rate is not None: + if len(rate) != 2: + raise ValueError("rate must only account for spatial dimensions.") + rate = [1, rate[0], rate[1], 1] # conv2d expects 4-element rate. + if not utils.is_data_format_channel_last(data_format): + raise ValueError("data_format must be channels-last.") -def _concat_along_batch_dim(tensor_list): - """Concatenate tensors along batch (first) dimension. + super(DepthwiseConvKFCBasicFB, self).__init__( + layer_collection=layer_collection, + params=params, + padding=padding, + strides=strides, + dilation_rate=rate, + data_format=data_format, + extract_patches_fn="extract_image_patches") + + # This is a hack to overwrite the same setting in ConvKFCBasicFB.__init__(). + filter_height, filter_width, in_channels, channel_multiplier = ( + params.shape.as_list()) + self._filter_shape = (filter_height, filter_width, in_channels, + in_channels * channel_multiplier) + + def multiply_matpower(self, vector, exp): + conv2d_vector = depthwise_conv2d_filter_to_conv2d_filter(vector) + conv2d_result = super(DepthwiseConvKFCBasicFB, self).multiply_matpower( + conv2d_vector, exp) + return conv2d_filter_to_depthwise_conv2d_filter(conv2d_result) + + +def depthwise_conv2d_filter_to_conv2d_filter(filter, name=None): # pylint: disable=redefined-builtin + """Converts a convolution filter for use with conv2d. + + Transforms a filter for use with tf.nn.depthwise_conv2d() to one that's + compatible with tf.nn.conv2d(). Args: - tensor_list: list of Tensors or list of tuples of Tensors. + filter: Tensor of shape [height, width, in_channels, channel_multiplier]. + name: None or str. Name of Op. Returns: - Tensor or tuple of Tensors. + Tensor of shape [height, width, in_channels, out_channels]. - Raises: - ValueError: If 'tensor_list' is empty. + """ + with ops.name_scope(name, "depthwise_conv2d_filter_to_conv2d_filter", + [filter]): + filter = ops.convert_to_tensor(filter) + filter_height, filter_width, in_channels, channel_multiplier = ( + filter.shape.as_list()) + + results = [] + for i in range(in_channels): + # Slice out one in_channel's filter. Insert zeros around it to force it + # to affect that channel and that channel alone. + elements = [] + if i > 0: + elements.append( + array_ops.zeros( + [filter_height, filter_width, i, channel_multiplier])) + elements.append(filter[:, :, i:(i + 1), :]) + if i + 1 < in_channels: + elements.append( + array_ops.zeros([ + filter_height, filter_width, in_channels - (i + 1), + channel_multiplier + ])) + + # Concat along in_channel. + results.append( + array_ops.concat(elements, axis=-2, name="in_channel_%d" % i)) + + # Concat along out_channel. + return array_ops.concat(results, axis=-1, name="out_channel") + + +def conv2d_filter_to_depthwise_conv2d_filter(filter, name=None): # pylint: disable=redefined-builtin + """Converts a convolution filter for use with depthwise_conv2d. + + Transforms a filter for use with tf.nn.conv2d() to one that's + compatible with tf.nn.depthwise_conv2d(). Ignores all filters but those along + the diagonal. + + Args: + filter: Tensor of shape [height, width, in_channels, out_channels]. + name: None or str. Name of Op. + Returns: + Tensor of shape, + [height, width, in_channels, channel_multiplier] + + Raises: + ValueError: if out_channels is not evenly divisible by in_channels. """ - if not tensor_list: - raise ValueError( - "Cannot concatenate Tensors if there are no Tensors to concatenate.") - - if isinstance(tensor_list[0], (tuple, list)): - # [(tensor1a, tensor1b), - # (tensor2a, tensor2b), ...] --> (tensor_a, tensor_b) - return tuple( - array_ops.concat(tensors, axis=0) for tensors in zip(*tensor_list)) - else: - # [tensor1, tensor2] --> tensor - return array_ops.concat(tensor_list, axis=0) + with ops.name_scope(name, "conv2d_filter_to_depthwise_conv2d_filter", + [filter]): + filter = ops.convert_to_tensor(filter) + filter_height, filter_width, in_channels, out_channels = ( + filter.shape.as_list()) + + if out_channels % in_channels != 0: + raise ValueError("out_channels must be evenly divisible by in_channels.") + channel_multiplier = out_channels // in_channels + + results = [] + filter = array_ops.reshape(filter, [ + filter_height, filter_width, in_channels, in_channels, + channel_multiplier + ]) + for i in range(in_channels): + # Slice out output corresponding to the correct filter. + filter_slice = array_ops.reshape( + filter[:, :, i, i, :], + [filter_height, filter_width, 1, channel_multiplier]) + results.append(filter_slice) + + # Concat along out_channel. + return array_ops.concat(results, axis=-2, name="in_channels") + + +def maybe_tuple(obj): + if not isinstance(obj, list): + return obj + return tuple(obj) 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. + input_shape: List of ints representing shape of inputs to + tf.nn.convolution(). + strides: List of ints representing strides along spatial dimensions as + passed in to tf.nn.convolution(). 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]) + spatial_input_locations = np.prod(input_shape[1:-1]) + if strides is None: + spatial_strides_divisor = 1 + else: + spatial_strides_divisor = np.prod(strides) + + return spatial_input_locations // spatial_strides_divisor -class FullyConnectedMultiIndepFB(KroneckerProductFB): + +class FullyConnectedMultiIndepFB(InputOutputMultiMinibatch, KroneckerProductFB): """FisherBlock for fully-connected layers that share parameters. """ - def __init__(self, layer_collection, inputs, outputs, has_bias=False): + def __init__(self, layer_collection, has_bias=False): """Creates a FullyConnectedMultiIndepFB block. Args: layer_collection: LayerCollection instance. - inputs: list or tuple of Tensors. Each Tensor has shape [batch_size, - inputs_size]. - outputs: list or tuple of Tensors. Each Tensor has shape [batch_size, - outputs_size]. has_bias: bool. If True, estimates Fisher with respect to a bias parameter as well as the layer's parameters. """ - - assert len(inputs) == len(outputs) - # We need to make sure inputs and outputs are tuples and not lists so that - # they get hashed by layer_collection.make_or_get_factor properly. - self._inputs = tuple(inputs) - self._outputs = tuple(outputs) self._has_bias = has_bias - self._num_uses = len(inputs) super(FullyConnectedMultiIndepFB, self).__init__(layer_collection) - @property - def num_registered_minibatches(self): - # TODO(b/69411207): Add support for registering additional minibatches. - return 1 - def instantiate_factors(self, grads_list, damping): + self._num_uses = float(len(self._inputs[0])) + inputs, grads_list = self._package_minibatches_multi(grads_list) + self._input_factor = self._layer_collection.make_or_get_factor( fisher_factors.FullyConnectedMultiKF, - ((self._inputs,), self._has_bias)) + ((inputs,), self._has_bias)) self._output_factor = self._layer_collection.make_or_get_factor( fisher_factors.FullyConnectedMultiKF, (grads_list,)) - damping = normalize_damping(damping, self._num_uses) - self._register_damped_input_and_output_inverses(damping) + self._setup_damping(damping, normalization=self._num_uses) @property def _renorm_coeff(self): @@ -826,9 +1174,6 @@ class FullyConnectedMultiIndepFB(KroneckerProductFB): def tensors_to_compute_grads(self): return self._outputs - def num_inputs(self): - return len(self._inputs) - class SeriesFBApproximation(enum.IntEnum): """See FullyConnectedSeriesFB.__init__ for description and usage.""" @@ -836,22 +1181,20 @@ class SeriesFBApproximation(enum.IntEnum): option2 = 2 -class FullyConnectedSeriesFB(FisherBlock): +class FullyConnectedSeriesFB(InputOutputMultiMinibatch, FisherBlock): """FisherBlock for fully-connected layers that share parameters across time. See the following preprint for details: https://openreview.net/pdf?id=HyMTkQZAb See the end of the appendix of the paper for a pseudo-code of the - algorithm being implemented by multiply_inverse here. Note that we are + algorithm being implemented by multiply_matpower here. Note that we are using pre-computed versions of certain matrix-matrix products to speed things up. This is explicitly explained wherever it is done. """ def __init__(self, layer_collection, - inputs, - outputs, has_bias=False, option=SeriesFBApproximation.option2): """Constructs a new `FullyConnectedSeriesFB`. @@ -859,10 +1202,6 @@ class FullyConnectedSeriesFB(FisherBlock): Args: layer_collection: The collection of all layers in the K-FAC approximate Fisher information matrix to which this FisherBlock belongs. - inputs: List of tensors of shape [batch_size, input_size]. - Inputs to the layer. - outputs: List of tensors of shape [batch_size, input_size]. - Outputs of the layer (before activations). has_bias: Whether the layer includes a bias parameter. option: A `SeriesFBApproximation` specifying the simplifying assumption to be used in this block. `option1` approximates the cross-covariance @@ -871,48 +1210,61 @@ class FullyConnectedSeriesFB(FisherBlock): 3.5 of the paper for more details. """ - assert len(inputs) == len(outputs) - # We need to make sure inputs and outputs are tuples and not lists so that - # they get hashed by layer_collection.make_or_get_factor properly. - self._inputs = tuple(inputs) - self._outputs = tuple(outputs) self._has_bias = has_bias - self._num_timesteps = len(inputs) self._option = option super(FullyConnectedSeriesFB, self).__init__(layer_collection) - @property - def num_registered_minibatches(self): - # TODO(b/69411207): Add support for registering additional minibatches. - return 1 - def instantiate_factors(self, grads_list, damping): + self._num_timesteps = len(self._inputs[0]) + inputs, grads_list = self._package_minibatches_multi(grads_list) + self._input_factor = self._layer_collection.make_or_get_factor( - fisher_factors.FullyConnectedMultiKF, ((self._inputs,), self._has_bias)) + fisher_factors.FullyConnectedMultiKF, ((inputs,), self._has_bias)) + self._input_factor.register_cov_dt1() self._output_factor = self._layer_collection.make_or_get_factor( fisher_factors.FullyConnectedMultiKF, (grads_list,)) - - 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) + self._output_factor.register_cov_dt1() + + def compute_damping(): + normalized_damping = normalize_damping(damping, self._num_timesteps) + return compute_pi_adjusted_damping(self._input_factor.get_cov(), + self._output_factor.get_cov(), + normalized_damping**0.5) + + damping_id = ("compute_pi_adjusted_damping", + "cov", self._input_factor.name, + "cov", self._output_factor.name, + "normalize_damping", + damping, self._num_timesteps, "power", 0.5) + self._input_damping_func = _package_func(lambda: compute_damping()[0], + damping_id + ("ref", 0)) + self._output_damping_func = _package_func(lambda: compute_damping()[1], + damping_id + ("ref", 1)) + + def register_matpower(self, exp): + if exp != -1: + raise NotImplementedError("FullyConnectedSeriesFB only supports inverse" + "multiplications.") if self._option == SeriesFBApproximation.option1: - self._input_factor.register_option1quants(self._damping_input) - self._output_factor.register_option1quants(self._damping_output) + self._input_factor.register_option1quants(self._input_damping_func) + self._output_factor.register_option1quants(self._output_damping_func) elif self._option == SeriesFBApproximation.option2: - self._input_factor.register_option2quants(self._damping_input) - self._output_factor.register_option2quants(self._damping_output) + self._input_factor.register_option2quants(self._input_damping_func) + self._output_factor.register_option2quants(self._output_damping_func) else: raise ValueError( "Unrecognized FullyConnectedSeriesFB approximation: {}".format( self._option)) - def multiply_inverse(self, vector): + def multiply_matpower(self, vector, exp): + if exp != -1: + raise NotImplementedError("FullyConnectedSeriesFB only supports inverse" + "multiplications.") + # pylint: disable=invalid-name Z = utils.layer_params_to_mat2d(vector) @@ -924,8 +1276,10 @@ class FullyConnectedSeriesFB(FisherBlock): if self._option == SeriesFBApproximation.option1: # Note that L_A = A0^(-1/2) * U_A and L_G = G0^(-1/2) * U_G. - L_A, psi_A = self._input_factor.get_option1quants(self._damping_input) - L_G, psi_G = self._output_factor.get_option1quants(self._damping_output) + L_A, psi_A = self._input_factor.get_option1quants( + self._input_damping_func) + L_G, psi_G = self._output_factor.get_option1quants( + self._output_damping_func) def gamma(x): # We are assuming that each case has the same number of time-steps. @@ -962,9 +1316,10 @@ class FullyConnectedSeriesFB(FisherBlock): # Note that P_A = A_1^T * A_0^(-1) and P_G = G_1^T * G_0^(-1), # and K_A = A_0^(-1/2) * E_A and K_G = G_0^(-1/2) * E_G. - P_A, K_A, mu_A = self._input_factor.get_option2quants(self._damping_input) + P_A, K_A, mu_A = self._input_factor.get_option2quants( + self._input_damping_func) P_G, K_G, mu_G = self._output_factor.get_option2quants( - self._damping_output) + self._output_damping_func) # Our approach differs superficially from the pseudo-code in the paper # in order to reduce the total number of matrix-matrix multiplies. @@ -1018,11 +1373,5 @@ class FullyConnectedSeriesFB(FisherBlock): # pylint: enable=invalid-name - def multiply(self, vector): - raise NotImplementedError - def tensors_to_compute_grads(self): return self._outputs - - def num_inputs(self): - return len(self._inputs) diff --git a/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py b/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py index ac396309206fe09af65c2b70840a513fb25b579b..c04cf727fa958160d61c7a3638ec65f6c93c2f24 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py @@ -29,6 +29,7 @@ _allowed_symbols = [ 'NaiveDiagonalFB', 'FullyConnectedDiagonalFB', 'KroneckerProductFB', + 'EmbeddingKFACFB', 'FullyConnectedKFACBasicFB', 'ConvKFCBasicFB', 'ConvDiagonalFB', @@ -36,7 +37,9 @@ _allowed_symbols = [ 'compute_pi_tracenorm', 'compute_pi_adjusted_damping', 'num_conv_locations', - 'normalize_damping' + '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 f59168cbc05fffd104ff5a44308eefd206beb9db..6fc163e2323666aca8489bf146ebc8582995cf06 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_factors.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_factors.py @@ -19,12 +19,12 @@ from __future__ import division from __future__ import print_function import abc -import contextlib 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 @@ -35,6 +35,7 @@ from tensorflow.python.ops import special_math_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.training import moving_averages +from tensorflow.python.util import nest # Whether to initialize covariance estimators at a zero matrix (or the identity # matrix). @@ -52,36 +53,16 @@ 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): - """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 - else: - with tf_ops.colocate_with(op): - yield - else: - yield - def set_global_constants(init_covariances_at_zero=None, zero_debias=None, eigenvalue_decomposition_threshold=None, - eigenvalue_clipping_threshold=None, - colocate_cov_ops_with_inputs=None): + eigenvalue_clipping_threshold=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 @@ -91,8 +72,6 @@ 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 @@ -180,7 +159,9 @@ def scope_string_from_params(params): name_parts = [] for param in params: - if isinstance(param, (tuple, list)): + if param is None: + name_parts.append("None") + elif isinstance(param, (tuple, list)): if all([isinstance(p, int) for p in param]): name_parts.append("-".join([str(p) for p in param])) else: @@ -189,6 +170,8 @@ def scope_string_from_params(params): name_parts.append(str(param)) elif isinstance(param, (tf_ops.Tensor, variables.Variable)): name_parts.append(scope_string_from_name(param)) + elif isinstance(param, utils.PartitionedTensor): + name_parts.append(scope_string_from_name(param.tensors)) else: raise ValueError("Encountered an unsupported param type {}".format( type(param))) @@ -206,33 +189,64 @@ def scalar_or_tensor_to_string(val): return repr(val) if np.isscalar(val) else scope_string_from_name(val) +def list_to_string(lst): + return "_".join(val if isinstance(val, six.string_types) + else scalar_or_tensor_to_string(val) for val in lst) + + +def graph_func_to_id(func): + """Returns a hashable object that represents func's computation.""" + # TODO(b/74201126): replace with Topohash of func's output + return func.func_id + + +def graph_func_to_string(func): + # TODO(b/74201126): replace with Topohash of func's output + return list_to_string(func.func_id) + + @six.add_metaclass(abc.ABCMeta) 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. """ def __init__(self): - self.instantiate_covariance() + self._cov = None @abc.abstractproperty def _var_scope(self): + """Variable scope for this FisherFactor instance. + + Returns: + string that unique identifies this FisherFactor instance. + """ pass + @property + def name(self): + return self._var_scope + @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 @@ -244,14 +258,17 @@ 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): - """Instantiates the covariance Variable as the instance member _cov.""" + def instantiate_cov_variables(self): + """Makes the internal cov variable(s).""" + assert self._cov is None with variable_scope.variable_scope(self._var_scope): self._cov = variable_scope.get_variable( "cov", @@ -262,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): @@ -274,34 +300,95 @@ class FisherFactor(object): """ 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) + 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 instantiate_inv_variables(self): + """Makes the internal "inverse" variable(s).""" + pass @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_matpower(self, x, exp, damping_func): + """Left multiplies 'x' by matrix power of this factor (w/ damping applied). + + This calculation is essentially: + (C + damping * I)**exp * x + where * is matrix-multiplication, ** is matrix power, I is the identity + matrix, and C is the matrix represented by this factor. + + x can represent either a matrix or a vector. For some factors, 'x' might + represent a vector but actually be stored as a 2D matrix for convenience. + + Args: + x: Tensor. Represents a single vector. Shape depends on implementation. + exp: float. The matrix exponent to use. + damping_func: A function that computes a 0-D Tensor or a float which will + be the damping value used. i.e. damping = damping_func(). + + Returns: + Tensor of same shape as 'x' representing the result of the multiplication. + """ + pass + + @abc.abstractmethod + def right_multiply_matpower(self, x, exp, damping_func): + """Right multiplies 'x' by matrix power of this factor (w/ damping applied). + + This calculation is essentially: + x * (C + damping * I)**exp + where * is matrix-multiplication, ** is matrix power, I is the identity + matrix, and C is the matrix represented by this factor. + + Unlike left_multiply_matpower, x will always be a matrix. + + Args: + x: Tensor. Represents a single vector. Shape depends on implementation. + exp: float. The matrix exponent to use. + damping_func: A function that computes a 0-D Tensor or a float which will + be the damping value used. i.e. damping = damping_func(). + + Returns: + Tensor of same shape as 'x' representing the result of the multiplication. + """ + 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. @@ -315,47 +402,52 @@ class InverseProvidingFactor(FisherFactor): # the latter. def __init__(self): - self._inverses_by_damping = {} - self._matpower_by_exp_and_damping = {} + self._matpower_by_exp_and_damping = {} # { (float, hashable): variable } + self._matpower_registrations = set() # { (float, hashable) } self._eigendecomp = None + self._damping_funcs_by_id = {} # {hashable: lambda} super(InverseProvidingFactor, self).__init__() - def register_damped_inverse(self, damping): - """Registers a damped inverse needed by a FisherBlock. - - This creates a variable and signals make_inverse_update_ops to make the - corresponding update op. The variable can be read via the method - get_inverse. + def _register_damping(self, damping_func): + damping_id = graph_func_to_id(damping_func) + if damping_id not in self._damping_funcs_by_id: + self._damping_funcs_by_id[damping_id] = damping_func + return damping_id - Args: - damping: The damping value (float or Tensor) for this factor. - """ - if damping not in self._inverses_by_damping: - damping_string = scalar_or_tensor_to_string(damping) - with variable_scope.variable_scope(self._var_scope): - inv = variable_scope.get_variable( - "inv_damp{}".format(damping_string), - initializer=inverse_initializer, - shape=self._cov_shape, - trainable=False, - dtype=self._dtype) - self._inverses_by_damping[damping] = inv + def register_inverse(self, damping_func): + # Just for backwards compatibility of some old code and tests + self.register_matpower(-1, damping_func) - def register_matpower(self, exp, damping): - """Registers a matrix power needed by a FisherBlock. + def register_matpower(self, exp, damping_func): + """Registers a matrix power to be maintained and served on demand. This creates a variable and signals make_inverse_update_ops to make the corresponding update op. The variable can be read via the method get_matpower. Args: - exp: The exponent (float or Tensor) to raise the matrix to. - damping: The damping value (float or Tensor). + exp: float. The exponent to use in the matrix power. + damping_func: A function that computes a 0-D Tensor or a float which will + be the damping value used. i.e. damping = damping_func(). """ - if (exp, damping) not in self._matpower_by_exp_and_damping: + if exp == 1.0: + # We don't register these. The user shouldn't even be calling this + # function with exp = 1.0. + return + + damping_id = self._register_damping(damping_func) + + if (exp, damping_id) not in self._matpower_registrations: + self._matpower_registrations.add((exp, damping_id)) + + def instantiate_inv_variables(self): + """Makes the internal "inverse" variable(s).""" + + for (exp, damping_id) in self._matpower_registrations: exp_string = scalar_or_tensor_to_string(exp) - damping_string = scalar_or_tensor_to_string(damping) + damping_func = self._damping_funcs_by_id[damping_id] + damping_string = graph_func_to_string(damping_func) with variable_scope.variable_scope(self._var_scope): matpower = variable_scope.get_variable( "matpower_exp{}_damp{}".format(exp_string, damping_string), @@ -363,34 +455,35 @@ class InverseProvidingFactor(FisherFactor): shape=self._cov_shape, trainable=False, dtype=self._dtype) - self._matpower_by_exp_and_damping[(exp, damping)] = matpower + assert (exp, damping_id) not in self._matpower_by_exp_and_damping + self._matpower_by_exp_and_damping[(exp, damping_id)] = matpower 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 = sum(1 for (exp, _) in self._matpower_by_exp_and_damping + if exp == -1) + + num_other_matpower = len(self._matpower_by_exp_and_damping) - num_inverses + + other_matrix_power_registered = num_other_matpower >= 1 - num_inverses = len(self._inverses_by_damping) - matrix_power_registered = bool(self._matpower_by_exp_and_damping) use_eig = ( - self._eigendecomp or matrix_power_registered or + self._eigendecomp or other_matrix_power_registered or num_inverses >= EIGENVALUE_DECOMPOSITION_THRESHOLD) + # We precompute these so we don't need to evaluate them multiple times (for + # each matrix power that uses them) + damping_value_by_id = {damping_id: self._damping_funcs_by_id[damping_id]() + for damping_id in self._damping_funcs_by_id} + if use_eig: eigenvalues, eigenvectors = self.get_eigendecomp() # pylint: disable=unpacking-non-sequence - for damping, inv in self._inverses_by_damping.items(): - ops.append( - inv.assign( - math_ops.matmul(eigenvectors / (eigenvalues + damping), - array_ops.transpose(eigenvectors)))) - - for (exp, damping), matpower in self._matpower_by_exp_and_damping.items(): + for (exp, damping_id), matpower in ( + self._matpower_by_exp_and_damping.items()): + damping = damping_value_by_id[damping_id] ops.append( matpower.assign( math_ops.matmul(eigenvectors * @@ -399,28 +492,31 @@ class InverseProvidingFactor(FisherFactor): # These ops share computation and should be run on a single device. ops = [control_flow_ops.group(*ops)] else: - for damping, inv in self._inverses_by_damping.items(): - ops.append(inv.assign(utils.posdef_inv(self._cov, damping))) + for (exp, damping_id), matpower in ( + self._matpower_by_exp_and_damping.items()): + assert exp == -1 + damping = damping_value_by_id[damping_id] + ops.append(matpower.assign(utils.posdef_inv(self._cov, damping))) + self._eigendecomp = False return ops - def get_damped_inverse(self, damping): - # Note that this function returns a variable which gets updated by the - # inverse ops. It may be stale / inconsistent with the latest value of - # get_cov(). - return self._inverses_by_damping[damping] + def get_inverse(self, damping_func): + # Just for backwards compatibility of some old code and tests + damping_id = graph_func_to_id(damping_func) + return self._matpower_by_exp_and_damping[(-1, damping_id)] - def get_matpower(self, exp, damping): + def get_matpower(self, exp, damping_func): # Note that this function returns a variable which gets updated by the # inverse ops. It may be stale / inconsistent with the latest value of # get_cov(). - return self._matpower_by_exp_and_damping[(exp, damping)] + damping_id = graph_func_to_id(damping_func) + return self._matpower_by_exp_and_damping[(exp, damping_id)] 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(). + # Unlike 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) @@ -433,8 +529,42 @@ class InverseProvidingFactor(FisherFactor): 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_matpower(self, x, exp, damping_func): + 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,)) + + if exp == 1: + return math_ops.matmul(self.get_cov(), x) + damping_func() * x + + return math_ops.matmul(self.get_matpower(exp, damping_func), x) + + def right_multiply_matpower(self, x, exp, damping_func): + if isinstance(x, tf_ops.IndexedSlices): + if exp == 1: + n = self.get_cov().shape[0] + damped_cov = self.get_cov() + damping_func() * array_ops.eye(n) + return utils.matmul_sparse_dense(x, damped_cov) + + return utils.matmul_sparse_dense(x, self.get_matpower(exp, damping_func)) + + if x.shape.ndims != 2: + raise ValueError( + "InverseProvidingFactors apply to matrix-shaped vectors. Found: %s." + % (x,)) + + if exp == 1: + return math_ops.matmul(x, self.get_cov()) + damping_func() * x + + return math_ops.matmul(x, self.get_matpower(exp, damping_func)) class FullFactor(InverseProvidingFactor): @@ -454,7 +584,7 @@ class FullFactor(InverseProvidingFactor): @property def _var_scope(self): - return "ff_full/" + scope_string_from_params( + return "ff_full_" + scope_string_from_params( [self._params_grads, self._batch_size]) @property @@ -473,17 +603,21 @@ class FullFactor(InverseProvidingFactor): def _compute_new_cov(self, idx=0): # This will be a very basic rank 1 estimate - 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)) + 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): + self._damping_funcs_by_id = {} # { hashable: lambda } super(DiagonalFactor, self).__init__() @property @@ -493,6 +627,32 @@ class DiagonalFactor(FisherFactor): def make_inverse_update_ops(self): return [] + def instantiate_inv_variables(self): + pass + + 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_matpower(self, x, exp, damping_func): + matpower = (self.get_cov_var() + damping_func())**exp + + if isinstance(x, tf_ops.IndexedSlices): + return utils.matmul_diag_sparse(array_ops.reshape(matpower, [-1]), x) + + if x.shape != matpower.shape: + raise ValueError("x (%s) and cov (%s) must have same shape." % + (x, matpower)) + return matpower * x + + def right_multiply_matpower(self, x, exp, damping_func): + raise NotImplementedError("Only left-multiply is currently supported.") + + def register_matpower(self, exp, damping_func): + pass + class NaiveDiagonalFactor(DiagonalFactor): """FisherFactor for a diagonal approximation of any type of param's Fisher. @@ -504,6 +664,14 @@ class NaiveDiagonalFactor(DiagonalFactor): def __init__(self, params_grads, 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 @@ -511,14 +679,14 @@ class NaiveDiagonalFactor(DiagonalFactor): @property def _var_scope(self): - return "ff_naivediag/" + scope_string_from_params( + return "ff_naivediag_" + scope_string_from_params( [self._params_grads, self._batch_size]) @property def _cov_shape(self): size = sum(param_grad.shape.num_elements() for param_grad in self._params_grads[0]) - return (size, 1) + return [size, 1] @property def _num_sources(self): @@ -529,10 +697,90 @@ class NaiveDiagonalFactor(DiagonalFactor): return self._params_grads[0][0].dtype def _compute_new_cov(self, idx=0): - 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)) + 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: Tensor 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 1 + + @property + def _dtype(self): + return self._cov_dtype + + def _compute_new_cov(self, idx=0): + if idx != 0: + raise ValueError("EmbeddingInputKroneckerFactor only supports idx = 0") + + input_ids = self._input_ids + + 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): @@ -553,29 +801,30 @@ class FullyConnectedDiagonalFactor(DiagonalFactor): """Instantiate FullyConnectedDiagonalFactor. Args: - inputs: Tensor of shape [batch_size, input_size]. Inputs to fully - connected layer. - outputs_grads: List of Tensors of shape [batch_size, output_size]. - Gradient of loss with respect to layer's preactivations. + inputs: Tensor of shape [batch_size, input_size]. Inputs to this layer. + outputs_grads: List of Tensors, each of shape [batch_size, output_size], + which are the gradients of the loss with respect to the layer's + outputs. One Tensor for each "source". + has_bias: bool. If True, append '1' to each input. """ self._inputs = inputs self._has_bias = has_bias self._outputs_grads = outputs_grads - self._batch_size = array_ops.shape(inputs)[0] self._squared_inputs = None super(FullyConnectedDiagonalFactor, self).__init__() @property def _var_scope(self): - return "ff_diagfc/" + scope_string_from_params( + return "ff_diagfc_" + scope_string_from_params( (self._inputs,) + tuple(self._outputs_grads)) @property def _cov_shape(self): - return [self._inputs.shape[1] + self._has_bias, - 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): @@ -585,25 +834,30 @@ class FullyConnectedDiagonalFactor(DiagonalFactor): def _dtype(self): return self._outputs_grads[0].dtype + def make_covariance_update_op(self, ema_decay): + inputs = self._inputs + + if self._has_bias: + inputs = append_homog(inputs) + self._squared_inputs = math_ops.square(inputs) + + return super(FullyConnectedDiagonalFactor, self).make_covariance_update_op( + ema_decay) + def _compute_new_cov(self, idx=0): + batch_size = array_ops.shape(self._squared_inputs)[0] + outputs_grad = self._outputs_grads[idx] + # The well-known special formula that uses the fact that the entry-wise # 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._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]), - transpose_a=True) - new_cov /= math_ops.cast(self._batch_size, new_cov.dtype) - return new_cov + new_cov = math_ops.matmul( + self._squared_inputs, + math_ops.square(outputs_grad), + transpose_a=True) + new_cov /= math_ops.cast(batch_size, new_cov.dtype) + return new_cov class ConvDiagonalFactor(DiagonalFactor): @@ -615,35 +869,64 @@ class ConvDiagonalFactor(DiagonalFactor): filter_shape, strides, padding, + data_format=None, + dilations=None, has_bias=False): """Creates a ConvDiagonalFactor object. Args: inputs: Tensor of shape [batch_size, height, width, in_channels]. Input activations to this layer. - outputs_grads: Tensor of shape [batch_size, height, width, out_channels]. - Per-example gradients to the loss with respect to the layer's output - preactivations. + outputs_grads: List of Tensors, each of shape [batch_size, + height, width, out_channels], which are the gradients of the loss + with respect to the layer's outputs. One Tensor for each "source". filter_shape: Tuple of 4 ints: (kernel_height, kernel_width, in_channels, out_channels). Represents shape of kernel used in this layer. strides: The stride size in this layer (1-D Tensor of length 4). padding: The padding in this layer (1-D of Tensor length 4). + data_format: None or str. Format of conv2d inputs. + dilations: None or tuple of 4 ints. has_bias: Python bool. If True, the layer is assumed to have a bias parameter in addition to its filter parameter. + + Raises: + ValueError: If inputs, output_grads, and filter_shape do not agree on + in_channels or out_channels. + ValueError: If strides, dilations are not length-4 lists of ints. + ValueError: If data_format does not put channel last. """ + if not utils.is_data_format_channel_last(data_format): + raise ValueError("Channel must be last.") + if inputs.shape.ndims != 4: + raise ValueError("inputs must be 4-D Tensor.") + if inputs.shape.as_list()[-1] != filter_shape[-2]: + raise ValueError("inputs and filter_shape must agree on in_channels.") + for i, outputs_grad in enumerate(outputs_grads): + if outputs_grad.shape.ndims != 4: + raise ValueError("outputs[%d] must be 4-D Tensor." % i) + if outputs_grad.shape.as_list()[-1] != filter_shape[-1]: + raise ValueError( + "outputs[%d] and filter_shape must agree on out_channels." % i) + if len(strides) != 4: + raise ValueError("strides must be length-4 list of ints.") + if dilations is not None and len(dilations) != 4: + raise ValueError("dilations must be length-4 list of ints.") + self._inputs = inputs + self._outputs_grads = outputs_grads self._filter_shape = filter_shape self._strides = strides self._padding = padding + self._data_format = data_format + self._dilations = dilations self._has_bias = has_bias - self._outputs_grads = outputs_grads self._patches = None super(ConvDiagonalFactor, self).__init__() @property def _var_scope(self): - return "ff_convdiag/" + scope_string_from_name( + return "ff_convdiag_" + scope_string_from_params( (self._inputs,) + tuple(self._outputs_grads)) @property @@ -663,38 +946,36 @@ class ConvDiagonalFactor(DiagonalFactor): 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 + 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) + # TODO(b/64144716): there is potential here for a big savings in terms + # of memory use. + if self._dilations is None: + rates = (1, 1, 1, 1) + else: + rates = tuple(self._dilations) + patches = array_ops.extract_image_patches( + self._inputs, + ksizes=[1, filter_height, filter_width, 1], + strides=self._strides, + rates=rates, + padding=self._padding) - self._patches = patches + if self._has_bias: + patches = append_homog(patches) - op = super(ConvDiagonalFactor, self).make_covariance_update_op(ema_decay) + self._patches = patches - self._patches = None - - return op + return super(ConvDiagonalFactor, self).make_covariance_update_op(ema_decay) def _compute_new_cov(self, idx=0): - with maybe_colocate_with(self._outputs_grads[idx]): - outputs_grad = self._outputs_grads[idx] - batch_size = array_ops.shape(self._patches)[0] + batch_size = array_ops.shape(self._patches)[0] + outputs_grad = self._outputs_grads[idx] - new_cov = self._convdiag_sum_of_squares(self._patches, outputs_grad) - new_cov /= math_ops.cast(batch_size, new_cov.dtype) + new_cov = self._convdiag_sum_of_squares(self._patches, outputs_grad) + new_cov /= math_ops.cast(batch_size, new_cov.dtype) - return new_cov + return new_cov def _convdiag_sum_of_squares(self, patches, outputs_grad): # This computes the sum of the squares of the per-training-case "gradients". @@ -715,8 +996,9 @@ class FullyConnectedKroneckerFactor(InverseProvidingFactor): """Instantiate FullyConnectedKroneckerFactor. Args: - tensors: List of Tensors of shape [batch_size, n]. Represents either a - layer's inputs or its output's gradients. + tensors: List of Tensors, each of shape [batch_size, n], one for each + source. The Tensors are typically either a layer's inputs or its + output's gradients. has_bias: bool. If True, append '1' to each row. """ # The tensor argument is either a tensor of input activations or a tensor of @@ -727,8 +1009,8 @@ class FullyConnectedKroneckerFactor(InverseProvidingFactor): @property def _var_scope(self): - return "ff_fckron/" + scope_string_from_params( - [self._tensors, self._has_bias]) + return "ff_fckron_" + scope_string_from_params( + tuple(self._tensors) + (self._has_bias,)) @property def _cov_shape(self): @@ -744,11 +1026,10 @@ class FullyConnectedKroneckerFactor(InverseProvidingFactor): return self._tensors[0].dtype def _compute_new_cov(self, idx=0): - with maybe_colocate_with(self._tensors[idx]): - tensor = self._tensors[idx] - if self._has_bias: - tensor = append_homog(tensor) - return compute_cov(tensor) + tensor = self._tensors[idx] + if self._has_bias: + tensor = append_homog(tensor) + return compute_cov(tensor) class ConvInputKroneckerFactor(InverseProvidingFactor): @@ -764,39 +1045,55 @@ class ConvInputKroneckerFactor(InverseProvidingFactor): def __init__(self, inputs, filter_shape, - strides, padding, + strides=None, + dilation_rate=None, + data_format=None, + extract_patches_fn=None, has_bias=False): """Initializes ConvInputKroneckerFactor. Args: - inputs: Tensor of shape [batch_size, height, width, in_channels]. Inputs - to layer. - filter_shape: 1-D Tensor of length 4. Contains [kernel_height, - kernel_width, in_channels, out_channels]. - strides: 1-D Tensor of length 4. Contains [batch_stride, height_stride, - width_stride, in_channel_stride]. + inputs: Tensor of shape [batch_size, ..spatial_input_size.., in_channels]. + Inputs to layer. + filter_shape: List of ints. Contains [..spatial_filter_size.., + in_channels, out_channels]. Shape of convolution kernel. padding: str. Padding method for layer. "SAME" or "VALID". + strides: List of ints or None. Contains [..spatial_filter_strides..] if + 'extract_patches_fn' is compatible with tf.nn.convolution(), else + [1, ..spatial_filter_strides, 1]. + dilation_rate: List of ints or None. Rate for dilation along each spatial + dimension if 'extract_patches_fn' is compatible with + tf.nn.convolution(), else [1, ..spatial_dilation_rates.., 1]. + data_format: str or None. Format of input data. + extract_patches_fn: str or None. Name of function that extracts image + patches. One of "extract_convolution_patches", "extract_image_patches", + "extract_pointwise_conv2d_patches". has_bias: bool. If True, append 1 to in_channel. """ + self._inputs = inputs self._filter_shape = filter_shape self._strides = strides self._padding = padding + self._dilation_rate = dilation_rate + self._data_format = data_format + self._extract_patches_fn = extract_patches_fn self._has_bias = has_bias - self._inputs = inputs + super(ConvInputKroneckerFactor, self).__init__() @property def _var_scope(self): - return "ff_convinkron/" + scope_string_from_params([ + return "ff_convinkron_" + scope_string_from_params([ self._inputs, self._filter_shape, self._strides, self._padding, - self._has_bias + self._dilation_rate, self._data_format, self._has_bias ]) @property def _cov_shape(self): - filter_height, filter_width, in_channels, _ = self._filter_shape - size = filter_height * filter_width * in_channels + self._has_bias + spatial_filter_shape = self._filter_shape[0:-2] + in_channels = self._filter_shape[-2] + size = np.prod(spatial_filter_shape) * in_channels + self._has_bias return [size, size] @property @@ -811,37 +1108,62 @@ class ConvInputKroneckerFactor(InverseProvidingFactor): if idx != 0: raise ValueError("ConvInputKroneckerFactor only supports idx = 0") - 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 - # memory use. + # TODO(b/64144716): there is potential here for a big savings in terms of + # memory use. + if self._extract_patches_fn in [None, "extract_convolution_patches"]: + patches = utils.extract_convolution_patches( + self._inputs, + self._filter_shape, + padding=self._padding, + strides=self._strides, + dilation_rate=self._dilation_rate, + data_format=self._data_format) + + elif self._extract_patches_fn == "extract_image_patches": + assert self._inputs.shape.ndims == 4 + assert len(self._filter_shape) == 4 + assert len(self._strides) == 4, self._strides + if self._dilation_rate is None: + rates = [1, 1, 1, 1] + else: + rates = self._dilation_rate + assert len(rates) == 4 + assert rates[0] == rates[-1] == 1 patches = array_ops.extract_image_patches( self._inputs, - ksizes=[1, filter_height, filter_width, 1], + ksizes=[1] + list(self._filter_shape[0:-2]) + [1], strides=self._strides, - rates=[1, 1, 1, 1], + rates=rates, padding=self._padding) - flatten_size = (filter_height * filter_width * in_channels) - # patches_flat below is the matrix [[A_l]] from the KFC paper (tilde - # omitted over A for clarity). It has shape M|T| x J|Delta| (eq. 14), - # where M = minibatch size, |T| = number of spatial locations, - # |Delta| = number of spatial offsets, and J = number of input maps - # for convolutional layer l. - patches_flat = array_ops.reshape(patches, [-1, flatten_size]) - # 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 - # 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) + elif self._extract_patches_fn == "extract_pointwise_conv2d_patches": + assert self._strides in [None, [1, 1, 1, 1], (1, 1, 1, 1)] + assert self._filter_shape[0] == self._filter_shape[1] == 1 + patches = utils.extract_pointwise_conv2d_patches( + self._inputs, self._filter_shape, data_format=None) + + else: + raise NotImplementedError(self._extract_patches_fn) + + flatten_size = np.prod(self._filter_shape[0:-1]) + # patches_flat below is the matrix [[A_l]] from the KFC paper (tilde + # omitted over A for clarity). It has shape M|T| x J|Delta| (eq. 14), + # where M = minibatch size, |T| = number of spatial locations, + # |Delta| = number of spatial offsets, and J = number of input maps + # for convolutional layer l. + patches_flat = array_ops.reshape(patches, [-1, flatten_size]) + # 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 + # 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) class ConvOutputKroneckerFactor(InverseProvidingFactor): @@ -855,20 +1177,27 @@ class ConvOutputKroneckerFactor(InverseProvidingFactor): Section 3.1 Estimating the factors. """ - def __init__(self, outputs_grads): + def __init__(self, outputs_grads, data_format=None): """Initializes ConvOutputKroneckerFactor. Args: outputs_grads: list of Tensors. Each Tensor is of shape - [batch_size, height, width, out_channels]. + [batch_size, ..spatial_input_size.., out_channels]. One Tensor per + source. + data_format: None or str. Format of outputs_grads. + + Raises: + ValueError: If channels are not final dimension. """ - self._out_channels = outputs_grads[0].shape.as_list()[3] + if not utils.is_data_format_channel_last(data_format): + raise ValueError("Channel must be last.") + self._out_channels = outputs_grads[0].shape.as_list()[-1] self._outputs_grads = outputs_grads super(ConvOutputKroneckerFactor, self).__init__() @property def _var_scope(self): - return "ff_convoutkron/" + scope_string_from_params(self._outputs_grads) + return "ff_convoutkron_" + scope_string_from_params(self._outputs_grads) @property def _cov_shape(self): @@ -884,22 +1213,22 @@ class ConvOutputKroneckerFactor(InverseProvidingFactor): return self._outputs_grads[0].dtype def _compute_new_cov(self, idx=0): - 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 - # I = number of output maps for convolutional layer l. - 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 - # as defined in the paper, with shape I x I. - # (Tilde omitted over S for clarity.) - return compute_cov(reshaped_tensor) + outputs_grad = 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 + # I = number of output maps for convolutional layer l. + reshaped_tensor = array_ops.reshape(outputs_grad, [-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 + # as defined in the paper, with shape I x I. + # (Tilde omitted over S for clarity.) + return compute_cov(reshaped_tensor) class FullyConnectedMultiKF(InverseProvidingFactor): - """Kronecker factor for a fully connected recurrent layer.""" + """Kronecker factor for a fully connected layer used multiple times.""" def __init__(self, tensor_lists, @@ -907,25 +1236,32 @@ class FullyConnectedMultiKF(InverseProvidingFactor): """Constructs a new `FullyConnectedMultiKF`. Args: - tensor_lists: List of lists of Tensors of shape [batch_size, n]. + tensor_lists: 2D array (list of lists) of Tensors of shape + [batch_size, n]. Each of these tensors is usually a layer's inputs or + its output's gradients. The first dimension of the array is the source, + and the second is the use in the graph (which is sometimes a + "time-step"). has_bias: bool. If True, '1' is appended to each row. """ 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]) self._tensors = [None] * len(tensor_lists) self._cov_dt1 = None + self._make_cov_dt1 = False self._option1quants_by_damping = {} self._option2quants_by_damping = {} + self._option1quants_registrations = set() + self._option2quants_registrations = set() super(FullyConnectedMultiKF, self).__init__() @property def _var_scope(self): - return "ff_fc_multi/" + scope_string_from_params(self._tensor_lists) + return "ff_fc_multi_" + scope_string_from_params( + tuple(nest.flatten(self._tensor_lists)) + (self._has_bias,)) @property def _num_sources(self): @@ -942,43 +1278,40 @@ class FullyConnectedMultiKF(InverseProvidingFactor): 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)) + 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) - 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) - - # TODO(b/69112164): - # It's important that _cov and _cov_dt1 remain consistent with each - # other while the inverse ops are happening. How can we ensure this? - # We will need to add explicit synchronization for this to - # work with asynchronous training. - op = control_flow_ops.group(op, op2) + # TODO(b/69112164): + # It's important that _cov and _cov_dt1 remain consistent with each + # other while the inverse ops are happening. How can we ensure this? + # We will need to add explicit synchronization for this to + # work with asynchronous training. + op = control_flow_ops.group(op, op2) return op def _compute_new_cov(self, idx=0): - 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): + # Concatenate across time/replications + 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): # pylint: disable=missing-docstring tensor = self._tensors[idx] - 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) + batch_size = array_ops.shape(self._tensor_lists[idx][0])[0] + # Is there a more elegant way to do this computation? + tensor_present = tensor[:-batch_size, :] + tensor_future = tensor[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 * batch_size + return compute_cov( + tensor_future, tensor_right=tensor_present, normalizer=normalizer) @property def _cov_shape(self): @@ -990,23 +1323,25 @@ class FullyConnectedMultiKF(InverseProvidingFactor): size = self._tensor_lists[0][0].shape[1] + self._has_bias return [size] - def get_option1quants(self, damping): - return self._option1quants_by_damping[damping] + def get_option1quants(self, damping_func): + damping_id = graph_func_to_id(damping_func) + return self._option1quants_by_damping[damping_id] - def get_option2quants(self, damping): - return self._option2quants_by_damping[damping] + def get_option2quants(self, damping_func): + damping_id = graph_func_to_id(damping_func) + return self._option2quants_by_damping[damping_id] def get_cov_dt1(self): assert self._cov_dt1 is not None return self._cov_dt1 def register_cov_dt1(self): - """Create a variable representing temporal cross-covariance. + self._make_cov_dt1 = True - (This is technically the second moment, not covariance, since it's - not mean subtracted.) - """ - if self._cov_dt1 is None: + def instantiate_cov_variables(self): + super(FullyConnectedMultiKF, self).instantiate_cov_variables() + assert self._cov_dt1 is None + if self._make_cov_dt1: with variable_scope.variable_scope(self._var_scope): self._cov_dt1 = variable_scope.get_variable( "cov_dt1", @@ -1015,15 +1350,25 @@ class FullyConnectedMultiKF(InverseProvidingFactor): trainable=False, dtype=self._dtype) - def register_option1quants(self, damping): + def register_option1quants(self, damping_func): + damping_id = self._register_damping(damping_func) + if damping_id not in self._option1quants_registrations: + self._option1quants_registrations.add(damping_id) + + def register_option2quants(self, damping_func): + damping_id = self._register_damping(damping_func) + if damping_id not in self._option2quants_registrations: + self._option2quants_registrations.add(damping_id) - self.register_cov_dt1() + def instantiate_inv_variables(self): + super(FullyConnectedMultiKF, self).instantiate_inv_variables() - if damping not in self._option1quants_by_damping: + for damping_id in self._option1quants_registrations: + damping_func = self._damping_funcs_by_id[damping_id] + damping_string = graph_func_to_string(damping_func) # It's questionable as to whether we should initialize with stuff like # this at all. Ideally these values should never be used until they are # updated at least once. - damping_string = scalar_or_tensor_to_string(damping) with variable_scope.variable_scope(self._var_scope): Lmat = variable_scope.get_variable( # pylint: disable=invalid-name "Lmat_damp{}".format(damping_string), @@ -1038,17 +1383,15 @@ class FullyConnectedMultiKF(InverseProvidingFactor): trainable=False, dtype=self._dtype) - self._option1quants_by_damping[damping] = (Lmat, psi) + assert damping_id not in self._option1quants_by_damping + self._option1quants_by_damping[damping_id] = (Lmat, psi) - def register_option2quants(self, damping): - - self.register_cov_dt1() - - if damping not in self._option2quants_by_damping: + for damping_id in self._option2quants_registrations: + damping_func = self._damping_funcs_by_id[damping_id] + damping_string = graph_func_to_string(damping_func) # It's questionable as to whether we should initialize with stuff like # this at all. Ideally these values should never be used until they are # updated at least once. - damping_string = scalar_or_tensor_to_string(damping) with variable_scope.variable_scope(self._var_scope): Pmat = variable_scope.get_variable( # pylint: disable=invalid-name "Lmat_damp{}".format(damping_string), @@ -1069,14 +1412,15 @@ class FullyConnectedMultiKF(InverseProvidingFactor): trainable=False, dtype=self._dtype) - self._option2quants_by_damping[damping] = (Pmat, Kmat, mu) + assert damping_id not in self._option2quants_by_damping + self._option2quants_by_damping[damping_id] = (Pmat, Kmat, mu) def make_inverse_update_ops(self): """Create and return update ops corresponding to registered computations.""" # TODO(b/69918258): Add correctness tests for this method. # pylint: disable=invalid-name - ops = super(FullyConnectedMultiKF, self).make_inverse_update_ops() + ops = [] if (len(self._option1quants_by_damping) + len(self._option2quants_by_damping)): @@ -1097,8 +1441,10 @@ class FullyConnectedMultiKF(InverseProvidingFactor): # consistently, or are somehow read between or during the cov updates. # Can this possibly happen? Is there a way to prevent it? - for damping, (Lmat_var, - psi_var) in self._option1quants_by_damping.items(): + for damping_id, (Lmat_var, + psi_var) in self._option1quants_by_damping.items(): + + damping = self._damping_funcs_by_id[damping_id]() invsqrtC0 = math_ops.matmul( eigen_V * (eigen_e + damping)**(-0.5), eigen_V, transpose_b=True) @@ -1123,8 +1469,10 @@ class FullyConnectedMultiKF(InverseProvidingFactor): ops.append(Lmat_var.assign(Lmat)) ops.append(psi_var.assign(psi)) - for damping, (Pmat_var, Kmat_var, - mu_var) in self._option2quants_by_damping.items(): + for damping_id, (Pmat_var, Kmat_var, + mu_var) in self._option2quants_by_damping.items(): + + damping = self._damping_funcs_by_id[damping_id]() # compute C0^(-1/2) invsqrtC0 = math_ops.matmul( @@ -1165,6 +1513,7 @@ class FullyConnectedMultiKF(InverseProvidingFactor): ops.append(Kmat_var.assign(Kmat)) ops.append(mu_var.assign(mu)) + ops += super(FullyConnectedMultiKF, self).make_inverse_update_ops() return [control_flow_ops.group(*ops)] # pylint: enable=invalid-name diff --git a/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py b/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py index ad93919149c287b1932dd2b6bd772c0dab26192d..2d8e378a932c16d48360bc4b15ff4f3239c0ed1f 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py @@ -24,26 +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", - "maybe_colocate_with", - "compute_cov", - "append_homog" + "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 8d450f04f379701e46a18b2e34bbbd6fcfcce2bb..4eb5e4c092b50ff4a908a22312330c40ca93cbee 100644 --- a/tensorflow/contrib/kfac/python/ops/layer_collection.py +++ b/tensorflow/contrib/kfac/python/ops/layer_collection.py @@ -26,6 +26,7 @@ from __future__ import print_function from collections import defaultdict from collections import OrderedDict +from contextlib import contextmanager from functools import partial import math @@ -75,6 +76,27 @@ _FULLY_CONNECTED_MULTI_APPROX_TO_BLOCK_TYPES = { # tf.get_variable_scope().reuse. VARIABLE_SCOPE = "VARIABLE_SCOPE" +_DEFAULT_LAYER_COLLECTION = None + + +def get_default_layer_collection(): + """Get default LayerCollection.""" + if _DEFAULT_LAYER_COLLECTION is None: + raise ValueError( + "Attempted to retrieve default LayerCollection when none is set. Use " + "LayerCollection.as_default().") + + return _DEFAULT_LAYER_COLLECTION + + +def set_default_layer_collection(layer_collection): + global _DEFAULT_LAYER_COLLECTION + + if _DEFAULT_LAYER_COLLECTION is not None and layer_collection is not None: + raise ValueError("Default LayerCollection is already set.") + + _DEFAULT_LAYER_COLLECTION = layer_collection + class LayerParametersDict(OrderedDict): """An OrderedDict where keys are Tensors or tuples of Tensors. @@ -130,6 +152,8 @@ class LayerCollection(object): fisher_factors: an OrderedDict mapping tuples to FisherFactor instances. losses: a list of LossFunction objects. The loss to be optimized is their sum. + loss_colocation_ops: ops to colocate loss function evaluations with. These + will typically be the inputs to the losses. """ def __init__(self, @@ -143,18 +167,26 @@ 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.loss_colocation_ops = {} + self._vars_to_uses = defaultdict(lambda: 0) with variable_scope.variable_scope(None, default_name=name) as scope: self._var_scope = scope.name @property def losses(self): - """LossFunctions registered with this LayerCollection.""" - return list(self._loss_dict.values()) + """Tuple of LossFunction objects registered with this LayerCollection.""" + return nest.flatten(self.towers_by_loss) + + @property + def towers_by_loss(self): + """Tuple across losses of LossFunction objects registered to each tower.""" + return tuple(tuple(lst) for lst in self._loss_dict.values()) @property def registered_variables(self): @@ -178,6 +210,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 @@ -278,23 +321,74 @@ class LayerCollection(object): self.fisher_blocks[layer_key] = fisher_block return fisher_block - def get_use_count_map(self): - """Returns a dict of variables to their number of registrations.""" - # TODO(b/70283403): Reimplement this in the old way, where each - # registration function would be responsible for incrementing the count. - # Also, this version has a bug: it won't do the right thing for generic - # registration for parameters that are shared. i.e. it won't set the use - # count to infinity. - vars_to_uses = defaultdict(int) - for key, block in six.iteritems(self.fisher_blocks): - n = ( - block.num_inputs()*block.num_registered_minibatches if isinstance( - block, (fb.FullyConnectedSeriesFB, fb.FullyConnectedMultiIndepFB)) - else block.num_registered_minibatches) - key = utils.ensure_sequence(key) - for k in key: - vars_to_uses[k] += n - return vars_to_uses + def register_loss_function(self, + loss, + colocation_op, + base_name, + name=None, + reuse=VARIABLE_SCOPE): + """Registers a LossFunction object. + + Args: + loss: The LossFunction object. + colocation_op: The op to colocate the loss function's computations with. + base_name: The name to derive a new unique name from is the name argument + is None. + name: (OPTIONAL) str or None. Unique name for this loss function. If None, + a new name is generated. (Default: None) + reuse: (OPTIONAL) 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: If reuse == True and name == None. + ValueError: If reuse == True and seed != None. + KeyError: If reuse == True and no existing LossFunction with 'name' found. + KeyError: If reuse == False and existing LossFunction with 'name' found. + """ + + name = name or self._graph.unique_name(base_name) + + if reuse == VARIABLE_SCOPE: + reuse = variable_scope.get_variable_scope().reuse + + if reuse: + if name is None: + raise ValueError( + "If reuse is enabled, loss function's name must be set.") + + loss_list = self._loss_dict.get(name, None) + + if loss_list is None: + raise KeyError( + "Unable to find loss function named {}. Register a new loss " + "function with reuse=False.".format(name)) + else: + if name in self._loss_dict: + raise KeyError( + "Loss function named {} already exists. Set reuse=True to append " + "another minibatch/tower.".format(name)) + + loss_list = [] + self._loss_dict[name] = loss_list + + loss_list.append(loss) + self.loss_colocation_ops[loss] = colocation_op + + def _get_use_count_map(self): + """Returns a dict mapping variables to their number of registrations.""" + return self._vars_to_uses + + def _add_uses(self, params, uses): + """Register additional uses by params in the graph. + + Args: + params: Variable or tuple of Variables. Parameters for a layer. + uses: int or float. Number of additional uses for these parameters. + """ + params = params if isinstance(params, (tuple, list)) else (params,) + for var in params: + self._vars_to_uses[var] += uses def check_registration(self, variables): """Checks that all variable uses have been registered properly. @@ -312,7 +406,7 @@ class LayerCollection(object): # Note that overlapping parameters (i.e. those that share variables) will # be caught by layer_collection.LayerParametersDict during registration. - reg_use_map = self.get_use_count_map() + reg_use_map = self._get_use_count_map() error_messages = [] @@ -402,12 +496,27 @@ class LayerCollection(object): inputs_to_losses = nest.flatten(tuple(loss.inputs for loss in self.losses)) self._subgraph = utils.SubGraph(inputs_to_losses) + def eval_losses(self): + """Return evaluated losses (colocated with inputs to losses).""" + evals = [] + for loss in self.losses: + with ops.colocate_with(self.loss_colocation_ops[loss]): + evals.append(loss.evaluate()) + return evals + + def eval_losses_on_samples(self): + """Return losses evaluated on samples (colocated with inputs to losses).""" + evals = [] + for loss in self.losses: + with ops.colocate_with(self.loss_colocation_ops[loss]): + evals.append(loss.evaluate_on_sample()) + return evals + def total_loss(self): - return math_ops.add_n(tuple(loss.evaluate() for loss in self.losses)) + return math_ops.add_n(self.eval_losses()) def total_sampled_loss(self): - return math_ops.add_n( - tuple(loss.evaluate_on_sample() for loss in self.losses)) + return math_ops.add_n(self.eval_losses_on_samples()) def _get_linked_approx(self, params): """If params were linked, return their specified approximation.""" @@ -417,6 +526,48 @@ class LayerCollection(object): 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) + + self._add_uses(params, 1) + def register_fully_connected(self, params, inputs, @@ -453,30 +604,37 @@ class LayerCollection(object): block_type = _FULLY_CONNECTED_APPROX_TO_BLOCK_TYPES[approx] has_bias = isinstance(params, (tuple, list)) - block = self.register_block(params, block_type(self, has_bias), reuse=reuse) + block = self.register_block(params, block_type(self, has_bias=has_bias), + reuse=reuse) block.register_additional_minibatch(inputs, outputs) + self._add_uses(params, 1) + def register_conv2d(self, params, strides, padding, inputs, outputs, + data_format=None, + dilations=None, approx=None, reuse=VARIABLE_SCOPE): - """Registers a convolutional layer. + """Registers a call to tf.nn.conv2d(). Args: params: Tensor or 2-tuple of Tensors corresponding to weight and bias of this layer. Weight matrix should have shape [kernel_height, kernel_width, in_channels, out_channels]. Bias should have shape [out_channels]. - strides: 1-D Tensor of length 4. Strides for convolution kernel. + strides: List of 4 ints. Strides for convolution kernel. padding: string. see tf.nn.conv2d for valid values. inputs: Tensor of shape [batch_size, height, width, in_channels]. Inputs to layer. outputs: Tensor of shape [batch_size, height, width, out_channels]. Output produced by layer. + data_format: str or None. Format of data. + dilations: List of 4 ints. Dilations along each dimension. approx: str. One of "kron" or "diagonal". reuse: bool or str. If True, reuse an existing FisherBlock. If False, create a new FisherBlock. If "VARIABLE_SCOPE", use @@ -497,10 +655,206 @@ class LayerCollection(object): raise ValueError("Bad value {} for approx.".format(approx)) block_type = _CONV2D_APPROX_TO_BLOCK_TYPES[approx] + if approx == APPROX_KRONECKER_NAME: + block = self.register_block( + params, + block_type( + layer_collection=self, + params=params, + padding=padding, + strides=strides, + data_format=data_format, + dilation_rate=dilations, + extract_patches_fn="extract_image_patches"), + reuse=reuse) + elif approx == APPROX_DIAGONAL_NAME: + assert strides[0] == strides[-1] == 1 + block = self.register_block( + params, + block_type( + layer_collection=self, + params=params, + padding=padding, + strides=strides, + dilations=dilations, + data_format=data_format), + reuse=reuse) + else: + raise NotImplementedError + + block.register_additional_minibatch(inputs, outputs) + + self._add_uses(params, 1) + + def register_convolution(self, + params, + inputs, + outputs, + padding, + strides=None, + dilation_rate=None, + data_format=None, + approx=None, + reuse=VARIABLE_SCOPE): + """Register a call to tf.nn.convolution(). + + Args: + params: Tensor or 2-tuple of Tensors corresponding to weight and bias of + this layer. Weight matrix should have shape [..filter_spatial_size.., + in_channels, out_channels]. Bias should have shape [out_channels]. + inputs: Tensor of shape [batch_size, ..input_spatial_size.., in_channels]. + Inputs to layer. + outputs: Tensor of shape [batch_size, ..output_spatial_size.., + out_channels]. Output produced by layer. + padding: string. see tf.nn.conv2d for valid values. + strides: List of ints of length len(..input_spatial_size..). Strides for + convolution kernel in spatial dimensions. + dilation_rate: List of ints of length len(..input_spatial_size..). + Dilations along spatial dimension. + data_format: str or None. Format of data. + approx: str. One of "kron" or "diagonal". + 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. + """ + assert approx is None or approx == APPROX_KRONECKER_NAME + block = self.register_block( - params, block_type(self, params, strides, padding), reuse=reuse) + params, + fb.ConvKFCBasicFB( + layer_collection=self, + params=params, + padding=padding, + strides=strides, + dilation_rate=dilation_rate, + data_format=data_format), + reuse=reuse) block.register_additional_minibatch(inputs, outputs) + self._add_uses(params, 1) + + def register_depthwise_conv2d(self, + params, + inputs, + outputs, + strides, + padding, + rate=None, + data_format=None, + approx=None, + reuse=VARIABLE_SCOPE): + """Register a call to tf.nn.depthwise_conv2d(). + + Args: + params: 4-D Tensor of shape [filter_height, filter_width, + in_channels, channel_multiplier]. Convolutional filter. + inputs: Tensor of shape [batch_size, input_height, input_width, + in_channels]. Inputs to layer. + outputs: Tensor of shape [batch_size, output_height, output_width, + in_channels * channel_multiplier]. Output produced by depthwise conv2d. + strides: List of ints of length 4. Strides along all dimensions. + padding: string. see tf.nn.conv2d for valid values. + rate: None or List of ints of length 2. Dilation rates in spatial + dimensions. + data_format: str or None. Format of data. + approx: None or str. Must be "diagonal" if non-None. + 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. + """ + assert approx is None or approx == APPROX_DIAGONAL_NAME + assert data_format in [None, "NHWC"] + + block = self.register_block( + params, + fb.DepthwiseConvDiagonalFB( + layer_collection=self, + params=params, + strides=strides, + padding=padding, + rate=rate, + data_format=data_format), + reuse=reuse) + block.register_additional_minibatch(inputs, outputs) + + self._add_uses(params, 1) + + def register_separable_conv2d(self, + depthwise_params, + pointwise_params, + inputs, + depthwise_outputs, + pointwise_outputs, + strides, + padding, + rate=None, + data_format=None, + approx=None, + reuse=VARIABLE_SCOPE): + """Register a call to tf.nn.separable_conv2d(). + + Note: This requires access to intermediate outputs betwee depthwise and + pointwise convolutions. + + Args: + depthwise_params: 4-D Tensor of shape [filter_height, filter_width, + in_channels, channel_multiplier]. Filter for depthwise conv2d. + pointwise_params: 4-D Tensor of shape [1, 1, in_channels * + channel_multiplier, out_channels]. Filter for pointwise conv2d. + inputs: Tensor of shape [batch_size, input_height, input_width, + in_channels]. Inputs to layer. + depthwise_outputs: Tensor of shape [batch_size, output_height, + output_width, in_channels * channel_multiplier]. Output produced by + depthwise conv2d. + pointwise_outputs: Tensor of shape [batch_size, output_height, + output_width, out_channels]. Output produced by pointwise conv2d. + strides: List of ints of length 4. Strides for depthwise conv2d kernel in + all dimensions. + padding: string. see tf.nn.conv2d for valid values. + rate: None or List of ints of length 2. Dilation rate of depthwise conv2d + kernel in spatial dimensions. + data_format: str or None. Format of data. + approx: None or str. Must be "kron" if non-None. + 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. + """ + self.register_depthwise_conv2d( + params=depthwise_params, + inputs=inputs, + outputs=depthwise_outputs, + strides=strides, + padding=padding, + rate=rate, + data_format=data_format, + approx=APPROX_DIAGONAL_NAME, + reuse=reuse) + + self.register_conv2d( + params=pointwise_params, + inputs=depthwise_outputs, + outputs=pointwise_outputs, + strides=[1, 1, 1, 1], + padding="VALID", + data_format=data_format, + approx=approx, + reuse=reuse) + def register_generic(self, params, batch_size, @@ -534,8 +888,10 @@ class LayerCollection(object): block = self.register_block(params, block_type(self, params), reuse=reuse) block.register_additional_minibatch(batch_size) + self._add_uses(params, float("inf")) + def register_fully_connected_multi(self, params, inputs, outputs, - approx=None): + approx=None, reuse=VARIABLE_SCOPE): """Register fully connected layers with shared parameters. This can handle general fully-connected layers with shared parameters, but @@ -552,6 +908,9 @@ class LayerCollection(object): [batch_size, output_size]. Outputs produced by layer. In the case of RNNs, one Tensor per time step. approx: str. One of "kron_indep", "kron_series_1", or "kron_series_2". + 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'. @@ -569,11 +928,14 @@ class LayerCollection(object): raise ValueError("Bad value {} for approx.".format(approx)) block_type = _FULLY_CONNECTED_MULTI_APPROX_TO_BLOCK_TYPES[approx] - # For now we don't support multiple minibatches for this type of layer, so - # we set reuse=False - self.register_block(params, - block_type(self, inputs, outputs, has_bias=has_bias), - reuse=False) + block = self.register_block(params, block_type(self, has_bias=has_bias), + reuse=reuse) + block.register_additional_minibatch(inputs, outputs) + self._add_uses(params, len(inputs)) + + # TODO(b/74108452): change the loss registration functions names to refer + # to "loss functions" instead of distributions. Following naming convention + # of the loss function classes themselves. def register_categorical_predictive_distribution(self, logits, @@ -596,50 +958,20 @@ class LayerCollection(object): reuse: (OPTIONAL) 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: If reuse == True and name == None. - ValueError: If reuse == True and seed != None. - KeyError: If reuse == True and no existing LossFunction with 'name' found. - KeyError: If reuse == False and existing LossFunction with 'name' found. """ - name = name or self._graph.unique_name( - "register_categorical_predictive_distribution") - - if reuse == VARIABLE_SCOPE: - reuse = variable_scope.get_variable_scope().reuse - - if reuse: - if name is None: - raise ValueError( - "If reuse is enabled, loss function's name must be set.") - if seed is not None: - raise ValueError( - "Seed can only be specified at LossFunction instantiation.") - - loss = self._loss_dict.get(name, None) - - if loss is None: - raise KeyError( - "Unable to find loss function named {}. Create a new LossFunction " - "with reuse=False.".format(name)) - - loss.register_additional_minibatch(logits, targets=targets) - else: - if name in self._loss_dict: - raise KeyError( - "Loss function named {} already exists. Set reuse=True to append " - "another minibatch.".format(name)) - loss = lf.CategoricalLogitsNegativeLogProbLoss( - logits, targets=targets, seed=seed) - self._loss_dict[name] = loss + loss = lf.CategoricalLogitsNegativeLogProbLoss(logits, targets=targets, + seed=seed) + self.register_loss_function(loss, logits, + "categorical_predictive_distribution", + name=name, reuse=reuse) def register_normal_predictive_distribution(self, mean, var=0.5, seed=None, targets=None, - name=None): + name=None, + reuse=VARIABLE_SCOPE): """Registers a normal predictive distribution. Args: @@ -656,21 +988,22 @@ class LayerCollection(object): (Default: None) name: (OPTIONAL) str or None. Unique name for this loss function. If None, a new name is generated. (Default: None) + reuse: (OPTIONAL) bool or str. If True, reuse an existing FisherBlock. + If False, create a new FisherBlock. If VARIABLE_SCOPE, use + tf.get_variable_scope().reuse. """ - name = name or self._graph.unique_name( - "register_normal_predictive_distribution") - if name in self._loss_dict: - raise NotImplementedError( - "Adding logits to an existing LossFunction not yet supported.") - loss = lf.NormalMeanNegativeLogProbLoss( - mean, var, targets=targets, seed=seed) - self._loss_dict[name] = loss + loss = lf.NormalMeanNegativeLogProbLoss(mean, var, targets=targets, + seed=seed) + self.register_loss_function(loss, mean, + "normal_predictive_distribution", + name=name, reuse=reuse) def register_multi_bernoulli_predictive_distribution(self, logits, seed=None, targets=None, - name=None): + name=None, + reuse=VARIABLE_SCOPE): """Registers a multi-Bernoulli predictive distribution. Args: @@ -683,15 +1016,15 @@ class LayerCollection(object): (Default: None) name: (OPTIONAL) str or None. Unique name for this loss function. If None, a new name is generated. (Default: None) + reuse: (OPTIONAL) bool or str. If True, reuse an existing FisherBlock. + If False, create a new FisherBlock. If VARIABLE_SCOPE, use + tf.get_variable_scope().reuse. """ - name = name or self._graph.unique_name( - "register_multi_bernoulli_predictive_distribution") - if name in self._loss_dict: - raise NotImplementedError( - "Adding logits to an existing LossFunction not yet supported.") - loss = lf.MultiBernoulliNegativeLogProbLoss( - logits, targets=targets, seed=seed) - self._loss_dict[name] = loss + loss = lf.MultiBernoulliNegativeLogProbLoss(logits, targets=targets, + seed=seed) + self.register_loss_function(loss, logits, + "multi_bernoulli_predictive_distribution", + name=name, reuse=reuse) def make_or_get_factor(self, cls, args): """Insert 'cls(args)' into 'self.fisher_factors' if not already present. @@ -720,3 +1053,10 @@ class LayerCollection(object): with variable_scope.variable_scope(self._var_scope): self.fisher_factors[key] = cls(*args) return self.fisher_factors[key] + + @contextmanager + def as_default(self): + """Sets this LayerCollection as the default.""" + set_default_layer_collection(self) + yield + set_default_layer_collection(None) diff --git a/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py b/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py index f8aa230d9ca1f542950f56b1e6cf1ab7ccd3d05f..9f4685380705bd409dbcd7e85d0e3bb4189a6adc 100644 --- a/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py +++ b/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py @@ -30,6 +30,8 @@ from tensorflow.python.util.all_util import remove_undocumented # pylint: enable=unused-import,line-too-long,wildcard-import _allowed_symbols = [ + "get_default_layer_collection", + "set_default_layer_collection", "LayerParametersDict", "LayerCollection", "APPROX_KRONECKER_NAME", diff --git a/tensorflow/contrib/kfac/python/ops/loss_functions.py b/tensorflow/contrib/kfac/python/ops/loss_functions.py index cb3e698b9ceab920785adf735f88bd8e535a628f..e7d4243fc3d1c2d860693f2f62447b1c9aeeee03 100644 --- a/tensorflow/contrib/kfac/python/ops/loss_functions.py +++ b/tensorflow/contrib/kfac/python/ops/loss_functions.py @@ -57,30 +57,6 @@ class LossFunction(object): """The inputs to the loss function (excluding the targets).""" pass - @property - def input_minibatches(self): - """A `list` of inputs to the loss function, separated by minibatch. - - Typically there will be one minibatch per tower in a multi-tower setup. - Returns a list consisting of `self.inputs` by default; `LossFunction`s - supporting registering multiple minibatches should override this method. - - Returns: - A `list` of `Tensor`s representing - """ - return [self.inputs] - - @property - def num_registered_minibatches(self): - """Number of minibatches registered for this LossFunction. - - Typically equal to the number of towers in a multi-tower setup. - - Returns: - An `int` representing the number of registered minibatches. - """ - return len(self.input_minibatches) - def evaluate(self): """Evaluate the loss function on the targets.""" if self.targets is not None: @@ -474,7 +450,6 @@ class NormalMeanVarianceNegativeLogProbLoss(DistributionNegativeLogProbLoss): assert len(variance.shape) == 2, "Expect 2D variance tensor." self._mean = mean self._variance = variance - self._scale = math_ops.sqrt(variance) self._targets = targets super(NormalMeanVarianceNegativeLogProbLoss, self).__init__(seed=seed) @@ -484,7 +459,7 @@ class NormalMeanVarianceNegativeLogProbLoss(DistributionNegativeLogProbLoss): @property def dist(self): - return normal.Normal(loc=self._mean, scale=self._scale) + return normal.Normal(loc=self._mean, scale=math_ops.sqrt(self._variance)) @property def params(self): @@ -502,7 +477,7 @@ class NormalMeanVarianceNegativeLogProbLoss(DistributionNegativeLogProbLoss): @property def _fisher_mean_factor(self): - return 1. / self._scale + return 1. / math_ops.sqrt(self._variance) @property def _fisher_var(self): @@ -611,36 +586,13 @@ class CategoricalLogitsNegativeLogProbLoss(DistributionNegativeLogProbLoss, index in [0, output_size). seed: int or None. Default random seed when sampling. """ - self._logits_components = [] - self._targets_components = [] - self.register_additional_minibatch(logits, targets=targets) + self._logits = logits + self._targets = targets super(CategoricalLogitsNegativeLogProbLoss, self).__init__(seed=seed) - def register_additional_minibatch(self, logits, targets=None): - """Register an additiona minibatch's worth of parameters. - - Args: - logits: Tensor of shape [batch_size, output_size]. Parameters for - underlying distribution. - targets: None or Tensor of shape [batch_size, output_size]. Each row must - be a one-hot vector. - """ - self._logits_components.append(logits) - self._targets_components.append(targets) - - @property - def _logits(self): - return array_ops.concat(self._logits_components, axis=0) - - @property - def input_minibatches(self): - return self._logits_components - @property def targets(self): - if all(target is None for target in self._targets_components): - return None - return array_ops.concat(self._targets_components, axis=0) + return self._targets @property def dist(self): 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 1974b07acfc879dc4bc844db9af88fd1043d6698..083da768ec97aca3e63995491bb579835bb5377f 100644 --- a/tensorflow/contrib/kfac/python/ops/optimizer.py +++ b/tensorflow/contrib/kfac/python/ops/optimizer.py @@ -18,16 +18,21 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import warnings + # pylint disable=long-line from tensorflow.contrib.kfac.python.ops import curvature_matrix_vector_products as cmvp from tensorflow.contrib.kfac.python.ops import estimator as est # pylint enable=long-line +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 state_ops +from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables as tf_variables from tensorflow.python.training import gradient_descent @@ -47,6 +52,7 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): name="KFAC", estimation_mode="gradients", colocate_gradients_with_ops=True, + batch_size=None, cov_devices=None, inv_devices=None): """Initializes the KFAC optimizer with the given settings. @@ -61,6 +67,8 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): damping: The damping factor used to stabilize training due to errors in the local approximation with the Fisher information matrix, and to regularize the update direction by making it closer to the gradient. + If damping is adapted during training then this value is used for + initializing damping varaible. (Higher damping means the update looks more like a standard gradient update - see Tikhonov regularization.) layer_collection: The layer collection object, which holds the fisher @@ -86,12 +94,16 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): colocate_gradients_with_ops: Whether we should request gradients we compute in the estimator be colocated with their respective ops. (Default: True) + batch_size: The size of the mini-batch. Only needed when momentum_type + == 'qmodel' or when automatic adjustment is used. (Default: None) 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. + Can be None, which means that no devices are specified. Only used + with (soon-to-be-depcrecated "convenience" properties). inv_devices: Iterable of device strings (e.g. '/gpu:0'). Inversion computations will be placed on these devices in a round-robin fashion. - Can be None, which means that no devices are specified. + Can be None, which means that no devices are specified. Only used + with (soon-to-be-depcrecated "convenience" properties). Raises: ValueError: If the momentum type is unsupported. @@ -105,15 +117,34 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): if variables is None: variables = tf_variables.trainable_variables() - self._fisher_est = est.FisherEstimator( - variables, - cov_ema_decay, - damping, - layer_collection, - estimation_mode=estimation_mode, - colocate_gradients_with_ops=colocate_gradients_with_ops, - cov_devices=cov_devices, - inv_devices=inv_devices) + # Parameters to be passed to the Fisher estimator: + self._variables = variables + self._cov_ema_decay = cov_ema_decay + self._layers = layer_collection + self._estimation_mode = estimation_mode + self._colocate_gradients_with_ops = colocate_gradients_with_ops + self._cov_devices = cov_devices + self._inv_devices = inv_devices + + # The below paramaters are required only if damping needs to be adapated. + # These parameters can be set by calling + # set_damping_adaptation_params() explicitly. + self._damping_adaptation_decay = 0.95 + self._damping_adaptation_interval = 5 + # Check section 6.5 KFAC paper. omega(1) = pow(damping decay, interval) + self._omega = ( + self._damping_adaptation_decay**self._damping_adaptation_interval) + self._adapt_damping = False + self._min_damping = 1e-5 + self._prev_train_batch = None + self._is_chief = False + self._loss_fn = None + self._damping_constant = damping + self._damping = None + self._rho = None + self._prev_loss = None + self._q_model_change = None + self._update_damping_op = None momentum_type = momentum_type.lower() legal_momentum_types = ["regular", "adam", "qmodel"] @@ -122,61 +153,290 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): raise ValueError("Unsupported momentum type {}. Must be one of {}." .format(momentum_type, legal_momentum_types)) if momentum_type != "regular" and norm_constraint is not None: - raise ValueError("Update clipping is only supported with momentum" + raise ValueError("Update clipping is only supported with momentum " "type 'regular'.") if momentum_type not in ["regular", "adam"] and momentum != 0: raise ValueError("Momentum must be unspecified if using a momentum_type " "other than 'regular' or 'adam'.") + # Extra parameters of the optimizer self._momentum = momentum self._momentum_type = momentum_type self._norm_constraint = norm_constraint - - # this is a bit of a hack - # TODO(duckworthd): Handle this in a better way (e.g. pass it in?) - self._batch_size = array_ops.shape(layer_collection.losses[0].inputs)[0] - self._losses = layer_collection.losses + self._batch_size = batch_size + + with variable_scope.variable_scope(name): + self._fisher_est = est.FisherEstimator( + self._variables, + self._cov_ema_decay, + self.damping, + self._layers, + exps=(-1,), + estimation_mode=self._estimation_mode, + colocate_gradients_with_ops=self._colocate_gradients_with_ops) super(KfacOptimizer, self).__init__(learning_rate, name=name) + def set_damping_adaptation_params(self, + is_chief, + prev_train_batch, + loss_fn, + min_damping=1e-5, + damping_adaptation_decay=0.99, + damping_adaptation_interval=5): + """Sets parameters required to adapt damping during training. + + When called, enables damping adaptation according to the Levenberg-Marquardt + style rule described in Section 6.5 of "Optimizing Neural Networks with + Kronecker-factored Approximate Curvature". + + Note that this function creates Tensorflow variables which store a few + scalars and are accessed by the ops which update the damping (as part + of the training op returned by the minimize() method). + + Args: + is_chief: `Boolean`, `True` if the worker is chief. + prev_train_batch: Training data used to minimize loss in the previous + step. This will be used to evaluate loss by calling + `loss_fn(prev_train_batch)`. + loss_fn: `function` that takes as input training data tensor and returns + a scalar loss. + min_damping: `float`(Optional), Minimum value the damping parameter + can take. Default value 1e-5. + damping_adaptation_decay: `float`(Optional), The `damping` parameter is + multipled by the `damping_adaptation_decay` every + `damping_adaptation_interval` number of iterations. Default value 0.99. + damping_adaptation_interval: `int`(Optional), Number of steps in between + updating the `damping` parameter. Default value 5. + + Raises: + ValueError: If `set_damping_adaptation_params` is already called and the + the `adapt_damping` is `True`. + """ + if self._adapt_damping: + raise ValueError("Damping adaptation parameters already set.") + + with variable_scope.variable_scope(self.get_name()): + self._adapt_damping = True + self._is_chief = is_chief + self._prev_train_batch = prev_train_batch + self._loss_fn = loss_fn + self._damping_adaptation_decay = damping_adaptation_decay + self._damping_adaptation_interval = damping_adaptation_interval + self._omega = ( + self._damping_adaptation_decay**self._damping_adaptation_interval) + self._min_damping = min_damping + + self._rho = variable_scope.get_variable( + "rho", shape=(), dtype=dtypes.float32, trainable=False) # LM ratio. + self._prev_loss = variable_scope.get_variable( + "prev_loss", shape=(), dtype=dtypes.float32, trainable=False) + self._q_model_change = variable_scope.get_variable( + "q_model_change", shape=(), dtype=dtypes.float32, trainable=False) + self._damping = variable_scope.get_variable( + "damping", initializer=self._damping_constant, trainable=False) + @property def cov_update_thunks(self): - return self._fisher_est.cov_update_thunks + self._maybe_make_and_save_everything() + return self._cov_update_thunks @property def cov_update_ops(self): - return self._fisher_est.cov_update_ops + self._maybe_make_and_save_everything() + return self._cov_update_ops @property def cov_update_op(self): - return self._fisher_est.cov_update_op + self._maybe_make_and_save_everything() + return self._cov_update_op @property def inv_update_thunks(self): - return self._fisher_est.inv_update_thunks + self._maybe_make_and_save_everything() + return self._inv_update_thunks @property def inv_update_ops(self): - return self._fisher_est.inv_update_ops + self._maybe_make_and_save_everything() + return self._inv_update_ops @property def inv_update_op(self): - return self._fisher_est.inv_update_op + self._maybe_make_and_save_everything() + return self._inv_update_op @property def variables(self): - return self._fisher_est.variables + return self._variables @property def damping(self): - return self._fisher_est.damping + if self._damping: + return self._damping + else: + return self._damping_constant + + @property + def damping_adaptation_interval(self): + return self._damping_adaptation_interval + + def _maybe_make_and_save_everything(self): + if not self._fisher_est.made_vars(): + warnings.warn("These convenience properties will be depcrecated soon. " + "Please use explicit op/thunk creation methods instead " + "(e.g. make_ops_and_vars_round_robin, etc).", + DeprecationWarning) + (self._cov_update_ops, self._cov_update_op, self._inv_update_ops, + self._inv_update_op, self._cov_update_thunks, + self._inv_update_thunks) = self.make_ops_and_vars_round_robin( + cov_devices=self._cov_devices, + inv_devices=self._inv_devices) + + def make_ops_and_vars(self): + """Make ops and vars with no specific device placement. + + See make_ops_and_vars_round_robin for details. + + Returns: + cov_update_ops: List of ops that compute the cov updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + cov_update_op: cov_update_ops grouped into a single op. + inv_update_ops: List of ops that compute the inv updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + cov_update_op: cov_update_ops grouped into a single op. + inv_update_op: inv_update_ops grouped into a single op. + """ + with variable_scope.variable_scope(self.get_name()): + return self._fisher_est.make_ops_and_vars() + + def make_ops_and_vars_round_robin(self, cov_devices=None, inv_devices=None): + """Make ops and vars with a round-robin device placement strategy. + + For each factor, all of that factor's cov variables and their associated + update ops will be placed on a particular device. A new device is chosen + for each factor by cycling through list of devices in the cov_devices + argument. If cov_devices is None then no explicit device placement occurs. + + An analogous strategy is followed for inverse update ops, with the list of + devices being given by the inv_devices argument. + + Inverse variables on the other hand are not placed on any specific device + (they will just use the current the device placement context, whatever + that happens to be). The idea is that the inverse variable belong where + they will be accessed most often, which is the device that actually applies + the preconditioner to the gradient. The user will be responsible for setting + the device context for this. + + Args: + 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. + inv_devices: Iterable of device strings (e.g. '/gpu:0'). Inversion + computations will be placed on these devices in a round-robin fashion. + Can be None, which means that no devices are specified. + + Returns: + cov_update_ops: List of ops that compute the cov updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + cov_update_op: cov_update_ops grouped into a single op. + inv_update_ops: List of ops that compute the inv updates. Corresponds + one-to-one with the list of factors given by the "factors" property. + cov_update_op: cov_update_ops grouped into a single op. + inv_update_op: inv_update_ops grouped into a single op. + cov_update_thunks: Thunks that make the ops in cov_update_ops. + inv_update_thunks: Thunks that make the ops in inv_update_ops. + """ + with variable_scope.variable_scope(self.get_name()): + return self._fisher_est.make_ops_and_vars_round_robin( + cov_devices=cov_devices, inv_devices=inv_devices) + + def make_vars_and_create_op_thunks_round_robin(self, + cov_devices=None, + inv_devices=None): + """Make vars and create op thunks w/ a round-robin device placement strat. + + For each factor, all of that factor's cov variables and their associated + update ops will be placed on a particular device. A new device is chosen + for each factor by cycling through list of devices in the cov_devices + argument. If cov_devices is None then no explicit device placement occurs. + + An analogous strategy is followed for inverse update ops, with the list of + devices being given by the inv_devices argument. + + Inverse variables on the other hand are not placed on any specific device + (they will just use the current the device placement context, whatever + that happens to be). The idea is that the inverse variable belong where + they will be accessed most often, which is the device that actually applies + the preconditioner to the gradient. The user will be responsible for setting + the device context for this. + + Args: + 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. + inv_devices: Iterable of device strings (e.g. '/gpu:0'). Inversion + computations will be placed on these devices in a round-robin fashion. + Can be None, which means that no devices are specified. + Returns: + cov_update_thunks: List of cov update thunks. Corresponds one-to-one with + the list of factors given by the "factors" property. + inv_update_thunks: List of inv update thunks. Corresponds one-to-one with + the list of factors given by the "factors" property. + """ + scope = self.get_name() + "/" + self._fisher_est.name + return self._fisher_est.make_vars_and_create_op_thunks_round_robin( + scope=scope, cov_devices=cov_devices, inv_devices=inv_devices) + + def ops_and_vars_thunks(self): + """Create thunks that make the ops and vars on demand. + + This function returns 4 lists of thunks: cov_variable_thunks, + cov_update_thunks, inv_variable_thunks, and inv_update_thunks. + + The length of each list is the number of factors and the i-th element of + each list corresponds to the i-th factor (given by the "factors" property). + + Note that the execution of these thunks must happen in a certain + partial order. The i-th element of cov_variable_thunks must execute + before the i-th element of cov_update_thunks (and also the i-th element + of inv_update_thunks). Similarly, the i-th element of inv_variable_thunks + must execute before the i-th element of inv_update_thunks. + + TL;DR (oversimplified): Execute the thunks according to the order that + they are returned. + + Returns: + cov_variable_thunks: A list of thunks that make the cov variables. + cov_update_thunks: A list of thunks that make the cov update ops. + inv_variable_thunks: A list of thunks that make the inv variables. + inv_update_thunks: A list of thunks that make the inv update ops. + """ + scope = self.get_name() + "/" + self._fisher_est.name + return self._fisher_est.ops_and_vars_thunks(scope=scope) def minimize(self, *args, **kwargs): - kwargs["var_list"] = kwargs.get("var_list") or self.variables - if set(kwargs["var_list"]) != set(self.variables): - raise ValueError("var_list doesn't match with set of Fisher-estimating " - "variables.") - return super(KfacOptimizer, self).minimize(*args, **kwargs) + # Should this variable scope encompass everything below? Or will the super- + # class make another copy of the same name scope? + with variable_scope.variable_scope(self.get_name()): + kwargs["var_list"] = kwargs.get("var_list") or self.variables + if set(kwargs["var_list"]) != set(self.variables): + raise ValueError("var_list doesn't match with set of Fisher-estimating " + "variables.") + if self._adapt_damping and self._is_chief: + global_step = kwargs.get("global_step", None) + if not global_step: + raise KeyError("global_step needs to be passed to optimizer.minimize " + "if damping parameter is adapted.") + update_damping_op = self._update_damping(self._prev_train_batch, + global_step) + with ops.control_dependencies([update_damping_op]): + loss = args[0] + loss_assign_op = state_ops.assign(self._prev_loss, loss) + train_op = super(KfacOptimizer, self).minimize(*args, **kwargs) + return control_flow_ops.group(loss_assign_op, train_op) + else: + return super(KfacOptimizer, self).minimize(*args, **kwargs) def compute_gradients(self, *args, **kwargs): # args[1] could be our var_list @@ -201,6 +461,8 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): Returns: An `Operation` that applies the specified gradients. """ + self._maybe_make_and_save_everything() + # In Python 3, grads_and_vars can be a zip() object which can only be # iterated over once. By converting it to a list, we ensure that it can be # iterated over more than once. @@ -296,6 +558,20 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): coeff = self._update_clip_coeff(grads_and_vars, precon_grads_and_vars) return [(pgrad * coeff, var) for pgrad, var in precon_grads_and_vars] + def _compute_prev_updates(self, variables): + """Computes previous updates as negative velocities scaled by learning rate. + + Args: + variables: List of variables in the graph that the update will be + applied to. + + Returns: + List of previous updates applied to the `variables`. + """ + return list( + -1 * self._learning_rate * self._zeros_slot(var, "velocity", self._name) + for var in variables) + def _compute_qmodel_hyperparams(self, precon_grads, prev_updates, grads, variables): """Compute optimal update hyperparameters from the quadratic model. @@ -336,7 +612,8 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): = qmodel(alpha*precon_grad + mu*prev_update) - L(theta). """ - cmvpc = cmvp.CurvatureMatrixVectorProductComputer(self._losses, variables) + cmvpc = cmvp.CurvatureMatrixVectorProductComputer(self._layers.losses, + variables) # compute the matrix-vector products with the transposed Fisher factor fft_precon_grads = cmvpc.multiply_fisher_factor_transpose(precon_grads) @@ -374,9 +651,9 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): c = ops.convert_to_tensor([[_inner_product_list(grads, precon_grads)], [_inner_product_list(grads, prev_updates)]]) - sol = _two_by_two_solve(m, c) - alpha = -sol[0] - mu = -sol[1] + sol = -1. * _two_by_two_solve(m, c) + alpha = sol[0] + mu = sol[1] qmodel_change = 0.5 * math_ops.reduce_sum(sol * c) return alpha, mu, qmodel_change @@ -404,6 +681,52 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): return control_flow_ops.cond( math_ops.equal(m_22, 0.0), zero_prevupd_case, non_zero_prevupd_case) + def _assign_q_model_change(self, q_model_change): + """Assigns `q_model_change` to `self._q_model_change` if damping is adapted. + + Note only the chief worker does the assignment. + + Args: + q_model_change: Scalar tensor of type `float32`. + + Returns: + If `adapt_damping` is `True` then returns an assign op, Otherwise returns + a no_op(). + """ + if self._adapt_damping and self._is_chief: + q_model_assign_op = state_ops.assign(self._q_model_change, q_model_change) + else: + q_model_assign_op = control_flow_ops.no_op() + return q_model_assign_op + + def _compute_qmodel_hyperparams_wrapper(self, grads_and_vars, + precon_grads_and_vars): + """Wrapper function for `self._compute_qmodel_hyperparams`. + + Constructs a list of preconditioned gradients and variables. Also creates a + op to asssign the computed q model change to `self._q_model_change`. + + Args: + grads_and_vars: List of (gradient, variable) pairs. + precon_grads_and_vars: List of (preconditioned gradients, variable) + pairs. + + Returns: + (alpha, mu, q_model_assign_op), where alpha and mu are chosen to optimize + the quadratic model, `q_model_assign_op` assigns the computed q model + change to `self._q_model_change`. + """ + precon_grads = list( + precon_grad for (precon_grad, _) in precon_grads_and_vars) + grads = list(grad for (grad, _) in grads_and_vars) + variables = list(var for (_, var) in grads_and_vars) + prev_updates = self._compute_prev_updates(variables) + # Compute optimal velocity update parameters according to quadratic model + alpha, mu, q_model_change = self._compute_qmodel_hyperparams( + precon_grads, prev_updates, grads, variables) + + return alpha, mu, self._assign_q_model_change(q_model_change) + def _compute_update_steps(self, grads_and_vars): """Computes the update steps for the variables given the gradients. @@ -411,8 +734,10 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): grads_and_vars: List of (gradient, variable) pairs. Returns: - An 'Operation that computes the update steps for the given variables. + A list of tuple (assign_op ,var) where `assign_op` assigns the update + steps to `var`. """ + if self._momentum_type == "regular": # Compute "preconditioned" gradient. precon_grads_and_vars = self._fisher_est.multiply_inverse(grads_and_vars) @@ -423,8 +748,13 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): precon_grads_and_vars) # Update the velocity with this and return it as the step. - return self._update_velocities(precon_grads_and_vars, self._momentum) - + if self._adapt_damping and self._is_chief: + _, _, q_model_assign_op = self._compute_qmodel_hyperparams_wrapper( + grads_and_vars, precon_grads_and_vars) + with ops.control_dependencies([q_model_assign_op]): + return self._update_velocities(precon_grads_and_vars, self._momentum) + else: + return self._update_velocities(precon_grads_and_vars, self._momentum) elif self._momentum_type == "adam": # Update velocity. velocities_and_vars = self._update_velocities(grads_and_vars, @@ -436,23 +766,13 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): # Compute "preconditioned" gradient. precon_grads_and_vars = self._fisher_est.multiply_inverse(grads_and_vars) - # Extract out singleton lists from the tuple-lists - precon_grads = list( - precon_grad for (precon_grad, _) in precon_grads_and_vars) - grads = list(grad for (grad, _) in grads_and_vars) - variables = list(var for (_, var) in grads_and_vars) - # previous updates are the negative velocities (up to scaling by LR) - prev_updates = list( - -self._zeros_slot(var, "velocity", self._name) for var in variables) - # Compute optimal velocity update parameters according to quadratic model - alpha, mu, _ = self._compute_qmodel_hyperparams( - precon_grads, prev_updates, grads, variables) + alpha, mu, q_model_assign_op = self._compute_qmodel_hyperparams_wrapper( + grads_and_vars, precon_grads_and_vars) - # Update the velocity with precon_grads according to these params - # and return it as the step. - return self._update_velocities( - precon_grads_and_vars, mu, vec_coeff=-alpha) + with ops.control_dependencies([q_model_assign_op]): + return self._update_velocities( + precon_grads_and_vars, mu, vec_coeff=-alpha) def _update_velocities(self, vecs_and_vars, decay, vec_coeff=1.0): """Updates the velocities of the variables with the given vectors. @@ -482,6 +802,51 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): # Go through variable and update its associated part of the velocity vector. return [_update_velocity(vec, var) for vec, var in vecs_and_vars] + # TODO(b/73448937): Move all update damping code to a separate class/function. + def _update_damping(self, prev_batch, global_step): + """Adapts damping parameter. Check KFAC (Section 6.5) for the details. + + The damping parameter is updated according to the Levenberg-Marquardt rule + every `self._damping_adaptation_interval` iterations. + + Args: + prev_batch: Tensor or tuple of tensors which can be passed to + `self._loss_fn` to evaluate loss. + global_step: `Variable` which keeps track of number of times the training + variables have been updated. + Returns: + A `tf.cond` op which updates the damping parameter. + """ + def compute_damping(): + """"Adapts damping parameter based on "reduction ratio". + + Reduction ratio captures how closely the quadratic approximation to the + loss function approximates the actual loss within a trust region. The + damping update tries to make the damping as small as possible while + maintaining the property that the quadratic model remains a good local + approximation to the loss function. + + Returns: + An Op to assign newly computed damping value to `self._damping`. + """ + prev_batch_loss = self._loss_fn(prev_batch) + with ops.control_dependencies([prev_batch_loss]): + rho_assign = self._rho.assign( + (prev_batch_loss - self._prev_loss) / self._q_model_change) + with ops.control_dependencies([rho_assign]): + new_damping = control_flow_ops.case( + [(self._rho < 0.25, lambda: self.damping / self._omega), + (self._rho > 0.75, lambda: self.damping * self._omega)], + lambda: self.damping) + with ops.control_dependencies([new_damping]): + new_damping_min = math_ops.maximum(new_damping, self._min_damping) + return control_flow_ops.group(self._damping.assign(new_damping_min)) + + return control_flow_ops.cond( + math_ops.equal( + math_ops.mod(global_step + 1, self._damping_adaptation_interval), + 0), compute_damping, control_flow_ops.no_op) + def _inner_product_list(list1, list2): return math_ops.add_n( diff --git a/tensorflow/contrib/kfac/python/ops/utils.py b/tensorflow/contrib/kfac/python/ops/utils.py index e89508fa46b6e2ce278e5373e6c9d17203ad1ef2..af26f5e56bf9bb22cc9bc2b409209d027477ed89 100644 --- a/tensorflow/contrib/kfac/python/ops/utils.py +++ b/tensorflow/contrib/kfac/python/ops/utils.py @@ -24,11 +24,13 @@ from tensorflow.contrib.tpu.python.ops import tpu_ops from tensorflow.contrib.tpu.python.tpu import tpu_function 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 control_flow_ops from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import linalg_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 resource_variable_ops from tensorflow.python.ops import variables @@ -144,7 +146,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]]) @@ -163,6 +167,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) @@ -234,19 +243,22 @@ class SubGraph(object): # Set of all ancestor Tensors, Ops to 'outputs'. self._members = set() - self._recurse_add(outputs) + self._iter_add(outputs) - def _recurse_add(self, nodes): - """Recursively adds all of nodes' ancestors.""" - for node in nodes: - if node in self._members: - continue - self._members.add(node) + def _iter_add(self, root): + """Iteratively adds all of nodes' ancestors using depth first search.""" + stack = [root] + while stack: + nodes = stack.pop() + for node in nodes: + if node in self._members: + continue + self._members.add(node) - if isinstance(node, ops.Tensor): - self._recurse_add((node.op,)) - elif isinstance(node, ops.Operation): - self._recurse_add(node.inputs) + if isinstance(node, ops.Tensor): + stack.append((node.op,)) + elif isinstance(node, ops.Operation): + stack.append(node.inputs) def is_member(self, node): """Check if 'node' is in this subgraph.""" @@ -420,5 +432,249 @@ def batch_execute(global_step, thunks, batch_size, name=None): return result +def extract_convolution_patches(inputs, + filter_shape, + padding, + strides=None, + dilation_rate=None, + name=None, + data_format=None): + """Extracts inputs to each output coordinate in tf.nn.convolution. + + This is a generalization of tf.extract_image_patches() to tf.nn.convolution(), + where the number of spatial dimensions may be something other than 2. + + Assumes, + - First dimension of inputs is batch_size + - Convolution filter is applied to all input channels. + + Args: + inputs: Tensor of shape [batch_size, ..spatial_image_shape.., + ..spatial_filter_shape.., in_channels]. Inputs to tf.nn.convolution(). + filter_shape: List of ints. Shape of filter passed to tf.nn.convolution(). + padding: string. Padding method. One of "VALID", "SAME". + strides: None or list of ints. Strides along spatial dimensions. + dilation_rate: None or list of ints. Dilation along spatial dimensions. + name: None or str. Name of Op. + data_format: None or str. Format of data. + + Returns: + Tensor of shape [batch_size, ..spatial_image_shape.., + ..spatial_filter_shape.., in_channels] + + Raises: + ValueError: If data_format does not put channel last. + ValueError: If inputs and filter disagree on in_channels. + """ + if not is_data_format_channel_last(data_format): + raise ValueError("Channel must be last dimension.") + with ops.name_scope(name, "extract_convolution_patches", + [inputs, filter_shape, padding, strides, dilation_rate]): + batch_size = inputs.shape.as_list()[0] + in_channels = inputs.shape.as_list()[-1] + + # filter_shape = spatial_filter_shape + [in_channels, out_channels] + spatial_filter_shape = filter_shape[:-2] + if in_channels != filter_shape[-2]: + raise ValueError("inputs and filter_shape must agree on in_channels.") + + # Map each input feature to a location in the output. + out_channels = np.prod(spatial_filter_shape) * in_channels + filters = linalg_ops.eye(out_channels) + filters = array_ops.reshape( + filters, + list(spatial_filter_shape) + [in_channels, out_channels]) + + result = nn_ops.convolution( + inputs, + filters, + padding=padding, + strides=strides, + dilation_rate=dilation_rate) + spatial_output_shape = result.shape.as_list()[1:-1] + result = array_ops.reshape(result, + [batch_size or -1] + spatial_output_shape + + list(spatial_filter_shape) + [in_channels]) + + return result + + +def extract_pointwise_conv2d_patches(inputs, + filter_shape, + name=None, + data_format=None): + """Extract patches for a 1x1 conv2d. + + Args: + inputs: 4-D Tensor of shape [batch_size, height, width, in_channels]. + filter_shape: List of 4 ints. Shape of filter to apply with conv2d() + name: None or str. Name for Op. + data_format: None or str. Format for data. See 'data_format' in + tf.nn.conv2d() for details. + + Returns: + Tensor of shape [batch_size, ..spatial_input_shape.., + ..spatial_filter_shape.., in_channels] + + Raises: + ValueError: if inputs is not 4-D. + ValueError: if filter_shape is not [1, 1, ?, ?] + ValueError: if data_format is not channels-last. + """ + if inputs.shape.ndims != 4: + raise ValueError("inputs must have 4 dims.") + if len(filter_shape) != 4: + raise ValueError("filter_shape must have 4 dims.") + if filter_shape[0] != 1 or filter_shape[1] != 1: + raise ValueError("filter_shape must have shape 1 along spatial dimensions.") + if not is_data_format_channel_last(data_format): + raise ValueError("data_format must be channels last.") + with ops.name_scope(name, "extract_pointwise_conv2d_patches", + [inputs, filter_shape]): + ksizes = [1, 1, 1, 1] # Spatial shape is 1x1. + strides = [1, 1, 1, 1] # Operate on all pixels. + rates = [1, 1, 1, 1] # Dilation has no meaning with spatial shape = 1. + padding = "VALID" # Doesn't matter. + result = array_ops.extract_image_patches(inputs, ksizes, strides, rates, + padding) + + batch_size, input_height, input_width, in_channels = inputs.shape.as_list() + filter_height, filter_width, in_channels, _ = filter_shape + return array_ops.reshape(result, [ + batch_size, input_height, input_width, filter_height, filter_width, + in_channels + ]) + + +def is_data_format_channel_last(data_format): + """True if data_format puts channel last.""" + if data_format is None: + return True + return data_format.endswith("C") + + +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) + + +class PartitionedTensor(object): + """A Tensor partitioned across its 0-th dimension.""" + + def __init__(self, tensors): + """Initializes PartitionedTensor. + + Args: + tensors: List of Tensors. All Tensors must agree on shape (excepting + batch dimension) and dtype. + + Raises: + ValueError: If 'tensors' has length zero. + ValueError: if contents of 'tensors' don't agree on shape or dtype. + """ + if not tensors: + raise ValueError("tensors must be a list of 1+ Tensors.") + + dtype = tensors[0].dtype + if not all(tensor.dtype == dtype for tensor in tensors): + raise ValueError("all tensors must have dtype = %s." % dtype) + + shape = tensors[0].shape[1:] + if not all(tensor.shape[1:] == shape for tensor in tensors): + raise ValueError("All tensors must have shape = %s (excluding batch " + "dimension)." % shape) + + self.tensors = tensors + self._concats = {} # {device: Tensor} + + @property + def shape(self): + feature_shape = self.tensors[0].shape[1:] + batch_size = sum([tensor.shape[0] for tensor in self.tensors], + tensor_shape.Dimension(0)) + return tensor_shape.TensorShape([batch_size]).concatenate(feature_shape) + + def get_shape(self): + return self.shape + + @property + def dtype(self): + return self.tensors[0].dtype + + def devices(self): + return set(tensor.device for tensor in self.tensors) + + def __str__(self): + return "PartitionedTensor([%s, ...], dtype=%s, shape=%s)" % ( + self.tensors[0].name, self.dtype.name, tuple(self.shape.as_list())) + + def __hash__(self): + return hash(tuple(self.tensors)) + + def as_tensor(self, dtype=None, name=None, as_ref=False): + with ops.name_scope(name, "PartitionedTensor.as_tensor", self.tensors): + assert not as_ref + assert dtype in [None, self.dtype] + result = array_ops.concat(self.tensors, axis=0) + + # Cache 'result' if we haven't already cached a value for this device. + if result.device not in self._concats: + self._concats[result.device] = result + return self._concats[result.device] + + +ops.register_tensor_conversion_function( + PartitionedTensor, + lambda val, dtype, name, as_ref: val.as_tensor(dtype, name, as_ref)) + + # 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 fe8e39c212c2c3381f9aa6fdb9fdf423ff958481..330d222dbf70fcfa02ffd47261c0513d9dd6e0e9 100644 --- a/tensorflow/contrib/kfac/python/ops/utils_lib.py +++ b/tensorflow/contrib/kfac/python/ops/utils_lib.py @@ -40,6 +40,11 @@ _allowed_symbols = [ "fwd_gradients", "ensure_sequence", "batch_execute", + "extract_convolution_patches", + "extract_pointwise_conv2d_patches", + "is_data_format_channel_last", + "matmul_sparse_dense", + "matmul_diag_sparse", ] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/labeled_tensor/BUILD b/tensorflow/contrib/labeled_tensor/BUILD index 894e6f6946bb59810a9da2d304cc0dd43d25201d..544065dac6a10094a376c18e84521b1a26401cdd 100644 --- a/tensorflow/contrib/labeled_tensor/BUILD +++ b/tensorflow/contrib/labeled_tensor/BUILD @@ -70,6 +70,7 @@ py_test( "python/ops/core_test.py", ], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":_typecheck", ":core", diff --git a/tensorflow/contrib/labeled_tensor/python/ops/core.py b/tensorflow/contrib/labeled_tensor/python/ops/core.py index abc18aa123bb4d40b54d22ec03257c5350118d13..0c6bba758b429a8c4112bc6abb2fae542b5dfc14 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/core.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/core.py @@ -361,6 +361,10 @@ class LabeledTensor(object): def dtype(self): return self._tensor.dtype + @property + def shape(self): + return self._tensor.shape + @property def name(self): return self._tensor.name diff --git a/tensorflow/contrib/labeled_tensor/python/ops/core_test.py b/tensorflow/contrib/labeled_tensor/python/ops/core_test.py index e70b4923749d89aba1bd0187857d762305daeb07..e378db56afb1d4f9463d2c9b0f1fa4c0feea8fb0 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/core_test.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/core_test.py @@ -244,6 +244,9 @@ class LabeledTensorTest(test_util.Base): def test_dtype(self): self.assertEqual(self.lt.dtype, self.lt.tensor.dtype) + def test_shape(self): + self.assertEqual(self.lt.shape, self.lt.tensor.shape) + def test_get_shape(self): self.assertEqual(self.lt.get_shape(), self.lt.tensor.get_shape()) diff --git a/tensorflow/contrib/labeled_tensor/python/ops/ops.py b/tensorflow/contrib/labeled_tensor/python/ops/ops.py index c957b41a49b292225e547ce17b0c5a247810325a..3ba1026383ef146adb32197ae41b5c251155bf46 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/ops.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/ops.py @@ -951,7 +951,7 @@ def define_reduce_op(op_name, reduce_fn): intermediate_axes.append(axis) reduce_op = reduce_fn( - labeled_tensor.tensor, reduction_dimensions, keep_dims=True) + labeled_tensor.tensor, reduction_dimensions, keepdims=True) reduce_lt = core.LabeledTensor(reduce_op, intermediate_axes) return squeeze(reduce_lt, axes_to_squeeze, name=scope) diff --git a/tensorflow/contrib/layers/BUILD b/tensorflow/contrib/layers/BUILD index 852d06e1e3cc8f8deecd15b7436cd4e4a393ad66..cc7bbabf210ded9a31eb789fa8b94e8bde62ea43 100644 --- a/tensorflow/contrib/layers/BUILD +++ b/tensorflow/contrib/layers/BUILD @@ -188,6 +188,7 @@ py_test( size = "small", srcs = ["python/layers/normalization_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":layers_py", "//tensorflow/contrib/framework:framework_py", @@ -353,6 +354,7 @@ py_test( size = "small", srcs = ["python/ops/sparse_ops_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":layers_py", "//tensorflow/python:array_ops", diff --git a/tensorflow/contrib/layers/__init__.py b/tensorflow/contrib/layers/__init__.py index ef419862b49f4d03d9b711c49155d4ae1252d5bc..337c9e06b870b2cca53fcdbf3d94225660e193c4 100644 --- a/tensorflow/contrib/layers/__init__.py +++ b/tensorflow/contrib/layers/__init__.py @@ -35,6 +35,7 @@ See the @{$python/contrib.layers} guide. @@fully_connected @@GDN @@gdn +@@images_to_sequence @@layer_norm @@linear @@max_pool2d @@ -50,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/python/layers/embedding_ops.py b/tensorflow/contrib/layers/python/layers/embedding_ops.py index b62e3050cd7003f1ba72061b133ff9b5d6b616da..ffa208540dae975cb139ad6d76dcf392678ba0ee 100644 --- a/tensorflow/contrib/layers/python/layers/embedding_ops.py +++ b/tensorflow/contrib/layers/python/layers/embedding_ops.py @@ -470,7 +470,7 @@ def embedding_lookup_unique(params, ids, name=None): ids = ops.convert_to_tensor(ids) shape = array_ops.shape(ids) ids_flat = array_ops.reshape( - ids, math_ops.reduce_prod(shape, keep_dims=True)) + ids, math_ops.reduce_prod(shape, keepdims=True)) unique_ids, idx = array_ops.unique(ids_flat) unique_embeddings = embedding_ops.embedding_lookup(params, unique_ids) embeds_flat = array_ops.gather(unique_embeddings, idx) diff --git a/tensorflow/contrib/layers/python/layers/encoders.py b/tensorflow/contrib/layers/python/layers/encoders.py index 89c9d37bd09cb6c43eebb91f3a16600eae9cb490..f42112206d0db9d2e42bd4cff19f6a6533951d46 100644 --- a/tensorflow/contrib/layers/python/layers/encoders.py +++ b/tensorflow/contrib/layers/python/layers/encoders.py @@ -125,7 +125,7 @@ def embed_sequence(ids, `reuse` is `None` or `False`. """ if not (reuse or (vocab_size and embed_dim)): - raise ValueError('Must specify vocab size and embedding dimension when not' + raise ValueError('Must specify vocab size and embedding dimension when not ' 'reusing. Got vocab_size=%s and embed_dim=%s' % ( vocab_size, embed_dim)) with variable_scope.variable_scope( diff --git a/tensorflow/contrib/layers/python/layers/feature_column.py b/tensorflow/contrib/layers/python/layers/feature_column.py index b7d34d6435789e54403926a342481971e854b449..9ccb589d698ad83c9654f5523ccdcb35b031b3da 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column.py +++ b/tensorflow/contrib/layers/python/layers/feature_column.py @@ -154,6 +154,7 @@ from tensorflow.python.ops import string_ops from tensorflow.python.ops import variables from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import deprecation +from tensorflow.python.util import nest # Imports the core `InputLayer` symbol in contrib during development. @@ -554,28 +555,70 @@ def sparse_column_with_integerized_feature(column_name, class _SparseColumnHashed(_SparseColumn): """See `sparse_column_with_hash_bucket`.""" + def __new__(cls, + column_name, + is_integerized=False, + bucket_size=None, + lookup_config=None, + combiner="sum", + dtype=dtypes.string, + hash_keys=None): + if hash_keys is not None: + if not isinstance(hash_keys, list) or not hash_keys: + raise ValueError("hash_keys must be a non-empty list.") + if (any([not isinstance(key_pair, list) for key_pair in hash_keys]) or + any([len(key_pair) != 2 for key_pair in hash_keys]) or + any([not isinstance(key, int) for key in nest.flatten(hash_keys)])): + raise ValueError( + "Each element of hash_keys must be a pair of integers.") + obj = super(_SparseColumnHashed, cls).__new__( + cls, + column_name, + is_integerized=is_integerized, + bucket_size=bucket_size, + lookup_config=lookup_config, + combiner=combiner, + dtype=dtype) + obj.hash_keys = hash_keys + return obj + def _do_transform(self, input_tensor): if self.dtype.is_integer: sparse_values = string_ops.as_string(input_tensor.values) else: sparse_values = input_tensor.values - sparse_id_values = string_ops.string_to_hash_bucket_fast( - sparse_values, self.bucket_size, name="lookup") - return sparse_tensor_py.SparseTensor(input_tensor.indices, sparse_id_values, - input_tensor.dense_shape) + if self.hash_keys: + result = [] + for key in self.hash_keys: + sparse_id_values = string_ops.string_to_hash_bucket_strong( + sparse_values, self.bucket_size, key) + result.append( + sparse_tensor_py.SparseTensor(input_tensor.indices, + sparse_id_values, + input_tensor.dense_shape)) + return sparse_ops.sparse_concat(axis=1, sp_inputs=result, name="lookup") + else: + sparse_id_values = string_ops.string_to_hash_bucket_fast( + sparse_values, self.bucket_size, name="lookup") + return sparse_tensor_py.SparseTensor( + input_tensor.indices, sparse_id_values, input_tensor.dense_shape) def sparse_column_with_hash_bucket(column_name, hash_bucket_size, combiner="sum", - dtype=dtypes.string): + dtype=dtypes.string, + hash_keys=None): """Creates a _SparseColumn with hashed bucket configuration. Use this when your sparse features are in string or integer format, but you don't have a vocab file that maps each value to an integer ID. output_id = Hash(input_feature_string) % bucket_size + When hash_keys is set, multiple integer IDs would be created with each key + pair in the `hash_keys`. This is useful to reduce the collision of hashed ids. + Args: column_name: A string defining sparse column name. hash_bucket_size: An int that is > 1. The number of buckets. @@ -588,6 +631,9 @@ def sparse_column_with_hash_bucket(column_name, * "sqrtn": do l2 normalization on features in the column For more information: `tf.embedding_lookup_sparse`. dtype: The type of features. Only string and integer types are supported. + hash_keys: The hash keys to use. It is a list of lists of two uint64s. If + None, simple and fast hashing algorithm is used. Otherwise, multiple + strong hash ids would be produced with each two unit64s in this argument. Returns: A _SparseColumn with hashed bucket configuration @@ -600,7 +646,8 @@ def sparse_column_with_hash_bucket(column_name, column_name, bucket_size=hash_bucket_size, combiner=combiner, - dtype=dtype) + dtype=dtype, + hash_keys=hash_keys) class _SparseColumnKeys(_SparseColumn): diff --git a/tensorflow/contrib/layers/python/layers/feature_column_test.py b/tensorflow/contrib/layers/python/layers/feature_column_test.py index fc8f153fe3abdc83aca5abfa9a4bb5f5d5531480..1de9ab705655db9863d9c7d2630f24283c83d44d 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column_test.py +++ b/tensorflow/contrib/layers/python/layers/feature_column_test.py @@ -329,6 +329,55 @@ class FeatureColumnTest(test.TestCase): self.assertEqual(one_hot.sparse_id_column.name, "ids_weighted_by_weights") self.assertEqual(one_hot.length, 3) + def testOneHotColumnWithSparseColumnWithHashKeys(self): + input_values = ["marlo", "unknown", "omar"] + inputs = constant_op.constant(input_values) + hash_keys = [[10, 20], [20, 30]] + hash_column = fc.sparse_column_with_hash_bucket( + column_name="ids", hash_bucket_size=10, hash_keys=hash_keys) + columns_to_tensors = {} + columns_to_tensors["ids"] = inputs + hash_column.insert_transformed_feature(columns_to_tensors) + self.assertEqual(len(columns_to_tensors), 2) + self.assertTrue(hash_column in columns_to_tensors) + + one_hot_column = fc.one_hot_column(hash_column) + one_hot_output = one_hot_column._to_dnn_input_layer( + columns_to_tensors[hash_column]) + + expected = np.array([[0., 1., 0., 0., 0., 0., 0., 1., 0., + 0.], [0., 1., 0., 0., 0., 0., 0., 0., 0., 1.], + [1., 0., 0., 0., 0., 0., 0., 0., 0., 1.]]) + with self.test_session() as sess: + one_hot_value = sess.run(one_hot_output) + self.assertTrue(np.array_equal(one_hot_value, expected)) + + def testSparseColumnWithHashKeysWithUnexpectedHashKeys(self): + with self.assertRaisesRegexp(ValueError, + "hash_keys must be a non-empty list."): + fc.sparse_column_with_hash_bucket( + column_name="ids", hash_bucket_size=100, hash_keys=[]) + + with self.assertRaisesRegexp(ValueError, + "hash_keys must be a non-empty list."): + fc.sparse_column_with_hash_bucket( + column_name="ids", hash_bucket_size=100, hash_keys=1) + + with self.assertRaisesRegexp( + ValueError, "Each element of hash_keys must be a pair of integers."): + fc.sparse_column_with_hash_bucket( + column_name="ids", hash_bucket_size=100, hash_keys=[1, 2]) + + with self.assertRaisesRegexp( + ValueError, "Each element of hash_keys must be a pair of integers."): + fc.sparse_column_with_hash_bucket( + column_name="ids", hash_bucket_size=100, hash_keys=["key"]) + + with self.assertRaisesRegexp( + ValueError, "Each element of hash_keys must be a pair of integers."): + fc.sparse_column_with_hash_bucket( + column_name="ids", hash_bucket_size=100, hash_keys=[[1, 2.0]]) + def testMissingValueInOneHotColumnForWeightedSparseColumn(self): # Github issue 12583 ids = fc.sparse_column_with_keys("ids", ["marlo", "omar", "stringer"]) diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py index c8e3307ee8b5ded30dc864c4e69452f58685b8f0..350bcb3bca11b4cad18ce863ab1496076477aa3c 100644 --- a/tensorflow/contrib/layers/python/layers/layers.py +++ b/tensorflow/contrib/layers/python/layers/layers.py @@ -51,7 +51,6 @@ from tensorflow.python.ops import standard_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables as tf_variables from tensorflow.python.training import moving_averages -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. @@ -60,12 +59,12 @@ __all__ = [ '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', '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' + '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' @@ -518,8 +517,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'. @@ -779,7 +778,7 @@ def batch_norm(inputs, else: if data_format == DATA_FORMAT_NCHW: mean, variance = nn.weighted_moments( - inputs, moments_axes, batch_weights, keep_dims=True) + inputs, moments_axes, batch_weights, keepdims=True) mean = array_ops.reshape(mean, [-1]) variance = array_ops.reshape(variance, [-1]) else: @@ -1415,10 +1414,11 @@ def dense_to_sparse(tensor, eos_token=0, outputs_collections=None, scope=None): 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: + 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))) + 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) @@ -2185,6 +2185,36 @@ 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, @@ -2664,6 +2694,39 @@ 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. @@ -2684,7 +2747,7 @@ def softmax(logits, scope=None): logits_2d = array_ops.reshape(logits, [-1, num_logits]) predictions = nn.softmax(logits_2d) predictions = array_ops.reshape(predictions, array_ops.shape(logits)) - if context.in_graph_mode(): + if not context.executing_eagerly(): predictions.set_shape(logits.get_shape()) return predictions @@ -2774,9 +2837,9 @@ def spatial_softmax(features, 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]) @@ -2877,6 +2940,53 @@ def unit_norm(inputs, dim, epsilon=1e-7, scope=None): return math_ops.div(inputs, array_ops.tile(lengths, multiples)) +@add_arg_scope +def maxout(inputs, num_units, axis=-1, scope=None): + """Adds a maxout op from https://arxiv.org/abs/1302.4389 + + "Maxout Networks" Ian J. Goodfellow, David Warde-Farley, Mehdi Mirza, Aaron + Courville, + Yoshua Bengio + + Usually the operation is performed in the filter/channel dimension. This can + also be + used after fully-connected layers to reduce number of features. + + Arguments: + inputs: Tensor input + num_units: Specifies how many features will remain after maxout + in the `axis` dimension (usually channel). + This must be multiple of number of `axis`. + axis: The dimension where max pooling will be performed. Default is the + last dimension. + scope: Optional scope for variable_scope. + + Returns: + A `Tensor` representing the results of the pooling operation. + + Raises: + ValueError: if num_units is not multiple of number of features. + """ + with variable_scope.variable_scope(scope, 'MaxOut', [inputs]): + inputs = ops.convert_to_tensor(inputs) + shape = inputs.get_shape().as_list() + num_channels = shape[axis] + if num_channels % num_units: + raise ValueError('number of features({}) is not ' + 'a multiple of num_units({})'.format( + num_channels, num_units)) + shape[axis] = -1 + shape += [num_channels // num_units] + + # Dealing with batches with arbitrary sizes + for i in range(len(shape)): + if shape[i] is None: + shape[i] = array_ops.shape(inputs)[i] + outputs = math_ops.reduce_max( + array_ops.reshape(inputs, shape), -1, keepdims=False) + return outputs + + def poincare_normalize(x, axis=1, epsilon=1e-5, name=None): """Project into the Poincare ball with norm <= 1.0 - epsilon. @@ -2909,7 +3019,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) diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py index f05f63224636967e99673c04fb9de60af34e4268..997f910a2a97567adbd7ffa3e81a31d2ae0bad7e 100644 --- a/tensorflow/contrib/layers/python/layers/layers_test.py +++ b/tensorflow/contrib/layers/python/layers/layers_test.py @@ -708,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 @@ -1308,11 +1308,13 @@ class DenseToSparseTest(test.TestCase): 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) + 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): @@ -2194,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) @@ -2948,6 +2950,28 @@ class GDNTest(test.TestCase): 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): def testInvalidDataFormat(self): @@ -3416,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): @@ -4084,5 +4135,31 @@ class LegacyFullyConnectedTest(test.TestCase): _layers.legacy_fully_connected(x, 2, activation_fn=nn_ops.softmax) +class MaxOutTest(test.TestCase): + + def test_simple(self): + inputs = random_ops.random_uniform((64, 10, 36), seed=1) + graph = _layers.maxout(inputs, num_units=3) + self.assertEqual(graph.get_shape().as_list(), [64, 10, 3]) + + def test_fully_connected(self): + inputs = random_ops.random_uniform((64, 50), seed=1) + graph = _layers.fully_connected(inputs, 50) + graph = _layers.maxout(graph, num_units=10) + self.assertEqual(graph.get_shape().as_list(), [64, 10]) + + def test_nchw(self): + inputs = random_ops.random_uniform((10, 100, 100, 3), seed=1) + graph = _layers.conv2d(inputs, 10, 3, padding='SAME') + graph = _layers.maxout(graph, num_units=1) + self.assertEqual(graph.get_shape().as_list(), [10, 100, 100, 1]) + + def test_invalid_shape(self): + inputs = random_ops.random_uniform((10, 100, 100, 3), seed=1) + graph = _layers.conv2d(inputs, 3, 10) + with self.assertRaisesRegexp(ValueError, 'number of features'): + graph = _layers.maxout(graph, num_units=2) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/layers/python/layers/optimizers.py b/tensorflow/contrib/layers/python/layers/optimizers.py index cdceea6fee5bdb5aeb6537ea55d25ccf107def4c..69d927e1b3001d14dd1af2f890b07c1a57ab2cfc 100644 --- a/tensorflow/contrib/layers/python/layers/optimizers.py +++ b/tensorflow/contrib/layers/python/layers/optimizers.py @@ -41,7 +41,7 @@ OPTIMIZER_CLS_NAMES = { "Adagrad": train.AdagradOptimizer, "Adam": train.AdamOptimizer, "Ftrl": train.FtrlOptimizer, - "Momentum": lambda lr: train.MomentumOptimizer(lr, momentum=0.9), + "Momentum": lambda learning_rate: train.MomentumOptimizer(learning_rate, momentum=0.9), # pylint: disable=line-too-long "RMSProp": train.RMSPropOptimizer, "SGD": train.GradientDescentOptimizer, } diff --git a/tensorflow/contrib/layers/python/layers/optimizers_test.py b/tensorflow/contrib/layers/python/layers/optimizers_test.py index 1ea25bd1a5685eb6f840e621b5739029a660aa0f..a4461a20e54c289886f1a1beb255de12fc054afe 100644 --- a/tensorflow/contrib/layers/python/layers/optimizers_test.py +++ b/tensorflow/contrib/layers/python/layers/optimizers_test.py @@ -61,7 +61,8 @@ class OptimizersTest(test.TestCase): optimizers = [ "SGD", gradient_descent.GradientDescentOptimizer, gradient_descent.GradientDescentOptimizer(learning_rate=0.1), - lambda lr: gradient_descent.GradientDescentOptimizer(learning_rate=lr) + lambda lr: gradient_descent.GradientDescentOptimizer(learning_rate=lr), + "Momentum" ] for optimizer in optimizers: with ops.Graph().as_default() as g: diff --git a/tensorflow/contrib/learn/BUILD b/tensorflow/contrib/learn/BUILD index 3c782b54a8559a6aac19d12ea11a9c76bffdb9c3..b05f5eeaeee8fb927970b608f65495f33d63f764 100644 --- a/tensorflow/contrib/learn/BUILD +++ b/tensorflow/contrib/learn/BUILD @@ -5,6 +5,8 @@ licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) +load("//tensorflow:tensorflow.bzl", "py_test") + package(default_visibility = [ "//engedu/ml/tf_from_scratch:__pkg__", "//tensorflow:internal", @@ -115,6 +117,7 @@ py_test( size = "small", srcs = ["python/learn/learn_io/data_feeder_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":learn", "//tensorflow/python:client_testlib", @@ -170,6 +173,7 @@ tf_py_test( "//tensorflow/python:variables", "//tensorflow/python/estimator", ], + tags = ["no_windows"], # TODO: needs investigation on Windows ) py_test( @@ -188,6 +192,7 @@ py_test( size = "small", srcs = ["python/learn/graph_actions_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":learn", "//tensorflow/contrib/framework:framework_py", @@ -388,6 +393,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", @@ -425,6 +431,10 @@ py_test( size = "medium", srcs = ["python/learn/estimators/kmeans_test.py"], srcs_version = "PY2AND3", + tags = [ + "noasan", # b/73741358 + "nomac", + ], deps = [ ":learn", "//tensorflow/python:array_ops", @@ -583,6 +593,7 @@ py_test( size = "small", srcs = ["python/learn/learn_io/io_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":learn", "//tensorflow/contrib/learn/python/learn/datasets", @@ -812,6 +823,7 @@ py_test( size = "small", srcs = ["python/learn/utils/saved_model_export_utils_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":learn", "//tensorflow/contrib/layers:layers_py", diff --git a/tensorflow/contrib/learn/README.md b/tensorflow/contrib/learn/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d516bffc5e0327a3400068b35de5503e5a925a54 --- /dev/null +++ b/tensorflow/contrib/learn/README.md @@ -0,0 +1,143 @@ +EVERYTHING IN THIS DIRECTORY IS DEPRECATED. + +Using functions or classes will result in warnings. + +Instructions for converting to current alternatives are included in the +warnings. A high-level overview is below. + +## Canned Estimators + +Many canned estimators (subclasses of `Estimator`) have equivalents in core: +`DNNClassifier`, `DNNRegressor`, `DNNEstimator`, `LinearClassifier`, +`LinearRegressor`, `DNNLinearCombinedClassifier` and +`DNNLinearCombinedRegressor`. They are exposed under `tf.estimator`. +`DNNEstimator`, `LinearEstimator` and `DNNLinearCombinedEstimator` +are exposed under `tf.contrib.estimator`. + +To migrate to the new api, users need to take the following steps: + +* Replace `tf.contrib.learn` with `tf.estimator`. +* If you subclass any of the estimators, stop doing that. You should be able to + write a factory method that returns a canned estimator instead. If this is not + possible (if you override methods from the canned estimator), consider writing + a custom estimator instead. See `tf.estimator.Estimator`. +* Set `loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE` to preserve loss + reduction as the average over batch. +* Some optimizer-related arguments are no longer passed in the estimator + constructor. Instead, we provide methods that perform the same job by wrapping + an optimizer. Specifically: + * `gradient_clip_norm`: Use `tf.contrib.estimator.clip_gradients_by_norm` + * `embedding_lr_multipliers`: Not supported. + Other arguments: + * `input_layer_min_slice_size`: Replaced by `input_layer_partitioner` + * `enable_centered_bias`: Not supported. Dropping this argument is unlikely to + harm your model. + * `feature_engineering_fn`: Not supported. You can call your + `feature_engineering_fn` inside your input_fn: + ```python + def new_input_fn(): + features, labels = old_input_fn() + return feature_engineering_fn(features, labels) + ``` +* Use `tf.reshape` to reshape labels in your `input_fn`. `tf.estimator` + classifiers and regressors expect labels as a 2D Tensor of shape + `[batch_size, 1]`, or `[batch_size, n_labels]`. In contrast, + `tf.contrib.learn` classifiers and regressors supported labels with shape + `[batch_size]`. +* If you pass custom metrics from the `evaluate()` method call, use + `tf.contrib.estimator.add_metrics`. +* Replace your `serving_input_fn` with a `serving_input_receiver_fn`. + Note this should be entirely distinct from your training `input_fn`, so if you + previously had one `input_fn` with different "modes", you should now factor + that apart. Where the former returned either a simple `(features, labels)` + tuple or `InputFnOps`, you should now return a `ServingInputReceiver`. + If you were generating your `serving_input_fn` using the + `build_parsing_serving_input_fn` helper, you can simply drop in the + replacement `build_parsing_serving_input_receiver_fn`. + +Some remaining estimators/classes: + +* `DynamicRnnEstimator`: Consider a custom `model_fn`. +* `KMeansClustering`: Use `tf.contrib.factorization.KMeansClustering`. +* `LogisticRegressor`: Not supported. Instead, use `binary_classification_head` + with a custom `model_fn`, or with `DNNEstimator`. +* `StateSavingRnnEstimator`: Consider a custom `model_fn`. +* SVM: Consider a custom `model_fn`. +* `LinearComposableModel` and `DNNComposableModel`: Not supported. + Consider `tf.contrib.estimator.DNNEstimator`, or write a custom model_fn. +* `MetricSpec`: Deprecated. For adding custom metrics to canned Estimators, use + `tf.contrib.estimator.add_metrics`. + +## Estimator +`tf.contrib.learn.Estimator` is migrated to `tf.estimator.Estimator`. + +To migrate, users need to take the following steps: + +* Replace `tf.contrib.learn.Estimator` with `tf.estimator.Estimator`. +* If you pass a `config` argument to `Estimator`, this must be + `tf.estimator.RunConfig`. You may need to edit your code accordingly. +* Edit your `model_fn` to return `tf.estimator.EstimatorSpec`. Refer to + `EstimatorSpec` for documentation of specific fields. +* If your `model_fn` uses the `mode` argument, use `tf.estimator.ModeKeys`. + +Some related classes: +* `Evaluable`, `Trainable`: Not supported, merged into `tf.estimator.Estimator`. +* ExportStrategy: Replaced by `tf.estimator.Exporter`. + +## Head/MultiHead +These classes are now supported under `tf.contrib.estimator`, e.g. +`tf.contrib.estimator.multi_class_head` and `tf.contrib.estimator.multi_head`. + +Some differences: + +* `multi_class_head`: If you use `tf.contrib.learn.multi_class_head` with + `n_classes=2`, switch to `tf.contrib.estimator.binary_classification_head`. +* `loss_only_head`: Not supported. +* `poisson_regression_head`: Not supported (yet). +* `binary_svm_head`: Not supported (yet). +* `no_op_train_fn`: Replace it with `tf.no_op`. + +Some arguments are renamed, please refer to documentation. In addition: + +* `loss_fn`: Supported for `multi_label_head`. If you need it for other heads, + please open an issue. +* `metric_class_ids`: Not supported (yet). +* `enable_centered_bias`: Not supported. Dropping this argument is unlikely to + harm your model. +* `label_name`: Not needed in `tf.estimator`. If you don’t use `multi_head`, + drop this argument. If you use `multi_head`, refer to + `tf.contrib.estimator.multi_head` documentation. + +## Experiment Class - Distributed Training Tooling + +Switch to `tf.estimator.train_and_evaluate`. Some differences: + +* Most of the constructor arguments, like `train_input_fn`, `eval_input_fn`, + should be wrapped into `tf.estimator.TrainSpec` and `tf.estimator.EvalSpec`. +* Remove the `experiment_fn`. Instead, create the `Estimator`, + `train_spec` and `eval_spec`, then call `tf.estimator.train_and_evaluate` + directly. +* Inside `tf.estimator.EvalSpec`, the `exporter` field is the replacement + for `export_strategy`. To be precise, `tf.estimator.LatestExporter` is the + replacement for `tf.contrib.learn.make_export_strategy`. If you want to export + only at the end of training use `tf.estimator.FinalExporter`. +* If the `TF_CONFIG` environment variable is constructed manually, please read + the `train_and_evaluate` documentation for the new requirementds (in + particular, the chief node and evaluator node). + +## Others Classes and Functions + +* `tf.contrib.learn.datasets` is deprecated. We are adding ready to use datasets + to tensorflow/models. Many smaller datasets are available from other sources, + such as scikits.learn. Some Python processing may have to be written, but this + is straightforward to implement using the standard modules. +* `tf.contrib.learn.preprocessing`: Deprecated. The python-only preprocessing + functions are not a good fit for TensorFlow. Please use `tf.data`, and + consider tensorflow/transform for more complex use cases. +* `tf.contrib.learn.models`: Not supported, use canned estimators instead. +* `tf.contrib.learn.monitors`: Implement `SessionRunHook` instead. Hook + implementations are in `tf.train`. +* `tf.contrib.learn.learn_io`: Use the methods in `tf.estimator.inputs`, such as + `tf.estimator.inputs.numpy_input_fn`. Some utility functions have no + equivalent, we encourage the use of `tf.data`. + diff --git a/tensorflow/contrib/learn/__init__.py b/tensorflow/contrib/learn/__init__.py index 3698af027e38f1063ad829c26eb179734968f813..79bd73faaf1301a2fc4999b64f88d30542577980 100644 --- a/tensorflow/contrib/learn/__init__.py +++ b/tensorflow/contrib/learn/__init__.py @@ -13,8 +13,11 @@ # limitations under the License. # ============================================================================== -# TODO(ptucker,ipolosukhin): Improve descriptions. -"""High level API for learning. +"""High level API for learning (DEPRECATED). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. See the @{$python/contrib.learn} guide. diff --git a/tensorflow/contrib/learn/python/__init__.py b/tensorflow/contrib/learn/python/__init__.py index bbebd5ab9792cb937219cf937f08c4d4e6e44a92..df23aeb2c433c2b4392f706730f715246ce01cea 100644 --- a/tensorflow/contrib/learn/python/__init__.py +++ b/tensorflow/contrib/learn/python/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""High level API for learning with TensorFlow.""" +"""High level API for learning with TensorFlow (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/__init__.py b/tensorflow/contrib/learn/python/learn/__init__.py index cdc67c77d5fd1df61016835dc75ba44feb458cf9..76e0e8ac8f19026086959f3b197cfd1a81e65a3e 100644 --- a/tensorflow/contrib/learn/python/learn/__init__.py +++ b/tensorflow/contrib/learn/python/learn/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""High level API for learning with TensorFlow.""" +"""High level API for learning with TensorFlow (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/basic_session_run_hooks.py b/tensorflow/contrib/learn/python/learn/basic_session_run_hooks.py index 2284ec46e971731af74f17678fc0d1d3888419e2..fed1c44d1970bf07c808ace817aa9972d7776d88 100644 --- a/tensorflow/contrib/learn/python/learn/basic_session_run_hooks.py +++ b/tensorflow/contrib/learn/python/learn/basic_session_run_hooks.py @@ -12,20 +12,47 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Some common SessionRunHook classes.""" +"""Some common SessionRunHook classes (deprected). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.python.training import basic_session_run_hooks +from tensorflow.python.util.deprecation import deprecated_alias # pylint: disable=invalid-name -LoggingTensorHook = basic_session_run_hooks.LoggingTensorHook -StopAtStepHook = basic_session_run_hooks.StopAtStepHook -CheckpointSaverHook = basic_session_run_hooks.CheckpointSaverHook -StepCounterHook = basic_session_run_hooks.StepCounterHook -NanLossDuringTrainingError = basic_session_run_hooks.NanLossDuringTrainingError -NanTensorHook = basic_session_run_hooks.NanTensorHook -SummarySaverHook = basic_session_run_hooks.SummarySaverHook +LoggingTensorHook = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.LoggingTensorHook', + 'tf.train.LoggingTensorHook', + basic_session_run_hooks.LoggingTensorHook) +StopAtStepHook = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.StopAtStepHook', + 'tf.train.StopAtStepHook', + basic_session_run_hooks.StopAtStepHook) +CheckpointSaverHook = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.CheckpointSaverHook', + 'tf.train.CheckpointSaverHook', + basic_session_run_hooks.CheckpointSaverHook) +StepCounterHook = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.StepCounterHook', + 'tf.train.StepCounterHook', + basic_session_run_hooks.StepCounterHook) +NanLossDuringTrainingError = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.NanLossDuringTrainingError', + 'tf.train.NanLossDuringTrainingError', + basic_session_run_hooks.NanLossDuringTrainingError) +NanTensorHook = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.NanTensorHook', + 'tf.train.NanTensorHook', + basic_session_run_hooks.NanTensorHook) +SummarySaverHook = deprecated_alias( + 'tf.contrib.learn.basic_session_run_hooks.SummarySaverHook', + 'tf.train.SummarySaverHook', + basic_session_run_hooks.SummarySaverHook) # pylint: enable=invalid-name diff --git a/tensorflow/contrib/learn/python/learn/datasets/__init__.py b/tensorflow/contrib/learn/python/learn/datasets/__init__.py index 7240b0de149051afa045a8113f9e9b212840c311..3c34712ac859d32f549468345950a93d2ed2aa56 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/__init__.py +++ b/tensorflow/contrib/learn/python/learn/datasets/__init__.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Dataset utilities and synthetic/reference datasets.""" +"""Dataset utilities and synthetic/reference datasets (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -27,6 +32,7 @@ from tensorflow.contrib.learn.python.learn.datasets import base from tensorflow.contrib.learn.python.learn.datasets import mnist from tensorflow.contrib.learn.python.learn.datasets import synthetic from tensorflow.contrib.learn.python.learn.datasets import text_datasets +from tensorflow.python.util.deprecation import deprecated # Export load_iris and load_boston. load_iris = base.load_iris @@ -51,6 +57,7 @@ SYNTHETIC = { } +@deprecated(None, 'Please use tf.data.') def load_dataset(name, size='small', test_with_fake_data=False): """Loads dataset by name. @@ -73,8 +80,9 @@ def load_dataset(name, size='small', test_with_fake_data=False): return DATASETS[name]() +@deprecated(None, 'Please use tf.data.') def make_dataset(name, n_samples=100, noise=None, seed=42, *args, **kwargs): - """Creates binary synthetic datasets + """Creates binary synthetic datasets. Args: name: str, name of the dataset to generate diff --git a/tensorflow/contrib/learn/python/learn/datasets/base.py b/tensorflow/contrib/learn/python/learn/datasets/base.py index 18bf16e246bcb6c0a6a4ce75bc5c28d4e0d045e5..3b5c9b97c08a388e1f35249967b6cab26861f100 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/base.py +++ b/tensorflow/contrib/learn/python/learn/datasets/base.py @@ -12,7 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Base utilities for loading datasets.""" + +"""Base utilities for loading datasets (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -23,19 +29,20 @@ 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 +from tensorflow.python.util.deprecation import deprecated + Dataset = collections.namedtuple('Dataset', ['data', 'target']) Datasets = collections.namedtuple('Datasets', ['train', 'validation', 'test']) +@deprecated(None, 'Use tf.data instead.') def load_csv_with_header(filename, target_dtype, features_dtype, @@ -55,6 +62,7 @@ def load_csv_with_header(filename, return Dataset(data=data, target=target) +@deprecated(None, 'Use tf.data instead.') def load_csv_without_header(filename, target_dtype, features_dtype, @@ -72,6 +80,7 @@ def load_csv_without_header(filename, return Dataset(data=data, target=target) +@deprecated(None, 'Use tf.data instead.') def shrink_csv(filename, ratio): """Create a smaller dataset of only 1/ratio of original data.""" filename_small = filename.replace('.', '_small.') @@ -86,6 +95,7 @@ def shrink_csv(filename, ratio): i += 1 +@deprecated(None, 'Use scikits.learn.datasets.') def load_iris(data_path=None): """Load Iris dataset. @@ -102,6 +112,7 @@ def load_iris(data_path=None): data_path, target_dtype=np.int, features_dtype=np.float) +@deprecated(None, 'Use scikits.learn.datasets.') def load_boston(data_path=None): """Load Boston housing dataset. @@ -118,7 +129,12 @@ def load_boston(data_path=None): 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): +@deprecated(None, 'Use the retry module or similar alternatives.') +def retry(initial_delay, + max_delay, + factor=2.0, + jitter=0.25, + is_retriable=None): """Simple decorator for wrapping retriable functions. Args: @@ -154,7 +170,7 @@ def retry(initial_delay, max_delay, factor=2.0, jitter=0.25, is_retriable=None): for delay in delays(): try: return fn(*args, **kwargs) - except Exception as e: # pylint: disable=broad-except) + except Exception as e: # pylint: disable=broad-except if is_retriable is None: continue @@ -178,11 +194,13 @@ def _is_retriable(e): return isinstance(e, IOError) and e.errno in _RETRIABLE_ERRNOS +@deprecated(None, 'Please use urllib or similar directly.') @retry(initial_delay=1.0, max_delay=16.0, is_retriable=_is_retriable) def urlretrieve_with_retry(url, filename=None): return urllib.request.urlretrieve(url, filename) +@deprecated(None, 'Please write your own downloading logic.') def maybe_download(filename, work_directory, source_url): """Download the data from source url, unless it's already here. diff --git a/tensorflow/contrib/learn/python/learn/datasets/mnist.py b/tensorflow/contrib/learn/python/learn/datasets/mnist.py index 37f9175015a239f763c7721cf36ab8063c0a3e32..abbb44c2f5b701829ce16f64eadd8ebc04c84e2c 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/mnist.py +++ b/tensorflow/contrib/learn/python/learn/datasets/mnist.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Functions for downloading and reading MNIST data.""" +"""Functions for downloading and reading MNIST data (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -27,6 +32,7 @@ from tensorflow.contrib.learn.python.learn.datasets import base from tensorflow.python.framework import dtypes from tensorflow.python.framework import random_seed from tensorflow.python.platform import gfile +from tensorflow.python.util.deprecation import deprecated # CVDF mirror of http://yann.lecun.com/exdb/mnist/ DEFAULT_SOURCE_URL = 'https://storage.googleapis.com/cvdf-datasets/mnist/' @@ -37,6 +43,7 @@ def _read32(bytestream): return numpy.frombuffer(bytestream.read(4), dtype=dt)[0] +@deprecated(None, 'Please use tf.data to implement this functionality.') def extract_images(f): """Extract the images into a 4D uint8 numpy array [index, y, x, depth]. @@ -65,6 +72,7 @@ def extract_images(f): return data +@deprecated(None, 'Please use tf.one_hot on tensors.') def dense_to_one_hot(labels_dense, num_classes): """Convert class labels from scalars to one-hot vectors.""" num_labels = labels_dense.shape[0] @@ -74,6 +82,7 @@ def dense_to_one_hot(labels_dense, num_classes): return labels_one_hot +@deprecated(None, 'Please use tf.data to implement this functionality.') def extract_labels(f, one_hot=False, num_classes=10): """Extract the labels into a 1D uint8 numpy array [index]. @@ -103,7 +112,15 @@ def extract_labels(f, one_hot=False, num_classes=10): class DataSet(object): + """Container class for a dataset (deprecated). + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'Please use alternatives such as official/mnist/dataset.py' + ' from tensorflow/models.') def __init__(self, images, labels, @@ -210,6 +227,8 @@ class DataSet(object): return self._images[start:end], self._labels[start:end] +@deprecated(None, 'Please use alternatives such as official/mnist/dataset.py' + ' from tensorflow/models.') def read_data_sets(train_dir, fake_data=False, one_hot=False, @@ -275,5 +294,7 @@ def read_data_sets(train_dir, return base.Datasets(train=train, validation=validation, test=test) +@deprecated(None, 'Please use alternatives such as official/mnist/dataset.py' + ' from tensorflow/models.') def load_mnist(train_dir='MNIST-data'): return read_data_sets(train_dir) diff --git a/tensorflow/contrib/learn/python/learn/datasets/produce_small_datasets.py b/tensorflow/contrib/learn/python/learn/datasets/produce_small_datasets.py index 6e0ba38941ce4650ede9f7210e284bde2ed8e6a9..a4848fa64a72f031ef35c0c3256e97a7326acd60 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/produce_small_datasets.py +++ b/tensorflow/contrib/learn/python/learn/datasets/produce_small_datasets.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Produce DBpedia datasets of a smaller size.""" +"""Produce DBpedia datasets of a smaller size (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/datasets/synthetic.py b/tensorflow/contrib/learn/python/learn/datasets/synthetic.py index 9a843168c27d9cae3f55efe4fe4c688d86c745f3..6a0e3350b3d1052249160a2a997a76de7a5040c3 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/synthetic.py +++ b/tensorflow/contrib/learn/python/learn/datasets/synthetic.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Synthetic dataset generators.""" +"""Synthetic dataset generators (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -21,8 +26,10 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.learn.python.learn.datasets.base import Dataset +from tensorflow.python.util.deprecation import deprecated +@deprecated(None, 'Consider using synthetic datasets from scikits.learn.') def circles(n_samples=100, noise=None, seed=None, @@ -93,6 +100,7 @@ def circles(n_samples=100, return Dataset(data=X[indices], target=y[indices]) +@deprecated(None, 'Consider using synthetic datasets from scikits.learn.') def spirals(n_samples=100, noise=None, seed=None, diff --git a/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py b/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py index 19791d7759e05f477f1b79220f29831fb2240e1b..5809995c8c7d8e72eb47ee88a72547bae7fd3594 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py +++ b/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py @@ -136,7 +136,7 @@ class SyntheticTest(test.TestCase): self.assertRaises(AssertionError, np.testing.assert_array_equal, spir0.data, spir1.data) - def test_spirals(self): + def test_spirals_synthetic(self): synthetic.spirals(3) diff --git a/tensorflow/contrib/learn/python/learn/datasets/text_datasets.py b/tensorflow/contrib/learn/python/learn/datasets/text_datasets.py index 2596a2ecaf1572506504831e8b08fab9b5dbc119..ce9466301728082f8e9d99c90989ba8fe623bcf0 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/text_datasets.py +++ b/tensorflow/contrib/learn/python/learn/datasets/text_datasets.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Text datasets.""" +"""Text datasets (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -26,10 +31,12 @@ import numpy as np from tensorflow.contrib.learn.python.learn.datasets import base from tensorflow.python.platform import gfile +from tensorflow.python.util.deprecation import deprecated DBPEDIA_URL = 'https://github.com/le-scientifique/torchDatasets/raw/master/dbpedia_csv.tar.gz' +@deprecated(None, 'See contrib/learn/README.md') def maybe_download_dbpedia(data_dir): """Download if DBpedia data is not present.""" train_path = os.path.join(data_dir, 'dbpedia_csv/train.csv') @@ -41,6 +48,7 @@ def maybe_download_dbpedia(data_dir): tfile.extractall(data_dir) +@deprecated(None, 'See contrib/learn/README.md') def load_dbpedia(size='small', test_with_fake_data=False): """Get DBpedia datasets from CSV files.""" if not test_with_fake_data: diff --git a/tensorflow/contrib/learn/python/learn/estimators/__init__.py b/tensorflow/contrib/learn/python/learn/estimators/__init__.py index 4981750c94c7ac31e23b7a3f71ca30e3c9573a20..3e64595f312bcc2a2e8dcba589fb993a249b684b 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/__init__.py +++ b/tensorflow/contrib/learn/python/learn/estimators/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""An estimator is a rule for calculating an estimate of a given quantity. +"""An estimator is a rule for calculating an estimate of a given quantity (deprecated). + +These classes are deprecated and replaced with `tf.estimator`. + +See [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. # Estimators diff --git a/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py b/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py index 15277415a1ce83dc1d4a334e60fe1933ba244df0..1f0e4663d060a3850e2002b27f809fde1db47e48 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py +++ b/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== -"""sklearn cross-support.""" +"""sklearn cross-support (deprecated).""" from __future__ import absolute_import from __future__ import division @@ -132,6 +132,8 @@ class _TransformerMixin(): class NotFittedError(ValueError, AttributeError): """Exception class to raise if estimator is used before fitting. + USE OF THIS EXCEPTION IS DEPRECATED. + This class inherits from both ValueError and AttributeError to help with exception handling and backward compatibility. diff --git a/tensorflow/contrib/learn/python/learn/estimators/composable_model.py b/tensorflow/contrib/learn/python/learn/estimators/composable_model.py index a02c726c74946d93b8e1726473db746220b00795..1fa58271e2b886cd143683a759145fd750791473 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/composable_model.py +++ b/tensorflow/contrib/learn/python/learn/estimators/composable_model.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""TensorFlow composable models used as building blocks for estimators.""" +"""TensorFlow composable models used as building blocks for estimators (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -34,6 +39,7 @@ from tensorflow.python.ops import nn from tensorflow.python.ops import partitioned_variables from tensorflow.python.ops import variable_scope from tensorflow.python.summary import summary +from tensorflow.python.util.deprecation import deprecated class _ComposableModel(object): @@ -46,6 +52,7 @@ class _ComposableModel(object): _ComposableModel and its subclasses are not part of the public tf.learn API. """ + @deprecated(None, "Please use model_fns in tf.estimator.") def __init__(self, num_label_columns, optimizer, @@ -141,6 +148,10 @@ class _ComposableModel(object): class LinearComposableModel(_ComposableModel): """A _ComposableModel that implements linear regression. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Instances of this class can be used to build estimators through the use of composition. """ @@ -252,6 +263,10 @@ class LinearComposableModel(_ComposableModel): class DNNComposableModel(_ComposableModel): """A _ComposableModel that implements a DNN. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Instances of this class can be used to build estimators through the use of composition. """ diff --git a/tensorflow/contrib/learn/python/learn/estimators/constants.py b/tensorflow/contrib/learn/python/learn/estimators/constants.py index fc69e810244a182b864be856e6720f8584f7aa65..d2548946bc77dea7c452d61c7e2b6e12c3d6239a 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/constants.py +++ b/tensorflow/contrib/learn/python/learn/estimators/constants.py @@ -13,9 +13,11 @@ # limitations under the License. # ============================================================================== -"""Constants regarding Estimators. +"""Constants regarding Estimators (deprecated). -This file is obsoleted in the move of Estimator to core. +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. """ from __future__ import absolute_import from __future__ import division @@ -25,6 +27,8 @@ from __future__ import print_function class ProblemType(object): """Enum-like values for the type of problem that the model solves. + THIS CLASS IS DEPRECATED. + These values are used when exporting the model to produce the appropriate signature function for serving. diff --git a/tensorflow/contrib/learn/python/learn/estimators/debug.py b/tensorflow/contrib/learn/python/learn/estimators/debug.py index 9d5f6c2bf969d7c85d251bf1b06a0307a41b2297..24b067b7e38b12df3d1d0c49f626344217218571 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/debug.py +++ b/tensorflow/contrib/learn/python/learn/estimators/debug.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Debug estimators. +"""Debug estimators (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. Debug estimators are bias-only estimators that can be used for debugging and as simple baselines. @@ -118,6 +122,10 @@ def debug_model_fn(features, labels, mode, params, config=None): class DebugClassifier(estimator.Estimator): """A classifier for TensorFlow Debug models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Example: ```python @@ -237,6 +245,10 @@ class DebugClassifier(estimator.Estimator): class DebugRegressor(estimator.Estimator): """A regressor for TensorFlow Debug models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Example: ```python diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn.py b/tensorflow/contrib/learn/python/learn/estimators/dnn.py index c17b41c0f767e19d9c3635a8f60347a49b297cfb..eabebb7e881558471c343c0573cc9a8f4a425312 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Deep Neural Network estimators.""" +"""Deep Neural Network estimators (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -212,6 +217,10 @@ def _dnn_model_fn(features, labels, mode, params, config=None): class DNNClassifier(estimator.Estimator): """A classifier for TensorFlow DNN models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Example: ```python @@ -521,6 +530,10 @@ class DNNClassifier(estimator.Estimator): class DNNRegressor(estimator.Estimator): """A regressor for TensorFlow DNN models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Example: ```python @@ -796,6 +809,10 @@ class DNNRegressor(estimator.Estimator): class DNNEstimator(estimator.Estimator): """A Estimator for TensorFlow DNN models with user specified _Head. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Example: ```python diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py index 726612235050def6e7addb503cc6646a25de0e42..3d85533d92d17095bae9a69f229171e1bf61ba10 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""TensorFlow estimators for Linear and DNN joined training models.""" +"""TensorFlow estimators for Linear and DNN joined training models (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -372,6 +377,10 @@ def _dnn_linear_combined_model_fn(features, labels, mode, params, config=None): class DNNLinearCombinedEstimator(estimator.Estimator): """An estimator for TensorFlow Linear and DNN joined training models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Note: New users must set `fix_global_step_increment_bug=True` when creating an estimator. @@ -490,6 +499,10 @@ class DNNLinearCombinedEstimator(estimator.Estimator): class DNNLinearCombinedClassifier(estimator.Estimator): """A classifier for TensorFlow Linear and DNN joined training models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Note: New users must set `fix_global_step_increment_bug=True` when creating an estimator. @@ -832,6 +845,10 @@ class DNNLinearCombinedClassifier(estimator.Estimator): class DNNLinearCombinedRegressor(estimator.Estimator): """A regressor for TensorFlow Linear and DNN joined training models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Note: New users must set `fix_global_step_increment_bug=True` when creating an estimator. 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/dynamic_rnn_estimator.py b/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator.py index 69440e823ef1ed2d739f28bc14587891f2de80bb..a703dc66e922d48ceb64edc2a979061b8e45db49 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Estimator for Dynamic RNNs.""" +"""Estimator for Dynamic RNNs (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -540,6 +545,12 @@ def _get_dynamic_rnn_model_fn( class DynamicRnnEstimator(estimator.Estimator): + """Dynamically unrolled RNN (deprecated). + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ def __init__(self, problem_type, diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator.py b/tensorflow/contrib/learn/python/learn/estimators/estimator.py index 8d59fe66d98b2ca7dc143cfdf05d29629e3bf616..7a026a15e4aeea0dde4ed9f7de053a757a0abb58 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Base Estimator class.""" +"""Base Estimator class (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -138,6 +143,7 @@ def _get_input_fn(x, y, input_fn, feed_fn, batch_size, shuffle=False, epochs=1): return df.input_builder, df.get_feed_dict_fn() +@deprecated(None, 'Please specify feature columns explicitly.') def infer_real_valued_columns_from_input_fn(input_fn): """Creates `FeatureColumn` objects for inputs defined by `input_fn`. @@ -158,6 +164,7 @@ def infer_real_valued_columns_from_input_fn(input_fn): return layers.infer_real_valued_columns(features) +@deprecated(None, 'Please specify feature columns explicitly.') def infer_real_valued_columns_from_input(x): """Creates `FeatureColumn` objects for inputs defined by input `x`. @@ -211,7 +218,7 @@ def _get_replica_device_setter(config): 'Variable', 'VariableV2', 'AutoReloadVariable', 'MutableHashTable', 'MutableHashTableV2', 'MutableHashTableOfTensors', 'MutableHashTableOfTensorsV2', 'MutableDenseHashTable', - 'MutableDenseHashTableV2' + 'MutableDenseHashTableV2', 'VarHandleOp' ] if config.task_type: @@ -389,6 +396,10 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, trainable.Trainable): """Abstract BaseEstimator class to train and evaluate TensorFlow models. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Users should not instantiate or subclass this class. Instead, use an `Estimator`. """ @@ -399,6 +410,8 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, # TODO(wicke): Remove this once launcher takes over config functionality _Config = run_config.RunConfig # pylint: disable=invalid-name + @deprecated(None, 'Please replace uses of any Estimator from tf.contrib.learn' + ' with an Estimator from tf.estimator.*') def __init__(self, model_dir=None, config=None): """Initializes a BaseEstimator instance. @@ -457,6 +470,20 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, # TODO(wicke): make RunConfig immutable, and then return it without a copy. return copy.deepcopy(self._config) + @property + def model_fn(self): + """Returns the model_fn which is bound to self.params. + + Returns: + The model_fn with the following signature: + `def model_fn(features, labels, mode, metrics)` + """ + + def public_model_fn(features, labels, mode, config): + return self._call_model_fn(features, labels, mode, config=config) + + return public_model_fn + @deprecated_args(SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, ('x', None), ('y', None), ('batch_size', None)) def fit(self, @@ -600,7 +627,8 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, input_fn=None, batch_size=None, outputs=None, - as_iterable=True): + as_iterable=True, + iterate_batches=False): """Returns predictions for given features. Args: @@ -616,6 +644,9 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, 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 @@ -635,7 +666,8 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, 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. @@ -885,8 +917,8 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, if feed_fn: hooks.append(basic_session_run_hooks.FeedFnHook(feed_fn)) if steps == 0: - logging.warning('evaluation steps are 0. If `input_fn` does not raise' - 'OutOfRangeError`, the evaluation will never stop.' + logging.warning('evaluation steps are 0. If `input_fn` does not raise ' + '`OutOfRangeError`, the evaluation will never stop. ' 'Use steps=None if intended.') if steps: hooks.append( @@ -1069,6 +1101,10 @@ def _identity_feature_engineering_fn(features, labels): class Estimator(BaseEstimator): """Estimator class is the basic TensorFlow model trainer/evaluator. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. """ def __init__(self, @@ -1157,7 +1193,7 @@ class Estimator(BaseEstimator): self._feature_engineering_fn = ( feature_engineering_fn or _identity_feature_engineering_fn) - def _call_model_fn(self, features, labels, mode, metrics=None): + def _call_model_fn(self, features, labels, mode, metrics=None, config=None): """Calls model function with support of 2, 3 or 4 arguments. Args: @@ -1165,6 +1201,7 @@ class Estimator(BaseEstimator): labels: labels dict. mode: ModeKeys metrics: Dict of metrics. + config: RunConfig. Returns: A `ModelFnOps` object. If model_fn returns a tuple, wraps them up in a @@ -1181,7 +1218,10 @@ class Estimator(BaseEstimator): if 'params' in model_fn_args: kwargs['params'] = self.params if 'config' in model_fn_args: - kwargs['config'] = self.config + if config: + kwargs['config'] = config + else: + kwargs['config'] = self.config if 'model_dir' in model_fn_args: kwargs['model_dir'] = self.model_dir model_fn_results = self._model_fn(features, labels, **kwargs) @@ -1453,8 +1493,14 @@ class Estimator(BaseEstimator): # For time of deprecation x,y from Estimator allow direct access. # pylint: disable=protected-access class SKCompat(sklearn.BaseEstimator): - """Scikit learn wrapper for TensorFlow Learn Estimator.""" + """Scikit learn wrapper for TensorFlow Learn Estimator. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'Please switch to the Estimator interface.') def __init__(self, estimator): self._estimator = estimator diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator_test_utils.py b/tensorflow/contrib/learn/python/learn/estimators/estimator_test_utils.py index fd47710e3015de9ae6a453f98978b0ef8f88968c..e4c31396baf8271c49395a2b87b454dbc77177e2 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator_test_utils.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator_test_utils.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utils for Estimator.""" +"""Utils for Estimator (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/estimators/head.py b/tensorflow/contrib/learn/python/learn/estimators/head.py index bc0e6fc0091c9b5419ab526855b404eb4a927e97..2b4b6eff39f4fc8a20a149edfc07d2f4f27a9bae 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head.py @@ -12,8 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Abstractions for the head(s) of a model. +"""Abstractions for the head(s) of a model (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. """ + from __future__ import absolute_import from __future__ import division from __future__ import print_function @@ -47,11 +52,16 @@ from tensorflow.python.summary import summary from tensorflow.python.training import training from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect +from tensorflow.python.util.deprecation import deprecated class Head(object): """Interface for the head/top of a model. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Given logits (or output of a hidden layer), a Head knows how to compute predictions, loss, default metric and export signature. It is meant to, @@ -177,11 +187,13 @@ class Head(object): raise NotImplementedError("Calling an abstract method.") +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") 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 +211,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,9 +224,10 @@ 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)) +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") def poisson_regression_head(label_name=None, weight_column_name=None, label_dimension=1, @@ -251,6 +266,7 @@ def poisson_regression_head(label_name=None, # TODO(zakaria): Consider adding a _RegressionHead for logistic_regression +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") def multi_class_head(n_classes, label_name=None, weight_column_name=None, @@ -332,6 +348,7 @@ def multi_class_head(n_classes, label_keys=label_keys) +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") def binary_svm_head( label_name=None, weight_column_name=None, @@ -367,6 +384,7 @@ def binary_svm_head( thresholds=thresholds) +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") def multi_label_head(n_classes, label_name=None, weight_column_name=None, @@ -427,6 +445,7 @@ def multi_label_head(n_classes, loss_fn=_wrap_custom_loss_fn(loss_fn) if loss_fn else None) +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") def loss_only_head(loss_fn, head_name=None): """Creates a Head that contains only loss terms. @@ -444,6 +463,7 @@ def loss_only_head(loss_fn, head_name=None): return _LossOnlyHead(loss_fn, head_name=head_name) +@deprecated(None, "Please switch to tf.contrib.estimator.*_head.") def multi_head(heads, loss_weights=None): """Creates a MultiHead stemming from same logits/hidden layer. @@ -476,6 +496,7 @@ def multi_head(heads, loss_weights=None): return _MultiHead(heads, loss_merger=_weighted_loss_merger) +@deprecated(None, "Use 'lambda _: tf.no_op()'.") def no_op_train_fn(loss): del loss return control_flow_ops.no_op() diff --git a/tensorflow/contrib/learn/python/learn/estimators/head_test.py b/tensorflow/contrib/learn/python/learn/estimators/head_test.py index 3881bf533d642bef68fa9ab4ba908bbb8f7f8091..7c2d9bb0767cb979dae9c84b5342d129225677ed 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(): diff --git a/tensorflow/contrib/learn/python/learn/estimators/kmeans.py b/tensorflow/contrib/learn/python/learn/estimators/kmeans.py index 8f9d6fc318a357853bdb8e3264f6691b410006b1..66ebcfd1d81904b9afe5be6bd1a648fe325e1e0b 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/kmeans.py +++ b/tensorflow/contrib/learn/python/learn/estimators/kmeans.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Implementation of k-means clustering on top of `Estimator` API. +"""Implementation of k-means clustering on top of `Estimator` API (deprecated). This module is deprecated. Please use @{tf.contrib.factorization.KMeansClustering} instead of @@ -153,7 +153,12 @@ def _kmeans_clustering_model_fn(features, labels, mode, params, config): # TODO(agarwal,ands): support sharded input. class KMeansClustering(estimator.Estimator): - """An Estimator for K-Means clustering.""" + """An Estimator for K-Means clustering. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ SQUARED_EUCLIDEAN_DISTANCE = clustering_ops.SQUARED_EUCLIDEAN_DISTANCE COSINE_DISTANCE = clustering_ops.COSINE_DISTANCE RANDOM_INIT = clustering_ops.RANDOM_INIT diff --git a/tensorflow/contrib/learn/python/learn/estimators/linear.py b/tensorflow/contrib/learn/python/learn/estimators/linear.py index 37aa8b339622415d082933cdf66d2472a4119b48..64d7ecc68e7abb1d36a3eb098fedd8184d6e9d77 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/linear.py +++ b/tensorflow/contrib/learn/python/learn/estimators/linear.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Linear Estimators.""" +"""Linear Estimators (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -305,6 +310,10 @@ class _SdcaUpdateWeightsHook(session_run_hook.SessionRunHook): class LinearClassifier(estimator.Estimator): """Linear classifier model. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Train a linear model to classify instances into one of multiple possible classes. When number of possible classes is 2, this is binary classification. @@ -625,6 +634,10 @@ class LinearClassifier(estimator.Estimator): class LinearRegressor(estimator.Estimator): """Linear regressor model. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Train a linear regression model to predict label value given observation of feature values. @@ -860,6 +873,10 @@ class LinearRegressor(estimator.Estimator): class LinearEstimator(estimator.Estimator): """Linear model with user specified head. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Train a generalized linear model to predict label value given observation of feature values. diff --git a/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor.py b/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor.py index fb339160d58e09d4ffd50090f2dbbcec08bebe47..3cbcc6e98de1c915c302617e4591c9baa33adeaf 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor.py +++ b/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Logistic regression (aka binary classifier) class. +"""Logistic regression (aka binary classifier) class (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. This defines some useful basic metrics for using logistic regression to classify a binary event (0 vs 1). @@ -75,6 +79,10 @@ def LogisticRegressor( # pylint: disable=invalid-name feature_engineering_fn=None): """Builds a logistic regression Estimator for binary classification. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + This method provides a basic Estimator with some additional metrics for custom binary classification models, including AUC, precision/recall and accuracy. diff --git a/tensorflow/contrib/learn/python/learn/estimators/metric_key.py b/tensorflow/contrib/learn/python/learn/estimators/metric_key.py index 99388f116b345bd038f2985606c6203011597ea2..f264248e44d9aa48f26ee32e36746bd4c3145a8d 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/metric_key.py +++ b/tensorflow/contrib/learn/python/learn/estimators/metric_key.py @@ -12,14 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Enum for metric keys.""" +"""Enum for metric keys (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function class MetricKey(object): - """Metric key strings.""" + """Metric key strings (deprecated).""" + LOSS = "loss" AUC = "auc" AUC_PR = "auc_precision_recall" diff --git a/tensorflow/contrib/learn/python/learn/estimators/model_fn.py b/tensorflow/contrib/learn/python/learn/estimators/model_fn.py index 44e6c7c52dac524a22e9099e33e2aef82f8fe7ba..dcb161180c99ce71195c820217e8bdaf79d70901 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/model_fn.py +++ b/tensorflow/contrib/learn/python/learn/estimators/model_fn.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Classes and methods related to model_fn.""" +"""Classes and methods related to model_fn (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -37,10 +42,13 @@ from tensorflow.python.ops import array_ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.saved_model import signature_constants from tensorflow.python.training import session_run_hook +from tensorflow.python.util.deprecation import deprecated class ModeKeys(object): - """Standard names for model modes. + """Standard names for model modes (deprecated). + + THIS CLASS IS DEPRECATED. The following standard keys are defined: @@ -65,8 +73,16 @@ class ModelFnOps( 'output_alternatives', 'training_chief_hooks', 'training_hooks', 'scaffold', 'mode' ])): - """Ops returned from a model_fn.""" + """Ops returned from a model_fn. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'When switching to tf.estimator.Estimator, use ' + 'tf.estimator.EstimatorSpec. You can use the `estimator_spec`' + ' method to create an equivalent one.') def __new__(cls, mode, predictions=None, diff --git a/tensorflow/contrib/learn/python/learn/estimators/prediction_key.py b/tensorflow/contrib/learn/python/learn/estimators/prediction_key.py index f8d87b8914307a86eb2f46123a28ff11eb925eda..6fd2fc9d592cef4e44a640e2f27cb28b367d44d5 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/prediction_key.py +++ b/tensorflow/contrib/learn/python/learn/estimators/prediction_key.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Enum for model prediction keys. +"""Enum for model prediction keys (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. This file is obsoleted in the move of Estimator to core. """ @@ -22,6 +26,8 @@ from __future__ import print_function class PredictionKey(object): + """THIS CLASS IS DEPRECATED.""" + CLASSES = "classes" PROBABILITIES = "probabilities" LOGITS = "logits" diff --git a/tensorflow/contrib/learn/python/learn/estimators/rnn_common.py b/tensorflow/contrib/learn/python/learn/estimators/rnn_common.py index 2752bc2d90ee0f51b2c40ccc4d24a4eb21cff38f..215022e5d9e5d3cd5d6a96583b325b19a1719568 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/rnn_common.py +++ b/tensorflow/contrib/learn/python/learn/estimators/rnn_common.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Common operations for RNN Estimators.""" +"""Common operations for RNN Estimators (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/estimators/run_config.py b/tensorflow/contrib/learn/python/learn/estimators/run_config.py index fd90fd1cc6277e7d80287aefdbab6134dac7c0d5..1d161093de01ef838d0c75ec9a39574c7529bd57 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/run_config.py +++ b/tensorflow/contrib/learn/python/learn/estimators/run_config.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Run Config.""" +"""Run Config (deprecated, use tf.estimator.RunConfig instead). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -29,11 +34,12 @@ from tensorflow.core.protobuf import config_pb2 from tensorflow.python.estimator import run_config as core_run_config from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import server_lib +from tensorflow.python.util.deprecation import deprecated # A list of the property names in RunConfig user allows to change. They will # not affect the execution framework, so when execution framework checks the -# `uid` of the RunConfig, it should be ingored. +# `uid` of the RunConfig, it should be ignored. _DEFAULT_UID_WHITE_LIST = [ 'tf_random_seed', 'save_summary_steps', @@ -47,6 +53,7 @@ _DEFAULT_UID_WHITE_LIST = [ class Environment(object): + """DEPRECATED CLASS.""" # For running general distributed training. CLOUD = 'cloud' # For running Google-internal distributed training. @@ -56,6 +63,7 @@ class Environment(object): class TaskType(object): + """DEPRECATED CLASS.""" MASTER = 'master' PS = 'ps' WORKER = 'worker' @@ -64,6 +72,8 @@ class TaskType(object): class ClusterConfig(object): """This class specifies the configurations for a distributed run. + THIS CLASS IS DEPRECATED. Use tf.estimator.RunConfig instead. + If you're using an `Estimator`, you should probably use the subclass RunConfig instead. """ @@ -211,10 +221,13 @@ class ClusterConfig(object): class RunConfig(ClusterConfig, core_run_config.RunConfig): """This class specifies the configurations for an `Estimator` run. - This class is the implementation of @{tf.estimator.RunConfig} interface. + This class is a deprecated implementation of @{tf.estimator.RunConfig} + interface. """ _USE_DEFAULT = 0 + @deprecated(None, 'When switching to tf.estimator.Estimator, use' + ' tf.estimator.RunConfig instead.') def __init__(self, master=None, num_cores=0, diff --git a/tensorflow/contrib/learn/python/learn/estimators/state_saving_rnn_estimator.py b/tensorflow/contrib/learn/python/learn/estimators/state_saving_rnn_estimator.py index 0cea35e219a4457417a161a3ac4ac4292fd24f53..de78c72c3ae3ef14f5f7c46b1d47f82e8266c7c6 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/state_saving_rnn_estimator.py +++ b/tensorflow/contrib/learn/python/learn/estimators/state_saving_rnn_estimator.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Estimator for State Saving RNNs.""" +"""Estimator for State Saving RNNs (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -528,6 +533,12 @@ def _get_rnn_model_fn(cell_type, class StateSavingRnnEstimator(estimator.Estimator): + """RNN with static unrolling and state saving (deprecated). + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ def __init__(self, problem_type, diff --git a/tensorflow/contrib/learn/python/learn/estimators/svm.py b/tensorflow/contrib/learn/python/learn/estimators/svm.py index 72920d73c0c92886e54f533ad7fe170fe27d9870..3459997baba16fc0d4045e50819ecdd0e7121657 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/svm.py +++ b/tensorflow/contrib/learn/python/learn/estimators/svm.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Support Vector Machine (SVM) Estimator.""" +"""Support Vector Machine (SVM) Estimator (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -36,6 +41,10 @@ def _as_iterable(preds, output): class SVM(estimator.Estimator): """Support Vector Machine (SVM) model for binary classification. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Currently, only linear SVMs are supported. For the underlying optimization problem, the `SDCAOptimizer` is used. For performance and convergence tuning, the num_loss_partitions parameter passed to `SDCAOptimizer` (see `__init__()` diff --git a/tensorflow/contrib/learn/python/learn/estimators/tensor_signature.py b/tensorflow/contrib/learn/python/learn/estimators/tensor_signature.py index a120bc6cc3975a3d4559d018c8aa74ff42a16d2d..71b5658dd174d2b47e33860844359f68e6768026 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/tensor_signature.py +++ b/tensorflow/contrib/learn/python/learn/estimators/tensor_signature.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""TensorSignature class and utilities.""" +"""TensorSignature class and utilities (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -33,6 +38,10 @@ class TensorSignature(collections.namedtuple( "TensorSignature", ["dtype", "shape", "is_sparse"])): """Signature of the `Tensor` object. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Useful to check compatibility of tensors. Example: diff --git a/tensorflow/contrib/learn/python/learn/estimators/test_data.py b/tensorflow/contrib/learn/python/learn/estimators/test_data.py index ed201bfc58f273e6587850032386c2686aea4148..e4b057b4f5a9e081c2d891bd9828ffc315e51e91 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/test_data.py +++ b/tensorflow/contrib/learn/python/learn/estimators/test_data.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Test data utilities.""" +"""Test data utilities (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/evaluable.py b/tensorflow/contrib/learn/python/learn/evaluable.py index 8f6cd39864b437f163dd7c1140dc88755ce98529..10881ca885599bc81386e15f814a2687d907f63b 100644 --- a/tensorflow/contrib/learn/python/learn/evaluable.py +++ b/tensorflow/contrib/learn/python/learn/evaluable.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""`Evaluable` interface.""" +"""`Evaluable` interface (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -23,6 +28,10 @@ import abc class Evaluable(object): """Interface for objects that are evaluatable by, e.g., `Experiment`. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. """ __metaclass__ = abc.ABCMeta diff --git a/tensorflow/contrib/learn/python/learn/experiment.py b/tensorflow/contrib/learn/python/learn/experiment.py index bec976afd2719138117976381669ca3292360480..3744abd860e7f460133873eb534fd75887182f78 100644 --- a/tensorflow/contrib/learn/python/learn/experiment.py +++ b/tensorflow/contrib/learn/python/learn/experiment.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Experiment class collecting information needed for a single training run.""" +"""Experiment class collecting information for a single training run (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -25,7 +30,6 @@ import os import time from tensorflow.contrib.framework import deprecated -from tensorflow.contrib.framework import deprecated_args from tensorflow.contrib.framework.python.framework import experimental from tensorflow.contrib.learn.python.learn import evaluable from tensorflow.contrib.learn.python.learn import export_strategy @@ -118,6 +122,10 @@ class _EvalAndExportListener(basic_session_run_hooks.CheckpointSaverListener): class Experiment(object): """Experiment is a class containing all information needed to train a model. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + After an experiment is created (by passing an Estimator and inputs for training and evaluation), an Experiment instance knows how to invoke training and eval loops in a sensible fashion for distributed training. @@ -125,16 +133,8 @@ class Experiment(object): # TODO(ispir): remove delay_workers_by_global_step and make global step based # waiting as only behavior. - @deprecated_args( - "2016-10-23", - "local_eval_frequency is deprecated as local_run will be renamed to " - "train_and_evaluate. Use min_eval_frequency and call train_and_evaluate " - "instead. Note, however, that the default for min_eval_frequency is 1, " - "meaning models will be evaluated every time a new checkpoint is " - "available. In contrast, the default for local_eval_frequency is None, " - "resulting in evaluation occurring only after training has completed. " - "min_eval_frequency is ignored when calling the deprecated local_run.", - "local_eval_frequency") + @deprecated(None, "Please switch to tf.estimator.train_and_evaluate. You will" + " also have to convert to a tf.estimator.Estimator.") def __init__(self, estimator, train_input_fn, @@ -152,7 +152,8 @@ class Experiment(object): export_strategies=None, train_steps_per_iteration=None, checkpoint_and_export=False, - saving_listeners=None): + saving_listeners=None, + check_interval_secs=5): """Constructor for `Experiment`. Creates an Experiment instance. None of the functions passed to this @@ -190,8 +191,9 @@ class Experiment(object): number of steps between evaluations. Of course, evaluation does not occur if no new snapshot is available, hence, this is the minimum. If 0, the evaluation will only happen after training. - If None, defaults to 1, unless model_dir is on GCS, in which case the - default is 1000. + If None, defaults to 1. To avoid checking for new checkpoints too + frequent, the interval is further limited to be at least + check_interval_secs between checks. delay_workers_by_global_step: if `True` delays training workers based on global step instead of time. export_strategies: Iterable of `ExportStrategy`s, or a single one, or @@ -215,7 +217,10 @@ class Experiment(object): saving_listeners: list of `CheckpointSaverListener` objects. Used by tf.estimator.Estimator for callbacks that run immediately before or after checkpoint savings. - + check_interval_secs: + Minimum time between subsequent checks for a new checkpoint. This + mostly applies if both min_eval_frequency and the time spent per + training step is low. Raises: ValueError: if `estimator` does not implement Estimator interface, or if export_strategies has the wrong type. @@ -261,13 +266,9 @@ class Experiment(object): self._continuous_eval_throttle_secs = continuous_eval_throttle_secs self._checkpoint_and_export = checkpoint_and_export self._saving_listeners = saving_listeners - # Using 1 on a non-cached file system requires a lot of overhead to - # read the checkpoint state file. This is particular bad on GCS, so - # we use a different default. This is a temporary band-aid, to be - # fixed holistically later (b/36498507). - default_min_eval_frequency = 1000 if _is_gcs(estimator.model_dir) else 1 self._min_eval_frequency = min_eval_frequency if ( - min_eval_frequency is not None) else default_min_eval_frequency + min_eval_frequency is not None) else 1 + self._check_interval_secs = check_interval_secs self._delay_workers_by_global_step = delay_workers_by_global_step self._train_monitors = train_monitors[:] if train_monitors else [] self._eval_hooks = eval_hooks[:] if eval_hooks else [] @@ -357,7 +358,7 @@ class Experiment(object): self._start_server() elif config.cluster_spec and config.master: raise ValueError( - "For distributed runtime, Experiment class only works with" + "For distributed runtime, Experiment class only works with " "tf.contrib.learn.RunConfig for now, but provided {}".format( type(config))) @@ -646,12 +647,19 @@ class Experiment(object): self._train_monitors += [saver_hook] else: if self._min_eval_frequency: + # Using low min_eval_frequency (default is 1) on a non-cached file + # system requires a lot of overhead to read the checkpoint state file. + # This is particular bad on GCS and CNS. See also b/36498507 for + # context. `check_interval_secs = 5` avoids polling a remote + # fileystem too often. + self._train_monitors += [ monitors.ValidationMonitor( input_fn=self._eval_input_fn, eval_steps=self._eval_steps, metrics=self._eval_metrics, every_n_steps=self._min_eval_frequency, + check_interval_secs=self._check_interval_secs, name=eval_dir_suffix, hooks=self._eval_hooks) ] @@ -928,7 +936,3 @@ def _new_attr_context(obj, attr): yield finally: setattr(obj, attr, saved) - - -def _is_gcs(model_dir): - return model_dir and model_dir.startswith("gs://") diff --git a/tensorflow/contrib/learn/python/learn/experiment_test.py b/tensorflow/contrib/learn/python/learn/experiment_test.py index 545d7d8924c0c10544e6113e2968b7ae3d2090fc..d10927a0cdd5c67c8d2a8e569153235ee175ec4d 100644 --- a/tensorflow/contrib/learn/python/learn/experiment_test.py +++ b/tensorflow/contrib/learn/python/learn/experiment_test.py @@ -674,37 +674,11 @@ class ExperimentTest(test.TestCase): def test_min_eval_frequency_defaults(self): def dummy_model_fn(features, labels): # pylint: disable=unused-argument pass - - # The default value when model_dir is on GCS is 1000 - estimator = core_estimator.Estimator(dummy_model_fn, 'gs://dummy_bucket') - ex = experiment.Experiment( - estimator, train_input_fn=None, eval_input_fn=None) - self.assertEquals(ex._min_eval_frequency, 1000) - - # The default value when model_dir is not on GCS is 1 estimator = core_estimator.Estimator(dummy_model_fn, '/tmp/dummy') ex = experiment.Experiment( estimator, train_input_fn=None, eval_input_fn=None) self.assertEquals(ex._min_eval_frequency, 1) - # Make sure default not used when explicitly set - estimator = core_estimator.Estimator(dummy_model_fn, 'gs://dummy_bucket') - ex = experiment.Experiment( - estimator, - min_eval_frequency=123, - train_input_fn=None, - eval_input_fn=None) - self.assertEquals(ex._min_eval_frequency, 123) - - # Make sure default not used when explicitly set as 0 - estimator = core_estimator.Estimator(dummy_model_fn, 'gs://dummy_bucket') - ex = experiment.Experiment( - estimator, - min_eval_frequency=0, - train_input_fn=None, - eval_input_fn=None) - self.assertEquals(ex._min_eval_frequency, 0) - def test_continuous_train_and_eval(self): for est in self._estimators_for_tests(eval_dict={'global_step': 100}): if isinstance(est, core_estimator.Estimator): diff --git a/tensorflow/contrib/learn/python/learn/export_strategy.py b/tensorflow/contrib/learn/python/learn/export_strategy.py index 55a8b824312b89e0ac66513242191f4201ac212a..075cab536ecb5279e7e6f23abb0b70c75043a7ec 100644 --- a/tensorflow/contrib/learn/python/learn/export_strategy.py +++ b/tensorflow/contrib/learn/python/learn/export_strategy.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""ExportStrategy class represents different flavors of model export.""" +"""ExportStrategy class represents different flavors of model export (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -21,6 +26,7 @@ from __future__ import print_function import collections from tensorflow.python.util import tf_inspect +from tensorflow.python.util.deprecation import deprecated __all__ = ['ExportStrategy'] @@ -30,6 +36,10 @@ class ExportStrategy( ['name', 'export_fn', 'strip_default_attrs'])): """A class representing a type of model export. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Typically constructed by a utility function specific to the exporter, such as `saved_model_export_utils.make_export_strategy()`. @@ -56,6 +66,8 @@ class ExportStrategy( forward compatibility of the resulting `SavedModel`. """ + @deprecated(None, 'Please switch to tf.estimator.train_and_evaluate, and use ' + 'tf.estimator.Exporter.') def __new__(cls, name, export_fn, strip_default_attrs=None): return super(ExportStrategy, cls).__new__( cls, name, export_fn, strip_default_attrs) diff --git a/tensorflow/contrib/learn/python/learn/graph_actions.py b/tensorflow/contrib/learn/python/learn/graph_actions.py index 98365c05f663e5d2a06703457fc5663d7135f7d9..a997fab723a16dddf150aa9397863605e4e77933 100644 --- a/tensorflow/contrib/learn/python/learn/graph_actions.py +++ b/tensorflow/contrib/learn/python/learn/graph_actions.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""High level operations on graphs.""" +"""High level operations on graphs (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -68,6 +73,7 @@ def clear_summary_writers(): return summary_io.SummaryWriterCache.clear() +@deprecated(None, 'Use `SummaryWriterCache.get` directly.') def get_summary_writer(logdir): """Returns single SummaryWriter per logdir in current run. diff --git a/tensorflow/contrib/learn/python/learn/learn_io/__init__.py b/tensorflow/contrib/learn/python/learn/learn_io/__init__.py index 06c3782a471537cf3879450e6bd20899a35d96ac..8b133a4440d8cbc19abca64f972791fc16ade6f8 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/__init__.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/__init__.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tools to allow different io formats.""" +"""Tools to allow different io formats (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/learn_io/dask_io.py b/tensorflow/contrib/learn/python/learn/learn_io/dask_io.py index 7d666391cea3c0a52a2cb7e324c00d5f480710d5..e0a1948d95a727675dac8ff3ce9f55c35d5f8d8d 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/dask_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/dask_io.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Methods to allow dask.DataFrame.""" +"""Methods to allow dask.DataFrame (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -21,6 +26,8 @@ from __future__ import print_function import numpy as np +from tensorflow.python.util.deprecation import deprecated + try: # pylint: disable=g-import-not-at-top import dask.dataframe as dd @@ -60,6 +67,7 @@ def _construct_dask_df_with_divisions(df): return dd.Series(merge(dsk, df.dask), name, df.name, divisions) +@deprecated(None, 'Please feed input to tf.data to support dask.') def extract_dask_data(data): """Extract data from dask.Series or dask.DataFrame for predictors. @@ -81,6 +89,7 @@ def extract_dask_data(data): return data +@deprecated(None, 'Please feed input to tf.data to support dask.') def extract_dask_labels(labels): """Extract data from dask.Series or dask.DataFrame for labels. 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 96be8b1bc402479d5611965f27abb197363cb939..c45b1d186471125776d6536112aebb66bb5ad558 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/data_feeder.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/data_feeder.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Implementations of different data feeders to provide data for TF trainer.""" +"""Implementations of different data feeders to provide data for TF trainer (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" # TODO(ipolosukhin): Replace this module with feed-dict queue runners & queues. @@ -31,6 +36,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.deprecation import deprecated # pylint: disable=g-multiple-import,g-bad-import-order from .pandas_io import HAS_PANDAS, extract_pandas_data, extract_pandas_matrix, extract_pandas_labels @@ -101,6 +107,7 @@ def _is_iterable(x): return hasattr(x, 'next') or hasattr(x, '__next__') +@deprecated(None, 'Please use tensorflow/transform or tf.data.') def setup_train_data_feeder(x, y, n_classes, @@ -188,6 +195,7 @@ def _batch_data(x, batch_size=None): yield np.matrix(chunk) +@deprecated(None, 'Please use tensorflow/transform or tf.data.') def setup_predict_data_feeder(x, batch_size=None): """Returns an iterable for feeding into predict step. @@ -219,6 +227,7 @@ def setup_predict_data_feeder(x, batch_size=None): return [x] +@deprecated(None, 'Please use tensorflow/transform or tf.data.') def setup_processor_data_feeder(x): """Sets up processor iterable. @@ -233,6 +242,7 @@ def setup_processor_data_feeder(x): return x +@deprecated(None, 'Please convert numpy dtypes explicitly.') def check_array(array, dtype): """Checks array on dtype and converts it if different. @@ -275,8 +285,14 @@ def _check_dtype(dtype): class DataFeeder(object): - """Data feeder is an example class to sample data for TF trainer.""" + """Data feeder is an example class to sample data for TF trainer. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'Please use tensorflow/transform or tf.data.') def __init__(self, x, y, @@ -563,6 +579,10 @@ class DataFeeder(object): class StreamingDataFeeder(DataFeeder): """Data feeder for TF trainer that reads data from iterator. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Streaming data feeder allows to read data as it comes it from disk or somewhere else. It's custom to have this iterators rotate infinetly over the dataset, to allow control of how much to learn on the trainer side. @@ -771,11 +791,16 @@ class StreamingDataFeeder(DataFeeder): class DaskDataFeeder(object): """Data feeder for that reads data from dask.Series and dask.DataFrame. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Numpy arrays can be serialized to disk and it's possible to do random seeks into them. DaskDataFeeder will remove requirement to have full dataset in the memory and still do random seeks for sampling of batches. """ + @deprecated(None, 'Please feed input to tf.data to support dask.') def __init__(self, x, y, diff --git a/tensorflow/contrib/learn/python/learn/learn_io/generator_io.py b/tensorflow/contrib/learn/python/learn/learn_io/generator_io.py index 884faf8335e2a3ca1d27d2d93b4c817131648774..f8aaa0c9e3e5b589a6ad47678dba3dc38de7c471 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/generator_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/generator_io.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Methods to allow generator of dict with numpy arrays.""" +"""Methods to allow generator of dict with numpy arrays (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -23,8 +28,10 @@ from types import FunctionType from types import GeneratorType from tensorflow.python.estimator.inputs.queues.feeding_functions import _enqueue_data as enqueue_data +from tensorflow.python.util.deprecation import deprecated +@deprecated(None, 'Please use tf.data.') def generator_input_fn(x, target_key=None, batch_size=128, diff --git a/tensorflow/contrib/learn/python/learn/learn_io/graph_io.py b/tensorflow/contrib/learn/python/learn/learn_io/graph_io.py index 3a46c239688017f9204d2c6182a6f81cd325a417..9e816f54b6cf8dee84c6d62406ab3db700054d06 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/graph_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/graph_io.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Methods to read data in the graph.""" +"""Methods to read data in the graph (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -34,11 +39,13 @@ from tensorflow.python.platform import gfile from tensorflow.python.summary import summary from tensorflow.python.training import input as input_ops from tensorflow.python.training import queue_runner +from tensorflow.python.util.deprecation import deprecated # Default name for key in the feature dict. KEY_FEATURE_NAME = '__key__' +@deprecated(None, 'Use tf.data.') def read_batch_examples(file_pattern, batch_size, reader, @@ -106,6 +113,7 @@ def read_batch_examples(file_pattern, return examples +@deprecated(None, 'Use tf.data.') def read_keyed_batch_examples(file_pattern, batch_size, reader, @@ -175,6 +183,7 @@ def read_keyed_batch_examples(file_pattern, seed=seed) +@deprecated(None, 'Use tf.data.') def read_keyed_batch_examples_shared_queue(file_pattern, batch_size, reader, @@ -452,6 +461,7 @@ def _read_keyed_batch_examples_helper(file_pattern, return queued_examples_with_keys +@deprecated(None, 'Use tf.data.') def read_keyed_batch_features(file_pattern, batch_size, features, @@ -540,6 +550,7 @@ def read_keyed_batch_features(file_pattern, name=scope) +@deprecated(None, 'Use tf.data.') def read_keyed_batch_features_shared_queue(file_pattern, batch_size, features, @@ -620,6 +631,7 @@ def read_keyed_batch_features_shared_queue(file_pattern, name=scope) +@deprecated(None, 'Use tf.data.') def queue_parsed_features(parsed_features, keys=None, feature_queue_capacity=100, @@ -742,6 +754,7 @@ def queue_parsed_features(parsed_features, return dequeued_keys, dequeued_parsed_features +@deprecated(None, 'Use tf.data.') def read_batch_features(file_pattern, batch_size, features, @@ -821,6 +834,7 @@ def read_batch_features(file_pattern, return features +@deprecated(None, 'Use tf.data.') def read_batch_record_features(file_pattern, batch_size, features, diff --git a/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py b/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py index 692438807fbd7febb156d4db73b5d3deba6c987d..29552d24f1eaa0d85a99c8b09f69d007e7e4fe9f 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py @@ -12,15 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Methods to allow dict of numpy arrays.""" +"""Methods to allow dict of numpy arrays (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.python.estimator.inputs.numpy_io import numpy_input_fn as core_numpy_input_fn +from tensorflow.python.util.deprecation import deprecated +@deprecated(None, 'Use tf.estimator.inputs.numpy_input_fn.') def numpy_input_fn(x, y=None, batch_size=128, diff --git a/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py b/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py index ede7558eafa9237dc63aa95a62e599c5e9755822..b4ef055f5ae484ec704ad42efcf2c00c4a7a4f56 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py @@ -13,13 +13,19 @@ # limitations under the License. # ============================================================================== -"""Methods to allow pandas.DataFrame.""" +"""Methods to allow pandas.DataFrame (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.python.estimator.inputs.pandas_io import pandas_input_fn as core_pandas_input_fn +from tensorflow.python.util.deprecation import deprecated try: # pylint: disable=g-import-not-at-top @@ -47,6 +53,7 @@ PANDAS_DTYPES = { } +@deprecated(None, 'Please use tf.estimator.inputs.pandas_input_fn') def pandas_input_fn(x, y=None, batch_size=128, @@ -66,6 +73,7 @@ def pandas_input_fn(x, target_column=target_column) +@deprecated(None, 'Please access pandas data directly.') def extract_pandas_data(data): """Extract data from pandas.DataFrame for predictors. @@ -96,6 +104,7 @@ def extract_pandas_data(data): 'float, or bool. Found: ' + ', '.join(error_report)) +@deprecated(None, 'Please access pandas data directly.') def extract_pandas_matrix(data): """Extracts numpy matrix from pandas DataFrame. @@ -111,6 +120,7 @@ def extract_pandas_matrix(data): return data.as_matrix() +@deprecated(None, 'Please access pandas data directly.') def extract_pandas_labels(labels): """Extract data from pandas.DataFrame for labels. diff --git a/tensorflow/contrib/learn/python/learn/learn_runner.py b/tensorflow/contrib/learn/python/learn/learn_runner.py index 2af723a0d64822e81fa0fbeb106ab812de6ab4e8..d719a3e488b9905ef7903e21d90dbaae0449735c 100644 --- a/tensorflow/contrib/learn/python/learn/learn_runner.py +++ b/tensorflow/contrib/learn/python/learn/learn_runner.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Runs an Experiment.""" +"""Runs an Experiment (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -22,6 +27,7 @@ from tensorflow.contrib.learn.python.learn.estimators import run_config as run_c from tensorflow.contrib.learn.python.learn.experiment import Experiment from tensorflow.contrib.training.python.training import hparam as hparam_lib from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.deprecation import deprecated # TODO(xiejw): Refactor the learn_runner to make code reusable. @@ -99,6 +105,7 @@ def _wrapped_experiment_fn_with_uid_check(experiment_fn, require_hparams=False): return wrapped_experiment_fn +@deprecated(None, 'Use tf.estimator.train_and_evaluate.') def run(experiment_fn, output_dir=None, schedule=None, run_config=None, hparams=None): """Make and run an experiment. @@ -218,6 +225,7 @@ def run(experiment_fn, output_dir=None, schedule=None, run_config=None, return _execute_schedule(experiment, schedule) +@deprecated(None, 'Use tf.estimator.train_and_evaluate.') def tune(experiment_fn, tuner): """Tune an experiment with hyper-parameters. diff --git a/tensorflow/contrib/learn/python/learn/learn_runner_lib.py b/tensorflow/contrib/learn/python/learn/learn_runner_lib.py index 7d9b1c7716f0ab1f2274ca53406175240b613027..ba2d067787c1dfd4e4820ecc916f1053e9f3cf60 100644 --- a/tensorflow/contrib/learn/python/learn/learn_runner_lib.py +++ b/tensorflow/contrib/learn/python/learn/learn_runner_lib.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utilities to run and tune an Experiment. +"""Utilities to run and tune an Experiment (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. @@run @@tune diff --git a/tensorflow/contrib/learn/python/learn/metric_spec.py b/tensorflow/contrib/learn/python/learn/metric_spec.py index 6440bc204b8e339ff51311dcc87b36f556b94092..97220365d5dddb82b602369f06bea021a86d584f 100644 --- a/tensorflow/contrib/learn/python/learn/metric_spec.py +++ b/tensorflow/contrib/learn/python/learn/metric_spec.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""The metric spec class to flexibly connect models and metrics.""" +"""The metric spec class to flexibly connect models and metrics (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -22,6 +27,7 @@ import six from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import tf_inspect +from tensorflow.python.util.deprecation import deprecated def _assert_named_args(sentinel): @@ -223,6 +229,10 @@ def _adapt_metric_fn( class MetricSpec(object): """MetricSpec connects a model to metric functions. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + The MetricSpec class contains all information necessary to connect the output of a `model_fn` to the metrics (usually, streaming metrics) that are used in evaluation. @@ -284,6 +294,7 @@ class MetricSpec(object): """ + @deprecated(None, 'Use tf.estimator.EstimatorSpec.eval_metric_ops.') def __init__(self, metric_fn, prediction_key=None, diff --git a/tensorflow/contrib/learn/python/learn/models.py b/tensorflow/contrib/learn/python/learn/models.py index 4283240d018c949bb35aeb12032d2ee8b75884a5..bd4bbf9f8c9ad7e8a0fc06d8c0dc24672536c158 100644 --- a/tensorflow/contrib/learn/python/learn/models.py +++ b/tensorflow/contrib/learn/python/learn/models.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Various high level TF models.""" +"""Various high level TF models (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -28,8 +33,10 @@ from tensorflow.python.ops import array_ops as array_ops_ from tensorflow.python.ops import init_ops from tensorflow.python.ops import variable_scope as vs from tensorflow.python.summary import summary +from tensorflow.python.util.deprecation import deprecated +@deprecated(None, 'Consider using a tf.estimator.LinearRegressor') def linear_regression_zero_init(x, y): """Linear regression subgraph with zero-value initial weights and bias. @@ -43,6 +50,7 @@ def linear_regression_zero_init(x, y): return linear_regression(x, y, init_mean=0.0, init_stddev=0.0) +@deprecated(None, 'Consider using a class from tf.estimator.LinearClassifier') def logistic_regression_zero_init(x, y): """Logistic regression subgraph with zero-value initial weights and bias. @@ -56,6 +64,7 @@ def logistic_regression_zero_init(x, y): return logistic_regression(x, y, init_mean=0.0, init_stddev=0.0) +@deprecated(None, 'Consider using a class from tf.estimator.') def linear_regression(x, y, init_mean=None, init_stddev=1.0): """Creates linear regression TensorFlow subgraph. @@ -107,6 +116,7 @@ def linear_regression(x, y, init_mean=None, init_stddev=1.0): return losses_ops.mean_squared_error_regressor(x, y, weights, bias) +@deprecated(None, 'Consider using a class from tf.estimator.') def logistic_regression(x, y, class_weight=None, @@ -203,6 +213,7 @@ def _reverse_seq(input_seq, lengths): return result +@deprecated(None, 'Please consider `tf.nn.bidirectional_dynamic_rnn`.') def bidirectional_rnn(cell_fw, cell_bw, inputs, @@ -283,6 +294,7 @@ def bidirectional_rnn(cell_fw, # End of TensorFlow 0.7 +@deprecated(None, 'Please consider tensorflow/tensor2tensor.') def get_rnn_model(rnn_size, cell_type, num_layers, input_op_fn, bidirectional, target_predictor_fn, sequence_length, initial_state, attn_length, attn_size, attn_vec_size): diff --git a/tensorflow/contrib/learn/python/learn/monitored_session.py b/tensorflow/contrib/learn/python/learn/monitored_session.py index 22602e9f69d972505d83a66a6f9183b5e4d15c44..ac0433f1775feeed2ec3cf49291da01500bef01b 100644 --- a/tensorflow/contrib/learn/python/learn/monitored_session.py +++ b/tensorflow/contrib/learn/python/learn/monitored_session.py @@ -13,7 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""A wrapper of Session API which runs hooks.""" +"""A wrapper of Session API which runs hooks (deprecated). + +These are deprecated aliases for classes and functions in `tf.train`. Please use +those directly. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/monitors.py b/tensorflow/contrib/learn/python/learn/monitors.py index 0948dee7e2fa1b1b3617abd08d2d43ebc5340f63..77f7c73d5412d40b338eaff4cf04d99fd0892723 100644 --- a/tensorflow/contrib/learn/python/learn/monitors.py +++ b/tensorflow/contrib/learn/python/learn/monitors.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Monitors instrument the training process. +"""Monitors instrument the training process (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. @@get_default_monitors @@BaseMonitor @@ -59,6 +63,10 @@ from tensorflow.python.util import tf_inspect class BaseMonitor(object): """Base class for Monitors. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Defines basic interfaces of Monitors. Monitors can either be run on all workers or, more commonly, restricted to run exclusively on the elected chief worker. @@ -229,6 +237,10 @@ def _extract_output(outputs, request): class EveryN(BaseMonitor): """Base class for monitors that execute callbacks every N steps. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + This class adds three new callbacks: - every_n_step_begin - every_n_step_end @@ -418,6 +430,10 @@ class StopAtStep(BaseMonitor): class PrintTensor(EveryN): """Prints given tensors every N steps. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + This is an `EveryN` monitor and has consistent semantic for `every_n` and `first_n`. @@ -455,9 +471,12 @@ class PrintTensor(EveryN): class LoggingTrainable(EveryN): """Writes trainable variable values into log every N steps. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Write the tensors in trainable variables `every_n` steps, starting with the `first_n`th step. - """ def __init__(self, scope=None, every_n=100, first_n=1): @@ -493,7 +512,12 @@ class LoggingTrainable(EveryN): class SummarySaver(EveryN): - """Saves summaries every N steps.""" + """Saves summaries every N steps. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ def __init__(self, summary_op, @@ -554,6 +578,10 @@ class SummarySaver(EveryN): class ValidationMonitor(EveryN): """Runs evaluation of a given estimator, at most every N steps. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Note that the evaluation is done based on the saved checkpoint, which will usually be older than the current step. @@ -573,7 +601,8 @@ class ValidationMonitor(EveryN): early_stopping_rounds=None, early_stopping_metric="loss", early_stopping_metric_minimize=True, - name=None): + name=None, + check_interval_secs=5): """Initializes a ValidationMonitor. Args: @@ -600,6 +629,9 @@ class ValidationMonitor(EveryN): loss metrics like mean squared error, and False for performance metrics like accuracy. name: See `BaseEstimator.evaluate`. + check_interval_secs: Only check for new checkpoint if at least + `check_interval_secs` have passed. Ignore if None. Default is 5 secs. + Raises: ValueError: If both x and input_fn are provided. @@ -626,6 +658,8 @@ class ValidationMonitor(EveryN): self._early_stopped = False self._latest_path = None self._latest_path_step = None + self._last_checkpoint_check_time = None + self._check_interval_secs = check_interval_secs @property def early_stopped(self): @@ -690,6 +724,16 @@ class ValidationMonitor(EveryN): # that's what is being evaluated. if self._estimator is None: raise ValueError("Missing call to set_estimator.") + current_time = time.time() + if (self._check_interval_secs is not None and + self._last_checkpoint_check_time is not None and + current_time - self._last_checkpoint_check_time <= + self._check_interval_secs): + logging.debug( + "Skipping evaluation since less than %d seconds have passed since " + "last check for a new checkpoint.", self._check_interval_secs) + return False + self._last_checkpoint_check_time = current_time # Check that we are not running evaluation on the same checkpoint. latest_path = saver_lib.latest_checkpoint(self._estimator.model_dir) if latest_path is None: @@ -740,6 +784,10 @@ class ValidationMonitor(EveryN): class CaptureVariable(EveryN): """Captures a variable's values into a collection. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + This monitor is useful for unit testing. You should exercise caution when using this monitor in production, since it never discards values. @@ -778,6 +826,7 @@ class CaptureVariable(EveryN): self._var_values[step] = _extract_output(outputs, self._var_name) +@deprecation.deprecated(None, "Use tf.train.MonitoredTrainingSession.") def get_default_monitors(loss_op=None, summary_op=None, save_summary_steps=100, @@ -812,6 +861,10 @@ def get_default_monitors(loss_op=None, class GraphDump(BaseMonitor): """Dumps almost all tensors in the graph at every step. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Note, this is very expensive, prefer `PrintTensor` in production. """ @@ -879,7 +932,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)) @@ -901,7 +954,12 @@ class GraphDump(BaseMonitor): class ExportMonitor(EveryN): - """Monitor that exports Estimator every N steps.""" + """Monitor that exports Estimator every N steps. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ @deprecation.deprecated("2017-03-25", "ExportMonitor is deprecated. Please pass an " @@ -1024,7 +1082,12 @@ class ExportMonitor(EveryN): class CheckpointSaver(BaseMonitor): - """Saves checkpoints every N steps or N seconds.""" + """Saves checkpoints every N steps or N seconds. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ def __init__(self, checkpoint_dir, @@ -1109,7 +1172,12 @@ class CheckpointSaver(BaseMonitor): class StepCounter(EveryN): - """Steps per second monitor.""" + """Steps per second monitor. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ def __init__(self, every_n_steps=100, output_dir=None, summary_writer=None): super(StepCounter, self).__init__(every_n_steps=every_n_steps) @@ -1149,6 +1217,10 @@ class NanLossDuringTrainingError(RuntimeError): class NanLoss(EveryN): """NaN Loss monitor. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Monitors loss and stops training if loss is NaN. Can either fail with exception or just stop training. """ diff --git a/tensorflow/contrib/learn/python/learn/monitors_test.py b/tensorflow/contrib/learn/python/learn/monitors_test.py index b2b24776c60183113a5f936dd276ff312d6d0079..5c34d0ddb01f3bcdc407e6926e7c5b73be1863b4 100644 --- a/tensorflow/contrib/learn/python/learn/monitors_test.py +++ b/tensorflow/contrib/learn/python/learn/monitors_test.py @@ -385,7 +385,11 @@ class MonitorsTest(test.TestCase): estimator.evaluate.return_value = validation_outputs monitor = learn.monitors.ValidationMonitor( - x=constant_op.constant(2.0), every_n_steps=0, early_stopping_rounds=2) + x=constant_op.constant(2.0), + every_n_steps=0, + early_stopping_rounds=2, + check_interval_secs=None) + self._assert_validation_monitor(monitor) monitor.set_estimator(estimator) with ops.Graph().as_default() as g, self.test_session(g): diff --git a/tensorflow/contrib/learn/python/learn/ops/__init__.py b/tensorflow/contrib/learn/python/learn/ops/__init__.py index 33962e34cc685ce2c830a7bbfd1b5c626bcd8b31..efb1f47cf5bb2dcd0fb37b7b85cd8f170d56e4d1 100644 --- a/tensorflow/contrib/learn/python/learn/ops/__init__.py +++ b/tensorflow/contrib/learn/python/learn/ops/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Various TensorFlow Ops.""" +"""Various TensorFlow Ops (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/ops/embeddings_ops.py b/tensorflow/contrib/learn/python/learn/ops/embeddings_ops.py index fa3b7323e343371e986b763d30a8a44620894549..8f9811cf251ae0af1e0055a56e1358c2771b1367 100644 --- a/tensorflow/contrib/learn/python/learn/ops/embeddings_ops.py +++ b/tensorflow/contrib/learn/python/learn/ops/embeddings_ops.py @@ -13,7 +13,11 @@ # limitations under the License. # ============================================================================== -"""TensorFlow Ops to work with embeddings. +"""TensorFlow Ops to work with embeddings (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. Note: categorical variables are handled via embeddings in many cases. For example, in case of words. @@ -57,7 +61,7 @@ def embedding_lookup(params, ids, name='embedding_lookup'): ids = ops.convert_to_tensor(ids) shape = array_ops_.shape(ids) ids_flat = array_ops_.reshape( - ids, math_ops.reduce_prod(shape, keep_dims=True)) + ids, math_ops.reduce_prod(shape, keepdims=True)) embeds_flat = nn.embedding_lookup(params, ids_flat, name) embed_shape = array_ops_.concat([shape, [-1]], 0) embeds = array_ops_.reshape(embeds_flat, embed_shape) diff --git a/tensorflow/contrib/learn/python/learn/ops/losses_ops.py b/tensorflow/contrib/learn/python/learn/ops/losses_ops.py index b040ab3bb6c516158589a8e30d56fff1f7728951..92976d1539c7ddc226b81f903beee82b798ec8db 100644 --- a/tensorflow/contrib/learn/python/learn/ops/losses_ops.py +++ b/tensorflow/contrib/learn/python/learn/ops/losses_ops.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""TensorFlow Ops for loss computation.""" +"""TensorFlow Ops for loss computation (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division 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/ops/seq2seq_ops.py b/tensorflow/contrib/learn/python/learn/ops/seq2seq_ops.py index 45727faab4362abeab18f77861353eb53976023a..aa37cb4a76e2a6157bf077d327248353bd516472 100644 --- a/tensorflow/contrib/learn/python/learn/ops/seq2seq_ops.py +++ b/tensorflow/contrib/learn/python/learn/ops/seq2seq_ops.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""TensorFlow Ops for Sequence to Sequence models.""" +"""TensorFlow Ops for Sequence to Sequence models (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -26,8 +31,10 @@ 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 variable_scope as vs +from tensorflow.python.util.deprecation import deprecated +@deprecated(None, 'Please use tf.nn/tf.layers directly.') def sequence_classifier(decoding, labels, sampling_decoding=None, name=None): """Returns predictions and loss for sequence of predictions. @@ -57,6 +64,7 @@ def sequence_classifier(decoding, labels, sampling_decoding=None, name=None): return array_ops.stack(predictions, axis=1), loss +@deprecated(None, 'Please use tf.nn/tf.layers directly.') def seq2seq_inputs(x, y, input_length, output_length, sentinel=None, name=None): """Processes inputs for Sequence to Sequence models. @@ -87,6 +95,7 @@ def seq2seq_inputs(x, y, input_length, output_length, sentinel=None, name=None): return in_x, in_y, out_y +@deprecated(None, 'Please use tf.nn/tf.layers directly.') def rnn_decoder(decoder_inputs, initial_state, cell, scope=None): """RNN Decoder that creates training and sampling sub-graphs. @@ -123,6 +132,7 @@ def rnn_decoder(decoder_inputs, initial_state, cell, scope=None): return outputs, states, sampling_outputs, sampling_states +@deprecated(None, 'Please use tf.nn/tf.layers directly.') def rnn_seq2seq(encoder_inputs, decoder_inputs, encoder_cell, diff --git a/tensorflow/contrib/learn/python/learn/preprocessing/__init__.py b/tensorflow/contrib/learn/python/learn/preprocessing/__init__.py index 7bcc177d4ea0ab57f092d68888a72de2b2fd5edc..e8c6e1acf80f0791421bee59aff30e67bccb44b2 100644 --- a/tensorflow/contrib/learn/python/learn/preprocessing/__init__.py +++ b/tensorflow/contrib/learn/python/learn/preprocessing/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Preprocessing tools useful for building models.""" +"""Preprocessing tools useful for building models (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/preprocessing/categorical.py b/tensorflow/contrib/learn/python/learn/preprocessing/categorical.py index 154739d497ec1029026eaca1e93b37cd225f1050..faba3b2025e8abb51d1989c3fafbd5e711d6559b 100644 --- a/tensorflow/contrib/learn/python/learn/preprocessing/categorical.py +++ b/tensorflow/contrib/learn/python/learn/preprocessing/categorical.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Implements preprocessing transformers for categorical variables.""" +"""Implements preprocessing transformers for categorical variables (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -22,6 +27,8 @@ from __future__ import print_function import math import numpy as np +from tensorflow.python.util.deprecation import deprecated + # pylint: disable=g-bad-import-order from . import categorical_vocabulary from ..learn_io.data_feeder import setup_processor_data_feeder @@ -31,10 +38,16 @@ from ..learn_io.data_feeder import setup_processor_data_feeder class CategoricalProcessor(object): """Maps documents to sequences of word ids. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + As a common convention, Nan values are handled as unknown tokens. Both float('nan') and np.nan are accepted. """ + @deprecated(None, 'Please use tensorflow/transform or tf.data for sequence ' + 'processing.') def __init__(self, min_frequency=0, share=False, vocabularies=None): """Initializes a CategoricalProcessor instance. diff --git a/tensorflow/contrib/learn/python/learn/preprocessing/categorical_vocabulary.py b/tensorflow/contrib/learn/python/learn/preprocessing/categorical_vocabulary.py index 5709955c49fba50ca4a299a443a2902bbd9c6b23..3ac370a6ab4423846e810900514445ad5269b680 100644 --- a/tensorflow/contrib/learn/python/learn/preprocessing/categorical_vocabulary.py +++ b/tensorflow/contrib/learn/python/learn/preprocessing/categorical_vocabulary.py @@ -13,7 +13,11 @@ # limitations under the License. # ============================================================================== -"""Categorical vocabulary classes to map categories to indexes. +"""Categorical vocabulary classes to map categories to indexes (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. Can be used for categorical variables, sparse variables and words. """ @@ -25,14 +29,21 @@ from __future__ import print_function import collections import six +from tensorflow.python.util.deprecation import deprecated + class CategoricalVocabulary(object): """Categorical variables vocabulary class. + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + Accumulates and provides mapping from classes to indexes. Can be easily used for words. """ + @deprecated(None, 'Please use tensorflow/transform or tf.data.') def __init__(self, unknown_token="", support_reverse=True): self._unknown_token = unknown_token self._mapping = {unknown_token: 0} diff --git a/tensorflow/contrib/learn/python/learn/preprocessing/text.py b/tensorflow/contrib/learn/python/learn/preprocessing/text.py index 3af2074c2a46f0258c04111fff0235ba8309625e..f2b6776be7789a9433bfe41eb9354b74347059ec 100644 --- a/tensorflow/contrib/learn/python/learn/preprocessing/text.py +++ b/tensorflow/contrib/learn/python/learn/preprocessing/text.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""Implements a number of text preprocessing utilities.""" +"""Implements a number of text preprocessing utilities (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -24,6 +29,7 @@ import numpy as np import six from tensorflow.python.platform import gfile +from tensorflow.python.util.deprecation import deprecated from .categorical_vocabulary import CategoricalVocabulary # pylint: disable=g-bad-import-order @@ -38,6 +44,7 @@ TOKENIZER_RE = re.compile(r"[A-Z]{2,}(?![a-z])|[A-Z][a-z]+(?=[A-Z])|[\'\w\-]+", re.UNICODE) +@deprecated(None, 'Please use tensorflow/transform or tf.data.') def tokenizer(iterator): """Tokenizer generator. @@ -51,9 +58,16 @@ def tokenizer(iterator): yield TOKENIZER_RE.findall(value) +@deprecated(None, 'Please use tensorflow/transform or tf.data.') class ByteProcessor(object): - """Maps documents into sequence of ids for bytes.""" + """Maps documents into sequence of ids for bytes. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'Please use tensorflow/transform or tf.data.') def __init__(self, max_document_length): self.max_document_length = max_document_length @@ -108,8 +122,14 @@ class ByteProcessor(object): class VocabularyProcessor(object): - """Maps documents to sequences of word ids.""" + """Maps documents to sequences of word ids. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'Please use tensorflow/transform or tf.data.') def __init__(self, max_document_length, min_frequency=0, diff --git a/tensorflow/contrib/learn/python/learn/session_run_hook.py b/tensorflow/contrib/learn/python/learn/session_run_hook.py index a8ba2be97206f2b974d256ad2c62c21a4e3e55d8..87edc9b720bdb3edcd5f2dcd1662d14da53c51cf 100644 --- a/tensorflow/contrib/learn/python/learn/session_run_hook.py +++ b/tensorflow/contrib/learn/python/learn/session_run_hook.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""This file is deprecated. Use tensorflow.python.training.session_run_hook.""" +"""This file is deprecated. Use `tensorflow.python.training.session_run_hook`. + +See [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/summary_writer_cache.py b/tensorflow/contrib/learn/python/learn/summary_writer_cache.py index 919d415c302b8ec17202aad34ff0bee69bfee2c7..d663cf5fb79c428b0e70d66b0f1305f0559a05c9 100644 --- a/tensorflow/contrib/learn/python/learn/summary_writer_cache.py +++ b/tensorflow/contrib/learn/python/learn/summary_writer_cache.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Wrapper for a Session-like object that handles threads and recovery. +"""Wrapper for a Session-like object that handles threads and recovery (deprecated). + +These are deprecated aliases for classes and functions in `tf.train`. Please use +those directly. Based on an original design of Illia Polosukhin. """ diff --git a/tensorflow/contrib/learn/python/learn/trainable.py b/tensorflow/contrib/learn/python/learn/trainable.py index 429b6040be21d8cbe1f2bba58090366552fdfbe7..a1a3f20dcd8cb5ff7baa559ac41d5e5c40780511 100644 --- a/tensorflow/contrib/learn/python/learn/trainable.py +++ b/tensorflow/contrib/learn/python/learn/trainable.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""`Trainable` interface.""" +"""`Trainable` interface (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division @@ -23,6 +28,8 @@ import abc class Trainable(object): """Interface for objects that are trainable by, e.g., `Experiment`. + + THIS CLASS IS DEPRECATED. """ __metaclass__ = abc.ABCMeta diff --git a/tensorflow/contrib/learn/python/learn/utils/__init__.py b/tensorflow/contrib/learn/python/learn/utils/__init__.py index 48978d0ac34cec2b18e6794dcf3b260bc3b683c4..66d8dc6fd43b383919a16515bc96be492a253bf6 100644 --- a/tensorflow/contrib/learn/python/learn/utils/__init__.py +++ b/tensorflow/contrib/learn/python/learn/utils/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. # ============================================================================== -"""TensorFlow Learn Utils.""" +"""TensorFlow Learn Utils (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/learn/python/learn/utils/export.py b/tensorflow/contrib/learn/python/learn/utils/export.py index cb34cb1d26b6812c7f3f39e9f965615de5a8ef07..3eacac7a3d3dcff4d39025fdee88e16e385b1b84 100644 --- a/tensorflow/contrib/learn/python/learn/utils/export.py +++ b/tensorflow/contrib/learn/python/learn/utils/export.py @@ -13,14 +13,18 @@ # limitations under the License. # ============================================================================== -"""Export utilities.""" +"""Export utilities (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.contrib.framework import deprecated -from tensorflow.python.training import training_util from tensorflow.contrib.session_bundle import exporter from tensorflow.contrib.session_bundle import gc from tensorflow.python.client import session as tf_session @@ -32,6 +36,7 @@ from tensorflow.python.ops import lookup_ops from tensorflow.python.ops import variables 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 @deprecated('2017-03-25', 'Please use Estimator.export_savedmodel() instead.') diff --git a/tensorflow/contrib/learn/python/learn/utils/gc.py b/tensorflow/contrib/learn/python/learn/utils/gc.py index 226915987a4934626066b12810f579ae675107b2..916aecbea88b10bbef316ffb89d4c4d89667cb29 100644 --- a/tensorflow/contrib/learn/python/learn/utils/gc.py +++ b/tensorflow/contrib/learn/python/learn/utils/gc.py @@ -13,7 +13,11 @@ # limitations under the License. # ============================================================================== -r"""System for specifying garbage collection (GC) of path based data. +r"""System for specifying garbage collection (GC) of path based data (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. This framework allows for GC of data specified by path names, for example files on disk. gc.Path objects each represent a single item stored at a path and may @@ -73,10 +77,12 @@ import os from tensorflow.python.platform import gfile from tensorflow.python.util import compat +from tensorflow.python.util.deprecation import deprecated Path = collections.namedtuple('Path', 'path export_version') +@deprecated(None, 'Please implement your own file management or use Saver.') def largest_export_versions(n): """Creates a filter that keeps the largest n export versions. @@ -97,6 +103,7 @@ def largest_export_versions(n): return keep +@deprecated(None, 'Please implement your own file management or use Saver.') def one_of_every_n_export_versions(n): """Creates a filter that keeps one of every n export versions. @@ -128,6 +135,7 @@ def one_of_every_n_export_versions(n): return keep +@deprecated(None, 'Please implement your own file management or use Saver.') def mod_export_version(n): """Creates a filter that keeps every export that is a multiple of n. @@ -146,6 +154,7 @@ def mod_export_version(n): return keep +@deprecated(None, 'Please implement your own file management or use Saver.') def union(lf, rf): """Creates a filter that keeps the union of two filters. @@ -163,6 +172,7 @@ def union(lf, rf): return keep +@deprecated(None, 'Please implement your own file management or use Saver.') def negation(f): """Negate a filter. @@ -179,6 +189,7 @@ def negation(f): return keep +@deprecated(None, 'Please implement your own file name management.') def get_paths(base_dir, parser): """Gets a list of Paths in a given directory. diff --git a/tensorflow/contrib/learn/python/learn/utils/input_fn_utils.py b/tensorflow/contrib/learn/python/learn/utils/input_fn_utils.py index b2521933e524e7ec24d73d4b5171f33e507dd88c..b92eb9fea8b7ccea56c781df74dcfa1cc5508e48 100644 --- a/tensorflow/contrib/learn/python/learn/utils/input_fn_utils.py +++ b/tensorflow/contrib/learn/python/learn/utils/input_fn_utils.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utilities for creating input_fns. +"""Utilities for creating input_fns (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. Contents of this file are moved to tensorflow/python/estimator/export.py. InputFnOps is renamed to ServingInputReceiver. @@ -32,13 +36,17 @@ 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 parsing_ops +from tensorflow.python.util.deprecation import deprecated class InputFnOps(collections.namedtuple('InputFnOps', ['features', 'labels', 'default_inputs'])): - """A return type for an input_fn. + """A return type for an input_fn (deprecated). + + THIS CLASS IS DEPRECATED. Please use tf.estimator.export.ServingInputReceiver + instead. This return type is currently only supported for serving input_fn. Training and eval input_fn should return a `(features, labels)` tuple. @@ -56,6 +64,8 @@ class InputFnOps(collections.namedtuple('InputFnOps', """ +@deprecated(None, 'Please use ' + 'tf.estimator.export.build_parsing_serving_input_receiver_fn.') def build_parsing_serving_input_fn(feature_spec, default_batch_size=None): """Build an input_fn appropriate for serving, expecting fed tf.Examples. @@ -84,6 +94,8 @@ def build_parsing_serving_input_fn(feature_spec, default_batch_size=None): return input_fn +@deprecated(None, 'Please use ' + 'tf.estimator.export.build_raw_serving_input_receiver_fn.') def build_default_serving_input_fn(features, default_batch_size=None): """Build an input_fn appropriate for serving, expecting feature Tensors. diff --git a/tensorflow/contrib/learn/python/learn/utils/inspect_checkpoint.py b/tensorflow/contrib/learn/python/learn/utils/inspect_checkpoint.py index 6a63fb545a56e6040b0b0c3bbb6a17cd96925895..6dbaa15f8391b0044be8e30ca191753beb88db93 100644 --- a/tensorflow/contrib/learn/python/learn/utils/inspect_checkpoint.py +++ b/tensorflow/contrib/learn/python/learn/utils/inspect_checkpoint.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""A simple script for inspect checkpoint files.""" +"""A simple script for inspect checkpoint files (deprecated).""" from __future__ import absolute_import from __future__ import division 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 1593380007b2799fb1d17e92408ab19a7b47fe1e..c7cdb4131215c388412407a008113de13bdd0934 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 @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utilities supporting export to SavedModel. +"""Utilities supporting export to SavedModel (deprecated). + +This module and all its submodules are deprecated. See +[contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) +for migration instructions. Some contents of this file are moved to tensorflow/python/estimator/export.py: @@ -52,8 +56,9 @@ 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 +from tensorflow.python.util.deprecation import deprecated + # A key for use in the input_alternatives dict indicating the default input. # This is the input that will be expected when a serving request does not @@ -77,6 +82,7 @@ FEATURES_INPUT_ALTERNATIVE_KEY = 'features_input_alternative' _FALLBACK_DEFAULT_OUTPUT_ALTERNATIVE_KEY = 'default_output_alternative' +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def build_standardized_signature_def(input_tensors, output_tensors, problem_type): """Build a SignatureDef using problem type and input and output Tensors. @@ -156,6 +162,7 @@ def _is_regression_problem(problem_type, input_tensors, output_tensors): len(input_tensors) == 1 and len(output_tensors) == 1) +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def get_input_alternatives(input_ops): """Obtain all input alternatives using the input_fn output and heuristics.""" input_alternatives = {} @@ -181,6 +188,7 @@ def get_input_alternatives(input_ops): return input_alternatives, features +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def get_output_alternatives(model_fn_ops, default_output_alternative_key=None): """Obtain all output alternatives using the model_fn output and heuristics. @@ -246,6 +254,7 @@ def get_output_alternatives(model_fn_ops, default_output_alternative_key=None): sorted(output_alternatives.keys()))) +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def build_all_signature_defs(input_alternatives, output_alternatives, actual_default_output_alternative_key): """Build `SignatureDef`s from all pairs of input and output alternatives.""" @@ -279,6 +288,7 @@ def build_all_signature_defs(input_alternatives, output_alternatives, MAX_DIRECTORY_CREATION_ATTEMPTS = 10 +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def get_timestamped_export_dir(export_dir_base): """Builds a path to a new subdirectory within the base directory. @@ -317,6 +327,7 @@ def get_timestamped_export_dir(export_dir_base): '{} attempts.'.format(MAX_DIRECTORY_CREATION_ATTEMPTS)) +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def get_temp_export_dir(timestamped_export_dir): """Builds a directory name based on the argument but starting with 'temp-'. @@ -344,6 +355,7 @@ def _export_version_parser(path): return path._replace(export_version=int(filename)) +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def get_most_recent_export(export_dir_base): """Locate the most recent SavedModel export in a directory of many exports. @@ -363,6 +375,7 @@ def get_most_recent_export(export_dir_base): return next(iter(results or []), None) +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def garbage_collect_exports(export_dir_base, exports_to_keep): """Deletes older exports, retaining only a given number of the most recent. @@ -387,6 +400,7 @@ def garbage_collect_exports(export_dir_base, exports_to_keep): logging.warn('Can not delete %s recursively: %s', p.path, e) +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def make_export_strategy(serving_input_fn, default_output_alternative_key=None, assets_extra=None, @@ -400,7 +414,7 @@ def make_export_strategy(serving_input_fn, `InputFnOps`. default_output_alternative_key: the name of the head to serve when an incoming serving request does not explicitly request a specific head. - Must be `None` if the estimator inherits from ${tf.estimator.Estimator} + Must be `None` if the estimator inherits from @{tf.estimator.Estimator} or for single-headed models. assets_extra: A dict specifying how to populate the assets.extra directory within the exported SavedModel. Each key should give the destination @@ -438,7 +452,7 @@ def make_export_strategy(serving_input_fn, The string path to the exported directory. Raises: - ValueError: If `estimator` is a ${tf.estimator.Estimator} instance + ValueError: If `estimator` is a @{tf.estimator.Estimator} instance and `default_output_alternative_key` was specified. """ if isinstance(estimator, core_estimator.Estimator): @@ -469,6 +483,8 @@ def make_export_strategy(serving_input_fn, return export_strategy.ExportStrategy('Servo', export_fn, strip_default_attrs) +@deprecated(None, + 'Use tf.estimator.export.build_parsing_serving_input_receiver_fn') def make_parsing_export_strategy(feature_columns, default_output_alternative_key=None, assets_extra=None, @@ -487,7 +503,7 @@ def make_parsing_export_strategy(feature_columns, that must be provided at serving time (excluding labels!). default_output_alternative_key: the name of the head to serve when an incoming serving request does not explicitly request a specific head. - Must be `None` if the estimator inherits from ${tf.estimator.Estimator} + Must be `None` if the estimator inherits from @{tf.estimator.Estimator} or for single-headed models. assets_extra: A dict specifying how to populate the assets.extra directory within the exported SavedModel. Each key should give the destination @@ -555,8 +571,14 @@ def _default_compare_fn(curr_best_eval_result, cand_eval_result): class BestModelSelector(object): - """A helper that keeps track of export selection candidates.""" + """A helper that keeps track of export selection candidates. + + THIS CLASS IS DEPRECATED. See + [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) + for general migration instructions. + """ + @deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def __init__(self, event_file_pattern=None, compare_fn=None): """Constructor of this class. @@ -622,6 +644,7 @@ class BestModelSelector(object): return best_eval_result +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def make_best_model_export_strategy( serving_input_fn, exports_to_keep=1, @@ -707,6 +730,7 @@ def make_best_model_export_strategy( # TODO(b/67013778): Revisit this approach when corresponding changes to # TF Core are finalized. +@deprecated(None, 'Switch to tf.estimator.Exporter and associated utilities.') def extend_export_strategy(base_export_strategy, post_export_fn, post_export_name=None): @@ -741,7 +765,7 @@ def extend_export_strategy(base_export_strategy, The string path to the SavedModel indicated by post_export_fn. Raises: - ValueError: If `estimator` is a ${tf.estimator.Estimator} instance + ValueError: If `estimator` is a @{tf.estimator.Estimator} instance and `default_output_alternative_key` was specified or if post_export_fn does not return a valid directory. RuntimeError: If unable to create temporary or final export directory. diff --git a/tensorflow/contrib/linalg/BUILD b/tensorflow/contrib/linalg/BUILD index 208e7bc69be76680868c766bc99429eea5870c80..3bc1427bd2f05eec75441f376b7ba41bde30f544 100644 --- a/tensorflow/contrib/linalg/BUILD +++ b/tensorflow/contrib/linalg/BUILD @@ -43,6 +43,24 @@ cuda_py_test( ], ) +cuda_py_test( + name = "linear_operator_block_diag_test", + size = "medium", + srcs = ["python/kernel_tests/linear_operator_block_diag_test.py"], + additional_deps = [ + ":linalg_py", + "//third_party/py/numpy", + "//tensorflow/python:array_ops", + "//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", + ], + shard_count = 4, +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/linalg/__init__.py b/tensorflow/contrib/linalg/__init__.py index 4720692c3384ba1bede1f486c1b1e0e69d10a63a..14cc3b2b4971de1a31960ee33c2f304154b1f411 100644 --- a/tensorflow/contrib/linalg/__init__.py +++ b/tensorflow/contrib/linalg/__init__.py @@ -17,6 +17,7 @@ See the @{$python/contrib.linalg} guide. @@LinearOperator +@@LinearOperatorBlockDiag @@LinearOperatorDiag @@LinearOperatorIdentity @@LinearOperatorScaledIdentity @@ -34,6 +35,7 @@ from __future__ import print_function # pylint: disable=unused-import,wildcard-import,line-too-long,g-importing-member from tensorflow.contrib.linalg.python.ops.linear_operator_addition import * +from tensorflow.contrib.linalg.python.ops.linear_operator_block_diag import * from tensorflow.python.ops.linalg.linear_operator import * from tensorflow.python.ops.linalg.linear_operator_composition import * from tensorflow.python.ops.linalg.linear_operator_diag import * @@ -45,4 +47,5 @@ from tensorflow.python.ops.linalg.linear_operator_lower_triangular import * # pylint: enable=unused-import,wildcard-import,line-too-long,g-importing-member from tensorflow.python.util.all_util import remove_undocumented + remove_undocumented(__name__) diff --git a/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_block_diag_test.py b/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_block_diag_test.py new file mode 100644 index 0000000000000000000000000000000000000000..cc1a047d6a2b6029080fad3f240aa00f50504f07 --- /dev/null +++ b/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_block_diag_test.py @@ -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. +# ============================================================================== + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.linalg.python.ops import linear_operator_block_diag as block_diag +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import random_seed +from tensorflow.python.ops import array_ops +from tensorflow.python.ops.linalg import linalg as linalg_lib +from tensorflow.python.ops.linalg import linear_operator_test_util +from tensorflow.python.ops.linalg import linear_operator_util +from tensorflow.python.platform import test + +linalg = linalg_lib +random_seed.set_random_seed(23) +rng = np.random.RandomState(0) + + +def _block_diag_dense(expected_shape, blocks): + """Convert a list of blocks, into a dense block diagonal matrix.""" + rows = [] + num_cols = 0 + for block in blocks: + # Get the batch shape for the block. + batch_row_shape = array_ops.shape(block)[:-1] + + zeros_to_pad_before_shape = array_ops.concat( + [batch_row_shape, [num_cols]], axis=-1) + zeros_to_pad_before = array_ops.zeros( + shape=zeros_to_pad_before_shape, dtype=block.dtype) + num_cols += array_ops.shape(block)[-1] + zeros_to_pad_after_shape = array_ops.concat( + [batch_row_shape, [expected_shape[-2] - num_cols]], axis=-1) + zeros_to_pad_after = array_ops.zeros( + zeros_to_pad_after_shape, dtype=block.dtype) + + rows.append(array_ops.concat( + [zeros_to_pad_before, block, zeros_to_pad_after], axis=-1)) + + return array_ops.concat(rows, axis=-2) + + +class SquareLinearOperatorBlockDiagTest( + linear_operator_test_util.SquareLinearOperatorDerivedClassTest): + """Most tests done in the base class LinearOperatorDerivedClassTest.""" + + def setUp(self): + # Increase from 1e-6 to 1e-4 + self._atol[dtypes.float32] = 1e-4 + self._atol[dtypes.complex64] = 1e-4 + self._rtol[dtypes.float32] = 1e-4 + self._rtol[dtypes.complex64] = 1e-4 + + @property + def _operator_build_infos(self): + build_info = linear_operator_test_util.OperatorBuildInfo + return [ + build_info((0, 0)), + build_info((1, 1)), + build_info((1, 3, 3)), + build_info((5, 5), blocks=[(2, 2), (3, 3)]), + ] + + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) + expected_blocks = ( + build_info.__dict__["blocks"] if "blocks" in build_info.__dict__ + else [shape]) + matrices = [ + linear_operator_test_util.random_positive_definite_matrix( + block_shape, dtype, force_well_conditioned=True) + for block_shape in expected_blocks + ] + + if use_placeholder: + matrices_ph = [ + array_ops.placeholder(dtype=dtype) for _ in expected_blocks + ] + # Evaluate here because (i) you cannot feed a tensor, and (ii) + # values are random and we want the same value used for both mat and + # feed_dict. + matrices = self.evaluate(matrices) + operator = block_diag.LinearOperatorBlockDiag( + [linalg.LinearOperatorFullMatrix( + m_ph, is_square=True) for m_ph in matrices_ph], + is_square=True) + feed_dict = {m_ph: m for (m_ph, m) in zip(matrices_ph, matrices)} + else: + operator = block_diag.LinearOperatorBlockDiag( + [linalg.LinearOperatorFullMatrix( + m, is_square=True) for m in matrices]) + feed_dict = None + # Should be auto-set. + self.assertTrue(operator.is_square) + + # Broadcast the shapes. + expected_shape = list(build_info.shape) + + matrices = linear_operator_util.broadcast_matrix_batch_dims(matrices) + + block_diag_dense = _block_diag_dense(expected_shape, matrices) + + if not use_placeholder: + block_diag_dense.set_shape( + expected_shape[:-2] + [expected_shape[-1], expected_shape[-1]]) + + return operator, block_diag_dense, feed_dict + + def test_is_x_flags(self): + # Matrix with two positive eigenvalues, 1, and 1. + # The matrix values do not effect auto-setting of the flags. + matrix = [[1., 0.], [1., 1.]] + operator = block_diag.LinearOperatorBlockDiag( + [linalg.LinearOperatorFullMatrix(matrix)], + is_positive_definite=True, + is_non_singular=True, + is_self_adjoint=False) + self.assertTrue(operator.is_positive_definite) + self.assertTrue(operator.is_non_singular) + self.assertFalse(operator.is_self_adjoint) + + def test_is_non_singular_auto_set(self): + # Matrix with two positive eigenvalues, 11 and 8. + # The matrix values do not effect auto-setting of the flags. + matrix = [[11., 0.], [1., 8.]] + operator_1 = linalg.LinearOperatorFullMatrix(matrix, is_non_singular=True) + operator_2 = linalg.LinearOperatorFullMatrix(matrix, is_non_singular=True) + + operator = block_diag.LinearOperatorBlockDiag( + [operator_1, operator_2], + is_positive_definite=False, # No reason it HAS to be False... + is_non_singular=None) + self.assertFalse(operator.is_positive_definite) + self.assertTrue(operator.is_non_singular) + + with self.assertRaisesRegexp(ValueError, "always non-singular"): + block_diag.LinearOperatorBlockDiag( + [operator_1, operator_2], is_non_singular=False) + + def test_name(self): + matrix = [[11., 0.], [1., 8.]] + operator_1 = linalg.LinearOperatorFullMatrix(matrix, name="left") + operator_2 = linalg.LinearOperatorFullMatrix(matrix, name="right") + + operator = block_diag.LinearOperatorBlockDiag([operator_1, operator_2]) + + self.assertEqual("left_ds_right", operator.name) + + def test_different_dtypes_raises(self): + operators = [ + linalg.LinearOperatorFullMatrix(rng.rand(2, 3, 3)), + linalg.LinearOperatorFullMatrix(rng.rand(2, 3, 3).astype(np.float32)) + ] + with self.assertRaisesRegexp(TypeError, "same dtype"): + block_diag.LinearOperatorBlockDiag(operators) + + def test_non_square_operator_raises(self): + operators = [ + linalg.LinearOperatorFullMatrix(rng.rand(3, 4), is_square=False), + linalg.LinearOperatorFullMatrix(rng.rand(3, 3)) + ] + with self.assertRaisesRegexp(ValueError, "square matrices"): + block_diag.LinearOperatorBlockDiag(operators) + + def test_empty_operators_raises(self): + with self.assertRaisesRegexp(ValueError, "non-empty"): + block_diag.LinearOperatorBlockDiag([]) + + +# This test is for blocks with different batch dimensions. +# LinearOperatorFullMatrix doesn't broadcast matmul/solve. +class SquareDiagLinearOperatorBlockDiagTest( + linear_operator_test_util.SquareLinearOperatorDerivedClassTest): + """Most tests done in the base class LinearOperatorDerivedClassTest.""" + + def setUp(self): + # Increase from 1e-6 to 1e-4 + self._atol[dtypes.float32] = 1e-4 + self._atol[dtypes.complex64] = 1e-4 + self._rtol[dtypes.float32] = 1e-4 + self._rtol[dtypes.complex64] = 1e-4 + + @property + def _operator_build_infos(self): + build_info = linear_operator_test_util.OperatorBuildInfo + return [ + build_info((3, 7, 7), blocks=[(1, 2, 2), (3, 2, 2), (1, 3, 3)]), + build_info((2, 1, 6, 6), blocks=[(2, 1, 2, 2), (1, 1, 4, 4)]), + ] + + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) + expected_blocks = ( + build_info.__dict__["blocks"] if "blocks" in build_info.__dict__ + else [shape]) + diag_matrices = [ + linear_operator_test_util.random_uniform( + shape=block_shape[:-1], minval=1., maxval=20., dtype=dtype) + for block_shape in expected_blocks + ] + + if use_placeholder: + diag_matrices_ph = [ + array_ops.placeholder(dtype=dtype) for _ in expected_blocks + ] + diag_matrices = self.evaluate(diag_matrices) + # Evaluate here because (i) you cannot feed a tensor, and (ii) + # values are random and we want the same value used for both mat and + # feed_dict. + operator = block_diag.LinearOperatorBlockDiag( + [linalg.LinearOperatorDiag(m_ph) for m_ph in diag_matrices_ph]) + feed_dict = {m_ph: m for (m_ph, m) in zip( + diag_matrices_ph, diag_matrices)} + else: + operator = block_diag.LinearOperatorBlockDiag( + [linalg.LinearOperatorDiag(m) for m in diag_matrices]) + feed_dict = None + # Should be auto-set. + self.assertTrue(operator.is_square) + + # Broadcast the shapes. + expected_shape = list(build_info.shape) + + matrices = linear_operator_util.broadcast_matrix_batch_dims( + [array_ops.matrix_diag(diag_block) for diag_block in diag_matrices]) + + block_diag_dense = _block_diag_dense(expected_shape, matrices) + if not use_placeholder: + block_diag_dense.set_shape( + expected_shape[:-2] + [expected_shape[-1], expected_shape[-1]]) + + return operator, block_diag_dense, feed_dict + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/linalg/python/ops/linear_operator_block_diag.py b/tensorflow/contrib/linalg/python/ops/linear_operator_block_diag.py new file mode 100644 index 0000000000000000000000000000000000000000..5d7a99664d38eca035bd5a86710050bce4b22c1e --- /dev/null +++ b/tensorflow/contrib/linalg/python/ops/linear_operator_block_diag.py @@ -0,0 +1,358 @@ +# 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. +# ============================================================================== +"""Create a Block Diagonal operator from one or more `LinearOperators`.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import common_shapes +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import check_ops +from tensorflow.python.ops.linalg import linear_operator +from tensorflow.python.ops.linalg import linear_operator_util + + +class LinearOperatorBlockDiag(linear_operator.LinearOperator): + """Combines one or more `LinearOperators` in to a Block Diagonal matrix. + + This operator combines one or more linear operators `[op1,...,opJ]`, + building a new `LinearOperator`, whose underlying matrix representation is + square and has each operator `opi` on the main diagonal, and zero's elsewhere. + + #### Shape compatibility + + If `opj` acts like a [batch] square matrix `Aj`, then `op_combined` acts like + the [batch] square matrix formed by having each matrix `Aj` on the main + diagonal. + + + Each `opj` is required to represent a square matrix, and hence will have + shape `batch_shape_j + [M_j, M_j]`. + + If `opj` has shape `batch_shape_j + [M_j, M_j]`, then the combined operator + has shape `broadcast_batch_shape + [sum M_j, sum M_j]`, where + `broadcast_batch_shape` is the mutual broadcast of `batch_shape_j`, + `j = 1,...,J`, assuming the intermediate batch shapes broadcast. + Even if the combined shape is well defined, the combined operator's + methods may fail due to lack of broadcasting ability in the defining + operators' methods. + + ```python + # Create a 4 x 4 linear operator combined of two 2 x 2 operators. + operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]]) + operator_2 = LinearOperatorFullMatrix([[1., 0.], [0., 1.]]) + operator = LinearOperatorBlockDiag([operator_1, operator_2]) + + operator.to_dense() + ==> [[1., 2., 0., 0.], + [3., 4., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]] + + operator.shape + ==> [4, 4] + + operator.log_abs_determinant() + ==> scalar Tensor + + x1 = ... # Shape [2, 2] Tensor + x2 = ... # Shape [2, 2] Tensor + x = tf.concat([x1, x2], 0) # Shape [2, 4] Tensor + operator.matmul(x) + ==> tf.concat([operator_1.matmul(x1), operator_2.matmul(x2)]) + + # Create a [2, 3] batch of 4 x 4 linear operators. + matrix_44 = tf.random_normal(shape=[2, 3, 4, 4]) + operator_44 = LinearOperatorFullMatrix(matrix) + + # Create a [1, 3] batch of 5 x 5 linear operators. + matrix_55 = tf.random_normal(shape=[1, 3, 5, 5]) + operator_55 = LinearOperatorFullMatrix(matrix_55) + + # Combine to create a [2, 3] batch of 9 x 9 operators. + operator_99 = LinearOperatorBlockDiag([operator_44, operator_55]) + + # Create a shape [2, 3, 9] vector. + x = tf.random_normal(shape=[2, 3, 9]) + operator_99.matmul(x) + ==> Shape [2, 3, 9] Tensor + ``` + + #### Performance + + The performance of `LinearOperatorBlockDiag` on any operation is equal to + the sum of the individual operators' operations. + + + #### Matrix property hints + + This `LinearOperator` is initialized with boolean flags of the form `is_X`, + for `X = non_singular, self_adjoint, positive_definite, square`. + These have the following meaning: + + * If `is_X == True`, callers should expect the operator to have the + property `X`. This is a promise that should be fulfilled, but is *not* a + runtime assert. For example, finite floating point precision may result + in these promises being violated. + * If `is_X == False`, callers should expect the operator to not have `X`. + * If `is_X == None` (the default), callers should have no expectation either + way. + """ + + def __init__(self, + operators, + is_non_singular=None, + is_self_adjoint=None, + is_positive_definite=None, + is_square=True, + name=None): + r"""Initialize a `LinearOperatorBlockDiag`. + + `LinearOperatorBlockDiag` is initialized with a list of operators + `[op_1,...,op_J]`. + + Args: + operators: Iterable of `LinearOperator` objects, each with + the same `dtype` and composable shape. + is_non_singular: Expect that this operator is non-singular. + is_self_adjoint: Expect that this operator is equal to its hermitian + transpose. + is_positive_definite: Expect that this operator is positive definite, + meaning the quadratic form `x^H A x` has positive real part for all + nonzero `x`. Note that we do not require the operator to be + self-adjoint to be positive-definite. See: + https://en.wikipedia.org/wiki/Positive-definite_matrix\ + #Extension_for_non_symmetric_matrices + is_square: Expect that this operator acts like square [batch] matrices. + This is true by default, and will raise a `ValueError` otherwise. + name: A name for this `LinearOperator`. Default is the individual + operators names joined with `_o_`. + + Raises: + TypeError: If all operators do not have the same `dtype`. + ValueError: If `operators` is empty or are non-square. + """ + # Validate operators. + check_ops.assert_proper_iterable(operators) + operators = list(operators) + if not operators: + raise ValueError( + "Expected a non-empty list of operators. Found: %s" % operators) + self._operators = operators + + # Validate dtype. + dtype = operators[0].dtype + for operator in operators: + if operator.dtype != dtype: + name_type = (str((o.name, o.dtype)) for o in operators) + raise TypeError( + "Expected all operators to have the same dtype. Found %s" + % " ".join(name_type)) + + # Auto-set and check hints. + if all(operator.is_non_singular for operator in operators): + if is_non_singular is False: + raise ValueError( + "The direct sum of non-singular operators is always non-singular.") + is_non_singular = True + + if all(operator.is_self_adjoint for operator in operators): + if is_self_adjoint is False: + raise ValueError( + "The direct sum of self-adjoint operators is always self-adjoint.") + is_self_adjoint = True + + if all(operator.is_positive_definite for operator in operators): + if is_positive_definite is False: + raise ValueError( + "The direct sum of positive definite operators is always " + "positive definite.") + is_positive_definite = True + + if not (is_square and all(operator.is_square for operator in operators)): + raise ValueError( + "Can only represent a block diagonal of square matrices.") + + # Initialization. + graph_parents = [] + for operator in operators: + graph_parents.extend(operator.graph_parents) + + if name is None: + # Using ds to mean direct sum. + name = "_ds_".join(operator.name for operator in operators) + with ops.name_scope(name, values=graph_parents): + super(LinearOperatorBlockDiag, self).__init__( + dtype=dtype, + graph_parents=graph_parents, + is_non_singular=is_non_singular, + is_self_adjoint=is_self_adjoint, + is_positive_definite=is_positive_definite, + is_square=True, + name=name) + + @property + def operators(self): + return self._operators + + def _shape(self): + # Get final matrix shape. + domain_dimension = self.operators[0].domain_dimension + range_dimension = self.operators[0].range_dimension + for operator in self.operators[1:]: + domain_dimension += operator.domain_dimension + range_dimension += operator.range_dimension + + matrix_shape = tensor_shape.TensorShape([domain_dimension, range_dimension]) + + # Get broadcast batch shape. + # broadcast_shape checks for compatibility. + batch_shape = self.operators[0].batch_shape + for operator in self.operators[1:]: + batch_shape = common_shapes.broadcast_shape( + batch_shape, operator.batch_shape) + + return batch_shape.concatenate(matrix_shape) + + def _shape_tensor(self): + # Avoid messy broadcasting if possible. + if self.shape.is_fully_defined(): + return ops.convert_to_tensor( + self.shape.as_list(), dtype=dtypes.int32, name="shape") + + domain_dimension = self.operators[0].domain_dimension_tensor() + range_dimension = self.operators[0].range_dimension_tensor() + for operator in self.operators[1:]: + domain_dimension += operator.domain_dimension_tensor() + range_dimension += operator.range_dimension_tensor() + + matrix_shape = array_ops.stack([domain_dimension, range_dimension]) + + # Dummy Tensor of zeros. Will never be materialized. + zeros = array_ops.zeros(shape=self.operators[0].batch_shape_tensor()) + for operator in self.operators[1:]: + zeros += array_ops.zeros(shape=operator.batch_shape_tensor()) + batch_shape = array_ops.shape(zeros) + + return array_ops.concat((batch_shape, matrix_shape), 0) + + def _matmul(self, x, adjoint=False, adjoint_arg=False): + split_dim = -1 if adjoint_arg else -2 + # Split input by rows normally, and otherwise columns. + split_x = self._split_input_into_blocks(x, axis=split_dim) + + result_list = [] + for index, operator in enumerate(self.operators): + result_list += [operator.matmul( + split_x[index], adjoint=adjoint, adjoint_arg=adjoint_arg)] + result_list = linear_operator_util.broadcast_matrix_batch_dims( + result_list) + return array_ops.concat(result_list, axis=-2) + + def _determinant(self): + result = self.operators[0].determinant() + for operator in self.operators[1:]: + result *= operator.determinant() + return result + + def _log_abs_determinant(self): + result = self.operators[0].log_abs_determinant() + for operator in self.operators[1:]: + result += operator.log_abs_determinant() + return result + + def _solve(self, rhs, adjoint=False, adjoint_arg=False): + split_dim = -1 if adjoint_arg else -2 + # Split input by rows normally, and otherwise columns. + split_rhs = self._split_input_into_blocks(rhs, axis=split_dim) + + solution_list = [] + for index, operator in enumerate(self.operators): + solution_list += [operator.solve( + split_rhs[index], adjoint=adjoint, adjoint_arg=adjoint_arg)] + + solution_list = linear_operator_util.broadcast_matrix_batch_dims( + solution_list) + return array_ops.concat(solution_list, axis=-2) + + def _diag_part(self): + diag_list = [] + for operator in self.operators: + # Extend the axis for broadcasting. + diag_list += [operator.diag_part()[..., array_ops.newaxis]] + diag_list = linear_operator_util.broadcast_matrix_batch_dims(diag_list) + diagonal = array_ops.concat(diag_list, axis=-2) + return array_ops.squeeze(diagonal, axis=-1) + + def _trace(self): + result = self.operators[0].trace() + for operator in self.operators[1:]: + result += operator.trace() + return result + + def _to_dense(self): + num_cols = 0 + rows = [] + broadcasted_blocks = [operator.to_dense() for operator in self.operators] + broadcasted_blocks = linear_operator_util.broadcast_matrix_batch_dims( + broadcasted_blocks) + for block in broadcasted_blocks: + batch_row_shape = array_ops.shape(block)[:-1] + + zeros_to_pad_before_shape = array_ops.concat( + [batch_row_shape, [num_cols]], axis=-1) + zeros_to_pad_before = array_ops.zeros( + shape=zeros_to_pad_before_shape, dtype=block.dtype) + num_cols += array_ops.shape(block)[-1] + zeros_to_pad_after_shape = array_ops.concat( + [batch_row_shape, + [self.domain_dimension_tensor() - num_cols]], axis=-1) + zeros_to_pad_after = array_ops.zeros( + shape=zeros_to_pad_after_shape, dtype=block.dtype) + + rows.append(array_ops.concat( + [zeros_to_pad_before, block, zeros_to_pad_after], axis=-1)) + + mat = array_ops.concat(rows, axis=-2) + mat.set_shape(self.shape) + return mat + + def _split_input_into_blocks(self, x, axis=-1): + """Split `x` into blocks matching `operators`'s `domain_dimension`. + + Specifically, if we have a block diagonal matrix, with block sizes + `[M_j, M_j] j = 1..J`, this method splits `x` on `axis` into `J` + tensors, whose shape at `axis` is `M_j`. + + Args: + x: `Tensor`. `x` is split into `J` tensors. + axis: Python `Integer` representing the axis to split `x` on. + + Returns: + A list of `Tensor`s. + """ + block_sizes = [] + if self.shape.is_fully_defined(): + for operator in self.operators: + block_sizes += [operator.domain_dimension.value] + else: + for operator in self.operators: + block_sizes += [operator.domain_dimension_tensor()] + + return array_ops.split(x, block_sizes, axis=axis) diff --git a/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py b/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py index 70f777f08bd5b8157e601f19019075d3e7543811..cfe62fac43b35d863eb559b95057ae62a41bed49 100644 --- a/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py +++ b/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py @@ -270,14 +270,14 @@ class SdcaWithLogisticLossTest(SdcaModelTest): train_op = lr.minimize() - def Minimize(): + def minimize(): with self._single_threaded_test_session(): for _ in range(_MAX_ITERATIONS): - train_op.run() + train_op.run() # pylint: disable=cell-var-from-loop threads = [] for _ in range(num_loss_partitions): - threads.append(threading.Thread(target=Minimize)) + threads.append(threading.Thread(target=minimize)) threads[-1].start() for t in threads: @@ -395,7 +395,7 @@ class SdcaWithLogisticLossTest(SdcaModelTest): predicted_labels = get_binary_predictions_for_logistic(predictions) self.assertAllClose([0, 1, 1, 1], predicted_labels.eval()) self.assertAllClose( - 0.01, lr.approximate_duality_gap().eval(), rtol=1e-2, atol=1e-2) + 0.0, lr.approximate_duality_gap().eval(), rtol=1e-2, atol=1e-2) def testFractionalExampleLabel(self): # Setup test data with 1 positive, and 1 mostly-negative example. @@ -407,7 +407,7 @@ class SdcaWithLogisticLossTest(SdcaModelTest): make_example_proto({ 'age': [1], 'gender': [1] - }, 1), + }, 0.9), ] example_weights = [1.0, 1.0] for num_shards in _SHARD_NUMBERS: 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/lite/BUILD b/tensorflow/contrib/lite/BUILD index 13350c5a438b75fe14e8753e5bb1bb77ec8f655b..5cfbb544b73991195a7bba9528ee9550104f3d78 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,11 +28,6 @@ config_setting( }, ) -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", -) - cc_library( name = "schema_fbs_version", hdrs = ["version.h"], @@ -53,6 +51,8 @@ cc_test( srcs = ["arena_planner_test.cc"], deps = [ ":arena_planner", + "//tensorflow/contrib/lite/testing:util", + "//tensorflow/core:lib", "@com_google_googletest//:gtest", ], ) @@ -107,6 +107,7 @@ cc_library( srcs = [ "allocation.cc", "error_reporter.cc", + "graph_info.cc", "interpreter.cc", "model.cc", "nnapi_delegate.cc", @@ -116,6 +117,7 @@ cc_library( "allocation.h", "context.h", "error_reporter.h", + "graph_info.h", "interpreter.h", "model.h", "nnapi_delegate.h", @@ -130,6 +132,7 @@ cc_library( ":memory_planner", ":schema_fbs_version", ":simple_memory_arena", + ":util", "//tensorflow/contrib/lite/kernels:gemm_support", "//tensorflow/contrib/lite/nnapi:nnapi_lib", "//tensorflow/contrib/lite/schema:schema_fbs", @@ -167,6 +170,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", ], ) @@ -214,22 +233,43 @@ cc_test( ], ) +cc_library( + name = "util", + srcs = ["util.cc"], + hdrs = ["util.h"], + deps = [ + ":context", + ], +) + +cc_test( + name = "util_test", + size = "small", + srcs = ["util_test.cc"], + deps = [ + ":context", + ":util", + "//tensorflow/contrib/lite/testing:util", + "@com_google_googletest//:gtest", + ], +) + # Test the serialization of a model with optional tensors. # 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 7f316292724ea0baaf034d4e914773ad97a957d4..b4504f246a0f806d35d8c3d659717a86d2f2a4f5 100644 --- a/tensorflow/contrib/lite/Makefile +++ b/tensorflow/contrib/lite/Makefile @@ -27,10 +27,10 @@ LIBDIR := $(MAKEFILE_DIR)/gen/lib/ GENDIR := $(MAKEFILE_DIR)/gen/obj/ # Settings for the host compiler. -CXX := $(CC_PREFIX) gcc +CXX := $(CC_PREFIX)gcc CXXFLAGS := --std=c++11 -O3 -DNDEBUG -CC := $(CC_PREFIX) gcc -CFLAGS := +CC := $(CC_PREFIX)gcc +CFLAGS := -O3 -DNDEBUG LDOPTS := LDOPTS += -L/usr/local/lib ARFLAGS := -r @@ -57,10 +57,11 @@ LIBS := \ # If we're on Linux, also link in the dl library. ifeq ($(HOST_OS),LINUX) - LIBS += -ldl -lpthread + LIBS += -ldl endif include $(MAKEFILE_DIR)/ios_makefile.inc +include $(MAKEFILE_DIR)/rpi_makefile.inc # This library is the main target for this makefile. It will contain a minimal # runtime that can be linked in to other programs. diff --git a/tensorflow/contrib/lite/README.md b/tensorflow/contrib/lite/README.md index 55a524b207b258e794f97e68a96cf01dc60efb7f..2680d515ebee941f6b80513b241cf520adc7e510 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 (299 * 299 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 `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. @@ -158,24 +165,24 @@ bazel-bin/tensorflow/python/tools/freeze_graph\ --input_graph=/tmp/mobilenet_v1_224.pb \ --input_checkpoint=/tmp/checkpoints/mobilenet-10202.ckpt \ --input_binary=true --output_graph=/tmp/frozen_mobilenet_v1_224.pb \ - --output_node_names=MobileNet/Predictions/Reshape_1 + --output_node_names=MobilenetV1/Predictions/Reshape_1 ``` 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 ``` @@ -211,7 +218,7 @@ 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/arena_planner.cc b/tensorflow/contrib/lite/arena_planner.cc index bf1bcdd1a7a7d3395c45ae95abd5980e9ffc0fc6..8e47e2375e2e306c345a2b6caa2411abd9b3ceb0 100644 --- a/tensorflow/contrib/lite/arena_planner.cc +++ b/tensorflow/contrib/lite/arena_planner.cc @@ -128,6 +128,11 @@ TfLiteStatus ArenaPlanner::PlanAllocations() { } TfLiteStatus ArenaPlanner::ExecuteAllocations(int first_node, int last_node) { + // Grow the size of `allocs_` if necessary. This allows allocating temporary + // tensors in op's `prepare` function. + TF_LITE_ENSURE(context_, graph_info_->num_tensors() >= allocs_.size()); + allocs_.resize(graph_info_->num_tensors()); + TF_LITE_ENSURE_STATUS(CalculateAllocations(first_node, last_node)); TF_LITE_ENSURE_STATUS(Commit()); @@ -185,8 +190,12 @@ TfLiteStatus ArenaPlanner::CalculateAllocations(int first_node, int last_node) { TfLiteStatus ArenaPlanner::ResolveTensorAllocation(int tensor_index) { TfLiteTensor& tensor = *graph_info_->tensor(tensor_index); if (tensor.allocation_type == kTfLiteArenaRw) { - TF_LITE_ENSURE_STATUS( - arena_.ResolveAlloc(context_, allocs_[tensor_index], &tensor.data.raw)); + // 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( diff --git a/tensorflow/contrib/lite/arena_planner.h b/tensorflow/contrib/lite/arena_planner.h index 58bc164619c2c053b9492e9a0e5de2da30e199af..f84b3dad9550e789237c8e45971002c7d336b9d3 100644 --- a/tensorflow/contrib/lite/arena_planner.h +++ b/tensorflow/contrib/lite/arena_planner.h @@ -33,7 +33,7 @@ class AllocationInfo; // 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 +// share some of the buffer 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 diff --git a/tensorflow/contrib/lite/arena_planner_test.cc b/tensorflow/contrib/lite/arena_planner_test.cc index c27c327abc63d7bd1e3912d368a1dacb62c50ca8..a8a8755e2c9e81474f2ff9cd2b85c0eb3d5c3441 100644 --- a/tensorflow/contrib/lite/arena_planner_test.cc +++ b/tensorflow/contrib/lite/arena_planner_test.cc @@ -18,6 +18,8 @@ limitations under the License. #include #include +#include "tensorflow/contrib/lite/testing/util.h" +#include "tensorflow/core/platform/logging.h" namespace tflite { namespace { @@ -191,8 +193,8 @@ TEST_F(ArenaPlannerTest, GraphWithNoOps) { EXPECT_EQ(GetOffset(10), GetOffsetAfter(0)); // The outputs are never allocated because they are not connected to any // inputs. - EXPECT_EQ(GetOffset(5), 0); - EXPECT_EQ(GetOffset(11), 0); + EXPECT_TRUE((*graph.tensors())[5].data.raw == nullptr); + EXPECT_TRUE((*graph.tensors())[11].data.raw == nullptr); } TEST_F(ArenaPlannerTest, GraphWithOneOp) { @@ -371,11 +373,7 @@ TEST_F(ArenaPlannerTest, LargerGraphAndStepwiseAllocation) { SetGraph(&graph); auto is_unallocated = [&](int tensor_index) { - // TODO(ahentz): We'd to use nullptr to represent unallocated tensors, but - // the current code still points them all to the beginning fo the alloc - // (that is, zero offset). - // return (*graph.tensors())[tensor_index].data.raw == nullptr; - return GetOffset(tensor_index) == 0; + return (*graph.tensors())[tensor_index].data.raw == nullptr; }; // The allocation plan is made at the beginning and is independent of @@ -464,9 +462,7 @@ TEST_F(ArenaPlannerTest, LargerGraphAndStepwiseAllocation) { } // namespace tflite int main(int argc, char** argv) { - // ::tflite::LogToStderr(); - FLAGS_logtostderr = true; - + ::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 19829e4991651111e13fc1805f97daef8bc016a7..2813d1c347163e67c70983d3dd49773f4a4b4544 100644 --- a/tensorflow/contrib/lite/build_def.bzl +++ b/tensorflow/contrib/lite/build_def.bzl @@ -104,7 +104,7 @@ def tflite_jni_binary(name, """Builds a jni binary for TFLite.""" linkopts = linkopts + [ "-Wl,--version-script", # Export only jni functions & classes. - linkscript, + "$(location {})".format(linkscript), ] native.cc_binary( name=name, diff --git a/tensorflow/contrib/ndlstm/__init__.py b/tensorflow/contrib/lite/build_rpi_lib.sh old mode 100644 new mode 100755 similarity index 76% rename from tensorflow/contrib/ndlstm/__init__.py rename to tensorflow/contrib/lite/build_rpi_lib.sh index da89bb4ab605e1570a119a88d481f3ad3d1fee4c..3824b16412ed26a6cab79df3242da6017c3322b0 --- a/tensorflow/contrib/ndlstm/__init__.py +++ b/tensorflow/contrib/lite/build_rpi_lib.sh @@ -1,3 +1,4 @@ +#!/bin/bash -x # Copyright 2017 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,9 +14,9 @@ # limitations under the License. # ============================================================================== -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function +set -e -from tensorflow.contrib.ndlstm.python import lstm2d -from tensorflow.contrib.ndlstm.python import lstm1d +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR/../../.." + +CC_PREFIX=arm-linux-gnueabihf- make -j 3 -f tensorflow/contrib/lite/Makefile TARGET=RPI TARGET_ARCH=armv7 diff --git a/tensorflow/contrib/lite/builtin_op_data.h b/tensorflow/contrib/lite/builtin_op_data.h index a1037a525c35b1e3e04e8b53f423696375a6ef84..5fc8954743e5b3b458e5c2004f4378cbad6056c0 100644 --- a/tensorflow/contrib/lite/builtin_op_data.h +++ b/tensorflow/contrib/lite/builtin_op_data.h @@ -151,6 +151,7 @@ typedef struct { } TfLiteLSTMParams; typedef struct { + bool align_corners; } TfLiteResizeBilinearParams; typedef struct { @@ -194,6 +195,10 @@ 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. diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h new file mode 100644 index 0000000000000000000000000000000000000000..ea3ae3489ecf07b22a02829c5235ad59264496af --- /dev/null +++ b/tensorflow/contrib/lite/builtin_ops.h @@ -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. +==============================================================================*/ + +#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 + +// The enum for builtin operators. +// Note: CUSTOM and DELEGATE are 2 special ops which are not real builtin +// ops. +typedef enum { + kTfLiteBuiltinAdd = 0, + kTfLiteBuiltinAveragePool2d = 1, + kTfLiteBuiltinConcatenation = 2, + kTfLiteBuiltinConv2d = 3, + kTfLiteBuiltinDepthwiseConv2d = 4, + kTfLiteBuiltinDequantize = 6, + 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, + kTfLiteBuiltinLogSoftmax = 50, + kTfLiteBuiltinDelegate = 51, + kTfLiteBuiltinBidirectionalSequenceLstm = 52, + kTfLiteBuiltinCast = 53, +} TfLiteBuiltinOperator; + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus +#endif // TENSORFLOW_CONTRIB_LITE_BUILTIN_OPS_H_ +} diff --git a/tensorflow/contrib/lite/context.c b/tensorflow/contrib/lite/context.c index c09e838c5c2e50e0f4a38eaf66e55246fd9a6f7f..5c6f5e72a47180cd98be46f60cfa8eaf28197806 100644 --- a/tensorflow/contrib/lite/context.c +++ b/tensorflow/contrib/lite/context.c @@ -17,9 +17,14 @@ limitations under the License. #include #include +int TfLiteIntArrayGetSizeInBytes(int size) { + static TfLiteIntArray dummy; + return sizeof(dummy) + sizeof(dummy.data[0]) * size; +} + TfLiteIntArray* TfLiteIntArrayCreate(int size) { TfLiteIntArray* ret = - (TfLiteIntArray*)malloc(sizeof(*ret) + sizeof(ret->data[0]) * size); + (TfLiteIntArray*)malloc(TfLiteIntArrayGetSizeInBytes(size)); ret->size = size; return ret; } @@ -55,12 +60,16 @@ TfLiteIntArray* TfLiteIntArrayCopy(TfLiteIntArray* src) { void TfLiteIntArrayFree(TfLiteIntArray* a) { free(a); } -void TfLiteTensorFree(TfLiteTensor* t) { +void TfLiteTensorDataFree(TfLiteTensor* t) { if (t->allocation_type == kTfLiteDynamic && t->data.raw) { free(t->data.raw); } - if (t->dims) TfLiteIntArrayFree(t->dims); t->data.raw = NULL; +} + +void TfLiteTensorFree(TfLiteTensor* t) { + TfLiteTensorDataFree(t); + if (t->dims) TfLiteIntArrayFree(t->dims); t->dims = NULL; } diff --git a/tensorflow/contrib/lite/context.h b/tensorflow/contrib/lite/context.h index d6dfc20ae829b13e9cb45efcf9e14af5d4b69b48..45184b05ecefb504c75815ae900f3b605359a443 100644 --- a/tensorflow/contrib/lite/context.h +++ b/tensorflow/contrib/lite/context.h @@ -29,6 +29,7 @@ limitations under the License. #ifndef TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ #define TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ +#include #include #include @@ -38,6 +39,10 @@ extern "C" { typedef enum { kTfLiteOk = 0, kTfLiteError = 1 } TfLiteStatus; +// Forward declare so GetNode can use this is in Context. +typedef struct _TfLiteRegistration TfLiteRegistration; +typedef struct _TfLiteDelegate TfLiteDelegate; + #define kOptionalTensor (-1) // Fixed size list of integers. Used for dimensions and inputs/outputs tensor @@ -54,6 +59,10 @@ typedef struct { #endif } TfLiteIntArray; +// Given the size (number of elements) in a TfLiteIntArray, calculate its size +// in bytes. +int TfLiteIntArrayGetSizeInBytes(int size); + // Create a array of a given `size` (uninitialized entries). // This returns a pointer, that you must free using TfLiteIntArrayFree(). TfLiteIntArray* TfLiteIntArrayCreate(int size); @@ -159,6 +168,11 @@ typedef enum { kTfLiteDynamic, } TfLiteAllocationType; +// The delegates should use zero or positive integers to represent handles. +// -1 is reserved from unallocated status. +typedef int TfLiteBufferHandle; +const TfLiteBufferHandle kTfLiteNullBufferHandle = -1; + // An tensor in the interpreter system which is a wrapper around a buffer of // data including a dimensionality (or NULL if not currently defined). typedef struct { @@ -191,8 +205,27 @@ typedef struct { // Null-terminated name of this tensor. const char* name; + + // The delegate which knows how to handle `buffer_handle`. + // WARNING: This is an experimental interface that is subject to change. + TfLiteDelegate* delegate; + + // An integer buffer handle that can be handled by `delegate`. + // The value is valid only when delegate is not null. + // WARNING: This is an experimental interface that is subject to change. + TfLiteBufferHandle buffer_handle; + + // If the delegate uses its own buffer (e.g. GPU memory), the delegate is + // responsible to set data_is_stale to true. + // `delegate->CopyFromBufferHandle` can be called to copy the data from + // delegate buffer. + // WARNING: This is an // experimental interface that is subject to change. + bool data_is_stale; } TfLiteTensor; +// Free data memory of tensor `t`; +void TfLiteTensorDataFree(TfLiteTensor* t); + // Free memory of tensor `t`; void TfLiteTensorFree(TfLiteTensor* t); @@ -205,10 +238,62 @@ 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; + + // The pointer to the delegate. This is non-null only when the node is + // created by calling `interpreter.ModifyGraphWithDelegate`. + // WARNING: This is an experimental interface that is subject to change. + TfLiteDelegate* delegate; +} TfLiteNode; + typedef struct TfLiteContext { // Number of tensors in the context. int tensors_size; - // An tensor of tensors in the interpreter context (of length `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 array of tensors in the interpreter context (of length `tensors_size`) TfLiteTensor* tensors; // opaque full context ptr (an opaque c++ data structure) @@ -227,34 +312,29 @@ 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 one or more stub delegate operations. This function + // does not take ownership of `nodes_to_replace`. + TfLiteStatus (*ReplaceSubgraphsWithDelegateKernels)( + struct TfLiteContext*, TfLiteRegistration registration, + const TfLiteIntArray* nodes_to_replace, TfLiteDelegate* delegate); + + // Number of threads that are recommended to subsystems like gemmlowp and + // eigen. + int recommended_num_threads; + // TODO(ahentz): we should create a more general mechanism for this sort of // library-global objects. void* gemm_context; + void* eigen_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*). @@ -291,8 +371,54 @@ 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 _TfLiteDelegate { + // 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, TfLiteDelegate* delegate); + + // Copy the data from delegate buffer handle to raw memory. + // This can be null if the delegate doesn't use its own buffer. + TfLiteStatus (*CopyFromBufferHandle)(TfLiteDelegate* delegate, + TfLiteBufferHandle buffer_handle, + void* data, int size); + + // Copy the data from raw memory to delegate buffer handle. + // This can be null if the delegate doesn't use its own buffer. + TfLiteStatus (*CopyToBufferHandle)(TfLiteDelegate* delegate, + TfLiteBufferHandle buffer_handle, + void* data, int size); + + // Free the Delegate Buffer Handle. Note: This only frees the handle, but + // this doesn't release the underlying resource (e.g. textures). The + // resources are either owned by application layer or the delegate. + // This can be null if the delegate doesn't use its own buffer. + void (*FreeBufferHandle)(TfLiteDelegate* delegate, + TfLiteBufferHandle* handle); +} TfLiteDelegate; + +// WARNING: This is an experimental interface that is subject to change. +typedef struct { + TfLiteDelegate* delegate; + TfLiteIntArray* nodes_to_replace; + TfLiteIntArray* input_tensors; + TfLiteIntArray* output_tensors; +} TfLiteDelegateParams; + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/tensorflow/contrib/lite/download_dependencies.sh b/tensorflow/contrib/lite/download_dependencies.sh index e1b7b3613a041287ff3cc4eeff8afd7cfcede174..a93ed201d647ddf2359a57254a959871c13fb94f 100755 --- a/tensorflow/contrib/lite/download_dependencies.sh +++ b/tensorflow/contrib/lite/download_dependencies.sh @@ -36,8 +36,6 @@ ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_ 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. @@ -93,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" @@ -103,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 da193d2586e9123341b9a41be049ee2a4382017a..3c5f805f12f6a1fb7185c140604f692ac282a143 100644 --- a/tensorflow/contrib/lite/error_reporter.h +++ b/tensorflow/contrib/lite/error_reporter.h @@ -30,7 +30,7 @@ namespace tflite { // va_list args; // foo.Report("test %d", args); // where args is va_list // -// Sublclass ErrorReporter to provide another reporting destination. +// Subclass ErrorReporter to provide another reporting destination. // For example, if you have a GUI program, you might redirect to a buffer // that drives a GUI error log box. class ErrorReporter { 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/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 index 476d85c0314e331d6d3bad382c331a8458fd01a1..959347b5491514ddc13af57ea6f7385a0d39e418 100644 --- a/tensorflow/contrib/lite/examples/label_image/BUILD +++ b/tensorflow/contrib/lite/examples/label_image/BUILD @@ -42,7 +42,15 @@ cc_library( "bitmap_helpers_impl.h", "label_image.h", ], - deps = ["//tensorflow/contrib/lite:string"], + 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 diff --git a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h index 860e27e5ba9cc9fe23d2a7f9f65dd53bbf76f7a3..97343dde6b31694e5b2de20b35a7083fb8fe4a0e 100644 --- a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h +++ b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h @@ -13,8 +13,8 @@ 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 +#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" @@ -26,15 +26,15 @@ uint8_t* read_bmp(const std::string& input_bmp_name, int* width, int* height, int* channels, Settings* s); template -void downsize(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); +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 downsize(uint8_t*, unsigned char*, int, int, int, int, - int, int, Settings*); -template void downsize(float*, unsigned char*, int, int, int, int, int, +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 diff --git a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h index 64a931082b0cbb4632ec3a814ce654d4f9106bc1..2a64c1de725b601e9b6e9325d9faacb37df0e626 100644 --- a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h +++ b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h @@ -13,8 +13,20 @@ 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 +#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" @@ -22,28 +34,70 @@ namespace tflite { namespace label_image { template -void downsize(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) { - for (int y = 0; y < wanted_height; ++y) { - const int in_y = (y * image_height) / wanted_height; - uint8_t* in_row = in + (in_y * image_width * image_channels); - T* out_row = out + (y * wanted_width * wanted_channels); - for (int x = 0; x < wanted_width; ++x) { - const int in_x = (x * image_width) / wanted_width; - uint8_t* in_pixel = in_row + (in_x * image_channels); - T* out_pixel = out_row + (x * wanted_channels); - for (int c = 0; c < wanted_channels; ++c) { - if (s->input_floating) - out_pixel[c] = (in_pixel[c] - s->input_mean) / s->input_std; - else - out_pixel[c] = in_pixel[c]; - } - } +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 +#endif // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H_ diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.cc b/tensorflow/contrib/lite/examples/label_image/label_image.cc index d7f49ad8757e8899fe9c23b985edff6ba7f68750..a91467d345fdce1268635a69a96939921dc170e8 100644 --- a/tensorflow/contrib/lite/examples/label_image/label_image.cc +++ b/tensorflow/contrib/lite/examples/label_image/label_image.cc @@ -151,14 +151,14 @@ void RunInference(Settings* s) { switch (interpreter->tensor(input)->type) { case kTfLiteFloat32: s->input_floating = true; - downsize(interpreter->typed_tensor(input), in, - image_height, image_width, image_channels, - wanted_height, wanted_width, wanted_channels, s); + resize(interpreter->typed_tensor(input), in, image_height, + image_width, image_channels, wanted_height, wanted_width, + wanted_channels, s); break; case kTfLiteUInt8: - downsize(interpreter->typed_tensor(input), in, - image_height, image_width, image_channels, - wanted_height, wanted_width, wanted_channels, s); + 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 " @@ -188,9 +188,8 @@ void RunInference(Settings* s) { 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); + 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), 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/ios.md b/tensorflow/contrib/lite/g3doc/ios.md index a359b8d4b481dbc15cc86db14eabda5433722b8b..e0358a444d6dffc377bf13ee72ba5477359d6e07 100644 --- a/tensorflow/contrib/lite/g3doc/ios.md +++ b/tensorflow/contrib/lite/g3doc/ios.md @@ -22,6 +22,15 @@ Then install brew install automake brew install libtool ``` +If you get an error where either automake or libtool install but do not link correctly, you'll first need to: +```bash +sudo chown -R $(whoami) /usr/local/* +``` +Then follow the instructions to perform the linking: +```bash +brew link automake +brew link libtool +``` Then you need to run a shell script to download the dependencies you need: diff --git a/tensorflow/contrib/lite/g3doc/rpi.md b/tensorflow/contrib/lite/g3doc/rpi.md new file mode 100644 index 0000000000000000000000000000000000000000..7a3a231626d0e1c71e474ff4ff16789ebe2901db --- /dev/null +++ b/tensorflow/contrib/lite/g3doc/rpi.md @@ -0,0 +1,50 @@ +# TensorFlow Lite for Raspberry Pi + +## Cross compiling +### Installing toolchian +This has been tested on Ubuntu 16.04.3 64bit and Tensorflow devel docker image [tensorflow/tensorflow:nightly-devel](https://hub.docker.com/r/tensorflow/tensorflow/tags/). + +To cross compiling TensorFlow Lite. First you should install the toolchain and libs. +```bash +sudo apt-get update +sudo apt-get install crossbuild-essential-armhf +``` +> If you are using docker, you may not use `sudo` + +### Building +Clone this Tensorflow repository, Run this script at the root of the repository to download all the dependencies: +> The Tensorflow repository is in `/tensorflow` if you are using `tensorflow/tensorflow:nightly-devel` docker image, just try it. +```bash +./tensorflow/contrib/lite/download_dependencies.sh +``` +Note than you only need to to this once. + +You should then be able to compile: +```bash +./tensorflow/contrib/lite/build_rpi_lib.sh +``` + +This should compile a static library in: +`tensorflow/contrib/lite/gen/lib/rpi_armv7/libtensorflow-lite.a`. + +## Native compiling +This has been tested on Raspberry Pi 3b, Raspbian GNU/Linux 9.1 (stretch), gcc version 6.3.0 20170516 (Raspbian 6.3.0-18+rpi1). + +Log in to you RPI, install the toolchain. +```bash +sudo apt-get instal build-essential +``` + +First, clone this TensorFlow repository. Run this at the root of the repository: +```bash +./tensorflow/contrib/lite/download_dependencies.sh +``` +Note than you only need to to this once. + +You should then be able to compile: +```bash +./tensorflow/contrib/lite/build_rpi_lib.sh +``` + +This should compile a static library in: +`tensorflow/contrib/lite/gen/lib/rpi_armv7/libtensorflow-lite.a`. diff --git a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md index 8e5e694a5cbe7f908572114db33c8257db6151f0..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 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 index 57690058c4630f75f8b23073f4ab44f27090c51b..313af5fb7574b42bcdd53b4baad06e4ccfb34053 100644 --- a/tensorflow/contrib/lite/graph_info.h +++ b/tensorflow/contrib/lite/graph_info.h @@ -48,6 +48,32 @@ class GraphInfo { 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 69a597dc5a219b55eced6ec8da5b388caf372b8e..cee57bba5ec8a35b0935c171774eb2ba62946afe 100644 --- a/tensorflow/contrib/lite/interpreter.cc +++ b/tensorflow/contrib/lite/interpreter.cc @@ -25,17 +25,37 @@ limitations under the License. #include "tensorflow/contrib/lite/kernels/gemm_support.h" #include "tensorflow/contrib/lite/memory_planner.h" #include "tensorflow/contrib/lite/nnapi_delegate.h" +#include "tensorflow/contrib/lite/schema/schema_generated.h" +#include "tensorflow/contrib/lite/util.h" + +namespace tflite { namespace { -// std::vector preallocation tuning. -constexpr const int kSlotsToReserve = 128; +// Stub method which returns kTfLiteError when the function is forbidden. +// We're registrating this function to several different function to save +// compiled binary size. Please note the restrictions: +// * The type of first parameter have to be `TfLiteContext*`. +// * All paramteters must be trivailly destructible. (E.g. No C++ class) +TfLiteStatus ForbiddenContextFunction(TfLiteContext* context, ...) { + context->ReportError(context, + "The function is forbidden if not calling in delegate."); + return kTfLiteError; +} -} // namespace +// Set the ForbiddenContextFunction to a compatible function pointer. +template +void SetForbiddenContextFunction(FunctionType* func) { + *func = reinterpret_cast(ForbiddenContextFunction); +} -namespace tflite { +} // namespace // 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) @@ -45,9 +65,12 @@ class InterpreterInfo : public GraphInfo { TfLiteTensor* tensor(size_t index) override { return interpreter_->tensor(index); } - size_t num_nodes() const override { return interpreter_->nodes_size(); } + size_t num_nodes() const override { + return interpreter_->execution_plan().size(); + } const TfLiteNode& node(size_t index) const override { - return interpreter_->node_and_registration(index)->first; + int node_index = interpreter_->execution_plan()[index]; + return interpreter_->node_and_registration(node_index)->first; } const std::vector& inputs() const override { return interpreter_->inputs(); @@ -69,11 +92,19 @@ Interpreter::Interpreter(ErrorReporter* error_reporter) context_.AddTensors = AddTensors; context_.tensors = nullptr; context_.tensors_size = 0; + context_.eigen_context = nullptr; context_.gemm_context = nullptr; + context_.recommended_num_threads = -1; + + // Invalid to call these these except from TfLiteDelegate + SetForbiddenContextFunction(&context_.GetNodeAndRegistration); + SetForbiddenContextFunction(&context_.ReplaceSubgraphsWithDelegateKernels); + SetForbiddenContextFunction(&context_.GetExecutionPlan); + // Reserve some space for the tensors to avoid excessive resizing. - tensors_.reserve(kSlotsToReserve); - nodes_and_registration_.reserve(kSlotsToReserve); - next_node_to_prepare_ = 0; + tensors_.reserve(kTensorsReservedCapacity); + nodes_and_registration_.reserve(kTensorsReservedCapacity); + next_execution_plan_index_to_prepare_ = 0; UseNNAPI(false); } @@ -89,10 +120,172 @@ Interpreter::~Interpreter() { } for (int i = 0; i < context_.tensors_size; i++) { - TfLiteTensorFree(&context_.tensors[i]); + TfLiteTensor* tensor = &context_.tensors[i]; + if (tensor->buffer_handle != kTfLiteNullBufferHandle) { + tensor->delegate->FreeBufferHandle(tensor->delegate, + &tensor->buffer_handle); + } + TfLiteTensorFree(tensor); } } +TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( + TfLiteContext* context, TfLiteRegistration registration, + const TfLiteIntArray* nodes_to_replace, TfLiteDelegate* delegate) { + return static_cast(context->impl_) + ->ReplaceSubgraphsWithDelegateKernels(registration, nodes_to_replace, + delegate); +} + +namespace { + +// Copy a std::vector to an existing TfLiteIntArray. +// This is a low-level data manipulation function, and it's caller's +// responsibility to ensure TfLiteIntArray has enough size. +void CopyVectorToTfLiteIntArray(const std::vector& vec, + TfLiteIntArray* arr) { + arr->size = vec.size(); + memcpy(arr->data, vec.data(), sizeof(int) * arr->size); +} + +// This function allocates a continuous memory space that contains a +// TfLiteDelegateParams followed by a several TfLiteIntArray. +// When calling `free` at TfLiteDelegateParams*, all the allocated space +// will be freed together. +// +// +-----------------------------------+ +// | TfLiteDelegateParams | +// | TfLiteDelegate* delegate; | +// | TfLiteIntArray* nodes_to_replace; |--\ +// | TfLiteIntArray* input_tensors; |--+--\ +// | TfLiteIntArray* output_tensors; |--+--+--\ +// +-----------------------------------+ | | | +// | TfLiteIntArray (variable size) |<-/ | | +// +-----------------------------------+ | | +// | TfLiteIntArray (variable size) |<----/ | +// +-----------------------------------+ | +// | TfLiteIntArray (variable size) |<-------/ +// +-----------------------------------+ +TfLiteDelegateParams* CreateDelegateParams(TfLiteDelegate* delegate, + const Subgraph& subgraph) { + // Step 1: Calculate the allocation size. + int allocation_size = sizeof(TfLiteDelegateParams); + + int nodes_to_replace_size = + TfLiteIntArrayGetSizeInBytes(subgraph.nodes.size()); + allocation_size += nodes_to_replace_size; + + int input_tensors_size = + TfLiteIntArrayGetSizeInBytes(subgraph.input_tensors.size()); + allocation_size += input_tensors_size; + + int output_tensors_size = + TfLiteIntArrayGetSizeInBytes(subgraph.output_tensors.size()); + allocation_size += output_tensors_size; + + // Step 2: Allocate the memory. + // Use `char*` for conveniently step through the allocated space by bytes. + char* allocation = reinterpret_cast(malloc(allocation_size)); + + // Step 3: Fill all data structures structures. + TfLiteDelegateParams* params = + reinterpret_cast(allocation); + params->delegate = delegate; + allocation += sizeof(TfLiteDelegateParams); + + params->nodes_to_replace = reinterpret_cast(allocation); + CopyVectorToTfLiteIntArray(subgraph.nodes, params->nodes_to_replace); + allocation += nodes_to_replace_size; + + params->input_tensors = reinterpret_cast(allocation); + CopyVectorToTfLiteIntArray(subgraph.input_tensors, params->input_tensors); + allocation += input_tensors_size; + + params->output_tensors = reinterpret_cast(allocation); + CopyVectorToTfLiteIntArray(subgraph.output_tensors, params->output_tensors); + allocation += output_tensors_size; + + return params; +} + +} // namespace + +TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( + TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace, + TfLiteDelegate* delegate) { + // Annotate the registration as DELEGATE op. + registration.builtin_code = BuiltinOperator_DELEGATE; + + // 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) { + // 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(); it != subgraph.nodes.end(); + ++it) { + execution_plan_.push_back(*it); + } + break; + case Subgraph::kTfPartition: { + int node_index; + + TfLiteDelegateParams* params = CreateDelegateParams(delegate, subgraph); + AddNodeWithParameters(subgraph.input_tensors, subgraph.output_tensors, + nullptr, 0, params, ®istration, &node_index); + + // Initialize the output tensors's delegate-related fields. + for (int tensor_index : subgraph.output_tensors) { + TfLiteTensor* tensor = &tensors_[tensor_index]; + TF_LITE_ENSURE_EQ(&context_, tensor->delegate, nullptr); + TF_LITE_ENSURE_EQ(&context_, tensor->buffer_handle, + kTfLiteNullBufferHandle); + // buffer_handle will be filled in delegate's `Prepare` + // function. + tensor->delegate = delegate; + } + + // Associate the node with the delegate. + TfLiteNode* node = &nodes_and_registration_[node_index].first; + node->delegate = delegate; + } 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."); + std::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())); @@ -151,16 +344,8 @@ TfLiteStatus Interpreter::BytesRequired(TfLiteType type, const int* dims, return kTfLiteOk; } -namespace { -TfLiteIntArray* convertVectorToTfLiteIntArray(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; -} -} // namespace - TfLiteStatus Interpreter::AllocateTensors() { - next_node_to_prepare_ = 0; + next_execution_plan_index_to_prepare_ = 0; if (memory_planner_) { TF_LITE_ENSURE_STATUS(memory_planner_->ResetAllocations()); } @@ -190,7 +375,8 @@ 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; @@ -201,8 +387,8 @@ TfLiteStatus Interpreter::AddNodeWithParameters( // NOTE, here we are not using move semantics yet, since our internal // representation isn't std::vector, but in the future we would like to avoid // copies, so we want the interface to take r-value references now. - node.inputs = convertVectorToTfLiteIntArray(inputs); - node.outputs = convertVectorToTfLiteIntArray(outputs); + node.inputs = ConvertVectorToTfLiteIntArray(inputs); + node.outputs = ConvertVectorToTfLiteIntArray(outputs); node.temporaries = TfLiteIntArrayCreate(0); if (init_data) { node.user_data = OpInit(*registration, init_data, init_data_size); @@ -211,8 +397,24 @@ TfLiteStatus Interpreter::AddNodeWithParameters( OpInit(*registration, reinterpret_cast(builtin_data_deleter.get()), 0); } + node.builtin_data = builtin_data_deleter.release(); + // TODO(ycling): Filling `custom_initial_data` and `custom_initial_data_size` + // properly for nodes generated by ReplaceSubgraphsWithDelegateKernels. + + if (registration->builtin_code == BuiltinOperator_CUSTOM) { + // When it's a CUSTOM op, the `custom_options` field in the Flatbuffer + // `Operator` table is passed in. + node.custom_initial_data = init_data; + node.custom_initial_data_size = init_data_size; + } else { + node.custom_initial_data = nullptr; + node.custom_initial_data_size = 0; + } + + node.delegate = nullptr; node_and_reg.second = *registration; + execution_plan_.push_back(new_node_index); return kTfLiteOk; } @@ -224,7 +426,7 @@ TfLiteStatus Interpreter::ResizeInputTensor(int tensor_index, TF_LITE_ENSURE(&context_, tensor_index < context_.tensors_size && tensor_index >= 0); invokable_ = false; - TfLiteIntArray* dims_lite = convertVectorToTfLiteIntArray(dims); + TfLiteIntArray* dims_lite = ConvertVectorToTfLiteIntArray(dims); return ResizeTensorImpl(&context_.tensors[tensor_index], dims_lite); } @@ -240,16 +442,20 @@ bool HasDynamicTensor(const TfLiteContext& context, return false; } -TfLiteStatus Interpreter::PrepareOpsStartingAt(int first_node, - int* last_node_prepared) { - for (int i = first_node; i < nodes_and_registration_.size(); i++) { - TfLiteNode& node = nodes_and_registration_[i].first; - const TfLiteRegistration& registration = nodes_and_registration_[i].second; +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; + EnsureTensorsVectorCapacity(); if (OpPrepare(registration, &node) == kTfLiteError) { return kTfLiteError; } - *last_node_prepared = i; + *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 @@ -268,14 +474,14 @@ TfLiteStatus Interpreter::PrepareOpsAndTensors() { memory_planner_->PlanAllocations(); } - int last_node_prepared = 0; + int last_exec_plan_index_prepared = 0; - TF_LITE_ENSURE_STATUS( - PrepareOpsStartingAt(next_node_to_prepare_, &last_node_prepared)); + TF_LITE_ENSURE_STATUS(PrepareOpsStartingAt( + next_execution_plan_index_to_prepare_, &last_exec_plan_index_prepared)); TF_LITE_ENSURE_STATUS(memory_planner_->ExecuteAllocations( - next_node_to_prepare_, last_node_prepared)); + next_execution_plan_index_to_prepare_, last_exec_plan_index_prepared)); - next_node_to_prepare_ = last_node_prepared + 1; + next_execution_plan_index_to_prepare_ = last_exec_plan_index_prepared + 1; return kTfLiteOk; } @@ -291,7 +497,7 @@ TfLiteStatus Interpreter::Invoke() { TfLiteStatus status = kTfLiteOk; if (nnapi_delegate_) { - if (next_node_to_prepare_ == 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 { @@ -311,17 +517,40 @@ TfLiteStatus Interpreter::Invoke() { // 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 i = 0; i < nodes_and_registration_.size(); i++) { - if (i == next_node_to_prepare_) { + 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_node_to_prepare_ >= i); + TF_LITE_ENSURE(&context_, next_execution_plan_index_to_prepare_ >= + 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; + + // TODO(ycling): This is an extra loop through inputs to check if the data + // need to be copied from Delegate buffer to raw memory, which is often not + // needed. We may want to cache this in prepare to know if this needs to be + // done for a node or not. + for (int i = 0; i < node.inputs->size; ++i) { + int tensor_index = node.inputs->data[i]; + if (tensor_index == kOptionalTensor) { + continue; + } + TfLiteTensor* tensor = &tensors_[tensor_index]; + if (tensor->delegate && tensor->delegate != node.delegate && + tensor->data_is_stale) { + EnsureTensorDataIsReadable(tensor_index); + } } - TfLiteNode& node = nodes_and_registration_[i].first; - const TfLiteRegistration& registration = nodes_and_registration_[i].second; + + EnsureTensorsVectorCapacity(); if (OpInvoke(registration, &node) == kTfLiteError) { status = kTfLiteError; } } + return status; } @@ -357,6 +586,7 @@ TfLiteStatus Interpreter::AddTensors(int tensors_to_add, tensors_.resize(tensors_.size() + tensors_to_add); for (int i = base_index; i < tensors_.size(); i++) { memset(&tensors_[i], 0, sizeof(tensors_[i])); + tensors_[i].buffer_handle = kTfLiteNullBufferHandle; } context_.tensors = tensors_.data(); context_.tensors_size = tensors_.size(); @@ -372,10 +602,26 @@ 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, - const char* buffer, size_t bytes, const Allocation* allocation) { + int tensor_index, TfLiteType type, const char* name, const int rank, + const int* dims, TfLiteQuantizationParams quantization, const char* buffer, + size_t bytes, const Allocation* allocation) { TF_LITE_ENSURE(&context_, tensor_index < context_.tensors_size && tensor_index >= 0); // For most tensors we know exactly how much memory is necessary so we can @@ -383,14 +629,27 @@ TfLiteStatus Interpreter::SetTensorParametersReadOnly( // because their sizes change with the contents of the individual strings. if (type != kTfLiteString) { size_t required_bytes; - TF_LITE_ENSURE_OK(&context_, BytesRequired(type, dims.data(), dims.size(), - &required_bytes)); + TF_LITE_ENSURE_OK(&context_, + BytesRequired(type, dims, rank, &required_bytes)); TF_LITE_ENSURE_EQ(&context_, required_bytes, bytes); } - invokable_ = false; - TfLiteTensorReset(type, name, convertVectorToTfLiteIntArray(dims), - quantization, const_cast(buffer), bytes, - kTfLiteMmapRo, allocation, &context_.tensors[tensor_index]); + + TfLiteTensor& tensor = context_.tensors[tensor_index]; + if (type == tensor.type && + EqualArrayAndTfLiteIntArray(tensor.dims, rank, dims)) { + // Fast path which does not invalidate the invokable property. + TfLiteTensorDataFree(&tensor); + tensor.data.raw = const_cast(buffer); + if (!tensor.dims) tensor.dims = ConvertArrayToTfLiteIntArray(rank, dims); + tensor.params = quantization; + tensor.allocation_type = kTfLiteMmapRo; + tensor.allocation = allocation; + } else { + invokable_ = false; + TfLiteTensorReset(type, name, ConvertArrayToTfLiteIntArray(rank, dims), + quantization, const_cast(buffer), bytes, + kTfLiteMmapRo, allocation, &tensor); + } return kTfLiteOk; } @@ -399,8 +658,8 @@ TfLiteStatus Interpreter::SetTensorParametersReadOnly( // bytes. The lifetime of buffer must be ensured to be greater or equal // to Interpreter. TfLiteStatus Interpreter::SetTensorParametersReadWrite( - int tensor_index, TfLiteType type, const char* name, - const std::vector& dims, TfLiteQuantizationParams quantization) { + int tensor_index, TfLiteType type, const char* name, const int rank, + const int* dims, TfLiteQuantizationParams quantization) { invokable_ = false; TF_LITE_ENSURE(&context_, tensor_index < context_.tensors_size && tensor_index >= 0); @@ -410,10 +669,10 @@ TfLiteStatus Interpreter::SetTensorParametersReadWrite( // many bytes we will need based on the dimensions. String tensors are // allocated dynamically and we can't know ahead of time how much space // they will require. - TF_LITE_ENSURE_OK(&context_, BytesRequired(type, dims.data(), dims.size(), - &required_bytes)); + TF_LITE_ENSURE_OK(&context_, + BytesRequired(type, dims, rank, &required_bytes)); } - TfLiteTensorReset(type, name, convertVectorToTfLiteIntArray(dims), + TfLiteTensorReset(type, name, ConvertArrayToTfLiteIntArray(rank, dims), quantization, /*buffer=*/nullptr, required_bytes, type == kTfLiteString ? kTfLiteDynamic : kTfLiteArenaRw, @@ -421,6 +680,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. @@ -434,6 +701,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); @@ -465,10 +735,54 @@ void Interpreter::UseNNAPI(bool enable) { } void Interpreter::SetNumThreads(int num_threads) { - // TODO(ahentz): this forces us to link against gemmlowp even when the ops - // don't use it. We should implement some dynamic mechanism for this sort of - // library-specific initialization. - tflite::gemm_support::SetMaxNumThreads(&context_, num_threads); + context_.recommended_num_threads = 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); + // Remove additional context info. + SetForbiddenContextFunction(&context_.GetNodeAndRegistration); + SetForbiddenContextFunction(&context_.ReplaceSubgraphsWithDelegateKernels); + SetForbiddenContextFunction(&context_.GetExecutionPlan); + return status; +} + +TfLiteStatus Interpreter::SetBufferHandle(int tensor_index, + TfLiteBufferHandle buffer_handle, + TfLiteDelegate* delegate) { + TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); + TfLiteTensor* tensor = &tensors_[tensor_index]; + + TF_LITE_ENSURE(&context_, + tensor->delegate == nullptr || tensor->delegate == delegate); + tensor->delegate = delegate; + if (tensor->buffer_handle != kTfLiteNullBufferHandle) { + TF_LITE_ENSURE(&context_, tensor->delegate->FreeBufferHandle != nullptr); + tensor->delegate->FreeBufferHandle(tensor->delegate, + &tensor->buffer_handle); + } + tensor->buffer_handle = buffer_handle; + + return kTfLiteOk; +} + +TfLiteStatus Interpreter::GetBufferHandle(int tensor_index, + TfLiteBufferHandle* buffer_handle, + TfLiteDelegate** delegate) { + TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); + TfLiteTensor* tensor = &tensors_[tensor_index]; + + *delegate = tensor->delegate; + *buffer_handle = tensor->buffer_handle; + + return kTfLiteOk; } } // namespace tflite diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h index 52e52df1b6aaaa511f7ac2d1edc33d3b21278996..3749869f58c4ce10bbedff6ff5c083087388c597 100644 --- a/tensorflow/contrib/lite/interpreter.h +++ b/tensorflow/contrib/lite/interpreter.h @@ -80,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 @@ -127,18 +133,34 @@ class Interpreter { // This variant assumes an external buffer has been allocated of size // bytes. The lifetime of buffer must be ensured to be greater or equal // to Interpreter. - TfLiteStatus SetTensorParametersReadOnly( + inline TfLiteStatus SetTensorParametersReadOnly( int tensor_index, TfLiteType type, const char* name, const std::vector& dims, TfLiteQuantizationParams quantization, + const char* buffer, size_t bytes, + const Allocation* allocation = nullptr) { + return SetTensorParametersReadOnly(tensor_index, type, name, dims.size(), + dims.data(), quantization, buffer, bytes, + allocation); + }; + + TfLiteStatus SetTensorParametersReadOnly( + int tensor_index, TfLiteType type, const char* name, const int rank, + const int* dims, TfLiteQuantizationParams quantization, const char* buffer, size_t bytes, const Allocation* allocation = nullptr); // Set description of inputs/outputs/data/fptrs for node `node_index`. // This variant assumes an external buffer has been allocated of size // bytes. The lifetime of buffer must be ensured to be greater or equal // to Interpreter. - TfLiteStatus SetTensorParametersReadWrite( + inline TfLiteStatus SetTensorParametersReadWrite( int tensor_index, TfLiteType type, const char* name, - const std::vector& dims, TfLiteQuantizationParams quantization); + const std::vector& dims, TfLiteQuantizationParams quantization) { + return SetTensorParametersReadWrite(tensor_index, type, name, dims.size(), + dims.data(), quantization); + } + TfLiteStatus SetTensorParametersReadWrite( + int tensor_index, TfLiteType type, const char* name, const int rank, + const int* dims, TfLiteQuantizationParams quantization); // Functions to access tensor data @@ -166,6 +188,13 @@ 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 @@ -240,6 +269,59 @@ 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); + + // Ensure the data in `tensor.data` is readable. In case delegate is used, + // it might require to copy the data from delegate buffer to raw memory. + TfLiteStatus EnsureTensorDataIsReadable(int tensor_index) { + TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); + TfLiteTensor* tensor = &tensors_[tensor_index]; + if (tensor->data_is_stale) { + TF_LITE_ENSURE(&context_, tensor->delegate != nullptr); + TF_LITE_ENSURE(&context_, + tensor->buffer_handle != kTfLiteNullBufferHandle); + // This can be null if the delegate doesn't use its own buffer. + TF_LITE_ENSURE(&context_, + tensor->delegate->CopyFromBufferHandle != nullptr); + tensor->delegate->CopyFromBufferHandle(tensor->delegate, + tensor->buffer_handle, + tensor->data.raw, tensor->bytes); + tensor->data_is_stale = false; + } + return kTfLiteOk; + } + + // Set the delegate buffer handle to a tensor. It can be called in the + // following cases: + // 1. Set the buffer handle to a tensor that's not being written by a + // delegate. For example, feeding an OpenGL texture as the input of the + // inference graph. + // 2. Set the buffer handle to a tensor that uses the same delegate. + // For example, set an OpenGL texture as the output of inference, while + // the node which produces output is an OpenGL delegate node. + // WARNING: This is an experimental API and subject to change. + TfLiteStatus SetBufferHandle(int tensor_index, + TfLiteBufferHandle buffer_handle, + TfLiteDelegate* delegate); + + // Get the delegate buffer handle, and the delegate which can process the + // buffer handle. + // WARNING: This is an experimental API and subject to change. + TfLiteStatus GetBufferHandle(int tensor_index, + TfLiteBufferHandle* buffer_handle, + TfLiteDelegate** delegate); + + // The default capacity of `tensors_` vector. + static constexpr int kTensorsReservedCapacity = 128; + // The capacity headroom of `tensors_` vector before calling ops' + // `prepare` and `invoke` function. In these functions, it's guaranteed + // allocating up to `kTensorsCapacityHeadroom` more tensors won't invalidate + // pointers to existing tensors. + static constexpr int kTensorsCapacityHeadroom = 16; + private: // Give 'op_reg' a chance to initialize itself using the contents of // 'buffer'. @@ -279,7 +361,8 @@ class Interpreter { // 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_node, int* last_node_prepared); + 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 @@ -299,7 +382,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). @@ -316,6 +400,53 @@ 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, TfLiteDelegate* delegate); + + // 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, + TfLiteDelegate* delegate); + + // 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); + + // Ensures that `tensors_` has at least `kTensorsCapacityHeadroom` extra + // capacity. Calling this function may invalidate existing pointers to + // tensors. After calling this function, adding `kTensorsCapacityHeadroom` + // more tensors won't invalidate the pointer to existing tensors. + void EnsureTensorsVectorCapacity() { + const int required_capacity = tensors_size() + kTensorsCapacityHeadroom; + if (required_capacity > tensors_.capacity()) { + tensors_.reserve(required_capacity); + context_.tensors = tensors_.data(); + } + } + // 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. @@ -350,11 +481,22 @@ class Interpreter { // 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. - // Intepreter will stop allocating tensors, set the value of next allocate + // Interpreter will stop allocating tensors, set the value of next allocate // 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_node_to_prepare_; + 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_; diff --git a/tensorflow/contrib/lite/interpreter_test.cc b/tensorflow/contrib/lite/interpreter_test.cc index edff2109430c6e1ec6c481619ed7772237a3301d..72d4acedbe72beea276c7c8fc1f7ddce01796d41 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 { @@ -38,7 +40,7 @@ TEST(BasicInterpreter, InvokeInvalidModel) { ASSERT_EQ(interpreter.Invoke(), kTfLiteOk); } -// Test size accesser functions. +// Test size accessor functions. TEST(BasicInterpreter, TestSizeFunctions) { Interpreter interpreter; int base_index; @@ -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,473 @@ TEST(BasicInterpreter, TestCustomErrorReporter) { ASSERT_EQ(reporter.calls, 1); } +TEST(BasicInterpreter, TestUnsupportedDelegateFunctions) { + Interpreter interpreter; + ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk); + TfLiteRegistration registration = { + .init = nullptr, .free = nullptr, .prepare = nullptr, .invoke = nullptr}; + // These functions are only supported inside Delegate's Prepare function. + // The test verifies that these functions returns `kTfLiteError`, but not + // `kTfLiteOk` or just crashes. + registration.prepare = [](TfLiteContext* context, TfLiteNode* node) { + { + TfLiteIntArray* execution_plan; + EXPECT_EQ(context->GetExecutionPlan(context, &execution_plan), + kTfLiteError); + } + { + TfLiteNode* node; + TfLiteRegistration* registration; + EXPECT_EQ( + context->GetNodeAndRegistration(context, 0, &node, ®istration), + kTfLiteError); + } + { + TfLiteRegistration delegate_registration = {nullptr, nullptr, nullptr, + nullptr}; + TfLiteIntArray nodes_to_replace; + nodes_to_replace.size = 0; + EXPECT_EQ(context->ReplaceSubgraphsWithDelegateKernels( + context, delegate_registration, &nodes_to_replace, nullptr), + kTfLiteError); + } + return kTfLiteError; + }; + ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk); + ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk); + ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, + ®istration), + kTfLiteOk); + EXPECT_EQ(interpreter.AllocateTensors(), kTfLiteError); +} + +TEST(InterpreterTensorsCapacityTest, TestWithinHeadroom) { + Interpreter interpreter; + ASSERT_EQ(interpreter.AddTensors(Interpreter::kTensorsReservedCapacity), + kTfLiteOk); + TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr}; + registration.prepare = [](TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* first_tensor = context->tensors; + + int new_tensor_index; + context->AddTensors(context, Interpreter::kTensorsCapacityHeadroom, + &new_tensor_index); + EXPECT_EQ(first_tensor, context->tensors); + return kTfLiteOk; + }; + ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, + ®istration), + kTfLiteOk); + ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk); +} + +TEST(InterpreterTensorsCapacityTest, TestExceedHeadroom) { + Interpreter interpreter; + ASSERT_EQ(interpreter.AddTensors(Interpreter::kTensorsReservedCapacity), + kTfLiteOk); + TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr}; + registration.prepare = [](TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* first_tensor = context->tensors; + + int new_tensor_index; + context->AddTensors(context, Interpreter::kTensorsCapacityHeadroom + 1, + &new_tensor_index); + EXPECT_NE(first_tensor, context->tensors); + return kTfLiteOk; + }; + ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, + ®istration), + kTfLiteOk); + ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk); +} + +// 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 { + protected: + void SetUp() override { + interpreter_.reset(new Interpreter); + 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, ®); + } + + void TearDown() override { + // Interpreter relies on delegate_ to free the resources properly. Thus + // the life cycle of delegate must be longer than interpreter. + interpreter_.reset(); + delegate_.reset(); + } + + TfLiteBufferHandle last_allocated_handle_ = kTfLiteNullBufferHandle; + + TfLiteBufferHandle AllocateBufferHandle() { return ++last_allocated_handle_; } + + 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, + TfLiteDelegate* delegate) -> TfLiteStatus { + auto* simple = reinterpret_cast(delegate->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, delegate); + TfLiteIntArrayFree(nodes_to_separate); + return kTfLiteOk; + }; + delegate_.CopyToBufferHandle = [](TfLiteDelegate* delegate, + TfLiteBufferHandle buffer_handle, + void* data, int size) -> TfLiteStatus { + // TODO(ycling): Implement tests to test buffer copying logic. + return kTfLiteOk; + }; + delegate_.CopyFromBufferHandle = + [](TfLiteDelegate* delegate, TfLiteBufferHandle buffer_handle, + void* data, int size) -> TfLiteStatus { + // TODO(ycling): Implement tests to test buffer copying logic. + return kTfLiteOk; + }; + delegate_.FreeBufferHandle = [](TfLiteDelegate* delegate, + TfLiteBufferHandle* handle) { + *handle = kTfLiteNullBufferHandle; + }; + // 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_; + }; + std::unique_ptr interpreter_; + std::unique_ptr delegate_; +}; + +TEST_F(TestDelegate, BasicDelegate) { + interpreter_->Invoke(); + delegate_ = std::unique_ptr(new SimpleDelegate({0, 1, 2})); + interpreter_->ModifyGraphWithDelegate(delegate_->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); + EXPECT_EQ(node_and_reg->second.custom_name, + SimpleDelegate::FakeFusedRegistration().custom_name); + + const TfLiteDelegateParams* params = + reinterpret_cast( + node_and_reg->first.builtin_data); + ASSERT_EQ(params->nodes_to_replace->size, 3); + EXPECT_EQ(params->nodes_to_replace->data[0], 0); + EXPECT_EQ(params->nodes_to_replace->data[1], 1); + EXPECT_EQ(params->nodes_to_replace->data[2], 2); + + ASSERT_EQ(params->input_tensors->size, 2); + EXPECT_EQ(params->input_tensors->data[0], 0); + EXPECT_EQ(params->input_tensors->data[1], 1); + + ASSERT_EQ(params->output_tensors->size, 2); + EXPECT_EQ(params->output_tensors->data[0], 3); + EXPECT_EQ(params->output_tensors->data[1], 4); +} + +TEST_F(TestDelegate, ComplexDeligate) { + interpreter_->Invoke(); + delegate_ = std::unique_ptr(new SimpleDelegate({1, 2})); + interpreter_->ModifyGraphWithDelegate(delegate_->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); +} + +TEST_F(TestDelegate, SetBufferHandleToInput) { + interpreter_->Invoke(); + delegate_ = std::unique_ptr(new SimpleDelegate({0, 1, 2})); + TfLiteDelegate* delegate = delegate_->get_tf_lite_delegate(); + interpreter_->ModifyGraphWithDelegate(delegate); + + constexpr int kOutputTensorIndex = 0; + TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex); + ASSERT_EQ(tensor->delegate, nullptr); + ASSERT_EQ(tensor->buffer_handle, kTfLiteNullBufferHandle); + + TfLiteBufferHandle handle = AllocateBufferHandle(); + TfLiteStatus status = + interpreter_->SetBufferHandle(kOutputTensorIndex, handle, delegate); + ASSERT_EQ(status, kTfLiteOk); + EXPECT_EQ(tensor->delegate, delegate); + EXPECT_EQ(tensor->buffer_handle, handle); +} + +TEST_F(TestDelegate, SetBufferHandleToOutput) { + interpreter_->Invoke(); + delegate_ = std::unique_ptr(new SimpleDelegate({0, 1, 2})); + TfLiteDelegate* delegate = delegate_->get_tf_lite_delegate(); + interpreter_->ModifyGraphWithDelegate(delegate); + + constexpr int kOutputTensorIndex = 3; + TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex); + // Before setting the buffer handle, the tensor's `delegate` is already set + // because it will be written by the delegate. + ASSERT_EQ(tensor->delegate, delegate); + ASSERT_EQ(tensor->buffer_handle, kTfLiteNullBufferHandle); + + TfLiteBufferHandle handle = AllocateBufferHandle(); + TfLiteStatus status = + interpreter_->SetBufferHandle(kOutputTensorIndex, handle, delegate); + ASSERT_EQ(status, kTfLiteOk); + EXPECT_EQ(tensor->delegate, delegate); + EXPECT_EQ(tensor->buffer_handle, handle); +} + +TEST_F(TestDelegate, SetInvalidHandleToTensor) { + interpreter_->Invoke(); + delegate_ = std::unique_ptr(new SimpleDelegate({0, 1, 2})); + TfLiteDelegate* delegate = delegate_->get_tf_lite_delegate(); + interpreter_->ModifyGraphWithDelegate(delegate); + + SimpleDelegate another_simple_delegate({0, 1, 2}); + + constexpr int kOutputTensorIndex = 3; + TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex); + // Before setting the buffer handle, the tensor's `delegate` is already set + // because it will be written by the delegate. + ASSERT_EQ(tensor->delegate, delegate); + ASSERT_EQ(tensor->buffer_handle, kTfLiteNullBufferHandle); + + TfLiteBufferHandle handle = AllocateBufferHandle(); + TfLiteStatus status = interpreter_->SetBufferHandle( + kOutputTensorIndex, handle, + another_simple_delegate.get_tf_lite_delegate()); + // Setting a buffer handle to a tensor with another delegate will fail. + ASSERT_EQ(status, kTfLiteError); + EXPECT_EQ(tensor->delegate, delegate); + EXPECT_EQ(tensor->buffer_handle, kTfLiteNullBufferHandle); +} + } // 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 26cfe6c3e286ed603c2183986c697562e846889c..079320586ffd01fc77818a81e0c5962f1d28c1f1 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,9 +31,6 @@ 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 \ @@ -42,6 +40,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/BUILD b/tensorflow/contrib/lite/java/BUILD index 9a1a888b93ff981b1d14faa7e847e80be1f167f2..f52d6ba6c5390e631d29e75f833aa4dd5bba1a68 100644 --- a/tensorflow/contrib/lite/java/BUILD +++ b/tensorflow/contrib/lite/java/BUILD @@ -29,7 +29,7 @@ android_library( visibility = ["//visibility:public"], deps = [ ":tflite_runtime", - "@javax_validation", + "@org_checkerframework_qual", ], ) @@ -42,7 +42,7 @@ android_library( ), visibility = ["//visibility:public"], deps = [ - "@javax_validation", + "@org_checkerframework_qual", ], ) @@ -58,7 +58,7 @@ java_library( deps = [ ":libtensorflowlite_jni.so", "//tensorflow/contrib/lite/java/src/main/native", - "@javax_validation", + "@org_checkerframework_qual", ], ) @@ -111,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/demo/app/src/main/BUILD b/tensorflow/contrib/lite/java/demo/app/src/main/BUILD index 654fa9d6d2799fc3cafa3e0e042cb2a5746bf2c5..5eb749aae6e224bec64b66832f116ebc3372c1ef 100644 --- a/tensorflow/contrib/lite/java/demo/app/src/main/BUILD +++ b/tensorflow/contrib/lite/java/demo/app/src/main/BUILD @@ -6,7 +6,7 @@ android_binary( name = "TfLiteCameraDemo", srcs = glob(["java/**/*.java"]), assets = [ - "@tflite_mobilenet//:labels.txt", + "//tensorflow/contrib/lite/java/demo/app/src/main/assets:labels_mobilenet_quant_v1_224.txt", "@tflite_mobilenet//:mobilenet_quant_v1_224.tflite", ], assets_dir = "", 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..300786c3ca01b12a46f7f9a6fe8fd720f97a79f4 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,9 +296,10 @@ 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."); + Log.e(TAG, "Failed to initialize an image classifier.", e); } startBackgroundThread(); } @@ -432,7 +433,7 @@ public class Camera2BasicFragment extends Fragment return; } } catch (CameraAccessException e) { - e.printStackTrace(); + Log.e(TAG, "Failed to access Camera", e); } catch (NullPointerException e) { // Currently an NPE is thrown when the Camera2API is used but not supported on the // device this code runs. @@ -477,7 +478,7 @@ public class Camera2BasicFragment extends Fragment } manager.openCamera(cameraId, stateCallback, backgroundHandler); } catch (CameraAccessException e) { - e.printStackTrace(); + Log.e(TAG, "Failed to open Camera", e); } catch (InterruptedException e) { throw new RuntimeException("Interrupted while trying to lock camera opening.", e); } @@ -544,7 +545,7 @@ public class Camera2BasicFragment extends Fragment runClassifier = false; } } catch (InterruptedException e) { - e.printStackTrace(); + Log.e(TAG, "Interrupted when stopping background thread", e); } } @@ -603,7 +604,7 @@ public class Camera2BasicFragment extends Fragment captureSession.setRepeatingRequest( previewRequest, captureCallback, backgroundHandler); } catch (CameraAccessException e) { - e.printStackTrace(); + Log.e(TAG, "Failed to set up config to capture Camera", e); } } @@ -614,7 +615,7 @@ public class Camera2BasicFragment extends Fragment }, null); } catch (CameraAccessException e) { - e.printStackTrace(); + Log.e(TAG, "Failed to preview Camera", e); } } @@ -658,8 +659,7 @@ public class Camera2BasicFragment extends Fragment showToast("Uninitialized Classifier or invalid context."); return; } - Bitmap bitmap = - textureView.getBitmap(ImageClassifier.DIM_IMG_SIZE_X, ImageClassifier.DIM_IMG_SIZE_Y); + Bitmap bitmap = 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 e44c5ae6b48eda187079dd3a0a1bc563276d816e..c57bb348c5b386a59327c7b1bc769717ca755269 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,18 @@ 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; - /** An array to hold inference results, to be feed into Tensorflow Lite as outputs. */ - private byte[][] labelProbArray = null; /** multi-stage low pass filter * */ private float[][] filterLabelProbArray = null; @@ -95,10 +88,13 @@ 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][labelList.size()]; + filterLabelProbArray = new float[FILTER_STAGES][getNumLabels()]; Log.d(TAG, "Created a Tensorflow Lite Image Classifier."); } @@ -111,7 +107,7 @@ 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)); @@ -125,12 +121,12 @@ public class ImageClassifier { } void applyFilter() { - int numLabels = labelList.size(); + 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 * (labelProbArray[0][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) { @@ -142,7 +138,7 @@ public class ImageClassifier { // Copy the last stage filter output back to `labelProbArray`. for (int j = 0; j < numLabels; ++j) { - labelProbArray[0][j] = (byte)filterLabelProbArray[FILTER_STAGES - 1][j]; + setProbability(j, filterLabelProbArray[FILTER_STAGES - 1][j]); } } @@ -156,7 +152,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); @@ -167,7 +163,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(); @@ -185,12 +181,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(); @@ -199,9 +193,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(); } @@ -214,4 +208,89 @@ public class ImageClassifier { } 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..e164ac75543ebab12e6b1c057c4ed487eb9accdf --- /dev/null +++ b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java @@ -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. +==============================================================================*/ + +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/proguard.flags b/tensorflow/contrib/lite/java/proguard.flags new file mode 100644 index 0000000000000000000000000000000000000000..8ee3d7e7ae728b27789336ac56208acdf13ee424 --- /dev/null +++ b/tensorflow/contrib/lite/java/proguard.flags @@ -0,0 +1,3 @@ +-keepclassmembers class org.tensorflow.lite.NativeInterpreterWrapper { + private long inferenceDurationNanoseconds; +} \ No newline at end of file diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/DataType.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/DataType.java index d63c299589d2e8ce1051a52d29b533ed126bbcf7..fc16488a6459eb227fde712055d3e8ccfcce0070 100644 --- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/DataType.java +++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/DataType.java @@ -71,6 +71,23 @@ enum DataType { throw new IllegalArgumentException("DataType " + this + " is not supported yet"); } + /** Gets string names of the data type. */ + String toStringName() { + switch (this) { + case FLOAT32: + return "float"; + case INT32: + return "int"; + case UINT8: + return "byte"; + case INT64: + return "long"; + case BYTEBUFFER: + return "ByteBuffer"; + } + throw new IllegalArgumentException("DataType " + this + " is not supported yet"); + } + // Cached to avoid copying it private static final DataType[] values = values(); } diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java index dd883d69d2065236ee29012b9bde99972aefbcf7..cc17b491f2a1a42d704c0262095ce817f3828764 100644 --- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java +++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java @@ -19,7 +19,7 @@ import java.io.File; import java.nio.MappedByteBuffer; import java.util.HashMap; import java.util.Map; -import javax.validation.constraints.NotNull; +import org.checkerframework.checker.nullness.qual.NonNull; /** * Driver class to drive model inference with TensorFlow Lite. @@ -60,7 +60,7 @@ public final class Interpreter implements AutoCloseable { * * @param modelFile: a File of a pre-trained TF Lite model. */ - public Interpreter(@NotNull File modelFile) { + public Interpreter(@NonNull File modelFile) { if (modelFile == null) { return; } @@ -73,20 +73,23 @@ public final class Interpreter implements AutoCloseable { *

The {@code MappedByteBuffer} should remain unchanged after the construction of a {@code * Interpreter}. */ - public Interpreter(@NotNull MappedByteBuffer mappedByteBuffer) { + public Interpreter(@NonNull MappedByteBuffer mappedByteBuffer) { wrapper = new NativeInterpreterWrapper(mappedByteBuffer); } /** * Runs model inference if the model takes only one input, and provides only one output. * + *

Warning: The API runs much faster if {@link ByteBuffer} is used as input data type. Please + * consider using {@link ByteBuffer} to feed input data for better performance. + * * @param input an array or multidimensional array, or a {@link ByteBuffer} of primitive types * including int, float, long, and byte. {@link ByteBuffer} is the preferred way to pass large * input data. When {@link ByteBuffer} is used, its content should remain unchanged until * model inference is done. * @param output a multidimensional array of output data. */ - public void run(@NotNull Object input, @NotNull Object output) { + public void run(@NonNull Object input, @NonNull Object output) { Object[] inputs = {input}; Map outputs = new HashMap<>(); outputs.put(0, output); @@ -96,6 +99,9 @@ public final class Interpreter implements AutoCloseable { /** * Runs model inference if the model takes multiple inputs, or returns multiple outputs. * + *

Warning: The API runs much faster if {@link ByteBuffer} is used as input data type. Please + * consider using {@link ByteBuffer} to feed input data for better performance. + * * @param inputs an array of input data. The inputs should be in the same order as inputs of the * model. Each input can be an array or multidimensional array, or a {@link ByteBuffer} of * primitive types including int, float, long, and byte. {@link ByteBuffer} is the preferred @@ -105,7 +111,7 @@ public final class Interpreter implements AutoCloseable { * needs to keep entries for the outputs to be used. */ public void runForMultipleInputsOutputs( - @NotNull Object[] inputs, @NotNull Map outputs) { + @NonNull Object[] inputs, @NonNull Map outputs) { if (wrapper == null) { throw new IllegalStateException("The Interpreter has already been closed."); } @@ -128,7 +134,7 @@ public final class Interpreter implements AutoCloseable { * *

IllegalArgumentException will be thrown if it fails to resize. */ - public void resizeInput(int idx, @NotNull int[] dims) { + public void resizeInput(int idx, @NonNull int[] dims) { if (wrapper == null) { throw new IllegalStateException("The Interpreter has already been closed."); } @@ -161,6 +167,27 @@ public final class Interpreter implements AutoCloseable { return wrapper.getOutputIndex(opName); } + /** + * Returns native inference timing. + *

IllegalArgumentException will be thrown if the model is not initialized by the + * {@link Interpreter}. + */ + public Long getLastNativeInferenceDurationNanoseconds() { + if (wrapper == null) { + throw new IllegalStateException("The interpreter has already been closed."); + } + return wrapper.getLastNativeInferenceDurationNanoseconds(); + } + + /** Turns on/off Android NNAPI for hardware acceleration when it is available. */ + public void setUseNNAPI(boolean useNNAPI) { + if (wrapper != null) { + wrapper.setUseNNAPI(useNNAPI); + } else { + throw new IllegalStateException("NativeInterpreterWrapper has already been closed."); + } + } + /** Release resources associated with the {@code Interpreter}. */ @Override public void close() { 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 5ee594dec492ad2fee22e603a6de311b3fed4cac..518e8b3a96389f327382481e82fb7b18a6ceb5cd 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 @@ -35,6 +35,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { errorHandle = createErrorReporter(ERROR_BUFFER_SIZE); modelHandle = createModel(modelPath, errorHandle); interpreterHandle = createInterpreter(modelHandle, errorHandle); + isMemoryAllocated = true; } /** @@ -47,6 +48,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { errorHandle = createErrorReporter(ERROR_BUFFER_SIZE); modelHandle = createModelWithBuffer(modelByteBuffer, errorHandle); interpreterHandle = createInterpreter(modelHandle, errorHandle); + isMemoryAllocated = true; } /** Releases resources associated with this {@code NativeInterpreterWrapper}. */ @@ -59,6 +61,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { modelByteBuffer = null; inputsIndexes = null; outputsIndexes = null; + isMemoryAllocated = false; } /** Sets inputs, runs model inference and returns outputs. */ @@ -91,11 +94,21 @@ final class NativeInterpreterWrapper implements AutoCloseable { i, inputs.length)); } } + inferenceDurationNanoseconds = -1; long[] outputsHandles = - run(interpreterHandle, errorHandle, sizes, dataTypes, numsOfBytes, inputs); + run( + interpreterHandle, + errorHandle, + sizes, + dataTypes, + numsOfBytes, + inputs, + this, + isMemoryAllocated); if (outputsHandles == null || outputsHandles.length == 0) { throw new IllegalStateException("Interpreter has no outputs."); } + isMemoryAllocated = true; Tensor[] outputs = new Tensor[outputsHandles.length]; for (int i = 0; i < outputsHandles.length; ++i) { outputs[i] = Tensor.fromHandle(outputsHandles[i]); @@ -109,20 +122,28 @@ final class NativeInterpreterWrapper implements AutoCloseable { Object[] sizes, int[] dtypes, int[] numsOfBytes, - Object[] values); + Object[] values, + NativeInterpreterWrapper wrapper, + boolean memoryAllocated); /** Resizes dimensions of a specific input. */ void resizeInput(int idx, int[] dims) { - resizeInput(interpreterHandle, errorHandle, idx, dims); + if (resizeInput(interpreterHandle, errorHandle, idx, dims)) { + isMemoryAllocated = false; + } } - private static native void resizeInput( + private static native boolean resizeInput( long interpreterHandle, long errorHandle, int inputIdx, int[] dims); void setUseNNAPI(boolean useNNAPI) { useNNAPI(interpreterHandle, useNNAPI); } + void setNumThreads(int numRecommendedThreads) { + numThreads(interpreterHandle, numRecommendedThreads); + } + /** Gets index of an input given its name. */ int getInputIndex(String name) { if (inputsIndexes == null) { @@ -236,6 +257,35 @@ final class NativeInterpreterWrapper implements AutoCloseable { } } + /** + * Gets the last inference duration in nanoseconds. It returns null if there is no previous + * inference run or the last inference run failed. + */ + Long getLastNativeInferenceDurationNanoseconds() { + return (inferenceDurationNanoseconds < 0) ? null : inferenceDurationNanoseconds; + } + + /** + * Gets the dimensions of an input. It throws IllegalArgumentException if input index is invalid. + */ + int[] getInputDims(int index) { + return getInputDims(interpreterHandle, index, -1); + } + + /** + * Gets the dimensions of an input. If numBytes >= 0, it will check whether num of bytes match the + * input. + */ + private static native int[] getInputDims(long interpreterHandle, int inputIdx, int numBytes); + + /** Gets the type of an output. It throws IllegalArgumentException if output index is invalid. */ + String getOutputDataType(int index) { + int type = getOutputDataType(interpreterHandle, index); + return DataType.fromNumber(type).toStringName(); + } + + private static native int getOutputDataType(long interpreterHandle, int outputIdx); + private static final int ERROR_BUFFER_SIZE = 512; private long errorHandle; @@ -246,18 +296,24 @@ final class NativeInterpreterWrapper implements AutoCloseable { private int inputSize; + private long inferenceDurationNanoseconds = -1; + private MappedByteBuffer modelByteBuffer; private Map inputsIndexes; private Map outputsIndexes; + private boolean isMemoryAllocated = false; + private static native String[] getInputNames(long interpreterHandle); private static native String[] getOutputNames(long interpreterHandle); private static native void useNNAPI(long interpreterHandle, boolean state); + private static native void numThreads(long interpreterHandle, int numRecommendedThreads); + private static native long createErrorReporter(int size); private static native long createModel(String modelPathOrBuffer, long errorHandle); @@ -268,8 +324,6 @@ final class NativeInterpreterWrapper implements AutoCloseable { private static native void delete(long errorHandle, long modelHandle, long interpreterHandle); - private static native int[] getInputDims(long interpreterHandle, int inputIdx, int numBytes); - static { TensorFlowLite.init(); } diff --git a/tensorflow/contrib/lite/java/src/main/native/BUILD b/tensorflow/contrib/lite/java/src/main/native/BUILD index 15806d57c8ed7a45d2db9b80e2aab8e22349ee3e..3571182ca92e959d54935cfdc76679ab69a8cfa9 100644 --- a/tensorflow/contrib/lite/java/src/main/native/BUILD +++ b/tensorflow/contrib/lite/java/src/main/native/BUILD @@ -11,6 +11,7 @@ licenses(["notice"]) # Apache 2.0 cc_library( name = "native_framework_only", srcs = [ + "duration_utils_jni.cc", "exception_jni.cc", "nativeinterpreterwrapper_jni.cc", "tensor_jni.cc", diff --git a/tensorflow/contrib/lite/java/src/main/native/duration_utils_jni.cc b/tensorflow/contrib/lite/java/src/main/native/duration_utils_jni.cc new file mode 100644 index 0000000000000000000000000000000000000000..0e08a04370592f6e3c92b5811fa7e163f808e03c --- /dev/null +++ b/tensorflow/contrib/lite/java/src/main/native/duration_utils_jni.cc @@ -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. +==============================================================================*/ + +#include +#include + +namespace tflite { + +// Gets the elapsed wall-clock timespec. +timespec getCurrentTime() { + timespec time; + clock_gettime(CLOCK_MONOTONIC, &time); + return time; +} + +// Computes the time diff from two timespecs. Returns '-1' if 'stop' is earlier +// than 'start'. +jlong timespec_diff_nanoseconds(struct timespec* start, struct timespec* stop) { + jlong result = stop->tv_sec - start->tv_sec; + if (result < 0) return -1; + result = 1000000000 * result + (stop->tv_nsec - start->tv_nsec); + if (result < 0) return -1; + return result; +} + +} // namespace tflite 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 f3f51b668f068ffcd02862a79b72dbae31d31c02..cc448b03c3d85c5ea2d986ce9e2e0e60ee2e7a6f 100644 --- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc +++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc @@ -79,6 +79,21 @@ TfLiteType resolveDataType(jint data_type) { } } +int getDataType(TfLiteType data_type) { + switch (data_type) { + case kTfLiteFloat32: + return 1; + case kTfLiteInt32: + return 2; + case kTfLiteUInt8: + return 3; + case kTfLiteInt64: + return 4; + default: + return -1; + } +} + void printDims(char* buffer, int max_size, int* dims, int num_dims) { if (max_size <= 0) return; buffer[0] = '?'; @@ -149,6 +164,45 @@ TfLiteStatus checkInputs(JNIEnv* env, tflite::Interpreter* interpreter, return kTfLiteOk; } +// Checks whether there is any difference between dimensions of a tensor and a +// given dimensions. Returns true if there is difference, else false. +bool areDimsDifferent(JNIEnv* env, TfLiteTensor* tensor, jintArray dims) { + int num_dims = static_cast(env->GetArrayLength(dims)); + jint* ptr = env->GetIntArrayElements(dims, nullptr); + if (ptr == nullptr) { + throwException(env, kIllegalArgumentException, + "Empty dimensions of input array."); + return true; + } + if (tensor->dims->size != num_dims) { + return true; + } + for (int i = 0; i < num_dims; ++i) { + if (ptr[i] != tensor->dims->data[i]) { + return true; + } + } + env->ReleaseIntArrayElements(dims, ptr, JNI_ABORT); + return false; +} + +bool areInputDimensionsTheSame(JNIEnv* env, tflite::Interpreter* interpreter, + int input_size, jobjectArray sizes) { + if (interpreter->inputs().size() != input_size) { + return false; + } + for (int i = 0; i < input_size; ++i) { + int input_idx = interpreter->inputs()[i]; + jintArray dims = + static_cast(env->GetObjectArrayElement(sizes, i)); + TfLiteTensor* target = interpreter->tensor(input_idx); + if (areDimsDifferent(env, target, dims)) return false; + env->DeleteLocalRef(dims); + if (env->ExceptionCheck()) return false; + } + return true; +} + TfLiteStatus resizeInputs(JNIEnv* env, tflite::Interpreter* interpreter, int input_size, jobjectArray sizes) { for (int i = 0; i < input_size; ++i) { @@ -200,6 +254,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 @@ -256,6 +316,16 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_useNNAPI(JNIEnv* env, interpreter->UseNNAPI(static_cast(state)); } +JNIEXPORT void JNICALL +Java_org_tensorflow_lite_NativeInterpreterWrapper_numThreads(JNIEnv* env, + jclass clazz, + jlong handle, + jint num_threads) { + tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle); + if (interpreter == nullptr) return; + interpreter->SetNumThreads(static_cast(num_threads)); +} + JNIEXPORT jlong JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_createErrorReporter( JNIEnv* env, jclass clazz, jint size) { @@ -264,6 +334,19 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createErrorReporter( return reinterpret_cast(error_reporter); } +// Verifies whether the model is a flatbuffer file. +class JNIFlatBufferVerifier : public tflite::TfLiteVerifier { + public: + bool Verify(const char* data, int length, + tflite::ErrorReporter* reporter) override { + if (!VerifyModel(data, length)) { + reporter->Report("The model is not a valid Flatbuffer file"); + return false; + } + return true; + } +}; + JNIEXPORT jlong JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_createModel( JNIEnv* env, jclass clazz, jstring model_file, jlong error_handle) { @@ -271,7 +354,12 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModel( convertLongToErrorReporter(env, error_handle); if (error_reporter == nullptr) return 0; const char* path = env->GetStringUTFChars(model_file, nullptr); - auto model = tflite::FlatBufferModel::BuildFromFile(path, error_reporter); + + std::unique_ptr verifier; + verifier.reset(new JNIFlatBufferVerifier()); + + auto model = tflite::FlatBufferModel::VerifyAndBuildFromFile( + path, verifier.get(), error_reporter); if (!model) { throwException(env, kIllegalArgumentException, "Contents of %s does not encode a valid TensorFlowLite " @@ -293,6 +381,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) { @@ -321,6 +415,15 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter( throwException(env, kIllegalArgumentException, "Cannot create interpreter: %s", error_reporter->CachedErrorMessage()); + return 0; + } + // allocates memory + status = interpreter->AllocateTensors(); + if (status != kTfLiteOk) { + throwException(env, kNullPointerException, + "Can not allocate memory for the interpreter", + error_reporter->CachedErrorMessage()); + return 0; } return reinterpret_cast(interpreter.release()); } @@ -330,7 +433,7 @@ JNIEXPORT jlongArray JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_run( JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle, jobjectArray sizes, jintArray data_types, jintArray nums_of_bytes, - jobjectArray values) { + jobjectArray values, jobject wrapper, jboolean memory_allocated) { tflite::Interpreter* interpreter = convertLongToInterpreter(env, interpreter_handle); if (interpreter == nullptr) return nullptr; @@ -342,25 +445,29 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_run( TfLiteStatus status = checkInputs(env, interpreter, input_size, data_types, nums_of_bytes, values, sizes); if (status != kTfLiteOk) return nullptr; - // resizes inputs - status = resizeInputs(env, interpreter, input_size, sizes); - if (status != kTfLiteOk) { - throwException(env, kNullPointerException, "Can not resize the input: %s", - error_reporter->CachedErrorMessage()); - return nullptr; - } - // allocates memory - status = interpreter->AllocateTensors(); - if (status != kTfLiteOk) { - throwException(env, kNullPointerException, - "Can not allocate memory for the given inputs: %s", - error_reporter->CachedErrorMessage()); - return nullptr; + if (!memory_allocated || + !areInputDimensionsTheSame(env, interpreter, input_size, sizes)) { + // resizes inputs + status = resizeInputs(env, interpreter, input_size, sizes); + if (status != kTfLiteOk) { + throwException(env, kNullPointerException, "Can not resize the input: %s", + error_reporter->CachedErrorMessage()); + return nullptr; + } + // allocates memory + status = interpreter->AllocateTensors(); + if (status != kTfLiteOk) { + throwException(env, kNullPointerException, + "Can not allocate memory for the given inputs: %s", + error_reporter->CachedErrorMessage()); + return nullptr; + } } // sets inputs status = setInputs(env, interpreter, input_size, data_types, nums_of_bytes, values); if (status != kTfLiteOk) return nullptr; + timespec beforeInference = ::tflite::getCurrentTime(); // runs inference if (interpreter->Invoke() != kTfLiteOk) { throwException(env, kIllegalArgumentException, @@ -368,6 +475,17 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_run( error_reporter->CachedErrorMessage()); return nullptr; } + timespec afterInference = ::tflite::getCurrentTime(); + jclass wrapper_clazz = env->GetObjectClass(wrapper); + jfieldID fid = + env->GetFieldID(wrapper_clazz, "inferenceDurationNanoseconds", "J"); + if (env->ExceptionCheck()) { + env->ExceptionClear(); + } else if (fid != nullptr) { + env->SetLongField( + wrapper, fid, + ::tflite::timespec_diff_nanoseconds(&beforeInference, &afterInference)); + } // returns outputs const std::vector& results = interpreter->outputs(); if (results.empty()) { @@ -391,7 +509,7 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims( tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle); if (interpreter == nullptr) return nullptr; const int idx = static_cast(input_idx); - if (input_idx >= interpreter->inputs().size()) { + if (input_idx < 0 || input_idx >= interpreter->inputs().size()) { throwException(env, kIllegalArgumentException, "Out of range: Failed to get %d-th input out of %d inputs", input_idx, interpreter->inputs().size()); @@ -399,45 +517,72 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims( } TfLiteTensor* target = interpreter->tensor(interpreter->inputs()[idx]); int size = target->dims->size; - int expected_num_bytes = elementByteSize(target->type); - for (int i = 0; i < size; ++i) { - expected_num_bytes *= target->dims->data[i]; - } - if (num_bytes != expected_num_bytes) { - throwException(env, kIllegalArgumentException, - "Failed to get input dimensions. %d-th input should have" - " %d bytes, but found %d bytes.", - idx, expected_num_bytes, num_bytes); - return nullptr; + if (num_bytes >= 0) { // verifies num of bytes matches if num_bytes if valid. + int expected_num_bytes = elementByteSize(target->type); + for (int i = 0; i < size; ++i) { + expected_num_bytes *= target->dims->data[i]; + } + if (num_bytes != expected_num_bytes) { + throwException(env, kIllegalArgumentException, + "Failed to get input dimensions. %d-th input should have" + " %d bytes, but found %d bytes.", + idx, expected_num_bytes, num_bytes); + return nullptr; + } } jintArray outputs = env->NewIntArray(size); env->SetIntArrayRegion(outputs, 0, size, &(target->dims->data[0])); return outputs; } -JNIEXPORT void JNICALL +JNIEXPORT jint JNICALL +Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputDataType( + JNIEnv* env, jclass clazz, jlong handle, jint output_idx) { + tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle); + if (interpreter == nullptr) return -1; + const int idx = static_cast(output_idx); + if (output_idx < 0 || output_idx >= interpreter->outputs().size()) { + throwException(env, kIllegalArgumentException, + "Out of range: Failed to get %d-th output out of %d outputs", + output_idx, interpreter->outputs().size()); + return -1; + } + TfLiteTensor* target = interpreter->tensor(interpreter->outputs()[idx]); + int type = getDataType(target->type); + return static_cast(type); +} + +JNIEXPORT jboolean JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_resizeInput( JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle, jint input_idx, jintArray dims) { BufferErrorReporter* error_reporter = convertLongToErrorReporter(env, error_handle); - if (error_reporter == nullptr) return; + if (error_reporter == nullptr) return JNI_FALSE; tflite::Interpreter* interpreter = convertLongToInterpreter(env, interpreter_handle); - if (interpreter == nullptr) return; + if (interpreter == nullptr) return JNI_FALSE; const int idx = static_cast(input_idx); if (idx < 0 || idx >= interpreter->inputs().size()) { throwException(env, kIllegalArgumentException, "Can not resize %d-th input for a model having %d inputs.", idx, interpreter->inputs().size()); + return JNI_FALSE; } - TfLiteStatus status = interpreter->ResizeInputTensor( - interpreter->inputs()[idx], convertJIntArrayToVector(env, dims)); - if (status != kTfLiteOk) { - throwException(env, kIllegalArgumentException, - "Failed to resize %d-th input: %s", idx, - error_reporter->CachedErrorMessage()); + // check whether it is resizing with the same dimensions. + TfLiteTensor* target = interpreter->tensor(input_idx); + bool is_changed = areDimsDifferent(env, target, dims); + if (is_changed) { + TfLiteStatus status = interpreter->ResizeInputTensor( + interpreter->inputs()[idx], convertJIntArrayToVector(env, dims)); + if (status != kTfLiteOk) { + throwException(env, kIllegalArgumentException, + "Failed to resize %d-th input: %s", idx, + error_reporter->CachedErrorMessage()); + return JNI_FALSE; + } } + return is_changed ? JNI_TRUE : JNI_FALSE; } JNIEXPORT void JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_delete( 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 c52a7e4e439936344be26d5761fb5747db64794a..fb76125471a11a7548d56c37354d38d4c2c83c06 100644 --- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h +++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/interpreter.h" @@ -28,6 +29,9 @@ limitations under the License. namespace tflite { // This is to be provided at link-time by a library. extern std::unique_ptr CreateOpResolver(); +extern timespec getCurrentTime(); +extern jlong timespec_diff_nanoseconds(struct timespec* start, + struct timespec* stop); } // namespace tflite #ifdef __cplusplus @@ -65,6 +69,17 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_useNNAPI(JNIEnv* env, jlong handle, jboolean state); +/* + * Class: org_tensorflow_lite_NativeInterpreterWrapper + * Method: + * Signature: (JI) + */ +JNIEXPORT void JNICALL +Java_org_tensorflow_lite_NativeInterpreterWrapper_numThreads(JNIEnv* env, + jclass clazz, + jlong handle, + jint num_threads); + /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: @@ -104,21 +119,23 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter( /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: - * Signature: (JJ[Ljava/lang/Object;[I[I[Ljava/lang/Object;)[J + * Signature: + * (JJ[Ljava/lang/Object;[I[I[Ljava/lang/Object;Ljava/lang/Object;Z)[J */ JNIEXPORT jlongArray JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_run( JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle, jobjectArray sizes, jintArray data_types, jintArray nums_of_bytes, - jobjectArray values); + jobjectArray values, jobject wrapper, jboolean memory_allocated); /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: * Signature: (JII)[I * - * It gets input dimensions if num_bytes matches number of bytes required by - * the input, else returns null and throws IllegalArgumentException. + * Gets input dimensions. If num_bytes is non-negative, it will check whether + * num_bytes matches num of bytes required by the input, and return null and + * throw IllegalArgumentException if not. */ JNIEXPORT jintArray JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims( @@ -127,11 +144,23 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims( /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: - * Signature: (JJI[I) + * Signature: (JI)I * - * It resizes dimensions of a input. + * Gets output dimensions. */ -JNIEXPORT void JNICALL +JNIEXPORT jint JNICALL +Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputDataType( + JNIEnv* env, jclass clazz, jlong handle, jint output_idx); + +/* + * Class: org_tensorflow_lite_NativeInterpreterWrapper + * Method: + * Signature: (JJI[I)Z + * + * It returns true if resizing input tensor to different dimensions, else return + * false. + */ +JNIEXPORT jboolean JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_resizeInput( JNIEnv* env, jclass clazz, jlong interpreter_handle, jlong error_handle, jint input_idx, jintArray dims); diff --git a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java index 424b3de6c97672e310c54230a7ac1204f46d9ac8..61d6c35ec86beebf78dd81e17e145863516802fa 100644 --- a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java +++ b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java @@ -218,4 +218,52 @@ public final class InterpreterTest { int index = interpreter.getOutputIndex("MobilenetV1/Predictions/Softmax"); assertThat(index).isEqualTo(0); } + + @Test + public void testTurnOffNNAPI() throws Exception { + Path path = MODEL_FILE.toPath(); + FileChannel fileChannel = + (FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ)); + MappedByteBuffer mappedByteBuffer = + fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size()); + Interpreter interpreter = new Interpreter(mappedByteBuffer); + interpreter.setUseNNAPI(true); + float[] oneD = {1.23f, 6.54f, 7.81f}; + float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; + float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; + float[][][][] fourD = {threeD, threeD}; + float[][][][] parsedOutputs = new float[2][8][8][3]; + interpreter.run(fourD, parsedOutputs); + float[] outputOneD = parsedOutputs[0][0][0]; + float[] expected = {3.69f, 19.62f, 23.43f}; + assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); + interpreter.setUseNNAPI(false); + interpreter.run(fourD, parsedOutputs); + outputOneD = parsedOutputs[0][0][0]; + assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); + interpreter.close(); + fileChannel.close(); + } + + @Test + public void testTurnOnNNAPI() throws Exception { + Path path = MODEL_FILE.toPath(); + FileChannel fileChannel = + (FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ)); + MappedByteBuffer mappedByteBuffer = + fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size()); + Interpreter interpreter = new Interpreter(mappedByteBuffer); + interpreter.setUseNNAPI(true); + float[] oneD = {1.23f, 6.54f, 7.81f}; + float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; + float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; + float[][][][] fourD = {threeD, threeD}; + float[][][][] parsedOutputs = new float[2][8][8][3]; + interpreter.run(fourD, parsedOutputs); + float[] outputOneD = parsedOutputs[0][0][0]; + float[] expected = {3.69f, 19.62f, 23.43f}; + assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); + interpreter.close(); + fileChannel.close(); + } } 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 473f73816fd3c0a414a2c2e232dec299579fcbb6..dbe45e5a05b8227b441de7ca6747f61d010ae210 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 @@ -47,6 +47,9 @@ public final class NativeInterpreterWrapperTest { private static final String MODEL_WITH_CUSTOM_OP_PATH = "tensorflow/contrib/lite/java/src/testdata/with_custom_op.lite"; + private static final String NONEXISTING_MODEL_PATH = + "tensorflow/contrib/lite/java/src/testdata/nonexisting_model.bin"; + @Test public void testConstructor() { NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); @@ -60,9 +63,18 @@ public final class NativeInterpreterWrapperTest { NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(INVALID_MODEL_PATH); fail(); } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains("Model provided has model identifier ' is ', should be 'TFL3'"); + assertThat(e).hasMessageThat().contains("The model is not a valid Flatbuffer file"); + } + } + + @Test + public void testConstructorWithNonexistingModel() { + try { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(NONEXISTING_MODEL_PATH); + fail(); + } catch (IllegalArgumentException e) { + assertThat(e).hasMessageThat().contains("The model is not a valid Flatbuffer file"); + assertThat(e).hasMessageThat().contains("Could not open"); } } @@ -96,6 +108,30 @@ public final class NativeInterpreterWrapperTest { wrapper.close(); } + @Test + public void testRunWithInputsOfSameDims() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + float[] oneD = {1.23f, -6.54f, 7.81f}; + float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; + float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; + float[][][][] fourD = {threeD, threeD}; + Object[] inputs = {fourD}; + Tensor[] outputs = wrapper.run(inputs); + assertThat(outputs.length).isEqualTo(1); + float[][][][] parsedOutputs = new float[2][8][8][3]; + outputs[0].copyTo(parsedOutputs); + float[] outputOneD = parsedOutputs[0][0][0]; + float[] expected = {3.69f, -19.62f, 23.43f}; + assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); + outputs = wrapper.run(inputs); + assertThat(outputs.length).isEqualTo(1); + parsedOutputs = new float[2][8][8][3]; + outputs[0].copyTo(parsedOutputs); + outputOneD = parsedOutputs[0][0][0]; + assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); + wrapper.close(); + } + @Test public void testRunWithInt() { NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(INT_MODEL_PATH); @@ -419,4 +455,87 @@ public final class NativeInterpreterWrapperTest { assertThat(shape[1]).isEqualTo(3); assertThat(shape[2]).isEqualTo(1); } + + @Test + public void testGetInferenceLatency() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + float[] oneD = {1.23f, 6.54f, 7.81f}; + float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; + float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; + float[][][][] fourD = {threeD, threeD}; + Object[] inputs = {fourD}; + Tensor[] outputs = wrapper.run(inputs); + assertThat(outputs.length).isEqualTo(1); + assertThat(wrapper.getLastNativeInferenceDurationNanoseconds()).isGreaterThan(0L); + wrapper.close(); + } + + @Test + public void testGetInferenceLatencyWithNewWrapper() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + assertThat(wrapper.getLastNativeInferenceDurationNanoseconds()).isNull(); + wrapper.close(); + } + + @Test + public void testGetLatencyAfterFailedInference() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + float[] oneD = {1.23f, 6.54f, 7.81f}; + float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD}; + float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; + float[][][][] fourD = {threeD, threeD}; + Object[] inputs = {fourD}; + try { + wrapper.run(inputs); + fail(); + } catch (IllegalArgumentException e) { + assertThat(e) + .hasMessageThat() + .contains("0-th input dimension should be [?,8,8,3], but found [?,8,7,3]"); + } + assertThat(wrapper.getLastNativeInferenceDurationNanoseconds()).isNull(); + wrapper.close(); + } + + @Test + public void testGetInputDims() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + int[] expectedDims = {1, 8, 8, 3}; + assertThat(wrapper.getInputDims(0)).isEqualTo(expectedDims); + wrapper.close(); + } + + @Test + public void testGetInputDimsOutOfRange() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + try { + wrapper.getInputDims(-1); + fail(); + } catch (IllegalArgumentException e) { + assertThat(e).hasMessageThat().contains("Out of range"); + } + try { + wrapper.getInputDims(1); + fail(); + } catch (IllegalArgumentException e) { + assertThat(e).hasMessageThat().contains("Out of range"); + } + wrapper.close(); + } + + @Test + public void testGetOutputDataType() { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); + assertThat(wrapper.getOutputDataType(0)).contains("float"); + wrapper.close(); + wrapper = new NativeInterpreterWrapper(LONG_MODEL_PATH); + assertThat(wrapper.getOutputDataType(0)).contains("long"); + wrapper.close(); + wrapper = new NativeInterpreterWrapper(INT_MODEL_PATH); + assertThat(wrapper.getOutputDataType(0)).contains("int"); + wrapper.close(); + wrapper = new NativeInterpreterWrapper(BYTE_MODEL_PATH); + assertThat(wrapper.getOutputDataType(0)).contains("byte"); + wrapper.close(); + } } diff --git a/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/TestHelper.java b/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/TestHelper.java index 8660cabf709e6531a5667a16e5cf43a93c7135bd..3722e51b3b7a2f8698783c12d275caf8d27be525 100644 --- a/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/TestHelper.java +++ b/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/TestHelper.java @@ -32,4 +32,69 @@ public class TestHelper { throw new IllegalArgumentException("Interpreter has not initialized; Failed to setUseNNAPI."); } } + + /** + * Sets the number of threads for an {@code Interpreter}. + * + * @param interpreter an instance of {@code Interpreter}. If it is not initialized, an {@code + * IllegalArgumentException} will be thrown. + * @param numRecommendedThreads an integer value indicating the number of recommended threads. + */ + public static void setNumThreads(Interpreter interpreter, int numRecommendedThreads) { + if (interpreter != null && interpreter.wrapper != null) { + interpreter.wrapper.setNumThreads(numRecommendedThreads); + } else { + throw new IllegalArgumentException("Interpreter has not initialized; Failed to setUseNNAPI."); + } + } + /** + * Gets the last inference duration in nanoseconds. It returns null if there is no previous + * inference run or the last inference run failed. + * + * @param interpreter an instance of {@code Interpreter}. If it is not initialized, an {@code + * IllegalArgumentException} will be thrown. + */ + public static Long getLastNativeInferenceDurationNanoseconds(Interpreter interpreter) { + if (interpreter != null && interpreter.wrapper != null) { + return interpreter.wrapper.getLastNativeInferenceDurationNanoseconds(); + } else { + throw new IllegalArgumentException("Interpreter has not initialized; Failed to get latency."); + } + } + + /** + * Gets the dimensions of an input. + * + * @param interpreter an instance of {@code Interpreter}. If it is not initialized, an {@code + * IllegalArgumentException} will be thrown. + * @param index an integer index of the input. If it is invalid, an {@code + * IllegalArgumentException} will be thrown. + */ + public static int[] getInputDims(Interpreter interpreter, int index) { + if (interpreter != null && interpreter.wrapper != null) { + return interpreter.wrapper.getInputDims(index); + } else { + throw new IllegalArgumentException( + "Interpreter has not initialized;" + " Failed to get input dimensions."); + } + } + + /** + * Gets the string name of the data type of an output. + * + * @param interpreter an instance of {@code Interpreter}. If it is not initialized, an {@code + * IllegalArgumentException} will be thrown. + * @param index an integer index of the output. If it is invalid, an {@code + * IllegalArgumentException} will be thrown. + * @return string name of the data type. Possible values include "float", "int", "byte", and + * "long". + */ + public static String getOutputDataType(Interpreter interpreter, int index) { + if (interpreter != null && interpreter.wrapper != null) { + return interpreter.wrapper.getOutputDataType(index); + } else { + throw new IllegalArgumentException( + "Interpreter has not initialized;" + " Failed to get output data type."); + } + } } diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD index d9051f3516367581d31f7477e0f5bb5fa0e979f8..200cb3075bd20a8fe3f4d9cb574a86ecdb63ab78 100644 --- a/tensorflow/contrib/lite/kernels/BUILD +++ b/tensorflow/contrib/lite/kernels/BUILD @@ -5,15 +5,17 @@ package(default_visibility = [ licenses(["notice"]) # Apache 2.0 load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", -) +load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") +load("//tensorflow:tensorflow.bzl", "tf_cc_test") tf_cc_test( name = "optional_tensor_test", size = "small", srcs = ["optional_tensor_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -33,11 +35,27 @@ cc_library( "//tensorflow/contrib/lite:schema_fbs_version", "//tensorflow/contrib/lite:string_util", "//tensorflow/contrib/lite/testing:util", - "//tensorflow/core:lib", + "//tensorflow/core:tflite_portable_logging", "@com_google_googletest//:gtest", ], ) +cc_library( + name = "eigen_support", + srcs = [ + "eigen_support.cc", + ], + hdrs = [ + "eigen_support.h", + ], + copts = tflite_copts(), + deps = [ + ":op_macros", + "//tensorflow/contrib/lite:context", + "//third_party/eigen3", + ], +) + cc_library( name = "gemm_support", srcs = [ @@ -90,6 +108,10 @@ tf_cc_test( name = "kernel_util_test", size = "small", srcs = ["kernel_util_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":kernel_util", "//tensorflow/contrib/lite/testing:util", @@ -97,20 +119,36 @@ tf_cc_test( ], ) +tf_cc_test( + name = "test_util_test", + size = "small", + srcs = ["test_util_test.cc"], + deps = [ + ":test_util", + "//tensorflow/contrib/lite/testing:util", + "@com_google_googletest//:gtest", + ], +) + cc_library( name = "builtin_ops", srcs = [ "activations.cc", "add.cc", + "audio_spectrogram.cc", "basic_rnn.cc", "batch_to_space_nd.cc", + "bidirectional_sequence_lstm.cc", "bidirectional_sequence_rnn.cc", + "cast.cc", "concatenation.cc", "conv.cc", "depthwise_conv.cc", + "dequantize.cc", "div.cc", "embedding_lookup.cc", "embedding_lookup_sparse.cc", + "exp.cc", "fully_connected.cc", "gather.cc", "hashtable_lookup.cc", @@ -119,6 +157,7 @@ cc_library( "lsh_projection.cc", "lstm.cc", "mean.cc", + "mfcc.cc", "mul.cc", "pad.cc", "pooling.cc", @@ -128,10 +167,12 @@ cc_library( "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", @@ -150,20 +191,49 @@ cc_library( }), deps = [ ":activation_functor", + ":eigen_support", ":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:audio_utils", + "//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", "//tensorflow/contrib/lite/kernels/internal:reference", "//tensorflow/contrib/lite/kernels/internal:reference_base", - "//tensorflow/contrib/lite/kernels/internal:round", "//tensorflow/contrib/lite/kernels/internal:tensor_utils", "@farmhash_archive//:farmhash", + "@flatbuffers", + ], +) + +tf_cc_test( + name = "audio_spectrogram_test", + size = "small", + srcs = ["audio_spectrogram_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + "@flatbuffers", + ], +) + +tf_cc_test( + name = "mfcc_test", + size = "small", + srcs = ["mfcc_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + "@flatbuffers", ], ) @@ -171,6 +241,10 @@ tf_cc_test( name = "activations_test", size = "small", srcs = ["activations_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -183,6 +257,42 @@ tf_cc_test( name = "add_test", size = "small", srcs = ["add_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "div_test", + size = "small", + srcs = ["div_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "sub_test", + size = "small", + srcs = ["sub_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -195,6 +305,10 @@ tf_cc_test( name = "transpose_test", size = "small", srcs = ["transpose_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -209,6 +323,10 @@ tf_cc_test( name = "space_to_batch_nd_test", size = "small", srcs = ["space_to_batch_nd_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -221,6 +339,22 @@ tf_cc_test( name = "batch_to_space_nd_test", size = "small", srcs = ["batch_to_space_nd_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "cast_test", + size = "small", + srcs = ["cast_test.cc"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -233,6 +367,10 @@ tf_cc_test( name = "concatenation_test", size = "small", srcs = ["concatenation_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -245,10 +383,15 @@ tf_cc_test( name = "conv_test", size = "small", srcs = ["conv_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_absl//absl/memory", "@com_google_googletest//:gtest", ], ) @@ -257,6 +400,10 @@ tf_cc_test( name = "depthwise_conv_test", size = "small", srcs = ["depthwise_conv_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -265,10 +412,43 @@ tf_cc_test( ], ) +tf_cc_test( + name = "dequantize_test", + size = "small", + srcs = ["dequantize_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_absl//absl/memory", + "@com_google_googletest//:gtest", + ], +) + tf_cc_test( name = "basic_rnn_test", size = "small", srcs = ["basic_rnn_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "bidirectional_sequence_lstm_test", + size = "small", + srcs = ["bidirectional_sequence_lstm_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -281,6 +461,10 @@ tf_cc_test( name = "unidirectional_sequence_lstm_test", size = "small", srcs = ["unidirectional_sequence_lstm_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -293,6 +477,9 @@ tf_cc_test( name = "bidirectional_sequence_rnn_test", size = "small", srcs = ["bidirectional_sequence_rnn_test.cc"], + tags = [ + "tflite_not_portable", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -305,6 +492,10 @@ tf_cc_test( name = "unidirectional_sequence_rnn_test", size = "small", srcs = ["unidirectional_sequence_rnn_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -317,6 +508,26 @@ tf_cc_test( name = "l2norm_test", size = "small", srcs = ["l2norm_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "exp_test", + size = "small", + srcs = ["exp_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -329,6 +540,10 @@ tf_cc_test( name = "mean_test", size = "small", srcs = ["mean_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -341,6 +556,10 @@ tf_cc_test( name = "mul_test", size = "small", srcs = ["mul_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -353,6 +572,10 @@ tf_cc_test( name = "pad_test", size = "small", srcs = ["pad_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -365,6 +588,10 @@ tf_cc_test( name = "reshape_test", size = "small", srcs = ["reshape_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -377,6 +604,27 @@ tf_cc_test( name = "gather_test", size = "small", srcs = ["gather_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + 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 = "topk_v2_test", + size = "small", + srcs = ["topk_v2_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:builtin_op_data", @@ -390,6 +638,10 @@ tf_cc_test( name = "resize_bilinear_test", size = "small", srcs = ["resize_bilinear_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -402,6 +654,10 @@ tf_cc_test( name = "svdf_test", size = "small", srcs = ["svdf_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -414,6 +670,10 @@ tf_cc_test( name = "embedding_lookup_test", size = "small", srcs = ["embedding_lookup_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -426,6 +686,10 @@ tf_cc_test( name = "embedding_lookup_sparse_test", size = "small", srcs = ["embedding_lookup_sparse_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -438,6 +702,10 @@ tf_cc_test( name = "fully_connected_test", size = "small", srcs = ["fully_connected_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -450,6 +718,10 @@ tf_cc_test( name = "local_response_norm_test", size = "small", srcs = ["local_response_norm_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -462,6 +734,10 @@ tf_cc_test( name = "pooling_test", size = "small", srcs = ["pooling_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -474,6 +750,27 @@ tf_cc_test( name = "softmax_test", size = "small", srcs = ["softmax_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "//tensorflow/contrib/lite/kernels/internal:reference_base", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "log_softmax_test", + size = "small", + srcs = ["log_softmax_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -487,6 +784,10 @@ tf_cc_test( name = "lsh_projection_test", size = "small", srcs = ["lsh_projection_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -499,6 +800,10 @@ tf_cc_test( name = "hashtable_lookup_test", size = "small", srcs = ["hashtable_lookup_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -512,6 +817,10 @@ tf_cc_test( name = "lstm_test", size = "small", srcs = ["lstm_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -524,6 +833,10 @@ tf_cc_test( name = "skip_gram_test", size = "small", srcs = ["skip_gram_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -537,6 +850,26 @@ tf_cc_test( name = "space_to_depth_test", size = "small", srcs = ["space_to_depth_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "split_test", + size = "small", + srcs = ["split_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -549,6 +882,10 @@ tf_cc_test( name = "squeeze_test", size = "small", srcs = ["squeeze_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -561,6 +898,10 @@ tf_cc_test( name = "strided_slice_test", size = "small", srcs = ["strided_slice_test.cc"], + tags = [ + "tflite_not_portable_ios_arm64", + "tflite_not_portable_ios_x86_64", + ], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -580,3 +921,5 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) + +tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/kernels/activations.cc b/tensorflow/contrib/lite/kernels/activations.cc index 3c5c77815d0f2592ab549152b4d77f45b967a660..093761c43c1cb41ddb2245da13c963014b51271c 100644 --- a/tensorflow/contrib/lite/kernels/activations.cc +++ b/tensorflow/contrib/lite/kernels/activations.cc @@ -63,6 +63,33 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) { TfLiteIntArrayCopy(input->dims)); } +TfLiteStatus TanhPrepare(TfLiteContext* context, TfLiteNode* node) { + OpData* data = reinterpret_cast(node->user_data); + + TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + TfLiteTensor* input = GetInput(context, node, 0); + TfLiteTensor* output = GetOutput(context, node, 0); + TF_LITE_ENSURE_EQ(context, input->type, output->type); + + if (input->type == kTfLiteUInt8) { + static constexpr int kInputIntegerBits = 4; + + const double input_real_multiplier = + input->params.scale * + static_cast(1 << (31 - kInputIntegerBits)); + + QuantizeMultiplierGreaterThanOne(input_real_multiplier, + &data->input_multiplier, + &data->input_left_shift); + data->input_range_radius = + CalculateInputRadius(kInputIntegerBits, data->input_left_shift); + } + + return context->ResizeTensor(context, output, + TfLiteIntArrayCopy(input->dims)); +} + TfLiteStatus SigmoidPrepare(TfLiteContext* context, TfLiteNode* node) { OpData* data = reinterpret_cast(node->user_data); @@ -180,6 +207,7 @@ TfLiteStatus Relu6Eval(TfLiteContext* context, TfLiteNode* node) { } TfLiteStatus TanhEval(TfLiteContext* context, TfLiteNode* node) { + OpData* data = reinterpret_cast(node->user_data); TfLiteTensor* input = GetInput(context, node, 0); TfLiteTensor* output = GetOutput(context, node, 0); switch (input->type) { @@ -191,6 +219,14 @@ TfLiteStatus TanhEval(TfLiteContext* context, TfLiteNode* node) { for (; in < in_end; in++, out++) *out = std::tanh(*in); return kTfLiteOk; } break; + case kTfLiteUInt8: { + optimized_ops::Tanh(GetTensorData(input), GetTensorDims(input), + input->params.zero_point, data->input_range_radius, + data->input_multiplier, data->input_left_shift, + GetTensorData(output), + GetTensorDims(output)); + return kTfLiteOk; + } break; default: context->ReportError(context, "Only float32 supported currently."); return kTfLiteError; @@ -337,6 +373,21 @@ TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) { } } +TfLiteStatus LogSoftmaxEval(TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* input = GetInput(context, node, 0); + TfLiteTensor* output = GetOutput(context, node, 0); + switch (input->type) { + case kTfLiteFloat32: + optimized_ops::LogSoftmax( + GetTensorData(input), GetTensorDims(input), + GetTensorData(output), GetTensorDims(output)); + return kTfLiteOk; + default: + context->ReportError(context, "Only float32 supported currently."); + return kTfLiteError; + } +} + } // namespace activations TfLiteRegistration* Register_RELU() { @@ -361,8 +412,8 @@ TfLiteRegistration* Register_RELU6() { } TfLiteRegistration* Register_TANH() { - static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - activations::GenericPrepare, + static TfLiteRegistration r = {activations::Init, activations::Free, + activations::TanhPrepare, activations::TanhEval}; return &r; } @@ -381,6 +432,13 @@ TfLiteRegistration* Register_SOFTMAX() { return &r; } +TfLiteRegistration* Register_LOG_SOFTMAX() { + static TfLiteRegistration r = {activations::Init, activations::Free, + activations::GenericPrepare, + activations::LogSoftmaxEval}; + return &r; +} + } // namespace builtin } // namespace ops } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/activations_test.cc b/tensorflow/contrib/lite/kernels/activations_test.cc index 68d49944e51b043b6b82aa1589d22f6ebed37574..b9a96e3f79677c5a94ade86a6b334abb4c265fa1 100644 --- a/tensorflow/contrib/lite/kernels/activations_test.cc +++ b/tensorflow/contrib/lite/kernels/activations_test.cc @@ -52,6 +52,14 @@ class BaseActivationsOpModel : public SingleOpModel { BuildInterpreter({GetShape(input_)}); } + BaseActivationsOpModel(BuiltinOperator type, const TensorData &input, + const TensorData &output) { + input_ = AddInput(input); + output_ = AddOutput(output); + SetBuiltinOp(type, BuiltinOptions_NONE, 0); + BuildInterpreter({GetShape(input_)}); + } + protected: int input_; int output_; @@ -143,6 +151,27 @@ TEST(FloatActivationsOpTest, Tanh) { }))); } +TEST(QuantizedActivationsOpTest, Tanh) { + QuantizedActivationsOpModel m( + BuiltinOperator_TANH, + /*input=*/{TensorType_UINT8, {1, 2, 4, 1}, -8, 8}, + /*output=*/{TensorType_UINT8, {1, 2, 4, 1}, -1, 1}); + m.SetInput({ + 0, -6, 2, 4, // + -4, -2, 8, 1, // + }); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(ArrayFloatNear( + { + 0.0, -0.999987, 0.964027, 0.999329, // + -0.996078, -0.96402, 0.99999, 0.76159, // + }, + 4 * (1. / 256)))); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray({128, 0, 251, 255, 0, 5, 255, 226})); +} + TEST(FloatActivationsOpTest, Sigmoid) { FloatActivationsOpModel m(BuiltinOperator_LOGISTIC, /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); @@ -313,6 +342,47 @@ TEST(QuantizedActivationsOpTest, Softmax2D) { kQuantizedTolerance))); } +// This contains the same test values as the Softmax test, but reference answer +// generated via the following snippet of python: +// logits1 = tf.constant([[0, -6, 2, 4],[3, -2, 10, 1]], dtype=tf.float32) +// logits2 = tf.constant([[0,-6],[2,4],[3,-2],[10,1]], dtype=tf.float32) +// lsm1 = tf.nn.log_softmax(logits1) +// lsm2 = tf.nn.log_softmax(logits2) +// with tf.Session() as sess: +// print('lsm1', sess.run(lsm1)) +// print('lsm2', sess.run(lsm2)) + +TEST(FloatActivationsOpTest, LogSoftmax) { + FloatActivationsOpModel m(BuiltinOperator_LOG_SOFTMAX, + /*input=*/{TensorType_FLOAT32, {2, 4}}); + m.SetInput({ + 0, -6, 2, 4, // + 3, -2, 10, 1, // + }); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ + -4.14297, -10.14297, -2.14297, -.142971, // + -7.00104, -12.00104, -.00104087, -9.00104, // + }))); + + // Same input, but a different shape. + FloatActivationsOpModel m2(BuiltinOperator_LOG_SOFTMAX, + /*input=*/{TensorType_FLOAT32, {4, 2}}); + m2.SetInput({ + 0, -6, // + 2, 4, // + 3, -2, // + 10, 1, // + }); + m2.Invoke(); + EXPECT_THAT(m2.GetOutput(), ElementsAreArray(ArrayFloatNear({ + -.00247565, -6.00247, // + -2.12692, -.126928, // + -.00671534, -5.00671, // + -.000123374, -9.00012, // + }))); +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/add.cc b/tensorflow/contrib/lite/kernels/add.cc index fb5764f280bd0dbe89e716a3f38ba5678d14581e..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; + + 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 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)) +#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) { - TF_LITE_ADD(reference_ops); + if (data->requires_broadcast) { + TF_LITE_ADD(reference_ops, BroadcastAdd); + } else { + TF_LITE_ADD(reference_ops, Add); + } } else { - TF_LITE_ADD(optimized_ops); + 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 306dfc3e803d3df34061767ba9ced032299bfa26..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}); @@ -78,9 +80,9 @@ TEST(FloatAddOpModel, NoActivation) { } TEST(FloatAddOpModel, ActivationRELU_N1_TO_1) { - FloatAddOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {}}, - ActivationFunctionType_RELU_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(); @@ -92,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}); @@ -102,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 = { @@ -112,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]); @@ -133,6 +154,7 @@ TEST(QuantizedAddOpModel, QuantizedTestsActivationRELU_N1_TO_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_RELU_N1_TO_1); m.QuantizeAndPopulate(m.input1(), inputs1[i]); @@ -150,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}); @@ -162,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/audio_spectrogram.cc b/tensorflow/contrib/lite/kernels/audio_spectrogram.cc new file mode 100644 index 0000000000000000000000000000000000000000..602f3888c10b3790dc0328c817bdd83276544b25 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/audio_spectrogram.cc @@ -0,0 +1,165 @@ +/* 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/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/spectrogram.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" + +#include "flatbuffers/flexbuffers.h" + +namespace tflite { +namespace ops { +namespace custom { +namespace audio_spectrogram { + +constexpr int kInputTensor = 0; +constexpr int kOutputTensor = 0; + +enum KernelType { + kReference, +}; + +typedef struct { + int window_size; + int stride; + bool magnitude_squared; + int output_height; + internal::Spectrogram* spectrogram; +} TfLiteAudioSpectrogramParams; + +void* Init(TfLiteContext* context, const char* buffer, size_t length) { + auto* data = new TfLiteAudioSpectrogramParams; + + const uint8_t* buffer_t = reinterpret_cast(buffer); + + const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap(); + data->window_size = m["window_size"].AsInt64(); + data->stride = m["stride"].AsInt64(); + data->magnitude_squared = m["magnitude_squared"].AsBool(); + + data->spectrogram = new internal::Spectrogram; + + return data; +} + +void Free(TfLiteContext* context, void* buffer) { + auto* params = reinterpret_cast(buffer); + delete params->spectrogram; + delete params; +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + auto* params = + reinterpret_cast(node->user_data); + + TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + TF_LITE_ENSURE_EQ(context, NumDimensions(input), 2); + + TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); + TF_LITE_ENSURE_EQ(context, input->type, output->type); + + TF_LITE_ENSURE(context, params->spectrogram->Initialize(params->window_size, + params->stride)); + const int64_t sample_count = input->dims->data[0]; + const int64_t length_minus_window = (sample_count - params->window_size); + if (length_minus_window < 0) { + params->output_height = 0; + } else { + params->output_height = 1 + (length_minus_window / params->stride); + } + TfLiteIntArray* output_size = TfLiteIntArrayCreate(3); + output_size->data[0] = input->dims->data[1]; + output_size->data[1] = params->output_height; + output_size->data[2] = params->spectrogram->output_frequency_channels(); + + return context->ResizeTensor(context, output, output_size); +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = + reinterpret_cast(node->user_data); + + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + TF_LITE_ENSURE(context, params->spectrogram->Initialize(params->window_size, + params->stride)); + + const float* input_data = GetTensorData(input); + + const int64_t sample_count = input->dims->data[0]; + const int64_t channel_count = input->dims->data[1]; + + const int64_t output_width = params->spectrogram->output_frequency_channels(); + + float* output_flat = GetTensorData(output); + + std::vector input_for_channel(sample_count); + for (int64_t channel = 0; channel < channel_count; ++channel) { + float* output_slice = + output_flat + (channel * params->output_height * output_width); + for (int i = 0; i < sample_count; ++i) { + input_for_channel[i] = input_data[i * channel_count + channel]; + } + std::vector> spectrogram_output; + TF_LITE_ENSURE(context, + params->spectrogram->ComputeSquaredMagnitudeSpectrogram( + input_for_channel, &spectrogram_output)); + TF_LITE_ENSURE_EQ(context, spectrogram_output.size(), + params->output_height); + TF_LITE_ENSURE(context, spectrogram_output.empty() || + (spectrogram_output[0].size() == output_width)); + for (int row_index = 0; row_index < params->output_height; ++row_index) { + const std::vector& spectrogram_row = spectrogram_output[row_index]; + TF_LITE_ENSURE_EQ(context, spectrogram_row.size(), output_width); + float* output_row = output_slice + (row_index * output_width); + if (params->magnitude_squared) { + for (int i = 0; i < output_width; ++i) { + output_row[i] = spectrogram_row[i]; + } + } else { + for (int i = 0; i < output_width; ++i) { + output_row[i] = sqrtf(spectrogram_row[i]); + } + } + } + } + return kTfLiteOk; +} + +} // namespace audio_spectrogram + +TfLiteRegistration* Register_AUDIO_SPECTROGRAM() { + static TfLiteRegistration r = { + audio_spectrogram::Init, audio_spectrogram::Free, + audio_spectrogram::Prepare, + audio_spectrogram::Eval}; + return &r; +} + +} // namespace custom +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/audio_spectrogram_test.cc b/tensorflow/contrib/lite/kernels/audio_spectrogram_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..8d460fdfc610ef9a867acd492ca0558fb6eab8c3 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/audio_spectrogram_test.cc @@ -0,0 +1,122 @@ +/* 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 "flatbuffers/flexbuffers.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 custom { + +TfLiteRegistration* Register_AUDIO_SPECTROGRAM(); + +namespace { + +using ::testing::ElementsAre; +using ::testing::ElementsAreArray; + +class BaseAudioSpectrogramOpModel : public SingleOpModel { + public: + BaseAudioSpectrogramOpModel(const TensorData& input1, + const TensorData& output, int window_size, + int stride, bool magnitude_squared) { + input1_ = AddInput(input1); + output_ = AddOutput(output); + + flexbuffers::Builder fbb; + fbb.Map([&]() { + fbb.Int("window_size", window_size); + fbb.Int("stride", stride); + fbb.Bool("magnitude_squared", magnitude_squared); + }); + fbb.Finish(); + SetCustomOp("AudioSpectrogram", fbb.GetBuffer(), + Register_AUDIO_SPECTROGRAM); + BuildInterpreter({GetShape(input1_)}); + } + + int input1() { return input1_; } + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input1_; + int output_; +}; + +TEST(BaseAudioSpectrogramOpModel, NonSquaredTest) { + BaseAudioSpectrogramOpModel m({TensorType_FLOAT32, {8, 1}}, + {TensorType_FLOAT32, {}}, 8, 1, false); + m.PopulateTensor(m.input1(), + {-1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f, 1.0f, 0.0f}); + + m.Invoke(); + + std::vector output_shape = m.GetOutputShape(); + EXPECT_EQ(3, output_shape.size()); + EXPECT_THAT(output_shape, ElementsAre(1, 1, 5)); + + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear( + {0.0f, 1.0f, 2.0f, 1.0f, 0.0f}, 1e-3))); +} + +TEST(SpectrogramOpTest, SquaredTest) { + BaseAudioSpectrogramOpModel m({TensorType_FLOAT32, {8, 1}}, + {TensorType_FLOAT32, {}}, 8, 1, true); + m.PopulateTensor(m.input1(), + {-1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f, 1.0f, 0.0f}); + + m.Invoke(); + + std::vector output_shape = m.GetOutputShape(); + EXPECT_EQ(3, output_shape.size()); + EXPECT_THAT(output_shape, ElementsAre(1, 1, 5)); + + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear( + {0.f, 1.f, 4.f, 1.f, 0.f}, 1e-3))); +} + +TEST(SpectrogramOpTest, StrideTest) { + BaseAudioSpectrogramOpModel m({TensorType_FLOAT32, {10, 1}}, + {TensorType_FLOAT32, {}}, 8, 2, true); + m.PopulateTensor(m.input1(), {-1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f, + 1.0f, 0.0f, 1.0f, 0.0f}); + + m.Invoke(); + + std::vector output_shape = m.GetOutputShape(); + EXPECT_THAT(output_shape, ElementsAre(1, 2, 5)); + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear( + {0, 1, 4, 1, 0, 1, 2, 1, 2, 1}, 1e-3))); +} + +} // namespace +} // namespace custom +} // namespace ops +} // 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/basic_rnn.cc b/tensorflow/contrib/lite/kernels/basic_rnn.cc index a0391e030f69e59ed0ecd76ab4ec287a0c415b7a..2c5074eca3176c7f33a6f051b492dc41333257ed 100644 --- a/tensorflow/contrib/lite/kernels/basic_rnn.cc +++ b/tensorflow/contrib/lite/kernels/basic_rnn.cc @@ -23,6 +23,7 @@ limitations under the License. #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 { @@ -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/batch_to_space_nd.cc b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc index d84a77039bfc065051d6439ae98656303dbe48ac..bc438f99c6a72fdbc2794dee03524db6a7523834 100644 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc +++ b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc @@ -57,6 +57,7 @@ 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); TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape), kBlockSizeDimensionNum); @@ -65,7 +66,13 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->crops), kSpatialDimensionNum); - // TODO(ycling): Add crops as part of calculation. + // 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, input_size->data[0] % (block_shape[0] * block_shape[1]), 0); @@ -109,7 +116,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* 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)); - TfLiteTensorRealloc(op_context.output->bytes, op_context.output); } #define TF_LITE_BATCH_TO_SPACE_ND(type, scalar) \ 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 c9152bf96729c07e4e06a606f193aac269aad260..8485cde1b40066f2070855bca91ea78a9f80e83c 100644 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc +++ b/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc @@ -119,6 +119,19 @@ TEST(BatchToSpaceNDOpTest, InvalidShapeTest) { "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_lstm.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc new file mode 100644 index 0000000000000000000000000000000000000000..a64ac42bc43336db928d2682e290f5263f3db0f4 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc @@ -0,0 +1,702 @@ +/* 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/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 bidirectional_sequence_lstm { + +// Input Tensors of size {max_time, n_batch, n_input} +constexpr int kInputTensor = 0; + +// Forward LSTM cell tensors. +// Input weight tensors of size: {n_cell, n_input} +constexpr int kFwInputToInputWeightsTensor = 1; // Optional +constexpr int kFwInputToForgetWeightsTensor = 2; +constexpr int kFwInputToCellWeightsTensor = 3; +constexpr int kFwInputToOutputWeightsTensor = 4; + +// Recurrent weight tensors of size {n_cell, n_output} +constexpr int kFwRecurrentToInputWeightsTensor = 5; // Optional +constexpr int kFwRecurrentToForgetWeightsTensor = 6; +constexpr int kFwRecurrentToCellWeightsTensor = 7; +constexpr int kFwRecurrentToOutputWeightsTensor = 8; + +// Peephole weights tensors of size {n_cell}, representing a diagonal matrix. +constexpr int kFwCellToInputWeightsTensor = 9; // Optional +constexpr int kFwCellToForgetWeightsTensor = 10; // Optional +constexpr int kFwCellToOutputWeightsTensor = 11; // Optional + +// Gates bias tensors of size {n_cell} +constexpr int kFwInputGateBiasTensor = 12; // Optional +constexpr int kFwForgetGateBiasTensor = 13; +constexpr int kFwCellGateBiasTensor = 14; +constexpr int kFwOutputGateBiasTensor = 15; + +// Projection weight tensor of size {n_output, n_cell} +constexpr int kFwProjectionWeightsTensor = 16; // Optional +// Projection bias tensor of size {n_output} +constexpr int kFwProjectionBiasTensor = 17; // Optional + +// Backward LSTM cell tensors. +// Input weight tensors of size: {n_cell, n_input} +constexpr int kBwInputToInputWeightsTensor = 18; // Optional +constexpr int kBwInputToForgetWeightsTensor = 19; +constexpr int kBwInputToCellWeightsTensor = 20; +constexpr int kBwInputToOutputWeightsTensor = 21; + +// Recurrent weight tensors of size {n_cell, n_output} +constexpr int kBwRecurrentToInputWeightsTensor = 22; // Optional +constexpr int kBwRecurrentToForgetWeightsTensor = 23; +constexpr int kBwRecurrentToCellWeightsTensor = 24; +constexpr int kBwRecurrentToOutputWeightsTensor = 25; + +// Peephole weights tensors of size {n_cell}, representing a diagonal matrix. +constexpr int kBwCellToInputWeightsTensor = 26; // Optional +constexpr int kBwCellToForgetWeightsTensor = 27; // Optional +constexpr int kBwCellToOutputWeightsTensor = 28; // Optional + +// Gates bias tensors of size {n_cell} +constexpr int kBwInputGateBiasTensor = 29; // Optional +constexpr int kBwForgetGateBiasTensor = 30; +constexpr int kBwCellGateBiasTensor = 31; +constexpr int kBwOutputGateBiasTensor = 32; + +// Projection weight tensor of size {n_output, n_cell} +constexpr int kBwProjectionWeightsTensor = 33; // Optional +// Projection bias tensor of size {n_output} +constexpr int kBwProjectionBiasTensor = 34; // Optional + +// Output tensors. +constexpr int kFwScratchBufferTensor = 0; +constexpr int kFwOutputStateTensor = 1; +constexpr int kFwCellStateTensor = 2; +constexpr int kFwOutputTensor = 3; + +constexpr int kBwScratchBufferTensor = 4; +constexpr int kBwOutputStateTensor = 5; +constexpr int kBwCellStateTensor = 6; +constexpr int kBwOutputTensor = 7; + +// Check that input tensor dimensions matches with each other. +TfLiteStatus CheckLstmTensorDimensions( + TfLiteContext* context, TfLiteNode* node, int n_input, int n_output, + int n_cell, int input_to_input_weights_tensor, + int input_to_forget_weights_tensor, int input_to_cell_weights_tensor, + int input_to_output_weights_tensor, int recurrent_to_input_weights_tensor, + int recurrent_to_forget_weights_tensor, + int recurrent_to_cell_weights_tensor, + int recurrent_to_output_weights_tensor, int cell_to_input_weights_tensor, + int cell_to_forget_weights_tensor, int cell_to_output_weights_tensor, + int input_gate_bias_tensor, int forget_gate_bias_tensor, + int cell_gate_bias_tensor, int output_gate_bias_tensor, + int projection_weights_tensor, int projection_bias_tensor) { + 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, input_to_input_weights_tensor); + 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, input_to_forget_weights_tensor); + 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, input_to_cell_weights_tensor); + 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, recurrent_to_input_weights_tensor); + 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, recurrent_to_forget_weights_tensor); + 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, recurrent_to_cell_weights_tensor); + 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, cell_to_input_weights_tensor); + 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, cell_to_forget_weights_tensor); + 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, cell_to_output_weights_tensor); + 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, input_gate_bias_tensor); + 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, forget_gate_bias_tensor); + 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, cell_gate_bias_tensor); + 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, output_gate_bias_tensor); + 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, projection_weights_tensor); + 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, projection_bias_tensor); + 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; +} + +TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context, + TfLiteNode* node, int n_input, + int n_output, int n_cell) { + CheckLstmTensorDimensions( + context, node, n_input, n_output, n_cell, kFwInputToInputWeightsTensor, + kFwInputToForgetWeightsTensor, kFwInputToCellWeightsTensor, + kFwInputToOutputWeightsTensor, kFwRecurrentToInputWeightsTensor, + kFwRecurrentToForgetWeightsTensor, kFwRecurrentToCellWeightsTensor, + kFwRecurrentToOutputWeightsTensor, kFwCellToInputWeightsTensor, + kFwCellToForgetWeightsTensor, kFwCellToOutputWeightsTensor, + kFwInputGateBiasTensor, kFwForgetGateBiasTensor, kFwCellGateBiasTensor, + kFwOutputGateBiasTensor, kFwProjectionWeightsTensor, + kFwProjectionBiasTensor); + + CheckLstmTensorDimensions( + context, node, n_input, n_output, n_cell, kBwInputToInputWeightsTensor, + kBwInputToForgetWeightsTensor, kBwInputToCellWeightsTensor, + kBwInputToOutputWeightsTensor, kBwRecurrentToInputWeightsTensor, + kBwRecurrentToForgetWeightsTensor, kBwRecurrentToCellWeightsTensor, + kBwRecurrentToOutputWeightsTensor, kBwCellToInputWeightsTensor, + kBwCellToForgetWeightsTensor, kBwCellToOutputWeightsTensor, + kBwInputGateBiasTensor, kBwForgetGateBiasTensor, kBwCellGateBiasTensor, + kBwOutputGateBiasTensor, kBwProjectionWeightsTensor, + kBwProjectionBiasTensor); + + // Check if Forward and Backward tensors match along required dimensions. + 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, 35); + TF_LITE_ENSURE_EQ(context, node->outputs->size, 8); + + // 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* fw_input_to_output_weights = + GetInput(context, node, kFwInputToOutputWeightsTensor); + const int n_fw_cell = fw_input_to_output_weights->dims->data[0]; + TF_LITE_ENSURE_EQ(context, fw_input_to_output_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, fw_input_to_output_weights->dims->data[1], + n_input); + + TfLiteTensor* fw_recurrent_to_output_weights = + GetInput(context, node, kFwRecurrentToOutputWeightsTensor); + TF_LITE_ENSURE_EQ(context, fw_recurrent_to_output_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, fw_recurrent_to_output_weights->dims->data[0], + n_fw_cell); + const int n_fw_output = fw_recurrent_to_output_weights->dims->data[1]; + + // Check that input tensor dimensions matches with each other. + CheckInputTensorDimensions(context, node, n_input, n_fw_output, n_fw_cell); + + // Get the pointer to output, state and scratch buffer tensors. + TfLiteTensor* fw_output = GetOutput(context, node, kFwOutputTensor); + TfLiteTensor* fw_output_state = + GetOutput(context, node, kFwOutputStateTensor); + TfLiteTensor* fw_cell_state = GetOutput(context, node, kFwCellStateTensor); + // TODO(ghodrat): Modify this as soon as we have a finalized method for + // scratch buffers. + TfLiteTensor* fw_scratch_buffer = + GetOutput(context, node, kFwScratchBufferTensor); + + // Resize the output and output_state tensors. + TfLiteIntArray* fw_output_size = TfLiteIntArrayCreate(3); + fw_output_size->data[0] = max_time; + fw_output_size->data[1] = n_batch; + fw_output_size->data[2] = n_fw_output; + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, fw_output, fw_output_size)); + + TfLiteIntArray* fw_output_state_size = TfLiteIntArrayCreate(2); + fw_output_state_size->data[0] = n_batch; + fw_output_state_size->data[1] = n_fw_output; + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, fw_output_state, + fw_output_state_size)); + + // Resize the scratch buffer tensor. + TfLiteIntArray* fw_cell_size = TfLiteIntArrayCreate(2); + fw_cell_size->data[0] = n_batch; + fw_cell_size->data[1] = n_fw_cell; + TF_LITE_ENSURE_OK( + context, context->ResizeTensor(context, fw_cell_state, fw_cell_size)); + + // Mark state tensors as persistent tensors. + fw_output_state->allocation_type = kTfLiteArenaRwPersistent; + fw_cell_state->allocation_type = kTfLiteArenaRwPersistent; + + TfLiteTensor* fw_input_to_input_weights = + GetOptionalInputTensor(context, node, kFwInputToInputWeightsTensor); + const bool fw_use_cifg = (fw_input_to_input_weights == nullptr); + TfLiteIntArray* fw_scratch_buffer_size = TfLiteIntArrayCreate(2); + fw_scratch_buffer_size->data[0] = n_batch; + if (fw_use_cifg) { + // Reserving space for Cell, Forget, Output gates + fw_scratch_buffer_size->data[1] = n_fw_cell * 3; + } else { + // Reserving space for Input, Cell, Forget, Output gates + fw_scratch_buffer_size->data[1] = n_fw_cell * 4; + } + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, fw_scratch_buffer, + fw_scratch_buffer_size)); + // Same for the backward cell. + TfLiteTensor* bw_input_to_output_weights = + GetInput(context, node, kBwInputToOutputWeightsTensor); + const int n_bw_cell = bw_input_to_output_weights->dims->data[0]; + TF_LITE_ENSURE_EQ(context, bw_input_to_output_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, bw_input_to_output_weights->dims->data[1], + n_input); + + TfLiteTensor* bw_recurrent_to_output_weights = + GetInput(context, node, kBwRecurrentToOutputWeightsTensor); + TF_LITE_ENSURE_EQ(context, bw_recurrent_to_output_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, bw_recurrent_to_output_weights->dims->data[0], + n_bw_cell); + const int n_bw_output = bw_recurrent_to_output_weights->dims->data[1]; + + // Check that input tensor dimensions matches with each other. + CheckInputTensorDimensions(context, node, n_input, n_bw_output, n_bw_cell); + + // Get the pointer to output, state and scratch buffer tensors. + TfLiteTensor* bw_output = GetOutput(context, node, kBwOutputTensor); + TfLiteTensor* bw_output_state = + GetOutput(context, node, kBwOutputStateTensor); + TfLiteTensor* bw_cell_state = GetOutput(context, node, kBwCellStateTensor); + // TODO(ghodrat): Modify this as soon as we have a finalized method for + // scratch buffers. + TfLiteTensor* bw_scratch_buffer = + GetOutput(context, node, kBwScratchBufferTensor); + + // Resize the output and output_state tensors. + TfLiteIntArray* bw_output_size = TfLiteIntArrayCreate(3); + bw_output_size->data[0] = max_time; + bw_output_size->data[1] = n_batch; + bw_output_size->data[2] = n_bw_output; + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, bw_output, bw_output_size)); + + TfLiteIntArray* bw_output_state_size = TfLiteIntArrayCreate(2); + bw_output_state_size->data[0] = n_batch; + bw_output_state_size->data[1] = n_bw_output; + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, bw_output_state, + bw_output_state_size)); + + // Resize the scratch buffer tensor. + TfLiteIntArray* bw_cell_size = TfLiteIntArrayCreate(2); + bw_cell_size->data[0] = n_batch; + bw_cell_size->data[1] = n_bw_cell; + TF_LITE_ENSURE_OK( + context, context->ResizeTensor(context, bw_cell_state, bw_cell_size)); + + // Mark state tensors as persistent tensors. + bw_output_state->allocation_type = kTfLiteArenaRwPersistent; + bw_cell_state->allocation_type = kTfLiteArenaRwPersistent; + + TfLiteTensor* bw_input_to_input_weights = + GetOptionalInputTensor(context, node, kBwInputToInputWeightsTensor); + const bool bw_use_cifg = (bw_input_to_input_weights == nullptr); + TfLiteIntArray* bw_scratch_buffer_size = TfLiteIntArrayCreate(2); + bw_scratch_buffer_size->data[0] = n_batch; + if (bw_use_cifg) { + // Reserving space for Cell, Forget, Output gates + bw_scratch_buffer_size->data[1] = n_bw_cell * 3; + } else { + // Reserving space for Input, Cell, Forget, Output gates + bw_scratch_buffer_size->data[1] = n_bw_cell * 4; + } + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, bw_scratch_buffer, + bw_scratch_buffer_size)); + return kTfLiteOk; +} + +// The LSTM Op engine. +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->builtin_data); + + // Input tensor. + TfLiteTensor* input = GetInput(context, node, kInputTensor); + const int max_time = input->dims->data[0]; + const int n_batch = input->dims->data[1]; + const int n_input = input->dims->data[2]; + + // Tensors for the forward cell. + TfLiteTensor* fw_input_to_input_weights = + GetOptionalInputTensor(context, node, kFwInputToInputWeightsTensor); + TfLiteTensor* fw_input_to_forget_weights = + GetInput(context, node, kFwInputToForgetWeightsTensor); + TfLiteTensor* fw_input_to_cell_weights = + GetInput(context, node, kFwInputToCellWeightsTensor); + TfLiteTensor* fw_input_to_output_weights = + GetInput(context, node, kFwInputToOutputWeightsTensor); + + TfLiteTensor* fw_recurrent_to_input_weights = + GetOptionalInputTensor(context, node, kFwRecurrentToInputWeightsTensor); + TfLiteTensor* fw_recurrent_to_forget_weights = + GetInput(context, node, kFwRecurrentToForgetWeightsTensor); + TfLiteTensor* fw_recurrent_to_cell_weights = + GetInput(context, node, kFwRecurrentToCellWeightsTensor); + TfLiteTensor* fw_recurrent_to_output_weights = + GetInput(context, node, kFwRecurrentToOutputWeightsTensor); + + TfLiteTensor* fw_cell_to_input_weights = + GetOptionalInputTensor(context, node, kFwCellToInputWeightsTensor); + TfLiteTensor* fw_cell_to_forget_weights = + GetOptionalInputTensor(context, node, kFwCellToForgetWeightsTensor); + TfLiteTensor* fw_cell_to_output_weights = + GetOptionalInputTensor(context, node, kFwCellToOutputWeightsTensor); + + TfLiteTensor* fw_input_gate_bias = + GetOptionalInputTensor(context, node, kFwInputGateBiasTensor); + TfLiteTensor* fw_forget_gate_bias = + GetInput(context, node, kFwForgetGateBiasTensor); + TfLiteTensor* fw_cell_bias = GetInput(context, node, kFwCellGateBiasTensor); + TfLiteTensor* fw_output_gate_bias = + GetInput(context, node, kFwOutputGateBiasTensor); + + TfLiteTensor* fw_projection_weights = + GetOptionalInputTensor(context, node, kFwProjectionWeightsTensor); + TfLiteTensor* fw_projection_bias = + GetOptionalInputTensor(context, node, kFwProjectionBiasTensor); + + TfLiteTensor* fw_output_state = + GetOutput(context, node, kFwOutputStateTensor); + TfLiteTensor* fw_cell_state = GetOutput(context, node, kFwCellStateTensor); + TfLiteTensor* fw_output = GetOutput(context, node, kFwOutputTensor); + + // Tensors for the backward cell. + TfLiteTensor* bw_input_to_input_weights = + GetOptionalInputTensor(context, node, kBwInputToInputWeightsTensor); + TfLiteTensor* bw_input_to_forget_weights = + GetInput(context, node, kBwInputToForgetWeightsTensor); + TfLiteTensor* bw_input_to_cell_weights = + GetInput(context, node, kBwInputToCellWeightsTensor); + TfLiteTensor* bw_input_to_output_weights = + GetInput(context, node, kBwInputToOutputWeightsTensor); + + TfLiteTensor* bw_recurrent_to_input_weights = + GetOptionalInputTensor(context, node, kBwRecurrentToInputWeightsTensor); + TfLiteTensor* bw_recurrent_to_forget_weights = + GetInput(context, node, kBwRecurrentToForgetWeightsTensor); + TfLiteTensor* bw_recurrent_to_cell_weights = + GetInput(context, node, kBwRecurrentToCellWeightsTensor); + TfLiteTensor* bw_recurrent_to_output_weights = + GetInput(context, node, kBwRecurrentToOutputWeightsTensor); + + TfLiteTensor* bw_cell_to_input_weights = + GetOptionalInputTensor(context, node, kBwCellToInputWeightsTensor); + TfLiteTensor* bw_cell_to_forget_weights = + GetOptionalInputTensor(context, node, kBwCellToForgetWeightsTensor); + TfLiteTensor* bw_cell_to_output_weights = + GetOptionalInputTensor(context, node, kBwCellToOutputWeightsTensor); + + TfLiteTensor* bw_input_gate_bias = + GetOptionalInputTensor(context, node, kBwInputGateBiasTensor); + TfLiteTensor* bw_forget_gate_bias = + GetInput(context, node, kBwForgetGateBiasTensor); + TfLiteTensor* bw_cell_bias = GetInput(context, node, kBwCellGateBiasTensor); + TfLiteTensor* bw_output_gate_bias = + GetInput(context, node, kBwOutputGateBiasTensor); + + TfLiteTensor* bw_projection_weights = + GetOptionalInputTensor(context, node, kBwProjectionWeightsTensor); + TfLiteTensor* bw_projection_bias = + GetOptionalInputTensor(context, node, kBwProjectionBiasTensor); + + TfLiteTensor* bw_output_state = + GetOutput(context, node, kBwOutputStateTensor); + TfLiteTensor* bw_cell_state = GetOutput(context, node, kBwCellStateTensor); + TfLiteTensor* bw_output = GetOutput(context, node, kBwOutputTensor); + + // n_cell and n_output will be the same size when there is no projection. + const int n_fw_cell = fw_input_to_output_weights->dims->data[0]; + const int n_fw_output = fw_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 fw_use_cifg = (fw_input_to_input_weights == nullptr); + const bool fw_use_peephole = (fw_cell_to_output_weights != nullptr); + + // Index the scratch buffers pointers to the global scratch buffer. + TfLiteTensor* fw_scratch_buffer = + GetOutput(context, node, kFwScratchBufferTensor); + float* fw_input_gate_scratch = nullptr; + float* fw_cell_scratch = nullptr; + float* fw_forget_gate_scratch = nullptr; + float* fw_output_gate_scratch = nullptr; + if (fw_use_cifg) { + fw_cell_scratch = fw_scratch_buffer->data.f; + fw_forget_gate_scratch = fw_scratch_buffer->data.f + n_fw_cell * n_batch; + fw_output_gate_scratch = + fw_scratch_buffer->data.f + 2 * n_fw_cell * n_batch; + } else { + fw_input_gate_scratch = fw_scratch_buffer->data.f; + fw_cell_scratch = fw_scratch_buffer->data.f + n_fw_cell * n_batch; + fw_forget_gate_scratch = + fw_scratch_buffer->data.f + 2 * n_fw_cell * n_batch; + fw_output_gate_scratch = + fw_scratch_buffer->data.f + 3 * n_fw_cell * n_batch; + } + + // Check optional tensors, the respective pointers can be null. + const float* fw_input_to_input_weights_ptr = + (fw_use_cifg) ? nullptr : fw_input_to_input_weights->data.f; + const float* fw_recurrent_to_input_weights_ptr = + (fw_use_cifg) ? nullptr : fw_recurrent_to_input_weights->data.f; + const float* fw_input_gate_bias_ptr = + (fw_use_cifg) ? nullptr : fw_input_gate_bias->data.f; + const float* fw_cell_to_input_weights_ptr = + (fw_use_peephole && !fw_use_cifg) ? fw_cell_to_input_weights->data.f + : nullptr; + const float* fw_cell_to_forget_weights_ptr = + (fw_use_peephole) ? fw_cell_to_forget_weights->data.f : nullptr; + const float* fw_cell_to_output_weights_ptr = + (fw_use_peephole) ? fw_cell_to_output_weights->data.f : nullptr; + const float* fw_projection_weights_ptr = (fw_projection_weights == nullptr) + ? nullptr + : fw_projection_weights->data.f; + const float* fw_projection_bias_ptr = + (fw_projection_bias == nullptr) ? nullptr : fw_projection_bias->data.f; + + // Loop through the sequence. + for (int t = 0; t < max_time; t++) { + const float* input_ptr_batch = input->data.f + t * n_batch * n_input; + float* output_ptr_time = fw_output->data.f + t * n_batch * n_fw_output; + + kernel_utils::LstmStep( + input_ptr_batch, fw_input_to_input_weights_ptr, + fw_input_to_forget_weights->data.f, fw_input_to_cell_weights->data.f, + fw_input_to_output_weights->data.f, fw_recurrent_to_input_weights_ptr, + fw_recurrent_to_forget_weights->data.f, + fw_recurrent_to_cell_weights->data.f, + fw_recurrent_to_output_weights->data.f, fw_cell_to_input_weights_ptr, + fw_cell_to_forget_weights_ptr, fw_cell_to_output_weights_ptr, + fw_input_gate_bias_ptr, fw_forget_gate_bias->data.f, + fw_cell_bias->data.f, fw_output_gate_bias->data.f, + fw_projection_weights_ptr, fw_projection_bias_ptr, params, n_batch, + n_fw_cell, n_input, n_fw_output, fw_output_state->data.f, + fw_cell_state->data.f, fw_input_gate_scratch, fw_forget_gate_scratch, + fw_cell_scratch, fw_output_gate_scratch, output_ptr_time); + } + + // n_cell and n_output will be the same size when there is no projection. + const int n_bw_cell = bw_input_to_output_weights->dims->data[0]; + const int n_bw_output = bw_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 bw_use_cifg = (bw_input_to_input_weights == nullptr); + const bool bw_use_peephole = (bw_cell_to_output_weights != nullptr); + + // Index the scratch buffers pointers to the global scratch buffer. + TfLiteTensor* bw_scratch_buffer = + GetOutput(context, node, kBwScratchBufferTensor); + float* bw_input_gate_scratch = nullptr; + float* bw_cell_scratch = nullptr; + float* bw_forget_gate_scratch = nullptr; + float* bw_output_gate_scratch = nullptr; + if (bw_use_cifg) { + bw_cell_scratch = bw_scratch_buffer->data.f; + bw_forget_gate_scratch = bw_scratch_buffer->data.f + n_bw_cell * n_batch; + bw_output_gate_scratch = + bw_scratch_buffer->data.f + 2 * n_bw_cell * n_batch; + } else { + bw_input_gate_scratch = bw_scratch_buffer->data.f; + bw_cell_scratch = bw_scratch_buffer->data.f + n_bw_cell * n_batch; + bw_forget_gate_scratch = + bw_scratch_buffer->data.f + 2 * n_bw_cell * n_batch; + bw_output_gate_scratch = + bw_scratch_buffer->data.f + 3 * n_bw_cell * n_batch; + } + + // Check optional tensors, the respective pointers can be null. + const float* bw_input_to_input_weights_ptr = + (bw_use_cifg) ? nullptr : bw_input_to_input_weights->data.f; + const float* bw_recurrent_to_input_weights_ptr = + (bw_use_cifg) ? nullptr : bw_recurrent_to_input_weights->data.f; + const float* bw_input_gate_bias_ptr = + (bw_use_cifg) ? nullptr : bw_input_gate_bias->data.f; + const float* bw_cell_to_input_weights_ptr = + (bw_use_peephole && !bw_use_cifg) ? bw_cell_to_input_weights->data.f + : nullptr; + const float* bw_cell_to_forget_weights_ptr = + (bw_use_peephole) ? bw_cell_to_forget_weights->data.f : nullptr; + const float* bw_cell_to_output_weights_ptr = + (bw_use_peephole) ? bw_cell_to_output_weights->data.f : nullptr; + const float* bw_projection_weights_ptr = (bw_projection_weights == nullptr) + ? nullptr + : bw_projection_weights->data.f; + const float* bw_projection_bias_ptr = + (bw_projection_bias == nullptr) ? nullptr : bw_projection_bias->data.f; + + // Loop through the sequence backwards. + for (int t = max_time - 1; t >= 0; t--) { + const float* input_ptr_batch = input->data.f + t * n_batch * n_input; + float* output_ptr_time = bw_output->data.f + t * n_batch * n_bw_output; + + kernel_utils::LstmStep( + input_ptr_batch, bw_input_to_input_weights_ptr, + bw_input_to_forget_weights->data.f, bw_input_to_cell_weights->data.f, + bw_input_to_output_weights->data.f, bw_recurrent_to_input_weights_ptr, + bw_recurrent_to_forget_weights->data.f, + bw_recurrent_to_cell_weights->data.f, + bw_recurrent_to_output_weights->data.f, bw_cell_to_input_weights_ptr, + bw_cell_to_forget_weights_ptr, bw_cell_to_output_weights_ptr, + bw_input_gate_bias_ptr, bw_forget_gate_bias->data.f, + bw_cell_bias->data.f, bw_output_gate_bias->data.f, + bw_projection_weights_ptr, bw_projection_bias_ptr, params, n_batch, + n_bw_cell, n_input, n_bw_output, bw_output_state->data.f, + bw_cell_state->data.f, bw_input_gate_scratch, bw_forget_gate_scratch, + bw_cell_scratch, bw_output_gate_scratch, output_ptr_time); + } + + // Backward step. + return kTfLiteOk; +} + +} // namespace bidirectional_sequence_lstm + +TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_LSTM() { + static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, + bidirectional_sequence_lstm::Prepare, + bidirectional_sequence_lstm::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..cca857bac0633ded01d40273d2e9e8dde488d61e --- /dev/null +++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc @@ -0,0 +1,1411 @@ +/* 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 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 BidirectionalLSTMOpModel : public SingleOpModel { + public: + BidirectionalLSTMOpModel(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_fw_cell_(n_cell), + n_bw_cell_(n_cell), + n_fw_output_(n_output), + n_bw_output_(n_output), + sequence_length_(sequence_length) { + input_ = AddInput(TensorType_FLOAT32); + + if (use_cifg) { + fw_input_to_input_weights_ = AddNullInput(); + } else { + fw_input_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + + fw_input_to_forget_weights_ = AddInput(TensorType_FLOAT32); + fw_input_to_cell_weights_ = AddInput(TensorType_FLOAT32); + fw_input_to_output_weights_ = AddInput(TensorType_FLOAT32); + + if (use_cifg) { + fw_recurrent_to_input_weights_ = AddNullInput(); + } else { + fw_recurrent_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + + fw_recurrent_to_forget_weights_ = AddInput(TensorType_FLOAT32); + fw_recurrent_to_cell_weights_ = AddInput(TensorType_FLOAT32); + fw_recurrent_to_output_weights_ = AddInput(TensorType_FLOAT32); + + if (use_peephole) { + if (use_cifg) { + fw_cell_to_input_weights_ = AddNullInput(); + } else { + fw_cell_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + fw_cell_to_forget_weights_ = AddInput(TensorType_FLOAT32); + fw_cell_to_output_weights_ = AddInput(TensorType_FLOAT32); + } else { + fw_cell_to_input_weights_ = AddNullInput(); + fw_cell_to_forget_weights_ = AddNullInput(); + fw_cell_to_output_weights_ = AddNullInput(); + } + + if (use_cifg) { + fw_input_gate_bias_ = AddNullInput(); + } else { + fw_input_gate_bias_ = AddInput(TensorType_FLOAT32); + } + fw_forget_gate_bias_ = AddInput(TensorType_FLOAT32); + fw_cell_bias_ = AddInput(TensorType_FLOAT32); + fw_output_gate_bias_ = AddInput(TensorType_FLOAT32); + + if (use_projection_weights) { + fw_projection_weights_ = AddInput(TensorType_FLOAT32); + if (use_projection_bias) { + fw_projection_bias_ = AddInput(TensorType_FLOAT32); + } else { + fw_projection_bias_ = AddNullInput(); + } + } else { + fw_projection_weights_ = AddNullInput(); + fw_projection_bias_ = AddNullInput(); + } + + fw_scratch_buffer_ = AddOutput(TensorType_FLOAT32); + // TODO(ghodrat): Modify these states when we have a permanent solution for + // persistent buffer. + fw_output_state_ = AddOutput(TensorType_FLOAT32); + fw_cell_state_ = AddOutput(TensorType_FLOAT32); + fw_output_ = AddOutput(TensorType_FLOAT32); + + if (use_cifg) { + bw_input_to_input_weights_ = AddNullInput(); + } else { + bw_input_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + + bw_input_to_forget_weights_ = AddInput(TensorType_FLOAT32); + bw_input_to_cell_weights_ = AddInput(TensorType_FLOAT32); + bw_input_to_output_weights_ = AddInput(TensorType_FLOAT32); + + if (use_cifg) { + bw_recurrent_to_input_weights_ = AddNullInput(); + } else { + bw_recurrent_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + + bw_recurrent_to_forget_weights_ = AddInput(TensorType_FLOAT32); + bw_recurrent_to_cell_weights_ = AddInput(TensorType_FLOAT32); + bw_recurrent_to_output_weights_ = AddInput(TensorType_FLOAT32); + + if (use_peephole) { + if (use_cifg) { + bw_cell_to_input_weights_ = AddNullInput(); + } else { + bw_cell_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + bw_cell_to_forget_weights_ = AddInput(TensorType_FLOAT32); + bw_cell_to_output_weights_ = AddInput(TensorType_FLOAT32); + } else { + bw_cell_to_input_weights_ = AddNullInput(); + bw_cell_to_forget_weights_ = AddNullInput(); + bw_cell_to_output_weights_ = AddNullInput(); + } + + if (use_cifg) { + bw_input_gate_bias_ = AddNullInput(); + } else { + bw_input_gate_bias_ = AddInput(TensorType_FLOAT32); + } + bw_forget_gate_bias_ = AddInput(TensorType_FLOAT32); + bw_cell_bias_ = AddInput(TensorType_FLOAT32); + bw_output_gate_bias_ = AddInput(TensorType_FLOAT32); + + if (use_projection_weights) { + bw_projection_weights_ = AddInput(TensorType_FLOAT32); + if (use_projection_bias) { + bw_projection_bias_ = AddInput(TensorType_FLOAT32); + } else { + bw_projection_bias_ = AddNullInput(); + } + } else { + bw_projection_weights_ = AddNullInput(); + bw_projection_bias_ = AddNullInput(); + } + + bw_scratch_buffer_ = AddOutput(TensorType_FLOAT32); + // TODO(ghodrat): Modify these states when we have a permanent solution for + // persistent buffer. + bw_output_state_ = AddOutput(TensorType_FLOAT32); + bw_cell_state_ = AddOutput(TensorType_FLOAT32); + bw_output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, + BuiltinOptions_LSTMOptions, + CreateLSTMOptions(builder_, ActivationFunctionType_TANH, + cell_clip, proj_clip) + .Union()); + BuildInterpreter(input_shapes); + } + + // Set weights in forward and backward cells to be the same. + void SetInputToInputWeights(std::initializer_list f) { + PopulateTensor(fw_input_to_input_weights_, f); + PopulateTensor(bw_input_to_input_weights_, f); + } + + void SetInputToForgetWeights(std::initializer_list f) { + PopulateTensor(fw_input_to_forget_weights_, f); + PopulateTensor(bw_input_to_forget_weights_, f); + } + + void SetInputToCellWeights(std::initializer_list f) { + PopulateTensor(fw_input_to_cell_weights_, f); + PopulateTensor(bw_input_to_cell_weights_, f); + } + + void SetInputToOutputWeights(std::initializer_list f) { + PopulateTensor(fw_input_to_output_weights_, f); + PopulateTensor(bw_input_to_output_weights_, f); + } + + void SetRecurrentToInputWeights(std::initializer_list f) { + PopulateTensor(fw_recurrent_to_input_weights_, f); + PopulateTensor(bw_recurrent_to_input_weights_, f); + } + + void SetRecurrentToForgetWeights(std::initializer_list f) { + PopulateTensor(fw_recurrent_to_forget_weights_, f); + PopulateTensor(bw_recurrent_to_forget_weights_, f); + } + + void SetRecurrentToCellWeights(std::initializer_list f) { + PopulateTensor(fw_recurrent_to_cell_weights_, f); + PopulateTensor(bw_recurrent_to_cell_weights_, f); + } + + void SetRecurrentToOutputWeights(std::initializer_list f) { + PopulateTensor(fw_recurrent_to_output_weights_, f); + PopulateTensor(bw_recurrent_to_output_weights_, f); + } + + void SetCellToInputWeights(std::initializer_list f) { + PopulateTensor(fw_cell_to_input_weights_, f); + PopulateTensor(bw_cell_to_input_weights_, f); + } + + void SetCellToForgetWeights(std::initializer_list f) { + PopulateTensor(fw_cell_to_forget_weights_, f); + PopulateTensor(bw_cell_to_forget_weights_, f); + } + + void SetCellToOutputWeights(std::initializer_list f) { + PopulateTensor(fw_cell_to_output_weights_, f); + PopulateTensor(bw_cell_to_output_weights_, f); + } + + void SetInputGateBias(std::initializer_list f) { + PopulateTensor(fw_input_gate_bias_, f); + PopulateTensor(bw_input_gate_bias_, f); + } + + void SetForgetGateBias(std::initializer_list f) { + PopulateTensor(fw_forget_gate_bias_, f); + PopulateTensor(bw_forget_gate_bias_, f); + } + + void SetCellBias(std::initializer_list f) { + PopulateTensor(fw_cell_bias_, f); + PopulateTensor(bw_cell_bias_, f); + } + + void SetOutputGateBias(std::initializer_list f) { + PopulateTensor(fw_output_gate_bias_, f); + PopulateTensor(bw_output_gate_bias_, f); + } + + void SetProjectionWeights(std::initializer_list f) { + PopulateTensor(fw_projection_weights_, f); + PopulateTensor(bw_projection_weights_, f); + } + + void SetProjectionBias(std::initializer_list f) { + PopulateTensor(fw_projection_bias_, f); + PopulateTensor(bw_projection_bias_, f); + } + + void ResetFwOutputAndCellStates() { + const int zero_buffer_size = n_fw_cell_ * n_batch_; + std::unique_ptr zero_buffer(new float[zero_buffer_size]); + memset(zero_buffer.get(), 0, zero_buffer_size * sizeof(float)); + PopulateTensor(fw_output_state_, 0, zero_buffer.get(), + zero_buffer.get() + zero_buffer_size); + PopulateTensor(fw_cell_state_, 0, zero_buffer.get(), + zero_buffer.get() + zero_buffer_size); + } + + void ResetBwOutputAndCellStates() { + const int zero_buffer_size = n_bw_cell_ * n_batch_; + std::unique_ptr zero_buffer(new float[zero_buffer_size]); + memset(zero_buffer.get(), 0, zero_buffer_size * sizeof(float)); + PopulateTensor(bw_output_state_, 0, zero_buffer.get(), + zero_buffer.get() + zero_buffer_size); + PopulateTensor(bw_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 GetFwOutput() { return ExtractVector(fw_output_); } + std::vector GetBwOutput() { return ExtractVector(bw_output_); } + + int num_inputs() { return n_input_; } + int num_fw_outputs() { return n_fw_output_; } + int num_bw_outputs() { return n_bw_output_; } + int num_fw_cells() { return n_fw_cell_; } + int num_bw_cells() { return n_bw_cell_; } + int num_batches() { return n_batch_; } + int sequence_length() { return sequence_length_; } + + private: + int input_; + int fw_input_to_input_weights_; + int fw_input_to_forget_weights_; + int fw_input_to_cell_weights_; + int fw_input_to_output_weights_; + + int fw_recurrent_to_input_weights_; + int fw_recurrent_to_forget_weights_; + int fw_recurrent_to_cell_weights_; + int fw_recurrent_to_output_weights_; + + int fw_cell_to_input_weights_; + int fw_cell_to_forget_weights_; + int fw_cell_to_output_weights_; + + int fw_input_gate_bias_; + int fw_forget_gate_bias_; + int fw_cell_bias_; + int fw_output_gate_bias_; + + int fw_projection_weights_; + int fw_projection_bias_; + + int bw_input_to_input_weights_; + int bw_input_to_forget_weights_; + int bw_input_to_cell_weights_; + int bw_input_to_output_weights_; + + int bw_recurrent_to_input_weights_; + int bw_recurrent_to_forget_weights_; + int bw_recurrent_to_cell_weights_; + int bw_recurrent_to_output_weights_; + + int bw_cell_to_input_weights_; + int bw_cell_to_forget_weights_; + int bw_cell_to_output_weights_; + + int bw_input_gate_bias_; + int bw_forget_gate_bias_; + int bw_cell_bias_; + int bw_output_gate_bias_; + + int bw_projection_weights_; + int bw_projection_bias_; + + int fw_output_; + int fw_output_state_; + int fw_cell_state_; + int fw_scratch_buffer_; + + int bw_output_; + int bw_output_state_; + int bw_cell_state_; + int bw_scratch_buffer_; + + int n_batch_; + int n_input_; + int n_fw_cell_; + int n_bw_cell_; + int n_fw_output_; + int n_bw_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; + + BidirectionalLSTMOpModel 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 + + // Forward cell + {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 + + // Backward cell + {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_fw_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}; + static float lstm_bw_golden_output[] = { + -0.0806187, 0.139077, 0.400476, -0.197842, + -0.0332076, 0.123838, 0.309777, -0.17621, + -0.0490733, 0.0739237, 0.067706, -0.0208124}; + + // Resetting cell_state and output_state + lstm.ResetFwOutputAndCellStates(); + lstm.ResetBwOutputAndCellStates(); + + 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* fw_golden_start = lstm_fw_golden_output; + float* fw_golden_end = + fw_golden_start + lstm.num_fw_outputs() * lstm.sequence_length(); + std::vector fw_expected; + fw_expected.insert(fw_expected.end(), fw_golden_start, fw_golden_end); + EXPECT_THAT(lstm.GetFwOutput(), + ElementsAreArray(ArrayFloatNear(fw_expected))); + + float* bw_golden_start = lstm_bw_golden_output; + float* bw_golden_end = + bw_golden_start + lstm.num_bw_outputs() * lstm.sequence_length(); + std::vector bw_expected; + bw_expected.insert(bw_expected.end(), bw_golden_start, bw_golden_end); + EXPECT_THAT(lstm.GetBwOutput(), + ElementsAreArray(ArrayFloatNear(bw_expected))); + + // Check reversed inputs. + static float lstm_input_reversed[] = {1., 1., 3., 4., 2., 3.}; + + // Resetting cell_state and output_state + lstm.ResetFwOutputAndCellStates(); + lstm.ResetBwOutputAndCellStates(); + + batch0_start = lstm_input_reversed; + batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); + + lstm.SetInput(0, batch0_start, batch0_end); + + lstm.Invoke(); + + fw_expected.clear(); + for (int s = 0; s < lstm.sequence_length(); s++) { + fw_golden_start = lstm_fw_golden_output + s * lstm.num_fw_outputs(); + fw_golden_end = fw_golden_start + lstm.num_fw_outputs(); + fw_expected.insert(fw_expected.begin(), fw_golden_start, fw_golden_end); + } + EXPECT_THAT(lstm.GetBwOutput(), + ElementsAreArray(ArrayFloatNear(fw_expected))); + + bw_expected.clear(); + for (int s = 0; s < lstm.sequence_length(); s++) { + bw_golden_start = lstm_bw_golden_output + s * lstm.num_bw_outputs(); + bw_golden_end = bw_golden_start + lstm.num_bw_outputs(); + bw_expected.insert(bw_expected.begin(), bw_golden_start, bw_golden_end); + } + EXPECT_THAT(lstm.GetFwOutput(), + ElementsAreArray(ArrayFloatNear(bw_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; + + BidirectionalLSTMOpModel 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 + + {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_fw_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}; + static float lstm_bw_golden_output[] = { + -0.401685, -0.0232794, 0.288642, -0.123074, -0.42915, -0.00871577, + 0.20912, -0.103567, -0.166398, -0.00486649, 0.0697471, -0.0537578}; + + // Resetting cell_state and output_state + lstm.ResetFwOutputAndCellStates(); + lstm.ResetBwOutputAndCellStates(); + + 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* fw_golden_start = lstm_fw_golden_output; + float* fw_golden_end = + fw_golden_start + lstm.num_fw_outputs() * lstm.sequence_length(); + std::vector fw_expected; + fw_expected.insert(fw_expected.end(), fw_golden_start, fw_golden_end); + EXPECT_THAT(lstm.GetFwOutput(), + ElementsAreArray(ArrayFloatNear(fw_expected))); + + float* bw_golden_start = lstm_bw_golden_output; + float* bw_golden_end = + bw_golden_start + lstm.num_bw_outputs() * lstm.sequence_length(); + std::vector bw_expected; + bw_expected.insert(bw_expected.end(), bw_golden_start, bw_golden_end); + EXPECT_THAT(lstm.GetBwOutput(), + ElementsAreArray(ArrayFloatNear(bw_expected))); + + // Check reversed inputs. + static float lstm_input_reversed[] = {1., 1., 3., 4., 2., 3.}; + + // Resetting cell_state and output_state + lstm.ResetFwOutputAndCellStates(); + lstm.ResetBwOutputAndCellStates(); + + batch0_start = lstm_input_reversed; + batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); + + lstm.SetInput(0, batch0_start, batch0_end); + + lstm.Invoke(); + + fw_expected.clear(); + for (int s = 0; s < lstm.sequence_length(); s++) { + fw_golden_start = lstm_fw_golden_output + s * lstm.num_fw_outputs(); + fw_golden_end = fw_golden_start + lstm.num_fw_outputs(); + fw_expected.insert(fw_expected.begin(), fw_golden_start, fw_golden_end); + } + EXPECT_THAT(lstm.GetBwOutput(), + ElementsAreArray(ArrayFloatNear(fw_expected))); + + bw_expected.clear(); + for (int s = 0; s < lstm.sequence_length(); s++) { + bw_golden_start = lstm_bw_golden_output + s * lstm.num_bw_outputs(); + bw_golden_end = bw_golden_start + lstm.num_bw_outputs(); + bw_expected.insert(bw_expected.begin(), bw_golden_start, bw_golden_end); + } + EXPECT_THAT(lstm.GetFwOutput(), + ElementsAreArray(ArrayFloatNear(bw_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; + + BidirectionalLSTMOpModel 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 + + {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_fw_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}}; + + static float lstm_combined_golden_output[][64] = { + { + -0.022014, 0.073544, -0.002235, 0.040068, -0.037136, -0.052788, + 0.075325, -0.029378, 0.024298, -0.07733 , -0.030674, -0.060229, + 0.040599, 0.011608, 0.042005, 0.045977, -0.039225, 0.076294, + 0.000735, 0.032852, -0.069869, -0.053312, 0.073527, -0.028136, + 0.021585, -0.102679, -0.004327, -0.043304, 0.072861, 0.027077, + 0.034558, 0.068292, -0.036292, 0.069832, -0.003032, 0.053829, + -0.043821, -0.072713, 0.085029, -0.040374, 0.020014, -0.104521, + -0.034504, -0.059759, 0.062569, 0.025652, 0.049306, 0.061189, + -0.025146, 0.079643, -0.005188, 0.033080, -0.048079, -0.048082, + 0.069369, -0.028900, 0.024572, -0.077547, -0.022517, -0.054477, + 0.038857, 0.013336, 0.043234, 0.044788}, + { + -0.039186, 0.070792, -0.005913, 0.02642, -0.068274, -0.05022, + 0.061444, -0.031241, 0.014996, -0.094544, -0.004146, -0.03464, + 0.058981, 0.026097, 0.039781, 0.058408, -0.031887, 0.069252, + 0.00576, 0.054062, -0.042801, -0.059974, 0.085272, -0.034453, + 0.026097, -0.0959, -0.031164, -0.058699, 0.06839, 0.020512, + 0.044727, 0.063609, -0.039863, 0.084819, -0.003909, 0.028666, + -0.075677, -0.045125, 0.070379, -0.033895, 0.022111, -0.097184, + -0.004921, -0.040851, 0.062316, 0.017435, 0.041437, 0.064568, + -0.039656, 0.060726, -0.003402, 0.036854, -0.056503, -0.058554, + 0.068588, -0.034879, 0.01352, -0.09962, -0.01434, -0.039505, + 0.065133, 0.024321, 0.038473, 0.062438 + }}; + + // Resetting cell_state and output_state + lstm.ResetFwOutputAndCellStates(); + lstm.ResetBwOutputAndCellStates(); + + 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_fw_golden_output[0] + i * lstm.num_fw_outputs(); + float* golden_end_batch0 = golden_start_batch0 + lstm.num_fw_outputs(); + float* golden_start_batch1 = + lstm_fw_golden_output[1] + i * lstm.num_fw_outputs(); + float* golden_end_batch1 = golden_start_batch1 + lstm.num_fw_outputs(); + expected.insert(expected.end(), golden_start_batch0, golden_end_batch0); + expected.insert(expected.end(), golden_start_batch1, golden_end_batch1); + } + EXPECT_THAT(lstm.GetFwOutput(), ElementsAreArray(ArrayFloatNear(expected))); + + // Check if the sum of forward backward matches the golden. + expected.clear(); + for (int i = 0; i < lstm.sequence_length(); i++) { + float* golden_start_batch0 = + lstm_combined_golden_output[0] + i * lstm.num_fw_outputs(); + float* golden_end_batch0 = golden_start_batch0 + lstm.num_fw_outputs(); + float* golden_start_batch1 = + lstm_combined_golden_output[1] + i * lstm.num_fw_outputs(); + float* golden_end_batch1 = golden_start_batch1 + lstm.num_fw_outputs(); + expected.insert(expected.end(), golden_start_batch0, golden_end_batch0); + expected.insert(expected.end(), golden_start_batch1, golden_end_batch1); + } + + std::vector combined; + for (int i = 0; i < lstm.GetFwOutput().size(); ++i) { + combined.push_back(lstm.GetFwOutput()[i] + lstm.GetBwOutput()[i]); + } + EXPECT_THAT(combined, 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/bidirectional_sequence_rnn.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn.cc index f54081623578a7b1f37de8d9f111d7950c9e2757..aa24c1f34cd1e8c02a6a75b62fbe5f3c629498ca 100644 --- a/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn.cc +++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn.cc @@ -23,6 +23,7 @@ limitations under the License. #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 { @@ -119,47 +120,6 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { return kTfLiteOk; } -namespace { -// Performs one RNN computation step for the input specified by input_ptr_batch. -// The RNN cell is specified by the pointers to its weights and biases, along -// with the input size, number of units, strides, activation. -// The pointers to the hidden state and the output are updated as a result. -// TODO(mirkov): factor out this function to a shared library. -void RnnStep(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 input_weights_stride, - int recurrent_weights_stride, TfLiteFusedActivation activation, - float* hidden_state_ptr_batch, float* output_ptr_batch) { - // 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(activation))(output_ptr_batch[o]); - hidden_state_ptr_batch[o] = output_ptr_batch[o]; - } -} -} // namespace - TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); @@ -189,15 +149,11 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { const int input_size = input->dims->data[2]; const int fw_num_units = fw_input_weights->dims->data[0]; - const int fw_input_weights_stride = fw_input_weights->dims->data[1]; - const int fw_recurrent_weights_stride = fw_recurrent_weights->dims->data[1]; 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 int bw_input_weights_stride = bw_input_weights->dims->data[1]; - const int bw_recurrent_weights_stride = bw_recurrent_weights->dims->data[1]; 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; @@ -212,10 +168,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { float* output_ptr_batch = fw_output->data.f + b * fw_num_units * max_time + s * fw_num_units; - RnnStep(input_ptr_batch, fw_input_weights_ptr, fw_recurrent_weights_ptr, - fw_bias_ptr, input_size, fw_num_units, fw_input_weights_stride, - fw_recurrent_weights_stride, params->activation, - fw_hidden_state_ptr_batch, output_ptr_batch); + 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 = @@ -226,10 +182,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { float* output_ptr_batch = bw_output->data.f + b * bw_num_units * max_time + s * bw_num_units; - RnnStep(input_ptr_batch, bw_input_weights_ptr, bw_recurrent_weights_ptr, - bw_bias_ptr, input_size, bw_num_units, bw_input_weights_stride, - bw_recurrent_weights_stride, params->activation, - bw_hidden_state_ptr_batch, output_ptr_batch); + 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; diff --git a/tensorflow/contrib/lite/kernels/cast.cc b/tensorflow/contrib/lite/kernels/cast.cc new file mode 100644 index 0000000000000000000000000000000000000000..19942de7bc0c083f192a4b337b224b778d991140 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/cast.cc @@ -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. +==============================================================================*/ +#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/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" +#include "tensorflow/contrib/lite/string_util.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace cast { +constexpr int kInputTensor = 0; +constexpr int kOutputTensor = 0; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + return context->ResizeTensor(context, output, + TfLiteIntArrayCopy(input->dims)); +} + +template +void copyCast(const FromT* in, ToT* out, int num_elements) { + std::transform(in, in + num_elements, out, + [](FromT a) { return static_cast(a); }); +} + +template +TfLiteStatus copyToTensor(const FromT* in, TfLiteTensor* out, + int num_elements) { + switch (out->type) { + case kTfLiteInt64: + copyCast(in, out->data.i64, num_elements); + break; + case kTfLiteInt32: + copyCast(in, out->data.i32, num_elements); + break; + case kTfLiteUInt8: + copyCast(in, out->data.uint8, num_elements); + break; + case kTfLiteFloat32: + copyCast(in, out->data.f, num_elements); + break; + default: + // Unsupported type. + return kTfLiteError; + } + return kTfLiteOk; +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + const int num_elements = NumElements(input); + TF_LITE_ENSURE_EQ(context, num_elements, NumElements(output)); + switch (input->type) { + case kTfLiteInt64: + return copyToTensor(input->data.i64, output, num_elements); + case kTfLiteInt32: + return copyToTensor(input->data.i32, output, num_elements); + case kTfLiteUInt8: + return copyToTensor(input->data.uint8, output, num_elements); + case kTfLiteFloat32: + return copyToTensor(input->data.f, output, num_elements); + default: + // Unsupported type. + return kTfLiteError; + } + return kTfLiteOk; +} +} // namespace cast + +TfLiteRegistration* Register_CAST() { + static TfLiteRegistration r = {nullptr, nullptr, cast::Prepare, cast::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/cast_test.cc b/tensorflow/contrib/lite/kernels/cast_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..4e56482a371550b6275a6380e2beebe3cef958ff --- /dev/null +++ b/tensorflow/contrib/lite/kernels/cast_test.cc @@ -0,0 +1,66 @@ +/* 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 CastOpModel : public SingleOpModel { + public: + CastOpModel(const TensorData& input, const TensorData& output) { + input_ = AddInput(input); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_CAST, BuiltinOptions_CastOptions, + CreateCastOptions(builder_).Union()); + BuildInterpreter({GetShape(input_)}); + } + + int input() const { return input_; } + int output() const { return output_; } + + protected: + int input_; + int output_; +}; + +TEST(CastOpModel, CastIntToFloat) { + CastOpModel m({TensorType_INT64, {2, 3}}, {TensorType_FLOAT32, {2, 3}}); + m.PopulateTensor(m.input(), {100, 200, 300, 400, 500, 600}); + m.Invoke(); + EXPECT_THAT(m.ExtractVector(m.output()), + ElementsAreArray({100.f, 200.f, 300.f, 400.f, 500.f, 600.f})); +} + +TEST(CastOpModel, CastFloatToInt) { + CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_INT32, {3, 2}}); + m.PopulateTensor(m.input(), {100.f, 20.f, 3.f, 0.4f, 0.999f, 1.1f}); + m.Invoke(); + EXPECT_THAT(m.ExtractVector(m.output()), + ElementsAreArray({100, 20, 3, 0, 0, 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/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 37f499a4d09a38765aa4b8db8aa91b708edd7823..e0cd12f1b4042c3d8b28159e288166bf1437e6ef 100644 --- a/tensorflow/contrib/lite/kernels/conv.cc +++ b/tensorflow/contrib/lite/kernels/conv.cc @@ -23,7 +23,9 @@ limitations under the License. #include "tensorflow/contrib/lite/builtin_op_data.h" #include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/eigen_support.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,22 +40,31 @@ 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 is a mixture of an Eigen-based kernel when threads + // are available and kGenericOptimized when we must use only one thread. + 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, }; +const int kTensorNotAllocated = -1; + struct OpData { // IDs are the arbitrary identifiers used by TF Lite to identify and access // memory buffers. - int im2col_id; - int hwcn_weights_id; + int im2col_id = kTensorNotAllocated; + int hwcn_weights_id = kTensorNotAllocated; TfLitePaddingValues padding; // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multipler plus a left shift. + // be represented as a fixed point multiplier plus a left shift. int32_t output_multiplier; int output_shift; // The range of the fused activation layer. For example for kNone and @@ -67,6 +78,8 @@ struct OpData { bool need_hwcn_weights; bool have_weights_been_transposed; bool need_im2col; + + bool run_multithreaded_kernel; }; void* Init(TfLiteContext* context, const char* buffer, size_t length) { @@ -74,13 +87,16 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) { // Instead, we allocate a new object to use as scratch space for im2col, and // to carry information from Prepare() to Eval(). auto* data = new OpData; - context->AddTensors(context, 1, &data->im2col_id); - context->AddTensors(context, 1, &data->hwcn_weights_id); gemm_support::IncrementUsageCounter(context); + eigen_support::IncrementUsageCounter(context); + + data->run_multithreaded_kernel = context->recommended_num_threads != 1; + return data; } void Free(TfLiteContext* context, void* buffer) { + eigen_support::DecrementUsageCounter(context); gemm_support::DecrementUsageCounter(context); delete reinterpret_cast(buffer); } @@ -101,10 +117,67 @@ void TransposeFloatTensor(TfLiteTensor* input, TfLiteTensor* output) { } } +// Allocate temporary tensors (`im2col`, `hwcn_weights` if necessary). +// Note: `context->AddTensors` might invalidate pointers to existing tensors. +// Therefore the logic to add tensors are isolated into this function. +static TfLiteStatus AllocateTemporaryTensorsIfRequired(TfLiteContext* context, + TfLiteNode* node) { + auto* params = reinterpret_cast(node->builtin_data); + OpData* data = reinterpret_cast(node->user_data); + + TF_LITE_ENSURE(context, node->inputs->size >= 2); + TfLiteTensor* input = &context->tensors[node->inputs->data[0]]; + TfLiteTensor* filter = &context->tensors[node->inputs->data[1]]; + + int filter_width = filter->dims->data[2]; + int filter_height = filter->dims->data[1]; + + // We don't always need to allocate im2col. It is only used in some versions + // of the optimized Conv. This test just mimics something that happens inside + // optimized_ops.h, in order to avoid a DCHECK(!im2col_data). + data->need_im2col = + (params->stride_width != 1 || params->stride_height != 1 || + filter_width != 1 || filter_height != 1); + // If we're using the optimized multithreaded EigenTensor implementation of + // convolution, it expects the filter weights to be transposed compared to + // the normal TF Lite buffer format. Typical TF Lite weights are + // [filter_count, filter_height, filter_width, input_depth], but for the float + // implementation we need them as [filter_height, filter_width, input_depth, + // filter_count]. We get to that format by transposing, and create a temporary + // buffer to store the results. + // This path is only used for float processing, so only create the buffer if + // we're running with that data type. + data->need_hwcn_weights = + (input->type == kTfLiteFloat32 && data->run_multithreaded_kernel); + + int temporaries_count = 0; + if (data->need_im2col) { + data->im2col_index = temporaries_count; + if (data->im2col_id == kTensorNotAllocated) { + context->AddTensors(context, 1, &data->im2col_id); + } + ++temporaries_count; + } + if (data->need_hwcn_weights) { + data->hwcn_weights_index = temporaries_count; + if (data->hwcn_weights_id == kTensorNotAllocated) { + context->AddTensors(context, 1, &data->hwcn_weights_id); + } + ++temporaries_count; + } + + TfLiteIntArrayFree(node->temporaries); + node->temporaries = TfLiteIntArrayCreate(temporaries_count); + + return kTfLiteOk; +} + TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); OpData* data = reinterpret_cast(node->user_data); + TF_LITE_ENSURE_STATUS(AllocateTemporaryTensorsIfRequired(context, node)); + bool hasBias = node->inputs->size == 3; // Check number of inputs/outputs TF_LITE_ENSURE(context, hasBias || node->inputs->size == 2); @@ -112,6 +185,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TfLiteTensor* output = &context->tensors[node->outputs->data[0]]; TfLiteTensor* input = &context->tensors[node->inputs->data[0]]; TfLiteTensor* filter = &context->tensors[node->inputs->data[1]]; + // Check dimensionality of input, filter TF_LITE_ENSURE_EQ(context, input->dims->size, 4); TF_LITE_ENSURE_EQ(context, filter->dims->size, 4); @@ -193,36 +267,6 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { if (output_status != kTfLiteOk) return output_status; - // We don't always need to allocate im2col. It is only used in some versions - // of the optimized Conv. This test just mimics something that happens inside - // optimized_ops.h, in order to avoid a DCHECK(!im2col_data). - data->need_im2col = - (params->stride_width != 1 || params->stride_height != 1 || - filter_width != 1 || filter_height != 1); - // If we're using the optimized multithreaded EigenTensor implementation of - // convolution, it expects the filter weights to be transposed compared to - // the normal TF Lite buffer format. Typical TF Lite weights are - // [filter_count, filter_height, filter_width, input_depth], but for the float - // implementation we need them as [filter_height, filter_width, input_depth, - // filter_count]. We get to that format by transposing, and create a temporary - // buffer to store the results. - // This path is only used for float processing, so only create the buffer if - // we're running with that data type. - data->need_hwcn_weights = (data_type == kTfLiteFloat32); - - int temporaries_count = 0; - if (data->need_im2col) { - data->im2col_index = temporaries_count; - ++temporaries_count; - } - if (data->need_hwcn_weights) { - data->hwcn_weights_index = temporaries_count; - ++temporaries_count; - } - - TfLiteIntArrayFree(node->temporaries); - node->temporaries = TfLiteIntArrayCreate(temporaries_count); - if (data->need_im2col) { node->temporaries->data[data->im2col_index] = data->im2col_id; @@ -265,10 +309,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 +337,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,31 +377,57 @@ void EvalFloat(TfLiteContext* context, TfLiteNode* node, CalculateActivationRangeFloat(params->activation, &output_activation_min, &output_activation_max); - if (kernel_type == 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)); - } else { - const float* filter_data; - if (data->need_hwcn_weights) { - filter_data = GetTensorData(hwcn_weights); - } else { - filter_data = GetTensorData(filter); + 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, 1, 1, + 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, 1, 1, + 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; } - - 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)); } } @@ -379,8 +460,13 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { // separate ops to avoid dispatch overhead here. switch (input->type) { // Already know in/outtypes are same. case kTfLiteFloat32: - EvalFloat(context, node, params, data, input, filter, bias, - im2col, hwcn_weights, output); + if (data->run_multithreaded_kernel) { + EvalFloat(context, node, params, data, input, filter, bias, + im2col, hwcn_weights, output); + } else { + EvalFloat(context, node, params, data, input, filter, + bias, im2col, hwcn_weights, output); + } break; case kTfLiteUInt8: EvalQuantized(context, node, params, data, input, filter, @@ -407,17 +493,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/depthwise_conv.cc b/tensorflow/contrib/lite/kernels/depthwise_conv.cc index 15dbfe08c82befcf001b9ed9a053528b5606053e..cad9ce114c8387047af2b63bee704035fd329330 100644 --- a/tensorflow/contrib/lite/kernels/depthwise_conv.cc +++ b/tensorflow/contrib/lite/kernels/depthwise_conv.cc @@ -52,7 +52,7 @@ enum KernelType { struct OpData { TfLitePaddingValues padding; // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multipler plus a left shift. + // be represented as a fixed point multiplier plus a left shift. int32_t output_multiplier; int output_shift; // The range of the fused activation layer. For example for kNone and diff --git a/tensorflow/contrib/lite/kernels/dequantize.cc b/tensorflow/contrib/lite/kernels/dequantize.cc new file mode 100644 index 0000000000000000000000000000000000000000..e685f2465f627cf30e02564e6f16e1ec69e208e2 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/dequantize.cc @@ -0,0 +1,77 @@ +/* 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/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 dequantize { + +struct OpContext { + OpContext(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); + + OpContext op_context(context, node); + + TF_LITE_ENSURE(context, op_context.input->type == kTfLiteUInt8); + + op_context.output->type = kTfLiteFloat32; + return context->ResizeTensor(context, op_context.output, + TfLiteIntArrayCopy(op_context.input->dims)); +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + OpContext op_context(context, node); + + auto zero_point = op_context.input->params.zero_point; + auto scale = op_context.input->params.scale; + + optimized_ops::Dequantize(GetTensorData(op_context.input), + GetTensorDims(op_context.input), zero_point, scale, + GetTensorData(op_context.output), + GetTensorDims(op_context.output)); + return kTfLiteOk; +} + +} // namespace dequantize + +TfLiteRegistration* Register_DEQUANTIZE_OPT() { + static TfLiteRegistration r = {nullptr, nullptr, dequantize::Prepare, + dequantize::Eval}; + return &r; +} + +TfLiteRegistration* Register_DEQUANTIZE() { return Register_DEQUANTIZE_OPT(); } + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/dequantize_test.cc b/tensorflow/contrib/lite/kernels/dequantize_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..fcd74206177a0a97db168338e3619d4b95c052a9 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/dequantize_test.cc @@ -0,0 +1,65 @@ +/* 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 DequantizeOpModel : public SingleOpModel { + public: + DequantizeOpModel(std::initializer_list shape, float min, float max) { + input_ = AddInput({TensorType_UINT8, shape, min, max}); + output_ = AddOutput({TensorType_FLOAT32, shape}); + SetBuiltinOp(BuiltinOperator_DEQUANTIZE, BuiltinOptions_DequantizeOptions, + CreateDequantizeOptions(builder_).Union()); + + BuildInterpreter({GetShape(input_)}); + } + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + std::vector GetOutput() { return ExtractVector(output_); } + + private: + int input_; + int output_; +}; + +TEST(SplitOpTest, FourDimensional) { + DequantizeOpModel m({2, 5}, -63.5, 64); + + m.SetInput({0, 1, 2, 3, 4, 251, 252, 253, 254, 255}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear( + {-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64}))); +} + +} // 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/div.cc b/tensorflow/contrib/lite/kernels/div.cc index 44bd0dc85d50c98ec6b6888e05064a8f2e2731c0..6dd243ad62ece3e094529d923ce80d1d4a0c19ca 100644 --- a/tensorflow/contrib/lite/kernels/div.cc +++ b/tensorflow/contrib/lite/kernels/div.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,35 +61,47 @@ 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 EvalDivFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteDivParams* params, TfLiteTensor* input1, - TfLiteTensor* input2, TfLiteTensor* output) { +void EvalFloat(TfLiteContext* context, TfLiteNode* node, + TfLiteDivParams* 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_DIV(type) \ - type::Div(GetTensorData(input1), GetTensorDims(input1), \ - GetTensorData(input2), GetTensorDims(input2), \ - output_activation_min, output_activation_max, \ - GetTensorData(output), GetTensorDims(output)) +#define TF_LITE_DIV(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_DIV(reference_ops); + if (data->requires_broadcast) { + TF_LITE_DIV(reference_ops, BroadcastDiv); + } else { + TF_LITE_DIV(reference_ops, Div); + } } else { - TF_LITE_DIV(optimized_ops); + if (data->requires_broadcast) { + TF_LITE_DIV(optimized_ops, BroadcastDiv); + } else { + TF_LITE_DIV(optimized_ops, Div); + } } #undef TF_LITE_DIV } @@ -81,13 +109,14 @@ void EvalDivFloat(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) { - EvalDivFloat(context, node, params, input1, input2, output); + EvalFloat(context, node, params, data, input1, input2, output); } else { context->ReportError(context, "Inputs and outputs not all float types."); return kTfLiteError; @@ -99,19 +128,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { } // namespace div TfLiteRegistration* Register_DIV_REF() { - static TfLiteRegistration r = {nullptr, nullptr, div::Prepare, + static TfLiteRegistration r = {div::Init, div::Free, div::Prepare, div::Eval}; return &r; } TfLiteRegistration* Register_DIV_GENERIC_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, div::Prepare, + static TfLiteRegistration r = {div::Init, div::Free, div::Prepare, div::Eval}; return &r; } TfLiteRegistration* Register_DIV_NEON_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, div::Prepare, + static TfLiteRegistration r = {div::Init, div::Free, div::Prepare, div::Eval}; return &r; } diff --git a/tensorflow/contrib/lite/kernels/div_test.cc b/tensorflow/contrib/lite/kernels/div_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..276b8289fbc1b4dcbf4624b76b854300d0fd4912 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/div_test.cc @@ -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. +==============================================================================*/ +#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 BaseDivOpModel : public SingleOpModel { + public: + BaseDivOpModel(const TensorData& input1, const TensorData& input2, + const TensorData& output, + ActivationFunctionType activation_type) { + input1_ = AddInput(input1); + input2_ = AddInput(input2); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_DIV, BuiltinOptions_DivOptions, + CreateDivOptions(builder_, activation_type).Union()); + BuildInterpreter({GetShape(input1_), GetShape(input2_)}); + } + + int input1() { return input1_; } + int input2() { return input2_; } + + protected: + int input1_; + int input2_; + int output_; +}; + +class FloatDivOpModel : public BaseDivOpModel { + public: + using BaseDivOpModel::BaseDivOpModel; + + std::vector GetOutput() { return ExtractVector(output_); } +}; + +TEST(FloatDivOpTest, NoActivation) { + FloatDivOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-0.2, 0.2, -1.2, 0.8}); + m.PopulateTensor(m.input2(), {0.5, 0.2, -1.5, 0.5}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-0.4, 1.0, 0.8, 1.6}))); +} + +TEST(FloatDivOpTest, ActivationRELU_N1_TO_1) { + FloatDivOpModel m( + {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {}}, ActivationFunctionType_RELU_N1_TO_1); + m.PopulateTensor(m.input1(), {-0.2, 0.2, -1.2, 0.8}); + m.PopulateTensor(m.input2(), {0.1, 0.2, -1.5, 0.5}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-1.0, 1.0, 0.8, 1.0}))); +} + +TEST(FloatDivOpTest, VariousInputShapes) { + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + FloatDivOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.3, 0.8, 1.1, -2.0}); + m.PopulateTensor(m.input2(), {0.1, 0.2, 0.6, 0.5, -1.1, -0.1}); + m.Invoke(); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-20.0, 1.0, 0.5, 1.6, -1.0, 20.0}))) + << "With shape number " << i; + } +} + +TEST(FloatDivOpTest, WithBroadcast) { + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + FloatDivOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, {}}, // always a scalar + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-0.2, 0.2, 0.07, 0.08, 0.11, -0.123}); + m.PopulateTensor(m.input2(), {0.1}); + m.Invoke(); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-2.0, 2.0, 0.7, 0.8, 1.1, -1.23}))) + << "With shape number " << i; + } +} + +} // 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/eigen_support.cc b/tensorflow/contrib/lite/kernels/eigen_support.cc new file mode 100644 index 0000000000000000000000000000000000000000..213e46555210102b8faeb2e4d9900f924a023366 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/eigen_support.cc @@ -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. +==============================================================================*/ +#include "tensorflow/contrib/lite/kernels/eigen_support.h" + +#include "third_party/eigen3/Eigen/Core" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace eigen_support { + +struct RefCountedEigenContext { + int num_references = 0; +}; + +void IncrementUsageCounter(TfLiteContext* context) { + auto* ptr = reinterpret_cast(context->eigen_context); + if (ptr == nullptr) { + if (context->recommended_num_threads != -1) { + Eigen::setNbThreads(context->recommended_num_threads); + } + ptr = new RefCountedEigenContext; + ptr->num_references = 0; + context->eigen_context = ptr; + } + ptr->num_references++; +} + +void DecrementUsageCounter(TfLiteContext* context) { + auto* ptr = reinterpret_cast(context->eigen_context); + if (ptr == nullptr) { + TF_LITE_FATAL( + "Call to DecrementUsageCounter() not preceded by " + "IncrementUsageCounter()"); + } + if (--ptr->num_references == 0) { + delete ptr; + } +} + +} // namespace eigen_support +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/eigen_support.h b/tensorflow/contrib/lite/kernels/eigen_support.h new file mode 100644 index 0000000000000000000000000000000000000000..d47e691123282a8a8cc53c29be1d95af037e3939 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/eigen_support.h @@ -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. +==============================================================================*/ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_EIGEN_SUPPORT_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_EIGEN_SUPPORT_H_ + +#include "tensorflow/contrib/lite/context.h" + +namespace tflite { +namespace eigen_support { + +// Let the framework know that the op will be using Eigen. If necessary a set of +// temporary Eigen objects might be created and placed in 'context'. +void IncrementUsageCounter(TfLiteContext* context); + +// Let the framework know that the op stopped using Eigen. If there are no more +// usages all temporary Eigen objects will be deleted. +void DecrementUsageCounter(TfLiteContext* context); + +} // namespace eigen_support +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_EIGEN_SUPPORT_H_ 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/fully_connected.cc b/tensorflow/contrib/lite/kernels/fully_connected.cc index a77fe94e499078bc2f0660e8e49fd557ed0f625d..888e67966c0a408257e763a405bf6e928310f4d9 100644 --- a/tensorflow/contrib/lite/kernels/fully_connected.cc +++ b/tensorflow/contrib/lite/kernels/fully_connected.cc @@ -48,7 +48,7 @@ enum KernelType { struct OpData { // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multipler plus a left shift. + // be represented as a fixed point multiplier plus a left shift. int32_t output_multiplier; int output_shift; // The range of the fused activation layer. For example for kNone and diff --git a/tensorflow/contrib/lite/kernels/gemm_support.cc b/tensorflow/contrib/lite/kernels/gemm_support.cc index eb2b0aacf7ecc3ed5dbde5ccce7a46dcda0a93b3..76a5165d148c6c1829580a47456cebce321d7c5a 100644 --- a/tensorflow/contrib/lite/kernels/gemm_support.cc +++ b/tensorflow/contrib/lite/kernels/gemm_support.cc @@ -29,6 +29,9 @@ void IncrementUsageCounter(TfLiteContext* context) { if (ptr == nullptr) { ptr = new RefCountedGemmContext; ptr->gemm_context_ = new gemmlowp::GemmContext(); + if (context->recommended_num_threads != -1) { + ptr->gemm_context_->set_max_num_threads(context->recommended_num_threads); + } ptr->num_references_ = 0; context->gemm_context = ptr; } @@ -58,11 +61,5 @@ gemmlowp::GemmContext* GetFromContext(TfLiteContext* context) { return ptr->gemm_context_; } -void SetMaxNumThreads(TfLiteContext* context, int num_threads) { - IncrementUsageCounter(context); - GetFromContext(context)->set_max_num_threads(num_threads); - DecrementUsageCounter(context); -} - } // namespace gemm_support } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/gemm_support.h b/tensorflow/contrib/lite/kernels/gemm_support.h index 466781cbcecc7fb851d9078c450cc6c12364d2bb..37af772c6846f2f8124faabf1a0f0987e2e9393d 100644 --- a/tensorflow/contrib/lite/kernels/gemm_support.h +++ b/tensorflow/contrib/lite/kernels/gemm_support.h @@ -45,9 +45,6 @@ void IncrementUsageCounter(TfLiteContext* context); // 'context'. If there are no more usages the GemmContext will be deleted. void DecrementUsageCounter(TfLiteContext* context); -// Set the maximum number threads available for gemmlowp operations. -void SetMaxNumThreads(TfLiteContext* context, int num_threads); - } // namespace gemm_support } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD index 38b032c6de7987ff5e3da3ba5fcf4e9fc8574c44..aa3957bee133c8b51a82e9c62884ce365e086d2e 100644 --- a/tensorflow/contrib/lite/kernels/internal/BUILD +++ b/tensorflow/contrib/lite/kernels/internal/BUILD @@ -10,21 +10,25 @@ tflite_deps_intel = [ "@arm_neon_2_x86_sse", ] +HARD_FP_FLAGS_IF_APPLICABLE = select({ + "//tensorflow:android_arm": ["-mfloat-abi=softfp"], + "//tensorflow:android_arm64": ["-mfloat-abi=softfp"], + "//tensorflow:android_armeabi": ["-mfloat-abi=softfp"], + "//conditions:default": [], +}) + NEON_FLAGS_IF_APPLICABLE = select({ ":arm": [ "-O3", "-mfpu=neon", - "-mfloat-abi=softfp", ], ":armeabi-v7a": [ "-O3", "-mfpu=neon", - "-mfloat-abi=softfp", ], ":armv7a": [ "-O3", "-mfpu=neon", - "-mfloat-abi=softfp", ], "//conditions:default": [ "-O3", @@ -124,6 +128,13 @@ config_setting( }, ) +config_setting( + name = "darwin_x86_64", + values = { + "cpu": "darwin_x86_64", + }, +) + config_setting( name = "freebsd", values = { @@ -138,6 +149,7 @@ cc_library( "common.h", "optimized/depthwiseconv_float.h", "optimized/depthwiseconv_uint8.h", + "optimized/depthwiseconv_uint8_3x3_filter.h", "optimized/optimized_ops.h", ], copts = tflite_copts(), @@ -154,6 +166,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 +175,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", @@ -198,7 +213,10 @@ cc_library( "compatibility.h", "quantization_util.h", ], - deps = [":round"], + deps = [ + ":round", + ":types", + ], ) cc_test( @@ -232,6 +250,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": [], }), @@ -272,7 +291,7 @@ cc_library( "optimized/neon_tensor_utils.h", "optimized/tensor_utils_impl.h", ], - copts = NEON_FLAGS_IF_APPLICABLE, + copts = NEON_FLAGS_IF_APPLICABLE + HARD_FP_FLAGS_IF_APPLICABLE, deps = [ ":cpu_check", ":portable_tensor_utils", @@ -280,7 +299,38 @@ cc_library( "//tensorflow/contrib/lite:builtin_op_data", "//tensorflow/contrib/lite/kernels:activation_functor", "@arm_neon_2_x86_sse", - "@gemmlowp//:gemmlowp", + "@gemmlowp", + ], +) + +cc_library( + name = "kernel_utils", + srcs = ["kernel_utils.cc"], + hdrs = ["kernel_utils.h"], + deps = [ + ":tensor_utils", + "//tensorflow/contrib/lite:builtin_op_data", + ], +) + +# Audio support classes imported directly from TensorFlow. +cc_library( + name = "audio_utils", + srcs = [ + "mfcc.cc", + "mfcc_dct.cc", + "mfcc_mel_filterbank.cc", + "spectrogram.cc", + ], + hdrs = [ + "mfcc.h", + "mfcc_dct.h", + "mfcc_mel_filterbank.h", + "spectrogram.h", + ], + deps = [ + "//third_party/fft2d:fft2d_headers", + "@fft2d", ], ) @@ -304,7 +354,7 @@ cc_library( "//tensorflow/contrib/lite/kernels:activation_functor", "//tensorflow/contrib/lite:builtin_op_data", "@arm_neon_2_x86_sse", - "@gemmlowp//:gemmlowp", + "@gemmlowp", ] + select({ ":arm": [ ":neon_tensor_utils", @@ -324,12 +374,18 @@ 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", ], diff --git a/tensorflow/contrib/lite/kernels/internal/common.h b/tensorflow/contrib/lite/kernels/internal/common.h index fdeacedace59972dd108f8443ffad3b84f1e7e88..18601df22c1894dea6ce51f46ba815cd12dab095 100644 --- a/tensorflow/contrib/lite/kernels/internal/common.h +++ b/tensorflow/contrib/lite/kernels/internal/common.h @@ -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 // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_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..f142374269606bdd3d4184af013749102666ab89 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/kernel_utils.cc @@ -0,0 +1,191 @@ +/* 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); +} + +void LstmStep( + const float* input_ptr_batch, const float* input_to_input_weights_ptr, + const float* input_to_forget_weights_ptr, + const float* input_to_cell_weights_ptr, + const float* input_to_output_weights_ptr, + const float* recurrent_to_input_weights_ptr, + const float* recurrent_to_forget_weights_ptr, + const float* recurrent_to_cell_weights_ptr, + const float* recurrent_to_output_weights_ptr, + const float* cell_to_input_weights_ptr, + const float* cell_to_forget_weights_ptr, + const float* cell_to_output_weights_ptr, const float* input_gate_bias_ptr, + const float* forget_gate_bias_ptr, const float* cell_bias_ptr, + const float* output_gate_bias_ptr, const float* projection_weights_ptr, + const float* projection_bias_ptr, const TfLiteLSTMParams* params, + int n_batch, int n_cell, int n_input, int n_output, float* output_state_ptr, + float* cell_state_ptr, float* input_gate_scratch, + float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch, + float* output_ptr_batch) { + // 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_ptr == nullptr); + const bool use_peephole = (cell_to_output_weights_ptr != nullptr); + // Initialize scratch buffers with bias. + if (!use_cifg) { + tensor_utils::VectorBatchVectorAssign(input_gate_bias_ptr, n_cell, n_batch, + input_gate_scratch); + } + tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch, + forget_gate_scratch); + tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch, + cell_scratch); + tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, 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_ptr, n_cell, n_input, input_ptr_batch, n_batch, + input_gate_scratch, /*result_stride=*/1); + } + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_forget_weights_ptr, n_cell, n_input, input_ptr_batch, n_batch, + forget_gate_scratch, /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_cell_weights_ptr, n_cell, n_input, input_ptr_batch, n_batch, + cell_scratch, /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_output_weights_ptr, n_cell, n_input, input_ptr_batch, 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_ptr, n_cell, n_output, output_state_ptr, + n_batch, input_gate_scratch, + /*result_stride=*/1); + } + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_forget_weights_ptr, n_cell, n_output, output_state_ptr, + n_batch, forget_gate_scratch, + /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_cell_weights_ptr, n_cell, n_output, output_state_ptr, + n_batch, cell_scratch, /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_output_weights_ptr, n_cell, n_output, output_state_ptr, + 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_ptr, n_cell, cell_state_ptr, 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_ptr, n_cell, cell_state_ptr, 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_ptr, + n_batch * n_cell, cell_state_ptr); + 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_ptr); + } else { + tensor_utils::VectorVectorCwiseProductAccumulate( + cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr); + } + if (params->cell_clip > 0.0) { + tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell, + params->cell_clip, cell_state_ptr); + } + + // For each batch and cell: update the output gate. + if (use_peephole) { + tensor_utils::VectorBatchVectorCwiseProductAccumulate( + cell_to_output_weights_ptr, n_cell, cell_state_ptr, n_batch, + output_gate_scratch); + } + tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell, + output_gate_scratch); + tensor_utils::ApplyActivationToVector(cell_state_ptr, 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_ptr != nullptr); + const bool use_projection_bias = (projection_bias_ptr != nullptr); + if (use_projection_weight) { + if (use_projection_bias) { + tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output, + n_batch, output_ptr_batch); + } else { + tensor_utils::ZeroVector(output_ptr_batch, n_batch * n_output); + } + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + projection_weights_ptr, n_output, n_cell, output_gate_scratch, n_batch, + output_ptr_batch, /*result_stride=*/1); + if (params->proj_clip > 0.0) { + tensor_utils::ClipVector(output_ptr_batch, n_batch * n_output, + params->proj_clip, output_ptr_batch); + } + } else { + tensor_utils::CopyVector(output_gate_scratch, n_batch * n_output, + output_ptr_batch); + } + tensor_utils::CopyVector(output_ptr_batch, n_batch * n_output, + output_state_ptr); +} + +} // 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..3ec60ee57a87833959a34ba95d32df15bea188a4 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/kernel_utils.h @@ -0,0 +1,76 @@ +/* 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); + +// Performs an LSTM batch inference step for input specified by input_ptr_batch. +// The LSTM cell is specified by the pointers to its weights (*_weights_ptr) and +// biases (*_bias_ptr), and buffers (*_scratch), along with additional +// parameters: +// - params: various LSTM params including activation, clipping, etc., +// - n_batch: size of batch, +// - n_cell: number of cells (or units), +// - n_input: the input size, +// - n_output: the output size. +// +// The pointers to the cell and output state and the output are updated. Unless +// projection is specified output and output state contain the same data. +// +// 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 cell and output states. +void LstmStep( + const float* input_ptr_batch, const float* input_to_input_weights_ptr, + const float* input_to_forget_weights_ptr, + const float* input_to_cell_weights_ptr, + const float* input_to_output_weights_ptr, + const float* recurrent_to_input_weights_ptr, + const float* recurrent_to_forget_weights_ptr, + const float* recurrent_to_cell_weights_ptr, + const float* recurrent_to_output_weights_ptr, + const float* cell_to_input_weights_ptr, + const float* cell_to_forget_weights_ptr, + const float* cell_to_output_weights_ptr, const float* input_gate_bias_ptr, + const float* forget_gate_bias_ptr, const float* cell_bias_ptr, + const float* output_gate_bias_ptr, const float* projection_weights_ptr, + const float* projection_bias_ptr, const TfLiteLSTMParams* params, + int n_batch, int n_cell, int n_input, int n_output, float* output_state_ptr, + float* cell_state_ptr, float* input_gate_scratch, + float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch, + 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/mfcc.cc b/tensorflow/contrib/lite/kernels/internal/mfcc.cc new file mode 100644 index 0000000000000000000000000000000000000000..eafe0c7afee6fabd5a4a258aa5176e23f5e8d62a --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/mfcc.cc @@ -0,0 +1,65 @@ +/* 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/kernels/internal/mfcc.h" + +namespace tflite { +namespace internal { + +const double kDefaultUpperFrequencyLimit = 4000; +const double kDefaultLowerFrequencyLimit = 20; +const double kFilterbankFloor = 1e-12; +const int kDefaultFilterbankChannelCount = 40; +const int kDefaultDCTCoefficientCount = 13; + +Mfcc::Mfcc() + : initialized_(false), + lower_frequency_limit_(kDefaultLowerFrequencyLimit), + upper_frequency_limit_(kDefaultUpperFrequencyLimit), + filterbank_channel_count_(kDefaultFilterbankChannelCount), + dct_coefficient_count_(kDefaultDCTCoefficientCount) {} + +bool Mfcc::Initialize(int input_length, double input_sample_rate) { + bool initialized = mel_filterbank_.Initialize( + input_length, input_sample_rate, filterbank_channel_count_, + lower_frequency_limit_, upper_frequency_limit_); + initialized &= + dct_.Initialize(filterbank_channel_count_, dct_coefficient_count_); + initialized_ = initialized; + return initialized; +} + +void Mfcc::Compute(const std::vector& spectrogram_frame, + std::vector* output) const { + if (!initialized_) { + // LOG(ERROR) << "Mfcc not initialized."; + return; + } + std::vector working; + mel_filterbank_.Compute(spectrogram_frame, &working); + for (int i = 0; i < working.size(); ++i) { + double val = working[i]; + if (val < kFilterbankFloor) { + val = kFilterbankFloor; + } + working[i] = log(val); + } + dct_.Compute(working, output); +} + +} // namespace internal +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/mfcc.h b/tensorflow/contrib/lite/kernels/internal/mfcc.h new file mode 100644 index 0000000000000000000000000000000000000000..d8500ecdcf38e5dcfe9eb89915501678455b3dd9 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/mfcc.h @@ -0,0 +1,78 @@ +/* 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. +==============================================================================*/ + +// Basic class for computing MFCCs from spectrogram slices. + +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_H_ + +#include + +#include "tensorflow/contrib/lite/kernels/internal/mfcc_dct.h" +#include "tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.h" + +namespace tflite { +namespace internal { + +class Mfcc { + public: + Mfcc(); + bool Initialize(int input_length, double input_sample_rate); + + // Input is a single squared-magnitude spectrogram frame. The input spectrum + // is converted to linear magnitude and weighted into bands using a + // triangular mel filterbank, and a discrete cosine transform (DCT) of the + // values is taken. Output is populated with the lowest dct_coefficient_count + // of these values. + void Compute(const std::vector& spectrogram_frame, + std::vector* output) const; + + void set_upper_frequency_limit(double upper_frequency_limit) { + // CHECK(!initialized_) << "Set frequency limits before calling + // Initialize."; + upper_frequency_limit_ = upper_frequency_limit; + } + + void set_lower_frequency_limit(double lower_frequency_limit) { + // CHECK(!initialized_) << "Set frequency limits before calling + // Initialize."; + lower_frequency_limit_ = lower_frequency_limit; + } + + void set_filterbank_channel_count(int filterbank_channel_count) { + /// CHECK(!initialized_) << "Set channel count before calling Initialize."; + filterbank_channel_count_ = filterbank_channel_count; + } + + void set_dct_coefficient_count(int dct_coefficient_count) { + // CHECK(!initialized_) << "Set coefficient count before calling + // Initialize."; + dct_coefficient_count_ = dct_coefficient_count; + } + + private: + MfccMelFilterbank mel_filterbank_; + MfccDct dct_; + bool initialized_; + double lower_frequency_limit_; + double upper_frequency_limit_; + int filterbank_channel_count_; + int dct_coefficient_count_; +}; + +} // namespace internal +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/mfcc_dct.cc b/tensorflow/contrib/lite/kernels/internal/mfcc_dct.cc new file mode 100644 index 0000000000000000000000000000000000000000..b0b7d181bdcf01688a387f33a3e64fc904324b50 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/mfcc_dct.cc @@ -0,0 +1,78 @@ +/* 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/mfcc_dct.h" + +#include + +namespace tflite { +namespace internal { + +MfccDct::MfccDct() : initialized_(false) {} + +bool MfccDct::Initialize(int input_length, int coefficient_count) { + coefficient_count_ = coefficient_count; + input_length_ = input_length; + + if (coefficient_count_ < 1) { + return false; + } + + if (input_length < 1) { + return false; + } + + if (coefficient_count_ > input_length_) { + return false; + } + + cosines_.resize(coefficient_count_); + double fnorm = sqrt(2.0 / input_length_); + // Some platforms don't have M_PI, so define a local constant here. + const double pi = atan(1) * 4; + double arg = pi / input_length_; + for (int i = 0; i < coefficient_count_; ++i) { + cosines_[i].resize(input_length_); + for (int j = 0; j < input_length_; ++j) { + cosines_[i][j] = fnorm * cos(i * arg * (j + 0.5)); + } + } + initialized_ = true; + return true; +} + +void MfccDct::Compute(const std::vector &input, + std::vector *output) const { + if (!initialized_) { + return; + } + + output->resize(coefficient_count_); + int length = input.size(); + if (length > input_length_) { + length = input_length_; + } + + for (int i = 0; i < coefficient_count_; ++i) { + double sum = 0.0; + for (int j = 0; j < length; ++j) { + sum += cosines_[i][j] * input[j]; + } + (*output)[i] = sum; + } +} + +} // namespace internal +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/mfcc_dct.h b/tensorflow/contrib/lite/kernels/internal/mfcc_dct.h new file mode 100644 index 0000000000000000000000000000000000000000..a53f5cbd9bb70c7c9dd49672681140bb9cbd2f4f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/mfcc_dct.h @@ -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. +==============================================================================*/ + +// Basic minimal DCT class for MFCC speech processing. + +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_DCT_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_DCT_H_ + +#include + +namespace tflite { +namespace internal { + +class MfccDct { + public: + MfccDct(); + bool Initialize(int input_length, int coefficient_count); + void Compute(const std::vector& input, + std::vector* output) const; + + private: + bool initialized_; + int coefficient_count_; + int input_length_; + std::vector > cosines_; +}; + +} // namespace internal +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_DCT_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.cc b/tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.cc new file mode 100644 index 0000000000000000000000000000000000000000..c3deb33d91a47bfe54b7c84d2a615df2422f90cc --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.cc @@ -0,0 +1,204 @@ +/* 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 code resamples the FFT bins, and smooths then with triangle-shaped +// weights to create a mel-frequency filter bank. For filter i centered at f_i, +// there is a triangular weighting of the FFT bins that extends from +// filter f_i-1 (with a value of zero at the left edge of the triangle) to f_i +// (where the filter value is 1) to f_i+1 (where the filter values returns to +// zero). + +// Note: this code fails if you ask for too many channels. The algorithm used +// here assumes that each FFT bin contributes to at most two channels: the +// right side of a triangle for channel i, and the left side of the triangle +// for channel i+1. If you ask for so many channels that some of the +// resulting mel triangle filters are smaller than a single FFT bin, these +// channels may end up with no contributing FFT bins. The resulting mel +// spectrum output will have some channels that are always zero. + +#include "tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.h" + +#include + +namespace tflite { +namespace internal { + +MfccMelFilterbank::MfccMelFilterbank() : initialized_(false) {} + +bool MfccMelFilterbank::Initialize(int input_length, double input_sample_rate, + int output_channel_count, + double lower_frequency_limit, + double upper_frequency_limit) { + num_channels_ = output_channel_count; + sample_rate_ = input_sample_rate; + input_length_ = input_length; + + if (num_channels_ < 1) { + // LOG(ERROR) << "Number of filterbank channels must be positive."; + return false; + } + + if (sample_rate_ <= 0) { + // LOG(ERROR) << "Sample rate must be positive."; + return false; + } + + if (input_length < 2) { + // LOG(ERROR) << "Input length must greater than 1."; + return false; + } + + if (lower_frequency_limit < 0) { + // LOG(ERROR) << "Lower frequency limit must be nonnegative."; + return false; + } + + if (upper_frequency_limit <= lower_frequency_limit) { + /// LOG(ERROR) << "Upper frequency limit must be greater than " + // << "lower frequency limit."; + return false; + } + + // An extra center frequency is computed at the top to get the upper + // limit on the high side of the final triangular filter. + center_frequencies_.resize(num_channels_ + 1); + const double mel_low = FreqToMel(lower_frequency_limit); + const double mel_hi = FreqToMel(upper_frequency_limit); + const double mel_span = mel_hi - mel_low; + const double mel_spacing = mel_span / static_cast(num_channels_ + 1); + for (int i = 0; i < num_channels_ + 1; ++i) { + center_frequencies_[i] = mel_low + (mel_spacing * (i + 1)); + } + + // Always exclude DC; emulate HTK. + const double hz_per_sbin = + 0.5 * sample_rate_ / static_cast(input_length_ - 1); + start_index_ = static_cast(1.5 + (lower_frequency_limit / hz_per_sbin)); + end_index_ = static_cast(upper_frequency_limit / hz_per_sbin); + + // Maps the input spectrum bin indices to filter bank channels/indices. For + // each FFT bin, band_mapper tells us which channel this bin contributes to + // on the right side of the triangle. Thus this bin also contributes to the + // left side of the next channel's triangle response. + band_mapper_.resize(input_length_); + int channel = 0; + for (int i = 0; i < input_length_; ++i) { + double melf = FreqToMel(i * hz_per_sbin); + if ((i < start_index_) || (i > end_index_)) { + band_mapper_[i] = -2; // Indicate an unused Fourier coefficient. + } else { + while ((center_frequencies_[channel] < melf) && + (channel < num_channels_)) { + ++channel; + } + band_mapper_[i] = channel - 1; // Can be == -1 + } + } + + // Create the weighting functions to taper the band edges. The contribution + // of any one FFT bin is based on its distance along the continuum between two + // mel-channel center frequencies. This bin contributes weights_[i] to the + // current channel and 1-weights_[i] to the next channel. + weights_.resize(input_length_); + for (int i = 0; i < input_length_; ++i) { + channel = band_mapper_[i]; + if ((i < start_index_) || (i > end_index_)) { + weights_[i] = 0.0; + } else { + if (channel >= 0) { + weights_[i] = + (center_frequencies_[channel + 1] - FreqToMel(i * hz_per_sbin)) / + (center_frequencies_[channel + 1] - center_frequencies_[channel]); + } else { + weights_[i] = (center_frequencies_[0] - FreqToMel(i * hz_per_sbin)) / + (center_frequencies_[0] - mel_low); + } + } + } + // Check the sum of FFT bin weights for every mel band to identify + // situations where the mel bands are so narrow that they don't get + // significant weight on enough (or any) FFT bins -- i.e., too many + // mel bands have been requested for the given FFT size. + std::vector bad_channels; + for (int c = 0; c < num_channels_; ++c) { + float band_weights_sum = 0.0; + for (int i = 0; i < input_length_; ++i) { + if (band_mapper_[i] == c - 1) { + band_weights_sum += (1.0 - weights_[i]); + } else if (band_mapper_[i] == c) { + band_weights_sum += weights_[i]; + } + } + // The lowest mel channels have the fewest FFT bins and the lowest + // weights sum. But given that the target gain at the center frequency + // is 1.0, if the total sum of weights is 0.5, we're in bad shape. + if (band_weights_sum < 0.5) { + bad_channels.push_back(c); + } + } + if (!bad_channels.empty()) { + /* + LOG(ERROR) << "Missing " << bad_channels.size() << " bands " + << " starting at " << bad_channels[0] + << " in mel-frequency design. " + << "Perhaps too many channels or " + << "not enough frequency resolution in spectrum. (" + << "input_length: " << input_length + << " input_sample_rate: " << input_sample_rate + << " output_channel_count: " << output_channel_count + << " lower_frequency_limit: " << lower_frequency_limit + << " upper_frequency_limit: " << upper_frequency_limit; + */ + } + initialized_ = true; + return true; +} + +// Compute the mel spectrum from the squared-magnitude FFT input by taking the +// square root, then summing FFT magnitudes under triangular integration windows +// whose widths increase with frequency. +void MfccMelFilterbank::Compute(const std::vector &input, + std::vector *output) const { + if (!initialized_) { + // LOG(ERROR) << "Mel Filterbank not initialized."; + return; + } + + if (input.size() <= end_index_) { + // LOG(ERROR) << "Input too short to compute filterbank"; + return; + } + + // Ensure output is right length and reset all values. + output->assign(num_channels_, 0.0); + + for (int i = start_index_; i <= end_index_; i++) { // For each FFT bin + double spec_val = sqrt(input[i]); + double weighted = spec_val * weights_[i]; + int channel = band_mapper_[i]; + if (channel >= 0) + (*output)[channel] += weighted; // Right side of triangle, downward slope + channel++; + if (channel < num_channels_) + (*output)[channel] += spec_val - weighted; // Left side of triangle + } +} + +double MfccMelFilterbank::FreqToMel(double freq) const { + return 1127.0 * log(1.0 + (freq / 700.0)); +} + +} // namespace internal +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.h b/tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.h new file mode 100644 index 0000000000000000000000000000000000000000..c1db28243eea39a694b7613ac7144dce9b294897 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.h @@ -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. +==============================================================================*/ + +// Basic class for applying a mel-scale mapping to a power spectrum. + +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_MEL_FILTERBANK_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_MEL_FILTERBANK_H_ + +#include + +namespace tflite { +namespace internal { + +class MfccMelFilterbank { + public: + MfccMelFilterbank(); + bool Initialize(int input_length, // Number of unique FFT bins fftsize/2+1. + double input_sample_rate, int output_channel_count, + double lower_frequency_limit, double upper_frequency_limit); + + // Takes a squared-magnitude spectrogram slice as input, computes a + // triangular-mel-weighted linear-magnitude filterbank, and places the result + // in output. + void Compute(const std::vector& input, + std::vector* output) const; + + private: + double FreqToMel(double freq) const; + bool initialized_; + int num_channels_; + double sample_rate_; + int input_length_; + std::vector center_frequencies_; // In mel, for each mel channel. + + // Each FFT bin b contributes to two triangular mel channels, with + // proportion weights_[b] going into mel channel band_mapper_[b], and + // proportion (1 - weights_[b]) going into channel band_mapper_[b] + 1. + // Thus, weights_ contains the weighting applied to each FFT bin for the + // upper-half of the triangular band. + std::vector weights_; // Right-side weight for this fft bin. + + // FFT bin i contributes to the upper side of mel channel band_mapper_[i] + std::vector band_mapper_; + int start_index_; // Lowest FFT bin used to calculate mel spectrum. + int end_index_; // Highest FFT bin used to calculate mel spectrum. +}; + +} // namespace internal +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MFCC_MEL_FILTERBANK_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 e0eca2e736be00ff09737325f06b0035e77e3103..3a53d3ab07faf63250fc18fc846e0b8f5a39d9c4 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/cpu_check.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/cpu_check.h @@ -34,7 +34,7 @@ inline bool TestCPUFeatureNeon() { #endif // __aarch64__ } -#elif defined USE_NEON || defined __ARM_NEON +#elif defined USE_NEON || defined __ARM_NEON inline bool TestCPUFeatureNeon() { return true; } diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h index e2c87df80bd927d823b150ed3799641796dfb4c7..7f6eea2d5d1cfd6f4e2a569760ecbe0d96f754c8 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.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) diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h index fc5897896477711c46b06f10003acb10783d12af..08674a6c59192a71bae48eac467fecc7b7dff7b7 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h @@ -18,6 +18,7 @@ limitations under the License. #include "fixedpoint/fixedpoint.h" #include "public/gemmlowp.h" #include "tensorflow/contrib/lite/kernels/internal/common.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h" #include "tensorflow/contrib/lite/kernels/internal/types.h" namespace tflite { @@ -1205,6 +1206,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, @@ -1643,6 +1693,22 @@ inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims, const int output_width = ArraySize(output_dims, 1); TFLITE_DCHECK(output_depth == input_depth * depth_multiplier); +#ifdef __aarch64__ + // Call kernel optimized for depthwise convolutions using 3x3 filters, + // stride = 1, no padding, depth_multiplier = 1 and depth a multiple of 16. + if (filter_width == 3 && filter_height == 3 && depth_multiplier == 1 && + stride_width == 1 && stride_height == 1 && pad_width == 0 && + pad_height == 0 && (input_depth % 16) == 0) { + DepthwiseConv3by3FilterDepth16( + input_data, input_dims, input_offset, filter_data, filter_dims, + filter_offset, bias_data, bias_dims, stride_width, stride_height, + pad_width, pad_height, depth_multiplier, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_data, output_dims); + return; + } +#endif + static const int kAccBufferMaxSize = 2048; int32 acc_buffer[kAccBufferMaxSize]; TFLITE_DCHECK_GE(kAccBufferMaxSize, output_depth); @@ -1691,6 +1757,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) diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h new file mode 100644 index 0000000000000000000000000000000000000000..e0335b2c74cbf6d9d4f17d816f5802286628964f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h @@ -0,0 +1,653 @@ +/* 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_DEPTHWISECONV_UINT8_3X3_FILTER_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_3X3_FILTER_H_ + +#include "fixedpoint/fixedpoint.h" +#include "public/gemmlowp.h" +#include "tensorflow/contrib/lite/kernels/internal/common.h" +#include "tensorflow/contrib/lite/kernels/internal/types.h" + +namespace tflite { +namespace optimized_ops { + +#ifdef __aarch64__ + +inline void preload_l1_keep(const uint8* ptr) { +#ifdef GEMMLOWP_ARM_64 + asm volatile("prfm pldl1keep, [%[ptr]]\n" ::[ptr] "r"(ptr) :); +#else + gemmlowp::Prefetch(ptr); +#endif +} + +// Implementation of quantized DepthwiseConv for 3x3 filters. + +// Below are helper structs to remove the use of arrays. +// There is an llvm bug that causes significant slowdown when using arrays for +// NEON intrinsics vector data types. +// See: https://bugs.llvm.org/show_bug.cgi?id=34945 + +struct Int32x16 { + int32x4_t v0, v1, v2, v3; +}; + +struct Int16x16 { + int16x8_t low, high; +}; + +struct Int16x16x3 { + Int16x16 v0, v1, v2; +}; + +struct Filter3x3x16 { + Int16x16x3 r0, r1, r2; +}; + +// Loads 3x3 filter of depth 16 and adds filter offsets. +inline Filter3x3x16 LoadFilterDepth16(const uint8* filter_ptr, + int32 filter_offset, int output_depth) { + Filter3x3x16 filter; + + uint8x8_t temp_u8_0, temp_u8_1, temp_u8_2, temp_u8_3, temp_u8_4, temp_u8_5, + temp_u8_6, temp_u8_7, temp_u8_8, temp_u8_9, temp_u8_10, temp_u8_11, + temp_u8_12, temp_u8_13, temp_u8_14, temp_u8_15, temp_u8_16, temp_u8_17; + int16x8_t filter_offset_vec = vdupq_n_s16(filter_offset); + + temp_u8_0 = vld1_u8(filter_ptr + 0 * output_depth); + temp_u8_1 = vld1_u8(filter_ptr + 0 * output_depth + 8); + temp_u8_2 = vld1_u8(filter_ptr + 1 * output_depth); + temp_u8_3 = vld1_u8(filter_ptr + 1 * output_depth + 8); + temp_u8_4 = vld1_u8(filter_ptr + 2 * output_depth); + temp_u8_5 = vld1_u8(filter_ptr + 2 * output_depth + 8); + + temp_u8_6 = vld1_u8(filter_ptr + 3 * output_depth); + temp_u8_7 = vld1_u8(filter_ptr + 3 * output_depth + 8); + temp_u8_8 = vld1_u8(filter_ptr + 4 * output_depth); + temp_u8_9 = vld1_u8(filter_ptr + 4 * output_depth + 8); + temp_u8_10 = vld1_u8(filter_ptr + 5 * output_depth); + temp_u8_11 = vld1_u8(filter_ptr + 5 * output_depth + 8); + + temp_u8_12 = vld1_u8(filter_ptr + 6 * output_depth); + temp_u8_13 = vld1_u8(filter_ptr + 6 * output_depth + 8); + temp_u8_14 = vld1_u8(filter_ptr + 7 * output_depth); + temp_u8_15 = vld1_u8(filter_ptr + 7 * output_depth + 8); + temp_u8_16 = vld1_u8(filter_ptr + 8 * output_depth); + temp_u8_17 = vld1_u8(filter_ptr + 8 * output_depth + 8); + + filter.r0.v0.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_0)); + filter.r0.v0.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_1)); + filter.r0.v1.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_2)); + filter.r0.v1.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_3)); + filter.r0.v2.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_4)); + filter.r0.v2.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_5)); + + filter.r1.v0.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_6)); + filter.r1.v0.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_7)); + filter.r1.v1.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_8)); + filter.r1.v1.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_9)); + filter.r1.v2.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_10)); + filter.r1.v2.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_11)); + + filter.r2.v0.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_12)); + filter.r2.v0.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_13)); + filter.r2.v1.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_14)); + filter.r2.v1.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_15)); + filter.r2.v2.low = vreinterpretq_s16_u16(vmovl_u8(temp_u8_16)); + filter.r2.v2.high = vreinterpretq_s16_u16(vmovl_u8(temp_u8_17)); + + filter.r0.v0.low = vaddq_s16(filter.r0.v0.low, filter_offset_vec); + filter.r0.v0.high = vaddq_s16(filter.r0.v0.high, filter_offset_vec); + filter.r0.v1.low = vaddq_s16(filter.r0.v1.low, filter_offset_vec); + filter.r0.v1.high = vaddq_s16(filter.r0.v1.high, filter_offset_vec); + filter.r0.v2.low = vaddq_s16(filter.r0.v2.low, filter_offset_vec); + filter.r0.v2.high = vaddq_s16(filter.r0.v2.high, filter_offset_vec); + + filter.r1.v0.low = vaddq_s16(filter.r1.v0.low, filter_offset_vec); + filter.r1.v0.high = vaddq_s16(filter.r1.v0.high, filter_offset_vec); + filter.r1.v1.low = vaddq_s16(filter.r1.v1.low, filter_offset_vec); + filter.r1.v1.high = vaddq_s16(filter.r1.v1.high, filter_offset_vec); + filter.r1.v2.low = vaddq_s16(filter.r1.v2.low, filter_offset_vec); + filter.r1.v2.high = vaddq_s16(filter.r1.v2.high, filter_offset_vec); + + filter.r2.v0.low = vaddq_s16(filter.r2.v0.low, filter_offset_vec); + filter.r2.v0.high = vaddq_s16(filter.r2.v0.high, filter_offset_vec); + filter.r2.v1.low = vaddq_s16(filter.r2.v1.low, filter_offset_vec); + filter.r2.v1.high = vaddq_s16(filter.r2.v1.high, filter_offset_vec); + filter.r2.v2.low = vaddq_s16(filter.r2.v2.low, filter_offset_vec); + filter.r2.v2.high = vaddq_s16(filter.r2.v2.high, filter_offset_vec); + + return filter; +} + +// Loads 3 input cells of depth 16 and adds input offsets. +inline Int16x16x3 LoadInputRowDepth16(const uint8* ptr, int input_depth, + int32 input_offset, + Int16x16x3 input_row) { + uint8x8_t temp_0, temp_1; + int16x8_t offset_vec = vdupq_n_s16(input_offset); + + temp_0 = vld1_u8(ptr + 0 * input_depth); + temp_1 = vld1_u8(ptr + 0 * input_depth + 8); + input_row.v0.low = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_row.v0.high = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_row.v0.low = vaddq_s16(input_row.v0.low, offset_vec); + input_row.v0.high = vaddq_s16(input_row.v0.high, offset_vec); + + temp_0 = vld1_u8(ptr + 1 * input_depth); + temp_1 = vld1_u8(ptr + 1 * input_depth + 8); + input_row.v1.low = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_row.v1.high = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_row.v1.low = vaddq_s16(input_row.v1.low, offset_vec); + input_row.v1.high = vaddq_s16(input_row.v1.high, offset_vec); + + temp_0 = vld1_u8(ptr + 2 * input_depth); + temp_1 = vld1_u8(ptr + 2 * input_depth + 8); + input_row.v2.low = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_row.v2.high = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_row.v2.low = vaddq_s16(input_row.v2.low, offset_vec); + input_row.v2.high = vaddq_s16(input_row.v2.high, offset_vec); + + return input_row; +} + +// Performs multiply accumulate on 3 inputs of depth 16. +inline Int32x16 MultiplyAccumulateRowDepth16(Int32x16 output, + const Int16x16x3& filter_row, + const Int16x16x3& input_row) { + output.v0 = vmlal_s16(output.v0, vget_low_s16(filter_row.v0.low), + vget_low_s16(input_row.v0.low)); + output.v1 = vmlal_s16(output.v1, vget_high_s16(filter_row.v0.low), + vget_high_s16(input_row.v0.low)); + output.v2 = vmlal_s16(output.v2, vget_low_s16(filter_row.v0.high), + vget_low_s16(input_row.v0.high)); + output.v3 = vmlal_s16(output.v3, vget_high_s16(filter_row.v0.high), + vget_high_s16(input_row.v0.high)); + + output.v0 = vmlal_s16(output.v0, vget_low_s16(filter_row.v1.low), + vget_low_s16(input_row.v1.low)); + output.v1 = vmlal_s16(output.v1, vget_high_s16(filter_row.v1.low), + vget_high_s16(input_row.v1.low)); + output.v2 = vmlal_s16(output.v2, vget_low_s16(filter_row.v1.high), + vget_low_s16(input_row.v1.high)); + output.v3 = vmlal_s16(output.v3, vget_high_s16(filter_row.v1.high), + vget_high_s16(input_row.v1.high)); + + output.v0 = vmlal_s16(output.v0, vget_low_s16(filter_row.v2.low), + vget_low_s16(input_row.v2.low)); + output.v1 = vmlal_s16(output.v1, vget_high_s16(filter_row.v2.low), + vget_high_s16(input_row.v2.low)); + output.v2 = vmlal_s16(output.v2, vget_low_s16(filter_row.v2.high), + vget_low_s16(input_row.v2.high)); + output.v3 = vmlal_s16(output.v3, vget_high_s16(filter_row.v2.high), + vget_high_s16(input_row.v2.high)); + + return output; +} + +// Applies activation, offset and downquantize on a set of accumulator +// registers of depth 16. Stores results to output. +inline void DownquantizeAndStoreDepth16(Int32x16 acc, int32 output_multiplier, + int output_shift, + int32x4_t output_offset_vec, + int32x4_t output_activation_min_vec, + int32x4_t output_activation_max_vec, + uint8* output_ptr) { + // Fixed-point multiplication. + acc.v0 = vqrdmulhq_n_s32(acc.v0, output_multiplier); + acc.v1 = vqrdmulhq_n_s32(acc.v1, output_multiplier); + acc.v2 = vqrdmulhq_n_s32(acc.v2, output_multiplier); + acc.v3 = vqrdmulhq_n_s32(acc.v3, output_multiplier); + + using gemmlowp::RoundingDivideByPOT; + acc.v0 = RoundingDivideByPOT(acc.v0, output_shift); + acc.v1 = RoundingDivideByPOT(acc.v1, output_shift); + acc.v2 = RoundingDivideByPOT(acc.v2, output_shift); + acc.v3 = RoundingDivideByPOT(acc.v3, output_shift); + + // Add the output offset. + acc.v0 = vaddq_s32(acc.v0, output_offset_vec); + acc.v1 = vaddq_s32(acc.v1, output_offset_vec); + acc.v2 = vaddq_s32(acc.v2, output_offset_vec); + acc.v3 = vaddq_s32(acc.v3, output_offset_vec); + + // Apply the activation function. + acc.v0 = vmaxq_s32(acc.v0, output_activation_min_vec); + acc.v1 = vmaxq_s32(acc.v1, output_activation_min_vec); + acc.v2 = vmaxq_s32(acc.v2, output_activation_min_vec); + acc.v3 = vmaxq_s32(acc.v3, output_activation_min_vec); + + acc.v0 = vminq_s32(acc.v0, output_activation_max_vec); + acc.v1 = vminq_s32(acc.v1, output_activation_max_vec); + acc.v2 = vminq_s32(acc.v2, output_activation_max_vec); + acc.v3 = vminq_s32(acc.v3, output_activation_max_vec); + + // Saturating cast to uint8 and store to destination. + int16x4_t acc_tlla_s16 = vqmovn_s32(acc.v0); + int16x4_t acc_tllb_s16 = vqmovn_s32(acc.v1); + int16x4_t acc_tlha_s16 = vqmovn_s32(acc.v2); + int16x4_t acc_tlhb_s16 = vqmovn_s32(acc.v3); + + int16x8_t res_s16_0 = vcombine_s16(acc_tlla_s16, acc_tllb_s16); + int16x8_t res_s16_1 = vcombine_s16(acc_tlha_s16, acc_tlhb_s16); + uint8x8_t res_u8_0 = vqmovun_s16(res_s16_0); + uint8x8_t res_u8_1 = vqmovun_s16(res_s16_1); + vst1q_u8(output_ptr, vcombine_u8(res_u8_0, res_u8_1)); +} + +// A kernel that is optimized on the number of output cells in the x and y +// direction, and the stride. Assumes 3x3 filters of 16 depth. +template +struct ConvKernel3x3FilterDepth16 {}; + +template <> +struct ConvKernel3x3FilterDepth16<1, 2, 1> { + static void Run(const Filter3x3x16& filter, const uint8* input_ptr, + int input_depth, int32 input_offset, int input_row_width, + const int32* bias_ptr, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, int32 output_activation_max, + uint8* output_ptr, int output_depth, int output_width) { + // 16 depth accumulators for the 2 outputs. + Int32x16 acc0, acc1; + + // Accumulators for top filter. + acc0.v0 = vld1q_s32(bias_ptr); + acc0.v1 = vld1q_s32(bias_ptr + 4); + acc0.v2 = vld1q_s32(bias_ptr + 8); + acc0.v3 = vld1q_s32(bias_ptr + 12); + // Accumulators for bottom filter. + acc1.v0 = vld1q_s32(bias_ptr); + acc1.v1 = vld1q_s32(bias_ptr + 4); + acc1.v2 = vld1q_s32(bias_ptr + 8); + acc1.v3 = vld1q_s32(bias_ptr + 12); + + // Main multiply accumulate work. + { + // Load inputs for one filter row at a time. + Int16x16x3 input; + + // Do first row of top filter. + input = LoadInputRowDepth16(input_ptr, input_depth, input_offset, input); + acc0 = MultiplyAccumulateRowDepth16(acc0, filter.r0, input); + + // Do second row of top filter. + input = LoadInputRowDepth16(input_ptr + input_row_width, input_depth, + input_offset, input); + acc0 = MultiplyAccumulateRowDepth16(acc0, filter.r1, input); + + // The inputs to second row of the top filter are also the inputs to the + // first row of the bottom filter. + acc1 = MultiplyAccumulateRowDepth16(acc1, filter.r0, input); + + // Do third row of top filter. + input = LoadInputRowDepth16(input_ptr + 2 * input_row_width, input_depth, + input_offset, input); + acc0 = MultiplyAccumulateRowDepth16(acc0, filter.r2, input); + + // The inputs to third row of the top filter are also the inputs to the + // second row of the bottom filter. + acc1 = MultiplyAccumulateRowDepth16(acc1, filter.r1, input); + + // Do third row of bottom filter. + input = LoadInputRowDepth16(input_ptr + 3 * input_row_width, input_depth, + input_offset, input); + acc1 = MultiplyAccumulateRowDepth16(acc1, filter.r2, input); + } + + // Apply activation, downquantize and store. + int32x4_t output_offset_vec = vdupq_n_s32(output_offset); + int32x4_t output_activation_min_vec = vdupq_n_s32(output_activation_min); + int32x4_t output_activation_max_vec = vdupq_n_s32(output_activation_max); + + DownquantizeAndStoreDepth16(acc0, output_multiplier, output_shift, + output_offset_vec, output_activation_min_vec, + output_activation_max_vec, output_ptr); + + DownquantizeAndStoreDepth16(acc1, output_multiplier, output_shift, + output_offset_vec, output_activation_min_vec, + output_activation_max_vec, + output_ptr + output_depth * output_width); + } +}; + +template <> +struct ConvKernel3x3FilterDepth16<1, 2, 2> { + static void Run(const Filter3x3x16& filter, const uint8* input_ptr, + int input_depth, int32 input_offset, int input_row_width, + const int32* bias_ptr, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, int32 output_activation_max, + uint8* output_ptr, int output_depth, int output_width) { + // 16 depth accumulators for the 2 outputs. + Int32x16 acc0, acc1; + + // Accumulators for top filter. + acc0.v0 = vld1q_s32(bias_ptr); + acc0.v1 = vld1q_s32(bias_ptr + 4); + acc0.v2 = vld1q_s32(bias_ptr + 8); + acc0.v3 = vld1q_s32(bias_ptr + 12); + // Accumulators for bottom filter. + acc1.v0 = vld1q_s32(bias_ptr); + acc1.v1 = vld1q_s32(bias_ptr + 4); + acc1.v2 = vld1q_s32(bias_ptr + 8); + acc1.v3 = vld1q_s32(bias_ptr + 12); + + // Main multiply accumulate work. + { + // Load inputs for one filter row at a time. + Int16x16x3 input; + + // Do first row of top filter. + input = LoadInputRowDepth16(input_ptr, input_depth, input_offset, input); + acc0 = MultiplyAccumulateRowDepth16(acc0, filter.r0, input); + + // Do second row of top filter. + input = LoadInputRowDepth16(input_ptr + input_row_width, input_depth, + input_offset, input); + acc0 = MultiplyAccumulateRowDepth16(acc0, filter.r1, input); + + // Do third row of top filter. + input = LoadInputRowDepth16(input_ptr + 2 * input_row_width, input_depth, + input_offset, input); + acc0 = MultiplyAccumulateRowDepth16(acc0, filter.r2, input); + + // The inputs to third row of the top filter are also the inputs + // to first row of the bottom filter. + acc1 = MultiplyAccumulateRowDepth16(acc1, filter.r0, input); + + // Do second row of bottom filter. + input = LoadInputRowDepth16(input_ptr + 3 * input_row_width, input_depth, + input_offset, input); + acc1 = MultiplyAccumulateRowDepth16(acc1, filter.r1, input); + + // Do third row of bottom filter. + input = LoadInputRowDepth16(input_ptr + 4 * input_row_width, input_depth, + input_offset, input); + acc1 = MultiplyAccumulateRowDepth16(acc1, filter.r2, input); + } + + // Apply activation, downquantize and store. + int32x4_t output_offset_vec = vdupq_n_s32(output_offset); + int32x4_t output_activation_min_vec = vdupq_n_s32(output_activation_min); + int32x4_t output_activation_max_vec = vdupq_n_s32(output_activation_max); + + DownquantizeAndStoreDepth16(acc0, output_multiplier, output_shift, + output_offset_vec, output_activation_min_vec, + output_activation_max_vec, output_ptr); + + DownquantizeAndStoreDepth16(acc1, output_multiplier, output_shift, + output_offset_vec, output_activation_min_vec, + output_activation_max_vec, + output_ptr + output_depth * output_width); + } +}; + +template <> +struct ConvKernel3x3FilterDepth16<1, 1> { + static void Run(const Filter3x3x16& filter, const uint8* input_ptr, + int input_depth, int32 input_offset, int input_row_width, + const int32* bias_ptr, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, int32 output_activation_max, + uint8* output_ptr, int output_depth, int output_width) { + Int32x16 acc; + acc.v0 = vld1q_s32(bias_ptr); + acc.v1 = vld1q_s32(bias_ptr + 4); + acc.v2 = vld1q_s32(bias_ptr + 8); + acc.v3 = vld1q_s32(bias_ptr + 12); + + // Main multiply accumulate work. + { + // Load inputs for one filter row at a time. + Int16x16x3 input; + + // Do first row. + input = LoadInputRowDepth16(input_ptr, input_depth, input_offset, input); + acc = MultiplyAccumulateRowDepth16(acc, filter.r0, input); + + // Do second row. + input = LoadInputRowDepth16(input_ptr + input_row_width, input_depth, + input_offset, input); + acc = MultiplyAccumulateRowDepth16(acc, filter.r1, input); + + // Do third row. + input = LoadInputRowDepth16(input_ptr + 2 * input_row_width, input_depth, + input_offset, input); + acc = MultiplyAccumulateRowDepth16(acc, filter.r2, input); + } + + // Apply activation, downquantize and store. + int32x4_t output_offset_vec = vdupq_n_s32(output_offset); + int32x4_t output_activation_min_vec = vdupq_n_s32(output_activation_min); + int32x4_t output_activation_max_vec = vdupq_n_s32(output_activation_max); + + DownquantizeAndStoreDepth16(acc, output_multiplier, output_shift, + output_offset_vec, output_activation_min_vec, + output_activation_max_vec, output_ptr); + } +}; + +inline void DepthwiseConv3by3FilterDepth16( + const uint8* input_data, const Dims<4>& input_dims, int32 input_offset, + const uint8* filter_data, const Dims<4>& filter_dims, int32 filter_offset, + const int32* bias_data, const Dims<4>& bias_dims, int stride_width, + int stride_height, int pad_width, int pad_height, int depth_multiplier, + int32 output_offset, int32 output_multiplier, int output_shift, + int32 output_activation_min, int32 output_activation_max, + uint8* output_data, const Dims<4>& output_dims) { + const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); + const int output_depth = MatchingArraySize(filter_dims, 0, output_dims, 0); + const int input_height = ArraySize(input_dims, 2); + const int input_width = ArraySize(input_dims, 1); + const int input_depth = ArraySize(input_dims, 0); + const int filter_height = ArraySize(filter_dims, 2); + const int filter_width = ArraySize(filter_dims, 1); + const int output_height = ArraySize(output_dims, 2); + const int output_width = ArraySize(output_dims, 1); + + // Algorithm assumes below constraints. It is optimized for depth multiplier + // of 1, 3x3 filter, no padding, strides 1 and 2. + TFLITE_DCHECK(output_depth == input_depth * depth_multiplier); + TFLITE_DCHECK(depth_multiplier == 1); + TFLITE_DCHECK(filter_height == 3); + TFLITE_DCHECK(filter_width == 3); + TFLITE_DCHECK(pad_height == 0); + TFLITE_DCHECK(pad_width == 0); + TFLITE_DCHECK(stride_width == 1); + TFLITE_DCHECK(stride_height == 1); + + // The number of outputs to process in the main loop. + const int num_x_outputs = 1; + const int num_y_outputs = 2; + + const int input_row_width = output_depth * (input_width + 2 * pad_width); + const int input_batch_size = + input_row_width * (input_height + 2 * pad_height); + const int output_batch_size = output_depth * output_width * output_height; + const int input_ptr_x_increment = input_depth * stride_width; + + // Calculate extents of non-boundary loop. + int out_x_start = 0; + for (; out_x_start < input_width; out_x_start++) { + int in_x = (out_x_start * stride_width) - pad_width; + if (in_x >= 0) { + break; + } + } + int out_x_end = output_width - 1; + for (; out_x_end >= 0; out_x_end--) { + int in_x = (out_x_end * stride_width) - pad_width; + int in_x_end = in_x + filter_width + (num_x_outputs - 1) * stride_width; + if (in_x_end <= input_width) { + out_x_end++; + break; + } + } + int out_y_start = 0; + for (; out_y_start < input_height; out_y_start++) { + int in_y = (out_y_start * stride_height) - pad_height; + if (in_y >= 0) { + break; + } + } + int out_y_end = output_height - 1; + for (; out_y_end >= 0; out_y_end--) { + int in_y = (out_y_end * stride_height) - pad_height; + int in_y_end = in_y + filter_height + (num_y_outputs - 1) * stride_height; + if (in_y_end <= input_height) { + out_y_end++; + break; + } + } + + // Offsets for preloading inputs. + const int i0 = 0; + const int i1 = input_depth; + const int i2 = 2 * input_depth; + const int i3 = input_row_width; + const int i4 = input_row_width + input_depth; + const int i5 = input_row_width + 2 * input_depth; + const int i6 = 2 * input_row_width; + const int i7 = 2 * input_row_width + input_depth; + const int i8 = 2 * input_row_width + 2 * input_depth; + const int i9 = 3 * input_row_width; + const int i10 = 3 * input_row_width + input_depth; + const int i11 = 3 * input_row_width + 2 * input_depth; + + for (int b = 0; b < batches; ++b) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const int in_batch_offset = b * input_batch_size; + const int out_batch_offset = b * output_batch_size; + + int depth = 0; + for (; depth <= output_depth - 16; depth += 16) { + Filter3x3x16 filter = + LoadFilterDepth16(filter_ptr, filter_offset, output_depth); + + // Handle 1x2 outputs. + int out_y = out_y_start; + for (; out_y < out_y_end; out_y += num_y_outputs) { + int out_x = out_x_start; + + int in_y_offset = + stride_height * input_row_width * (out_y + pad_height); + int in_x_offset = stride_width * input_depth * (out_x + pad_width); + + const uint8* input_ptr = + input_data + depth + in_x_offset + in_y_offset + in_batch_offset; + + uint8* output_ptr = output_data + depth + (out_x * output_depth) + + (output_depth * output_width * out_y) + + out_batch_offset; + + // Preload inputs. If input depth is large, preload every value of the + // input for this depth range. Otherwise, preload only the first values + // of each row. + if (input_depth >= 32) { + preload_l1_keep(input_ptr + i0); + preload_l1_keep(input_ptr + i1); + preload_l1_keep(input_ptr + i2); + preload_l1_keep(input_ptr + i3); + preload_l1_keep(input_ptr + i4); + preload_l1_keep(input_ptr + i5); + preload_l1_keep(input_ptr + i6); + preload_l1_keep(input_ptr + i7); + preload_l1_keep(input_ptr + i8); + preload_l1_keep(input_ptr + i9); + preload_l1_keep(input_ptr + i10); + preload_l1_keep(input_ptr + i11); + } else { + preload_l1_keep(input_ptr + i0); + preload_l1_keep(input_ptr + i3); + preload_l1_keep(input_ptr + i6); + preload_l1_keep(input_ptr + i9); + } + + for (; out_x < out_x_end; out_x += num_x_outputs) { + ConvKernel3x3FilterDepth16<1, 2, 1>::Run( + filter, input_ptr, input_depth, input_offset, input_row_width, + bias_ptr, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, + output_depth, output_width); + + input_ptr += input_ptr_x_increment * num_x_outputs; + output_ptr += output_depth * num_x_outputs; + + // Preload the next inputs depending on stride. + if (stride_width == 1) { + preload_l1_keep(input_ptr + i2); + preload_l1_keep(input_ptr + i5); + preload_l1_keep(input_ptr + i8); + preload_l1_keep(input_ptr + i11); + } else if (stride_width == 2) { + preload_l1_keep(input_ptr + i1); + preload_l1_keep(input_ptr + i2); + preload_l1_keep(input_ptr + i4); + preload_l1_keep(input_ptr + i5); + preload_l1_keep(input_ptr + i7); + preload_l1_keep(input_ptr + i8); + preload_l1_keep(input_ptr + i10); + preload_l1_keep(input_ptr + i11); + } + } + + // Handle the rest of the right side. + for (; out_x < output_width; out_x++) { + // This code path can only be reached if we're handling >1 x outputs + // at a time or support padding. + } + } + + // Handle the rest of the bottom side. + for (; out_y < output_height; out_y++) { + int out_x = out_x_start; + + int in_y_offset = + stride_height * input_row_width * (out_y + pad_height); + int in_x_offset = stride_width * input_depth * (out_x + pad_width); + + const uint8* input_ptr = + input_data + depth + in_x_offset + in_y_offset + in_batch_offset; + + uint8* output_ptr = output_data + depth + (out_x * output_depth) + + (output_depth * output_width * out_y) + + out_batch_offset; + + for (; out_x < output_width; out_x++) { + ConvKernel3x3FilterDepth16<1, 1>::Run( + filter, input_ptr, input_depth, input_offset, input_row_width, + bias_ptr, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, + output_depth, output_width); + + input_ptr += input_ptr_x_increment; + output_ptr += output_depth; + } + } + filter_ptr += 16; + bias_ptr += 16; + } + } +} + +#endif // __aarch64__ + +} // namespace optimized_ops +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_3X3_FILTER_H_ 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 ea8502ae33bc2ee5965e3be30a7d9ee36604abe3..780401e052733cccae0cc34f495df090c1530624 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc +++ b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc @@ -17,6 +17,7 @@ limitations under the License. #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 diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h index 8163c76cfd2eb9b320fe65e54c6b88f3d694a598..edd65c9170b0f7646acbbae666ac72f2ffb8b4de 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h @@ -610,6 +610,58 @@ inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, input_offset, output_pipeline); } +inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, + int32 input_offset, const uint8* filter_data, + const Dims<4>& filter_dims, int32 filter_offset, + const int32* bias_data, const Dims<4>& bias_dims, + int32 output_offset, int32 output_multiplier, + int output_shift, int32 output_activation_min, + int32 output_activation_max, int16* output_data, + const Dims<4>& output_dims, + gemmlowp::GemmContext* gemm_context) { + gemmlowp::ScopedProfilingLabel label("FullyConnected/Uint8Int16"); + // This is a copy of the reference implementation. We do not currently have a + // properly optimized version. + (void)gemm_context; // only used in properly optimized code. + TFLITE_DCHECK_LE(output_activation_min, output_activation_max); + TFLITE_DCHECK_EQ(output_offset, 0); + // TODO(benoitjacob): This really should be: + // const int batches = ArraySize(output_dims, 1); + // but the current --variable_batch hack consists in overwriting the 3rd + // dimension with the runtime batch size, as we don't keep track for each + // array of which dimension is the batch dimension in it. + const int batches = ArraySize(output_dims, 1) * ArraySize(output_dims, 2) * + ArraySize(output_dims, 3); + const int output_depth = MatchingArraySize(filter_dims, 1, output_dims, 0); + const int accum_depth = ArraySize(filter_dims, 0); + TFLITE_DCHECK(IsPackedWithoutStrides(input_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(filter_dims)); + for (int b = 0; b < batches; ++b) { + for (int out_c = 0; out_c < output_depth; ++out_c) { + // Internal accumulation. + // Initialize accumulator with the bias-value. + int32 accum = bias_data[out_c]; + // Accumulation loop. + for (int d = 0; d < accum_depth; ++d) { + int16 input_val = input_data[b * accum_depth + d] + input_offset; + int16 filter_val = filter_data[out_c * accum_depth + d] + filter_offset; + accum += filter_val * input_val; + } + // Down-scale the final int32 accumulator to the scale used by our + // (16-bit, typically 3 integer bits) fixed-point format. The quantized + // multiplier and shift here have been pre-computed offline + // (e.g. by toco). + accum = MultiplyByQuantizedMultiplier(accum, output_multiplier, + -output_shift); + // Saturate, cast to int16, and store to output array. + accum = std::max(accum, output_activation_min - output_offset); + accum = std::min(accum, output_activation_max - output_offset); + accum += output_offset; + output_data[out_c + output_depth * b] = accum; + } + } +} + // legacy, for compatibility with old checked-in code template void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, @@ -758,14 +810,90 @@ void Im2col(const T* input_data, const Dims<4>& input_dims, int stride, kwidth, byte_zero, output_data, output_dims); } +inline void DilatedConv(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 dilation_width_factor, int dilation_height_factor, + 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("DilatedConv"); + // This is a copy of the reference Conv implementation. We do not currently + // have an optimized path for dilation. + (void)im2col_data; // only used in optimized code. + (void)im2col_dims; // only used in optimized code. + const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); + const int input_depth = MatchingArraySize(input_dims, 0, filter_dims, 0); + const int output_depth = MatchingArraySize(filter_dims, 3, output_dims, 0); + if (bias_data) { + TFLITE_DCHECK_EQ(ArraySize(filter_dims, 3), ArraySize(bias_dims, 0)); + } + const int input_height = ArraySize(input_dims, 2); + const int input_width = ArraySize(input_dims, 1); + const int filter_height = ArraySize(filter_dims, 2); + const int filter_width = ArraySize(filter_dims, 1); + const int output_height = ArraySize(output_dims, 2); + const int output_width = ArraySize(output_dims, 1); + for (int batch = 0; batch < batches; ++batch) { + for (int out_y = 0; out_y < output_height; ++out_y) { + for (int out_x = 0; out_x < output_width; ++out_x) { + for (int out_channel = 0; out_channel < output_depth; ++out_channel) { + const int in_x_origin = (out_x * stride_width) - pad_width; + const int in_y_origin = (out_y * stride_height) - pad_height; + float total = 0.f; + for (int filter_y = 0; filter_y < filter_height; ++filter_y) { + for (int filter_x = 0; filter_x < filter_width; ++filter_x) { + for (int in_channel = 0; in_channel < input_depth; ++in_channel) { + const int in_x = in_x_origin + dilation_width_factor * filter_x; + const int in_y = + in_y_origin + dilation_height_factor * filter_y; + // If the location is outside the bounds of the input image, + // use zero as a default value. + if ((in_x >= 0) && (in_x < input_width) && (in_y >= 0) && + (in_y < input_height)) { + float input_value = input_data[Offset(input_dims, in_channel, + in_x, in_y, batch)]; + float filter_value = + filter_data[Offset(filter_dims, in_channel, filter_x, + filter_y, out_channel)]; + total += (input_value * filter_value); + } + } + } + } + float bias_value = 0.0f; + if (bias_data) { + bias_value = bias_data[Offset(bias_dims, out_channel, 0, 0, 0)]; + } + output_data[Offset(output_dims, out_channel, out_x, out_y, batch)] = + ActivationFunctionWithMinMax(total + bias_value, + output_activation_min, + output_activation_max); + } + } + } + } +} + 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) { + int stride_width, int stride_height, int dilation_width_factor, + int dilation_height_factor, 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) { + if ((dilation_width_factor != 1) || (dilation_height_factor != 1)) { + return DilatedConv(input_data, input_dims, filter_data, filter_dims, + bias_data, bias_dims, stride_width, stride_height, + dilation_width_factor, dilation_height_factor, pad_width, + pad_height, output_activation_min, output_activation_max, + output_data, output_dims, im2col_data, im2col_dims); + } + (void)im2col_data; (void)im2col_dims; gemmlowp::ScopedProfilingLabel label("Conv"); @@ -805,6 +933,23 @@ inline void Conv(const float* input_data, const Dims<4>& input_dims, output_activation_max); } +template +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 dilation_width_factor, + int dilation_height_factor, int pad_width, int pad_height, + float* output_data, const Dims<4>& output_dims, float* im2col_data, + const Dims<4>& im2col_dims) { + float output_activation_min, output_activation_max; + GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); + Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims, + stride_width, stride_height, dilation_width_factor, + dilation_height_factor, pad_width, pad_height, output_activation_min, + output_activation_max, output_data, output_dims, im2col_data, + im2col_dims); +} + // legacy, for compatibility with old checked-in code template void Conv(const float* input_data, const Dims<4>& input_dims, @@ -816,7 +961,7 @@ void Conv(const float* input_data, const Dims<4>& input_dims, float output_activation_min, output_activation_max; GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims, - stride_width, stride_height, pad_width, pad_height, + stride_width, stride_height, 1, 1, pad_width, pad_height, output_activation_min, output_activation_max, output_data, output_dims, im2col_data, im2col_dims); } @@ -830,7 +975,7 @@ void Conv(const float* input_data, const Dims<4>& input_dims, const Dims<4>& output_dims, float* im2col_data, const Dims<4>& im2col_dims) { Conv(input_data, input_dims, filter_data, filter_dims, bias_data, - bias_dims, stride, stride, pad_width, pad_height, output_data, + bias_dims, stride, stride, 1, 1, pad_width, pad_height, output_data, output_dims, im2col_data, im2col_dims); } @@ -1506,6 +1651,39 @@ inline void Add(int left_shift, const uint8* input1_data, } } +template +inline void Add(const int16* input1_data, const Dims<4>& input1_dims, + int input1_shift, const int16* input2_data, + const Dims<4>& input2_dims, int input2_shift, + int16* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Add/Int16"); + // This is a copy of the reference implementation. We do not currently have a + // properly optimized version. + static_assert(Ac == FusedActivationFunctionType::kNone, ""); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input1_dims), flat_size); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input2_dims), flat_size); + + TFLITE_DCHECK(input1_shift == 0 || input2_shift == 0); + TFLITE_DCHECK_GE(input1_shift, 0); + TFLITE_DCHECK_GE(input2_shift, 0); + const int16* not_shift_input = input1_shift == 0 ? input1_data : input2_data; + const int16* shift_input = input1_shift == 0 ? input2_data : input1_data; + const int input_shift = input1_shift == 0 ? input2_shift : input1_shift; + + for (int i = 0; i < flat_size; i++) { + // F0 uses 0 integer bits, range [-1, 1]. + using F0 = gemmlowp::FixedPoint; + + F0 input_ready_scaled = F0::FromRaw(not_shift_input[i]); + F0 scaled_input = + F0::FromRaw(gemmlowp::RoundingDivideByPOT(shift_input[i], input_shift)); + F0 result = gemmlowp::SaturatingAdd(scaled_input, input_ready_scaled); + output_data[i] = result.raw(); + } +} + template void Add(const int32* input1_data, const Dims<4>& input1_dims, const int32* input2_data, const Dims<4>& input2_dims, @@ -1780,6 +1958,57 @@ void Mul(const int32* input1_data, const Dims<4>& input1_dims, } } +inline void Mul(const int16* input1_data, const Dims<4>& input1_dims, + const int16* input2_data, const Dims<4>& input2_dims, + int16* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Mul/Int16"); + // This is a copy of the reference implementation. We do not currently have a + // properly optimized version. + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input1_dims), flat_size); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input2_dims), flat_size); + + for (int i = 0; i < flat_size; i++) { + // F0 uses 0 integer bits, range [-1, 1]. + using F0 = gemmlowp::FixedPoint; + + F0 unclamped_result = + F0::FromRaw(input1_data[i]) * F0::FromRaw(input2_data[i]); + output_data[i] = unclamped_result.raw(); + } +} + +inline void Mul(const int16* input1_data, const Dims<4>& input1_dims, + const int16* input2_data, const Dims<4>& input2_dims, + int32 output_offset, int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Mul/Int16Uint8"); + // This is a copy of the reference implementation. We do not currently have a + // properly optimized version. + TFLITE_DCHECK_LE(output_activation_min, output_activation_max); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input1_dims), flat_size); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input2_dims), flat_size); + + for (int i = 0; i < flat_size; i++) { + // F0 uses 0 integer bits, range [-1, 1]. + using F0 = gemmlowp::FixedPoint; + + F0 unclamped_result = + F0::FromRaw(input1_data[i]) * F0::FromRaw(input2_data[i]); + int16 rescaled_result = + gemmlowp::RoundingDivideByPOT(unclamped_result.raw(), 8); + int16 clamped_result = + std::min(output_activation_max - output_offset, rescaled_result); + clamped_result = + std::max(output_activation_min - output_offset, clamped_result); + output_data[i] = output_offset + clamped_result; + } +} + // TODO(jiawen): We can implement BroadcastMul on buffers of arbitrary // dimensionality if the runtime code does a single loop over one dimension // that handles broadcasting as the base case. The code generator would then @@ -1928,6 +2157,51 @@ inline void Div(const float* input1_data, const Dims<4>& input1_dims, } } +// TODO(jiawen): We can implement BroadcastDiv on buffers of arbitrary +// 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. +// TODO(benoitjacob): BroadcastDiv is intentionally duplicated from +// 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 +void BroadcastDiv(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("BroadcastDiv"); + + NdArrayDesc<4> desc1; + NdArrayDesc<4> desc2; + NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2); + + // In Tensorflow, the dimensions are canonically named (batch_number, row, + // col, channel), with extents (batches, height, width, depth), with the + // trailing dimension changing most rapidly (channels has the smallest stride, + // typically 1 element). + // + // In generated C code, we store arrays with the dimensions reversed. The + // first dimension has smallest stride. + // + // We name our variables by their Tensorflow convention, but generate C code + // nesting loops such that the innermost loop has the smallest stride for the + // best cache behavior. + for (int b = 0; b < ArraySize(output_dims, 3); ++b) { + 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)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] / + input2_data[SubscriptToIndex(desc2, 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, @@ -1955,6 +2229,52 @@ inline void Sub(const float* input1_data, const Dims<4>& input1_dims, } } } + +// TODO(jiawen): We can implement BroadcastSub on buffers of arbitrary +// 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. +// TODO(benoitjacob): BroadcastSub is intentionally duplicated from +// 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 +void BroadcastSub(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("BroadcastSub"); + + NdArrayDesc<4> desc1; + NdArrayDesc<4> desc2; + NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2); + + // In Tensorflow, the dimensions are canonically named (batch_number, row, + // col, channel), with extents (batches, height, width, depth), with the + // trailing dimension changing most rapidly (channels has the smallest stride, + // typically 1 element). + // + // In generated C code, we store arrays with the dimensions reversed. The + // first dimension has smallest stride. + // + // We name our variables by their Tensorflow convention, but generate C code + // nesting loops such that the innermost loop has the smallest stride for the + // best cache behavior. + for (int b = 0; b < ArraySize(output_dims, 3); ++b) { + 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)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] - + input2_data[SubscriptToIndex(desc2, 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, @@ -2081,6 +2401,458 @@ inline void LstmCell(const float* input_data, const Dims<4>& input_dims, output_state_map.tanh(); } +#ifdef GEMMLOWP_NEON +// In the common case of batch size 1, a fully-connected node degenerates +// to a matrix*vector product. LSTM cells contain a fully-connected node; +// when quantized, this becomes a special type of GEMV operation where +// the output is 16bit-quantized, thus needs its own special path. +inline void GEMVForLstmCell(const uint8* input_data, const Dims<4>& input_dims, + const uint8* weights_data, + const Dims<4>& weights_dims, + uint8 weights_zero_point, const int32* bias_data, + const Dims<4>& bias_dims, int32 accum_multiplier, + int accum_shift, int16* output_data, + const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("GEMVForLstmCell"); + TFLITE_DCHECK(IsPackedWithoutStrides(input_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(weights_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(bias_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(output_dims)); + TFLITE_DCHECK_EQ(ArraySize(output_dims, 1) * ArraySize(output_dims, 2) * + ArraySize(output_dims, 3), + 1); + const int input_size = input_dims.strides[3]; + const int output_size = MatchingArraySize(weights_dims, 1, output_dims, 0); + // This special fast path for quantized LSTM cells does not try to support + // odd sizes that we haven't encountered in any LSTM cell, that would + // require special code (that would go untested until any LSTM cell + // exercises it). We just guard our assumptions about size evenness with + // the following assertions. + TFLITE_DCHECK(!(output_size % 4)); + TFLITE_DCHECK(!(input_size % 8)); + const int32* bias_ptr = bias_data; + int16* output_ptr = output_data; + for (int out = 0; out < output_size; out += 4) { + int32x4_t acc_0 = vdupq_n_s32(0); + int32x4_t acc_1 = vdupq_n_s32(0); + int32x4_t acc_2 = vdupq_n_s32(0); + int32x4_t acc_3 = vdupq_n_s32(0); + const int16x8_t input_offset_vec = vdupq_n_s16(-128); + const int16x8_t weights_offset_vec = vdupq_n_s16(-weights_zero_point); + int in = 0; + // Handle 16 levels of depth at a time. + for (; in <= input_size - 16; in += 16) { + const uint8x16_t input_val_u8 = vld1q_u8(input_data + in); + const uint8* weights_ptr = weights_data + in + out * input_size; + uint8x16_t weights_val_u8_0 = vld1q_u8(weights_ptr + 0 * input_size); + uint8x16_t weights_val_u8_1 = vld1q_u8(weights_ptr + 1 * input_size); + uint8x16_t weights_val_u8_2 = vld1q_u8(weights_ptr + 2 * input_size); + uint8x16_t weights_val_u8_3 = vld1q_u8(weights_ptr + 3 * input_size); + int16x8_t input_val_0, input_val_1; + const uint8x8_t low = vget_low_u8(input_val_u8); + const uint8x8_t high = vget_high_u8(input_val_u8); + input_val_0 = vreinterpretq_s16_u16(vmovl_u8(low)); + input_val_1 = vreinterpretq_s16_u16(vmovl_u8(high)); + input_val_0 = vaddq_s16(input_val_0, input_offset_vec); + input_val_1 = vaddq_s16(input_val_1, input_offset_vec); + int16x8_t weights_val_0_0, weights_val_1_0, weights_val_2_0, + weights_val_3_0; + int16x8_t weights_val_0_1, weights_val_1_1, weights_val_2_1, + weights_val_3_1; + weights_val_0_0 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(weights_val_u8_0))), + weights_offset_vec); + weights_val_0_1 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(weights_val_u8_0))), + weights_offset_vec); + weights_val_1_0 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(weights_val_u8_1))), + weights_offset_vec); + weights_val_1_1 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(weights_val_u8_1))), + weights_offset_vec); + weights_val_2_0 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(weights_val_u8_2))), + weights_offset_vec); + weights_val_2_1 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(weights_val_u8_2))), + weights_offset_vec); + weights_val_3_0 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(weights_val_u8_3))), + weights_offset_vec); + weights_val_3_1 = vaddq_s16( + vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(weights_val_u8_3))), + weights_offset_vec); + acc_0 = vmlal_s16(acc_0, vget_low_s16(weights_val_0_0), + vget_low_s16(input_val_0)); + acc_1 = vmlal_s16(acc_1, vget_low_s16(weights_val_1_0), + vget_low_s16(input_val_0)); + acc_2 = vmlal_s16(acc_2, vget_low_s16(weights_val_2_0), + vget_low_s16(input_val_0)); + acc_3 = vmlal_s16(acc_3, vget_low_s16(weights_val_3_0), + vget_low_s16(input_val_0)); + acc_0 = vmlal_s16(acc_0, vget_high_s16(weights_val_0_0), + vget_high_s16(input_val_0)); + acc_1 = vmlal_s16(acc_1, vget_high_s16(weights_val_1_0), + vget_high_s16(input_val_0)); + acc_2 = vmlal_s16(acc_2, vget_high_s16(weights_val_2_0), + vget_high_s16(input_val_0)); + acc_3 = vmlal_s16(acc_3, vget_high_s16(weights_val_3_0), + vget_high_s16(input_val_0)); + acc_0 = vmlal_s16(acc_0, vget_low_s16(weights_val_0_1), + vget_low_s16(input_val_1)); + acc_1 = vmlal_s16(acc_1, vget_low_s16(weights_val_1_1), + vget_low_s16(input_val_1)); + acc_2 = vmlal_s16(acc_2, vget_low_s16(weights_val_2_1), + vget_low_s16(input_val_1)); + acc_3 = vmlal_s16(acc_3, vget_low_s16(weights_val_3_1), + vget_low_s16(input_val_1)); + acc_0 = vmlal_s16(acc_0, vget_high_s16(weights_val_0_1), + vget_high_s16(input_val_1)); + acc_1 = vmlal_s16(acc_1, vget_high_s16(weights_val_1_1), + vget_high_s16(input_val_1)); + acc_2 = vmlal_s16(acc_2, vget_high_s16(weights_val_2_1), + vget_high_s16(input_val_1)); + acc_3 = vmlal_s16(acc_3, vget_high_s16(weights_val_3_1), + vget_high_s16(input_val_1)); + } + // Handle 8 levels of depth at a time. + for (; in < input_size; in += 8) { + const uint8x8_t input_val_u8 = vld1_u8(input_data + in); + const uint8* weights_ptr = weights_data + in + out * input_size; + uint8x8_t weights_val_u8_0 = vld1_u8(weights_ptr + 0 * input_size); + uint8x8_t weights_val_u8_1 = vld1_u8(weights_ptr + 1 * input_size); + uint8x8_t weights_val_u8_2 = vld1_u8(weights_ptr + 2 * input_size); + uint8x8_t weights_val_u8_3 = vld1_u8(weights_ptr + 3 * input_size); + int16x8_t input_val; + input_val = vreinterpretq_s16_u16(vmovl_u8(input_val_u8)); + input_val = vaddq_s16(input_val, input_offset_vec); + int16x8_t weights_val_0, weights_val_1, weights_val_2, weights_val_3; + weights_val_0 = + vaddq_s16(vreinterpretq_s16_u16(vmovl_u8(weights_val_u8_0)), + weights_offset_vec); + weights_val_1 = + vaddq_s16(vreinterpretq_s16_u16(vmovl_u8(weights_val_u8_1)), + weights_offset_vec); + weights_val_2 = + vaddq_s16(vreinterpretq_s16_u16(vmovl_u8(weights_val_u8_2)), + weights_offset_vec); + weights_val_3 = + vaddq_s16(vreinterpretq_s16_u16(vmovl_u8(weights_val_u8_3)), + weights_offset_vec); + acc_0 = vmlal_s16(acc_0, vget_low_s16(weights_val_0), + vget_low_s16(input_val)); + acc_1 = vmlal_s16(acc_1, vget_low_s16(weights_val_1), + vget_low_s16(input_val)); + acc_2 = vmlal_s16(acc_2, vget_low_s16(weights_val_2), + vget_low_s16(input_val)); + acc_3 = vmlal_s16(acc_3, vget_low_s16(weights_val_3), + vget_low_s16(input_val)); + acc_0 = vmlal_s16(acc_0, vget_high_s16(weights_val_0), + vget_high_s16(input_val)); + acc_1 = vmlal_s16(acc_1, vget_high_s16(weights_val_1), + vget_high_s16(input_val)); + acc_2 = vmlal_s16(acc_2, vget_high_s16(weights_val_2), + vget_high_s16(input_val)); + acc_3 = vmlal_s16(acc_3, vget_high_s16(weights_val_3), + vget_high_s16(input_val)); + } + // Horizontally reduce accumulators + int32x2_t pairwise_reduced_acc_0, pairwise_reduced_acc_1, + pairwise_reduced_acc_2, pairwise_reduced_acc_3; + pairwise_reduced_acc_0 = + vpadd_s32(vget_low_s32(acc_0), vget_high_s32(acc_0)); + pairwise_reduced_acc_1 = + vpadd_s32(vget_low_s32(acc_1), vget_high_s32(acc_1)); + pairwise_reduced_acc_2 = + vpadd_s32(vget_low_s32(acc_2), vget_high_s32(acc_2)); + pairwise_reduced_acc_3 = + vpadd_s32(vget_low_s32(acc_3), vget_high_s32(acc_3)); + const int32x2_t reduced_lo = + vpadd_s32(pairwise_reduced_acc_0, pairwise_reduced_acc_1); + const int32x2_t reduced_hi = + vpadd_s32(pairwise_reduced_acc_2, pairwise_reduced_acc_3); + int32x4_t reduced = vcombine_s32(reduced_lo, reduced_hi); + // Add bias values. + int32x4_t bias_vec = vld1q_s32(bias_ptr); + bias_ptr += 4; + reduced = vaddq_s32(reduced, bias_vec); + int left_shift = accum_shift > 0 ? accum_shift : 0; + int right_shift = accum_shift > 0 ? 0 : -accum_shift; + reduced = vshlq_s32(reduced, vdupq_n_s32(left_shift)); + // Multiply by the fixed-point multiplier. + reduced = vqrdmulhq_n_s32(reduced, accum_multiplier); + // Rounding-shift-right. + using gemmlowp::RoundingDivideByPOT; + reduced = RoundingDivideByPOT(reduced, right_shift); + // Narrow values down to 16 bit signed. + const int16x4_t res16 = vqmovn_s32(reduced); + vst1_s16(output_ptr, res16); + output_ptr += 4; + } +} +#endif + +// 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::GemmContext* gemm_context) { + 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. + bool gemm_already_performed = false; +#ifdef GEMMLOWP_NEON + if (fc_batches == 1 && !(fc_output_depth % 4) && !(fc_accum_depth % 8)) { + GEMVForLstmCell(concat_temp_data_uint8, concat_temp_dims, + weights_data_uint8, weights_dims, weights_zero_point, + bias_data_int32, bias_dims, accum_multiplier, accum_shift, + activ_temp_data_int16, activ_temp_dims); + gemm_already_performed = true; + } +#endif + if (!gemm_already_performed) { + gemmlowp::MatrixMap + weights_matrix(weights_data_uint8, fc_output_depth, fc_accum_depth); + gemmlowp::MatrixMap input_matrix( + concat_temp_data_uint8, fc_accum_depth, fc_batches); + gemmlowp::MatrixMap output_matrix( + activ_temp_data_int16, fc_output_depth, fc_batches); + typedef gemmlowp::VectorMap + ColVectorMap; + ColVectorMap bias_vector(bias_data_int32, fc_output_depth); + gemmlowp::OutputStageBiasAddition bias_addition_stage; + bias_addition_stage.bias_vector = bias_vector; + gemmlowp::OutputStageScaleInt32ByFixedPointAndExponent scale_stage; + scale_stage.result_offset_after_shift = 0; + scale_stage.result_fixedpoint_multiplier = accum_multiplier; + scale_stage.result_exponent = accum_shift; + gemmlowp::OutputStageSaturatingCastToInt16 saturating_cast_int16_stage; + auto output_pipeline = std::make_tuple(bias_addition_stage, scale_stage, + saturating_cast_int16_stage); + gemmlowp::GemmWithOutputPipeline< + uint8, int16, gemmlowp::L8R8WithLhsNonzeroBitDepthParams>( + gemm_context, weights_matrix, input_matrix, &output_matrix, + -weights_zero_point, -128, output_pipeline); + } + + // 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; + const int16* input_gate_input_ptr = activ_temp_data_int16; + const int16* input_modulation_gate_input_ptr = + activ_temp_data_int16 + output_depth; + const int16* forget_gate_input_ptr = activ_temp_data_int16 + 2 * output_depth; + const int16* output_gate_input_ptr = activ_temp_data_int16 + 3 * output_depth; + const int16* prev_state_ptr = prev_state_data_int16; + int16* output_state_data_ptr = output_state_data_int16; + uint8* output_activ_data_ptr = output_activ_data_uint8; + + for (int b = 0; b < outer_size; ++b) { + int c = 0; +#ifdef GEMMLOWP_NEON + for (; c <= output_depth - 8; c += 8) { + // 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(vld1q_s16(input_gate_input_ptr)); + input_gate_input_ptr += 8; + 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(vld1q_s16(input_modulation_gate_input_ptr)); + input_modulation_gate_input_ptr += 8; + 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(vld1q_s16(forget_gate_input_ptr)); + forget_gate_input_ptr += 8; + F0 forget_gate_output = gemmlowp::logistic(forget_gate_input); + // Implementation of output gate, using fixed-point logistic function. + F3 output_gate_input = F3::FromRaw(vld1q_s16(output_gate_input_ptr)); + output_gate_input_ptr += 8; + 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(vld1q_s16(prev_state_ptr)); + prev_state_ptr += 8; + 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. + // Since a Tanh fixed-point implementation is specialized for a given + // number or integer bits, and each specialization can have a substantial + // code size, and we already used above a Tanh on an input with 3 integer + // bits, and per the table in the above function comment there is no + // significant accuracy to be lost by clamping to [-8, +8] for a + // 3-integer-bits representation, let us just do that. This helps people + // porting this to targets where code footprint must be minimized. + F3 new_state_f3 = gemmlowp::Rescale<3>(new_state); + F0 output_activ_int16 = output_gate_output * gemmlowp::tanh(new_state_f3); + // Store the new internal state back to memory, as 16-bit integers. + // Note: here we store the original value with StateIntegerBits, not + // the rescaled 3-integer-bits value fed to tanh. + vst1q_s16(output_state_data_ptr, new_state.raw()); + output_state_data_ptr += 8; + // Down-scale the output activations to 8-bit integers, saturating, + // and store back to memory. + int16x8_t rescaled_output_activ = + gemmlowp::RoundingDivideByPOT(output_activ_int16.raw(), 8); + int8x8_t int8_output_activ = vqmovn_s16(rescaled_output_activ); + uint8x8_t uint8_output_activ = + vadd_u8(vdup_n_u8(128), vreinterpret_u8_s8(int8_output_activ)); + vst1_u8(output_activ_data_ptr, uint8_output_activ); + output_activ_data_ptr += 8; + } +#endif + for (; 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(*input_gate_input_ptr++); + 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(*input_modulation_gate_input_ptr++); + 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(*forget_gate_input_ptr++); + F0 forget_gate_output = gemmlowp::logistic(forget_gate_input); + // Implementation of output gate, using fixed-point logistic function. + F3 output_gate_input = F3::FromRaw(*output_gate_input_ptr++); + 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_ptr++); + 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. + // Since a Tanh fixed-point implementation is specialized for a given + // number or integer bits, and each specialization can have a substantial + // code size, and we already used above a Tanh on an input with 3 integer + // bits, and per the table in the above function comment there is no + // significant accuracy to be lost by clamping to [-8, +8] for a + // 3-integer-bits representation, let us just do that. This helps people + // porting this to targets where code footprint must be minimized. + F3 new_state_f3 = gemmlowp::Rescale<3>(new_state); + F0 output_activ_int16 = output_gate_output * gemmlowp::tanh(new_state_f3); + // Store the new internal state back to memory, as 16-bit integers. + // Note: here we store the original value with StateIntegerBits, not + // the rescaled 3-integer-bits value fed to tanh. + *output_state_data_ptr++ = 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_ptr++ = 128 + clamped_output_activ; + } + input_gate_input_ptr += 3 * output_depth; + input_modulation_gate_input_ptr += 3 * output_depth; + forget_gate_input_ptr += 3 * output_depth; + output_gate_input_ptr += 3 * output_depth; + } +} + template void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims, int outputs_count, Scalar* const* output_data, @@ -2706,74 +3478,231 @@ 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); + + const int outer_size = batches * height * width; + + 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); + + 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_ptr[c] = std::max(std::min(unsat_output, 255), 0); + + } else { + output_data_ptr[c] = 0; + } + } + } + } +} + +// TODO(myenik): This is the same as the reference implementation, not actually +// optimized yet. +inline void LogSoftmax(const float* input_data, const Dims<4>& input_dims, + float* output_data, const Dims<4>& output_dims) { 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 y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + // Find max element value which we'll use to ensure numerical stability + // taking advantage of the following equality: + // log(exp(x[i])/sum(exp(x[i]))) == log(exp(x[i]+C)/sum(exp(x[i]+C))) + float max = std::numeric_limits::lowest(); for (int c = 0; c < depth; ++c) { - max_in_row = - std::max(max_in_row, input_data[Offset(input_dims, c, x, y, b)]); + max = std::max(max, input_data[Offset(input_dims, c, x, y, b)]); } - FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); + // Compute sum. + float sum = 0.f; 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)); - } + sum += std::exp(input_data[Offset(input_dims, c, x, y, b)] - max); } - 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)); - + // Compute result. + const float log_sum = std::log(sum); 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[Offset(output_dims, c, x, y, b)] = + input_data[Offset(input_dims, c, x, y, b)] - max - log_sum; } } } @@ -2930,6 +3859,28 @@ inline void Logistic(const uint8* input_data, const Dims<4>& input_dims, } } +inline void Logistic(const int16* input_data, const Dims<4>& input_dims, + int16* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Logistic/Int16"); + // This is a copy of the reference implementation. We do not currently have a + // properly optimized version. + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input_dims), flat_size); + + for (int i = 0; i < flat_size; i++) { + // 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], the input range expected here. + using F3 = gemmlowp::FixedPoint; + + const F3 input = F3::FromRaw(input_data[i]); + F0 output = gemmlowp::logistic(input); + output_data[i] = output.raw(); + } +} + inline void Tanh(const float* input_data, const Dims<4>& input_dims, float* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("Tanh"); @@ -2938,6 +3889,195 @@ 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 Tanh(const int16* input_data, const Dims<4>& input_dims, + int input_left_shift, int16* output_data, + const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Tanh/Int16"); + // This is a copy of the reference implementation. We do not currently have a + // properly optimized version. + + // Support for shifts is limited until we have a parameterized version of + // SaturatingRoundingMultiplyByPOT(). + TFLITE_DCHECK_GE(input_left_shift, 0); + TFLITE_DCHECK_LE(input_left_shift, 1); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input_dims), flat_size); + + // 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], the input range expected here. + using F3 = gemmlowp::FixedPoint; + + if (input_left_shift == 0) { + for (int i = 0; i < flat_size; i++) { + F3 input = F3::FromRaw(input_data[i]); + F0 output = gemmlowp::tanh(input); + output_data[i] = output.raw(); + } + } else { + for (int i = 0; i < flat_size; i++) { + F3 input = F3::FromRaw( + gemmlowp::SaturatingRoundingMultiplyByPOT<1>(input_data[i])); + F0 output = gemmlowp::tanh(input); + output_data[i] = output.raw(); + } + } +} + 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) { @@ -3410,7 +4550,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); @@ -3425,13 +4565,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, @@ -3440,6 +4587,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, @@ -3829,6 +4985,107 @@ 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, const 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)]; + } + } + } + } +} + +inline void TransposeConv(const float* input_data, const Dims<4>& input_dims, + const float* filter_data, const Dims<4>& filter_dims, + int stride_width, int stride_height, int pad_width, + int pad_height, float* output_data, + const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("TransposeConv"); + // THIS FUNCTION IS A COPY FROM reference_ops.h. + // To optimize, start by using the conv code with transposed weights for the + // case of stride_height = stride_width = 1. + const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); + const int input_depth = MatchingArraySize(input_dims, 0, filter_dims, 3); + const int output_depth = MatchingArraySize(filter_dims, 0, output_dims, 0); + const int input_height = ArraySize(input_dims, 2); + const int input_width = ArraySize(input_dims, 1); + const int filter_height = ArraySize(filter_dims, 2); + const int filter_width = ArraySize(filter_dims, 1); + const int output_height = ArraySize(output_dims, 2); + const int output_width = ArraySize(output_dims, 1); + + // Although transpose convolution simplifies to convolution with transposed + // weights for strides of 1, non-unitary striding complicates matters. To + // keep this reference implementation as clear as possible, we use a "scatter" + // access pattern, where we loop through all the input elements, computing + // their influence on the output, rather than looping through the output + // elements in the typical "gather" access pattern of a conv. We therefore + // must initialize the output array to zero. + for (int batch = 0; batch < batches; ++batch) { + for (int out_y = 0; out_y < output_height; ++out_y) { + for (int out_x = 0; out_x < output_width; ++out_x) { + for (int out_channel = 0; out_channel < output_depth; ++out_channel) { + output_data[Offset(output_dims, out_channel, out_x, out_y, batch)] = + 0.0f; + } + } + } + } + + // Loop through input elements one at a time. + for (int batch = 0; batch < batches; ++batch) { + for (int in_y = 0; in_y < input_height; ++in_y) { + for (int in_x = 0; in_x < input_width; ++in_x) { + for (int in_channel = 0; in_channel < input_depth; ++in_channel) { + // Loop through the output elements it will influence + const int out_x_origin = (in_x * stride_width) - pad_width; + const int out_y_origin = (in_y * stride_height) - pad_height; + for (int filter_y = 0; filter_y < filter_height; ++filter_y) { + for (int filter_x = 0; filter_x < filter_width; ++filter_x) { + for (int out_channel = 0; out_channel < input_depth; + ++out_channel) { + // Compute output element location + const int out_x = out_x_origin + filter_x; + const int out_y = out_y_origin + filter_y; + // We cannot accumulate out of bounds + if ((out_x >= 0) && (out_x < output_width) && (out_y >= 0) && + (out_y < output_height)) { + float input_value = input_data[Offset(input_dims, in_channel, + in_x, in_y, batch)]; + float filter_value = + filter_data[Offset(filter_dims, out_channel, filter_x, + filter_y, in_channel)]; + output_data[Offset(output_dims, out_channel, out_x, out_y, + batch)] += input_value * filter_value; + } + } + } + } + } + } + } + } +} + } // namespace optimized_ops } // namespace tflite 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..f7706c793883f012d2b66cf3d3167a59afe31f91 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.h +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util.h @@ -12,15 +12,94 @@ WITHOUT 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 PHOTOS_VISION_LEARNING_TENSORFLOW_MINI_QUANTIZATION_UTIL_H_ -#define PHOTOS_VISION_LEARNING_TENSORFLOW_MINI_QUANTIZATION_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ +#include #include +#include + +#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" +#include "tensorflow/contrib/lite/kernels/internal/round.h" +#include "tensorflow/contrib/lite/kernels/internal/types.h" namespace tflite { +// Given the min and max values of a float array, return +// reasonable quantization parameters to use for this array. +template +QuantizationParams ChooseQuantizationParams(double rmin, double rmax) { + const T qmin = std::numeric_limits::min(); + const T qmax = std::numeric_limits::max(); + const double qmin_double = qmin; + const double qmax_double = qmax; + // 0 should always be a representable value. Let's assume that the initial + // min,max range contains 0. + TFLITE_CHECK_LE(rmin, 0.); + TFLITE_CHECK_GE(rmax, 0.); + if (rmin == rmax) { + // Special case where the min,max range is a point. Should be {0}. + TFLITE_CHECK_EQ(rmin, 0.); + TFLITE_CHECK_EQ(rmax, 0.); + QuantizationParams quantization_params; + quantization_params.zero_point = 0; + quantization_params.scale = 0.; + return quantization_params; + } + + // General case. + // + // First determine the scale. + const double scale = (rmax - rmin) / (qmax_double - qmin_double); + + // Zero-point computation. + // First the initial floating-point computation. The zero-point can be + // determined from solving an affine equation for any known pair + // (real value, corresponding quantized value). + // We know two such pairs: (rmin, qmin) and (rmax, qmax). + // The arithmetic error on the zero point computed from either pair + // will be roughly machine_epsilon * (sum of absolute values of terms) + // so we want to use the variant that adds the smaller terms. + const double zero_point_from_min = qmin_double - rmin / scale; + const double zero_point_from_max = qmax_double - rmax / scale; + const double zero_point_from_min_error = + std::abs(qmin_double) + std::abs(rmin / scale); + const double zero_point_from_max_error = + std::abs(qmax_double) + std::abs(rmax / scale); + + const double zero_point_double = + zero_point_from_min_error < zero_point_from_max_error + ? zero_point_from_min + : zero_point_from_max; + + // Now we need to nudge the zero point to be an integer + // (our zero points are integer, and this is motivated by the requirement + // to be able to represent the real value "0" exactly as a quantized value, + // which is required in multiple places, for example in Im2col with SAME + // padding). + T nudged_zero_point = 0; + if (zero_point_double < qmin_double) { + nudged_zero_point = qmin; + } else if (zero_point_double > qmax_double) { + nudged_zero_point = qmax; + } else { + nudged_zero_point = static_cast(round(zero_point_double)); + } + // The zero point should always be in the range of quantized value, + // [qmin, qmax]. + TFLITE_CHECK_GE(nudged_zero_point, qmin); + TFLITE_CHECK_LE(nudged_zero_point, qmax); + + // Finally, store the result nudged quantization params. + QuantizationParams quantization_params; + quantization_params.zero_point = nudged_zero_point; + quantization_params.scale = scale; + return quantization_params; +} + // 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 +114,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 @@ -52,4 +141,4 @@ int CalculateInputRadius(int input_integer_bits, int input_left_shift); } // namespace tflite -#endif // PHOTOS_VISION_LEARNING_TENSORFLOW_MINI_QUANTIZATION_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc b/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc index d6f306e2cbae3c780b3d773638ba46cd2abf02f5..4ae2085c30ed08790eaff27c7921909d47687707 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc @@ -22,6 +22,51 @@ namespace { using ::testing::Pair; +// Example taken from http://www.tensorflow.org/performance/quantization +// +// Quantized | Float +// --------- | ----- +// 0 | -10.0 +// 255 | 30.0 +// 128 | 10.0 +TEST(QuantizationUtilTest, ChooseQuantizationParams) { + QuantizationParams qp = ChooseQuantizationParams(-10.0, 30.0); + EXPECT_NEAR(qp.scale, 0.156863, 1e-5); + EXPECT_EQ(qp.zero_point, 64); +} + +TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroPointOnMinBoundary) { + QuantizationParams qp = ChooseQuantizationParams(0.0, 30.0); + EXPECT_NEAR(qp.scale, 0.117647, 1e-5); + EXPECT_EQ(qp.zero_point, 0); +} + +TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroNotInRange) { + // Assumption is that zero is within the range. + EXPECT_DEATH(ChooseQuantizationParams(10.0, 30.0), ""); +} + +TEST(QuantizationUtilTest, ChooseQuantizationParamsEmptyRangePositive) { + // Assumption is that zero is within the range. + EXPECT_DEATH(ChooseQuantizationParams(30.0, 30.0), ""); +} + +TEST(QuantizationUtilTest, ChooseQuantizationParamsEmptyRangeZero) { + QuantizationParams qp = ChooseQuantizationParams(0.0, 0.0); + EXPECT_NEAR(qp.scale, 0.0, 1e-5); + EXPECT_EQ(qp.zero_point, 0); +} + +TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroPointOnMaxBoundary) { + QuantizationParams qp = ChooseQuantizationParams(-10.0, 0.0); + EXPECT_NEAR(qp.scale, 0.039216, 1e-5); + EXPECT_EQ(qp.zero_point, 255); +} + +TEST(QuantizationUtilTest, ChooseQuantizationParamsInvalidRange) { + EXPECT_DEATH(ChooseQuantizationParams(10.0, -30.0), ""); +} + TEST(QuantizationUtilTest, QuantizeMultiplierSmallerThanOne) { auto quantize = [](double d) { int32_t q; @@ -31,7 +76,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/portable_tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h index afc3e26e7988a369fb777ae99c08c4e98f26ebb8..c05c21b472b05f2cbe133adf94d91ab0c6d9ef40 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h @@ -15,7 +15,7 @@ limitations under the License. #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" diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h index 4bcf4993e92136388b3febe7e04030943f4b54f3..527276f7bdaa98608592cb20131c287d4b39886b 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h @@ -157,11 +157,11 @@ inline void NdArrayDescsForElementwiseBroadcast(const Dims& input0_dims, 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) { + int stride_width, int stride_height, int dilation_width_factor, + int dilation_height_factor, 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) { (void)im2col_data; // only used in optimized code. (void)im2col_dims; // only used in optimized code. const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); @@ -186,8 +186,9 @@ inline void Conv(const float* input_data, const Dims<4>& input_dims, for (int filter_y = 0; filter_y < filter_height; ++filter_y) { for (int filter_x = 0; filter_x < filter_width; ++filter_x) { for (int in_channel = 0; in_channel < input_depth; ++in_channel) { - const int in_x = in_x_origin + filter_x; - const int in_y = in_y_origin + filter_y; + const int in_x = in_x_origin + dilation_width_factor * filter_x; + const int in_y = + in_y_origin + dilation_height_factor * filter_y; // If the location is outside the bounds of the input image, // use zero as a default value. if ((in_x >= 0) && (in_x < input_width) && (in_y >= 0) && @@ -216,6 +217,23 @@ inline void Conv(const float* input_data, const Dims<4>& input_dims, } } +template +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 dilation_width_factor, + int dilation_height_factor, int pad_width, int pad_height, + float* output_data, const Dims<4>& output_dims, float* im2col_data, + const Dims<4>& im2col_dims) { + float output_activation_min, output_activation_max; + GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); + Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims, + stride_width, stride_height, dilation_width_factor, + dilation_height_factor, pad_width, pad_height, output_activation_min, + output_activation_max, output_data, output_dims, im2col_data, + im2col_dims); +} + // legacy, for compatibility with old checked-in code template void Conv(const float* input_data, const Dims<4>& input_dims, @@ -227,7 +245,7 @@ void Conv(const float* input_data, const Dims<4>& input_dims, float output_activation_min, output_activation_max; GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims, - stride_width, stride_height, pad_width, pad_height, + stride_width, stride_height, 1, 1, pad_width, pad_height, output_activation_min, output_activation_max, output_data, output_dims, im2col_data, im2col_dims); } @@ -241,7 +259,7 @@ void Conv(const float* input_data, const Dims<4>& input_dims, const Dims<4>& output_dims, float* im2col_data, const Dims<4>& im2col_dims) { Conv(input_data, input_dims, filter_data, filter_dims, bias_data, - bias_dims, stride, stride, pad_width, pad_height, output_data, + bias_dims, stride, stride, 1, 1, pad_width, pad_height, output_data, output_dims, im2col_data, im2col_dims); } @@ -533,6 +551,55 @@ inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, } } +inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, + int32 input_offset, const uint8* filter_data, + const Dims<4>& filter_dims, int32 filter_offset, + const int32* bias_data, const Dims<4>& bias_dims, + int32 output_offset, int32 output_multiplier, + int output_shift, int32 output_activation_min, + int32 output_activation_max, int16* output_data, + const Dims<4>& output_dims, + gemmlowp::GemmContext* gemm_context) { + (void)gemm_context; // only used in optimized code. + TFLITE_DCHECK_LE(output_activation_min, output_activation_max); + TFLITE_DCHECK_EQ(output_offset, 0); + // TODO(benoitjacob): This really should be: + // const int batches = ArraySize(output_dims, 1); + // but the current --variable_batch hack consists in overwriting the 3rd + // dimension with the runtime batch size, as we don't keep track for each + // array of which dimension is the batch dimension in it. + const int batches = ArraySize(output_dims, 1) * ArraySize(output_dims, 2) * + ArraySize(output_dims, 3); + const int output_depth = MatchingArraySize(filter_dims, 1, output_dims, 0); + const int accum_depth = ArraySize(filter_dims, 0); + TFLITE_DCHECK(IsPackedWithoutStrides(input_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(filter_dims)); + for (int b = 0; b < batches; ++b) { + for (int out_c = 0; out_c < output_depth; ++out_c) { + // Internal accumulation. + // Initialize accumulator with the bias-value. + int32 accum = bias_data[out_c]; + // Accumulation loop. + for (int d = 0; d < accum_depth; ++d) { + int16 input_val = input_data[b * accum_depth + d] + input_offset; + int16 filter_val = filter_data[out_c * accum_depth + d] + filter_offset; + accum += filter_val * input_val; + } + // Down-scale the final int32 accumulator to the scale used by our + // (16-bit, typically 3 integer bits) fixed-point format. The quantized + // multiplier and shift here have been pre-computed offline + // (e.g. by toco). + accum = MultiplyByQuantizedMultiplier(accum, output_multiplier, + -output_shift); + // Saturate, cast to int16, and store to output array. + accum = std::max(accum, output_activation_min - output_offset); + accum = std::min(accum, output_activation_max - output_offset); + accum += output_offset; + output_data[out_c + output_depth * b] = accum; + } + } +} + // legacy, for compatibility with old checked-in code template void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, @@ -885,6 +952,36 @@ inline void Add(int left_shift, const uint8* input1_data, } } +template +inline void Add(const int16* input1_data, const Dims<4>& input1_dims, + int input1_shift, const int16* input2_data, + const Dims<4>& input2_dims, int input2_shift, + int16* output_data, const Dims<4>& output_dims) { + static_assert(Ac == FusedActivationFunctionType::kNone, ""); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input1_dims), flat_size); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input2_dims), flat_size); + + TFLITE_DCHECK(input1_shift == 0 || input2_shift == 0); + TFLITE_DCHECK_GE(input1_shift, 0); + TFLITE_DCHECK_GE(input2_shift, 0); + const int16* not_shift_input = input1_shift == 0 ? input1_data : input2_data; + const int16* shift_input = input1_shift == 0 ? input2_data : input1_data; + const int input_shift = input1_shift == 0 ? input2_shift : input1_shift; + + for (int i = 0; i < flat_size; i++) { + // F0 uses 0 integer bits, range [-1, 1]. + using F0 = gemmlowp::FixedPoint; + + F0 input_ready_scaled = F0::FromRaw(not_shift_input[i]); + F0 scaled_input = + F0::FromRaw(gemmlowp::RoundingDivideByPOT(shift_input[i], input_shift)); + F0 result = gemmlowp::SaturatingAdd(scaled_input, input_ready_scaled); + output_data[i] = result.raw(); + } +} + // TODO(jiawen): We can implement BroadcastAdd on buffers of arbitrary // dimensionality if the runtime code does a single loop over one dimension // that handles broadcasting as the base case. The code generator would then @@ -1166,6 +1263,53 @@ inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, } } +inline void Mul(const int16* input1_data, const Dims<4>& input1_dims, + const int16* input2_data, const Dims<4>& input2_dims, + int16* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Mul/Int16"); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input1_dims), flat_size); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input2_dims), flat_size); + + for (int i = 0; i < flat_size; i++) { + // F0 uses 0 integer bits, range [-1, 1]. + using F0 = gemmlowp::FixedPoint; + + F0 unclamped_result = + F0::FromRaw(input1_data[i]) * F0::FromRaw(input2_data[i]); + output_data[i] = unclamped_result.raw(); + } +} + +inline void Mul(const int16* input1_data, const Dims<4>& input1_dims, + const int16* input2_data, const Dims<4>& input2_dims, + int32 output_offset, int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("Mul/Int16Uint8"); + TFLITE_DCHECK_LE(output_activation_min, output_activation_max); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input1_dims), flat_size); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input2_dims), flat_size); + + for (int i = 0; i < flat_size; i++) { + // F0 uses 0 integer bits, range [-1, 1]. + using F0 = gemmlowp::FixedPoint; + + F0 unclamped_result = + F0::FromRaw(input1_data[i]) * F0::FromRaw(input2_data[i]); + int16 rescaled_result = + gemmlowp::RoundingDivideByPOT(unclamped_result.raw(), 8); + int16 clamped_result = + std::min(output_activation_max - output_offset, rescaled_result); + clamped_result = + std::max(output_activation_min - output_offset, clamped_result); + output_data[i] = output_offset + clamped_result; + } +} + // legacy, for compatibility with old checked-in code template inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, @@ -1181,6 +1325,47 @@ inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, output_data, output_dims); } +// TODO(jiawen): We can implement BroadcastDiv on buffers of arbitrary +// 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 BroadcastDiv(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("BroadcastDiv"); + + NdArrayDesc<4> desc1; + NdArrayDesc<4> desc2; + NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2); + + // In Tensorflow, the dimensions are canonically named (batch_number, row, + // col, channel), with extents (batches, height, width, depth), with the + // trailing dimension changing most rapidly (channels has the smallest + // stride, typically 1 element). + // + // In generated C code, we store arrays with the dimensions reversed. The + // first dimension has smallest stride. + // + // We name our variables by their Tensorflow convention, but generate C code + // nesting loops such that the innermost loop has the smallest stride for + // the best cache behavior. + for (int b = 0; b < ArraySize(output_dims, 3); ++b) { + 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)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] / + input2_data[SubscriptToIndex(desc2, c, x, y, b)], + output_activation_min, output_activation_max); + } + } + } + } +} + 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, @@ -1235,6 +1420,47 @@ inline void Sub(const float* input1_data, const Dims<4>& input1_dims, } } +// TODO(jiawen): We can implement BroadcastSub on buffers of arbitrary +// 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 BroadcastSub(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("BroadcastSub"); + + NdArrayDesc<4> desc1; + NdArrayDesc<4> desc2; + NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2); + + // In Tensorflow, the dimensions are canonically named (batch_number, row, + // col, channel), with extents (batches, height, width, depth), with the + // trailing dimension changing most rapidly (channels has the smallest stride, + // typically 1 element). + // + // In generated C code, we store arrays with the dimensions reversed. The + // first dimension has smallest stride. + // + // We name our variables by their Tensorflow convention, but generate C code + // nesting loops such that the innermost loop has the smallest stride for the + // best cache behavior. + for (int b = 0; b < ArraySize(output_dims, 3); ++b) { + 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)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] - + input2_data[SubscriptToIndex(desc2, 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, @@ -1358,6 +1584,278 @@ 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, + gemmlowp::GemmContext* gemm_context) { + (void)gemm_context; // only used in optimized code. + + // 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. + // Since a Tanh fixed-point implementation is specialized for a given + // number or integer bits, and each specialization can have a substantial + // code size, and we already used above a Tanh on an input with 3 integer + // bits, and per the table in the above function comment there is no + // significant accuracy to be lost by clamping to [-8, +8] for a + // 3-integer-bits representation, let us just do that. This helps people + // porting this to targets where code footprint must be minimized. + F3 new_state_f3 = gemmlowp::Rescale<3>(new_state); + F0 output_activ_int16 = output_gate_output * gemmlowp::tanh(new_state_f3); + // Store the new internal state back to memory, as 16-bit integers. + // Note: here we store the original value with StateIntegerBits, not + // the rescaled 3-integer-bits value fed to tanh. + 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, @@ -1368,28 +1866,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! @@ -1960,6 +2442,41 @@ inline void Softmax(const uint8* input_data, const Dims<4>& input_dims, } } +inline void LogSoftmax(const float* input_data, const Dims<4>& input_dims, + float* output_data, const Dims<4>& output_dims) { + 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) { + // Find max element value which we'll use to ensure numerical stability + // taking advantage of the following equality: + // log(exp(x[i])/sum(exp(x[i]))) == log(exp(x[i]+C)/sum(exp(x[i]+C))) + float max = std::numeric_limits::lowest(); + for (int c = 0; c < depth; ++c) { + max = std::max(max, input_data[Offset(input_dims, c, x, y, b)]); + } + + // Compute sum. + float sum = 0.f; + for (int c = 0; c < depth; ++c) { + sum += std::exp(input_data[Offset(input_dims, c, x, y, b)] - max); + } + + // Compute result. + const float log_sum = std::log(sum); + for (int c = 0; c < depth; ++c) { + output_data[Offset(output_dims, c, x, y, b)] = + input_data[Offset(input_dims, c, x, y, b)] - max - log_sum; + } + } + } + } +} + inline void Logistic(const float* input_data, const Dims<4>& input_dims, float* output_data, const Dims<4>& output_dims) { const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); @@ -2008,11 +2525,13 @@ inline void Logistic(const uint8* input_data, const Dims<4>& input_dims, const FixedPoint4 input_val_f4 = FixedPoint4::FromRaw(input_val_rescaled); const FixedPoint0 output_val_f0 = gemmlowp::logistic(input_val_f4); + // Convert from Q0.31 to Q23.8. using gemmlowp::RoundingDivideByPOT; int32 output_val_s32 = RoundingDivideByPOT(output_val_f0.raw(), 23); if (output_val_s32 == 256) { output_val_s32 = 255; } + // Reinterpret as U0.8. TFLITE_DCHECK_GE(output_val_s32, 0); TFLITE_DCHECK_LE(output_val_s32, 255); output_val = static_cast(output_val_s32); @@ -2024,6 +2543,25 @@ inline void Logistic(const uint8* input_data, const Dims<4>& input_dims, } } +inline void Logistic(const int16* input_data, const Dims<4>& input_dims, + int16* output_data, const Dims<4>& output_dims) { + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input_dims), flat_size); + + for (int i = 0; i < flat_size; i++) { + // 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], the input range expected here. + using F3 = gemmlowp::FixedPoint; + + const F3 input = F3::FromRaw(input_data[i]); + F0 output = gemmlowp::logistic(input); + output_data[i] = output.raw(); + } +} + inline void Tanh(const float* input_data, const Dims<4>& input_dims, float* output_data, const Dims<4>& output_dims) { const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); @@ -2043,6 +2581,89 @@ 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); + // Convert from Q0.31 to Q24.7. + 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; + } + // Reinterpret as Q0.7, encoded in uint8. + 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 Tanh(const int16* input_data, const Dims<4>& input_dims, + int input_left_shift, int16* output_data, + const Dims<4>& output_dims) { + // Support for shifts is limited until we have a parameterized version of + // SaturatingRoundingMultiplyByPOT(). + TFLITE_DCHECK_GE(input_left_shift, 0); + TFLITE_DCHECK_LE(input_left_shift, 1); + + const int flat_size = RequiredBufferSizeForDims(output_dims); + TFLITE_DCHECK_EQ(RequiredBufferSizeForDims(input_dims), flat_size); + + // 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], the input range expected here. + using F3 = gemmlowp::FixedPoint; + + if (input_left_shift == 0) { + for (int i = 0; i < flat_size; i++) { + F3 input = F3::FromRaw(input_data[i]); + F0 output = gemmlowp::tanh(input); + output_data[i] = output.raw(); + } + } else { + for (int i = 0; i < flat_size; i++) { + F3 input = F3::FromRaw( + gemmlowp::SaturatingRoundingMultiplyByPOT<1>(input_data[i])); + F0 output = gemmlowp::tanh(input); + output_data[i] = output.raw(); + } + } +} + 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) { @@ -2202,7 +2823,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); @@ -2216,6 +2837,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) { @@ -2243,6 +2870,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, @@ -2467,6 +3103,14 @@ 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, @@ -2519,9 +3163,11 @@ inline void Mean(T* input_data, const int* input_dims, const int input_num_dims, 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); + if (num_elements_in_axis > 0) { + for (size_t idx = 0; idx < num_outputs; ++idx) { + output_data[idx] = static_cast(static_cast(output_data[idx]) / + num_elements_in_axis); + } } } @@ -2702,6 +3348,67 @@ void Transpose(const T* input, const Dims<4>& input_dims, T* output, } } +inline void TransposeConv(const float* input_data, const Dims<4>& input_dims, + const float* filter_data, const Dims<4>& filter_dims, + int stride_width, int stride_height, int pad_width, + int pad_height, float* output_data, + const Dims<4>& output_dims) { + const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); + const int input_depth = MatchingArraySize(input_dims, 0, filter_dims, 3); + const int output_depth = MatchingArraySize(filter_dims, 0, output_dims, 0); + const int input_height = ArraySize(input_dims, 2); + const int input_width = ArraySize(input_dims, 1); + const int filter_height = ArraySize(filter_dims, 2); + const int filter_width = ArraySize(filter_dims, 1); + const int output_height = ArraySize(output_dims, 2); + const int output_width = ArraySize(output_dims, 1); + + // Although transpose convolution simplifies to convolution with transposed + // weights for strides of 1, non-unitary striding complicates matters. To + // keep this reference implementation as clear as possible, we use a "scatter" + // access pattern, where we loop through all the input elements, computing + // their influence on the output, rather than looping through the output + // elements in the typical "gather" access pattern of a conv. We therefore + // must initialize the output array to zero. + for (int i = 0; i < RequiredBufferSizeForDims(output_dims); i++) { + output_data[i] = 0.0f; + } + + // Loop through input elements one at a time. + for (int batch = 0; batch < batches; ++batch) { + for (int in_y = 0; in_y < input_height; ++in_y) { + for (int in_x = 0; in_x < input_width; ++in_x) { + for (int in_channel = 0; in_channel < input_depth; ++in_channel) { + // Loop through the output elements it will influence + const int out_x_origin = (in_x * stride_width) - pad_width; + const int out_y_origin = (in_y * stride_height) - pad_height; + for (int filter_y = 0; filter_y < filter_height; ++filter_y) { + for (int filter_x = 0; filter_x < filter_width; ++filter_x) { + for (int out_channel = 0; out_channel < output_depth; + ++out_channel) { + // Compute output element location + const int out_x = out_x_origin + filter_x; + const int out_y = out_y_origin + filter_y; + // We cannot accumulate out of bounds + if ((out_x >= 0) && (out_x < output_width) && (out_y >= 0) && + (out_y < output_height)) { + float input_value = input_data[Offset(input_dims, in_channel, + in_x, in_y, batch)]; + float filter_value = + filter_data[Offset(filter_dims, out_channel, filter_x, + filter_y, in_channel)]; + output_data[Offset(output_dims, out_channel, out_x, out_y, + batch)] += input_value * filter_value; + } + } + } + } + } + } + } + } +} + } // namespace reference_ops } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/spectrogram.cc b/tensorflow/contrib/lite/kernels/internal/spectrogram.cc new file mode 100644 index 0000000000000000000000000000000000000000..4eddf7bf0a2cbca695dae20ba8ba56a9cd72e4ba --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/spectrogram.cc @@ -0,0 +1,244 @@ +/* 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/spectrogram.h" + +#include +#include + +#include "third_party/fft2d/fft.h" + +namespace tflite { +namespace internal { + +using std::complex; + +namespace { +// Returns the default Hann window function for the spectrogram. +void GetPeriodicHann(int window_length, std::vector* window) { + // Some platforms don't have M_PI, so define a local constant here. + const double pi = std::atan(1) * 4; + window->resize(window_length); + for (int i = 0; i < window_length; ++i) { + (*window)[i] = 0.5 - 0.5 * cos((2 * pi * i) / window_length); + } +} +} // namespace + +bool Spectrogram::Initialize(int window_length, int step_length) { + std::vector window; + GetPeriodicHann(window_length, &window); + return Initialize(window, step_length); +} + +inline int Log2Floor(uint n) { + if (n == 0) return -1; + int log = 0; + uint value = n; + for (int i = 4; i >= 0; --i) { + int shift = (1 << i); + uint x = value >> shift; + if (x != 0) { + value = x; + log += shift; + } + } + return log; +} + +inline int Log2Ceiling(uint n) { + int floor = Log2Floor(n); + if (n == (n & ~(n - 1))) // zero or a power of two + return floor; + else + return floor + 1; +} + +inline uint NextPowerOfTwo(uint value) { + int exponent = Log2Ceiling(value); + // DCHECK_LT(exponent, std::numeric_limits::digits); + return 1 << exponent; +} + +bool Spectrogram::Initialize(const std::vector& window, + int step_length) { + window_length_ = window.size(); + window_ = window; // Copy window. + if (window_length_ < 2) { + // LOG(ERROR) << "Window length too short."; + initialized_ = false; + return false; + } + + step_length_ = step_length; + if (step_length_ < 1) { + // LOG(ERROR) << "Step length must be positive."; + initialized_ = false; + return false; + } + + fft_length_ = NextPowerOfTwo(window_length_); + // CHECK(fft_length_ >= window_length_); + output_frequency_channels_ = 1 + fft_length_ / 2; + + // Allocate 2 more than what rdft needs, so we can rationalize the layout. + fft_input_output_.assign(fft_length_ + 2, 0.0); + + int half_fft_length = fft_length_ / 2; + fft_double_working_area_.assign(half_fft_length, 0.0); + fft_integer_working_area_.assign(2 + static_cast(sqrt(half_fft_length)), + 0); + // Set flag element to ensure that the working areas are initialized + // on the first call to cdft. It's redundant given the assign above, + // but keep it as a reminder. + fft_integer_working_area_[0] = 0; + input_queue_.clear(); + samples_to_next_step_ = window_length_; + initialized_ = true; + return true; +} + +template +bool Spectrogram::ComputeComplexSpectrogram( + const std::vector& input, + std::vector>>* output) { + if (!initialized_) { + // LOG(ERROR) << "ComputeComplexSpectrogram() called before successful call + // " + // << "to Initialize()."; + return false; + } + // CHECK(output); + output->clear(); + int input_start = 0; + while (GetNextWindowOfSamples(input, &input_start)) { + // DCHECK_EQ(input_queue_.size(), window_length_); + ProcessCoreFFT(); // Processes input_queue_ to fft_input_output_. + // Add a new slice vector onto the output, to save new result to. + output->resize(output->size() + 1); + // Get a reference to the newly added slice to fill in. + auto& spectrogram_slice = output->back(); + spectrogram_slice.resize(output_frequency_channels_); + for (int i = 0; i < output_frequency_channels_; ++i) { + // This will convert double to float if it needs to. + spectrogram_slice[i] = complex( + fft_input_output_[2 * i], fft_input_output_[2 * i + 1]); + } + } + return true; +} +// Instantiate it four ways: +template bool Spectrogram::ComputeComplexSpectrogram( + const std::vector& input, std::vector>>*); +template bool Spectrogram::ComputeComplexSpectrogram( + const std::vector& input, + std::vector>>*); +template bool Spectrogram::ComputeComplexSpectrogram( + const std::vector& input, + std::vector>>*); +template bool Spectrogram::ComputeComplexSpectrogram( + const std::vector& input, + std::vector>>*); + +template +bool Spectrogram::ComputeSquaredMagnitudeSpectrogram( + const std::vector& input, + std::vector>* output) { + if (!initialized_) { + // LOG(ERROR) << "ComputeSquaredMagnitudeSpectrogram() called before " + // << "successful call to Initialize()."; + return false; + } + // CHECK(output); + output->clear(); + int input_start = 0; + while (GetNextWindowOfSamples(input, &input_start)) { + // DCHECK_EQ(input_queue_.size(), window_length_); + ProcessCoreFFT(); // Processes input_queue_ to fft_input_output_. + // Add a new slice vector onto the output, to save new result to. + output->resize(output->size() + 1); + // Get a reference to the newly added slice to fill in. + auto& spectrogram_slice = output->back(); + spectrogram_slice.resize(output_frequency_channels_); + for (int i = 0; i < output_frequency_channels_; ++i) { + // Similar to the Complex case, except storing the norm. + // But the norm function is known to be a performance killer, + // so do it this way with explicit real and imagninary temps. + const double re = fft_input_output_[2 * i]; + const double im = fft_input_output_[2 * i + 1]; + // Which finally converts double to float if it needs to. + spectrogram_slice[i] = re * re + im * im; + } + } + return true; +} +// Instantiate it four ways: +template bool Spectrogram::ComputeSquaredMagnitudeSpectrogram( + const std::vector& input, std::vector>*); +template bool Spectrogram::ComputeSquaredMagnitudeSpectrogram( + const std::vector& input, std::vector>*); +template bool Spectrogram::ComputeSquaredMagnitudeSpectrogram( + const std::vector& input, std::vector>*); +template bool Spectrogram::ComputeSquaredMagnitudeSpectrogram( + const std::vector& input, std::vector>*); + +// Return true if a full window of samples is prepared; manage the queue. +template +bool Spectrogram::GetNextWindowOfSamples(const std::vector& input, + int* input_start) { + auto input_it = input.begin() + *input_start; + int input_remaining = input.end() - input_it; + if (samples_to_next_step_ > input_remaining) { + // Copy in as many samples are left and return false, no full window. + input_queue_.insert(input_queue_.end(), input_it, input.end()); + *input_start += input_remaining; // Increases it to input.size(). + samples_to_next_step_ -= input_remaining; + return false; // Not enough for a full window. + } else { + // Copy just enough into queue to make a new window, then trim the + // front off the queue to make it window-sized. + input_queue_.insert(input_queue_.end(), input_it, + input_it + samples_to_next_step_); + *input_start += samples_to_next_step_; + input_queue_.erase( + input_queue_.begin(), + input_queue_.begin() + input_queue_.size() - window_length_); + // DCHECK_EQ(window_length_, input_queue_.size()); + samples_to_next_step_ = step_length_; // Be ready for next time. + return true; // Yes, input_queue_ now contains exactly a window-full. + } +} + +void Spectrogram::ProcessCoreFFT() { + for (int j = 0; j < window_length_; ++j) { + fft_input_output_[j] = input_queue_[j] * window_[j]; + } + // Zero-pad the rest of the input buffer. + for (int j = window_length_; j < fft_length_; ++j) { + fft_input_output_[j] = 0.0; + } + const int kForwardFFT = 1; // 1 means forward; -1 reverse. + // This real FFT is a fair amount faster than using cdft here. + rdft(fft_length_, kForwardFFT, &fft_input_output_[0], + &fft_integer_working_area_[0], &fft_double_working_area_[0]); + // Make rdft result look like cdft result; + // unpack the last real value from the first position's imag slot. + fft_input_output_[fft_length_] = fft_input_output_[1]; + fft_input_output_[fft_length_ + 1] = 0; + fft_input_output_[1] = 0; +} + +} // namespace internal +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/spectrogram.h b/tensorflow/contrib/lite/kernels/internal/spectrogram.h new file mode 100644 index 0000000000000000000000000000000000000000..b77a68f7dfe6edb07ec4e5db540c673b2d6f6d6e --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/spectrogram.h @@ -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. +==============================================================================*/ + +// Class for generating spectrogram slices from a waveform. +// Initialize() should be called before calls to other functions. Once +// Initialize() has been called and returned true, The Compute*() functions can +// be called repeatedly with sequential input data (ie. the first element of the +// next input vector directly follows the last element of the previous input +// vector). Whenever enough audio samples are buffered to produce a +// new frame, it will be placed in output. Output is cleared on each +// call to Compute*(). This class is thread-unsafe, and should only be +// called from one thread at a time. +// With the default parameters, the output of this class should be very +// close to the results of the following MATLAB code: +// overlap_samples = window_length_samples - step_samples; +// window = hann(window_length_samples, 'periodic'); +// S = abs(spectrogram(audio, window, overlap_samples)).^2; + +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_SPECTROGRAM_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_SPECTROGRAM_H_ + +#include +#include +#include + +#include "third_party/fft2d/fft.h" + +namespace tflite { +namespace internal { + +class Spectrogram { + public: + Spectrogram() : initialized_(false) {} + ~Spectrogram() {} + + // Initializes the class with a given window length and step length + // (both in samples). Internally a Hann window is used as the window + // function. Returns true on success, after which calls to Process() + // are possible. window_length must be greater than 1 and step + // length must be greater than 0. + bool Initialize(int window_length, int step_length); + + // Initialize with an explicit window instead of a length. + bool Initialize(const std::vector& window, int step_length); + + // Processes an arbitrary amount of audio data (contained in input) + // to yield complex spectrogram frames. After a successful call to + // Initialize(), Process() may be called repeatedly with new input data + // each time. The audio input is buffered internally, and the output + // vector is populated with as many temporally-ordered spectral slices + // as it is possible to generate from the input. The output is cleared + // on each call before the new frames (if any) are added. + // + // The template parameters can be float or double. + template + bool ComputeComplexSpectrogram( + const std::vector& input, + std::vector>>* output); + + // This function works as the one above, but returns the power + // (the L2 norm, or the squared magnitude) of each complex value. + template + bool ComputeSquaredMagnitudeSpectrogram( + const std::vector& input, + std::vector>* output); + + // Return reference to the window function used internally. + const std::vector& GetWindow() const { return window_; } + + // Return the number of frequency channels in the spectrogram. + int output_frequency_channels() const { return output_frequency_channels_; } + + private: + template + bool GetNextWindowOfSamples(const std::vector& input, + int* input_start); + void ProcessCoreFFT(); + + int fft_length_; + int output_frequency_channels_; + int window_length_; + int step_length_; + bool initialized_; + int samples_to_next_step_; + + std::vector window_; + std::vector fft_input_output_; + std::deque input_queue_; + + // Working data areas for the FFT routines. + std::vector fft_integer_working_area_; + std::vector fft_double_working_area_; +}; + +} // namespace internal +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_SPECTROGRAM_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/tensor.h b/tensorflow/contrib/lite/kernels/internal/tensor.h index dfe76c2afd40c692063710a4d98464b55e40feb9..62e38e0d4c3e023d0ed2242fc9438b096b86dc59 100644 --- a/tensorflow/contrib/lite/kernels/internal/tensor.h +++ b/tensorflow/contrib/lite/kernels/internal/tensor.h @@ -81,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 // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/types.h b/tensorflow/contrib/lite/kernels/internal/types.h index afe131b06ec41201395e80aa5415fd7db990f8d4..293538fcbb6406d6065d8efd25adb3b163638c92 100644 --- a/tensorflow/contrib/lite/kernels/internal/types.h +++ b/tensorflow/contrib/lite/kernels/internal/types.h @@ -21,6 +21,22 @@ namespace tflite { enum class FusedActivationFunctionType : uint8 { kNone, kRelu6, kRelu1, kRelu }; +// Quantization parameters, determining the mapping of quantized values +// to real values (i.e. determining how quantized values are mathematically +// interpreted). +// +// The correspondence is as follows: +// +// real_value = scale * (quantized_value - zero_point); +// +// In other words, zero_point designates which quantized value corresponds to +// the real 0 value, and scale designates the difference between the real values +// corresponding to consecutive quantized values differing by 1. +struct QuantizationParams { + int32 zero_point = 0; + double scale = 0.0; +}; + template struct Dims { int sizes[N]; diff --git a/tensorflow/contrib/lite/kernels/kernel_util.h b/tensorflow/contrib/lite/kernels/kernel_util.h index 3cfa72615a95d6f215ef9d35f2572026ec90fad8..21da1daff7783bc98bc36814f1f9605a83eb65cf 100644 --- a/tensorflow/contrib/lite/kernels/kernel_util.h +++ b/tensorflow/contrib/lite/kernels/kernel_util.h @@ -58,14 +58,17 @@ inline bool IsConstantTensor(TfLiteTensor* tensor) { } // Determines whether tensor is dynamic. Note that a tensor can be non-const and -// not dynamic. This function specificially checks for a dynamic tensor. +// not dynamic. This function specifically checks for a dynamic tensor. inline bool IsDynamicTensor(TfLiteTensor* tensor) { return tensor->allocation_type == kTfLiteDynamic; } // Sets tensor to dynamic. inline void SetTensorToDynamic(TfLiteTensor* tensor) { - tensor->allocation_type = kTfLiteDynamic; + if (tensor->allocation_type != kTfLiteDynamic) { + tensor->allocation_type = kTfLiteDynamic; + tensor->data.raw = nullptr; + } } // Calculates the multiplication factor for a quantized convolution (or diff --git a/tensorflow/contrib/lite/kernels/log_softmax_test.cc b/tensorflow/contrib/lite/kernels/log_softmax_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..62820a2f5113cb6ae252386aaf3842135383b79f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/log_softmax_test.cc @@ -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. +==============================================================================*/ +// Unit test for TFLite LOG_SOFTMAX op. + +#include +#include +#include + +#include +#include +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.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 { + +class LogSoftmaxOpModel : public SingleOpModel { + public: + LogSoftmaxOpModel(int batches, int size) + : batches_(batches), input_size_(size) { + input_ = AddInput(TensorType_FLOAT32); + output_ = AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_LOG_SOFTMAX, BuiltinOptions_LogSoftmaxOptions, + CreateLogSoftmaxOptions(builder_).Union()); + BuildInterpreter({{batches_, input_size_}}); + } + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetInput(int offset, float* begin, float* end) { + PopulateTensor(input_, offset, begin, end); + } + + std::vector GetOutput() { return ExtractVector(output_); } + + private: + int input_; + int output_; + + int batches_; + int input_size_; +}; + +TEST(LogSoftmaxOpTest, SimpleTest) { + LogSoftmaxOpModel m(/*batches=*/2, /*size=*/5); + m.SetInput({ + 1.0, 2.0, 3.0, 4.0, 5.0, // b = 0 + -1.0, -2.0, -3.0, -4.0, -5.0, // b = 1 + }); + + m.Invoke(); + + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear( + {-4.45191431, -3.45191431, -2.45191431, -1.45191443, -0.4519144, + -0.4519144, -1.45191443, -2.45191431, -3.45191431, -4.45191431}, + 1e-6))); +} + +TEST(LogSoftmaxOpTest, CompareWithTFmini) { + const int batch_size = 2; + const int input_size = 5; + static float input_buffer[] = { + 1.0, 2.0, 3.0, 4.0, 5.0, // b = 0 + -1.0, -2.0, -3.0, -4.0, -5.0, // b = 1 + }; + + LogSoftmaxOpModel m(batch_size, input_size); + + m.SetInput(0, input_buffer, input_buffer + input_size * batch_size); + + m.Invoke(); + + std::unique_ptr output_buffer(new float[input_size * batch_size]); + static tflite::Dims<4> input_dims = {{input_size, 1, 1, batch_size}, + {1, 0, 0, input_size}}; + tflite::reference_ops::LogSoftmax(input_buffer, input_dims, + output_buffer.get(), input_dims); + + std::vector expected; + expected.insert(expected.end(), output_buffer.get(), + output_buffer.get() + input_size * batch_size); + + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(expected, 1e-6))); +} + +} // 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/lsh_projection.cc b/tensorflow/contrib/lite/kernels/lsh_projection.cc index 5f73b56ed9790b216adc788490faebaabd2bc756..0ee35775d50b8750455572f789d7b92481655a95 100644 --- a/tensorflow/contrib/lite/kernels/lsh_projection.cc +++ b/tensorflow/contrib/lite/kernels/lsh_projection.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// LSH Projection projects an input to a bit vector via locality senstive +// LSH Projection projects an input to a bit vector via locality sensitive // hashing. // // Options: diff --git a/tensorflow/contrib/lite/kernels/lstm.cc b/tensorflow/contrib/lite/kernels/lstm.cc index 6c06264d845c24e71647b6fd2374734be32383ef..8cf1165135bdb0d4669bb97fd2d98e3dc044b4d9 100644 --- a/tensorflow/contrib/lite/kernels/lstm.cc +++ b/tensorflow/contrib/lite/kernels/lstm.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/activation_functor.h" +#include "tensorflow/contrib/lite/kernels/internal/kernel_utils.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" @@ -212,9 +213,9 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context, // present. // 2) If projection weight is present, then projection bias is optional. // TODO(ghodrat): make sure this is correct. - const bool projecton_tensors_consistent = + const bool projection_tensors_consistent = ((projection_weights != nullptr) || (projection_bias == nullptr)); - TF_LITE_ENSURE(context, projecton_tensors_consistent == true); + TF_LITE_ENSURE(context, projection_tensors_consistent == true); return kTfLiteOk; } @@ -356,7 +357,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { 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. + // check the existence of only one to get the condition. const bool use_cifg = (input_to_input_weights == nullptr); const bool use_peephole = (cell_to_output_weights != nullptr); @@ -377,127 +378,54 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { output_gate_scratch = scratch_buffer->data.f + 3 * n_cell * n_batch; } - // 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->data.f, n_batch, - input_gate_scratch, /*result_stride=*/1); - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_forget_weights->data.f, n_cell, n_input, input->data.f, n_batch, - forget_gate_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_cell_weights->data.f, n_cell, n_input, input->data.f, n_batch, - cell_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_output_weights->data.f, n_cell, n_input, input->data.f, 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); - if (use_projection_weight) { - if (use_projection_bias) { - tensor_utils::VectorBatchVectorAssign(projection_bias->data.f, n_output, - n_batch, output->data.f); - } else { - tensor_utils::ZeroVector(output->data.f, n_batch * n_output); - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - projection_weights->data.f, n_output, n_cell, output_gate_scratch, - n_batch, output->data.f, /*result_stride=*/1); - if (params->proj_clip > 0.0) { - tensor_utils::ClipVector(output->data.f, n_batch * n_output, - params->proj_clip, output->data.f); - } - } else { - tensor_utils::CopyVector(output_gate_scratch, n_batch * n_output, - output->data.f); - } - tensor_utils::CopyVector(output->data.f, n_batch * n_output, - output_state->data.f); + // Check optional tensors, the respective pointers can be null. + const float* input_to_input_weights_ptr = + (use_cifg) ? nullptr : input_to_input_weights->data.f; + const float* recurrent_to_input_weights_ptr = + (use_cifg) ? nullptr : recurrent_to_input_weights->data.f; + const float* input_gate_bias_ptr = + (use_cifg) ? nullptr : input_gate_bias->data.f; + const float* cell_to_input_weights_ptr = + (use_peephole && !use_cifg) ? cell_to_input_weights->data.f : nullptr; + const float* cell_to_forget_weights_ptr = + (use_peephole) ? cell_to_forget_weights->data.f : nullptr; + const float* cell_to_output_weights_ptr = + (use_peephole) ? cell_to_output_weights->data.f : nullptr; + const float* projection_weights_ptr = + (projection_weights == nullptr) ? nullptr : projection_weights->data.f; + const float* projection_bias_ptr = + (projection_bias == nullptr) ? nullptr : projection_bias->data.f; + + // Required tensors, pointers are non-null. + const float* input_ptr_batch = input->data.f; + const float* input_to_forget_weights_ptr = input_to_forget_weights->data.f; + const float* input_to_cell_weights_ptr = input_to_cell_weights->data.f; + const float* input_to_output_weights_ptr = input_to_output_weights->data.f; + const float* recurrent_to_forget_weights_ptr = + recurrent_to_forget_weights->data.f; + const float* recurrent_to_cell_weights_ptr = + recurrent_to_cell_weights->data.f; + const float* recurrent_to_output_weights_ptr = + recurrent_to_output_weights->data.f; + const float* forget_gate_bias_ptr = forget_gate_bias->data.f; + const float* cell_bias_ptr = cell_bias->data.f; + const float* output_gate_bias_ptr = output_gate_bias->data.f; + + float* output_state_ptr = output_state->data.f; + float* cell_state_ptr = cell_state->data.f; + float* output_ptr_batch = output->data.f; + + kernel_utils::LstmStep( + input_ptr_batch, input_to_input_weights_ptr, input_to_forget_weights_ptr, + input_to_cell_weights_ptr, input_to_output_weights_ptr, + recurrent_to_input_weights_ptr, recurrent_to_forget_weights_ptr, + recurrent_to_cell_weights_ptr, recurrent_to_output_weights_ptr, + cell_to_input_weights_ptr, cell_to_forget_weights_ptr, + cell_to_output_weights_ptr, input_gate_bias_ptr, forget_gate_bias_ptr, + cell_bias_ptr, output_gate_bias_ptr, projection_weights_ptr, + projection_bias_ptr, params, n_batch, n_cell, n_input, n_output, + output_state_ptr, cell_state_ptr, input_gate_scratch, forget_gate_scratch, + cell_scratch, output_gate_scratch, output_ptr_batch); return kTfLiteOk; } diff --git a/tensorflow/contrib/lite/kernels/mean.cc b/tensorflow/contrib/lite/kernels/mean.cc index ec1c40202761e3789462a4740e5547eba654b0f9..aff19581ea56f94c08638b7b388ae181f566cf4f 100644 --- a/tensorflow/contrib/lite/kernels/mean.cc +++ b/tensorflow/contrib/lite/kernels/mean.cc @@ -183,8 +183,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_OK(context, ResizeTempAxis(context, &op_context, resolved_axis)); TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); - TfLiteTensorRealloc(resolved_axis->bytes, resolved_axis); - TfLiteTensorRealloc(op_context.output->bytes, op_context.output); } #define TF_LITE_MEAN(kernel_type, data_type) \ diff --git a/tensorflow/contrib/lite/kernels/mean_test.cc b/tensorflow/contrib/lite/kernels/mean_test.cc index c4c53c2ded351849e7c458fc754c36395a25ebd0..2d6d4bc2da4b75289ee27c3f2a12787216716d44 100644 --- a/tensorflow/contrib/lite/kernels/mean_test.cc +++ b/tensorflow/contrib/lite/kernels/mean_test.cc @@ -74,7 +74,7 @@ class MeanOpDynamicModel : public BaseMeanOpModel { } }; -TEST(ConstMeanOpTest, NotKeepDims) { +TEST(ConstFloatMeanOpTest, 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}; @@ -86,7 +86,7 @@ TEST(ConstMeanOpTest, NotKeepDims) { EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({12, 13}))); } -TEST(ConstMeanOpTest, KeepDims) { +TEST(ConstFloatMeanOpTest, 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}; @@ -99,7 +99,7 @@ TEST(ConstMeanOpTest, KeepDims) { ElementsAreArray(ArrayFloatNear({10.5, 12.5, 14.5}))); } -TEST(DynamicMeanOpTest, NotKeepDims) { +TEST(DynamicFloatMeanOpTest, 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}; @@ -114,7 +114,7 @@ TEST(DynamicMeanOpTest, NotKeepDims) { EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({12, 13}))); } -TEST(DynamicMeanOpTest, KeepDims) { +TEST(DynamicFloatMeanOpTest, 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}; @@ -130,6 +130,70 @@ TEST(DynamicMeanOpTest, KeepDims) { ElementsAreArray(ArrayFloatNear({10.5, 12.5, 14.5}))); } +TEST(DynamicFloatMeanOpTest, Scale) { + std::initializer_list data = {9.527}; + MeanOpDynamicModel m({TensorType_FLOAT32, {1}}, {TensorType_FLOAT32, {1}}, + {TensorType_INT32, {1}}, true); + std::initializer_list axis = {0}; + m.SetAxis(axis); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({9.527}))); +} + +TEST(ConstUint8MeanOpTest, NotKeepDims) { + std::initializer_list data = {1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24}; + MeanOpConstModel m({TensorType_UINT8, {4, 3, 2}}, {TensorType_UINT8, {2}}, + {4}, {1, 0, -3, -3}, false); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({12, 13})); +} + +TEST(ConstUint8MeanOpTest, KeepDims) { + std::initializer_list data = {1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24}; + MeanOpConstModel m({TensorType_UINT8, {4, 3, 2}}, {TensorType_UINT8, {3}}, + {2}, {0, 2}, true); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({10, 12, 14})); +} + +TEST(DynamicUint8MeanOpTest, NotKeepDims) { + std::initializer_list data = {1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24}; + MeanOpDynamicModel m({TensorType_UINT8, {4, 3, 2}}, {TensorType_UINT8, {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({12, 13})); +} + +TEST(DynamicUint8MeanOpTest, KeepDims) { + std::initializer_list data = {1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24}; + MeanOpDynamicModel m({TensorType_UINT8, {4, 3, 2}}, {TensorType_UINT8, {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({10, 12, 14})); +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/mfcc.cc b/tensorflow/contrib/lite/kernels/mfcc.cc new file mode 100644 index 0000000000000000000000000000000000000000..018db0dc54c5d281bf3fb3ff8a1f111b427fe76b --- /dev/null +++ b/tensorflow/contrib/lite/kernels/mfcc.cc @@ -0,0 +1,154 @@ +/* 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/mfcc.h" +#include "flatbuffers/flexbuffers.h" +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/mfcc_dct.h" +#include "tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.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 custom { +namespace mfcc { + +enum KernelType { + kReference, +}; + +typedef struct { + float upper_frequency_limit; + float lower_frequency_limit; + int filterbank_channel_count; + int dct_coefficient_count; +} TfLiteMfccParams; + +constexpr int kInputTensorWav = 0; +constexpr int kInputTensorRate = 1; +constexpr int kOutputTensor = 0; + +void* Init(TfLiteContext* context, const char* buffer, size_t length) { + auto* data = new TfLiteMfccParams; + + const uint8_t* buffer_t = reinterpret_cast(buffer); + + const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap(); + data->upper_frequency_limit = m["upper_frequency_limit"].AsInt64(); + data->lower_frequency_limit = m["lower_frequency_limit"].AsInt64(); + data->filterbank_channel_count = m["filterbank_channel_count"].AsInt64(); + data->dct_coefficient_count = m["dct_coefficient_count"].AsInt64(); + return data; +} + +void Free(TfLiteContext* context, void* buffer) { + delete reinterpret_cast(buffer); +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->user_data); + + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TfLiteTensor* inputWav = GetInput(context, node, kInputTensorWav); + TfLiteTensor* inputRate = GetInput(context, node, kInputTensorRate); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + TF_LITE_ENSURE_EQ(context, NumDimensions(inputWav), 3); + TF_LITE_ENSURE_EQ(context, NumDimensions(inputRate), 1); + + TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); + TF_LITE_ENSURE_EQ(context, inputWav->type, output->type); + + TfLiteIntArray* output_size = TfLiteIntArrayCreate(3); + output_size->data[0] = inputWav->dims->data[0]; + output_size->data[1] = inputWav->dims->data[1]; + output_size->data[2] = params->dct_coefficient_count; + + return context->ResizeTensor(context, output, output_size); +} + +// Input is a single squared-magnitude spectrogram frame. The input spectrum +// is converted to linear magnitude and weighted into bands using a +// triangular mel filterbank, and a discrete cosine transform (DCT) of the +// values is taken. Output is populated with the lowest dct_coefficient_count +// of these values. +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->user_data); + + TfLiteTensor* inputWav = GetInput(context, node, kInputTensorWav); + TfLiteTensor* inputRate = GetInput(context, node, kInputTensorRate); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + const int32 sample_rate = *GetTensorData(inputRate); + + const int spectrogram_channels = inputWav->dims->data[2]; + const int spectrogram_samples = inputWav->dims->data[1]; + const int audio_channels = inputWav->dims->data[0]; + + internal::Mfcc mfcc; + mfcc.set_upper_frequency_limit(params->upper_frequency_limit); + mfcc.set_lower_frequency_limit(params->lower_frequency_limit); + mfcc.set_filterbank_channel_count(params->filterbank_channel_count); + mfcc.set_dct_coefficient_count(params->dct_coefficient_count); + + mfcc.Initialize(spectrogram_channels, sample_rate); + + const float* spectrogram_flat = GetTensorData(inputWav); + float* output_flat = GetTensorData(output); + + for (int audio_channel = 0; audio_channel < audio_channels; ++audio_channel) { + for (int spectrogram_sample = 0; spectrogram_sample < spectrogram_samples; + ++spectrogram_sample) { + const float* sample_data = + spectrogram_flat + + (audio_channel * spectrogram_samples * spectrogram_channels) + + (spectrogram_sample * spectrogram_channels); + std::vector mfcc_input(sample_data, + sample_data + spectrogram_channels); + std::vector mfcc_output; + mfcc.Compute(mfcc_input, &mfcc_output); + TF_LITE_ENSURE_EQ(context, params->dct_coefficient_count, + mfcc_output.size()); + float* output_data = output_flat + + (audio_channel * spectrogram_samples * + params->dct_coefficient_count) + + (spectrogram_sample * params->dct_coefficient_count); + for (int i = 0; i < params->dct_coefficient_count; ++i) { + output_data[i] = mfcc_output[i]; + } + } + } + + return kTfLiteOk; +} + +} // namespace mfcc + +TfLiteRegistration* Register_MFCC() { + static TfLiteRegistration r = {mfcc::Init, mfcc::Free, mfcc::Prepare, + mfcc::Eval}; + return &r; +} + +} // namespace custom +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/mfcc_test.cc b/tensorflow/contrib/lite/kernels/mfcc_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0291ca8c1c58ea6ab3bb7c22bc436ed3404cba74 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/mfcc_test.cc @@ -0,0 +1,104 @@ +/* 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 "flatbuffers/flexbuffers.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 custom { + +TfLiteRegistration* Register_MFCC(); + +namespace { + +using ::testing::ElementsAre; +using ::testing::ElementsAreArray; + +class BaseMfccOpModel : public SingleOpModel { + public: + BaseMfccOpModel(const TensorData& input1, const TensorData& input2, + const TensorData& output) { + input1_ = AddInput(input1); + input2_ = AddInput(input2); + output_ = AddOutput(output); + + flexbuffers::Builder fbb; + fbb.Map([&]() { + fbb.Int("upper_frequency_limit", 4000); + fbb.Int("lower_frequency_limit", 20); + fbb.Int("filterbank_channel_count", 40); + fbb.Int("dct_coefficient_count", 13); + }); + fbb.Finish(); + SetCustomOp("Mfcc", fbb.GetBuffer(), Register_MFCC); + + BuildInterpreter({GetShape(input1_), GetShape(input2_)}); + } + + int input1() { return input1_; } + int input2() { return input2_; } + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input1_; + int input2_; + int output_; +}; + +TEST(MfccOpTest, SimpleTest) { + BaseMfccOpModel m({TensorType_FLOAT32, {1, 1, 513}}, {TensorType_INT32, {1}}, + {TensorType_FLOAT32, {}}); + + std::vector data(513); + for (int i = 0; i < data.size(); ++i) { + data[i] = i + 1; + } + m.PopulateTensor(m.input1(), 0, data.data(), + data.data() + data.size()); + m.PopulateTensor(m.input2(), {22050}); + + m.Invoke(); + + std::vector output_shape = m.GetOutputShape(); + EXPECT_THAT(output_shape, ElementsAre(1, 1, 13)); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear( + {29.13970072, -6.41568601, -0.61903012, -0.96778652, -0.26819878, + -0.40907028, -0.15614748, -0.23203119, -0.10481487, -0.1543029, + -0.0769791, -0.10806114, -0.06047613}, + 1e-3))); +} + +} // namespace +} // namespace custom +} // namespace ops +} // 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 8838b300c0af167bf2ffcf944fc7c31d6173f462..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}); @@ -79,9 +81,9 @@ TEST(FloatMulOpTest, NoActivation) { } TEST(FloatMulOpTest, ActivationRELU_N1_TO_1) { - FloatMulOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {}}, - ActivationFunctionType_RELU_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(); @@ -94,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}); @@ -105,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}); @@ -117,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/pad.cc b/tensorflow/contrib/lite/kernels/pad.cc index 48114e5a4069abf864a996141c7b0906301d9809..c29da3862e84d6756bf5ef34b2ca06307b0a065d 100644 --- a/tensorflow/contrib/lite/kernels/pad.cc +++ b/tensorflow/contrib/lite/kernels/pad.cc @@ -101,7 +101,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* 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)); - TfLiteTensorRealloc(op_context.output->bytes, op_context.output); } // TODO(nupurgarg): Change kernel implementation to take in int* instead of diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc index 1fb779fd5174a255b7d34322c57dc084f68d8c3f..369d3b9886ad80837e0e042735b3be72c7f0e33f 100644 --- a/tensorflow/contrib/lite/kernels/register.cc +++ b/tensorflow/contrib/lite/kernels/register.cc @@ -17,6 +17,14 @@ limitations under the License. namespace tflite { namespace ops { + +namespace custom { + +TfLiteRegistration* Register_AUDIO_SPECTROGRAM(); +TfLiteRegistration* Register_MFCC(); + +} // namespace custom + namespace builtin { TfLiteRegistration* Register_RELU(); @@ -49,6 +57,7 @@ TfLiteRegistration* Register_MUL(); TfLiteRegistration* Register_L2_NORMALIZATION(); TfLiteRegistration* Register_LOCAL_RESPONSE_NORMALIZATION(); TfLiteRegistration* Register_LSTM(); +TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_LSTM(); TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_LSTM(); TfLiteRegistration* Register_PAD(); TfLiteRegistration* Register_RESHAPE(); @@ -58,8 +67,14 @@ 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(); +TfLiteRegistration* Register_LOG_SOFTMAX(); +TfLiteRegistration* Register_CAST(); +TfLiteRegistration* Register_DEQUANTIZE(); BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_RELU, Register_RELU()); @@ -94,6 +109,8 @@ BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION, Register_LOCAL_RESPONSE_NORMALIZATION()); AddBuiltin(BuiltinOperator_LSTM, Register_LSTM()); + AddBuiltin(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, + Register_BIDIRECTIONAL_SEQUENCE_LSTM()); AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, Register_UNIDIRECTIONAL_SEQUENCE_LSTM()); AddBuiltin(BuiltinOperator_PAD, Register_PAD()); @@ -106,8 +123,20 @@ BuiltinOpResolver::BuiltinOpResolver() { 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()); + AddBuiltin(BuiltinOperator_LOG_SOFTMAX, Register_LOG_SOFTMAX()); + AddBuiltin(BuiltinOperator_CAST, Register_CAST()); + AddBuiltin(BuiltinOperator_DEQUANTIZE, Register_DEQUANTIZE()); + + // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that + // custom ops aren't always included by default. + AddCustom("Mfcc", tflite::ops::custom::Register_MFCC()); + AddCustom("AudioSpectrogram", + tflite::ops::custom::Register_AUDIO_SPECTROGRAM()); } TfLiteRegistration* BuiltinOpResolver::FindOp( diff --git a/tensorflow/contrib/lite/kernels/reshape.cc b/tensorflow/contrib/lite/kernels/reshape.cc index f3e6ddc9f480e3863cac52157ae28b7329ee2088..438f70d3115130efe477a3ceeccd2e77108c979a 100644 --- a/tensorflow/contrib/lite/kernels/reshape.cc +++ b/tensorflow/contrib/lite/kernels/reshape.cc @@ -49,20 +49,20 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TfLiteIntArray* output_size = TfLiteIntArrayCreate(params->num_dimensions); int num_output_elements = 1; - int strech_dim = -1; + int stretch_dim = -1; for (int i = 0; i < params->num_dimensions; ++i) { int value = params->shape[i]; if (value == -1) { - TF_LITE_ENSURE_EQ(context, strech_dim, -1); - strech_dim = i; + TF_LITE_ENSURE_EQ(context, stretch_dim, -1); + stretch_dim = i; } else { num_output_elements *= value; output_size->data[i] = value; } } - if (strech_dim != -1) { - output_size->data[strech_dim] = num_input_elements / num_output_elements; - num_output_elements *= output_size->data[strech_dim]; + if (stretch_dim != -1) { + output_size->data[stretch_dim] = num_input_elements / num_output_elements; + num_output_elements *= output_size->data[stretch_dim]; } TF_LITE_ENSURE_EQ(context, num_input_elements, num_output_elements); diff --git a/tensorflow/contrib/lite/kernels/reshape_test.cc b/tensorflow/contrib/lite/kernels/reshape_test.cc index 0fbcf6e6aa311d2cac491336ee54ccf58bbda8fd..aecbd0399f7454045e8189072f45b695b0525204 100644 --- a/tensorflow/contrib/lite/kernels/reshape_test.cc +++ b/tensorflow/contrib/lite/kernels/reshape_test.cc @@ -60,7 +60,7 @@ TEST(ReshapeOpTest, TooManyDimensions) { TEST(ReshapeOpTest, TooManySpecialDimensions) { EXPECT_DEATH(ReshapeOpModel({1, 2, 4, 1}, {-1, -1, 2, 4}), - "strech_dim != -1"); + "stretch_dim != -1"); } TEST(ReshapeOpTest, SimpleTest) { diff --git a/tensorflow/contrib/lite/kernels/resize_bilinear.cc b/tensorflow/contrib/lite/kernels/resize_bilinear.cc index 9a419af0238e1a25e4b9e81f109b54de6b49097b..9e3e19c09a4012ebdadbc2a7c2ba06c4bfefd206 100644 --- a/tensorflow/contrib/lite/kernels/resize_bilinear.cc +++ b/tensorflow/contrib/lite/kernels/resize_bilinear.cc @@ -36,6 +36,17 @@ constexpr int kInputTensor = 0; constexpr int kSizeTensor = 1; constexpr int kOutputTensor = 0; +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); +} + TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); @@ -55,32 +66,33 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // integers. output->type = kTfLiteFloat32; - // TODO(ahentz): if the input is constant, we can allocate here. - output->allocation_type = kTfLiteDynamic; - return kTfLiteOk; + if (!IsConstantTensor(size)) { + SetTensorToDynamic(output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, input, size, output); } template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = + reinterpret_cast(node->builtin_data); + TfLiteTensor* input = GetInput(context, node, kInputTensor); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); TfLiteTensor* size = GetInput(context, node, kSizeTensor); - // TODO(ahentz): we only need to do this here if it wasn't done in Eval(). - 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]; - context->ResizeTensor(context, output, output_size); - TfLiteTensorRealloc(output->bytes, output); + 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), \ - GetTensorData(size), GetTensorDims(size), \ - 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 2b1aaf654f87f435ec464b2cc1a63c77ba86ae5b..4e03f3820a5c14ee1692c553db61e385716b1723 100644 --- a/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc +++ b/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc @@ -25,14 +25,24 @@ using ::testing::ElementsAreArray; class ResizeBilinearOpModel : public SingleOpModel { public: - ResizeBilinearOpModel(std::initializer_list input_shape) { - input_ = AddInput(TensorType_FLOAT32); - size_ = AddInput(TensorType_INT32); - output_ = AddOutput(TensorType_FLOAT32); + 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()); - BuildInterpreter({input_shape, {2}}); + if (const_size) { + BuildInterpreter({GetShape(input_)}); + } else { + BuildInterpreter({GetShape(input_), GetShape(size_)}); + } } void SetInput(std::initializer_list data) { @@ -49,23 +59,33 @@ class ResizeBilinearOpModel : public SingleOpModel { }; TEST(ResizeBilinearOpTest, HorizontalResize) { - ResizeBilinearOpModel m({1, 1, 2, 1}); + 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}); + 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}); + ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}); m.SetInput({ 3, 6, // 9, 12 // @@ -77,10 +97,22 @@ TEST(ResizeBilinearOpTest, TwoDimensionalResize) { 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}); + ResizeBilinearOpModel m({TensorType_FLOAT32, {2, 2, 2, 1}}); m.SetInput({ 3, 6, // 9, 12, // @@ -97,10 +129,27 @@ 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}); + ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 2, 2}}); m.SetInput({ 3, 4, 6, 10, // 9, 10, 12, 16, // @@ -112,6 +161,18 @@ TEST(ResizeBilinearOpTest, ThreeDimensionalResize) { 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 index e2e1873f770fad889137b43d87585602162819f7..d8c9e352f00627eee45ae836b720f2af77140538 100644 --- a/tensorflow/contrib/lite/kernels/space_to_batch_nd.cc +++ b/tensorflow/contrib/lite/kernels/space_to_batch_nd.cc @@ -111,7 +111,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* 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)); - TfLiteTensorRealloc(op_context.output->bytes, op_context.output); } #define TF_LITE_SPACE_TO_BATCH_ND(type, scalar) \ 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/strided_slice.cc b/tensorflow/contrib/lite/kernels/strided_slice.cc index c510ee3b9f4555c79a49df5932b49ef735d2feef..eb374d903182f46b40f5c80bfd769a19a5594742 100644 --- a/tensorflow/contrib/lite/kernels/strided_slice.cc +++ b/tensorflow/contrib/lite/kernels/strided_slice.cc @@ -48,7 +48,7 @@ struct StridedSliceContext { output = GetOutput(context, node, kOutputTensor); dims = NumDimensions(input); } - TfLiteStridedSliceParams* params; + const TfLiteStridedSliceParams* params; TfLiteTensor* input; TfLiteTensor* begin; TfLiteTensor* end; @@ -57,63 +57,6 @@ struct StridedSliceContext { int dims; }; -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."); - - // TODO(soroosh): optimize for constant tensors to do allocation in Prepare - op_context.output->allocation_type = kTfLiteDynamic; - return kTfLiteOk; -} // namespace strided_slice - -// TODO(soroosh): consolidate with BytesRequired in interpreter.h -TfLiteStatus BytesRequired(TfLiteContext* context, TfLiteType type, - const int* dims, int dims_size, size_t* bytes) { - // TODO(aselle): Check for overflow here using overflow.h in TensorFlow - // MultiplyWithoutOverflow. - TF_LITE_ENSURE(context, bytes != nullptr); - size_t count = 1; - for (int k = 0; k < dims_size; k++) count *= dims[k]; - switch (type) { - case kTfLiteFloat32: - *bytes = sizeof(float) * count; - break; - case kTfLiteInt32: - *bytes = sizeof(int32_t) * count; - break; - case kTfLiteUInt8: - *bytes = sizeof(uint8_t) * count; - break; - case kTfLiteInt64: - *bytes = sizeof(int64_t) * count; - break; - default: - return kTfLiteError; - } - return kTfLiteOk; -} - // 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; @@ -144,43 +87,44 @@ inline int32_t ClampedIndex(int32_t index, int dim, bool pos_stride) { std::min(std::max(index, -dim), dim - 1), dim)); } -template -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - StridedSliceContext op_context(context, node); +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); +} - std::vector starts; - std::vector stops; - std::vector strides; +// 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) { - int dim = op_context.input->dims->data[idx]; - int32_t stride = GetTensorData(op_context.strides)[idx]; + 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"); - bool pos_stride = stride > 0; - - int32_t begin = - op_context.params->begin_mask & (1 << idx) - ? pos_stride ? 0 : dim - 1 - : ClampedIndex(GetTensorData(op_context.begin)[idx], dim, - pos_stride); - int32_t end = - op_context.params->end_mask & (1 << idx) - ? pos_stride ? dim : -1 - : ClampedIndex(GetTensorData(op_context.end)[idx], dim, - pos_stride); + + 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))) { + if (!(op_context->params->shrink_axis_mask & (1 << idx))) { output_shape_vector.push_back(dim_shape); } - - starts.emplace_back(begin); - stops.emplace_back(end); - strides.emplace_back(stride); } TfLiteIntArray* output_shape = @@ -189,35 +133,83 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { 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); } - TF_LITE_ENSURE_STATUS( - context->ResizeTensor(context, op_context.output, output_shape)); - - size_t required_bytes; - TF_LITE_ENSURE_OK( - context, - BytesRequired(context, op_context.output->type, output_shape->data, - output_shape->size, &required_bytes)); - TfLiteTensorRealloc(required_bytes, op_context.output); - - op_context.params->begin_mask = + int 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 = + int end_mask = ReverseMaskBits(op_context.params->end_mask, op_context.dims); + int 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), \ +#define TF_LITE_STRIDED_SLICE(kernel_type, data_type) \ + kernel_type::StridedSlice( \ + GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), begin_mask, end_mask, shrink_axis_mask, \ + starts, stops, strides, GetTensorData(op_context.output), \ GetTensorDims(op_context.output)) switch (op_context.input->type) { diff --git a/tensorflow/contrib/lite/kernels/strided_slice_test.cc b/tensorflow/contrib/lite/kernels/strided_slice_test.cc index 5bc7dc353b4904bdb182ce029b9b7c654e4a5f33..5c98c5f43181fe75f35716dae5682113bde883ec 100644 --- a/tensorflow/contrib/lite/kernels/strided_slice_test.cc +++ b/tensorflow/contrib/lite/kernels/strided_slice_test.cc @@ -21,6 +21,7 @@ limitations under the License. namespace tflite { namespace { +using ::int32; using ::testing::ElementsAreArray; class StridedSliceOpModel : public SingleOpModel { @@ -521,6 +522,28 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis7) { EXPECT_TRUE(m.GetOutputShape().empty()); EXPECT_THAT(m.GetOutput(), ElementsAreArray({1})); } + +// This tests catches a very subtle bug that was fixed by cl/188403234. +TEST(StridedSliceOpTest, RunTwice) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 1, 0, 0, 0, 0); + + auto setup_inputs = [&m]() { + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, 0}); + m.SetEnd({2, 2}); + m.SetStrides({1, 1}); + }; + + setup_inputs(); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 4, 5})); + + setup_inputs(); + m.Invoke(); + // Prior to cl/188403234 this was {4, 5}. + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 4, 5})); +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/sub.cc b/tensorflow/contrib/lite/kernels/sub.cc index ddaf498d5bac0109429224e7cf66cb3debcabc22..c15a7a50a4c5e6584f1329ffc9ecc155b3675219 100644 --- a/tensorflow/contrib/lite/kernels/sub.cc +++ b/tensorflow/contrib/lite/kernels/sub.cc @@ -26,7 +26,7 @@ namespace ops { namespace builtin { namespace sub { -// This file has three implementation of Div. +// This file has three implementation of Sub. enum KernelType { kReference, kGenericOptimized, // Neon-free @@ -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,49 +61,62 @@ 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 EvalSubFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteSubParams* params, TfLiteTensor* input1, - TfLiteTensor* input2, TfLiteTensor* output) { +void EvalFloat(TfLiteContext* context, TfLiteNode* node, + TfLiteSubParams* 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_Sub(type) \ - type::Sub(GetTensorData(input1), GetTensorDims(input1), \ - GetTensorData(input2), GetTensorDims(input2), \ - output_activation_min, output_activation_max, \ - GetTensorData(output), GetTensorDims(output)) +#define TF_LITE_SUB(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_Sub(reference_ops); + if (data->requires_broadcast) { + TF_LITE_SUB(reference_ops, BroadcastSub); + } else { + TF_LITE_SUB(reference_ops, Sub); + } } else { - TF_LITE_Sub(optimized_ops); + if (data->requires_broadcast) { + TF_LITE_SUB(optimized_ops, BroadcastSub); + } else { + TF_LITE_SUB(optimized_ops, Sub); + } } -#undef TF_LITE_Sub +#undef TF_LITE_SUB } 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) { - EvalSubFloat(context, node, params, input1, input2, output); + EvalFloat(context, node, params, data, input1, input2, output); } else { context->ReportError(context, "Inputs and outputs not all float types."); return kTfLiteError; @@ -99,19 +128,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { } // namespace sub TfLiteRegistration* Register_SUB_REF() { - static TfLiteRegistration r = {nullptr, nullptr, sub::Prepare, + static TfLiteRegistration r = {sub::Init, sub::Free, sub::Prepare, sub::Eval}; return &r; } TfLiteRegistration* Register_SUB_GENERIC_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, sub::Prepare, + static TfLiteRegistration r = {sub::Init, sub::Free, sub::Prepare, sub::Eval}; return &r; } TfLiteRegistration* Register_SUB_NEON_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, sub::Prepare, + static TfLiteRegistration r = {sub::Init, sub::Free, sub::Prepare, sub::Eval}; return &r; } diff --git a/tensorflow/contrib/lite/kernels/sub_test.cc b/tensorflow/contrib/lite/kernels/sub_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..fdbb4243bb5dbbbd92ca09f730f39d0413c8f133 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/sub_test.cc @@ -0,0 +1,117 @@ +/* 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 BaseSubOpModel : public SingleOpModel { + public: + BaseSubOpModel(const TensorData& input1, const TensorData& input2, + const TensorData& output, + ActivationFunctionType activation_type) { + input1_ = AddInput(input1); + input2_ = AddInput(input2); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_SUB, BuiltinOptions_SubOptions, + CreateSubOptions(builder_, activation_type).Union()); + BuildInterpreter({GetShape(input1_), GetShape(input2_)}); + } + + int input1() { return input1_; } + int input2() { return input2_; } + + protected: + int input1_; + int input2_; + int output_; +}; + +class FloatSubOpModel : public BaseSubOpModel { + public: + using BaseSubOpModel::BaseSubOpModel; + + std::vector GetOutput() { return ExtractVector(output_); } +}; + +TEST(FloatSubOpModel, NoActivation) { + FloatSubOpModel 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, 1.7, 0.5}); + m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.8}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-2.1, 0.0, 1.4, -0.3}))); +} + +TEST(FloatSubOpModel, ActivationRELU_N1_TO_1) { + FloatSubOpModel 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, 1.7, 0.5}); + m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.8}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-1.0, 0.0, 1.0, -0.3}))); +} + +TEST(FloatSubOpModel, VariousInputShapes) { + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + FloatSubOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-2.0, 0.2, 1.7, 0.5, -1.1, 2.0}); + m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.8, -1.1, 0.1}); + m.Invoke(); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-2.1, 0.0, 1.4, -0.3, 0.0, 1.9}))) + << "With shape number " << i; + } +} + +TEST(FloatSubOpModel, WithBroadcast) { + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + FloatSubOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, {}}, // always a scalar + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-2.0, 0.2, 1.7, 0.5, -1.1, 2.0}); + m.PopulateTensor(m.input2(), {0.5}); + m.Invoke(); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-2.5, -0.3, 1.2, 0.0, -1.6, 1.5}))) + << "With shape number " << i; + } +} + +} // 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/test_util.cc b/tensorflow/contrib/lite/kernels/test_util.cc index 3a58e7ec321f649a6cae4cc0969807c2c74c6529..0bb28b50b2a5e5a9fd803ecf1b0928026f63881e 100644 --- a/tensorflow/contrib/lite/kernels/test_util.cc +++ b/tensorflow/contrib/lite/kernels/test_util.cc @@ -141,8 +141,8 @@ void SingleOpModel::SetBuiltinOp(BuiltinOperator type, void SingleOpModel::SetCustomOp( const string& name, const std::vector& custom_option, - const std::function& registeration) { - custom_registrations_[name] = registeration; + const std::function& registration) { + custom_registrations_[name] = registration; opcodes_.push_back( CreateOperatorCodeDirect(builder_, BuiltinOperator_CUSTOM, name.data())); operators_.push_back(CreateOperator( @@ -172,11 +172,14 @@ void SingleOpModel::BuildInterpreter( 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); @@ -184,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 cc445299ff9f0b75610c7ff38f28facbbbe5587d..a9064d54e7704d52eefa34f6bf446ec1cfe68fe1 100644 --- a/tensorflow/contrib/lite/kernels/test_util.h +++ b/tensorflow/contrib/lite/kernels/test_util.h @@ -39,10 +39,10 @@ inline std::vector Quantize(const std::vector& data, float scale, int32_t zero_point) { std::vector q; for (float f : data) { - q.push_back(std::max( + q.push_back(static_cast(std::max( std::numeric_limits::min(), - std::min(std::numeric_limits::max(), - static_cast(std::round(zero_point + (f / scale)))))); + std::min(std::numeric_limits::max(), + std::round(zero_point + (f / scale)))))); } return q; } @@ -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() {} @@ -178,11 +195,16 @@ 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, std::initializer_list data); @@ -197,6 +219,36 @@ class SingleOpModel { 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); diff --git a/tensorflow/contrib/lite/kernels/test_util_test.cc b/tensorflow/contrib/lite/kernels/test_util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..1e10e89061213b6fcabd404310893dd97a51d83f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/test_util_test.cc @@ -0,0 +1,51 @@ +/* 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/test_util.h" +#include + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +TEST(TestUtilTest, QuantizeVector) { + std::vector data = {-1.0, -0.5, 0.0, 0.5, 1.0, 1000.0}; + auto q_data = Quantize(data, /*scale=*/1.0, /*zero_point=*/0); + std::vector expected = {0, 0, 0, 1, 1, 255}; + EXPECT_THAT(q_data, ElementsAreArray(expected)); +} + +TEST(TestUtilTest, QuantizeVectorScalingDown) { + std::vector data = {-1.0, -0.5, 0.0, 0.5, 1.0, 1000.0}; + auto q_data = Quantize(data, /*scale=*/10.0, /*zero_point=*/0); + std::vector expected = {0, 0, 0, 0, 0, 100}; + EXPECT_THAT(q_data, ElementsAreArray(expected)); +} + +TEST(TestUtilTest, QuantizeVectorScalingUp) { + std::vector data = {-1.0, -0.5, 0.0, 0.5, 1.0, 1000.0}; + auto q_data = Quantize(data, /*scale=*/0.1, /*zero_point=*/0); + std::vector expected = {0, 0, 0, 5, 10, 255}; + EXPECT_THAT(q_data, ElementsAreArray(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/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 index 093814bc4494041809d6e82965519e38a8997f49..d3c10a9bb7b07404ccd8cfe2636473a622b91787 100644 --- a/tensorflow/contrib/lite/kernels/transpose.cc +++ b/tensorflow/contrib/lite/kernels/transpose.cc @@ -90,7 +90,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* 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)); - TfLiteTensorRealloc(op_context.output->bytes, op_context.output); } // Reverse the permuted axes and convert to 4D due to the way Dims are diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc index 9cdb58714edb5fee771fc45f3c53a570f8fb28d1..42941a97db70adb37c20500c8f9438adfea25389 100644 --- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc +++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.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/activation_functor.h" +#include "tensorflow/contrib/lite/kernels/internal/kernel_utils.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" @@ -359,7 +360,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { 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. + // check the existence of only one to get the condition. const bool use_cifg = (input_to_input_weights == nullptr); const bool use_peephole = (cell_to_output_weights != nullptr); @@ -380,135 +381,57 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { output_gate_scratch = scratch_buffer->data.f + 3 * n_cell * n_batch; } + // Check optional tensors, the respective pointers can be null. + const float* input_to_input_weights_ptr = + (use_cifg) ? nullptr : input_to_input_weights->data.f; + const float* recurrent_to_input_weights_ptr = + (use_cifg) ? nullptr : recurrent_to_input_weights->data.f; + const float* input_gate_bias_ptr = + (use_cifg) ? nullptr : input_gate_bias->data.f; + const float* cell_to_input_weights_ptr = + (use_peephole && !use_cifg) ? cell_to_input_weights->data.f : nullptr; + const float* cell_to_forget_weights_ptr = + (use_peephole) ? cell_to_forget_weights->data.f : nullptr; + const float* cell_to_output_weights_ptr = + (use_peephole) ? cell_to_output_weights->data.f : nullptr; + const float* projection_weights_ptr = + (projection_weights == nullptr) ? nullptr : projection_weights->data.f; + const float* projection_bias_ptr = + (projection_bias == nullptr) ? nullptr : projection_bias->data.f; + + // Required tensors, pointers are non-null. + const float* input_to_forget_weights_ptr = input_to_forget_weights->data.f; + const float* input_to_cell_weights_ptr = input_to_cell_weights->data.f; + const float* input_to_output_weights_ptr = input_to_output_weights->data.f; + const float* recurrent_to_forget_weights_ptr = + recurrent_to_forget_weights->data.f; + const float* recurrent_to_cell_weights_ptr = + recurrent_to_cell_weights->data.f; + const float* recurrent_to_output_weights_ptr = + recurrent_to_output_weights->data.f; + const float* forget_gate_bias_ptr = forget_gate_bias->data.f; + const float* cell_bias_ptr = cell_bias->data.f; + const float* output_gate_bias_ptr = output_gate_bias->data.f; + + float* output_state_ptr = output_state->data.f; + float* cell_state_ptr = cell_state->data.f; + 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); + const float* input_ptr_batch = input->data.f + t * n_batch * n_input; + float* output_ptr_batch = output->data.f + t * n_batch * n_output; + + kernel_utils::LstmStep( + input_ptr_batch, input_to_input_weights_ptr, + input_to_forget_weights_ptr, input_to_cell_weights_ptr, + input_to_output_weights_ptr, recurrent_to_input_weights_ptr, + recurrent_to_forget_weights_ptr, recurrent_to_cell_weights_ptr, + recurrent_to_output_weights_ptr, cell_to_input_weights_ptr, + cell_to_forget_weights_ptr, cell_to_output_weights_ptr, + input_gate_bias_ptr, forget_gate_bias_ptr, cell_bias_ptr, + output_gate_bias_ptr, projection_weights_ptr, projection_bias_ptr, + params, n_batch, n_cell, n_input, n_output, output_state_ptr, + cell_state_ptr, input_gate_scratch, forget_gate_scratch, cell_scratch, + output_gate_scratch, output_ptr_batch); } return kTfLiteOk; } diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc index 7ce87e4deb81b95ae08b0c55397873ae4582eec6..ac00c37b67dcbe77023a2495a698967ca555b1d5 100644 --- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc +++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc @@ -23,6 +23,7 @@ limitations under the License. #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 { @@ -88,42 +89,6 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { return kTfLiteOk; } -namespace { -void RnnStep(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 input_weights_stride, - int recurrent_weights_stride, TfLiteFusedActivation activation, - float* hidden_state_ptr_batch, float* output_ptr_batch) { - // 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(activation))(output_ptr_batch[o]); - hidden_state_ptr_batch[o] = output_ptr_batch[o]; - } -} -} // namespace - TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); @@ -147,30 +112,25 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { (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]; // 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) { - // Unroll the sequence + // Initialize the pointer to hidden state. + 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++) { - 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; - // Initialize the pointer to input and output. - const float* input_ptr_batch = - input->data.f + s * input_size * batch_size + b * input_size; - float* output_ptr_batch = - output->data.f + s * num_units * batch_size + b * num_units; - - RnnStep(input_ptr_batch, input_weights_ptr, recurrent_weights_ptr, - bias_ptr, input_size, num_units, input_weights_stride, - recurrent_weights_stride, params->activation, - hidden_state_ptr_batch, output_ptr_batch); - } + // Initialize the pointer to input and output. + const float* input_ptr_batch = + input->data.f + s * input_size * batch_size; + float* output_ptr_batch = output->data.f + s * num_units * batch_size; + + 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 @@ -184,10 +144,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { float* output_ptr_batch = output->data.f + b * num_units * max_time + s * num_units; - RnnStep(input_ptr_batch, input_weights_ptr, recurrent_weights_ptr, - bias_ptr, input_size, num_units, input_weights_stride, - recurrent_weights_stride, params->activation, - hidden_state_ptr_batch, output_ptr_batch); + 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/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 index 5cd6c208500f3ea84ab8146f7f136e8b7851ff03..0294ec815c4820d41361b8cd4a814b74c3c1d770 100644 --- a/tensorflow/contrib/lite/memory_planner.h +++ b/tensorflow/contrib/lite/memory_planner.h @@ -34,8 +34,8 @@ class MemoryPlanner { // [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 + // Invalidates allocations made earlier. This is called when tensors sizes + // have changed. All planned allocations remain, but can't be used until // ExecuteAllocations() is called. virtual TfLiteStatus ResetAllocations() = 0; }; diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc index b36bfcef843cc3b7a9647192c18dfb0dbf6307e9..f28d56af674982859c5776c5cc1dc48f7b38a6f9 100644 --- a/tensorflow/contrib/lite/model.cc +++ b/tensorflow/contrib/lite/model.cc @@ -32,11 +32,46 @@ namespace tflite { const char* kEmptyTensorName = ""; +// Loads a model from `filename`. If `mmap_file` is true then use mmap, +// otherwise make a copy of the model in a buffer. +std::unique_ptr GetAllocationFromFile(const char* filename, + bool mmap_file, + ErrorReporter* error_reporter, + bool use_nnapi) { + std::unique_ptr allocation; + if (mmap_file) { + if (use_nnapi && NNAPIExists()) + allocation.reset(new NNAPIAllocation(filename, error_reporter)); + else + allocation.reset(new MMAPAllocation(filename, error_reporter)); + } else { + allocation.reset(new FileCopyAllocation(filename, error_reporter)); + } + return allocation; +} + std::unique_ptr FlatBufferModel::BuildFromFile( const char* filename, ErrorReporter* error_reporter) { std::unique_ptr model; - model.reset(new FlatBufferModel(filename, /*mmap_file=*/true, error_reporter, - /*use_nnapi=*/true)); + auto allocation = GetAllocationFromFile(filename, /*mmap_file=*/true, + error_reporter, /*use_nnapi=*/true); + model.reset(new FlatBufferModel(allocation.release(), error_reporter)); + if (!model->initialized()) model.reset(); + return model; +} + +std::unique_ptr FlatBufferModel::VerifyAndBuildFromFile( + const char* filename, TfLiteVerifier* verifier, + ErrorReporter* error_reporter) { + std::unique_ptr model; + auto allocation = GetAllocationFromFile(filename, /*mmap_file=*/true, + error_reporter, /*use_nnapi=*/true); + if (verifier && + !verifier->Verify(static_cast(allocation->base()), + allocation->bytes(), error_reporter)) { + return model; + } + model.reset(new FlatBufferModel(allocation.release(), error_reporter)); if (!model->initialized()) model.reset(); return model; } @@ -44,7 +79,9 @@ std::unique_ptr FlatBufferModel::BuildFromFile( std::unique_ptr FlatBufferModel::BuildFromBuffer( const char* buffer, size_t buffer_size, ErrorReporter* error_reporter) { std::unique_ptr model; - model.reset(new FlatBufferModel(buffer, buffer_size, error_reporter)); + Allocation* allocation = + new MemoryAllocation(buffer, buffer_size, error_reporter); + model.reset(new FlatBufferModel(allocation, error_reporter)); if (!model->initialized()) model.reset(); return model; } @@ -57,23 +94,6 @@ std::unique_ptr FlatBufferModel::BuildFromModel( return model; } -FlatBufferModel::FlatBufferModel(const char* filename, bool mmap_file, - ErrorReporter* error_reporter, bool use_nnapi) - : error_reporter_(error_reporter ? error_reporter - : DefaultErrorReporter()) { - if (mmap_file) { - if (use_nnapi && NNAPIExists()) - allocation_ = new NNAPIAllocation(filename, error_reporter); - else - allocation_ = new MMAPAllocation(filename, error_reporter); - } else { - allocation_ = new FileCopyAllocation(filename, error_reporter); - } - if (!allocation_->valid() || !CheckModelIdentifier()) return; - - model_ = ::tflite::GetModel(allocation_->base()); -} - bool FlatBufferModel::CheckModelIdentifier() const { if (!tflite::ModelBufferHasIdentifier(allocation_->base())) { const char* ident = flatbuffers::GetBufferIdentifier(allocation_->base()); @@ -85,21 +105,21 @@ bool FlatBufferModel::CheckModelIdentifier() const { return true; } -FlatBufferModel::FlatBufferModel(const char* ptr, size_t num_bytes, +FlatBufferModel::FlatBufferModel(const Model* model, ErrorReporter* error_reporter) : error_reporter_(error_reporter ? error_reporter : DefaultErrorReporter()) { - allocation_ = new MemoryAllocation(ptr, num_bytes, error_reporter); - if (!allocation_->valid()) return; - - model_ = ::tflite::GetModel(allocation_->base()); + model_ = model; } -FlatBufferModel::FlatBufferModel(const Model* model, +FlatBufferModel::FlatBufferModel(Allocation* allocation, ErrorReporter* error_reporter) : error_reporter_(error_reporter ? error_reporter : DefaultErrorReporter()) { - model_ = model; + allocation_ = allocation; + if (!allocation_->valid() || !CheckModelIdentifier()) return; + + model_ = ::tflite::GetModel(allocation_->base()); } FlatBufferModel::~FlatBufferModel() { delete allocation_; } @@ -124,19 +144,25 @@ TfLiteStatus InterpreterBuilder::BuildLocalIndexToRegistrationMapping() { auto opcodes = model_->operator_codes(); for (const OperatorCode* opcode : *opcodes) { TfLiteRegistration* registration = nullptr; - - if (opcode->builtin_code() != BuiltinOperator_CUSTOM) { - auto x = opcode->builtin_code(); - flatbuffer_op_index_to_registration_types_.push_back(x); - registration = op_resolver_.FindOp(x); + auto builtin_code = opcode->builtin_code(); + if (builtin_code > BuiltinOperator_MAX || + builtin_code < BuiltinOperator_MIN) { + error_reporter_->Report( + "Op builtin_code out or range: %d. Are you using old TFLite binary " + "with newer model?", + builtin_code); + status = kTfLiteError; + } else if (builtin_code != BuiltinOperator_CUSTOM) { + flatbuffer_op_index_to_registration_types_.push_back(builtin_code); + registration = op_resolver_.FindOp(builtin_code); if (registration == nullptr) { error_reporter_->Report("Didn't find op for builtin opcode '%s'\n", - EnumNameBuiltinOperator(x)); + EnumNameBuiltinOperator(builtin_code)); status = kTfLiteError; } } 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(); @@ -278,6 +304,11 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, case BuiltinOperator_RELU_N1_TO_1: case BuiltinOperator_RELU6: case BuiltinOperator_CONCAT_EMBEDDINGS: + case BuiltinOperator_EXP: + case BuiltinOperator_TOPK_V2: + case BuiltinOperator_LOG_SOFTMAX: + case BuiltinOperator_CAST: + case BuiltinOperator_DEQUANTIZE: break; case BuiltinOperator_LSH_PROJECTION: { TfLiteLSHProjectionParams* params = @@ -453,6 +484,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, builtin_data = reinterpret_cast(params); break; } + case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM: case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM: case BuiltinOperator_LSTM: { TfLiteLSTMParams* params = MallocPOD(); @@ -469,6 +501,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, auto* params = MallocPOD(); if (auto* schema_params = op->builtin_options_as_ResizeBilinearOptions()) { + params->align_corners = schema_params->align_corners(); } builtin_data = reinterpret_cast(params); break; @@ -532,6 +565,14 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, 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()) { @@ -555,6 +596,11 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, builtin_data = reinterpret_cast(params); break; } + case BuiltinOperator_DELEGATE: { + // TODO(ycling): Revisit when supporting saving delegated models. + error_reporter->Report("DELEGATE op shouldn't exist in model."); + break; + } } return builtin_data; } diff --git a/tensorflow/contrib/lite/model.h b/tensorflow/contrib/lite/model.h index a467df5bb4eee3f6ce814512cb8b74bf09a6a4e7..38eea0e26bc97c09dd84094972b5d64770b003e2 100644 --- a/tensorflow/contrib/lite/model.h +++ b/tensorflow/contrib/lite/model.h @@ -41,6 +41,17 @@ limitations under the License. namespace tflite { +// Abstract interface that verifies whether a given model is legit. +// It facilitates the use-case to verify and build a model without loading it +// twice. +class TfLiteVerifier { + public: + // Returns true if the model is legit. + virtual bool Verify(const char* data, int length, + ErrorReporter* reporter) = 0; + virtual ~TfLiteVerifier() {} +}; + // An RAII object that represents a read-only tflite model, copied from disk, // or mmapped. This uses flatbuffers as the serialization format. class FlatBufferModel { @@ -50,6 +61,12 @@ class FlatBufferModel { const char* filename, ErrorReporter* error_reporter = DefaultErrorReporter()); + // Verifies whether the content of the file is legit, then builds a model + // based on the file. Returns a nullptr in case of failure. + static std::unique_ptr VerifyAndBuildFromFile( + const char* filename, TfLiteVerifier* verifier = nullptr, + ErrorReporter* error_reporter = DefaultErrorReporter()); + // Builds a model based on a pre-loaded flatbuffer. The caller retains // ownership of the buffer and should keep it alive until the returned object // is destroyed. Returns a nullptr in case of failure. @@ -64,7 +81,7 @@ class FlatBufferModel { const tflite::Model* model_spec, ErrorReporter* error_reporter = DefaultErrorReporter()); - // Releases memory or unmaps mmaped meory. + // Releases memory or unmaps mmaped memory. ~FlatBufferModel(); // Copying or assignment is disallowed to simplify ownership semantics. @@ -82,23 +99,9 @@ class FlatBufferModel { bool CheckModelIdentifier() const; private: - // Loads a model from `filename`. If `mmap_file` is true then use mmap, - // otherwise make a copy of the model in a buffer. - // - // Note, if `error_reporter` is null, then a DefaultErrorReporter() will be - // used. - explicit FlatBufferModel( - const char* filename, bool mmap_file = true, - ErrorReporter* error_reporter = DefaultErrorReporter(), - bool use_nnapi = false); - - // Loads a model from `ptr` and `num_bytes` of the model file. The `ptr` has - // to remain alive and unchanged until the end of this flatbuffermodel's - // lifetime. - // - // Note, if `error_reporter` is null, then a DefaultErrorReporter() will be - // used. - FlatBufferModel(const char* ptr, size_t num_bytes, + // Loads a model from a given allocation. FlatBufferModel will take over the + // ownership of `allocation`, and delete it in desctructor. + FlatBufferModel(Allocation* allocation, ErrorReporter* error_reporter = DefaultErrorReporter()); // Loads a model from Model flatbuffer. The `model` has to remain alive and diff --git a/tensorflow/contrib/lite/model_test.cc b/tensorflow/contrib/lite/model_test.cc index 66f22fd66a9ae0d35553a1f780ef73a5c5994c99..ae6c1ece18963f11f48a6f07bea4065ce39687e0 100644 --- a/tensorflow/contrib/lite/model_test.cc +++ b/tensorflow/contrib/lite/model_test.cc @@ -209,6 +209,38 @@ TEST(BasicFlatBufferModel, TestNullModel) { ASSERT_EQ(interpreter.get(), nullptr); } +// Mocks the verifier by setting the result in ctor. +class FakeVerifier : public tflite::TfLiteVerifier { + public: + explicit FakeVerifier(bool result) : result_(result) {} + bool Verify(const char* data, int length, + tflite::ErrorReporter* reporter) override { + return result_; + } + + private: + bool result_; +}; + +TEST(BasicFlatBufferModel, TestWithTrueVerifier) { + FakeVerifier verifier(true); + ASSERT_TRUE(FlatBufferModel::VerifyAndBuildFromFile( + "tensorflow/contrib/lite/testdata/test_model.bin", + &verifier)); +} + +TEST(BasicFlatBufferModel, TestWithFalseVerifier) { + FakeVerifier verifier(false); + ASSERT_FALSE(FlatBufferModel::VerifyAndBuildFromFile( + "tensorflow/contrib/lite/testdata/test_model.bin", + &verifier)); +} + +TEST(BasicFlatBufferModel, TestWithNullVerifier) { + ASSERT_TRUE(FlatBufferModel::VerifyAndBuildFromFile( + "tensorflow/contrib/lite/testdata/test_model.bin", nullptr)); +} + struct TestErrorReporter : public ErrorReporter { int Report(const char* format, va_list args) override { calls++; 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/speech_test.cc b/tensorflow/contrib/lite/models/speech_test.cc index daa8c3100b64e9290256aa14a6ab641f19174a0a..a354179a9480c136d65f83836d81f69c2089fdbe 100644 --- a/tensorflow/contrib/lite/models/speech_test.cc +++ b/tensorflow/contrib/lite/models/speech_test.cc @@ -97,7 +97,12 @@ bool ConvertCsvData(const string& model_name, const string& in_name, return true; } -TEST(SpeechTest, HotwordOkGoogleRank1Test) { +class SpeechTest : public ::testing::TestWithParam { + protected: + int GetMaxInvocations() { return GetParam(); } +}; + +TEST_P(SpeechTest, HotwordOkGoogleRank1Test) { std::stringstream os; ASSERT_TRUE(ConvertCsvData( "speech_hotword_model_rank1.tflite", "speech_hotword_model_in.csv", @@ -105,11 +110,11 @@ TEST(SpeechTest, HotwordOkGoogleRank1Test) { /*output_tensor=*/"18", /*persistent_tensors=*/"4", /*sequence_size=*/40, &os)); testing::TfLiteDriver test_driver(/*use_nnapi=*/false); - ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } -TEST(SpeechTest, HotwordOkGoogleRank2Test) { +TEST_P(SpeechTest, HotwordOkGoogleRank2Test) { std::stringstream os; ASSERT_TRUE(ConvertCsvData( "speech_hotword_model_rank2.tflite", "speech_hotword_model_in.csv", @@ -117,11 +122,11 @@ TEST(SpeechTest, HotwordOkGoogleRank2Test) { /*output_tensor=*/"18", /*persistent_tensors=*/"1", /*sequence_size=*/40, &os)); testing::TfLiteDriver test_driver(/*use_nnapi=*/false); - ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } -TEST(SpeechTest, SpeakerIdOkGoogleTest) { +TEST_P(SpeechTest, SpeakerIdOkGoogleTest) { std::stringstream os; ASSERT_TRUE(ConvertCsvData( "speech_speakerid_model.tflite", "speech_speakerid_model_in.csv", @@ -130,11 +135,11 @@ TEST(SpeechTest, SpeakerIdOkGoogleTest) { /*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)) + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } -TEST(SpeechTest, AsrAmTest) { +TEST_P(SpeechTest, AsrAmTest) { std::stringstream os; ASSERT_TRUE( ConvertCsvData("speech_asr_am_model.tflite", "speech_asr_am_model_in.csv", @@ -143,7 +148,7 @@ TEST(SpeechTest, AsrAmTest) { /*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)) + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } @@ -151,15 +156,16 @@ TEST(SpeechTest, AsrAmTest) { // 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) { +TEST_P(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)) + ASSERT_TRUE( + testing::ParseAndRunTests(&in_file, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } -TEST(SpeechTest, EndpointerTest) { +TEST_P(SpeechTest, EndpointerTest) { std::stringstream os; ASSERT_TRUE(ConvertCsvData( "speech_endpointer_model.tflite", "speech_endpointer_model_in.csv", @@ -168,11 +174,11 @@ TEST(SpeechTest, EndpointerTest) { /*persistent_tensors=*/"28,29,49,50", /*sequence_size=*/320, &os)); testing::TfLiteDriver test_driver(/*use_nnapi=*/false); - ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } -TEST(SpeechTest, TtsTest) { +TEST_P(SpeechTest, TtsTest) { std::stringstream os; ASSERT_TRUE(ConvertCsvData("speech_tts_model.tflite", "speech_tts_model_in.csv", @@ -181,9 +187,19 @@ TEST(SpeechTest, TtsTest) { /*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)) + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations())) << test_driver.GetErrorMessage(); } +// Define two instantiations. The "ShortTests" instantiations is used when +// running the tests on Android, in order to prevent timeouts (It takes about +// 200s just to bring up the Android emulator.) +static const int kAllInvocations = -1; +static const int kFirstFewInvocations = 10; +INSTANTIATE_TEST_CASE_P(LongTests, SpeechTest, + ::testing::Values(kAllInvocations)); +INSTANTIATE_TEST_CASE_P(ShortTests, SpeechTest, + ::testing::Values(kFirstFewInvocations)); + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h index 76032771af2c8e099aed498b2071816646f3b606..bd49d327c995ef53dc6cf9f8301ab749c925b2c7 100644 --- a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h +++ b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h @@ -569,7 +569,7 @@ enum { ANEURALNETWORKS_LOGISTIC = 14, /** - * Projects an input to a bit vector via locality senstive hashing. + * Projects an input to a bit vector via locality sensitive hashing. * * Inputs: * * 0: Hash functions. Dim.size == 2, DataType: Float. diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc index da9ceec2f1401745ba477824bf494ee5b0ee1187..9d00d965d397d3af0dd454609901685ba965ded6 100644 --- a/tensorflow/contrib/lite/nnapi_delegate.cc +++ b/tensorflow/contrib/lite/nnapi_delegate.cc @@ -323,6 +323,7 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN: case tflite::BuiltinOperator_EMBEDDING_LOOKUP: case tflite::BuiltinOperator_EMBEDDING_LOOKUP_SPARSE: + case tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM: case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM: case tflite::BuiltinOperator_L2_NORMALIZATION: case tflite::BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION: @@ -335,12 +336,19 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, 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: + case tflite::BuiltinOperator_LOG_SOFTMAX: + case tflite::BuiltinOperator_DEQUANTIZE: + case tflite::BuiltinOperator_DELEGATE: + case tflite::BuiltinOperator_CAST: 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/python/BUILD b/tensorflow/contrib/lite/python/BUILD index 2d8c49b7d7a5ae5c180f100a399a1870679c455f..d6f39219ed93718af767e7ad1e1de2d843062c61 100644 --- a/tensorflow/contrib/lite/python/BUILD +++ b/tensorflow/contrib/lite/python/BUILD @@ -4,6 +4,38 @@ package(default_visibility = ["//tensorflow:internal"]) load("//tensorflow:tensorflow.bzl", "py_test") +filegroup( + name = "interpreter_test_data", + srcs = glob(["**/testdata/*"]), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "interpreter", + srcs = [ + "interpreter.py", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/contrib/lite/python/interpreter_wrapper:tensorflow_wrap_interpreter_wrapper", + ], +) + +py_test( + name = "interpreter_test", + srcs = ["interpreter_test.py"], + data = [":interpreter_test_data"], + srcs_version = "PY2AND3", + tags = ["no_oss"], + deps = [ + ":interpreter", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:platform_test", + ], +) + py_library( name = "lite", srcs = ["lite.py"], @@ -28,6 +60,7 @@ py_library( visibility = ["//visibility:public"], deps = [ "//tensorflow/contrib/framework:framework_py", + "//tensorflow/core:protos_all_py", "//tensorflow/python:platform", ], ) @@ -36,7 +69,10 @@ py_test( name = "lite_test", srcs = ["lite_test.py"], srcs_version = "PY2AND3", - tags = ["no_oss"], + tags = [ + "no-internal-py3", + "no_oss", + ], deps = [ ":lite", ":op_hint", diff --git a/tensorflow/contrib/lite/python/interpreter.py b/tensorflow/contrib/lite/python/interpreter.py new file mode 100644 index 0000000000000000000000000000000000000000..5b5a7c3199a480d943aa4a8ce34710a703017d28 --- /dev/null +++ b/tensorflow/contrib/lite/python/interpreter.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. +# ============================================================================== +"""Python TF-Lite interpreter.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.lite.python.interpreter_wrapper import tensorflow_wrap_interpreter_wrapper as interpreter_wrapper + + +class Interpreter(object): + """Interpreter inferace for TF-Lite Models.""" + + def __init__(self, model_path): + """Constructor. + + Args: + model_path: Path to TF-Lite Flatbuffer file. + + Raises: + ValueError: If the interpreter was unable to open the model. + """ + self._interpreter = ( + interpreter_wrapper.InterpreterWrapper_CreateWrapperCPP(model_path)) + if not self._interpreter: + raise ValueError('Failed to open {}'.format(model_path)) + + def allocate_tensors(self): + if not self._interpreter.AllocateTensors(): + raise ValueError('Failed to allocate tensors') + + def _get_tensor_details(self, tensor_index): + """Gets tensor details. + + Args: + tensor_index: Tensor index of tensor to query. + + Returns: + a dictionary containing the name, index, shape and type of the tensor. + + Raises: + ValueError: If tensor_index is invalid. + """ + tensor_index = int(tensor_index) + tensor_name = self._interpreter.TensorName(tensor_index) + tensor_size = self._interpreter.TensorSize(tensor_index) + tensor_type = self._interpreter.TensorType(tensor_index) + + if not tensor_name or not tensor_type: + raise ValueError('Could not get tensor details') + + details = { + 'name': tensor_name, + 'index': tensor_index, + 'shape': tensor_size, + 'dtype': tensor_type, + } + + return details + + def get_input_details(self): + """Gets model input details. + + Returns: + A list of input details. + """ + return [ + self._get_tensor_details(i) for i in self._interpreter.InputIndices() + ] + + def set_tensor(self, tensor_index, value): + """Sets the value of the input. + + Args: + tensor_index: Tensor index of tensor to set. This value can be gotten from + the 'index' field in get_input_details. + value: Value of tensor to set. + + Raises: + ValueError: If the interpreter could not set the tensor. + """ + if not self._interpreter.SetTensor(tensor_index, value): + raise ValueError('Failed to set tensor') + + def resize_tensor_input(self, input_index, tensor_size): + """Resizes an input tensor. + + Args: + input_index: Tensor index of input to set. This value can be gotten from + the 'index' field in get_input_details. + tensor_size: The tensor_shape to resize the input to. + + Raises: + ValueError: If the interpreter could not resize the input tensor. + """ + if not self.ResizeInputTensor.SetTensor(input_index, tensor_size): + raise ValueError('Failed to set input') + + def get_output_details(self): + """Gets model output details. + + Returns: + A list of output details. + """ + return [ + self._get_tensor_details(i) for i in self._interpreter.OutputIndices() + ] + + def get_tensor(self, tensor_index): + """Sets the value of the input. + + Args: + tensor_index: Tensor index of tensor to get. This value can be gotten from + the 'index' field in get_output_details. + + Returns: + a numpy array. + """ + return self._interpreter.GetTensor(tensor_index) + + def invoke(self): + if not self._interpreter.Invoke(): + raise ValueError('Failed to invoke TFLite model') diff --git a/tensorflow/contrib/lite/python/interpreter_test.py b/tensorflow/contrib/lite/python/interpreter_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e0215b721c1a7627ced1080def0fbd426c94c2f4 --- /dev/null +++ b/tensorflow/contrib/lite/python/interpreter_test.py @@ -0,0 +1,82 @@ +# 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. +# ============================================================================== +"""TensorFlow Lite Python Interface: Sanity check.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.lite.python import interpreter as interpreter_wrapper +from tensorflow.python.framework import test_util +from tensorflow.python.platform import resource_loader +from tensorflow.python.platform import test + + +class InterpreterTest(test_util.TensorFlowTestCase): + + def testFloat(self): + interpreter = interpreter_wrapper.Interpreter( + resource_loader.get_path_to_datafile('testdata/permute_float.tflite')) + interpreter.allocate_tensors() + + input_details = interpreter.get_input_details() + self.assertEqual(1, len(input_details)) + self.assertEqual('input', input_details[0]['name']) + self.assertEqual(np.float32, input_details[0]['dtype']) + self.assertTrue(([1, 4] == input_details[0]['shape']).all()) + + output_details = interpreter.get_output_details() + self.assertEqual(1, len(output_details)) + self.assertEqual('output', output_details[0]['name']) + self.assertEqual(np.float32, output_details[0]['dtype']) + self.assertTrue(([1, 4] == output_details[0]['shape']).all()) + + test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32) + expected_output = np.array([[4.0, 3.0, 2.0, 1.0]], dtype=np.float32) + interpreter.set_tensor(input_details[0]['index'], test_input) + interpreter.invoke() + + output_data = interpreter.get_tensor(output_details[0]['index']) + self.assertTrue((expected_output == output_data).all()) + + def testUint8(self): + interpreter = interpreter_wrapper.Interpreter( + resource_loader.get_path_to_datafile('testdata/permute_uint8.tflite')) + interpreter.allocate_tensors() + + input_details = interpreter.get_input_details() + self.assertEqual(1, len(input_details)) + self.assertEqual('input', input_details[0]['name']) + self.assertEqual(np.uint8, input_details[0]['dtype']) + self.assertTrue(([1, 4] == input_details[0]['shape']).all()) + + output_details = interpreter.get_output_details() + self.assertEqual(1, len(output_details)) + self.assertEqual('output', output_details[0]['name']) + self.assertEqual(np.uint8, output_details[0]['dtype']) + self.assertTrue(([1, 4] == output_details[0]['shape']).all()) + + test_input = np.array([[1, 2, 3, 4]], dtype=np.uint8) + expected_output = np.array([[4, 3, 2, 1]], dtype=np.uint8) + interpreter.set_tensor(input_details[0]['index'], test_input) + interpreter.invoke() + + output_data = interpreter.get_tensor(output_details[0]['index']) + self.assertTrue((expected_output == output_data).all()) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/BUILD b/tensorflow/contrib/lite/python/interpreter_wrapper/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..453eda6e7345762666917fd501b69c7181c349e8 --- /dev/null +++ b/tensorflow/contrib/lite/python/interpreter_wrapper/BUILD @@ -0,0 +1,32 @@ +package( + default_visibility = ["//visibility:public"], +) + +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") + +cc_library( + name = "interpreter_wrapper_lib", + srcs = ["interpreter_wrapper.cc"], + hdrs = ["interpreter_wrapper.h"], + deps = [ + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:builtin_ops", + "//tensorflow/core:lib", + "//tensorflow/python:numpy_lib", + "//util/python:python_headers", + "@com_google_absl//absl/memory", + ], +) + +tf_py_wrap_cc( + name = "tensorflow_wrap_interpreter_wrapper", + srcs = [ + "interpreter_wrapper.i", + ], + deps = [ + ":interpreter_wrapper_lib", + "//util/python:python_headers", + ], +) diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc new file mode 100644 index 0000000000000000000000000000000000000000..f30067de94cf631260d42ca1f0179ee4f65582c8 --- /dev/null +++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc @@ -0,0 +1,313 @@ +/* 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/python/interpreter_wrapper/interpreter_wrapper.h" + +#include + +#include "absl/memory/memory.h" +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/model.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/python/lib/core/numpy.h" + +#if PY_MAJOR_VERSION >= 3 +#define PY_TO_CPPSTRING PyBytes_AsStringAndSize +#define CPP_TO_PYSTRING PyBytes_FromStringAndSize +#else +#define PY_TO_CPPSTRING PyString_AsStringAndSize +#define CPP_TO_PYSTRING PyString_FromStringAndSize +#endif + +namespace tflite { +namespace interpreter_wrapper { + +namespace { +std::unique_ptr CreateInterpreter( + const tflite::FlatBufferModel* model, + const tflite::ops::builtin::BuiltinOpResolver& resolver) { + if (!model) { + return nullptr; + } + + std::unique_ptr interpreter; + tflite::InterpreterBuilder(*model, resolver)(&interpreter); + if (interpreter) { + for (const int input_index : interpreter->inputs()) { + const TfLiteTensor* tensor = interpreter->tensor(input_index); + CHECK(tensor); + const TfLiteIntArray* dims = tensor->dims; + if (!dims) { + continue; + } + + std::vector input_dims(dims->data, dims->data + dims->size); + interpreter->ResizeInputTensor(input_index, input_dims); + } + } + return interpreter; +} + +int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) { + switch (tf_lite_type) { + case kTfLiteFloat32: + return NPY_FLOAT32; + case kTfLiteInt32: + return NPY_INT32; + case kTfLiteUInt8: + return NPY_UINT8; + case kTfLiteInt64: + return NPY_INT64; + case kTfLiteString: + return NPY_OBJECT; + case kTfLiteNoType: + return -1; + } + LOG(ERROR) << "Unknown TfLiteType " << tf_lite_type; + return -1; +} + +TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) { + int pyarray_type = PyArray_TYPE(array); + switch (pyarray_type) { + case NPY_FLOAT32: + return kTfLiteFloat32; + case NPY_INT32: + return kTfLiteInt32; + case NPY_UINT8: + return kTfLiteUInt8; + case NPY_INT64: + return kTfLiteInt64; + case NPY_OBJECT: + case NPY_STRING: + case NPY_UNICODE: + return kTfLiteString; + } + LOG(ERROR) << "Unknown PyArray dtype " << pyarray_type; + return kTfLiteNoType; +} + +struct PyDecrefDeleter { + void operator()(PyObject* p) const { Py_DECREF(p); } +}; + +PyObject* PyArrayFromIntVector(const int* data, npy_intp size) { + void* pydata = malloc(size * sizeof(int)); + memcpy(pydata, data, size * sizeof(int)); + return PyArray_SimpleNewFromData(1, &size, NPY_INT32, pydata); +} + +} // namespace + +InterpreterWrapper::InterpreterWrapper( + std::unique_ptr model) + : model_(std::move(model)), + resolver_(absl::make_unique()), + interpreter_(CreateInterpreter(model_.get(), *resolver_)) {} + +InterpreterWrapper::~InterpreterWrapper() {} + +bool InterpreterWrapper::AllocateTensors() { + if (!interpreter_) { + LOG(ERROR) << "Cannot allocate tensors: invalid interpreter."; + return false; + } + + if (interpreter_->AllocateTensors() != kTfLiteOk) { + LOG(ERROR) << "Unable to allocate tensors."; + return false; + } + + return true; +} + +bool InterpreterWrapper::Invoke() { + return interpreter_ ? (interpreter_->Invoke() == kTfLiteOk) : false; +} + +PyObject* InterpreterWrapper::InputIndices() const { + PyObject* np_array = PyArrayFromIntVector(interpreter_->inputs().data(), + interpreter_->inputs().size()); + + return PyArray_Return(reinterpret_cast(np_array)); +} + +PyObject* InterpreterWrapper::OutputIndices() const { + PyObject* np_array = PyArrayFromIntVector(interpreter_->outputs().data(), + interpreter_->outputs().size()); + + return PyArray_Return(reinterpret_cast(np_array)); +} + +bool InterpreterWrapper::ResizeInputTensor(int i, PyObject* value) { + if (!interpreter_) { + LOG(ERROR) << "Invalid interpreter."; + return false; + } + + std::unique_ptr array_safe( + PyArray_FromAny(value, nullptr, 0, 0, NPY_ARRAY_CARRAY, nullptr)); + if (!array_safe) { + LOG(ERROR) << "Failed to convert value into readable tensor."; + return false; + } + + PyArrayObject* array = reinterpret_cast(array_safe.get()); + + if (PyArray_NDIM(array) != 1) { + LOG(ERROR) << "Expected 1-D defining input shape."; + return false; + } + + if (PyArray_TYPE(array) != NPY_INT32) { + LOG(ERROR) << "Shape must be an int32 array"; + return false; + } + + std::vector dims(PyArray_SHAPE(array)[0]); + memcpy(dims.data(), PyArray_BYTES(array), dims.size() * sizeof(int)); + + return interpreter_->ResizeInputTensor(i, dims); +} + +std::string InterpreterWrapper::TensorName(int i) const { + if (!interpreter_ || i >= interpreter_->tensors_size() || i < 0) { + return ""; + } + + const TfLiteTensor* tensor = interpreter_->tensor(i); + return tensor->name; +} + +PyObject* InterpreterWrapper::TensorType(int i) const { + if (!interpreter_ || i >= interpreter_->tensors_size() || i < 0) { + return nullptr; + } + + const TfLiteTensor* tensor = interpreter_->tensor(i); + int typenum = TfLiteTypeToPyArrayType(tensor->type); + return PyArray_TypeObjectFromType(typenum); +} + +PyObject* InterpreterWrapper::TensorSize(int i) const { + if (!interpreter_ || i >= interpreter_->tensors_size() || i < 0) { + Py_INCREF(Py_None); + return Py_None; + } + + const TfLiteTensor* tensor = interpreter_->tensor(i); + PyObject* np_array = + PyArrayFromIntVector(tensor->dims->data, tensor->dims->size); + + return PyArray_Return(reinterpret_cast(np_array)); +} + +bool InterpreterWrapper::SetTensor(int i, PyObject* value) { + if (!interpreter_) { + LOG(ERROR) << "Invalid interpreter."; + return false; + } + + if (i >= interpreter_->tensors_size()) { + LOG(ERROR) << "Invalid tensor index: " << i << " exceeds max tensor index " + << interpreter_->tensors_size(); + return false; + } + + std::unique_ptr array_safe( + PyArray_FromAny(value, nullptr, 0, 0, NPY_ARRAY_CARRAY, nullptr)); + if (!array_safe) { + LOG(ERROR) << "Failed to convert value into readable tensor."; + return false; + } + + PyArrayObject* array = reinterpret_cast(array_safe.get()); + const TfLiteTensor* tensor = interpreter_->tensor(i); + + if (TfLiteTypeFromPyArray(array) != tensor->type) { + LOG(ERROR) << "Cannot set tensor:" + << " Got tensor of type " << TfLiteTypeFromPyArray(array) + << " but expected type " << tensor->type << " for input " << i; + return false; + } + + if (PyArray_NDIM(array) != tensor->dims->size) { + LOG(ERROR) << "Cannot set tensor: Dimension mismatch"; + return false; + } + + for (int j = 0; j < PyArray_NDIM(array); j++) { + if (tensor->dims->data[j] != PyArray_SHAPE(array)[j]) { + LOG(ERROR) << "Cannot set tensor: Dimension mismatch"; + return false; + } + } + + size_t size = PyArray_NBYTES(array); + DCHECK_EQ(size, tensor->bytes); + memcpy(tensor->data.raw, PyArray_DATA(array), size); + return true; +} + +PyObject* InterpreterWrapper::GetTensor(int i) const { + if (!interpreter_) { + LOG(ERROR) << "Invalid interpreter."; + Py_INCREF(Py_None); + return Py_None; + } + + if (i >= interpreter_->tensors_size()) { + LOG(ERROR) << "Invalid tensor index: " << i << " exceeds max tensor index " + << interpreter_->inputs().size(); + Py_INCREF(Py_None); + return Py_None; + } + + const TfLiteTensor* output_tensor = interpreter_->tensor(i); + const int tensor_size = output_tensor->bytes; + if (tensor_size <= 0) { + LOG(ERROR) << "Invalid tensor size"; + Py_INCREF(Py_None); + return Py_None; + } + + int type_num = TfLiteTypeToPyArrayType(output_tensor->type); + if (type_num == -1) { + LOG(ERROR) << "Unknown tensor type " << output_tensor->type; + Py_INCREF(Py_None); + return Py_None; + } + + void* data = malloc(tensor_size); + memcpy(data, output_tensor->data.raw, tensor_size); + + const TfLiteIntArray* output_dims = output_tensor->dims; + std::vector dims(output_dims->data, + output_dims->data + output_dims->size); + PyObject* np_array = + PyArray_SimpleNewFromData(dims.size(), dims.data(), type_num, data); + + return PyArray_Return(reinterpret_cast(np_array)); +} + +InterpreterWrapper* InterpreterWrapper::CreateWrapperCPP( + const char* model_path) { + std::unique_ptr model = + tflite::FlatBufferModel::BuildFromFile(model_path); + return model ? new InterpreterWrapper(std::move(model)) : nullptr; +} + +} // namespace interpreter_wrapper +} // namespace tflite diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h new file mode 100644 index 0000000000000000000000000000000000000000..dea71ca8794da3aea4a4266721df0a9ebbe6bb58 --- /dev/null +++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.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_CONTRIB_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_ +#define TENSORFLOW_CONTRIB_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_ + +#include +#include +#include + +#include + +// We forward declare TFLite classes here to avoid exposing them to SWIG. +namespace tflite { +namespace ops { +namespace builtin { +class BuiltinOpResolver; +} // namespace builtin +} // namespace ops + +class FlatBufferModel; +class Interpreter; + +namespace interpreter_wrapper { + +class InterpreterWrapper { + public: + // SWIG caller takes ownership of pointer. + static InterpreterWrapper* CreateWrapperCPP(const char* model_path); + + ~InterpreterWrapper(); + bool AllocateTensors(); + bool Invoke(); + + PyObject* InputIndices() const; + PyObject* OutputIndices() const; + bool ResizeInputTensor(int i, PyObject* value); + + std::string TensorName(int i) const; + PyObject* TensorType(int i) const; + PyObject* TensorSize(int i) const; + bool SetTensor(int i, PyObject* value); + PyObject* GetTensor(int i) const; + + private: + InterpreterWrapper(std::unique_ptr model); + + // InterpreterWrapper is not copyable or assignable. We avoid the use of + // InterpreterWrapper() = delete here for SWIG compatibility. + InterpreterWrapper(); + InterpreterWrapper(const InterpreterWrapper& rhs); + + const std::unique_ptr model_; + const std::unique_ptr resolver_; + const std::unique_ptr interpreter_; +}; + +} // namespace interpreter_wrapper +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_ diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.i b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.i new file mode 100644 index 0000000000000000000000000000000000000000..7f51f9f00d1b2fe057052f7b7bd52bcb65231164 --- /dev/null +++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.i @@ -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 "std_string.i" + + +%{ +#define SWIG_FILE_WITH_INIT +#include "tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h" +%} + + +%include "tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h" diff --git a/tensorflow/contrib/lite/python/lite.py b/tensorflow/contrib/lite/python/lite.py index 5d2f21653762a405a57288a7ba38323e5e42b3e1..35d224924ee4d8cd94543e10c082afee25b7630e 100644 --- a/tensorflow/contrib/lite/python/lite.py +++ b/tensorflow/contrib/lite/python/lite.py @@ -202,11 +202,12 @@ def toco_convert(input_data, input_array.name = _tensor_name(input_tensor) input_array.shape.dims.extend(map(int, input_tensor.get_shape())) - toco.inference_input_type = tflite_input_type for output_tensor in output_tensors: model.output_arrays.append(_tensor_name(output_tensor)) + # TODO(aselle): Consider handling the case of allowing quantized + # inputs to be converted to float (via the toco.inference_input_type field). data = toco_convert_protos(model.SerializeToString(), toco.SerializeToString(), input_data.SerializeToString()) diff --git a/tensorflow/contrib/lite/python/op_hint.py b/tensorflow/contrib/lite/python/op_hint.py index 7c587e38b16dc3011fc7c8bef4eec4d0ea99ec21..7908689ce4a719ab15bd49a368a87f9cad7c6d61 100644 --- a/tensorflow/contrib/lite/python/op_hint.py +++ b/tensorflow/contrib/lite/python/op_hint.py @@ -73,6 +73,7 @@ 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 @@ -118,8 +119,10 @@ class OpHint(object): 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) + # pylint: disable=protected-access + dest_op.op._set_attr(name, _attr_value_pb2.AttrValue( + tensor=tensor_value.op.node_def.attr["value"].tensor)) + # pylint: enable=protected-access def add_inputs(self, *args): """Add a sequence of inputs to the function invocation. @@ -133,10 +136,17 @@ class OpHint(object): def augmented_identity(arg): identity_op = _array_ops.identity(arg) - attr = identity_op.op.node_def.attr - attr[OpHint.FUNCTION_NAME_ATTR].s = self._function_name - attr[OpHint.FUNCTION_UUID_ATTR].s = self._unique_function_id - attr[OpHint.FUNCTION_INPUT_INDEX_ATTR].i = self._curr_input_index + # 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 @@ -154,10 +164,17 @@ class OpHint(object): def augmented_identity(arg): identity_op = _array_ops.identity(arg) - attr = identity_op.op.node_def.attr - attr[OpHint.FUNCTION_NAME_ATTR].s = self._function_name - attr[OpHint.FUNCTION_UUID_ATTR].s = self._unique_function_id - attr[OpHint.FUNCTION_OUTPUT_INDEX_ATTR].i = self._curr_output_index + # 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 diff --git a/tensorflow/contrib/lite/rpi_makefile.inc b/tensorflow/contrib/lite/rpi_makefile.inc new file mode 100644 index 0000000000000000000000000000000000000000..832ef5824bea86a368184bd7e3d17915739e9d46 --- /dev/null +++ b/tensorflow/contrib/lite/rpi_makefile.inc @@ -0,0 +1,33 @@ +# Settings for Raspberry Pi. +ifeq ($(TARGET), RPI) + ifeq ($(TARGET_ARCH), armv7) + CXXFLAGS += \ + -march=armv7-a \ + -mfpu=neon-vfpv4 \ + -funsafe-math-optimizations \ + -ftree-vectorize + + CCFLAGS += \ + -march=armv7-a \ + -mfpu=neon-vfpv4 \ + -funsafe-math-optimizations \ + -ftree-vectorize + + LDFLAGS := \ + -Wl,--no-export-dynamic \ + -Wl,--exclude-libs,ALL \ + -Wl,--gc-sections \ + -Wl,--as-needed + endif + + LIBS := \ + -lstdc++ \ + -lpthread \ + -lm \ + -ldl + + OBJDIR := $(OBJDIR)rpi_$(TARGET_ARCH)/ + LIBDIR := $(LIBDIR)rpi_$(TARGET_ARCH)/ + BINDIR := $(BINDIR)rpi_$(TARGET_ARCH)/ + DEPDIR := $(DEPDIR)rpi_$(TARGET_ARCH)/ +endif diff --git a/tensorflow/contrib/lite/schema/BUILD b/tensorflow/contrib/lite/schema/BUILD index 54167ddd9a5a003d0ff21e6627a1dbe94afa3e87..a758c5e7e1434447e83c993034d6496bff8b2a48 100644 --- a/tensorflow/contrib/lite/schema/BUILD +++ b/tensorflow/contrib/lite/schema/BUILD @@ -1,10 +1,13 @@ -package(default_visibility = [ - "//visibility:public", -]) +package( + default_visibility = [ + "//visibility:public", + ], +) licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "py_test") +load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") py_binary( name = "upgrade_schema", @@ -80,3 +83,5 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) + +tflite_portable_test_suite() 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..08bcfe451685f488be2c3bc180f2dfc43dfe4f05 --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/generator.cc @@ -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. +==============================================================================*/ +#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 + +// The enum for builtin operators. +// Note: CUSTOM and DELEGATE are 2 special ops which are not real biultin +// ops. +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 c0b220e872176af46035f2ace8defca17e59ea1d..04387fed33fb8d99cd862a5c5fe0c8465f84337e 100644 --- a/tensorflow/contrib/lite/schema/schema.fbs +++ b/tensorflow/contrib/lite/schema/schema.fbs @@ -75,7 +75,7 @@ enum BuiltinOperator : byte { CONV_2D = 3, DEPTHWISE_CONV_2D = 4, // DEPTH_TO_SPACE = 5, - // DEQUANTIZE = 6, + DEQUANTIZE = 6, EMBEDDING_LOOKUP = 7, // FLOOR = 8, FULLY_CONNECTED = 9, @@ -120,6 +120,16 @@ enum BuiltinOperator : byte { UNIDIRECTIONAL_SEQUENCE_LSTM = 44, STRIDED_SLICE = 45, BIDIRECTIONAL_SEQUENCE_RNN = 46, + EXP = 47, + TOPK_V2 = 48, + SPLIT = 49, + LOG_SOFTMAX = 50, + // DELEGATE is a special op type for the operations which are delegated to + // other backends. + // WARNING: Experimental interface, subject to change + DELEGATE = 51, + BIDIRECTIONAL_SEQUENCE_LSTM = 52, + CAST = 53, } // Options for the builtin operators. @@ -156,6 +166,12 @@ union BuiltinOptions { SqueezeOptions, SequenceRNNOptions, StridedSliceOptions, + ExpOptions, + TopKV2Options, + SplitOptions, + LogSoftmaxOptions, + CastOptions, + DequantizeOptions, } enum Padding : byte { SAME, VALID } @@ -273,6 +289,9 @@ table LSTMOptions { } table ResizeBilinearOptions { + new_height: int (deprecated); + new_width: int (deprecated); + align_corners: bool; } // A call operation options @@ -312,6 +331,9 @@ table DivOptions { fused_activation_function:ActivationFunctionType; } +table TopKV2Options { +} + enum CombinerType : byte { SUM = 0, MEAN = 1, @@ -329,6 +351,9 @@ table GatherOptions { table TransposeOptions { } +table ExpOptions { +} + table MeanOptions { keep_dims: bool; } @@ -337,6 +362,10 @@ table SqueezeOptions { squeeze_dims:[int]; } +table SplitOptions { + num_splits: int; +} + table StridedSliceOptions { begin_mask: int; end_mask: int; @@ -345,6 +374,15 @@ table StridedSliceOptions { shrink_axis_mask: int; } +table LogSoftmaxOptions { +} + +table CastOptions { +} + +table DequantizeOptions { +} + // 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 29f3a17be7aef71fbd19fdf2dd390421b724b53a..b922de20816bafc6e01c4bb66ddeef839be2bbd8 100755 --- a/tensorflow/contrib/lite/schema/schema_generated.h +++ b/tensorflow/contrib/lite/schema/schema_generated.h @@ -14,6 +14,7 @@ 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_ @@ -108,6 +109,9 @@ struct SubOptionsT; struct DivOptions; struct DivOptionsT; +struct TopKV2Options; +struct TopKV2OptionsT; + struct EmbeddingLookupSparseOptions; struct EmbeddingLookupSparseOptionsT; @@ -117,15 +121,30 @@ 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 LogSoftmaxOptions; +struct LogSoftmaxOptionsT; + +struct CastOptions; +struct CastOptionsT; + +struct DequantizeOptions; +struct DequantizeOptionsT; + struct OperatorCode; struct OperatorCodeT; @@ -153,15 +172,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; } @@ -176,6 +207,7 @@ enum BuiltinOperator { BuiltinOperator_CONCATENATION = 2, BuiltinOperator_CONV_2D = 3, BuiltinOperator_DEPTHWISE_CONV_2D = 4, + BuiltinOperator_DEQUANTIZE = 6, BuiltinOperator_EMBEDDING_LOOKUP = 7, BuiltinOperator_FULLY_CONNECTED = 9, BuiltinOperator_HASHTABLE_LOOKUP = 10, @@ -215,108 +247,133 @@ enum BuiltinOperator { 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_LOG_SOFTMAX = 50, + BuiltinOperator_DELEGATE = 51, + BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM = 52, + BuiltinOperator_CAST = 53, BuiltinOperator_MIN = BuiltinOperator_ADD, - BuiltinOperator_MAX = BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN + BuiltinOperator_MAX = BuiltinOperator_CAST }; -inline BuiltinOperator (&EnumValuesBuiltinOperator())[44] { +inline BuiltinOperator (&EnumValuesBuiltinOperator())[52] { 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_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_ADD, + BuiltinOperator_AVERAGE_POOL_2D, + BuiltinOperator_CONCATENATION, + BuiltinOperator_CONV_2D, + BuiltinOperator_DEPTHWISE_CONV_2D, + BuiltinOperator_DEQUANTIZE, + 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, + BuiltinOperator_LOG_SOFTMAX, + BuiltinOperator_DELEGATE, + BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, + BuiltinOperator_CAST + }; 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", - "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", - nullptr}; + static const char *names[] = { + "ADD", + "AVERAGE_POOL_2D", + "CONCATENATION", + "CONV_2D", + "DEPTHWISE_CONV_2D", + "", + "DEQUANTIZE", + "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", + "LOG_SOFTMAX", + "DELEGATE", + "BIDIRECTIONAL_SEQUENCE_LSTM", + "CAST", + nullptr + }; return names; } @@ -359,83 +416,104 @@ enum BuiltinOptions { BuiltinOptions_SqueezeOptions = 30, BuiltinOptions_SequenceRNNOptions = 31, BuiltinOptions_StridedSliceOptions = 32, + BuiltinOptions_ExpOptions = 33, + BuiltinOptions_TopKV2Options = 34, + BuiltinOptions_SplitOptions = 35, + BuiltinOptions_LogSoftmaxOptions = 36, + BuiltinOptions_CastOptions = 37, + BuiltinOptions_DequantizeOptions = 38, BuiltinOptions_MIN = BuiltinOptions_NONE, - BuiltinOptions_MAX = BuiltinOptions_StridedSliceOptions + BuiltinOptions_MAX = BuiltinOptions_DequantizeOptions }; -inline BuiltinOptions (&EnumValuesBuiltinOptions())[33] { +inline BuiltinOptions (&EnumValuesBuiltinOptions())[39] { 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_SpaceToBatchNDOptions, - BuiltinOptions_TransposeOptions, - BuiltinOptions_MeanOptions, - BuiltinOptions_SubOptions, - BuiltinOptions_DivOptions, - BuiltinOptions_SqueezeOptions, - BuiltinOptions_SequenceRNNOptions, - BuiltinOptions_StridedSliceOptions}; + 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, + BuiltinOptions_LogSoftmaxOptions, + BuiltinOptions_CastOptions, + BuiltinOptions_DequantizeOptions + }; 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", - "SpaceToBatchNDOptions", - "TransposeOptions", - "MeanOptions", - "SubOptions", - "DivOptions", - "SqueezeOptions", - "SequenceRNNOptions", - "StridedSliceOptions", - 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", + "LogSoftmaxOptions", + "CastOptions", + "DequantizeOptions", + nullptr + }; return names; } @@ -444,206 +522,182 @@ 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 { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SpaceToBatchNDOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_TransposeOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_MeanOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SubOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_DivOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SqueezeOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SequenceRNNOptions; }; -template <> -struct BuiltinOptionsTraits { +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; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_LogSoftmaxOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_CastOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_DequantizeOptions; +}; + 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) { @@ -652,342 +706,317 @@ 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; + return type == BuiltinOptions_SpaceToBatchNDOptions ? + reinterpret_cast(value) : nullptr; } const SpaceToBatchNDOptionsT *AsSpaceToBatchNDOptions() const { - return type == BuiltinOptions_SpaceToBatchNDOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SpaceToBatchNDOptions ? + reinterpret_cast(value) : nullptr; } TransposeOptionsT *AsTransposeOptions() { - return type == BuiltinOptions_TransposeOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_TransposeOptions ? + reinterpret_cast(value) : nullptr; } const TransposeOptionsT *AsTransposeOptions() const { - return type == BuiltinOptions_TransposeOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_TransposeOptions ? + reinterpret_cast(value) : nullptr; } MeanOptionsT *AsMeanOptions() { - return type == BuiltinOptions_MeanOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_MeanOptions ? + reinterpret_cast(value) : nullptr; } const MeanOptionsT *AsMeanOptions() const { - return type == BuiltinOptions_MeanOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_MeanOptions ? + reinterpret_cast(value) : nullptr; } SubOptionsT *AsSubOptions() { - return type == BuiltinOptions_SubOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SubOptions ? + reinterpret_cast(value) : nullptr; } const SubOptionsT *AsSubOptions() const { - return type == BuiltinOptions_SubOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SubOptions ? + reinterpret_cast(value) : nullptr; } DivOptionsT *AsDivOptions() { - return type == BuiltinOptions_DivOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_DivOptions ? + reinterpret_cast(value) : nullptr; } const DivOptionsT *AsDivOptions() const { - return type == BuiltinOptions_DivOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_DivOptions ? + reinterpret_cast(value) : nullptr; } SqueezeOptionsT *AsSqueezeOptions() { - return type == BuiltinOptions_SqueezeOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SqueezeOptions ? + reinterpret_cast(value) : nullptr; } const SqueezeOptionsT *AsSqueezeOptions() const { - return type == BuiltinOptions_SqueezeOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SqueezeOptions ? + reinterpret_cast(value) : nullptr; } SequenceRNNOptionsT *AsSequenceRNNOptions() { - return type == BuiltinOptions_SequenceRNNOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SequenceRNNOptions ? + reinterpret_cast(value) : nullptr; } const SequenceRNNOptionsT *AsSequenceRNNOptions() const { - return type == BuiltinOptions_SequenceRNNOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SequenceRNNOptions ? + reinterpret_cast(value) : nullptr; } StridedSliceOptionsT *AsStridedSliceOptions() { - return type == BuiltinOptions_StridedSliceOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_StridedSliceOptions ? + reinterpret_cast(value) : nullptr; } const StridedSliceOptionsT *AsStridedSliceOptions() const { - return type == BuiltinOptions_StridedSliceOptions - ? reinterpret_cast(value) - : nullptr; + 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; + } + LogSoftmaxOptionsT *AsLogSoftmaxOptions() { + return type == BuiltinOptions_LogSoftmaxOptions ? + reinterpret_cast(value) : nullptr; + } + const LogSoftmaxOptionsT *AsLogSoftmaxOptions() const { + return type == BuiltinOptions_LogSoftmaxOptions ? + reinterpret_cast(value) : nullptr; + } + CastOptionsT *AsCastOptions() { + return type == BuiltinOptions_CastOptions ? + reinterpret_cast(value) : nullptr; + } + const CastOptionsT *AsCastOptions() const { + return type == BuiltinOptions_CastOptions ? + reinterpret_cast(value) : nullptr; + } + DequantizeOptionsT *AsDequantizeOptions() { + return type == BuiltinOptions_DequantizeOptions ? + reinterpret_cast(value) : nullptr; + } + const DequantizeOptionsT *AsDequantizeOptions() const { + return type == BuiltinOptions_DequantizeOptions ? + 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, @@ -997,12 +1026,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; } @@ -1024,15 +1060,26 @@ enum ActivationFunctionType { inline ActivationFunctionType (&EnumValuesActivationFunctionType())[6] { static ActivationFunctionType values[] = { - ActivationFunctionType_NONE, ActivationFunctionType_RELU, - ActivationFunctionType_RELU_N1_TO_1, 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", "RELU_N1_TO_1", "RELU6", - "TANH", "SIGN_BIT", nullptr}; + static const char *names[] = { + "NONE", + "RELU", + "RELU_N1_TO_1", + "RELU6", + "TANH", + "SIGN_BIT", + nullptr + }; return names; } @@ -1050,14 +1097,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; } @@ -1075,13 +1129,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; } @@ -1097,12 +1159,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; } @@ -1117,13 +1184,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); } @@ -1137,20 +1209,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 { @@ -1165,16 +1237,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); @@ -1196,23 +1266,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; @@ -1221,7 +1289,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 { @@ -1239,7 +1310,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); } @@ -1247,20 +1320,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 { @@ -1278,11 +1351,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 &); @@ -1296,7 +1369,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); @@ -1311,17 +1385,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; @@ -1333,7 +1410,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 { @@ -1347,11 +1425,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) && @@ -1361,22 +1442,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); @@ -1384,13 +1459,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 &); @@ -1402,10 +1475,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); @@ -1414,9 +1488,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; @@ -1432,7 +1504,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 { @@ -1448,15 +1521,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) && @@ -1468,22 +1546,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); @@ -1497,13 +1569,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 &); @@ -1515,11 +1585,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); @@ -1530,9 +1602,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; @@ -1546,11 +1616,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, @@ -1562,14 +1632,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) && @@ -1580,22 +1653,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); @@ -1604,21 +1671,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); @@ -1627,10 +1689,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); @@ -1640,34 +1704,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) && @@ -1675,43 +1738,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); @@ -1719,13 +1770,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); @@ -1733,61 +1782,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 &); @@ -1806,25 +1848,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) && @@ -1832,14 +1878,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 { @@ -1848,13 +1889,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 &); @@ -1866,57 +1905,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 &); @@ -1929,16 +1962,13 @@ 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; @@ -1946,16 +1976,21 @@ struct SequenceRNNOptionsT : public flatbuffers::NativeTable { ActivationFunctionType fused_activation_function; SequenceRNNOptionsT() : time_major(false), - fused_activation_function(ActivationFunctionType_NONE) {} + 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; } + 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)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1963,30 +1998,22 @@ struct SequenceRNNOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { 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); + 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); + 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); + 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) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SequenceRNNOptionsBuilder &operator=(const SequenceRNNOptionsBuilder &); @@ -1998,18 +2025,16 @@ struct SequenceRNNOptionsBuilder { }; inline flatbuffers::Offset CreateSequenceRNNOptions( - flatbuffers::FlatBufferBuilder &_fbb, bool time_major = false, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + 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); +flatbuffers::Offset CreateSequenceRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const SequenceRNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct BidirectionalSequenceRNNOptionsT : public flatbuffers::NativeTable { typedef BidirectionalSequenceRNNOptions TableType; @@ -2017,17 +2042,21 @@ struct BidirectionalSequenceRNNOptionsT : public flatbuffers::NativeTable { ActivationFunctionType fused_activation_function; BidirectionalSequenceRNNOptionsT() : time_major(false), - fused_activation_function(ActivationFunctionType_NONE) {} + fused_activation_function(ActivationFunctionType_NONE) { + } }; -struct BidirectionalSequenceRNNOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +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; } + 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)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -2035,37 +2064,25 @@ struct BidirectionalSequenceRNNOptions FLATBUFFERS_FINAL_CLASS 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); + 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) { + 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 &); + BidirectionalSequenceRNNOptionsBuilder &operator=(const BidirectionalSequenceRNNOptionsBuilder &); flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); auto o = flatbuffers::Offset(end); @@ -2073,63 +2090,52 @@ struct BidirectionalSequenceRNNOptionsBuilder { } }; -inline flatbuffers::Offset -CreateBidirectionalSequenceRNNOptions( - flatbuffers::FlatBufferBuilder &_fbb, bool time_major = false, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { +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); +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 &); @@ -2142,39 +2148,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 { @@ -2184,7 +2189,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 &); @@ -2196,32 +2201,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) && @@ -2229,14 +2238,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 { @@ -2245,13 +2249,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 &); @@ -2263,57 +2265,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 &); @@ -2326,55 +2322,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 &); @@ -2387,55 +2376,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 &); @@ -2448,16 +2430,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; @@ -2466,61 +2445,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); @@ -2528,10 +2512,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); @@ -2540,11 +2526,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; @@ -2554,41 +2536,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); @@ -2597,7 +2581,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 &); @@ -2610,9 +2594,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); @@ -2620,36 +2604,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; - ResizeBilinearOptionsT() {} + 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_ALIGN_CORNERS = 8 + }; + bool align_corners() const { + return GetField(VT_ALIGN_CORNERS, 0) != 0; + } bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && verifier.EndTable(); + return VerifyTableStart(verifier) && + 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_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 &); @@ -2661,37 +2651,39 @@ struct ResizeBilinearOptionsBuilder { }; inline flatbuffers::Offset CreateResizeBilinearOptions( - flatbuffers::FlatBufferBuilder &_fbb) { + flatbuffers::FlatBufferBuilder &_fbb, + bool align_corners = false) { ResizeBilinearOptionsBuilder builder_(_fbb); + 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 { @@ -2701,7 +2693,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 &); @@ -2713,41 +2705,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; - PadOptionsT() {} + PadOptionsT() { + } }; struct PadOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef PadOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && verifier.EndTable(); + return VerifyTableStart(verifier) && + 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_; explicit PadOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } PadOptionsBuilder &operator=(const PadOptionsBuilder &); @@ -2764,45 +2752,42 @@ inline flatbuffers::Offset CreatePadOptions( return builder_.Finish(); } -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 &); @@ -2825,39 +2810,34 @@ 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 SpaceToBatchNDOptionsT : public flatbuffers::NativeTable { typedef SpaceToBatchNDOptions TableType; - SpaceToBatchNDOptionsT() {} + SpaceToBatchNDOptionsT() { + } }; -struct SpaceToBatchNDOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct SpaceToBatchNDOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SpaceToBatchNDOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && verifier.EndTable(); + 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); + 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) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SpaceToBatchNDOptionsBuilder &operator=(const SpaceToBatchNDOptionsBuilder &); @@ -2874,36 +2854,30 @@ inline flatbuffers::Offset CreateSpaceToBatchNDOptions( return builder_.Finish(); } -flatbuffers::Offset CreateSpaceToBatchNDOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +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() {} + BatchToSpaceNDOptionsT() { + } }; -struct BatchToSpaceNDOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct BatchToSpaceNDOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef BatchToSpaceNDOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && verifier.EndTable(); + return VerifyTableStart(verifier) && + 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); + 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_; explicit BatchToSpaceNDOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } BatchToSpaceNDOptionsBuilder &operator=(const BatchToSpaceNDOptionsBuilder &); @@ -2920,9 +2894,7 @@ inline flatbuffers::Offset CreateBatchToSpaceNDOptions( return builder_.Finish(); } -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; @@ -2930,13 +2902,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); } @@ -2950,14 +2931,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 { @@ -2967,15 +2943,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 &); @@ -2987,8 +2961,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); @@ -2996,33 +2972,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 { @@ -3032,7 +3007,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 &); @@ -3044,54 +3019,49 @@ 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) {} + SubOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } }; struct SubOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SubOptionsT 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(); } - 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); + 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); + 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) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SubOptionsBuilder &operator=(const SubOptionsBuilder &); @@ -3104,55 +3074,48 @@ struct SubOptionsBuilder { inline flatbuffers::Offset CreateSubOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + 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); +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) {} + DivOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } }; struct DivOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef DivOptionsT 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(); } - 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); + 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); + 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) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } DivOptionsBuilder &operator=(const DivOptionsBuilder &); @@ -3165,59 +3128,91 @@ struct DivOptionsBuilder { inline flatbuffers::Offset CreateDivOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + 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); +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); @@ -3225,42 +3220,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 { @@ -3270,7 +3263,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 &); @@ -3282,41 +3275,37 @@ 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 TransposeOptionsT : public flatbuffers::NativeTable { typedef TransposeOptions TableType; - TransposeOptionsT() {} + TransposeOptionsT() { + } }; struct TransposeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef TransposeOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && verifier.EndTable(); + return VerifyTableStart(verifier) && + verifier.EndTable(); } - 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); + 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 TransposeOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; explicit TransposeOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } TransposeOptionsBuilder &operator=(const TransposeOptionsBuilder &); @@ -3333,100 +3322,136 @@ inline flatbuffers::Offset CreateTransposeOptions( return builder_.Finish(); } -flatbuffers::Offset CreateTransposeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateTransposeOptions(flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); -struct MeanOptionsT : public flatbuffers::NativeTable { - typedef MeanOptions TableType; - bool keep_dims; - MeanOptionsT() : keep_dims(false) {} +struct ExpOptionsT : public flatbuffers::NativeTable { + typedef ExpOptions TableType; + ExpOptionsT() { + } }; -struct MeanOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { - typedef MeanOptionsT NativeTableType; - enum { VT_KEEP_DIMS = 4 }; - bool keep_dims() const { return GetField(VT_KEEP_DIMS, 0) != 0; } +struct ExpOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef ExpOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_KEEP_DIMS) && verifier.EndTable(); + verifier.EndTable(); } - 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); + 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 MeanOptionsBuilder { +struct ExpOptionsBuilder { 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); - } - explicit MeanOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + explicit ExpOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - MeanOptionsBuilder &operator=(const MeanOptionsBuilder &); - flatbuffers::Offset Finish() { + ExpOptionsBuilder &operator=(const ExpOptionsBuilder &); + flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); - auto o = flatbuffers::Offset(end); + auto o = flatbuffers::Offset(end); return o; } }; -inline flatbuffers::Offset CreateMeanOptions( - flatbuffers::FlatBufferBuilder &_fbb, bool keep_dims = false) { - MeanOptionsBuilder builder_(_fbb); - builder_.add_keep_dims(keep_dims); +inline flatbuffers::Offset CreateExpOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + ExpOptionsBuilder builder_(_fbb); return builder_.Finish(); } -flatbuffers::Offset CreateMeanOptions( - flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT *_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 SqueezeOptionsT : public flatbuffers::NativeTable { - typedef SqueezeOptions TableType; - std::vector squeeze_dims; - SqueezeOptionsT() {} +struct MeanOptionsT : public flatbuffers::NativeTable { + typedef MeanOptions TableType; + bool keep_dims; + MeanOptionsT() + : keep_dims(false) { + } }; -struct SqueezeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { +struct MeanOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef MeanOptionsT NativeTableType; + enum { + VT_KEEP_DIMS = 4 + }; + bool keep_dims() const { + return GetField(VT_KEEP_DIMS, 0) != 0; + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_KEEP_DIMS) && + verifier.EndTable(); + } + 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); + } + explicit MeanOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + MeanOptionsBuilder &operator=(const MeanOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +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() { + } +}; + +struct SqueezeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SqueezeOptionsT NativeTableType; - enum { VT_SQUEEZE_DIMS = 4 }; + 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(); + 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); + 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) { + void add_squeeze_dims(flatbuffers::Offset> squeeze_dims) { fbb_.AddOffset(SqueezeOptions::VT_SQUEEZE_DIMS, squeeze_dims); } explicit SqueezeOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SqueezeOptionsBuilder &operator=(const SqueezeOptionsBuilder &); @@ -3449,12 +3474,65 @@ inline flatbuffers::Offset CreateSqueezeOptionsDirect( flatbuffers::FlatBufferBuilder &_fbb, const std::vector *squeeze_dims = nullptr) { return tflite::CreateSqueezeOptions( - _fbb, squeeze_dims ? _fbb.CreateVector(*squeeze_dims) : 0); + _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 CreateSqueezeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateSplitOptions(flatbuffers::FlatBufferBuilder &_fbb, const SplitOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct StridedSliceOptionsT : public flatbuffers::NativeTable { typedef StridedSliceOptions TableType; @@ -3468,11 +3546,11 @@ struct StridedSliceOptionsT : public flatbuffers::NativeTable { end_mask(0), ellipsis_mask(0), new_axis_mask(0), - shrink_axis_mask(0) {} + shrink_axis_mask(0) { + } }; -struct StridedSliceOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct StridedSliceOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef StridedSliceOptionsT NativeTableType; enum { VT_BEGIN_MASK = 4, @@ -3481,8 +3559,12 @@ struct StridedSliceOptions FLATBUFFERS_FINAL_CLASS 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 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); } @@ -3501,14 +3583,9 @@ struct StridedSliceOptions FLATBUFFERS_FINAL_CLASS 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); + 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 { @@ -3521,19 +3598,16 @@ struct StridedSliceOptionsBuilder { 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); + 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); + 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); + fbb_.AddElement(StridedSliceOptions::VT_SHRINK_AXIS_MASK, shrink_axis_mask, 0); } explicit StridedSliceOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } StridedSliceOptionsBuilder &operator=(const StridedSliceOptionsBuilder &); @@ -3545,8 +3619,11 @@ struct StridedSliceOptionsBuilder { }; 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, + 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); @@ -3557,20 +3634,143 @@ inline flatbuffers::Offset CreateStridedSliceOptions( return builder_.Finish(); } -flatbuffers::Offset CreateStridedSliceOptions( - flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateStridedSliceOptions(flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct LogSoftmaxOptionsT : public flatbuffers::NativeTable { + typedef LogSoftmaxOptions TableType; + LogSoftmaxOptionsT() { + } +}; + +struct LogSoftmaxOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef LogSoftmaxOptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + LogSoftmaxOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(LogSoftmaxOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const LogSoftmaxOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct LogSoftmaxOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit LogSoftmaxOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + LogSoftmaxOptionsBuilder &operator=(const LogSoftmaxOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateLogSoftmaxOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + LogSoftmaxOptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateLogSoftmaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const LogSoftmaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct CastOptionsT : public flatbuffers::NativeTable { + typedef CastOptions TableType; + CastOptionsT() { + } +}; + +struct CastOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef CastOptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + CastOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(CastOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const CastOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct CastOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit CastOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + CastOptionsBuilder &operator=(const CastOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateCastOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + CastOptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateCastOptions(flatbuffers::FlatBufferBuilder &_fbb, const CastOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct DequantizeOptionsT : public flatbuffers::NativeTable { + typedef DequantizeOptions TableType; + DequantizeOptionsT() { + } +}; + +struct DequantizeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef DequantizeOptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + DequantizeOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(DequantizeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const DequantizeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct DequantizeOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit DequantizeOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + DequantizeOptionsBuilder &operator=(const DequantizeOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateDequantizeOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + DequantizeOptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateDequantizeOptions(flatbuffers::FlatBufferBuilder &_fbb, const DequantizeOptionsT *_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) {} + 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 }; + enum { + VT_BUILTIN_CODE = 4, + VT_CUSTOM_CODE = 6 + }; BuiltinOperator builtin_code() const { return static_cast(GetField(VT_BUILTIN_CODE, 0)); } @@ -3581,30 +3781,25 @@ struct OperatorCode FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { return VerifyTableStart(verifier) && VerifyField(verifier, VT_BUILTIN_CODE) && VerifyOffset(verifier, VT_CUSTOM_CODE) && - verifier.Verify(custom_code()) && verifier.EndTable(); + 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); + 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); + 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) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } OperatorCodeBuilder &operator=(const OperatorCodeBuilder &); @@ -3630,12 +3825,12 @@ inline flatbuffers::Offset CreateOperatorCodeDirect( BuiltinOperator builtin_code = BuiltinOperator_ADD, const char *custom_code = nullptr) { return tflite::CreateOperatorCode( - _fbb, builtin_code, custom_code ? _fbb.CreateString(custom_code) : 0); + _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); +flatbuffers::Offset CreateOperatorCode(flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct OperatorT : public flatbuffers::NativeTable { typedef Operator TableType; @@ -3647,7 +3842,8 @@ struct OperatorT : public flatbuffers::NativeTable { CustomOptionsFormat custom_options_format; OperatorT() : opcode_index(0), - custom_options_format(CustomOptionsFormat_FLEXBUFFERS) {} + custom_options_format(CustomOptionsFormat_FLEXBUFFERS) { + } }; struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { @@ -3671,398 +3867,304 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { return GetPointer *>(VT_OUTPUTS); } BuiltinOptions builtin_options_type() const { - return static_cast( - GetField(VT_BUILTIN_OPTIONS_TYPE, 0)); + 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; + 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; + 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; + 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; + 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; - } - 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; + 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; + 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; + return builtin_options_type() == BuiltinOptions_MeanOptions ? static_cast(builtin_options()) : nullptr; } const SubOptions *builtin_options_as_SubOptions() const { - return builtin_options_type() == BuiltinOptions_SubOptions - ? static_cast(builtin_options()) - : nullptr; + 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; + 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; + 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; + 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; + 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 LogSoftmaxOptions *builtin_options_as_LogSoftmaxOptions() const { + return builtin_options_type() == BuiltinOptions_LogSoftmaxOptions ? static_cast(builtin_options()) : nullptr; + } + const CastOptions *builtin_options_as_CastOptions() const { + return builtin_options_type() == BuiltinOptions_CastOptions ? static_cast(builtin_options()) : nullptr; + } + const DequantizeOptions *builtin_options_as_DequantizeOptions() const { + return builtin_options_type() == BuiltinOptions_DequantizeOptions ? 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 { +template<> inline const SpaceToBatchNDOptions *Operator::builtin_options_as() const { return builtin_options_as_SpaceToBatchNDOptions(); } -template <> -inline const TransposeOptions *Operator::builtin_options_as() - const { +template<> inline const TransposeOptions *Operator::builtin_options_as() const { return builtin_options_as_TransposeOptions(); } -template <> -inline const MeanOptions *Operator::builtin_options_as() const { +template<> inline const MeanOptions *Operator::builtin_options_as() const { return builtin_options_as_MeanOptions(); } -template <> -inline const SubOptions *Operator::builtin_options_as() const { +template<> inline const SubOptions *Operator::builtin_options_as() const { return builtin_options_as_SubOptions(); } -template <> -inline const DivOptions *Operator::builtin_options_as() const { +template<> inline const DivOptions *Operator::builtin_options_as() const { return builtin_options_as_DivOptions(); } -template <> -inline const SqueezeOptions *Operator::builtin_options_as() - const { +template<> inline const SqueezeOptions *Operator::builtin_options_as() const { return builtin_options_as_SqueezeOptions(); } -template <> -inline const SequenceRNNOptions * -Operator::builtin_options_as() const { +template<> inline const SequenceRNNOptions *Operator::builtin_options_as() const { return builtin_options_as_SequenceRNNOptions(); } -template <> -inline const StridedSliceOptions * -Operator::builtin_options_as() const { +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(); +} + +template<> inline const LogSoftmaxOptions *Operator::builtin_options_as() const { + return builtin_options_as_LogSoftmaxOptions(); +} + +template<> inline const CastOptions *Operator::builtin_options_as() const { + return builtin_options_as_CastOptions(); +} + +template<> inline const DequantizeOptions *Operator::builtin_options_as() const { + return builtin_options_as_DequantizeOptions(); +} + struct OperatorBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; @@ -4076,21 +4178,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 &); @@ -4102,14 +4202,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); @@ -4122,25 +4222,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; @@ -4149,7 +4250,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 { @@ -4162,8 +4264,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); @@ -4172,41 +4273,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) { @@ -4215,15 +4311,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 &); @@ -4236,12 +4331,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); @@ -4264,38 +4357,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 { @@ -4304,7 +4395,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 &); @@ -4326,13 +4418,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; @@ -4341,7 +4432,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 { @@ -4353,24 +4446,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) && @@ -4383,16 +4472,14 @@ 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 { @@ -4401,26 +4488,20 @@ struct ModelBuilder { void add_version(uint32_t version) { fbb_.AddElement(Model::VT_VERSION, version, 0); } - void add_operator_codes( - flatbuffers::Offset< - flatbuffers::Vector>> - operator_codes) { + void add_operator_codes(flatbuffers::Offset>> operator_codes) { fbb_.AddOffset(Model::VT_OPERATOR_CODES, operator_codes); } - void add_subgraphs( - flatbuffers::Offset>> - subgraphs) { + 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) { + void add_buffers(flatbuffers::Offset>> buffers) { fbb_.AddOffset(Model::VT_BUFFERS, buffers); } - explicit ModelBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { + explicit ModelBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } ModelBuilder &operator=(const ModelBuilder &); @@ -4432,14 +4513,12 @@ struct ModelBuilder { }; inline flatbuffers::Offset CreateModel( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t version = 0, - flatbuffers::Offset>> - operator_codes = 0, - flatbuffers::Offset>> - subgraphs = 0, + flatbuffers::FlatBufferBuilder &_fbb, + uint32_t version = 0, + flatbuffers::Offset>> operator_codes = 0, + flatbuffers::Offset>> subgraphs = 0, flatbuffers::Offset description = 0, - flatbuffers::Offset>> - buffers = 0) { + flatbuffers::Offset>> buffers = 0) { ModelBuilder builder_(_fbb); builder_.add_buffers(buffers); builder_.add_description(description); @@ -4450,2010 +4529,1346 @@ inline flatbuffers::Offset CreateModel( } inline flatbuffers::Offset CreateModelDirect( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t version = 0, - const std::vector> *operator_codes = - nullptr, + 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, + _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); +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 { +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 { +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); - } - } - }; + { 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) { +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) { +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; + 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); + 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 { +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 { +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)); - }; + { 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) { +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) { +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; + 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); + 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 { +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 { +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; - }; + { 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) { +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) { +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; + 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); + return tflite::CreateConv2DOptions( + _fbb, + _padding, + _stride_w, + _stride_h, + _fused_activation_function); } -inline Pool2DOptionsT *Pool2DOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +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 { +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; - }; + { 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) { +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) { +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; + 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); + 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 { +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 { +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; - }; + { 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) { +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) { +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; + 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); + 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 { +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 { +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); - } - } - }; + { 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) { +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) { +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; + 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 _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 { +inline void LSHProjectionOptions::UnPackTo(LSHProjectionOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = type(); - _o->type = _e; - }; + { auto _e = type(); _o->type = _e; }; } -inline flatbuffers::Offset LSHProjectionOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +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) { +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; + 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); + return tflite::CreateLSHProjectionOptions( + _fbb, + _type); } -inline SVDFOptionsT *SVDFOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +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 { +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; - }; + { 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) { +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) { +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; + 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); + return tflite::CreateSVDFOptions( + _fbb, + _rank, + _fused_activation_function); } -inline RNNOptionsT *RNNOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +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 { +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; - }; + { 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) { +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) { +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; + 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); + return tflite::CreateRNNOptions( + _fbb, + _fused_activation_function); } -inline SequenceRNNOptionsT *SequenceRNNOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +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 { +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; - }; + { 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) { +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) { +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; + 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); + return tflite::CreateSequenceRNNOptions( + _fbb, + _time_major, + _fused_activation_function); } -inline BidirectionalSequenceRNNOptionsT * -BidirectionalSequenceRNNOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline BidirectionalSequenceRNNOptionsT *BidirectionalSequenceRNNOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new BidirectionalSequenceRNNOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void BidirectionalSequenceRNNOptions::UnPackTo( - BidirectionalSequenceRNNOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +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; - }; + { auto _e = time_major(); _o->time_major = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset -BidirectionalSequenceRNNOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const BidirectionalSequenceRNNOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset BidirectionalSequenceRNNOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const BidirectionalSequenceRNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { return CreateBidirectionalSequenceRNNOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset -CreateBidirectionalSequenceRNNOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const BidirectionalSequenceRNNOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +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; + 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); + _fbb, + _time_major, + _fused_activation_function); } -inline FullyConnectedOptionsT *FullyConnectedOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline FullyConnectedOptionsT *FullyConnectedOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new FullyConnectedOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void FullyConnectedOptions::UnPackTo( - FullyConnectedOptionsT *_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 = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset FullyConnectedOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateFullyConnectedOptions( - flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_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 FullyConnectedOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateFullyConnectedOptions( + _fbb, + _fused_activation_function); } -inline SoftmaxOptionsT *SoftmaxOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline SoftmaxOptionsT *SoftmaxOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new SoftmaxOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SoftmaxOptions::UnPackTo( - SoftmaxOptionsT *_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 = beta(); - _o->beta = _e; - }; + { auto _e = beta(); _o->beta = _e; }; } -inline flatbuffers::Offset SoftmaxOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateSoftmaxOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_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 SoftmaxOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateSoftmaxOptions( + _fbb, + _beta); } -inline ConcatenationOptionsT *ConcatenationOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline ConcatenationOptionsT *ConcatenationOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new ConcatenationOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ConcatenationOptions::UnPackTo( - ConcatenationOptionsT *_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 = axis(); - _o->axis = _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 ConcatenationOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateConcatenationOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_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 ConcatenationOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateConcatenationOptions( + _fbb, + _axis, + _fused_activation_function); } -inline AddOptionsT *AddOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline AddOptionsT *AddOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new AddOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void AddOptions::UnPackTo( - AddOptionsT *_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 = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { 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) { +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 CreateAddOptions( - flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_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 AddOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateAddOptions( + _fbb, + _fused_activation_function); } -inline MulOptionsT *MulOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline MulOptionsT *MulOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new MulOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void MulOptions::UnPackTo( - MulOptionsT *_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 = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { 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) { +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 CreateMulOptions( - flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_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 MulOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateMulOptions( + _fbb, + _fused_activation_function); } -inline L2NormOptionsT *L2NormOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline L2NormOptionsT *L2NormOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new L2NormOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void L2NormOptions::UnPackTo( - L2NormOptionsT *_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 = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset L2NormOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateL2NormOptions( - flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_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 L2NormOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateL2NormOptions( + _fbb, + _fused_activation_function); } -inline LocalResponseNormalizationOptionsT * -LocalResponseNormalizationOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline LocalResponseNormalizationOptionsT *LocalResponseNormalizationOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new LocalResponseNormalizationOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void LocalResponseNormalizationOptions::UnPackTo( - LocalResponseNormalizationOptionsT *_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 = radius(); - _o->radius = _e; - }; - { - auto _e = bias(); - _o->bias = _e; - }; - { - auto _e = alpha(); - _o->alpha = _e; - }; - { - auto _e = beta(); - _o->beta = _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 -LocalResponseNormalizationOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const LocalResponseNormalizationOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 -CreateLocalResponseNormalizationOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const LocalResponseNormalizationOptionsT *_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 LocalResponseNormalizationOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateLocalResponseNormalizationOptions( + _fbb, + _radius, + _bias, + _alpha, + _beta); } -inline LSTMOptionsT *LSTMOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline LSTMOptionsT *LSTMOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new LSTMOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void LSTMOptions::UnPackTo( - LSTMOptionsT *_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 = 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; }; + { auto _e = cell_clip(); _o->cell_clip = _e; }; + { auto _e = proj_clip(); _o->proj_clip = _e; }; } -inline flatbuffers::Offset LSTMOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateLSTMOptions( - flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_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 LSTMOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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; 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::CreateLSTMOptions( + _fbb, + _fused_activation_function, + _cell_clip, + _proj_clip); } -inline ResizeBilinearOptionsT *ResizeBilinearOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline ResizeBilinearOptionsT *ResizeBilinearOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new ResizeBilinearOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ResizeBilinearOptions::UnPackTo( - ResizeBilinearOptionsT *_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 = align_corners(); _o->align_corners = _e; }; } -inline flatbuffers::Offset ResizeBilinearOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateResizeBilinearOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_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 ResizeBilinearOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - return tflite::CreateResizeBilinearOptions(_fbb); + 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 CallOptionsT *CallOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline CallOptionsT *CallOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new CallOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void CallOptions::UnPackTo( - CallOptionsT *_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 = subgraph(); - _o->subgraph = _e; - }; + { auto _e = subgraph(); _o->subgraph = _e; }; } -inline flatbuffers::Offset CallOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateCallOptions( - flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_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 CallOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateCallOptions( + _fbb, + _subgraph); } -inline PadOptionsT *PadOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline PadOptionsT *PadOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new PadOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void PadOptions::UnPackTo( - PadOptionsT *_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; } -inline flatbuffers::Offset PadOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreatePadOptions( - flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_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 PadOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - return tflite::CreatePadOptions(_fbb); + 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 ReshapeOptionsT *ReshapeOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline ReshapeOptionsT *ReshapeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new ReshapeOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ReshapeOptions::UnPackTo( - ReshapeOptionsT *_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 = 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); - } - } - }; + { 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) { +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 CreateReshapeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_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 ReshapeOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateReshapeOptions( + _fbb, + _new_shape); } -inline SpaceToBatchNDOptionsT *SpaceToBatchNDOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline SpaceToBatchNDOptionsT *SpaceToBatchNDOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new SpaceToBatchNDOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SpaceToBatchNDOptions::UnPackTo( - SpaceToBatchNDOptionsT *_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; } -inline flatbuffers::Offset SpaceToBatchNDOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateSpaceToBatchNDOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT *_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 SpaceToBatchNDOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - return tflite::CreateSpaceToBatchNDOptions(_fbb); + 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 BatchToSpaceNDOptionsT *BatchToSpaceNDOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline BatchToSpaceNDOptionsT *BatchToSpaceNDOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new BatchToSpaceNDOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void BatchToSpaceNDOptions::UnPackTo( - BatchToSpaceNDOptionsT *_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; } -inline flatbuffers::Offset BatchToSpaceNDOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateBatchToSpaceNDOptions( - flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_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 BatchToSpaceNDOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - return tflite::CreateBatchToSpaceNDOptions(_fbb); + 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 SkipGramOptionsT *SkipGramOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline SkipGramOptionsT *SkipGramOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new SkipGramOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SkipGramOptions::UnPackTo( - SkipGramOptionsT *_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 = 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 = 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 SkipGramOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateSkipGramOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_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 SkipGramOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateSkipGramOptions( + _fbb, + _ngram_size, + _max_skip_size, + _include_all_ngrams); } -inline SpaceToDepthOptionsT *SpaceToDepthOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline SpaceToDepthOptionsT *SpaceToDepthOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new SpaceToDepthOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SpaceToDepthOptions::UnPackTo( - SpaceToDepthOptionsT *_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 = block_size(); - _o->block_size = _e; - }; + { auto _e = block_size(); _o->block_size = _e; }; } -inline flatbuffers::Offset SpaceToDepthOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, - const flatbuffers::rehasher_function_t *_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) { +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; + 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); + return tflite::CreateSpaceToDepthOptions( + _fbb, + _block_size); } -inline SubOptionsT *SubOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +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 { +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; - }; + { 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) { +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 CreateSubOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SubOptionsT *_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 SubOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateSubOptions( + _fbb, + _fused_activation_function); } -inline DivOptionsT *DivOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline DivOptionsT *DivOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new DivOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void DivOptions::UnPackTo( - DivOptionsT *_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 = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset DivOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const DivOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateDivOptions( - flatbuffers::FlatBufferBuilder &_fbb, const DivOptionsT *_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 DivOptionsT *__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; - return tflite::CreateDivOptions(_fbb, _fused_activation_function); + return tflite::CreateDivOptions( + _fbb, + _fused_activation_function); } -inline EmbeddingLookupSparseOptionsT *EmbeddingLookupSparseOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline TopKV2OptionsT *TopKV2Options::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new TopKV2OptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void TopKV2Options::UnPackTo(TopKV2OptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; +} + +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 CreateTopKV2Options(flatbuffers::FlatBufferBuilder &_fbb, const TopKV2OptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + 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 EmbeddingLookupSparseOptions::UnPackTo( - EmbeddingLookupSparseOptionsT *_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 = combiner(); - _o->combiner = _e; - }; + { auto _e = combiner(); _o->combiner = _e; }; } -inline flatbuffers::Offset -EmbeddingLookupSparseOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 -CreateEmbeddingLookupSparseOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const EmbeddingLookupSparseOptionsT *_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 EmbeddingLookupSparseOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateEmbeddingLookupSparseOptions( + _fbb, + _combiner); } -inline GatherOptionsT *GatherOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline GatherOptionsT *GatherOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new GatherOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void GatherOptions::UnPackTo( - GatherOptionsT *_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 = axis(); - _o->axis = _e; - }; + { auto _e = axis(); _o->axis = _e; }; } -inline flatbuffers::Offset GatherOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateGatherOptions( - flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_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 GatherOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateGatherOptions( + _fbb, + _axis); } -inline TransposeOptionsT *TransposeOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline TransposeOptionsT *TransposeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new TransposeOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void TransposeOptions::UnPackTo( - TransposeOptionsT *_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; } -inline flatbuffers::Offset TransposeOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateTransposeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT *_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 TransposeOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateTransposeOptions( + _fbb); +} + +inline ExpOptionsT *ExpOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ExpOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void ExpOptions::UnPackTo(ExpOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; +} + +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 CreateExpOptions(flatbuffers::FlatBufferBuilder &_fbb, const ExpOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const TransposeOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - return tflite::CreateTransposeOptions(_fbb); + 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 { +inline MeanOptionsT *MeanOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new MeanOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void MeanOptions::UnPackTo( - MeanOptionsT *_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 = keep_dims(); - _o->keep_dims = _e; - }; + { auto _e = keep_dims(); _o->keep_dims = _e; }; } -inline flatbuffers::Offset MeanOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateMeanOptions( - flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT *_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 MeanOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateMeanOptions( + _fbb, + _keep_dims); } -inline SqueezeOptionsT *SqueezeOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline SqueezeOptionsT *SqueezeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new SqueezeOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SqueezeOptions::UnPackTo( - SqueezeOptionsT *_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 = 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); - } - } - }; + { 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 SqueezeOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateSqueezeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT *_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 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 SplitOptionsT *SplitOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SplitOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void SplitOptions::UnPackTo(SplitOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = num_splits(); _o->num_splits = _e; }; +} + +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 CreateSplitOptions(flatbuffers::FlatBufferBuilder &_fbb, const SplitOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - 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 StridedSliceOptionsT *StridedSliceOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { + 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 StridedSliceOptionsT *StridedSliceOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new StridedSliceOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void StridedSliceOptions::UnPackTo( - StridedSliceOptionsT *_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 = 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; - }; + { 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 StridedSliceOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT *_o, - const flatbuffers::rehasher_function_t *_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 CreateStridedSliceOptions( - flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT *_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 StridedSliceOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + 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); + return tflite::CreateStridedSliceOptions( + _fbb, + _begin_mask, + _end_mask, + _ellipsis_mask, + _new_axis_mask, + _shrink_axis_mask); +} + +inline LogSoftmaxOptionsT *LogSoftmaxOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new LogSoftmaxOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void LogSoftmaxOptions::UnPackTo(LogSoftmaxOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; +} + +inline flatbuffers::Offset LogSoftmaxOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const LogSoftmaxOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateLogSoftmaxOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateLogSoftmaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const LogSoftmaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const LogSoftmaxOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateLogSoftmaxOptions( + _fbb); +} + +inline CastOptionsT *CastOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new CastOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void CastOptions::UnPackTo(CastOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; +} + +inline flatbuffers::Offset CastOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const CastOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateCastOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateCastOptions(flatbuffers::FlatBufferBuilder &_fbb, const CastOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const CastOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateCastOptions( + _fbb); } -inline OperatorCodeT *OperatorCode::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline DequantizeOptionsT *DequantizeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new DequantizeOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void DequantizeOptions::UnPackTo(DequantizeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; +} + +inline flatbuffers::Offset DequantizeOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const DequantizeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateDequantizeOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateDequantizeOptions(flatbuffers::FlatBufferBuilder &_fbb, const DequantizeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const DequantizeOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateDequantizeOptions( + _fbb); +} + +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; @@ -6507,8 +5922,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: { @@ -6587,28 +6001,47 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, auto ptr = reinterpret_cast(obj); return verifier.VerifyTable(ptr); } - default: - return false; + 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); + } + case BuiltinOptions_LogSoftmaxOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_CastOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_DequantizeOptions: { + 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); @@ -6659,8 +6092,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: { @@ -6739,14 +6171,35 @@ inline void *BuiltinOptionsUnion::UnPack( auto ptr = reinterpret_cast(obj); return ptr->UnPack(resolver); } - default: - return nullptr; + 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); + } + case BuiltinOptions_LogSoftmaxOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_CastOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_DequantizeOptions: { + 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); @@ -6797,10 +6250,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); @@ -6878,32 +6329,50 @@ inline flatbuffers::Offset BuiltinOptionsUnion::Pack( auto ptr = reinterpret_cast(value); return CreateStridedSliceOptions(_fbb, ptr, _rehasher).Union(); } - default: - return 0; + 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(); + } + case BuiltinOptions_LogSoftmaxOptions: { + auto ptr = reinterpret_cast(value); + return CreateLogSoftmaxOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_CastOptions: { + auto ptr = reinterpret_cast(value); + return CreateCastOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_DequantizeOptions: { + auto ptr = reinterpret_cast(value); + return CreateDequantizeOptions(_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: { @@ -6919,18 +6388,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: { @@ -6942,8 +6408,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: { @@ -6951,8 +6416,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: { @@ -6960,23 +6424,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: { @@ -6992,18 +6452,15 @@ 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)); + value = new SpaceToBatchNDOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_TransposeOptions: { - value = new TransposeOptionsT( - *reinterpret_cast(u.value)); + value = new TransposeOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_MeanOptions: { @@ -7019,18 +6476,39 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) break; } case BuiltinOptions_SqueezeOptions: { - value = - new SqueezeOptionsT(*reinterpret_cast(u.value)); + value = new SqueezeOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_SequenceRNNOptions: { - value = new SequenceRNNOptionsT( - *reinterpret_cast(u.value)); + value = new SequenceRNNOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_StridedSliceOptions: { - value = new StridedSliceOptionsT( - *reinterpret_cast(u.value)); + 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; + } + case BuiltinOptions_LogSoftmaxOptions: { + value = new LogSoftmaxOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_CastOptions: { + value = new CastOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_DequantizeOptions: { + value = new DequantizeOptionsT(*reinterpret_cast(u.value)); break; } default: @@ -7200,8 +6678,37 @@ inline void BuiltinOptionsUnion::Reset() { delete ptr; break; } - default: + 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; + } + case BuiltinOptions_LogSoftmaxOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_CastOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_DequantizeOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; break; + } + default: break; } value = nullptr; type = BuiltinOptions_NONE; @@ -7211,25 +6718,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/schema/upgrade_schema.py b/tensorflow/contrib/lite/schema/upgrade_schema.py index 94f5730be5d991ae13fb019e4d035e23f76fe441..e0b36d3d3ee94b00cccd3968d14c63fe19c3c27c 100644 --- a/tensorflow/contrib/lite/schema/upgrade_schema.py +++ b/tensorflow/contrib/lite/schema/upgrade_schema.py @@ -39,8 +39,8 @@ import tensorflow as tf from tensorflow.python.platform import resource_loader parser = argparse.ArgumentParser( - description="Script to move TFLite models from pre-release schema to" - " new schema.") + description="Script to move TFLite models from pre-release schema to " + "new schema.") parser.add_argument( "input", type=str, @@ -48,7 +48,7 @@ parser.add_argument( parser.add_argument( "output", type=str, - help="Output json or bin TensorFlow lite model compliant with" + help="Output json or bin TensorFlow lite model compliant with " "the new schema. Extension must be `.json`, `.bin` or `.tflite`.") @@ -258,7 +258,7 @@ class Converter(object): # Check if builtin_code is the appropriate string type # use type("") instead of str or unicode. for py2and3 if not isinstance(operator_code["builtin_code"], type(u"")): - raise ValueError("builtin_code %r is non-string. this usually means" + raise ValueError("builtin_code %r is non-string. this usually means " "your model has consistency problems." % (operator_code["builtin_code"])) operator_code["builtin_code"] = (RemapOperator( diff --git a/tensorflow/contrib/lite/simple_memory_arena.cc b/tensorflow/contrib/lite/simple_memory_arena.cc index 4aab244989ca5300fbe74162e03deaac89af60ad..2f2004f56bcad5b56f9dd6d4bc824ec14d79e795 100644 --- a/tensorflow/contrib/lite/simple_memory_arena.cc +++ b/tensorflow/contrib/lite/simple_memory_arena.cc @@ -113,21 +113,21 @@ TfLiteStatus SimpleMemoryArena::Commit(TfLiteContext* context) { underlying_buffer_size_ = required_size; underlying_buffer_aligned_ptr_ = new_underlying_buffer_aligned_ptr; } - commited_ = true; + committed_ = true; return underlying_buffer_ != nullptr ? kTfLiteOk : kTfLiteError; } TfLiteStatus SimpleMemoryArena::ResolveAlloc(TfLiteContext* context, const ArenaAlloc& alloc, char** output_ptr) { - TF_LITE_ENSURE(context, commited_); + TF_LITE_ENSURE(context, committed_); TF_LITE_ENSURE(context, output_ptr != nullptr); *output_ptr = underlying_buffer_aligned_ptr_ + alloc.offset; return kTfLiteOk; } TfLiteStatus SimpleMemoryArena::Clear() { - commited_ = false; + committed_ = false; high_water_mark_ = 0; allocs_.clear(); return kTfLiteOk; diff --git a/tensorflow/contrib/lite/simple_memory_arena.h b/tensorflow/contrib/lite/simple_memory_arena.h index 0c5e00a1f2e6a3303556ec54d8e50e8398644bf5..5faf78b59e3755d22e4e866d433e622baa6c66c1 100644 --- a/tensorflow/contrib/lite/simple_memory_arena.h +++ b/tensorflow/contrib/lite/simple_memory_arena.h @@ -22,7 +22,7 @@ limitations under the License. namespace tflite { // This little structure holds the offset and the size for a dynamic memory -// allocation in the memory arena. When the arena is commited and the +// allocation in the memory arena. When the arena is committed and the // underlying buffer is set, the alloc can be resolved into an actual memory // pointer. struct ArenaAlloc { @@ -36,14 +36,14 @@ 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: explicit SimpleMemoryArena(size_t arena_alignment) - : commited_(false), + : committed_(false), arena_alignment_(arena_alignment), high_water_mark_(0), underlying_buffer_size_(0), @@ -73,7 +73,7 @@ class SimpleMemoryArena { } private: - bool commited_; + bool committed_; size_t arena_alignment_; size_t high_water_mark_; std::unique_ptr underlying_buffer_; diff --git a/tensorflow/contrib/lite/special_rules.bzl b/tensorflow/contrib/lite/special_rules.bzl new file mode 100644 index 0000000000000000000000000000000000000000..54083c49182c707620cbd231b957405cfe24be92 --- /dev/null +++ b/tensorflow/contrib/lite/special_rules.bzl @@ -0,0 +1,6 @@ +"""External versions of build rules that differ outside of Google.""" + +def tflite_portable_test_suite(**kwargs): + """This is a no-op outside of Google.""" + _ignore = [kwargs] + pass 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 7f84a0ab9bb5b9fbacba1728f8bfe4df25f13f86..631601656dbad59b7cf0b9b81276d42ea9322cb6 100644 --- a/tensorflow/contrib/lite/testing/BUILD +++ b/tensorflow/contrib/lite/testing/BUILD @@ -8,6 +8,7 @@ load( "//tensorflow/contrib/lite:build_def.bzl", "gen_zipped_test_files", ) +load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") load( "//tensorflow:tensorflow.bzl", "tf_cc_test", @@ -25,6 +26,7 @@ gen_zipped_test_files( "conv.zip", "depthwiseconv.zip", "div.zip", + "exp.zip", "fully_connected.zip", "fused_batch_norm.zip", "gather.zip", @@ -32,6 +34,8 @@ gen_zipped_test_files( "l2_pool.zip", "l2norm.zip", "local_response_norm.zip", + "log_softmax.zip", + "lstm.zip", "max_pool.zip", "mean.zip", "mul.zip", @@ -45,9 +49,11 @@ gen_zipped_test_files( "softmax.zip", "space_to_batch_nd.zip", "space_to_depth.zip", + "split.zip", "squeeze.zip", "strided_slice.zip", "sub.zip", + "topk.zip", "transpose.zip", ], ) @@ -121,6 +127,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"], @@ -195,6 +216,125 @@ cc_binary( ], ) +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"], + tags = [ + "tflite_not_portable", + ], + deps = [ + ":tf_driver", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "generate_testspec", + testonly = 1, + srcs = ["generate_testspec.cc"], + hdrs = ["generate_testspec.h"], + deps = [ + ":join", + ":split", + ":tf_driver", + "//tensorflow/core:framework", + ], +) + +cc_test( + name = "generate_testspec_test", + size = "small", + srcs = ["generate_testspec_test.cc"], + tags = [ + "tflite_not_portable", + ], + deps = [ + ":generate_testspec", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "tflite_diff_util", + testonly = 1, + srcs = ["tflite_diff_util.cc"], + hdrs = ["tflite_diff_util.h"], + deps = [ + ":generate_testspec", + ":parse_testdata_lib", + ":split", + ":tflite_driver", + ":util", + "//tensorflow/contrib/lite:builtin_op_data", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:string", + "//tensorflow/contrib/lite/kernels:builtin_ops", + ], +) + +cc_library( + name = "tflite_diff_flags", + testonly = 1, + hdrs = ["tflite_diff_flags.h"], + deps = [ + ":split", + ":tflite_diff_util", + ] + select({ + "//conditions:default": [ + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + ], + "//tensorflow:android": [ + "//tensorflow/core:android_tensorflow_lib", + ], + }), +) + +tf_cc_test( + name = "tflite_diff_example_test", + size = "medium", + srcs = ["tflite_diff_example_test.cc"], + args = [ + "--tensorflow_model=third_party/tensorflow/contrib/lite/testdata/multi_add.pb", + "--tflite_model=third_party/tensorflow/contrib/lite/testdata/multi_add.bin", + "--input_layer=a,b,c,d", + "--input_layer_type=float,float,float,float", + "--input_layer_shape=1,3,4,3:1,3,4,3:1,3,4,3:1,3,4,3", + "--output_layer=x,y", + ], + data = [ + "//tensorflow/contrib/lite:testdata/multi_add.bin", + "//tensorflow/contrib/lite:testdata/multi_add.pb", + ], + tags = [ + "no_cuda_on_cpu_tap", + "no_oss", + "tflite_not_portable", + ], + deps = [ + ":tflite_diff_flags", + ":tflite_diff_util", + ], +) + tf_cc_test( name = "generated_examples_zip_test", size = "large", @@ -206,8 +346,11 @@ tf_cc_test( "--unzip_binary_path=/usr/bin/unzip", ], data = [":optest"], - shard_count = 10, - tags = ["no_oss"], + shard_count = 20, + tags = [ + "no_oss", + "tflite_not_portable", + ], deps = [ ":parse_testdata_lib", ":tflite_driver", @@ -241,3 +384,5 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) + +tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py index e7606eecc462ae2c3935475cc53c91f039980223..420bdb41f140d2ef35c66f47f386b016eb9e60c0 100644 --- a/tensorflow/contrib/lite/testing/generate_examples.py +++ b/tensorflow/contrib/lite/testing/generate_examples.py @@ -36,6 +36,7 @@ import traceback import zipfile import numpy as np from six import StringIO +from six.moves import xrange # TODO(aselle): Disable GPU for now os.environ["CUDA_VISIBLE_DEVICES"] = "-1" @@ -46,6 +47,7 @@ from google.protobuf import text_format # TODO(aselle): switch to TensorFlow's resource_loader from tensorflow.contrib.lite.testing import generate_examples_report as report_lib from tensorflow.python.framework import graph_util as tf_graph_util +from tensorflow.python.ops import rnn parser = argparse.ArgumentParser(description="Script to generate TFLite tests.") parser.add_argument("output_path", @@ -99,16 +101,32 @@ KNOWN_BUGS = { 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" + # 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", } +class ExtraTocoOptions(object): + """Additonal toco options besides input, output, shape.""" + + def __init__(self): + # Whether to ignore control dependency nodes. + self.drop_control_dependency = False + # Allow custom ops in the toco conversion. + self.allow_custom_ops = False + # Rnn states that are used to support rnn / lstm cells. + self.rnn_states = None + + def toco_options(data_types, input_arrays, output_arrays, shapes, - drop_control_dependency): + extra_toco_options=ExtraTocoOptions()): """Create TOCO options to process a model. Args: @@ -116,8 +134,7 @@ def toco_options(data_types, input_arrays: names of the input tensors output_arrays: name of the output tensors shapes: shapes of the input tensors - drop_control_dependency: whether to ignore control dependency nodes. - + extra_toco_options: additional toco options Returns: the options in a string. """ @@ -133,37 +150,15 @@ def toco_options(data_types, " --input_arrays=%s" % ",".join(input_arrays) + " --input_shapes=%s" % shape_str + " --output_arrays=%s" % ",".join(output_arrays)) - if drop_control_dependency: + if extra_toco_options.drop_control_dependency: s += " --drop_control_dependency" + if extra_toco_options.allow_custom_ops: + s += " --allow_custom_ops" + if extra_toco_options.rnn_states: + s += (" --rnn_states='" + extra_toco_options.rnn_states + "'") return s -def write_toco_options(filename, - data_types, - input_arrays, - output_arrays, - shapes, - drop_control_dependency=False): - """Create TOCO options to process a model. - - Args: - filename: Filename to write the options to. - data_types: input and inference types used by TOCO. - input_arrays: names of the input tensors - output_arrays: names of the output tensors - shapes: shapes of the input tensors - drop_control_dependency: whether to ignore control dependency nodes. - """ - with open(filename, "w") as fp: - fp.write( - toco_options( - data_types=data_types, - input_arrays=input_arrays, - output_arrays=output_arrays, - shapes=shapes, - drop_control_dependency=drop_control_dependency)) - - def write_examples(fp, examples): """Given a list `examples`, write a text format representation. @@ -241,7 +236,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) @@ -281,12 +276,14 @@ def make_control_dep_tests(zip_path): return [input_values], sess.run( outputs, feed_dict=dict(zip(inputs, [input_values]))) + extra_toco_options = ExtraTocoOptions() + extra_toco_options.drop_control_dependency = True make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs, - drop_control_dependency=True) + extra_toco_options) def toco_convert(graph_def_str, input_tensors, output_tensors, - drop_control_dependency=False): + extra_toco_options): """Convert a model's graph def into a tflite model. NOTE: this currently shells out to the toco binary, but we would like @@ -294,9 +291,9 @@ def toco_convert(graph_def_str, input_tensors, output_tensors, Args: graph_def_str: Graph def proto in serialized string format. - input_tensors: List of input tensor tuples `(name, shape, type)` - output_tensors: List of output tensors (names) - drop_control_dependency: whether to ignore control dependency nodes. + input_tensors: List of input tensor tuples `(name, shape, type)`. + output_tensors: List of output tensors (names). + extra_toco_options: Additional toco options. Returns: output tflite model, log_txt from conversion @@ -308,7 +305,7 @@ def toco_convert(graph_def_str, input_tensors, output_tensors, input_arrays=[x[0] for x in input_tensors], shapes=[x[1] for x in input_tensors], output_arrays=output_tensors, - drop_control_dependency=drop_control_dependency) + extra_toco_options=extra_toco_options) with tempfile.NamedTemporaryFile() as graphdef_file, \ tempfile.NamedTemporaryFile() as output_file, \ @@ -327,11 +324,18 @@ 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, make_test_inputs, - drop_control_dependency=False): + extra_toco_options=ExtraTocoOptions(), + use_frozen_graph=False): """Helper to make a zip file of a bunch of TensorFlow models. This does a cartestian product of the dictionary of test_parameters and @@ -349,7 +353,9 @@ def make_zip_of_tests(zip_path, `[input1, input2, ...], [output1, output2, ...]` make_test_inputs: function taking `curr_params`, `session`, `input_tensors`, `output_tensors` and returns tuple `(input_values, output_values)`. - drop_control_dependency: whether to ignore control dependency nodes. + extra_toco_options: Additional toco options. + use_frozen_graph: Whether or not freeze graph before toco converter. + Raises: RuntimeError: if there are toco errors that can't be ignored. """ @@ -409,21 +415,25 @@ def make_zip_of_tests(zip_path, return None, report report["toco"] = report_lib.FAILED report["tf"] = report_lib.SUCCESS - # Convert graph to toco + input_tensors = [(input_tensor.name.split(":")[0], + input_tensor.get_shape(), input_tensor.dtype) + for input_tensor in inputs] + output_tensors = [normalize_output_name(out.name) for out in outputs] + graph_def = freeze_graph( + sess, + tf.global_variables() + inputs + + outputs) if use_frozen_graph else sess.graph_def tflite_model_binary, toco_log = toco_convert( - 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) + graph_def.SerializeToString(), input_tensors, output_tensors, + extra_toco_options) report["toco"] = (report_lib.SUCCESS if tflite_model_binary is not None else report_lib.FAILED) report["toco_log"] = toco_log if FLAGS.save_graphdefs: archive.writestr(label + ".pb", - text_format.MessageToString(sess.graph_def), + text_format.MessageToString(graph_def), zipfile.ZIP_DEFLATED) if tflite_model_binary: @@ -619,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 @@ -745,6 +755,65 @@ def make_mean_tests(zip_path): 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_log_softmax_tests(zip_path): + """Make a set of tests to do log_softmax.""" + + test_parameters = [{ + "input_dtype": [tf.float32], + "input_shape": [[1, 100], [4, 2], [5, 224]], + }] + + def build_graph(parameters): + """Build the log_softmax op testing graph.""" + input_tensor = tf.placeholder( + dtype=parameters["input_dtype"], + name="input", + shape=parameters["input_shape"]) + + out = tf.nn.log_softmax(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) @@ -995,19 +1064,44 @@ def make_depthwiseconv_tests(zip_path): 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 @@ -1477,7 +1571,7 @@ def make_transpose_tests(zip_path): }, { "dtype": [tf.float32], "input_shape": [[1, 2, 3, 4, 5]], - "perm": [[0, 1, 2, 3, 4]], + "perm": [[4, 3, 2, 1, 0]], "constant_perm": [True, False], }] @@ -1560,10 +1654,24 @@ def make_strided_slice_tests(zip_path): "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, [1, 1, 1, 1], [2, 1, 3, 1]], - "begin_mask": [None, 1, 2, 8], - "end_mask": [None, 1, 2, 8], - "shrink_axis_mask": [None, 1, 2, 4, 8, 11, 15, -1], + "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 { @@ -1572,14 +1680,15 @@ def make_strided_slice_tests(zip_path): "input_shape": [[2, 3]], "begin": [[0, 0], [1, 0]], "end": [[2, 3], [2, 2]], - "strides": [None, [1, 1], [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, tf.int32, tf.int64], + "dtype": [tf.float32], "index_type": [tf.int32], "input_shape": [[2, 3]], "begin": [[0, -1]], @@ -1588,6 +1697,7 @@ def make_strided_slice_tests(zip_path): "begin_mask": [None, 1, 2], "end_mask": [None, 1, 2], "shrink_axis_mask": [None, 1, 2, 3, -1], + "constant_indices": [False], }, ] @@ -1597,23 +1707,29 @@ def make_strided_slice_tests(zip_path): dtype=parameters["dtype"], name="input", shape=parameters["input_shape"]) - 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) + 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, @@ -1628,20 +1744,101 @@ def make_strided_slice_tests(zip_path): input_values = create_tensor_data(parameters["dtype"], parameters["input_shape"]) index_type = _TF_TYPE_INFO[parameters["index_type"]][0] - 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 = [input_values, begin_values, end_values] - if stride_values is not None: - values.append(stride_values) + 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_lstm_tests(zip_path): + """Make a set of tests to do basic Lstm cell.""" + + test_parameters = [ + { + "dtype": [tf.float32], + "num_batchs": [1], + "time_step_size": [1], + "input_vec_size": [3], + "num_cells": [4], + }, + ] + + def build_graph(parameters): + """Build a simple graph with BasicLSTMCell.""" + + num_batchs = parameters["num_batchs"] + time_step_size = parameters["time_step_size"] + input_vec_size = parameters["input_vec_size"] + num_cells = parameters["num_cells"] + inputs_after_split = [] + for i in xrange(time_step_size): + one_timestamp_input = tf.placeholder( + dtype=parameters["dtype"], + name="split_{}".format(i), + shape=[num_batchs, input_vec_size]) + inputs_after_split.append(one_timestamp_input) + # Currently lstm identifier has a few limitations: only supports + # forget_bias == 0, inner state activiation == tanh. + # TODO(zhixianyan): Add another test with forget_bias == 1. + # TODO(zhixianyan): Add another test with relu as activation. + lstm_cell = tf.contrib.rnn.BasicLSTMCell( + num_cells, forget_bias=0.0, state_is_tuple=True) + cell_outputs, _ = rnn.static_rnn( + lstm_cell, inputs_after_split, dtype=tf.float32) + out = cell_outputs[-1] + return inputs_after_split, [out] + + def build_inputs(parameters, sess, inputs, outputs): + """Feed inputs, assign vairables, and freeze graph.""" + + with tf.variable_scope("", reuse=True): + kernel = tf.get_variable("rnn/basic_lstm_cell/kernel") + bias = tf.get_variable("rnn/basic_lstm_cell/bias") + kernel_values = create_tensor_data( + parameters["dtype"], [kernel.shape[0], kernel.shape[1]], -1, 1) + bias_values = create_tensor_data(parameters["dtype"], [bias.shape[0]], 0, + 1) + sess.run(tf.group(kernel.assign(kernel_values), bias.assign(bias_values))) + + num_batchs = parameters["num_batchs"] + time_step_size = parameters["time_step_size"] + input_vec_size = parameters["input_vec_size"] + input_values = [] + for _ in xrange(time_step_size): + tensor_data = create_tensor_data(parameters["dtype"], + [num_batchs, input_vec_size], 0, 1) + input_values.append(tensor_data) + out = sess.run(outputs, feed_dict=dict(zip(inputs, input_values))) + return input_values, out + + # TODO(zhixianyan): Automatically generate rnn_states for lstm cell. + extra_toco_options = ExtraTocoOptions() + extra_toco_options.rnn_states = ( + "{state_array:rnn/BasicLSTMCellZeroState/zeros," + "back_edge_source_array:rnn/basic_lstm_cell/Add_1,size:4}," + "{state_array:rnn/BasicLSTMCellZeroState/zeros_1," + "back_edge_source_array:rnn/basic_lstm_cell/Mul_2,size:4}") + + make_zip_of_tests( + zip_path, + test_parameters, + build_graph, + build_inputs, + extra_toco_options, + use_frozen_graph=True) + + 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( @@ -1649,6 +1846,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 @@ -1697,10 +1920,15 @@ 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, + "log_softmax.zip": make_log_softmax_tests, + "lstm.zip": make_lstm_tests, } out = FLAGS.zip_to_output bin_path = FLAGS.toco diff --git a/tensorflow/contrib/lite/testing/generate_testspec.cc b/tensorflow/contrib/lite/testing/generate_testspec.cc new file mode 100644 index 0000000000000000000000000000000000000000..eb3deafb6986e877f0a553a8b6f712102af4caca --- /dev/null +++ b/tensorflow/contrib/lite/testing/generate_testspec.cc @@ -0,0 +1,88 @@ +/* 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/generate_testspec.h" +#include "tensorflow/contrib/lite/testing/join.h" +#include "tensorflow/contrib/lite/testing/split.h" +#include "tensorflow/contrib/lite/testing/tf_driver.h" +#include "tensorflow/core/framework/types.h" + +namespace tflite { +namespace testing { + +void GenerateTestSpecFromTensorflowModel( + std::iostream& stream, const string& tensorflow_model_path, + const string& tflite_model_path, const std::vector& input_layer, + const std::vector& input_layer_type, + const std::vector& input_layer_shape, + const std::vector& output_layer) { + CHECK_EQ(input_layer.size(), input_layer_type.size()); + CHECK_EQ(input_layer.size(), input_layer_shape.size()); + + // Initialize random functions. + static unsigned int seed = 0; + std::function float_rand = [](int idx) { + return static_cast(rand_r(&seed)) / RAND_MAX - 0.5f; + }; + + // Generate inputs. + std::vector input_values; + input_values.resize(input_layer.size()); + for (int i = 0; i < input_layer.size(); i++) { + tensorflow::DataType type; + CHECK(DataTypeFromString(input_layer_type[i], &type)); + auto shape = Split(input_layer_shape[i], ","); + + switch (type) { + case tensorflow::DT_FLOAT: { + const auto& data = GenerateRandomTensor(shape, float_rand); + input_values[i] = Join(data.data(), data.size(), ","); + break; + } + default: + + fprintf(stderr, "Unsupported type %d when generating testspec\n", type); + return; + } + } + + // Invoke tensorflow model. + TfDriver runner(input_layer, input_layer_type, input_layer_shape, + output_layer); + runner.LoadModel(tensorflow_model_path); + for (int i = 0; i < input_values.size(); i++) { + runner.SetInput(i, input_values[i]); + } + runner.Invoke(); + + // Write test spec. + stream << "load_model: " << tflite_model_path << "\n"; + stream << "reshape {\n"; + for (const auto& shape : input_layer_shape) { + stream << " input: \"" << shape << "\"\n"; + } + stream << "}\n"; + stream << "invoke {\n"; + for (const auto& value : input_values) { + stream << " input: \"" << value << "\"\n"; + } + for (int i = 0; i < output_layer.size(); i++) { + stream << " output: \"" << runner.ReadOutput(i) << "\"\n"; + } + stream << "}\n"; +} + +} // namespace testing +} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/generate_testspec.h b/tensorflow/contrib/lite/testing/generate_testspec.h new file mode 100644 index 0000000000000000000000000000000000000000..3529ee709b66625fff6e2a35b78e47f3778f0fe7 --- /dev/null +++ b/tensorflow/contrib/lite/testing/generate_testspec.h @@ -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. +==============================================================================*/ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_GENERATE_TESTSPEC_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_GENERATE_TESTSPEC_H_ + +#include +#include +#include + +namespace tflite { +namespace testing { + +// Generate test spec by executing TensorFlow model on random inputs. +// The test spec can be consumed by ParseAndRunTests. +// See test spec format in parse_testdata.h +// +// Inputs: +// stream: mutable iostream that contains the contents of test spec. +// tensorflow_model_path: path to TensorFlow model. +// tflite_model_path: path to tflite_model_path that the test spec runs +// against. input_layer: names of input tensors. Example: input1 +// input_layer_type: datatypes of input tensors. Example: float +// input_layer_shape: shapes of input tensors, separated by comma. example: +// 1,3,4 output_layer: names of output tensors. Example: output +void GenerateTestSpecFromTensorflowModel( + std::iostream& stream, const string& tensorflow_model_path, + const string& tflite_model_path, const std::vector& input_layer, + const std::vector& input_layer_type, + const std::vector& input_layer_shape, + const std::vector& output_layer); + +// Generates random values that are filled into the tensor. +// random_func returns the generated random element at given index. +template +std::vector GenerateRandomTensor(const std::vector& shape, + const std::function& random_func) { + int64_t num_elements = 1; + for (const int dim : shape) { + num_elements *= dim; + } + + std::vector result(num_elements); + for (int i = 0; i < num_elements; i++) { + result[i] = random_func(i); + } + return result; +} + +} // namespace testing +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_GENERATE_TESTSPEC_H_ diff --git a/tensorflow/contrib/lite/testing/generate_testspec_test.cc b/tensorflow/contrib/lite/testing/generate_testspec_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..2a97b757a413246c9ad9b5f453741b13e381c903 --- /dev/null +++ b/tensorflow/contrib/lite/testing/generate_testspec_test.cc @@ -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. +==============================================================================*/ +#include "tensorflow/contrib/lite/testing/generate_testspec.h" + +#include +#include + +namespace tflite { +namespace testing { +namespace { + +TEST(GenerateRandomTensor, FloatValue) { + static unsigned int seed = 0; + std::function float_rand = [](int idx) { + return static_cast(rand_r(&seed)) / RAND_MAX - 0.5f; + }; + + std::set values; + float sum_x_square = 0.0f; + float sum_x = 0.0f; + for (int i = 0; i < 100; i++) { + const auto& data = GenerateRandomTensor({1, 3, 4}, float_rand); + for (float value : data) { + values.insert(value); + sum_x_square += value * value; + sum_x += value; + } + } + + // Eech round, generated tensor has different values. + EXPECT_GT(values.size(), 200); + int num = 1 * 3 * 4 * 100; + float stddev = sum_x_square / num - (sum_x / num) * (sum_x / num); + + // Stddev is greater than 1/2 stddev of uniform distribution: (B-A)^2 / 12 + float minstddev = 1.0f / 12 / 2; + EXPECT_GT(stddev, minstddev); +} + +} // namespace +} // namespace testing +} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc index d73c9937cec19634c8ede620d5f55af526188eca..88c5aaa099e79ccd9dcef6a6f1886c3ddfc73e4d 100644 --- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc +++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc @@ -47,12 +47,6 @@ 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"(^\/adda.*input_shape_1=\[1,3,4,3\],input_shape_2=\[3\])", "68500195"}, - {R"(^\/mula.*input_shape_1=\[1,3,4,3\],input_shape_2=\[3\])", "68500195"}, - {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"(^\/adda.*int32)", "68808744"}, {R"(^\/constant.*int32)", "68808744"}, @@ -91,12 +85,12 @@ std::map kBrokenTests = { // ResizeBilinear looks completely incompatible with Tensorflow {R"(^\/resize_bilinear.*dtype=tf.int32)", "72401107"}, - {R"(^\/resize_bilinearalign_corners=True,.*,size=\[2,2\])", "72401483"}, - {R"(^\/resize_bilinearalign_corners=True,.*,size=\[4,3\])", "72401483"}, - {R"(^\/resize_bilinearalign_corners=True,.*,size=\[5,6\])", "72401483"}, // Transpose only supports 1D-4D input tensors. {R"(^\/transpose.*input_shape=\[.,.,.,.,.\])", "71545879"}, + + // Lstm kernel gets different results on tsan, asan, msan. + {R"(^\/lstmdtype=tf.float32.*)", "73830845"}, }; // Allows test data to be unzipped into a temporary directory and makes @@ -240,21 +234,25 @@ 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(div) +INSTANTIATE_TESTS(exp) INSTANTIATE_TESTS(fully_connected) INSTANTIATE_TESTS(fused_batch_norm) INSTANTIATE_TESTS(gather) INSTANTIATE_TESTS(global_batch_norm) -INSTANTIATE_TESTS(l2norm) INSTANTIATE_TESTS(l2_pool) +INSTANTIATE_TESTS(l2norm) INSTANTIATE_TESTS(local_response_norm) +INSTANTIATE_TESTS(log_softmax) +INSTANTIATE_TESTS(lstm) INSTANTIATE_TESTS(max_pool) +INSTANTIATE_TESTS(mean) INSTANTIATE_TESTS(mul) INSTANTIATE_TESTS(pad) INSTANTIATE_TESTS(relu) @@ -264,13 +262,13 @@ INSTANTIATE_TESTS(reshape) INSTANTIATE_TESTS(resize_bilinear) INSTANTIATE_TESTS(sigmoid) INSTANTIATE_TESTS(softmax) +INSTANTIATE_TESTS(space_to_batch_nd) INSTANTIATE_TESTS(space_to_depth) -INSTANTIATE_TESTS(sub) -INSTANTIATE_TESTS(div) -INSTANTIATE_TESTS(transpose) -INSTANTIATE_TESTS(mean) +INSTANTIATE_TESTS(split) INSTANTIATE_TESTS(squeeze) INSTANTIATE_TESTS(strided_slice) +INSTANTIATE_TESTS(sub) +INSTANTIATE_TESTS(transpose) } // namespace testing } // namespace tflite diff --git a/tensorflow/compiler/xla/array2d.cc b/tensorflow/contrib/lite/testing/join.h similarity index 51% rename from tensorflow/compiler/xla/array2d.cc rename to tensorflow/contrib/lite/testing/join.h index 418587c1f75c7249f92e925455d40685d870c57a..ce8c072a21c6e61e8ab8ae12ba52418e6144009a 100644 --- a/tensorflow/compiler/xla/array2d.cc +++ b/tensorflow/contrib/lite/testing/join.h @@ -12,25 +12,31 @@ WITHOUT 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 "tensorflow/compiler/xla/array2d.h" -#include "tensorflow/compiler/xla/ptr_util.h" - -namespace xla { - -std::unique_ptr> MakeLinspaceArray2D(float from, float to, - int64 n1, int64 n2) { - auto array = MakeUnique>(n1, n2); - int64 count = n1 * n2; - float step = (count > 1) ? (to - from) / (count - 1) : 0.0f; - auto set = [&array, n1, n2](int64 index, float value) { - (*array)(index / n2, index % n2) = value; - }; - for (int64 i = 0; i < count - 1; ++i) { - set(i, from + i * step); +#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 ""; } - set(count - 1, to); - return array; + std::stringstream result; + result << data[0]; + for (int i = 1; i < len; i++) { + result << delimiter << data[i]; + } + return result.str(); } -} // namespace xla +} // 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/parse_testdata.cc b/tensorflow/contrib/lite/testing/parse_testdata.cc index 0caef0fe2201a668b2235a98304eb353072a3c2f..389688d552051ea735ce71533943af33df5059ef 100644 --- a/tensorflow/contrib/lite/testing/parse_testdata.cc +++ b/tensorflow/contrib/lite/testing/parse_testdata.cc @@ -192,27 +192,25 @@ TfLiteStatus CheckOutputs(tflite::Interpreter* interpreter, int model_outputs = interpreter->outputs().size(); TF_LITE_ENSURE_EQ(context, model_outputs, example.outputs.size()); for (size_t i = 0; i < interpreter->outputs().size(); i++) { + bool tensors_differ = false; int output_index = interpreter->outputs()[i]; if (const float* data = interpreter->typed_tensor(output_index)) { for (size_t idx = 0; idx < example.outputs[i].flat_data.size(); idx++) { float computed = data[idx]; float reference = example.outputs[0].flat_data[idx]; 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) { + bool local_tensors_differ = + std::abs(reference) < kRelativeThreshold + ? diff > kAbsoluteThreshold + : diff > kRelativeThreshold * std::abs(reference); + if (local_tensors_differ) { fprintf(stdout, "output[%zu][%zu] did not match %f vs reference %f\n", i, idx, data[idx], reference); - return kTfLiteError; + tensors_differ = local_tensors_differ; } } - fprintf(stderr, "\n"); } else if (const int32_t* data = interpreter->typed_tensor(output_index)) { for (size_t idx = 0; idx < example.outputs[i].flat_data.size(); idx++) { @@ -221,10 +219,9 @@ TfLiteStatus CheckOutputs(tflite::Interpreter* interpreter, if (std::abs(computed - reference) > 0) { fprintf(stderr, "output[%zu][%zu] did not match %d vs reference %d\n", i, idx, computed, reference); - return kTfLiteError; + tensors_differ = true; } } - 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++) { @@ -235,14 +232,15 @@ TfLiteStatus CheckOutputs(tflite::Interpreter* interpreter, "output[%zu][%zu] did not match %" PRId64 " vs reference %" PRId64 "\n", i, idx, computed, reference); - return kTfLiteError; + tensors_differ = true; } } - fprintf(stderr, "\n"); } else { fprintf(stderr, "output[%zu] was not float or int data\n", i); return kTfLiteError; } + fprintf(stderr, "\n"); + if (tensors_differ) return kTfLiteError; } return kTfLiteOk; } @@ -319,8 +317,9 @@ class Reshape : public Message { // This is the top-level message in a test file. class TestData : public Message { public: - explicit TestData(TestRunner* test_runner) : test_runner_(test_runner) {} - + explicit TestData(TestRunner* test_runner) + : test_runner_(test_runner), num_invocations_(0), max_invocations_(-1) {} + void SetMaxInvocations(int max) { max_invocations_ = max; } void SetField(const std::string& name, const std::string& value) override { if (name == "load_model") { test_runner_->LoadModel(value); @@ -334,7 +333,12 @@ class TestData : public Message { Message* AddChild(const std::string& s) override { if (s == "invoke") { test_runner_->AllocateTensors(); - return Store(new Invoke(test_runner_)); + if (max_invocations_ == -1 || num_invocations_ < max_invocations_) { + ++num_invocations_; + return Store(new Invoke(test_runner_)); + } else { + return nullptr; + } } else if (s == "reshape") { return Store(new Reshape(test_runner_)); } @@ -343,10 +347,14 @@ class TestData : public Message { private: TestRunner* test_runner_; + int num_invocations_; + int max_invocations_; }; -bool ParseAndRunTests(std::istream* input, TestRunner* test_runner) { +bool ParseAndRunTests(std::istream* input, TestRunner* test_runner, + int max_invocations) { TestData test_data(test_runner); + test_data.SetMaxInvocations(max_invocations); Message::Read(input, &test_data); return test_runner->IsValid() && test_runner->GetOverallSuccess(); } diff --git a/tensorflow/contrib/lite/testing/parse_testdata.h b/tensorflow/contrib/lite/testing/parse_testdata.h index 7ebf362eb99c5f4cf6ea3654cf71e13ff1de99b3..d94361d735e2be8dc130dc8d6bf0bb5c822ebb7c 100644 --- a/tensorflow/contrib/lite/testing/parse_testdata.h +++ b/tensorflow/contrib/lite/testing/parse_testdata.h @@ -66,7 +66,8 @@ TfLiteStatus CheckOutputs(tflite::Interpreter* interpreter, const Example&); // output: "12,3,4,545,3" // output: "0.01,0.02" // } -bool ParseAndRunTests(std::istream* input, TestRunner* test_runner); +bool ParseAndRunTests(std::istream* input, TestRunner* test_runner, + int max_invocations = -1); } // namespace testing } // namespace tflite diff --git a/tensorflow/contrib/lite/testing/test_runner.h b/tensorflow/contrib/lite/testing/test_runner.h index 60eaafa474a01887bee12b031b1f59cc5c91f173..05770beee23275ebe210606dbfd2b33eea17612d 100644 --- a/tensorflow/contrib/lite/testing/test_runner.h +++ b/tensorflow/contrib/lite/testing/test_runner.h @@ -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; 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_diff_example_test.cc b/tensorflow/contrib/lite/testing/tflite_diff_example_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..3817e68111dbaaf2a38ceff9fbc38f30f303cb5f --- /dev/null +++ b/tensorflow/contrib/lite/testing/tflite_diff_example_test.cc @@ -0,0 +1,28 @@ +/* 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/tflite_diff_flags.h" +#include "tensorflow/contrib/lite/testing/tflite_diff_util.h" + +int main(int argc, char** argv) { + ::tflite::testing::DiffOptions options = + ::tflite::testing::ParseTfliteDiffFlags(&argc, argv); + for (int i = 0; i < 100; i++) { + if (!tflite::testing::RunDiffTest(options)) { + return 1; + } + } + return 0; +} diff --git a/tensorflow/contrib/lite/testing/tflite_diff_flags.h b/tensorflow/contrib/lite/testing/tflite_diff_flags.h new file mode 100644 index 0000000000000000000000000000000000000000..5f1129d501b7235f1202b704cf36904e07b8720e --- /dev/null +++ b/tensorflow/contrib/lite/testing/tflite_diff_flags.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_TESTING_TFLITE_DIFF_FLAGS_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DIFF_FLAGS_H_ + +#include "tensorflow/contrib/lite/testing/split.h" +#include "tensorflow/contrib/lite/testing/tflite_diff_util.h" +#include "tensorflow/core/util/command_line_flags.h" + +namespace tflite { +namespace testing { + +DiffOptions ParseTfliteDiffFlags(int* argc, char** argv) { + struct { + string tensorflow_model; + string tflite_model; + string input_layer; + string input_layer_type; + string input_layer_shape; + string output_layer; + } values; + + std::vector flags = { + tensorflow::Flag("tensorflow_model", &values.tensorflow_model, + "Path of tensorflow model."), + tensorflow::Flag("tflite_model", &values.tflite_model, + "Path of tensorflow lite model."), + tensorflow::Flag("input_layer", &values.input_layer, + "Names of input tensors, separated by comma. Example: " + "input_1,input_2"), + tensorflow::Flag("input_layer_type", &values.input_layer_type, + "Data types of input tensors, separated by comma. " + "Example: float,int"), + tensorflow::Flag( + "input_layer_shape", &values.input_layer_shape, + "Shapes of input tensors, separated by colon. Example: 1,3,4,1:2"), + tensorflow::Flag("output_layer", &values.output_layer, + "Names of output tensors, separated by comma. Example " + "output_1,output_2"), + }; + + 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 {values.tensorflow_model, + values.tflite_model, + Split(values.input_layer, ","), + Split(values.input_layer_type, ","), + Split(values.input_layer_shape, ":"), + Split(values.output_layer, ",")}; +} + +} // namespace testing +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DIFF_FLAGS_H_ diff --git a/tensorflow/contrib/lite/testing/tflite_diff_util.cc b/tensorflow/contrib/lite/testing/tflite_diff_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..9ef4e1f66c7d31c746c18d63495e760585d4af9e --- /dev/null +++ b/tensorflow/contrib/lite/testing/tflite_diff_util.cc @@ -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. +==============================================================================*/ +#include +#include +#include +#include + +#include "tensorflow/contrib/lite/testing/generate_testspec.h" +#include "tensorflow/contrib/lite/testing/parse_testdata.h" +#include "tensorflow/contrib/lite/testing/tflite_diff_util.h" +#include "tensorflow/contrib/lite/testing/tflite_driver.h" + +namespace tflite { +namespace testing { + +bool RunDiffTest(const DiffOptions& options) { + std::stringstream tflite_stream; + GenerateTestSpecFromTensorflowModel( + tflite_stream, options.tensorflow_model, options.tflite_model, + options.input_layer, options.input_layer_type, options.input_layer_shape, + options.output_layer); + TfLiteDriver tflite_driver(/*use_nnapi=*/true); + tflite_driver.LoadModel(options.tflite_model); + std::cout << tflite_stream.str(); + return tflite::testing::ParseAndRunTests(&tflite_stream, &tflite_driver); +} +} // namespace testing + +} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tflite_diff_util.h b/tensorflow/contrib/lite/testing/tflite_diff_util.h new file mode 100644 index 0000000000000000000000000000000000000000..326fa6c3e28000dee9b6eb9cc5b3a6c5c87e28d0 --- /dev/null +++ b/tensorflow/contrib/lite/testing/tflite_diff_util.h @@ -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. +==============================================================================*/ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DIFF_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DIFF_UTIL_H_ + +#include + +#include "tensorflow/contrib/lite/string.h" + +namespace tflite { +namespace testing { + +// Configurations to run Tflite diff test. +struct DiffOptions { + // Path of tensorflow model. + string tensorflow_model; + // Path of tensorflow lite model. + string tflite_model; + // Names of input tensors. + // Example: input_1,input_2 + std::vector input_layer; + // Data types of input tensors. + // Example: float,int + std::vector input_layer_type; + // Shapes of input tensors, separated by comma. + // Example: 1,3,4,1 + std::vector input_layer_shape; + // Names of output tensors. + // Example output_1,output_2 + std::vector output_layer; +}; + +// Run a single TensorFLow Lite diff test with a given options. +bool RunDiffTest(const DiffOptions& options); + +} // namespace testing +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DIFF_UTIL_H_ diff --git a/tensorflow/contrib/lite/testing/tflite_driver.cc b/tensorflow/contrib/lite/testing/tflite_driver.cc index bae639ea95318a16c963269de5e55afcb681d4c5..613223f3d4ff212cb8672494243b2d7a1d06b3db 100644 --- a/tensorflow/contrib/lite/testing/tflite_driver.cc +++ b/tensorflow/contrib/lite/testing/tflite_driver.cc @@ -106,8 +106,8 @@ class TfLiteDriver::Expectation { if (error_is_large) { good_output = false; if (verbose) { - std::cerr << " index " << i << ": " << reference - << " != " << computed << std::endl; + std::cerr << " index " << i << ": got " << computed + << ", but expected " << reference << std::endl; } } } @@ -203,6 +203,10 @@ 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: diff --git a/tensorflow/contrib/lite/testing/tflite_driver.h b/tensorflow/contrib/lite/testing/tflite_driver.h index 25689a9fb42c06fa3f8f2f92064cf59e8c331637..02b7de1534e648734d7bc53154afa42f2ef256b4 100644 --- a/tensorflow/contrib/lite/testing/tflite_driver.h +++ b/tensorflow/contrib/lite/testing/tflite_driver.h @@ -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; diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD index 20c156a93262568cf0c6c349b44fbf3d3afa5bc4..395abc532607d3cab49be6e995648f2839c1555c 100644 --- a/tensorflow/contrib/lite/toco/BUILD +++ b/tensorflow/contrib/lite/toco/BUILD @@ -173,7 +173,9 @@ cc_library( "graph_transformations/convert_expanddims_to_reshape.cc", "graph_transformations/convert_pure_conv_to_depthwise.cc", "graph_transformations/convert_reorder_axes.cc", + "graph_transformations/convert_squeeze_to_reshape.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", @@ -185,11 +187,16 @@ cc_library( "graph_transformations/fuse_binary_into_preceding_affine.cc", "graph_transformations/graph_transformations.cc", "graph_transformations/hardcode_min_max.cc", + "graph_transformations/identify_dilated_conv.cc", "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_activation_function_into_constants.cc", "graph_transformations/propagate_array_data_types.cc", "graph_transformations/propagate_fixed_sizes.cc", "graph_transformations/quantize.cc", @@ -204,6 +211,7 @@ 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", @@ -212,12 +220,15 @@ cc_library( "graph_transformations/resolve_constant_concatenation.cc", "graph_transformations/resolve_constant_fake_quant.cc", "graph_transformations/resolve_constant_fill.cc", + "graph_transformations/resolve_constant_gather.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", @@ -232,9 +243,12 @@ cc_library( "graph_transformations/resolve_tensorflow_tile.cc", "graph_transformations/resolve_transpose_attributes.cc", "graph_transformations/unfuse_activation_functions.cc", + "graph_transformations/unpartition_embedding_lookup.cc", + "graph_transformations/unroll_batch_matmul.cc", ], hdrs = [ "graph_transformations/graph_transformations.h", + "graph_transformations/lstm_utils.h", ], visibility = ["//visibility:public"], deps = [ @@ -245,6 +259,7 @@ cc_library( ":tooling_util", ":types_proto_cc", "//tensorflow/core:lib", + "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], ) @@ -317,6 +332,7 @@ cc_library( ":toco_graphviz_dump_options", ":toco_port", ":types_proto_cc", + "//tensorflow/contrib/lite/kernels/internal:quantization_util", "//tensorflow/core:lib", "@com_google_absl//absl/strings", "@protobuf_archive//:protobuf_headers", diff --git a/tensorflow/contrib/lite/toco/args.h b/tensorflow/contrib/lite/toco/args.h index b97a4720a7c4e69f8b69574475d19e0522cfe86d..59a6115920614d38900c0370708324c122384420 100644 --- a/tensorflow/contrib/lite/toco/args.h +++ b/tensorflow/contrib/lite/toco/args.h @@ -229,6 +229,7 @@ struct ParsedTocoFlags { // Deprecated flags Arg input_type; Arg input_types; + Arg debug_disable_recurrent_cell_fusion = Arg(false); Arg drop_control_dependency = Arg(false); }; diff --git a/tensorflow/contrib/lite/toco/dump_graphviz.cc b/tensorflow/contrib/lite/toco/dump_graphviz.cc index c726eb6d8678e2703f5acba8b3d8d740186939f5..c8352741b44cd627ff9edb9c4677b994c4cb9a09 100644 --- a/tensorflow/contrib/lite/toco/dump_graphviz.cc +++ b/tensorflow/contrib/lite/toco/dump_graphviz.cc @@ -142,14 +142,8 @@ NodeProperties GetPropertiesForArray(const Model& model, // Append array shape to the label. auto& array = model.GetArray(array_name); - - if (array.data_type == ArrayDataType::kFloat) { - AppendF(&node_properties.label, "\\nType: float"); - } else if (array.data_type == ArrayDataType::kInt32) { - AppendF(&node_properties.label, "\\nType: int32"); - } else if (array.data_type == ArrayDataType::kUint8) { - AppendF(&node_properties.label, "\\nType: uint8"); - } + AppendF(&node_properties.label, "\\nType: %s", + ArrayDataTypeName(array.data_type)); if (array.has_shape()) { auto& array_shape = array.shape(); @@ -199,12 +193,12 @@ NodeProperties GetPropertiesForArray(const Model& model, } if (array.minmax) { - AppendF(&node_properties.label, "\\nMinMax: [%.3g, %.3g]", + AppendF(&node_properties.label, "\\nMinMax: [%.7g, %.7g]", array.minmax->min, array.minmax->max); } if (array.quantization_params) { - AppendF(&node_properties.label, "\\nQuantization: %.3g * (x - %d)", + AppendF(&node_properties.label, "\\nQuantization: %7g * (x - %d)", array.quantization_params->scale, array.quantization_params->zero_point); } diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.cc b/tensorflow/contrib/lite/toco/export_tensorflow.cc index 529df3cd2e56f1888f3d431ddcd7dc7051a98355..22a23357b36c16ea937e726f1e49aa95d7f964e3 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. @@ -212,6 +238,30 @@ void ConvertIntTensorConst(const Model& model, const string& name, } } +void CreateIntTensorConst(const string& name, const std::vector& data, + const std::vector& shape, + 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* tensor_shape = tensor->mutable_tensor_shape(); + int num_elements = 1; + for (int size : shape) { + tensor_shape->add_dim()->set_size(size); + num_elements *= size; + } + CHECK_EQ(num_elements, data.size()); +} + void CreateMatrixShapeTensorConst(const string& name, int rows, int cols, GraphDef* tensorflow_graph) { if (HasAlreadyExportedConst(name, *tensorflow_graph)) { @@ -341,6 +391,84 @@ void ConvertConvOperator(const Model& model, const ConvOperator& src_op, } } +void ConvertDilatedConvOperator(const Model& model, const ConvOperator& src_op, + GraphDef* tensorflow_graph) { + CHECK((src_op.dilation_width_factor > 1) || + (src_op.dilation_height_factor > 1)) + << "Conv operator must have height or width dilation factor > 1. " + "Otherwise, use regular conv op."; + CHECK_EQ(src_op.stride_width, 1) + << "Dilated AND strided convolution is unsupported"; + CHECK_EQ(src_op.stride_height, 1) + << "Dilated AND strided convolution is unsupported"; + + // Emulate dilated convolution with a chain of SpaceToBatchND -> Conv -> + // BatchToSpaceND ops. + + // Compute padding + const auto& input_array = model.GetArray(src_op.inputs[0]); + const auto& input_shape = input_array.shape(); + CHECK_EQ(input_shape.dimensions_count(), 4); + int height_mod_dilation = input_shape.dims(1) % src_op.dilation_height_factor; + int pad_height; + if (height_mod_dilation) { + pad_height = src_op.dilation_height_factor - height_mod_dilation; + } else { + pad_height = 0; + } + int pad_width; + int width_mod_dilation = input_shape.dims(2) % src_op.dilation_width_factor; + if (width_mod_dilation) { + pad_width = src_op.dilation_width_factor - width_mod_dilation; + } else { + pad_width = 0; + } + + // SpaceToBatchND op "collapses" the spatially separated elements together + string stb_output = src_op.outputs[0] + "/dilated_conv_SpaceToBatch"; + auto* stb_op = tensorflow_graph->add_node(); + stb_op->set_op("SpaceToBatchND"); + stb_op->set_name(stb_output); + *stb_op->add_input() = src_op.inputs[0]; + (*stb_op->mutable_attr())["T"].set_type(DT_FLOAT); + string block_shape = src_op.outputs[0] + "/dilated_conv_block_shape"; + CreateIntTensorConst( + block_shape, + {src_op.dilation_height_factor, src_op.dilation_width_factor}, {2}, + tensorflow_graph); + *stb_op->add_input() = block_shape; + (*stb_op->mutable_attr())["Tblock_shape"].set_type(DT_INT32); + string stb_paddings = src_op.outputs[0] + "/dilated_conv_paddings"; + CreateIntTensorConst(stb_paddings, {0, pad_height, pad_width, 0}, {2, 2}, + tensorflow_graph); + *stb_op->add_input() = stb_paddings; + (*stb_op->mutable_attr())["Tpaddings"].set_type(DT_INT32); + + // Perform a regular conv on the "collapsed" elements + ConvOperator conv_op; + string conv_output = src_op.outputs[0] + "/dilated_conv_Conv2D"; + conv_op.inputs = src_op.inputs; + conv_op.inputs[0] = stb_output; + conv_op.outputs = {conv_output}; + conv_op.padding.type = src_op.padding.type; + conv_op.stride_width = src_op.stride_width; + conv_op.stride_height = src_op.stride_height; + conv_op.dilation_width_factor = 1; + conv_op.dilation_height_factor = 1; + ConvertConvOperator(model, conv_op, tensorflow_graph); + + // BatchToSpaceND op restores elements to their original layout + auto* bts_op = tensorflow_graph->add_node(); + bts_op->set_op("BatchToSpaceND"); + bts_op->set_name(src_op.outputs[0]); + *bts_op->add_input() = conv_output; + (*bts_op->mutable_attr())["T"].set_type(DT_FLOAT); + *bts_op->add_input() = block_shape; + (*bts_op->mutable_attr())["Tblock_shape"].set_type(DT_INT32); + *bts_op->add_input() = stb_paddings; + (*bts_op->mutable_attr())["Tcrops"].set_type(DT_INT32); +} + void ConvertDepthwiseConvOperator(const Model& model, const DepthwiseConvOperator& src_op, GraphDef* tensorflow_graph) { @@ -420,6 +548,38 @@ void ConvertDepthwiseConvOperator(const Model& model, } } +void ConvertTransposeConvOperator(const Model& model, + const TransposeConvOperator& src_op, + GraphDef* tensorflow_graph) { + auto* conv2d_op = tensorflow_graph->add_node(); + conv2d_op->set_op("Conv2DBackpropInput"); + conv2d_op->set_name(src_op.outputs[0]); + *conv2d_op->add_input() = src_op.inputs[0]; + *conv2d_op->add_input() = src_op.inputs[1]; + *conv2d_op->add_input() = src_op.inputs[2]; + (*conv2d_op->mutable_attr())["T"].set_type(DT_FLOAT); + const string& weights_array_name = WalkUpToConstantArray( + model, src_op.inputs[TransposeConvOperator::WEIGHTS]); + const auto& weights_array = model.GetArray(weights_array_name); + CHECK(weights_array.buffer->type == ArrayDataType::kFloat); + ConvertFloatTensorConst(model, weights_array_name, AxesOrder::kOHWI, + AxesOrder::kHWIO, tensorflow_graph); + auto& strides = (*conv2d_op->mutable_attr())["strides"]; + strides.mutable_list()->add_i(1); + strides.mutable_list()->add_i(src_op.stride_height); + strides.mutable_list()->add_i(src_op.stride_width); + strides.mutable_list()->add_i(1); + string padding; + if (src_op.padding.type == PaddingType::kSame) { + padding = "SAME"; + } else if (src_op.padding.type == PaddingType::kValid) { + padding = "VALID"; + } else { + LOG(FATAL) << "Bad padding (only SAME and VALID are supported)"; + } + (*conv2d_op->mutable_attr())["padding"].set_s(padding); +} + void ConvertDepthToSpaceOperator(const Model& model, const DepthToSpaceOperator& src_op, GraphDef* tensorflow_graph) { @@ -445,14 +605,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,38 +629,43 @@ 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}, {2}, 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.HasArray(src_op.inputs[1])); - const string& fc_weights_name = - WalkUpToConstantArray(model, src_op.inputs[1]); - const auto& fc_weights_array = model.GetArray(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); + // 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); + (*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. @@ -621,7 +795,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 @@ -656,6 +831,46 @@ 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 != 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 + // 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"; @@ -798,7 +1013,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()); } @@ -812,7 +1028,8 @@ 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]); QCHECK(shape_array.data_type == ArrayDataType::kInt32) << "Only int32 shape is supported."; @@ -909,24 +1126,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(); @@ -981,6 +1180,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) { @@ -991,6 +1297,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 { @@ -1046,8 +1353,9 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op, // Write weights const string weights_output = base + "weights"; CHECK(model.HasArray(src_op.inputs[LstmCellOperator::WEIGHTS_INPUT])); - const auto& weights_array = - model.GetArray(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); @@ -1072,8 +1380,9 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op, // Write biases const string biases_output = base + "biases"; CHECK(model.HasArray(src_op.inputs[LstmCellOperator::BIASES_INPUT])); - const auto& bias_array = - model.GetArray(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); @@ -1345,9 +1654,11 @@ void ConvertSqueezeOperator(const Model& model, const SqueezeOperator& src_op, const auto params_type = GetTensorFlowDataType(model, src_op.inputs[0]); (*new_op->mutable_attr())["T"].set_type(params_type); - auto& squeeze_dims = (*new_op->mutable_attr())["squeeze_dims"]; - for (int i : src_op.squeeze_dims) { - squeeze_dims.mutable_list()->add_i(i); + if (!src_op.squeeze_dims.empty()) { + auto& squeeze_dims = (*new_op->mutable_attr())["squeeze_dims"]; + for (int i : src_op.squeeze_dims) { + squeeze_dims.mutable_list()->add_i(i); + } } } @@ -1389,6 +1700,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) { @@ -1397,8 +1719,13 @@ void ConvertOperator(const Model& model, const Operator& src_op, } if (src_op.type == OperatorType::kConv) { - ConvertConvOperator(model, static_cast(src_op), - tensorflow_graph); + const ConvOperator& conv_op = static_cast(src_op); + if ((conv_op.dilation_width_factor != 1) || + (conv_op.dilation_height_factor != 1)) { + return ConvertDilatedConvOperator(model, conv_op, tensorflow_graph); + } else { + ConvertConvOperator(model, conv_op, tensorflow_graph); + } } else if (src_op.type == OperatorType::kDepthwiseConv) { ConvertDepthwiseConvOperator( model, static_cast(src_op), @@ -1445,6 +1772,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), @@ -1533,6 +1864,39 @@ 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 if (src_op.type == OperatorType::kTransposeConv) { + ConvertTransposeConvOperator( + model, static_cast(src_op), + tensorflow_graph); } else { LOG(FATAL) << "Unhandled operator type " << OperatorTypeName(src_op.type); } @@ -1622,6 +1986,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 79682153a8fd143c4934095567764b886bd776af..d7310bb75f258cde25236da2a9269f18234784e4 100644 --- a/tensorflow/contrib/lite/toco/export_tensorflow.h +++ b/tensorflow/contrib/lite/toco/export_tensorflow.h @@ -22,6 +22,8 @@ namespace toco { void ExportTensorFlowGraphDef(const Model& model, string* output_file_contents); +void EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(Model* model); + } // namespace toco #endif // TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_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/graph_transformations/convert_squeeze_to_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_squeeze_to_reshape.cc new file mode 100644 index 0000000000000000000000000000000000000000..81cedb5dad751aacbbb32326db73de386aba282d --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_squeeze_to_reshape.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 +#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 { + +// Replaces a tf.squeeze operator with a reshape. +// Squeeze removes dimensions == 1 (if in the list of squeeze_dims). This +// means that the data layout will never change with this op, just the shape. +// By converting these to reshapes once we have run shape propagation we allow +// standard reshape optimization transforms to do their magic. +bool ConvertSqueezeToReshape::Run(Model* model, std::size_t op_index) { + auto squeeze_it = model->operators.begin() + op_index; + if (squeeze_it->get()->type != OperatorType::kSqueeze) { + return false; + } + auto squeeze_op = static_cast(squeeze_it->get()); + CHECK_EQ(squeeze_op->inputs.size(), 1); + CHECK_EQ(squeeze_op->outputs.size(), 1); + + const auto& input_array = model->GetArray(squeeze_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. + return false; + } + if (!model->HasArray(squeeze_op->outputs[0]) || + !model->GetArray(squeeze_op->outputs[0]).has_shape()) { + // Yield until shape propagation has set the output shape for us. + return false; + } + + // We use the output shape that has been calculated by shape propagation. + const auto& output_shape = model->GetArray(squeeze_op->outputs[0]).shape(); + + // Empty shapes will not work as empty data arrays. + if (output_shape.dimensions_count() == 0) { + return false; + } + + auto* reshape_op = new TensorFlowReshapeOperator; + reshape_op->inputs = { + squeeze_op->inputs[0], + CreateInt32Array(model, squeeze_op->outputs[0] + "_shape", + output_shape.dims()), + }; + reshape_op->outputs = squeeze_op->outputs; + + AddMessageF("Replacing %s with %s", LogName(*squeeze_op), + LogName(*reshape_op)); + + // Replace the operator in the graph. + const auto reshape_it = model->operators.emplace(squeeze_it, reshape_op); + squeeze_it = reshape_it + 1; + CHECK_EQ(squeeze_it->get(), squeeze_op); + model->operators.erase(squeeze_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 index c2b166033c33b777bad88cb712adf8517be1762a..5a36a90b3841504d6f018832777e50bac95218d7 100644 --- 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 @@ -21,6 +21,33 @@ limitations under the License. namespace toco { +namespace { + +bool TransposeAffectsMemoryOrder(std::vector perm, + std::vector in_shape) { + CHECK_EQ(perm.size(), in_shape.size()); + // See what the ordering of the non-unary columns are before and after + // transpose permutation. If the major indices stay in the same order (not + // just the shape) then the flat buffer representation shouldn't change. + std::vector old_major_index_ordering; + std::vector new_major_index_ordering; + for (int i = 0; i < in_shape.size(); i++) { + if (in_shape[i] != 1) { + old_major_index_ordering.push_back(i); + } + + if (in_shape[perm[i]] != 1) { + new_major_index_ordering.push_back(perm[i]); + } + } + + CHECK_EQ(new_major_index_ordering.size(), old_major_index_ordering.size()); + + return old_major_index_ordering != new_major_index_ordering; +} + +} // namespace + 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) { @@ -29,23 +56,26 @@ bool ConvertTrivialTransposeToReshape::Run(Model* model, std::size_t op_index) { TransposeOperator* transpose_op = static_cast(transpose_it->get()); + const auto& input_array = model->GetArray(transpose_op->inputs[0]); const auto& output_array = model->GetArray(transpose_op->outputs[0]); - if (!output_array.has_shape()) { + if (!input_array.has_shape() || !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++; - } + // Check that the permutation has propogated. + std::vector const& perm = transpose_op->perm; + if (perm.empty()) { + return false; } - if (non_unitary_axis_count > 1) { - // Transpose is not trivial + + // This transpose is trivial if non-unitary dimensions remain in the same + // order. + std::vector const& input_dims = input_array.shape().dims(); + std::vector const& output_dims = output_array.shape().dims(); + + if (TransposeAffectsMemoryOrder(perm, input_dims)) { return false; } @@ -61,11 +91,11 @@ bool ConvertTrivialTransposeToReshape::Run(Model* model, std::size_t op_index) { 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())}; + 1, static_cast(output_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; + shape_buffer.data = output_dims; // Delete perm array if unused if (IsDiscardableArray(*model, perm_array_name) && 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 ab943f72d1dd87ae9ff4bd53a807cd4923a88c38..c5ce3fcd95eb0aaf63dcc7f43b96d8a13ed93929 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/fuse_activation_functions.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/fuse_activation_functions.cc @@ -42,9 +42,9 @@ bool FuseActivationFunctions::Run(Model* model, std::size_t op_index) { if (CountTrueOutputs(*model, *op) > 1) { AddMessageF( - "Not fusing activation function into %s because it has more than one " - " consumed output", - LogName(*op)); + "Not fusing activation function %s into %s because it has more than " + "one consumed output", + LogName(*ac_op), LogName(*op)); return false; } @@ -56,22 +56,31 @@ bool FuseActivationFunctions::Run(Model* model, std::size_t op_index) { AddMessageF( "Not fusing activation function into %s because it is consumed by more " "than 1 other operator", - LogName(*op)); + LogName(*ac_op), LogName(*op)); + return false; + } + + if (!IsDiscardableArray(*model, op->outputs[0])) { + AddMessageF( + "Not fusing activation function %s into %s because output %s it is not " + "discardable", + LogName(*ac_op), LogName(*op), op->outputs[0]); return false; } if (op->fused_activation_function != FusedActivationFunctionType::kNone) { AddMessageF( - "Not fusing activation function into %s because it already has a fused " - "activation function", - LogName(*op)); + "Not fusing activation function %s into %s because it already has a " + "fused activation function", + LogName(*ac_op), LogName(*op)); return false; } if (!OperatorSupportsFusedActivation(op->type)) { AddMessageF( - "Not fusing activation function because the %s op doesn't support it", - LogName(*op)); + "Not fusing activation function %s because the %s op doesn't support " + "it", + LogName(*ac_op), LogName(*op)); return false; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h index cf90ebe99697fe8a40b4c707e70fdc5318123854..11e5e19f508ecfb2131b391af9b4e51bb9cf0fb4 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h +++ b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h @@ -114,7 +114,9 @@ void RunGraphTransformations(Model* model, const string& message, // List of all graph transformations DECLARE_GRAPH_TRANSFORMATION(ConvertExpandDimsToReshape) DECLARE_GRAPH_TRANSFORMATION(ConvertPureConvToDepthwise) +DECLARE_GRAPH_TRANSFORMATION(ConvertSqueezeToReshape) DECLARE_GRAPH_TRANSFORMATION(ConvertTrivialAddNToAdd) +DECLARE_GRAPH_TRANSFORMATION(ConvertTrivialStackToReshape) DECLARE_GRAPH_TRANSFORMATION(ConvertTrivialTransposeToReshape) DECLARE_GRAPH_TRANSFORMATION(ConvertReorderAxes) DECLARE_GRAPH_TRANSFORMATION(EnsureBiasVectors) @@ -124,8 +126,12 @@ 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(IdentifyDilatedConv) DECLARE_GRAPH_TRANSFORMATION(MakeInitialDequantizeOperator) +DECLARE_GRAPH_TRANSFORMATION(PropagateActivationFunctionIntoConstants) DECLARE_GRAPH_TRANSFORMATION(PropagateArrayDataTypes) DECLARE_GRAPH_TRANSFORMATION(PropagateFixedSizes) DECLARE_GRAPH_TRANSFORMATION(HardcodeMinMax) @@ -136,6 +142,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) @@ -154,8 +161,10 @@ 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) @@ -168,7 +177,10 @@ DECLARE_GRAPH_TRANSFORMATION(ResolveConstantShapeOrRank) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantStack) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantStridedSlice) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantFill) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantGather) +DECLARE_GRAPH_TRANSFORMATION(ResolveMultiplyByZero) DECLARE_GRAPH_TRANSFORMATION(Dequantize) +DECLARE_GRAPH_TRANSFORMATION(UnpartitionEmbeddingLookup) class ResolveReshapeAttributes : public GraphTransformation { public: 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..5cc82da5d544846cc095046ceccf0664525aae41 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/hardcode_min_max.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/hardcode_min_max.cc @@ -125,6 +125,27 @@ bool HardcodeMinMaxForConcatenation(Model* model, Operator* op) { return changed; } +bool HardcodeMinMaxForSplit(Model* model, Operator* op) { + for (const auto& output : op->outputs) { + if (model->GetArray(output).minmax) { + LOG(WARNING) << "Skipping min-max setting for " << LogName(*op) + << " because output " << output << " already has min-max."; + return false; + } + } + // Data is in second input. + auto& input_array = model->GetArray(op->inputs[1]); + if (!input_array.minmax) { + return false; + } else { + for (const auto& output : op->outputs) { + auto& array = model->GetArray(output); + array.GetOrCreateMinMax() = *input_array.minmax; + } + return true; + } +} + // The output of average or max pooling is within the same range as its input. bool HardcodeMinMaxForAverageOrMaxPool(Model* model, Operator* op) { auto& output_array = model->GetArray(op->outputs[0]); @@ -177,6 +198,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) { @@ -196,14 +317,21 @@ bool HardcodeMinMax::Run(Model* model, std::size_t op_index) { changed = HardcodeMinMaxForConcatenation(model, op); break; + case OperatorType::kTensorFlowSplit: + changed = HardcodeMinMaxForSplit(model, op); + break; + case OperatorType::kAveragePool: case OperatorType::kMaxPool: changed = HardcodeMinMaxForAverageOrMaxPool(model, op); break; + case OperatorType::kStridedSlice: case OperatorType::kSqueeze: case OperatorType::kTensorFlowReshape: case OperatorType::kPad: + case OperatorType::kGather: + case OperatorType::kTranspose: changed = HardcodeMinMaxFromFirstInput(model, op); break; @@ -219,6 +347,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_dilated_conv.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_dilated_conv.cc new file mode 100644 index 0000000000000000000000000000000000000000..ae3301f467de5714230e731b4bab87ddc1637201 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_dilated_conv.cc @@ -0,0 +1,213 @@ +/* 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/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 { + +// A dilated convolution can be emulated with a regular convolution by chaining +// SpaceToBatch and BatchToSpace ops before and after it: +// +// SpaceToBatchND -> Conv2D -> BatchToSpaceND +// +// This method was common before Conv2D fully supported dilated convolution in +// TensorFlow. This transformation detects this "emulation", and replaces it +// with a true dilated convolution, eliminating the SpaceToBatch and +// BatchtoSpace ops. +// +// Detecting this alone would be relatively easy. However, in practice some +// extra ops are used, so we detect the following patterns: +// +// +// SpaceToBatchND -> Expand -> Conv2D -> Squeeze -> BatchToSpaceND -> BiasAdd +// +// SpaceToBatchND -> Expand -> Conv2D -> Squeeze -> Pad -> BatchToSpaceND -> +// BiasAdd +// +// SpaceToBatchND -> Expand -> Conv2D -> Squeeze -> BiasAdd -> BatchToSpaceND +// +// SpaceToBatchND -> Conv2D -> Pad -> BatchToSpaceND -> BiasAdd +// +// SpaceToBatchND -> Conv2D -> BatchToSpaceND -> BiasAdd +// +// +// The Expand/Squeeze combination is used to adapt a 3D array (such as in +// WaveNet) to the 4D arrays that Conv2D requires. Padding and BiasAdd are +// thrown in just for the extra headache. Padding adapts non-conforming input +// sizes, and can be discarded. The bias is necessary, so is kept. + +bool IdentifyDilatedConv::Run(Model* model, std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + auto* stb_op = it->get(); + + // 1. IDENTIFY OPERATORS + // *************************************************************************** + // SpaceToBatch Op. + if (stb_op->type != OperatorType::kSpaceToBatchND) { + return false; + } + if (stb_op->inputs.size() != 3) { + return false; + } + CHECK_EQ(stb_op->outputs.size(), 1); + // Extract the dilation factor from Input[1] of SpaceToBatch + // TODO(mjmatthews): Support 2D dilation factors. + const auto& block_shape_array = model->GetArray(stb_op->inputs[1]); + if (!block_shape_array.buffer) { + return false; + } + CHECK_EQ(block_shape_array.shape().dimensions_count(), 1); + int dilation_factor = + block_shape_array.Array::GetBuffer().data[0]; + + // Expand Op + auto* post_stb_op = GetOpWithInput(*model, stb_op->outputs[0]); + if (!post_stb_op) { + return false; + } + bool has_expand_op = false; + if (post_stb_op->type == OperatorType::kExpandDims) { + has_expand_op = true; + CHECK_EQ(post_stb_op->inputs.size(), 2); + CHECK_EQ(post_stb_op->outputs.size(), 1); + } + + // Conv Op + ConvOperator* conv_op = dynamic_cast( + has_expand_op ? GetOpWithInput(*model, post_stb_op->outputs[0]) + : GetOpWithInput(*model, stb_op->outputs[0])); + if (!conv_op || conv_op->type != OperatorType::kConv) { + return false; + } + if (conv_op->inputs.size() != 2) { + // The conv op must only have weights, no bias. + return false; + } + CHECK_EQ(conv_op->outputs.size(), 1); + + // Squeeze Op + auto* post_conv_op = GetOpWithInput(*model, conv_op->outputs[0]); + if (!post_conv_op) { + return false; + } + if (has_expand_op) { + if (post_conv_op->type != OperatorType::kSqueeze) { + // If an expand op was used, the post-conv op must be a squeeze op + return false; + } + CHECK_EQ(post_conv_op->inputs.size(), 1); + CHECK_EQ(post_conv_op->outputs.size(), 1); + } + + // Pad Op + const auto* pad_op = has_expand_op + ? GetOpWithInput(*model, post_conv_op->outputs[0]) + : GetOpWithInput(*model, conv_op->outputs[0]); + bool has_pad_op = false; + if (pad_op->type == OperatorType::kPad) { + has_pad_op = true; + CHECK_EQ(pad_op->inputs.size(), 2); + CHECK_EQ(pad_op->outputs.size(), 1); + } + // TODO(mjmatthews): Perform validity checking on padding dimensions. + + // Pre-BatchToSpace Bias Op + auto* next_op = has_pad_op + ? GetOpWithInput(*model, pad_op->outputs[0]) + : has_expand_op + ? GetOpWithInput(*model, post_conv_op->outputs[0]) + : GetOpWithInput(*model, conv_op->outputs[0]); + bool has_bias_before_bts = false; + if (next_op->type == OperatorType::kAdd) { + has_bias_before_bts = true; + } + auto final_op = GetOpWithInput(*model, next_op->outputs[0]); + + // BatchToSpace Op + const auto* bts_op = has_bias_before_bts ? final_op : next_op; + if (bts_op->type != OperatorType::kBatchToSpaceND) { + return false; + } + CHECK_EQ(bts_op->inputs.size(), 3); + CHECK_EQ(bts_op->outputs.size(), 1); + + // Post-BatchToSpace Bias Op + Operator* bias_add_op = !has_bias_before_bts ? final_op : next_op; + if (bias_add_op->type != OperatorType::kAdd) { + // Bias op is required before or after BatchToSpace + return false; + } + CHECK_EQ(bias_add_op->inputs.size(), 2); + CHECK_EQ(bias_add_op->outputs.size(), 1); + + LOG(INFO) << "Identified sub-network emulating dilated convolution."; + + // 2. RE-WIRE OPERATORS + // *************************************************************************** + // Re-use the existing Conv2D op. + conv_op->dilation_width_factor = dilation_factor; + conv_op->dilation_height_factor = dilation_factor; + conv_op->padding.type = PaddingType::kSame; + + // Rewire the ops to bypass SpaceToBatch, BatchToSpace, and Pad. + bias_add_op->outputs[0] = final_op->outputs[0]; + if (has_expand_op) { + bias_add_op->inputs[0] = post_conv_op->outputs[0]; + post_conv_op->inputs[0] = conv_op->outputs[0]; + conv_op->inputs[0] = post_stb_op->outputs[0]; + post_stb_op->inputs[0] = stb_op->inputs[0]; + } else { + bias_add_op->inputs[0] = conv_op->outputs[0]; + conv_op->inputs[0] = stb_op->inputs[0]; + } + // TODO(mjmatthews): Connect bias directly into the Conv2D? + + // 3. DELETE LEFTOVER OPERATORS + // *************************************************************************** + // Order is important. Delete the output array first, then the op, then it's + // redundant inputs. + // BatchToSpace Op + DeleteArrayIfUnused(bts_op->outputs[0], model); + std::vector bts_op_inputs = bts_op->inputs; + model->operators.erase(FindOp(*model, bts_op)); + DeleteArrayIfUnused(bts_op_inputs[1], model); + DeleteArrayIfUnused(bts_op_inputs[2], model); + + // Pad Op if present + if (has_pad_op) { + DeleteArrayIfUnused(pad_op->outputs[0], model); + std::vector pad_op_inputs = pad_op->inputs; + model->operators.erase(FindOp(*model, pad_op)); + DeleteArrayIfUnused(pad_op_inputs[1], model); + } + + // SpaceToBatch Op + DeleteArrayIfUnused(stb_op->outputs[0], model); + std::vector stb_op_inputs = stb_op->inputs; + model->operators.erase(FindOp(*model, stb_op)); + DeleteArrayIfUnused(stb_op_inputs[1], model); + DeleteArrayIfUnused(stb_op_inputs[2], model); + + LOG(INFO) << "Replaced with Dilated Conv2D op outputting \"" + << conv_op->outputs[0] << "\"."; + return true; +} + +} // namespace toco 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 d36e95060937d6af0789766bcb29ae70cef4569d..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->EraseArray(maximum_op->inputs[scalar_input_index]); - model->operators.erase(FindOperator(model, maximum_op)); - - // Erase Minimum inputs & operator - model->EraseArray(minimum_op->inputs[0]); - model->EraseArray(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/make_initial_dequantize_operator.cc b/tensorflow/contrib/lite/toco/graph_transformations/make_initial_dequantize_operator.cc index d83603e9a2c59ae74a5e5fda5b11178740336bfb..935da9f966ca63095faa17476be3a559d1a0193a 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/make_initial_dequantize_operator.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/make_initial_dequantize_operator.cc @@ -85,8 +85,8 @@ bool AddDequantizeOperatorToInput(const string& input_name, const Operator* op, auto& dequantized_input_minmax = dequantized_input_array.GetOrCreateMinMax(); dequantized_input_minmax = input_minmax; auto& input_qparams = input_array.GetOrCreateQuantizationParams(); - GetQuantizationParamsFromMinMax( - model->flags, input_minmax, &input_qparams); + GetQuantizationParamsFromMinMax(input_minmax, + &input_qparams); transformation->AddMessageF( "Created %s" diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_activation_function_into_constants.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_activation_function_into_constants.cc new file mode 100644 index 0000000000000000000000000000000000000000..cf17c49b1098d02468935aa72d1d1e73b4addbe1 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_activation_function_into_constants.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 +#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/runtime/types.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool PropagateActivationFunctionIntoConstants::Run(Model* model, + std::size_t op_index) { + const auto ac_it = model->operators.begin() + op_index; + const auto* ac_op = ac_it->get(); + if (ac_op->type != OperatorType::kRelu6 && + ac_op->type != OperatorType::kRelu1 && + ac_op->type != OperatorType::kRelu) { + return false; + } + + // Find the op producing the array passed to this activation function. + auto* src_op = GetOpWithOutput(*model, ac_op->inputs[0]); + if (!src_op) { + return false; + } + + // Ensure the src_op is not used without the activation function applied. + if (CountTrueOutputs(*model, *src_op) > 1) { + AddMessageF( + "Not propagating activation function %s into %s because it has more " + "than one consumed output", + LogName(*ac_op), LogName(*src_op)); + } + + // Filter to the list of supported ops. + string src_op_input; + switch (src_op->type) { + case OperatorType::kGather: + src_op_input = src_op->inputs[0]; + break; + default: + return false; + } + CHECK_EQ(src_op->outputs[0], ac_op->inputs[0]); + + // Ensure the input is constant as otherwise this needs to happen at runtime. + // If we bail here, it's still possible that FuseActivationFunctions will fuse + // the activation if it's supported by the op. + if (!IsConstantParameterArray(*model, src_op_input)) { + AddMessageF( + "Not propagating activation function %s into %s:%s because it is not " + "constant", + LogName(*ac_op), LogName(*src_op), src_op_input); + return false; + } + + // Get the array we'll be working with and ensure it's a compatible type. + auto& const_array = model->GetArray(src_op_input); + if (const_array.data_type != ArrayDataType::kFloat) { + AddMessageF( + "Not propagating activation function %s into %s:%s because it is " + "non-float data", + LogName(*ac_op), LogName(*src_op), src_op_input); + return false; + } + auto& const_array_data = + const_array.GetMutableBuffer().data; + + // Perform the activation function directly into the constant data array. + for (size_t i = 0; i < const_array_data.size(); ++i) { + const float value = const_array_data[i]; + float new_value = value; + switch (ac_op->type) { + case OperatorType::kRelu: { + static constexpr float kLower = 0; + new_value = value < kLower ? kLower : value; + break; + } + case OperatorType::kRelu1: { + static constexpr float kUpper = 1; + static constexpr float kLower = -1; + new_value = value > kUpper ? kUpper : value < kLower ? kLower : value; + break; + } + case OperatorType::kRelu6: { + static constexpr float kUpper = 6; + static constexpr float kLower = 0; + new_value = value > kUpper ? kUpper : value < kLower ? kLower : value; + break; + } + default: + LOG(FATAL) << "Unsupported activation function " << LogName(*ac_op); + return false; + } + const_array_data[i] = new_value; + } + + AddMessageF("Propagated activation function %s into %s:%s", LogName(*ac_op), + LogName(*src_op), src_op_input); + return RemoveTrivialPassthroughOp(this, model, op_index); +} + +} // 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 f0d107232b4517115aa3f64b39b825dbaffb83ce..778da39bf13563cbbdbe54f1140595b057253ae3 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 @@ -71,6 +71,11 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { CHECK_GE(op->inputs.size(), 2); const ArrayDataType data_type = model->GetArray(op->inputs[1]).data_type; SetDataTypeForAllOutputs(model, op, data_type); + } else if (op->type == OperatorType::kTransposeConv) { + // These operators produce an output with the same type as their 3rd input + CHECK_GE(op->inputs.size(), 3); + const ArrayDataType data_type = model->GetArray(op->inputs[2]).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); @@ -97,10 +102,13 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { 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()) { + // Some output tensors from the op could be eliminated by optimization. + // This can make unsupported_op->output_data_types have more elements than + // op->outputs. + if (unsupported_op->output_data_types.size() < op->outputs.size()) { return false; } - for (int i = 0; i < unsupported_op->output_data_types.size(); ++i) { + for (int i = 0; i < op->outputs.size(); ++i) { auto output = op->outputs[i]; auto data_type = unsupported_op->output_data_types[i]; model->GetArray(output).data_type = data_type; 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 4fb3b6ae7a5fc5bfc2719b978331c67ae799eb54..375848a7d408b58eab8692a15376b5b2e04b34af 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc @@ -31,17 +31,22 @@ namespace { void ComputeConvSizes(const Shape& input_shape, int output_depth, int kwidth, int kheight, int stride_width, int stride_height, + int dilation_width_factor, int dilation_height_factor, PaddingType padding_type, Shape* output_shape, FixedPadding* fixed_padding) { const int input_width = input_shape.dims(2); const int input_height = input_shape.dims(1); const int batch = input_shape.dims(0); + int dilated_kwidth = dilation_width_factor * (kwidth - 1) + 1; + int dilated_kheight = dilation_height_factor * (kheight - 1) + 1; + int output_height = 0; int output_width = 0; if (padding_type == PaddingType::kValid) { - output_height = (input_height + stride_height - kheight) / stride_height; - output_width = (input_width + stride_width - kwidth) / stride_width; + output_height = + (input_height + stride_height - dilated_kheight) / stride_height; + output_width = (input_width + stride_width - dilated_kwidth) / stride_width; } else if (padding_type == PaddingType::kSame) { output_height = (input_height + stride_height - 1) / stride_height; output_width = (input_width + stride_width - 1) / stride_width; @@ -49,10 +54,12 @@ void ComputeConvSizes(const Shape& input_shape, int output_depth, int kwidth, LOG(FATAL) << "Only supporting SAME or VALID padding"; } - fixed_padding->height = std::max( - 0, ((output_height - 1) * stride_height + kheight - input_height) / 2); + fixed_padding->height = std::max(0, ((output_height - 1) * stride_height + + dilated_kheight - input_height) / + 2); fixed_padding->width = std::max( - 0, ((output_width - 1) * stride_width + kwidth - input_width) / 2); + 0, + ((output_width - 1) * stride_width + dilated_kwidth - input_width) / 2); // Actually had to debug a situation where those were negative due to bad // propagation of placeholder -1 sizes in TensorFlowReshape. @@ -61,23 +68,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()); @@ -147,7 +173,8 @@ void ProcessConvOperator(Model* model, ConvOperator* op) { const int kheight = weights_shape.dims(1); const int kwidth = weights_shape.dims(2); ComputeConvSizes(input_shape, output_depth, kwidth, kheight, op->stride_width, - op->stride_height, op->padding.type, + op->stride_height, op->dilation_width_factor, + op->dilation_height_factor, op->padding.type, output_array.mutable_shape(), &op->padding.GetOrCreateFixedPadding()); CHECK_EQ(output_array.shape().dimensions_count(), 4); @@ -163,6 +190,116 @@ void ProcessConvOperator(Model* model, ConvOperator* op) { } } +void ProcessTransposeConvOperator(Model* model, TransposeConvOperator* op) { + // TransposeConv is unique in that it is specifically given the output shape + // as a 1D array on it's 1st input. Theoretically then, resolving the output + // shape is as easy as waiting for this input to be resolved. However, we also + // have to calculate the padding which requires the weights shape. So, we + // might as well calculate the output shape and ensure it matches the + // specified one + + // Check if we have already run. + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.has_shape()) { + return; + } + + // SPECIFIED OUTPUT SHAPE + // The below is the specified, or prescribed output shape, _given_ to the + // operator as an input. + auto& specified_output_shape_array = + model->GetArray(op->inputs[TransposeConvOperator::OUTPUT_SHAPE]); + if (!specified_output_shape_array.has_shape() || + !specified_output_shape_array.buffer) { + // Yield until the specified output shape is resolved as a constant + return; + } + + CHECK(specified_output_shape_array.data_type == ArrayDataType::kInt32) + << "TransposeConv input_dims must be int32"; + + CHECK(specified_output_shape_array.shape().dimensions_count() == 1 && + specified_output_shape_array.shape().dims(0) == 4) + << "TransposeConv requires a 1D, 4 element array on it's 0th input " + "specifying the output shape. \"" + << op->inputs[TransposeConvOperator::OUTPUT_SHAPE] << "\" had shape " + << toco::ShapeToString(specified_output_shape_array.shape()); + + // COMPUTE PADDING + // We require the weights shape to calculate padding. + const auto& weights_array = + model->GetArray(op->inputs[TransposeConvOperator::WEIGHTS]); + if (!weights_array.has_shape()) { + // Yield until weights dims have been resolved. + return; + } + const auto& weights_shape = weights_array.shape(); + CHECK_EQ(weights_shape.dimensions_count(), 4) + << "TransposeConv weights must have 4 input dimensions. Input weights \"" + << op->inputs[TransposeConvOperator::WEIGHTS] << "\" had shape " + << toco::ShapeToString(weights_shape) << "."; + + CHECK(weights_shape.dims(0) == 1 && weights_shape.dims(3) == 1) + << "TransposeConv weights dimensions must begin and end with 1. Input " + "weights \"" + << op->inputs[TransposeConvOperator::WEIGHTS] << "\" had shape " + << toco::ShapeToString(weights_shape) << "."; + + // Compute padding + const int kheight = weights_shape.dims(1); + const int kwidth = weights_shape.dims(2); + op->padding.GetOrCreateFixedPadding(); + if (op->padding.type == PaddingType::kValid) { + op->padding.fixed->height = 0; + op->padding.fixed->width = 0; + } else if (op->padding.type == PaddingType::kSame) { + op->padding.fixed->height = (kheight - 1) / 2; + op->padding.fixed->width = (kwidth - 1) / 2; + } else { + LOG(FATAL) << "TransposeConv only supports SAME or VALID padding"; + } + + // VALIDATE OUTPUT SHAPE + // Compute the output shape from the input and weights shapes to verify it + // agrees with the specified output shape. + const auto& input_array = + model->GetArray(op->inputs[TransposeConvOperator::DATA_INPUT]); + if (!input_array.has_shape()) { + // Yield until input dims have been resolved. + return; + } + const auto& input_shape = input_array.shape(); + CHECK_EQ(input_shape.dimensions_count(), 4) + << "TransposeConv input shape must have 4 dimensions. Input \"" + << op->inputs[TransposeConvOperator::WEIGHTS] << "\" had shape " + << toco::ShapeToString(weights_shape) << "."; + + // Compute output shape + const int input_width = input_shape.dims(2); + const int input_height = input_shape.dims(1); + int output_height = op->stride_height * (input_height - 1); + int output_width = op->stride_width * (input_width - 1); + if (op->padding.type == PaddingType::kValid) { + output_height += kheight; + output_width += kwidth; + } else if (op->padding.type == PaddingType::kSame) { + output_height += 1; + output_width += 1; + } + + CHECK(specified_output_shape_array.GetBuffer().data == + std::vector({input_shape.dims(0), output_height, output_width, + weights_shape.dims(3)})) + << "Specified output shape: " << ShapeToString(output_array.shape()) + << ", does not agree with shape computed from input data and weights: [" + << input_shape.dims(0) << ", " << output_height << ", " << output_width + << ", " << weights_shape.dims(3) << "]."; + + // SUCCESS: Set the op's output shape according to the specified output shape. + *(output_array.mutable_shape()->mutable_dims()) = + specified_output_shape_array.GetBuffer().data; +} + void ProcessDepthwiseConvOperator(Model* model, DepthwiseConvOperator* op) { if (!EnsureBiasVectorShape(model, op)) { return; @@ -203,7 +340,7 @@ void ProcessDepthwiseConvOperator(Model* model, DepthwiseConvOperator* op) { const int kheight = weights_shape.dims(1); const int kwidth = weights_shape.dims(2); ComputeConvSizes(input_shape, output_depth, kwidth, kheight, op->stride_width, - op->stride_height, op->padding.type, + op->stride_height, 1, 1, op->padding.type, model->GetArray(output_name).mutable_shape(), &op->padding.GetOrCreateFixedPadding()); } @@ -546,6 +683,9 @@ void ProcessConcatenationOperator(Model* model, ConcatenationOperator* op) { // Use 0 input as basis for output dimensions. 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; @@ -558,14 +698,14 @@ 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) { @@ -628,15 +768,34 @@ void ProcessTensorFlowSplitOperator(Model* model, TensorFlowSplitOperator* op) { } 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) { @@ -656,7 +815,7 @@ void ProcessAveragePoolOperator(Model* model, AveragePoolOperator* op) { const string& output_name = op->outputs[0]; const int output_depth = input_shape.dims(3); ComputeConvSizes(input_shape, output_depth, op->kwidth, op->kheight, - op->stride_width, op->stride_height, op->padding.type, + op->stride_width, op->stride_height, 1, 1, op->padding.type, model->GetArray(output_name).mutable_shape(), &op->padding.GetOrCreateFixedPadding()); } @@ -673,7 +832,7 @@ void ProcessMaxPoolOperator(Model* model, MaxPoolOperator* op) { const string& output_name = op->outputs[0]; const int output_depth = input_shape.dims(3); ComputeConvSizes(input_shape, output_depth, op->kwidth, op->kheight, - op->stride_width, op->stride_height, op->padding.type, + op->stride_width, op->stride_height, 1, 1, op->padding.type, model->GetArray(output_name).mutable_shape(), &op->padding.GetOrCreateFixedPadding()); } @@ -692,7 +851,7 @@ void ProcessL2PoolOperator(Model* model, L2PoolOperator* op) { const string& output_name = op->outputs[0]; const int output_depth = input_shape.dims(3); ComputeConvSizes(input_shape, output_depth, op->kwidth, op->kheight, - op->stride_width, op->stride_height, op->padding.type, + op->stride_width, op->stride_height, 1, 1, op->padding.type, model->GetArray(output_name).mutable_shape(), &op->padding.GetOrCreateFixedPadding()); } @@ -725,9 +884,8 @@ void ProcessResizeBilinearOperator(Model* model, ResizeBilinearOperator* op) { } 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->GetArray(op->inputs[LstmCellOperator::DATA_INPUT]); @@ -942,6 +1100,43 @@ 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); @@ -1120,7 +1315,8 @@ void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) { stop += input_array.shape().dims(i); } - int dim_size = (stop - start) / op->strides[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"; @@ -1214,7 +1410,8 @@ void ProcessTransposeOperator(Model* model, TransposeOperator* op) { std::vector const& perm = perm_array.GetBuffer().data; CHECK_EQ(perm.size(), input_shape.dimensions_count()) - << "Transpose permutation input must be same length as input dimensions"; + << "Transpose permutation input " << op->inputs[1] + << " 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]; @@ -1271,6 +1468,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: @@ -1284,12 +1482,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: @@ -1311,8 +1512,8 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { ProcessConvOperator(model, static_cast(op)); break; case OperatorType::kTransposeConv: - // Unimplemented, hopefully another graph transformation will drop it or - // rewrite it. + ProcessTransposeConvOperator(model, + static_cast(op)); break; case OperatorType::kDepthwiseConv: ProcessDepthwiseConvOperator(model, @@ -1420,6 +1621,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. @@ -1450,6 +1652,12 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { case OperatorType::kTranspose: ProcessTransposeOperator(model, static_cast(op)); break; + case OperatorType::kDynamicPartition: + case OperatorType::kDynamicStitch: + // DynamicPartition/DynamicStitch are currently only supported for + // transforms that remove them, so we avoid propagating shapes through + // them and let things settle once they've been removed. + break; default: // Unimplemented, another graph transformation should drop it. LOG(FATAL) << "Unhandled operator type " << OperatorTypeName(op->type); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc index b973b2b813147cc580d2e87cea7d395f180f5aa1..ad3f05274b59e3019726f1b7a7080e74d1934c89 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc @@ -41,11 +41,18 @@ 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::kStridedSlice || + type == OperatorType::kDepthToSpace || + type == OperatorType::kLstmCell || type == OperatorType::kGather || + type == OperatorType::kTranspose; } template @@ -96,10 +103,19 @@ void QuantizeArray(GraphTransformation* transformation, Model* model, void QuantizeArray(GraphTransformation* transformation, Model* model, const string& name, ArrayDataType quantized_data_type, const QuantizationParams& quantization_params) { - switch (quantized_data_type) { + ArrayDataType adjusted_data_type = quantized_data_type; + auto& array = model->GetArray(name); + if (array.final_data_type == ArrayDataType::kInt16) { + adjusted_data_type = array.final_data_type; + } + + switch (adjusted_data_type) { 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); @@ -159,6 +175,102 @@ const MinMax& GetOrComputeMinMax(Model* model, const string& array_name) { "proceed with quantization."; } +struct QuantizationPoints { + int64 min_value; + int64 max_value; + int64 central_value; +}; + +template +QuantizationPoints GetQuantizationPoints() { + QuantizationPoints qp; + using Integer = DataType; + qp.min_value = std::numeric_limits::min(); + qp.max_value = std::numeric_limits::max(); + // eg [-128,127]... + qp.central_value = (qp.min_value / 2 + // -128 -> -64. + (qp.max_value - 1) / 2 + // 127 -> 63. + 1); + return qp; +} + +QuantizationPoints GetQuantizationPoints(ArrayDataType data_type) { + switch (data_type) { + case ArrayDataType::kUint8: + return GetQuantizationPoints(); + case ArrayDataType::kInt16: + return GetQuantizationPoints(); + case ArrayDataType::kInt32: + return GetQuantizationPoints(); + default: + LOG(FATAL) << "Unhandled case."; + } +} + +ArrayDataType GetQuantizedDataType(const Array& array, + ArrayDataType default_type) { + switch (array.final_data_type) { + case ArrayDataType::kInt8: + case ArrayDataType::kUint8: + case ArrayDataType::kInt16: + case ArrayDataType::kUint16: + case ArrayDataType::kInt32: + case ArrayDataType::kUint32: + case ArrayDataType::kInt64: + case ArrayDataType::kUint64: + return array.final_data_type; + case ArrayDataType::kFloat: + case ArrayDataType::kNone: + return default_type; + default: + LOG(FATAL) << "Unhandled final quantization type " + << static_cast(array.final_data_type); + } +} + +void GetQuantizationParams(ArrayDataType data_type, const MinMax& minmax, + QuantizationParams* quantization_params) { + switch (data_type) { + case ArrayDataType::kInt8: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kUint8: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kInt16: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kUint16: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kInt32: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kUint32: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kInt64: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kUint64: + GetQuantizationParamsFromMinMax( + minmax, quantization_params); + break; + case ArrayDataType::kFloat: + case ArrayDataType::kNone: + default: + LOG(FATAL) << "Unhandled final quantization type " + << static_cast(data_type); + } +} + bool ChooseQuantizationForOperatorInput( GraphTransformation* transformation, Model* model, const Operator& op, std::size_t input_index, ArrayDataType* quantized_data_type, @@ -168,46 +280,70 @@ 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 = GetQuantizedDataType(array, ArrayDataType::kInt32); + transformation->AddMessageF( + "Input array %s is a bias vector. Choosing quantization params " + "accordingly.", + input); + return true; } const MinMax& minmax = GetOrComputeMinMax(model, input); - GetQuantizationParamsFromMinMax(model->flags, minmax, - quantization_params); + + if (op.type == OperatorType::kLstmCell) { + if (input_index == LstmCellOperator::PREV_STATE_INPUT) { + *quantized_data_type = ArrayDataType::kInt16; + GetQuantizationParams(*quantized_data_type, minmax, quantization_params); + return true; + } + } + + *quantized_data_type = GetQuantizedDataType(array, ArrayDataType::kUint8); + GetQuantizationParams(*quantized_data_type, minmax, quantization_params); transformation->AddMessageF( "For input array %s with min=%g" ", max=%g" - ", chose to quantize as uint8 with zero_point=%d" + ", chose to quantize as %s with zero_point=%d" ", scale=%g", - input, minmax.min, minmax.max, quantization_params->zero_point, - quantization_params->scale); - *quantized_data_type = ArrayDataType::kUint8; + input, minmax.min, minmax.max, ArrayDataTypeName(*quantized_data_type), + quantization_params->zero_point, quantization_params->scale); return true; } @@ -229,16 +365,18 @@ bool IsExactlyRepresentable(double real_value, ArrayDataType data_type, return true; } +// Quantized data type is preset to the type of the input before this function. bool ChooseHardcodedQuantizationForOperatorOutput( - const Operator& op, ArrayDataType* quantized_data_type, + const Operator& op, const Array& array, ArrayDataType* quantized_data_type, QuantizationParams* quantization_params) { if (op.type == OperatorType::kL2Normalization) { // L2Normalization has range: [-1, 1]. // 0 should be exactly representable, as values will typically be centered // around 0, with many values near 0. - *quantized_data_type = ArrayDataType::kUint8; - quantization_params->zero_point = 128; - quantization_params->scale = 1. / 128.; + *quantized_data_type = GetQuantizedDataType(array, *quantized_data_type); + const QuantizationPoints qp = GetQuantizationPoints(*quantized_data_type); + quantization_params->zero_point = qp.central_value; + quantization_params->scale = 1. / (qp.central_value - qp.min_value); CHECK( IsExactlyRepresentable(0., *quantized_data_type, *quantization_params)); return true; @@ -251,13 +389,26 @@ bool ChooseHardcodedQuantizationForOperatorOutput( // will typically exploit the symmetry logistic(-x) = 1 - logistic(x), and // the glueing of the two halves of the graph will only be seamless if we // are accurately representing logistic(0) == 0.5. - *quantized_data_type = ArrayDataType::kUint8; + *quantized_data_type = GetQuantizedDataType(array, *quantized_data_type); + const QuantizationPoints qp = GetQuantizationPoints(*quantized_data_type); quantization_params->zero_point = 0; - quantization_params->scale = 1. / 256.; + quantization_params->scale = 1. / (qp.max_value + 1); CHECK(IsExactlyRepresentable(0.5, *quantized_data_type, *quantization_params)); return true; } + if (op.type == OperatorType::kTanh) { + // Tanh has the range: [-1, 1]. + *quantized_data_type = GetQuantizedDataType(array, *quantized_data_type); + const QuantizationPoints qp = GetQuantizationPoints(*quantized_data_type); + quantization_params->zero_point = qp.central_value; + quantization_params->scale = 1. / (qp.central_value - qp.min_value); + // 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; } @@ -270,8 +421,9 @@ bool ChooseQuantizationForOperatorOutput( if (array.data_type != ArrayDataType::kFloat) { return false; } - if (ChooseHardcodedQuantizationForOperatorOutput(op, quantized_data_type, - quantization_params)) { + *quantized_data_type = model->GetArray(op.inputs[0]).data_type; + if (ChooseHardcodedQuantizationForOperatorOutput( + op, array, quantized_data_type, quantization_params)) { transformation->AddMessageF( "Output array %s is produced by a %s operator. Choosing fixed " "quantization params accordingly.", @@ -279,12 +431,21 @@ bool ChooseQuantizationForOperatorOutput( return true; } if ((op.type == OperatorType::kDepthToSpace) || - (op.type == OperatorType::kSpaceToDepth)) { - // DepthToSpace and SpaceToDepth should preserve the quantization parameters - // of the input array, as these are simple reshape operations. - const auto& input_quantization_params = - model->GetArray(op.inputs[0]).GetQuantizationParams(); - *quantized_data_type = ArrayDataType::kUint8; + (op.type == OperatorType::kSpaceToDepth) || + (op.type == OperatorType::kTensorFlowReshape) || + (op.type == OperatorType::kTensorFlowSplit) || + (op.type == OperatorType::kConcatenation)) { + int data_input_index = 0; + if (op.type == OperatorType::kTensorFlowSplit) { + data_input_index = 1; + } + // Copying and rearrangement ops should preserve the quantization parameters + // of the input array. + const auto& input_array = model->GetArray(op.inputs[data_input_index]); + const auto& input_quantization_params = input_array.GetQuantizationParams(); + *quantized_data_type = + GetQuantizedDataType(input_array, ArrayDataType::kUint8); + *quantized_data_type = GetQuantizedDataType(array, *quantized_data_type); quantization_params->zero_point = input_quantization_params.zero_point; quantization_params->scale = input_quantization_params.scale; @@ -295,15 +456,22 @@ bool ChooseQuantizationForOperatorOutput( return true; } const MinMax& minmax = GetOrComputeMinMax(model, output); - GetQuantizationParamsFromMinMax(model->flags, minmax, - quantization_params); - *quantized_data_type = ArrayDataType::kUint8; + if (op.type == OperatorType::kLstmCell) { + if (output_index == LstmCellOperator::STATE_OUTPUT || + output_index == LstmCellOperator::ACTIV_TEMP) { + *quantized_data_type = ArrayDataType::kInt16; + GetQuantizationParams(*quantized_data_type, minmax, quantization_params); + return true; + } + } + *quantized_data_type = GetQuantizedDataType(array, ArrayDataType::kUint8); + GetQuantizationParams(*quantized_data_type, minmax, quantization_params); transformation->AddMessageF( "For output array %s with min=%g, max=%g" - ", chose to quantize as uint8 with zero_point=%d" + ", chose to quantize as %s with zero_point=%d" ", scale=%g", - output, minmax.min, minmax.max, quantization_params->zero_point, - quantization_params->scale); + output, minmax.min, minmax.max, ArrayDataTypeName(*quantized_data_type), + quantization_params->zero_point, quantization_params->scale); return true; } @@ -345,9 +513,11 @@ bool Quantize::Run(Model* model, std::size_t op_index) { // // Let us just guard this assumption by the following assertion: for (const auto& input : op.inputs) { - if (IsInputArray(*model, input)) { - const auto& input_array = model->GetArray(input); - CHECK(input_array.quantization_params); + const auto& input_array = model->GetArray(input); + if (IsInputArray(*model, input) && + input_array.data_type == ArrayDataType::kFloat) { + CHECK(input_array.quantization_params) + << "Input array " << input << " is missing quantization_params"; } } if (!SupportsQuantization(op)) { @@ -390,30 +560,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->EraseArray(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/remove_trivial_passthrough.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc index 587f171bbf823408a45083c36d52f1d38c300123..aa93ace03af300f9cbd3f9c6620a6a58b9329aa4 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc @@ -60,7 +60,9 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation, for (int i = 0; i < passthru_op->inputs.size(); i++) { if (!model->GetArray(passthru_op->inputs[i]).buffer) { count_nonconstant_input_arrays++; - main_input_array_index = i; + if (count_nonconstant_input_arrays == 1) { + main_input_array_index = i; + } } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_activation_func.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_activation_func.cc index 28f76c9d36d6f68c8997fa0cf620c8aec4273619..9b65feaa6443cd32ac1bef961600ff225d52d4b2 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_activation_func.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_activation_func.cc @@ -12,6 +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 #include #include #include @@ -30,6 +31,7 @@ bool RemoveTrivialQuantizedActivationFunc::Run(Model* model, const auto it = model->operators.begin() + op_index; auto* op = it->get(); if (op->fused_activation_function != FusedActivationFunctionType::kRelu && + op->fused_activation_function != FusedActivationFunctionType::kRelu1 && op->fused_activation_function != FusedActivationFunctionType::kRelu6) { return false; } @@ -42,33 +44,49 @@ bool RemoveTrivialQuantizedActivationFunc::Run(Model* model, } const auto& quantization_params = output_array.GetQuantizationParams(); + double clamp_min; + double clamp_max; + switch (op->fused_activation_function) { + case FusedActivationFunctionType::kRelu: + clamp_min = 0.0; + clamp_max = std::numeric_limits::infinity(); + break; + case FusedActivationFunctionType::kRelu1: + clamp_min = -1.0; + clamp_max = 1.0; + break; + case FusedActivationFunctionType::kRelu6: + clamp_min = 0.0; + clamp_max = 6.0; + break; + default: + LOG(FATAL) << "Unsupported fused activation type: " + << static_cast(op->fused_activation_function); + return false; + } + bool has_nontrivial_min_bound = false; bool has_nontrivial_max_bound = false; - if (op->fused_activation_function == FusedActivationFunctionType::kRelu || - op->fused_activation_function == FusedActivationFunctionType::kRelu6) { - double lowest_representable_output = - (0. - quantization_params.zero_point) * quantization_params.scale; - if (lowest_representable_output < 0.) { - has_nontrivial_min_bound = true; - AddMessageF( - "Quantized activation function is not trivial: " - "the lowest representable output value %g" - " less than the clamp min bound.", - lowest_representable_output); - } + double lowest_representable_output = + (0. - quantization_params.zero_point) * quantization_params.scale; + if (lowest_representable_output < clamp_min) { + has_nontrivial_min_bound = true; + AddMessageF( + "Quantized activation function is not trivial: " + "the lowest representable output value %g" + " less than the clamp min bound %g.", + lowest_representable_output, clamp_min); } - if (op->fused_activation_function == FusedActivationFunctionType::kRelu6) { - double highest_representable_output = - (255. - quantization_params.zero_point) * quantization_params.scale; - if (highest_representable_output > 6.) { - has_nontrivial_max_bound = true; - AddMessageF( - "Quantized activation function is not trivial: " - "the highest representable output value %g" - " is greater than the clamp max bound.", - highest_representable_output); - } + double highest_representable_output = + (255. - quantization_params.zero_point) * quantization_params.scale; + if (highest_representable_output > clamp_max) { + has_nontrivial_max_bound = true; + AddMessageF( + "Quantized activation function is not trivial: " + "the highest representable output value %g" + " is greater than the clamp max bound %g.", + highest_representable_output, clamp_max); } if (has_nontrivial_min_bound || has_nontrivial_max_bound) { diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_reshape.cc index 90f9381ec154f145cda826ff9730ff332cd96701..61477d59aea2f11c6347b84d8863763a86c43558 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_reshape.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_reshape.cc @@ -61,8 +61,8 @@ bool IsReshapeTrivial(const Model& model, const Operator& op, if (next_op->type == OperatorType::kTensorFlowReshape) { transformation->AddMessageF( "%s is trivial because its output is only consumed by another " - "Reshape op", - LogName(op)); + "Reshape op %s", + LogName(op), LogName(*next_op)); return true; } } 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/reorder_activation_functions.cc b/tensorflow/contrib/lite/toco/graph_transformations/reorder_activation_functions.cc index cabbc4d313be3069053f056eb0de45c37ba2e7a4..9852c86c21b9a0714bc728e60b5d9dfe61ff52d1 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/reorder_activation_functions.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/reorder_activation_functions.cc @@ -42,14 +42,22 @@ bool ReorderActivationFunctions::Run(Model* model, std::size_t op_index) { std::unique_ptr& exchange_op = *exchange_it; DCHECK(exchange_op); - if (exchange_op->type != OperatorType::kTensorFlowReshape) { - return false; + // Allow activation functions to move up over any operator that does not + // change the values. + switch (exchange_op->type) { + case OperatorType::kExpandDims: + case OperatorType::kSqueeze: + case OperatorType::kTensorFlowReshape: + case OperatorType::kTranspose: + break; + default: + 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]; + 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); @@ -62,18 +70,62 @@ bool ReorderActivationFunctions::Run(Model* model, std::size_t op_index) { 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; + // If the ac_op was originally producing an output_array we can't trivially + // reorder as otherwise the output array name would change and break + // downstream assumptions. To work around that we perform some renaming below + // in that case at the cost of a bit more confusing array names in this rare + // case. + bool is_ac_op_output = + std::find(model->flags.output_arrays().begin(), + model->flags.output_arrays().end(), + ac_op_output) != model->flags.output_arrays().end(); + if (is_ac_op_output) { + // To preserve the output array name of the activation function we need to + // create a temporary to use to pass between ac->ex. + // + // Original: + // (a) -> EX -> (b) -> AC -> (c) + // Now: + // (a) -> AC -> (c') -> EX -> (c) + AddMessageF( + "Exchanging activation function %s with %s but renaming to preserve " + "output array %s", + LogName(*ac_op), LogName(*exchange_op), ac_op->outputs[0]); + + auto renamed_ac_op_output = + AvailableArrayName(*model, ac_op_output + "_exchange"); + ac_op->inputs[0] = exchange_op_input; + ac_op->outputs[0] = renamed_ac_op_output; + model->EraseArray(exchange_op->outputs[0]); + exchange_op->inputs[0] = renamed_ac_op_output; + exchange_op->outputs[0] = ac_op_output; + } else { + // Simply swap the order and update consumers to use the exchange_op output + // array (b). + // + // Original: + // (a) -> EX -> (b) -> AC -> (c) + // Now: + // (a) -> AC -> (c) -> EX -> (b) + AddMessageF("Exchanging activation function %s with %s", LogName(*ac_op), + LogName(*exchange_op)); + + 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); } - consumer = GetFirstOpWithInput(*model, ac_op_output); + ac_op->inputs[0] = exchange_op_input; + exchange_op->inputs[0] = 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. 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 5ac449749adbc9b5422f996eeccb72575dca8722..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, @@ -190,7 +190,7 @@ bool ResolveConstantConcatenation::Run(Model* model, std::size_t op_index) { // Remove all the resolved arrays. for (const string& input_name : concat_op->inputs) { // Check to prevent removal of shared tensors - if(CountOpsWithInput(*model, input_name) == 1) { + if (CountOpsWithInput(*model, input_name) == 1) { model->EraseArray(input_name); } } 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 81fe37d7e017c6e2440de34cc2daedf7fb2a422e..625d90205a801ad7c3fc1026c9cedc9b509f920d 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,12 +50,13 @@ 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); QuantizationParams qparams; - GetQuantizationParamsFromMinMax( - model->flags, *fakequant_op->minmax, &qparams); + GetQuantizationParamsFromMinMax(*fakequant_op->minmax, + &qparams); for (int i = 0; i < size; i++) { const double src_val = input_buffer.data[i]; const double unclamped_quantized_val = diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc new file mode 100644 index 0000000000000000000000000000000000000000..d999c2df9483e096f333c6af83e1d9fee873d4d6 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc @@ -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. +==============================================================================*/ +#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 { + +// Gathers data from axis 0. +template +inline void Gather(const Array& input_array, int input_rank, + const Array& coords_array, Array* output_array) { + const Shape& input_shape = input_array.shape(); + const std::vector>& input_data = + input_array.GetBuffer().data; + const Shape& coords_shape = coords_array.shape(); + const std::vector& coords_data = + coords_array.GetBuffer().data; + + const Shape& output_shape = output_array->shape(); + std::vector>& output_data = + output_array->GetMutableBuffer().data; + output_data.resize(RequiredBufferSizeForShape(output_shape)); + + int rev_input_rank = input_shape.dimensions_count() - 1 - (input_rank - 1); + CHECK_EQ(coords_shape.dims(0), output_array->shape().dims(rev_input_rank)); + + int stride = 1; + for (int i = input_shape.dimensions_count() - 1; i >= input_rank - 1; --i) { + stride *= input_shape.dims(i); + } + + for (int i = 0; i < coords_shape.dims(0); ++i) { + DCHECK_GE(coords_data[i], 0); + DCHECK_LT(coords_data[i], input_shape.dims(rev_input_rank)); + DataType* out = output_data.data() + i * stride; + const DataType* in = input_data.data() + coords_data[i] * stride; + memcpy(out, in, sizeof(DataType) * stride); + } +} + +} // namespace + +// Resolves a constant Gather operation. +// This simply performs the gather and produces the output array with the +// appropriate values. +bool ResolveConstantGather::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::kGather) { + 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; + } + + // Only handling axis=0 for now. + if (op->axis != 0) { + AddMessageF("%s has axis %d; only axis=0 is supported", LogName(*op), + op->axis); + 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]); + const Array& coords_array = model->GetArray(op->inputs[1]); + CHECK(coords_array.data_type == ArrayDataType::kInt32) + << "Only int32 indices are supported"; + + CHECK(!output_array.buffer); + switch (output_array.data_type) { + case ArrayDataType::kFloat: + Gather(input_array, op->input_rank, coords_array, + &output_array); + break; + case ArrayDataType::kUint8: + Gather(input_array, op->input_rank, coords_array, + &output_array); + break; + case ArrayDataType::kInt32: + Gather(input_array, op->input_rank, coords_array, + &output_array); + break; + case ArrayDataType::kInt64: + Gather(input_array, op->input_rank, coords_array, + &output_array); + break; + default: + LOG(FATAL) << "Unsupported data type given to Gather op with output \"" + << op->outputs[0] << "\""; + break; + } + + // Erase input arrays if no longer used after we remove the op. + DeleteArrayIfUsedOnce(op->inputs[0], model); + DeleteArrayIfUsedOnce(op->inputs[1], model); + + // Erase the operator. + model->operators.erase(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 1cd2aff28c68eaba4e9b18d8e2c2803834328696..d4db6f1c009cd19515655fb31974a2e97cfa42e8 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_unary.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_unary.cc @@ -28,21 +28,45 @@ limitations under the License. namespace toco { +bool CopyMinMaxFromFirstInput(const Operator& op, Model* model) { + auto& output_array = model->GetArray(op.outputs[0]); + if (output_array.minmax) { + return false; + } + const auto& input_array = model->GetArray(op.inputs[0]); + if (!input_array.minmax) { + return false; + } + const auto& input_minmax = input_array.GetMinMax(); + CHECK(!output_array.minmax); + auto& output_minmax = output_array.GetOrCreateMinMax(); + output_minmax.min = input_minmax.min; + output_minmax.max = input_minmax.max; + return true; +} + bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) { const auto unary_it = model->operators.begin() + op_index; const auto* unary_op = unary_it->get(); - // Test for unary ops of types that we know how to resolve - if (unary_op->type != OperatorType::kCast && - unary_op->type != OperatorType::kNeg && - unary_op->type != OperatorType::kTensorFlowRsqrt && - unary_op->type != OperatorType::kTensorFlowSqrt && - unary_op->type != OperatorType::kTensorFlowSquare && - unary_op->type != OperatorType::kTensorFlowSum && - unary_op->type != OperatorType::kTensorFlowMin && - unary_op->type != OperatorType::kTensorFlowMax && - unary_op->type != OperatorType::kTensorFlowReshape) { - return false; + // Test for unary ops of types that we know how to resolve. + switch (unary_op->type) { + case OperatorType::kCast: + case OperatorType::kNeg: + case OperatorType::kTensorFlowRsqrt: + case OperatorType::kTensorFlowSqrt: + case OperatorType::kTensorFlowSquare: + case OperatorType::kTensorFlowSum: + case OperatorType::kTensorFlowMin: + case OperatorType::kTensorFlowMax: + case OperatorType::kTensorFlowReshape: + case OperatorType::kRelu6: + case OperatorType::kRelu1: + case OperatorType::kRelu: + break; + default: + return false; } + // Check if the input is a constant parameter. if (!IsConstantParameterArray(*model, unary_op->inputs[0])) { return false; @@ -76,6 +100,12 @@ bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) { return false; } + // The min-max is only copied for ops that copy data without arithmetic. + // In future trivial transpose, etc, can be handled here. + if (unary_op->type == OperatorType::kTensorFlowReshape) { + CopyMinMaxFromFirstInput(*unary_op, model); + } + const auto& input_array = model->GetArray(unary_op->inputs[0]); // We have already tested above for existence of buffers (synonymous to being // a constant param). @@ -135,18 +165,36 @@ bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) { } } else if (unary_op->type == OperatorType::kTensorFlowReshape) { CHECK(input_buffer_size == output_buffer_size); - memcpy(output_float_data.data(), (*input_float_data).data(), - output_buffer_size * sizeof(output_float_data[0])); + output_float_data = *input_float_data; } 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); + CHECK_EQ(unary_op->inputs.size(), 2) << "Sum needs 2 inputs"; + if (!IsConstantParameterArray(*model, unary_op->inputs[1])) { + AddMessageF("Axis input is non-constant"); + return false; } - float sum = 0.f; - for (int i = 0; i < input_buffer_size; i++) { - sum += (*input_float_data)[i]; + auto& axis_array = model->GetArray(unary_op->inputs[1]); + CHECK(axis_array.data_type == ArrayDataType::kInt32); + int axis = axis_array.GetBuffer().data[0]; + CHECK_LT(axis, input_shape.dimensions_count()) << "Axis out of bounds"; + + // We currently only handle reduction on axis 0. + CHECK_EQ(axis, 0) << "Only reduction along axis 0 is supported"; + // We currently only handle 1-D and 2-D input tensors. + CHECK_LE(input_shape.dimensions_count(), 2) << "Rank >2 not yet supported"; + // We only support keep_dims=true; shape prop will need to change otherwise. + auto sum_op = static_cast(unary_op); + CHECK(sum_op->keep_dims) << "Only keep_dims=true is supported"; + + std::vector indices(input_shape.dimensions_count()); + for (int i = 0; i < input_shape.dims(1); ++i) { + indices[1] = i; + float sum = 0.f; + for (int j = 0; j < input_shape.dims(0); ++j) { + indices[0] = j; + sum += (*input_float_data)[Offset(input_shape, indices)]; + } + output_float_data[i] = sum; } - output_float_data[0] = 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. @@ -194,6 +242,37 @@ bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) { } output_float_data[i] = outval; } + } else if (unary_op->type == OperatorType::kRelu6 && + unary_op->type == OperatorType::kRelu1 && + unary_op->type == OperatorType::kRelu) { + for (size_t i = 0; i < output_buffer_size; ++i) { + const float value = (*input_float_data)[i]; + float new_value = 0.0f; + switch (unary_op->type) { + case OperatorType::kRelu: { + static constexpr float kLower = 0; + new_value = value < kLower ? kLower : value; + break; + } + case OperatorType::kRelu1: { + static constexpr float kUpper = 1; + static constexpr float kLower = -1; + new_value = value > kUpper ? kUpper : value < kLower ? kLower : value; + break; + } + case OperatorType::kRelu6: { + static constexpr float kUpper = 6; + static constexpr float kLower = 0; + new_value = value > kUpper ? kUpper : value < kLower ? kLower : value; + break; + } + default: + LOG(FATAL) << "Unsupported activation function " + << LogName(*unary_op); + return false; + } + output_float_data[i] = new_value; + } } else { LOG(FATAL) << "should not get here."; } 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_reorder_axes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc index 5c68f87f6ccd912a94213c95a59a78076b0e768b..bc70db0bd8c26319fa140616de96452260a01058 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc @@ -60,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. @@ -77,14 +68,14 @@ bool ResolveReorderAxes::Run(Model* model, std::size_t op_index) { return false; } // Reorder the input array dims and buffer data - if (constant_input_array.buffer->type == ArrayDataType::kFloat) { - ReorderAxes( - reorder_op->input_axes_order, reorder_op->output_axes_order, - &constant_input_array, &output_array); - } else if (constant_input_array.buffer->type == ArrayDataType::kInt32) { - ReorderAxes( - reorder_op->input_axes_order, reorder_op->output_axes_order, - &constant_input_array, &output_array); + 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."; } 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 ad1e56888e53133c5a84cc0e3d5e76b7ef3b29b4..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,22 +76,26 @@ 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), @@ -72,7 +105,7 @@ bool ResolveTensorFlowMatMul::Run(Model* model, std::size_t op_index) { 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]; @@ -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/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/unpartition_embedding_lookup.cc b/tensorflow/contrib/lite/toco/graph_transformations/unpartition_embedding_lookup.cc new file mode 100644 index 0000000000000000000000000000000000000000..48c326651f3201b4f7a31ac2440b171841e8ed7b --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/unpartition_embedding_lookup.cc @@ -0,0 +1,240 @@ +/* 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/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" + +namespace toco { + +bool UnpartitionEmbeddingLookup::Run(Model* model, std::size_t op_index) { + // Collapses a partitioned tf.nn.embedding_lookup back into a single Gather. + // https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup + // This transform attempts to identify the len(params) > 1 case and collapse + // it to the len(params) = 1 case by concatenating the original params and + // reversing the partitioning. + // + // If len(params) to the tf.nn.embedding_lookup == 1, the whole op becomes + // simply a gather: + // https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/python/ops/embedding_ops.py#L150 + // + // Notes on this implementation: + // - only supports partition_strategy='mod' + // + // A rough graph of a partitioned embedding_lookup looks like: + // (ids)--+-->FloorDiv--+-->DynamicPartition-->[[Gather]]--\ + // \-->FloorMod--/ | + // V | + // Range-->DynamicPartition-------->DynamicStitch<---------/ + // (const) V + // (embeddings) + + // First look for the final DynamicStitch. + auto op_it = model->operators.begin() + op_index; + if (op_it->get()->type != OperatorType::kDynamicStitch) { + return false; + } + auto* stitch_op = static_cast(op_it->get()); + + // Split up the DynamicStitch inputs into the indices and data. + std::vector stitch_indices_inputs; + std::vector stitch_data_inputs; + for (size_t i = 0; i < stitch_op->num_partitions; ++i) { + stitch_indices_inputs.push_back(stitch_op->inputs[i]); + } + for (size_t i = stitch_op->num_partitions; i < stitch_op->num_partitions * 2; + ++i) { + stitch_data_inputs.push_back(stitch_op->inputs[i]); + } + + // Validate all indices come from the same DynamicPartition. + DynamicPartitionOperator* indices_partition_op = nullptr; + for (const string& indices_partition_output_name : stitch_indices_inputs) { + auto* op = GetOpWithOutput(*model, indices_partition_output_name); + CHECK(op) << "Source of " << indices_partition_output_name << " not found"; + if (op->type != OperatorType::kDynamicPartition) { + AddMessageF( + "Skipping because indices input %s into " + "%s is unexpected", + LogName(*op), LogName(*stitch_op)); + return false; + } + if (!indices_partition_op) { + indices_partition_op = static_cast(op); + } else { + // Ensure this is the same op as previous ones. + if (op != indices_partition_op) { + AddMessageF( + "Skipping because indices input %s into " + "%s is from a different source op than others", + LogName(*op), LogName(*stitch_op)); + return false; + } + } + } + CHECK(indices_partition_op) << "No indices inputs"; + + // The data for the indices must be a constant range of the array shape. + if (!IsConstantParameterArray(*model, indices_partition_op->inputs[0])) { + AddMessageF("Skipping because indices partition data is non-constant"); + return false; + } + auto& indices_data_array = model->GetArray(indices_partition_op->inputs[0]); + if (indices_data_array.data_type == ArrayDataType::kNone) { + // Yield until data types are propagated. + return false; + } + CHECK(indices_data_array.data_type == ArrayDataType::kInt32) + << "Indices partition inputs must be int32"; + const auto& indices_data_buffer = + indices_data_array.GetBuffer().data; + for (size_t i = 0; i < indices_data_buffer.size(); ++i) { + CHECK_EQ(indices_data_buffer[i], i) << "Indices range must be identity"; + } + + // Find all of the gathers used for the data inputs. + std::vector gather_ops; + for (const string& gather_output_name : stitch_data_inputs) { + auto* op = GetOpWithOutput(*model, gather_output_name); + CHECK(op) << "Source of " << gather_output_name << " not found"; + if (op->type != OperatorType::kGather) { + AddMessageF( + "Skipping because data input %s into %s " + "is unexpected", + LogName(*op), LogName(*stitch_op)); + return false; + } + gather_ops.push_back(static_cast(op)); + } + + // Validate all gathers come from the same DynamicPartition. + DynamicPartitionOperator* data_partition_op = nullptr; + for (auto* gather_op : gather_ops) { + auto* op = GetOpWithOutput(*model, gather_op->inputs[1]); + CHECK(op) << "Source of " << gather_op->inputs[1] << " not found"; + if (op->type != OperatorType::kDynamicPartition) { + AddMessageF( + "Skipping because data input %s into " + "%s is unexpected", + LogName(*op), LogName(*gather_op)); + return false; + } + if (!data_partition_op) { + data_partition_op = static_cast(op); + } else { + // Ensure this is the same op as previous ones. + if (op != data_partition_op) { + AddMessageF( + "Skipping because data input %s into " + "%s is from a different source op than others", + LogName(*op), LogName(*gather_op)); + return false; + } + } + } + CHECK(data_partition_op) << "No data inputs"; + + // Validate the partition ops have the same sizes. + CHECK_EQ(indices_partition_op->num_partitions, + data_partition_op->num_partitions) + << "Indices and data partition ops have differing dimensions"; + int num_partitions = indices_partition_op->num_partitions; + + // Partition strategy of 'mod' gives us a FloorMod and FloorDiv. + // The gather partition uses the FloorDiv as the data and FloorMod as the + // partitions and the indices use the FloorMod as their partitions. + Operator* div_op = GetOpWithOutput(*model, data_partition_op->inputs[0]); + Operator* mod_op = GetOpWithOutput(*model, data_partition_op->inputs[1]); + CHECK(div_op && div_op->type == OperatorType::kFloorDiv) + << "Unsupported partition strategy"; + CHECK(mod_op && mod_op->type == OperatorType::kFloorMod) + << "Unsupported partition strategy"; + CHECK_EQ(mod_op, GetOpWithOutput(*model, indices_partition_op->inputs[1])) + << "Indices and data parition ops require the same partition strategy " + "and inputs"; + + // Glob together all of the gather data. This is not yet in the correct order. + auto* gather_params_concat_op = new ConcatenationOperator; + for (const auto& gather_op : gather_ops) { + gather_params_concat_op->inputs.push_back(gather_op->inputs[0]); + } + gather_params_concat_op->outputs.push_back( + AvailableArrayName(*model, gather_ops[0]->inputs[0] + "_unpartitioned")); + op_it = model->operators.emplace(op_it, gather_params_concat_op) + 1; + model->GetOrCreateArray(gather_params_concat_op->outputs[0]); + + // Permute the gather params to undo the partitioning that was originally + // done. + auto* gather_params_permute_op = new GatherOperator; + gather_params_permute_op->inputs.push_back( + gather_params_concat_op->outputs[0]); + gather_params_permute_op->inputs.push_back( + AvailableArrayName(*model, gather_ops[0]->inputs[0] + "_permuted/perm")); + gather_params_permute_op->outputs.push_back( + AvailableArrayName(*model, gather_ops[0]->inputs[0] + "_permuted")); + op_it = model->operators.emplace(op_it, gather_params_permute_op) + 1; + model->GetOrCreateArray(gather_params_permute_op->outputs[0]); + const auto& partition_array = model->GetArray(gather_ops[0]->inputs[0]); + const auto& partition_array_dims = partition_array.shape().dims(); + gather_params_permute_op->input_rank = + partition_array.shape().dimensions_count(); + auto& perm_array = + model->GetOrCreateArray(gather_params_permute_op->inputs[1]); + perm_array.data_type = ArrayDataType::kInt32; + perm_array.mutable_shape()->ReplaceDims( + {num_partitions * partition_array_dims[0]}); + auto& perm_data = perm_array.GetMutableBuffer().data; + perm_data.resize(RequiredBufferSizeForShape(perm_array.shape())); + // NOTE: this is what relies on the partition_strategy. + for (int i = 0; i < num_partitions * partition_array_dims[0]; ++i) { + int p = i % num_partitions; + perm_data[i] = p * partition_array_dims[0] + i / num_partitions; + } + + // Insert the new unpartitioned gather op. + auto* merged_gather_op = new GatherOperator; + merged_gather_op->inputs = {gather_params_permute_op->outputs[0], + mod_op->inputs[0]}; + merged_gather_op->outputs = {stitch_op->outputs[0]}; + merged_gather_op->input_rank = partition_array.shape().dimensions_count(); + model->operators.emplace(op_it, merged_gather_op); + + AddMessageF( + "Replacing suspected partitioned tf.nn.embedding_lookup (starting at %s " + "+ %s and ending at %s) with a single unpartitioned gather %s", + LogName(*div_op), LogName(*mod_op), LogName(*stitch_op), + LogName(*merged_gather_op)); + + // Ensure the stitch output array is dead, as we don't want whatever was in it + // previously now that we've redefined it. It'll be recreated when needed. + model->EraseArray(stitch_op->outputs[0]); + model->GetOrCreateArray(merged_gather_op->outputs[0]); + + // Erase all the original ops. + DeleteOpAndArraysIfUnused(model, div_op); + DeleteOpAndArraysIfUnused(model, mod_op); + for (auto* gather_op : gather_ops) { + DeleteOpAndArraysIfUnused(model, gather_op); + } + DeleteOpAndArraysIfUnused(model, indices_partition_op); + DeleteOpAndArraysIfUnused(model, data_partition_op); + DeleteOpAndArraysIfUnused(model, stitch_op); + return true; +} + +} // namespace toco 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 ca378af4c5c1e1b8cf42a10d3820db3feeb49a05..a7a50e6fc9326338d69cd0334c83b60e2fa50402 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" @@ -173,7 +174,8 @@ void ImportFloatArray(const TensorProto& input_tensor, Array* output_array) { } auto& output_float_data = output_array->GetMutableBuffer().data; - output_float_data.resize(input_flat_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); @@ -203,7 +205,7 @@ void ImportQuint8Array(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); @@ -229,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); @@ -255,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); @@ -270,6 +272,39 @@ void ImportInt64Array(const TensorProto& input_tensor, Array* output_array) { } } +void ImportBoolArray(const TensorProto& input_tensor, Array* output_array) { + CHECK_EQ(input_tensor.dtype(), DT_BOOL); + 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_bool_data = + output_array->GetMutableBuffer().data; + output_bool_data.resize(RequiredBufferSizeForShape(output_array->shape()), + false); + if (input_tensor.bool_val_size()) { + for (int i = 0; i < input_tensor.bool_val_size(); i++) { + output_bool_data[i] = input_tensor.bool_val(i); + } + } else if (input_tensor.tensor_content().size() == input_flat_size) { + std::vector buf(input_tensor.tensor_content().size()); + toco::port::CopyToBuffer(input_tensor.tensor_content(), buf.data()); + for (int i = 0; i < input_tensor.tensor_content().size(); i++) { + output_bool_data[i] = static_cast(buf[i]); + } + } else { + // Some graphs have bool const nodes without actual value... + // assuming that 'false' is implied. + // So far only encountered that in an array with 1 entry, let's + // require that until we encounter a graph where that's not the case. + CHECK_EQ(output_bool_data.size(), 1); + output_bool_data[0] = false; + } +} + void ImportStringArray(const TensorProto& input_tensor, Array* output_array) { CHECK_EQ(input_tensor.dtype(), DT_STRING); const auto& input_shape = input_tensor.tensor_shape(); @@ -281,7 +316,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."; @@ -316,6 +351,18 @@ void CheckInputsCount(const NodeDef& node, << " input(s) other than control dependencies: " << node.DebugString(); } +template +string CreateConstArray(Model* model, string const& name, + std::vector > const& data) { + // Utility function to create a const 1D array, useful for input parameters. + string array_name = toco::AvailableArrayName(*model, name); + auto& array = model->GetOrCreateArray(array_name); + array.data_type = T; + array.mutable_shape()->mutable_dims()->emplace_back(data.size()); + array.GetMutableBuffer().data = data; + return array_name; +} + void ConvertConstOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { @@ -345,6 +392,10 @@ void ConvertConstOperator(const NodeDef& node, array.data_type = ArrayDataType::kString; ImportStringArray(tensor, &array); break; + case DT_BOOL: + array.data_type = ArrayDataType::kBool; + ImportBoolArray(tensor, &array); + break; default: array.data_type = ArrayDataType::kNone; // do nothing, silently ignore the Const data. @@ -363,7 +414,7 @@ void ConvertConvOperator(const NodeDef& node, // We only support NHWC, which is the default data_format. // So if data_format is not defined, we're all good. - if (node.attr().count("data_format")) { + if (HasAttr(node, "data_format")) { CHECK_EQ(GetStringAttr(node, "data_format"), "NHWC"); } CHECK_EQ(GetDataTypeAttr(node, "T"), DT_FLOAT); @@ -397,6 +448,17 @@ void ConvertConvOperator(const NodeDef& node, CHECK_EQ(strides.i(3), 1); conv->stride_height = strides.i(1); conv->stride_width = strides.i(2); + if (HasAttr(node, "dilations")) { + const auto& dilations = GetListAttr(node, "dilations"); + CHECK_EQ(dilations.i_size(), 4); + CHECK_EQ(dilations.i(0), 1); + CHECK_EQ(dilations.i(3), 1); + conv->dilation_height_factor = dilations.i(1); + conv->dilation_width_factor = dilations.i(2); + } else { + conv->dilation_height_factor = 1; + conv->dilation_width_factor = 1; + } const auto& padding = GetStringAttr(node, "padding"); if (padding == "SAME") { conv->padding.type = PaddingType::kSame; @@ -416,7 +478,7 @@ void ConvertDepthwiseConvOperator(const NodeDef& node, // We only support NHWC, which is the default data_format. // So if data_format is not defined, we're all good. - if (node.attr().count("data_format")) { + if (HasAttr(node, "data_format")) { CHECK_EQ(GetStringAttr(node, "data_format"), "NHWC"); } CHECK_EQ(GetDataTypeAttr(node, "T"), DT_FLOAT); @@ -665,9 +727,12 @@ void ConvertSqueezeOperator(const NodeDef& node, op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); - const auto& squeeze_dims = GetListAttr(node, "squeeze_dims"); - for (int i = 0; i < squeeze_dims.i_size(); ++i) { - op->squeeze_dims.push_back(squeeze_dims.i(i)); + // When omitted we are to squeeze all dimensions == 1. + if (HasAttr(node, "squeeze_dims")) { + const auto& squeeze_dims = GetListAttr(node, "squeeze_dims"); + for (int i = 0; i < squeeze_dims.i_size(); ++i) { + op->squeeze_dims.push_back(squeeze_dims.i(i)); + } } model->operators.emplace_back(op); @@ -838,6 +903,7 @@ 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) { @@ -853,6 +919,18 @@ 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) { @@ -961,49 +1039,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) { CheckInputsCount(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'"; - } - 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); - } + // 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)); + 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); } @@ -1311,6 +1377,12 @@ void ConvertResizeBilinearOperator(const NodeDef& node, CHECK_EQ(node.op(), "ResizeBilinear"); 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()); @@ -1379,12 +1451,8 @@ void ConvertFusedBatchNormOperator(const NodeDef& node, const string& moving_variance_input = node.input(4); // Create an array holding the epsilon value (typically, 0.001). - const string epsilon_array_name = node.name() + "_epsilon_array"; - auto& epsilon_array = model->GetOrCreateArray(epsilon_array_name); - epsilon_array.data_type = ArrayDataType::kFloat; - *epsilon_array.mutable_shape()->mutable_dims() = {1}; - epsilon_array.GetMutableBuffer().data.push_back( - GetFloatAttr(node, "epsilon")); + const string epsilon_array_name = CreateConstArray( + model, node.name() + "_epsilon_array", {GetFloatAttr(node, "epsilon")}); // Add epsilon to the moving variance. const string epsilon_add_op_name = node.name() + "_epsilon"; @@ -1452,6 +1520,17 @@ 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) { @@ -1501,16 +1580,56 @@ void ConvertTransposeConvOperator(const NodeDef& node, CHECK_EQ(node.op(), "Conv2DBackpropInput"); CheckInputsCount(node, tf_import_flags, 3); auto* op = new TransposeConvOperator; - op->inputs.push_back(node.input(2)); - op->inputs.push_back(node.input(1)); op->inputs.push_back(node.input(0)); + op->inputs.push_back(node.input(1)); + op->inputs.push_back(node.input(2)); op->outputs.push_back(node.name()); const auto& strides = GetListAttr(node, "strides"); - CHECK_EQ(strides.i_size(), 4); - CHECK_EQ(strides.i(0), 1); op->stride_height = strides.i(1); op->stride_width = strides.i(2); - CHECK_EQ(strides.i(3), 1); + CHECK_EQ(strides.i_size(), 4) + << "Can only import TransposeConv ops with 4D strides. TensorFlow op \"" + << node.name() << "\" has " << strides.i_size() << "D strides."; + CHECK((strides.i(0) == 1) && (strides.i(3) == 1)) + << "Can only import TransposeConv ops with striding along the height " + "(1st) or width (2nd) axis. TensorFlow op \"" + << node.name() << "\" had strides:[ " << strides.i(0) << ", " + << strides.i(1) << ", " << strides.i(2) << ", " << strides.i(3) << "]."; + op->stride_height = strides.i(1); + op->stride_width = strides.i(2); + if (HasAttr(node, "dilations")) { + const auto& dilations = GetListAttr(node, "dilations"); + CHECK_EQ(dilations.i_size(), 4) + << "Dilation unsupported in TransposeConv. TensorFlow op \"" + << node.name() << "\" had dilations"; + CHECK((dilations.i(0) == 1) && (dilations.i(1) == 1) && + (dilations.i(1) == 1) && (dilations.i(3) == 1)) + << "Dilation unsupported in TransposeConv. TensorFlow op \"" + << node.name() << "\" had dilations:[ " << dilations.i(0) << ", " + << dilations.i(1) << ", " << dilations.i(2) << ", " << dilations.i(3) + << "]."; + } + + const string& weights_name = node.input(TransposeConvOperator::WEIGHTS); + const string& transposed_weights_name = weights_name + "_transposed"; + // Check if a TransposeOperator was already created for these weights + // (can happen when multiple layers share the same weights). + const Operator* existing_transpose = + GetOpWithOutput(*model, transposed_weights_name); + if (existing_transpose) { + CHECK(existing_transpose->type == OperatorType::kTranspose); + } else { + // Transpose weights from HWIO order to OHWI order, which is more efficient + // for computation + TransposeOperator* transpose = new TransposeOperator; + string perm_array = CreateConstArray( + model, node.name() + "_transpose_perm", {3, 0, 1, 2}); + transpose->inputs = {weights_name, perm_array}; + transpose->outputs = {transposed_weights_name}; + model->operators.emplace_back(transpose); + } + op->inputs[1] = transposed_weights_name; + auto const& padding = GetStringAttr(node, "padding"); if (padding == "SAME") { op->padding.type = PaddingType::kSame; @@ -1562,7 +1681,7 @@ void ConvertFloorDivOperator(const NodeDef& node, void ConvertFloorModOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { - CHECK(node.op() == "FloorMod"); + CHECK_EQ(node.op(), "FloorMod"); CheckInputsCount(node, tf_import_flags, 2); auto* op = new FloorModOperator; op->inputs.push_back(node.input(0)); @@ -1797,6 +1916,63 @@ 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")) { + string k_array = CreateConstArray( + model, node.name() + "k", {GetIntAttr(node, "k")}); + op->inputs.push_back(k_array); + } 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()); +} + +void ConvertDynamicPartitionOperator( + const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, + Model* model) { + auto op = absl::make_unique(); + CHECK(HasAttr(node, "num_partitions")); + op->num_partitions = GetIntAttr(node, "num_partitions"); + CheckInputsCount(node, tf_import_flags, 2); + op->inputs.push_back(node.input(0)); + op->inputs.push_back(node.input(1)); + CHECK_GT(op->num_partitions, 1); + op->outputs.push_back(node.name()); // Implicit :0. + for (int i = 1; i < op->num_partitions; ++i) { + op->outputs.push_back(node.name() + ":" + std::to_string(i)); + } + model->operators.emplace_back(op.release()); +} + +void ConvertDynamicStitchOperator(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + Model* model) { + // The parallel and non-parallel variants are the same besides whether they + // have a parallel loop; there are no behavioral differences. + CHECK(node.op() == "DynamicStitch" || node.op() == "ParallelDynamicStitch"); + auto op = absl::make_unique(); + CHECK(HasAttr(node, "N")); + op->num_partitions = GetIntAttr(node, "N"); + // Expect all ID partitions + all value partitions. + CheckInputsCount(node, tf_import_flags, op->num_partitions * 2); + for (int i = 0; i < op->num_partitions * 2; ++i) { + op->inputs.push_back(node.input(i)); + } + op->outputs.push_back(node.name()); + model->operators.emplace_back(op.release()); +} + } // namespace std::unique_ptr ImportTensorFlowGraphDef( @@ -1852,7 +2028,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); @@ -1891,6 +2069,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") { @@ -1974,6 +2154,15 @@ 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 if (node.op() == "DynamicPartition") { + ConvertDynamicPartitionOperator(node, tf_import_flags, model); + } else if (node.op() == "DynamicStitch" || + node.op() == "ParallelDynamicStitch") { + ConvertDynamicStitchOperator(node, tf_import_flags, model); } else { ConvertUnsupportedOperator(node, tf_import_flags, model); } diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h index 6fba8f2629f785ffeb3ae37b80ec1d24c29d9d56..3fa0089cba15fbfed869b5c33883741c35493cda 100644 --- a/tensorflow/contrib/lite/toco/model.h +++ b/tensorflow/contrib/lite/toco/model.h @@ -29,12 +29,15 @@ limitations under the License. namespace toco { +using tflite::QuantizationParams; + enum class OperatorType { kNone, // General-purpose neural network operators. kAdd, kAddN, kAveragePool, + kBatchMatMul, kBatchNormalization, kConv, kConcatenation, @@ -43,6 +46,7 @@ enum class OperatorType { kSpaceToDepth, kDequantize, kDiv, + kExp, kExpandDims, kFill, kFloorDiv, @@ -62,6 +66,7 @@ enum class OperatorType { kRelu1, kRelu6, kSoftmax, + kLogSoftmax, kSub, kTanh, kTransposeConv, @@ -111,6 +116,9 @@ enum class OperatorType { kTensorFlowSwitch, kTensorFlowTile, kTranspose, + kTopK_V2, + kDynamicPartition, + kDynamicStitch, // 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, @@ -156,12 +164,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 }; @@ -181,18 +194,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; }; @@ -215,6 +248,8 @@ struct GenericBuffer { // in containers and have the containers call the right subclass destructor. virtual ~GenericBuffer() {} + virtual int Length() const = 0; + const ArrayDataType type; protected: @@ -227,6 +262,8 @@ template struct Buffer : GenericBuffer { Buffer() : GenericBuffer(A) {} + int Length() const override { return data.size(); } + std::vector> data; }; @@ -330,7 +367,8 @@ struct ConvOperator : Operator { // A dilation_rate of 0 is invalid and this field is an optional attribute. // Thus initializing it to 1 to allow default conv behavior when the // attribute is not present. - int dilation_rate = 1; + int dilation_width_factor = 1; + int dilation_height_factor = 1; }; // Depthwise-separable convolution operator. @@ -712,6 +750,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. @@ -797,19 +848,40 @@ struct SqueezeOperator : Operator { }; // Inputs: -// inputs[0]: required: the input activations array -// inputs[1]: required: the Conv weights -// channel. +// inputs[0]: required: the output shape +// inputs[1]: required: the weights +// inputs[2]: required: the input activations array +// NOTE: The input activations is NOT the first input. +// // // Outputs: // outputs[0]: required: the output activations array // // TensorFlow equivalent: Conv2DBackpropInput struct TransposeConvOperator : Operator { + enum Inputs { + OUTPUT_SHAPE = 0, + WEIGHTS = 1, + DATA_INPUT = 2, + }; + TransposeConvOperator() : Operator(OperatorType::kTransposeConv) {} Padding padding; int stride_width = 0; int stride_height = 0; + // Dilation is possible with transpose convolution, but Tensorflow does not + // currently support it, so we omit it. +}; + +// 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 @@ -1216,6 +1288,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: @@ -1273,6 +1355,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 @@ -1336,6 +1420,38 @@ struct SvdfOperator : Operator { int rank; }; +// TopKV2 operator. +// +// Inputs: +// input tensor and top_k scalar. +struct TopKV2Operator : Operator { + TopKV2Operator() : Operator(OperatorType::kTopK_V2) {} +}; + +// DynamicPartition operator: +// +// Inputs: +// inputs[0]: required: data. +// inputs[1]: required: partitions. +// +// TensorFlow equivalent: DynamicPartition +struct DynamicPartitionOperator : Operator { + DynamicPartitionOperator() : Operator(OperatorType::kDynamicPartition) {} + int num_partitions; +}; + +// DynamicStitch operator: +// +// Inputs: +// inputs[0,N): required: indices. +// inputs[N,2N): required: data. +// +// TensorFlow equivalent: DynamicStitch/ParallelDynamicStitch +struct DynamicStitchOperator : Operator { + DynamicStitchOperator() : Operator(OperatorType::kDynamicStitch) {} + int num_partitions; +}; + // 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 @@ -1349,22 +1465,6 @@ inline bool operator<(const Alloc& a, const Alloc& b) { return a.start < b.start; } -// Quantization parameters, determining the mapping of quantized values -// to real values (i.e. determining how quantized values are mathematically -// interpreted). -// -// The correspondence is as follows: -// -// real_value = scale * (quantized_value - zero_point); -// -// In other words, zero_point designates which quantized value corresponds to -// the real 0 value, and scale designates the difference between the real values -// corresponding to consecutive quantized values differing by 1. -struct QuantizationParams { - int32 zero_point = 0; - double scale = 0.; -}; - class Shape { public: // For Shape, we stick to half-way encapsulation for now: @@ -1542,7 +1642,7 @@ class Model { bool HasArray(const string& name) const { return arrays.count(name) > 0; } Array& GetArray(const string& name) const { - DCHECK(HasArray(name)); + DCHECK(HasArray(name)) << "Array not found: " << name; return *arrays.at(name); } Array& GetOrCreateArray(const string& name) { diff --git a/tensorflow/contrib/lite/toco/model_flags.proto b/tensorflow/contrib/lite/toco/model_flags.proto index e4b39b34e85e4d703c1b41cb68f8139abd1f6279..867b86f31d16b502a7aeb92cb3d8c96117630cd2 100644 --- a/tensorflow/contrib/lite/toco/model_flags.proto +++ b/tensorflow/contrib/lite/toco/model_flags.proto @@ -96,9 +96,11 @@ message RnnState { // model that does not already contain such MinMax information. message ArraysExtraInfo { message Entry { + // Next ID to use: 5. optional string name = 1; optional float min = 2; optional float max = 3; + optional IODataType data_type = 4; } repeated Entry entries = 1; } diff --git a/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_cluster.cc b/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_cluster.cc index fddf6cc83686632033f31496ec42b33e2ea15f20..5e421ba944cccd9746c66bc33e986b4406dd3bf5 100644 --- a/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_cluster.cc +++ b/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_cluster.cc @@ -144,7 +144,9 @@ std::unique_ptr MaybeReplaceCompositeSubgraph( MaybeResolveClusters(tf_graph, cluster_factories); // Copy function definitions - *(pruned_graph->mutable_library()) = tf_graph.library(); + if (pruned_graph) { + *(pruned_graph->mutable_library()) = tf_graph.library(); + } return pruned_graph; } 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/tflite/import.cc b/tensorflow/contrib/lite/toco/tflite/import.cc index 5b1ab514b23248cd98e66847185d0e8b9fe2d6aa..e16784fd21dd5e6eeac5f03ba6250341dcdfbdb4 100644 --- a/tensorflow/contrib/lite/toco/tflite/import.cc +++ b/tensorflow/contrib/lite/toco/tflite/import.cc @@ -64,6 +64,9 @@ void ImportTensors(const ::tflite::Model& input_model, Model* model) { auto shape = input_tensor->shape(); if (shape) { + // If the shape is 0-dimensional, make sure to record it as such, + // as oppose to leaving the array without a shape. + array.mutable_shape()->mutable_dims()->clear(); for (int i = 0; i < shape->Length(); ++i) { auto d = shape->Get(i); array.mutable_shape()->mutable_dims()->push_back(d); @@ -167,8 +170,8 @@ std::unique_ptr Import(const ModelFlags& model_flags, // Full list of all known operators. const auto ops_by_name = BuildOperatorByNameMap(); - if (input_model->subgraphs()->size() != 1) { - LOG(FATAL) << "# of subgraphs in tflite should be exactly 1 for now."; + if (!input_model->subgraphs() || input_model->subgraphs()->size() != 1) { + LOG(FATAL) << "Number of subgraphs in tflite should be exactly 1."; } std::unique_ptr model; model.reset(new Model); diff --git a/tensorflow/contrib/lite/toco/tflite/import_test.cc b/tensorflow/contrib/lite/toco/tflite/import_test.cc index aad6e780d5eb5c3dbc880906df5053ad231ffd54..f25b17087639d1b04de4aaf371726fa594427ff1 100644 --- a/tensorflow/contrib/lite/toco/tflite/import_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/import_test.cc @@ -27,33 +27,32 @@ namespace { using ::testing::ElementsAre; +using flatbuffers::Offset; +using flatbuffers::Vector; class ImportTest : public ::testing::Test { protected: template - flatbuffers::Offset> CreateDataVector( - const std::vector& data) { + Offset> CreateDataVector(const std::vector& data) { return builder_.CreateVector(reinterpret_cast(data.data()), sizeof(T) * data.size()); } - // This is a very simplistic model. We are not interested in testing all the - // details here, since tf.mini's testing framework will be exercising all the - // conversions multiple times, and the conversion of operators is tested by - // separate unittests. - void BuildTestModel() { - // The tensors + Offset>> BuildBuffers() { + auto buf0 = ::tflite::CreateBuffer(builder_, CreateDataVector({})); + auto buf1 = + ::tflite::CreateBuffer(builder_, CreateDataVector({1.0f, 2.0f})); + auto buf2 = + ::tflite::CreateBuffer(builder_, CreateDataVector({3.0f})); + return builder_.CreateVector( + std::vector>({buf0, buf1, buf2})); + } + + Offset>> BuildTensors() { auto q = ::tflite::CreateQuantizationParameters( builder_, /*min=*/builder_.CreateVector({0.1f}), /*max=*/builder_.CreateVector({0.2f}), /*scale=*/builder_.CreateVector({0.3f}), /*zero_point=*/builder_.CreateVector({100ll})); - auto buf0 = ::tflite::CreateBuffer(builder_, CreateDataVector({})); - auto buf1 = - ::tflite::CreateBuffer(builder_, CreateDataVector({1.0f, 2.0f})); - auto buf2 = - ::tflite::CreateBuffer(builder_, CreateDataVector({3.0f})); - auto buffers = builder_.CreateVector( - std::vector>({buf0, buf1, buf2})); auto t1 = ::tflite::CreateTensor(builder_, builder_.CreateVector({1, 2, 3, 4}), ::tflite::TensorType_FLOAT32, 1, @@ -62,17 +61,28 @@ class ImportTest : public ::testing::Test { ::tflite::CreateTensor(builder_, builder_.CreateVector({2, 1}), ::tflite::TensorType_FLOAT32, 2, builder_.CreateString("tensor_two"), q); - auto tensors = builder_.CreateVector( - std::vector>({t1, t2})); + return builder_.CreateVector( + std::vector>({t1, t2})); + } - // The operator codes. + Offset>> BuildOpCodes() { auto c1 = ::tflite::CreateOperatorCode(builder_, ::tflite::BuiltinOperator_CUSTOM, builder_.CreateString("custom_op_one")); auto c2 = ::tflite::CreateOperatorCode( builder_, ::tflite::BuiltinOperator_CONV_2D, 0); - auto opcodes = builder_.CreateVector( - std::vector>({c1, c2})); + return builder_.CreateVector( + std::vector>({c1, c2})); + } + + // This is a very simplistic model. We are not interested in testing all the + // details here, since tf.mini's testing framework will be exercising all the + // conversions multiple times, and the conversion of operators is tested by + // separate unittests. + void BuildTestModel() { + auto buffers = BuildBuffers(); + auto tensors = BuildTensors(); + auto opcodes = BuildOpCodes(); auto subgraph = ::tflite::CreateSubGraph(builder_, tensors, 0, 0, 0); std::vector> subgraph_vector( @@ -133,6 +143,19 @@ TEST_F(ImportTest, Tensors) { EXPECT_EQ(100, q->zero_point); } +TEST_F(ImportTest, NoSubGraphs) { + auto buffers = BuildBuffers(); + auto opcodes = BuildOpCodes(); + auto subgraphs = 0; // no subgraphs in this model + auto comment = builder_.CreateString(""); + builder_.Finish(::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION, + opcodes, subgraphs, comment, buffers)); + input_model_ = ::tflite::GetModel(builder_.GetBufferPointer()); + + EXPECT_DEATH(Import(ModelFlags(), InputModelAsString()), + "Number of subgraphs in tflite should be exactly 1."); +} + // TODO(ahentz): still need tests for Operators and IOTensors. } // namespace diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc index 461494fd99e8d09b3b9fa0add87be64a89faee26..f2cc4ef71f71902e363ac4cddd3695446af30c7d 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator.cc +++ b/tensorflow/contrib/lite/toco/tflite/operator.cc @@ -524,6 +524,28 @@ class Transpose 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: @@ -540,6 +562,24 @@ class Mean : 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 { @@ -561,15 +601,21 @@ class Squeeze } }; -class Split : public CustomOperator { +class Split + : 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::CreateSplitOptions(*builder, op.num_split); } - 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 { + op->num_split = options.num_splits(); } }; @@ -597,6 +643,20 @@ class StridedSlice } }; +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; @@ -755,17 +815,24 @@ std::vector> BuildOperatorList() { 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)); @@ -787,12 +854,13 @@ std::vector> BuildOperatorList() { 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)); + ops.emplace_back(new SimpleOperator( + "LOG_SOFTMAX", OperatorType::kLogSoftmax)); return ops; } diff --git a/tensorflow/contrib/lite/toco/tflite/operator_test.cc b/tensorflow/contrib/lite/toco/tflite/operator_test.cc index 6daa296282e381b8ff22cf75eb5e8fb87c68be9e..9c19f8d4649acf40fdd85b78874f7b18798533f2 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/operator_test.cc @@ -104,10 +104,11 @@ TEST_F(OperatorTest, SimpleOperators) { CheckSimpleOperator("RELU", OperatorType::kRelu); 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); + CheckSimpleOperator("LOG_SOFTMAX", + OperatorType::kLogSoftmax); } TEST_F(OperatorTest, BuiltinAdd) { @@ -331,6 +332,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; @@ -373,6 +382,13 @@ TEST_F(OperatorTest, StridedSlice) { 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/types.cc b/tensorflow/contrib/lite/toco/tflite/types.cc index b4c2851502a40a1ca36965d4ddd2c8a15b8fe60f..0afd2f3df57caf3214dd198bfa2ee75fa7a8fd7b 100644 --- a/tensorflow/contrib/lite/toco/tflite/types.cc +++ b/tensorflow/contrib/lite/toco/tflite/types.cc @@ -90,6 +90,8 @@ flatbuffers::Offset> DataBuffer::Serialize( return CopyBuffer(array, builder); case ArrayDataType::kInt32: return CopyBuffer(array, builder); + case ArrayDataType::kInt64: + return CopyBuffer(array, builder); case ArrayDataType::kString: return CopyBuffer(array, builder); case ArrayDataType::kUint8: diff --git a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc index c5a62fdb620ee7d6b7195f6e8e2bc3cb208feb10..0f67c2de728532b5b8101b3514811a78a3b3bc38 100644 --- a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc +++ b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc @@ -112,6 +112,11 @@ bool ParseTocoFlagsFromCommandLineFlags( "If true, ignore control dependency requirements in input TensorFlow " "GraphDef. Otherwise an error will be raised upon control dependency " "inputs."), + Flag("debug_disable_recurrent_cell_fusion", + parsed_flags.debug_disable_recurrent_cell_fusion.bind(), + parsed_flags.debug_disable_recurrent_cell_fusion.default_value(), + "If true, disable fusion of known identifiable cell subgraphs into " + "cells. This includes, for example, specific forms of LSTM cell."), }; bool asked_for_help = *argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-help")); diff --git a/tensorflow/contrib/lite/toco/toco_flags.proto b/tensorflow/contrib/lite/toco/toco_flags.proto index 3b9d7e22570b66aef2c9fc819e5ab4ec38e179f5..3237147a736f97f65953ca965420fcea934820a4 100644 --- a/tensorflow/contrib/lite/toco/toco_flags.proto +++ b/tensorflow/contrib/lite/toco/toco_flags.proto @@ -36,7 +36,8 @@ enum FileFormat { // are not normally encoded in model files and in general may not be thought // of as properties of models, instead describing how models are to be // processed in the context of the present tooling job. -// Next Id: 13 +// +// Next ID to use: 14. message TocoFlags { // Input file format optional FileFormat input_format = 1; @@ -136,4 +137,8 @@ message TocoFlags { // - Default to false if the output format is TENSORFLOW_GRAPHDEF. // - Default to true in all other cases. optional bool drop_control_dependency = 12; + + // Disables transformations that fuse subgraphs such as known LSTMs (not all + // LSTMs are identified). + optional bool debug_disable_recurrent_cell_fusion = 13; } diff --git a/tensorflow/contrib/lite/toco/toco_tooling.cc b/tensorflow/contrib/lite/toco/toco_tooling.cc index b715881774bf00f6cf2a50452a5b5c59c647ade6..024335b5e4ca4547eba24b0944dc539eddc6bdf7 100644 --- a/tensorflow/contrib/lite/toco/toco_tooling.cc +++ b/tensorflow/contrib/lite/toco/toco_tooling.cc @@ -52,19 +52,25 @@ void MakeGeneralGraphTransformationsSet( GraphTransformationsSet* transformations) { CHECK(transformations->empty()); transformations->Add(new ConvertExpandDimsToReshape); + transformations->Add(new ConvertSqueezeToReshape); 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 PropagateActivationFunctionIntoConstants); 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); @@ -72,15 +78,19 @@ void MakeGeneralGraphTransformationsSet( transformations->Add(new ResolveBatchNormalization); transformations->Add(new ResolveConstantBinaryOperator); transformations->Add(new ResolveConstantFill); + transformations->Add(new ResolveConstantGather); 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 ResolveSqueezeAttributes); transformations->Add(new ResolveTensorFlowSwitch); transformations->Add(new ResolveTensorFlowTile); transformations->Add(new ResolveTensorFlowConcat); + transformations->Add(new ResolveMultiplyByZero); + transformations->Add(new IdentifyDilatedConv); transformations->Add(new IdentifyL2Normalization); transformations->Add(new IdentifyL2Pool); transformations->Add(new IdentifyRelu1); @@ -92,9 +102,10 @@ void MakeGeneralGraphTransformationsSet( transformations->Add(new ResolveStridedSliceAttributes); transformations->Add(new ResolveSliceAttributes); transformations->Add(new ResolveMeanAttributes); - transformations->Add(new ResolveTransposeAttributes); transformations->Add(new ResolveConstantShapeOrRank); transformations->Add(new MakeInitialDequantizeOperator); + transformations->Add(new ResolveConstantFakeQuant); + transformations->Add(new UnpartitionEmbeddingLookup); } bool SupportsQuantization(FileFormat format) { @@ -106,7 +117,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) { @@ -182,20 +194,23 @@ 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(); const bool quantize_output = - SupportsQuantization(output_format) && inference_type == QUANTIZED_UINT8; + SupportsQuantization(output_format) && + (inference_type == QUANTIZED_UINT8 || inference_type == QUANTIZED_INT16); if (quantize_output) { QCHECK_NE(toco_flags.inference_input_type(), FLOAT) << "Quantized inference is not allowed with float inputs."; } - SetFinalDataTypeOnInputs(toco_flags, model); - UseArraysExtraInfo(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 @@ -212,9 +227,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 { @@ -227,9 +239,15 @@ 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 (!toco_flags.debug_disable_recurrent_cell_fusion()) { + 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", @@ -267,6 +285,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/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc index d0bf66c542bfde849018d6302e900eb45c1cfdd1..e70291ad0ef4bb4c1969cfc3c87bf0f8e50ecef5 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" @@ -61,6 +62,37 @@ string LogName(const Operator& op) { } } +string ArrayDataTypeName(ArrayDataType data_type) { + switch (data_type) { + case ArrayDataType::kFloat: + return "Float"; + case ArrayDataType::kInt8: + return "Int8"; + case ArrayDataType::kUint8: + return "Uint8"; + case ArrayDataType::kInt16: + return "Int16"; + case ArrayDataType::kUint16: + return "Uint16"; + case ArrayDataType::kInt32: + return "Int32"; + case ArrayDataType::kUint32: + return "Uint32"; + case ArrayDataType::kInt64: + return "Int64"; + case ArrayDataType::kUint64: + return "Uint64"; + case ArrayDataType::kString: + return "String"; + case ArrayDataType::kBool: + return "Bool"; + case ArrayDataType::kNone: + return "None"; + default: + LOG(FATAL) << "Unhandled array data type " << static_cast(data_type); + } +} + bool IsInputArray(const Model& model, const string& name) { for (const auto& input_array : model.flags.input_arrays()) { if (input_array.name() == name) { @@ -110,13 +142,32 @@ int CountOpsWithInput(const Model& model, const string& array_name) { } bool DeleteArrayIfUnused(const string& array_name, Model* model) { - if (CountOpsWithInput(*model, array_name) == 0) { + 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; +} + +void DeleteOpAndArraysIfUnused(Model* model, Operator* op) { + for (const string& array_name : op->inputs) { + DeleteArrayIfUsedOnce(array_name, model); + } + auto op_it = FindOp(*model, op); + CHECK(op_it != model->operators.end()); + model->operators.erase(op_it); +} + std::vector>::const_iterator FindOpWithOutput( const Model& model, const string& array_name) { for (auto it = model.operators.begin(); it != model.operators.end(); ++it) { @@ -159,6 +210,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) { @@ -217,6 +280,7 @@ const char* OperatorTypeName(OperatorType type) { 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) @@ -238,6 +302,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) @@ -288,7 +353,11 @@ 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) + HANDLE_OPERATORTYPENAME_CASE(DynamicPartition) + HANDLE_OPERATORTYPENAME_CASE(DynamicStitch) default: LOG(FATAL) << "Unhandled op type"; #undef HANDLE_OPERATORTYPENAME_CASE @@ -307,6 +376,7 @@ string HelpfulOperatorTypeName(const Operator& op) { bool OperatorSupportsFusedActivation(OperatorType type) { switch (type) { case OperatorType::kConcatenation: + case OperatorType::kGather: case OperatorType::kSlice: case OperatorType::kSqueeze: case OperatorType::kTensorFlowReshape: @@ -335,48 +405,9 @@ void LogSummary(int log_level, const Model& model) { void LogArray(int log_level, const Model& model, const string& name) { const auto& array = model.GetArray(name); VLOG(log_level) << "Array: " << name; - switch (array.data_type) { - case ArrayDataType::kNone: - VLOG(log_level) << " Data type:"; - break; - case ArrayDataType::kFloat: - VLOG(log_level) << " Data type: kFloat"; - break; - case ArrayDataType::kInt32: - VLOG(log_level) << " Data type: kInt32"; - break; - case ArrayDataType::kUint8: - VLOG(log_level) << " Data type: kUint8"; - break; - case ArrayDataType::kString: - VLOG(log_level) << " Data type: kString"; - break; - default: - VLOG(log_level) << " Data type: other (numerical value: " - << static_cast(array.data_type) << ")"; - break; - } - switch (array.final_data_type) { - case ArrayDataType::kNone: - VLOG(log_level) << " Final type:"; - break; - case ArrayDataType::kFloat: - VLOG(log_level) << " Final type: kFloat"; - break; - case ArrayDataType::kInt32: - VLOG(log_level) << " Final type: kInt32"; - break; - case ArrayDataType::kUint8: - VLOG(log_level) << " Final type: kUint8"; - break; - case ArrayDataType::kString: - VLOG(log_level) << " Final type: kString"; - break; - default: - VLOG(log_level) << " Final type: other (numerical value: " - << static_cast(array.data_type) << ")"; - break; - } + VLOG(log_level) << " Data type: " << ArrayDataTypeName(array.data_type); + VLOG(log_level) << " Final type: " + << ArrayDataTypeName(array.final_data_type); if (array.buffer) { VLOG(log_level) << " Constant Buffer"; } @@ -406,7 +437,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; } } @@ -606,6 +637,14 @@ bool IsConstantParameterArray(const Model& model, const string& name) { } 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()) { @@ -685,12 +724,10 @@ void CheckNoMissingArray(const Model& model) { for (const auto& op : model.operators) { for (const auto& input : op->inputs) { CHECK(model.HasArray(input) || model.optional_arrays.count(input)) - << "Input: " << input << " missing for op: " - << op->outputs[0] << "."; + << "Input: " << input << " missing for op: " << op->outputs[0] << "."; } for (const auto& output : op->outputs) { - CHECK(model.HasArray(output)) << "Output: " << output - << " missing."; + CHECK(model.HasArray(output)) << "Output: " << output << " missing."; } } CheckNonExistentIOArrays(model); @@ -771,7 +808,10 @@ void FixNoOrphanedArray(Model* model) { } } -void CheckArrayFieldsConsistent(const Model& model) { +// 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()) { @@ -782,10 +822,28 @@ void CheckArrayFieldsConsistent(const Model& model) { // It's OK to have a buffer or an alloc, but not both. // (Since allocs are for transient arrays without a buffer). CHECK(!array->buffer || !array->alloc); - // If there is a buffer, its type should be consistent with data_type. if (array->buffer) { + // If there is a buffer, its type should be consistent with data_type. CHECK(array->buffer->type == array->data_type); + // The presence of a fixed buffer should imply the presence of a fixed + // shape. + CHECK(array->has_shape()); + // The shape flat-size should agree with the buffer length. + CHECK_EQ(array->buffer->Length(), + RequiredBufferSizeForShape(array->shape())); + } + + // 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."; } } @@ -934,7 +992,7 @@ void CheckInvariants(const Model& model) { CheckNonAsciiIOArrays(model.flags); CheckNoMissingArray(model); CheckNoOrphanedArray(model); - CheckArrayFieldsConsistent(model); + CheckEachArray(model); CheckOperatorOrdering(model); } @@ -1026,9 +1084,6 @@ void CreateOrCheckRnnStateArray(const string& name, int size, Model* model) { if (array.has_shape()) { num_dims = array.shape().dimensions_count(); } - CHECK(array.data_type == ArrayDataType::kFloat || - array.data_type == ArrayDataType::kNone); - array.data_type = ArrayDataType::kFloat; if (!array.has_shape() && num_dims >= 0) { Shape* shape = array.mutable_shape(); std::vector dims; @@ -1052,7 +1107,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 && @@ -1155,7 +1210,7 @@ void ResolveModelFlags(const ModelFlags& model_flags, Model* model) { << "This model does not define output arrays, so a " "--output_arrays flag must be given on the command-line."; - for (const auto& input_array_proto : model->flags.input_arrays()) { + for (auto& input_array_proto : *model->flags.mutable_input_arrays()) { auto& input_array = model->GetOrCreateArray(input_array_proto.name()); if (input_array_proto.has_data_type()) { const ArrayDataType specified_type = @@ -1199,6 +1254,11 @@ void ResolveModelFlags(const ModelFlags& model_flags, Model* model) { for (int i = 0; i < input_array_dims.size(); i++) { CHECK_EQ(input_array_dims[i], input_array_proto.shape().dims(i)); } + } else { + for (int i = 0; i < input_array.shape().dimensions_count(); i++) { + input_array_proto.mutable_shape()->add_dims( + input_array.shape().dims(i)); + } } } @@ -1294,21 +1354,34 @@ void UseDefaultMinMaxRangeValues(Model* model, double default_ranges_min, int ElementSize(ArrayDataType data_type) { switch (data_type) { + case ArrayDataType::kBool: + return sizeof(bool); 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: LOG(FATAL) << "Transient arrays with strings are not supported yet"; return 0; default: - LOG(FATAL) << "Should not get here."; + LOG(FATAL) << "Unknown data_type = " << static_cast(data_type); return 0; } } @@ -1348,19 +1421,23 @@ bool IsAllocatableTransientArray(const Model& model, const string& array_name) { } string AvailableArrayName(const Model& model, const string& name) { - if (!model.HasArray(name) && !model.IsOptionalArray(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); + 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 ""; } @@ -1399,6 +1476,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) { @@ -1446,6 +1538,7 @@ bool EstimateArithmeticOpsCount(const Model& model, int64* result) { } 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()) { @@ -1545,10 +1638,6 @@ void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order, } } -namespace { - -// 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; @@ -1563,8 +1652,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 && @@ -1724,7 +1811,10 @@ bool IsDiscardableArray(const Model& model, const string& array_name) { void CheckFinalDataTypesSatisfied(const Model& model) { for (const auto& array_entry : model.GetArrayMap()) { const auto& array = *array_entry.second; - if (array.final_data_type != ArrayDataType::kNone) { + // If the final data type is int16, the data type may be float, for example + // after dequantization. + if (array.final_data_type != ArrayDataType::kNone && + array.final_data_type != ArrayDataType::kInt16) { CHECK(array.final_data_type == array.data_type) << "Array \"" << array_entry.first << "\" has mis-matching actual and final data types (" @@ -1740,6 +1830,8 @@ ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type) { return ArrayDataType::kFloat; case QUANTIZED_UINT8: return ArrayDataType::kUint8; + case QUANTIZED_INT16: + return ArrayDataType::kInt16; case INT32: return ArrayDataType::kInt32; case INT64: @@ -1749,14 +1841,39 @@ 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(); + auto& array = model->GetArray(entry.name()); + auto& minmax = array.GetOrCreateMinMax(); + if (entry.has_min() || entry.has_max()) { + CHECK_EQ(entry.has_min(), entry.has_max()); + minmax.min = entry.min(); + minmax.max = entry.max(); + } + if (entry.has_data_type()) { + array.final_data_type = + ConvertIODataTypeToArrayDataType(entry.data_type()); + } } } diff --git a/tensorflow/contrib/lite/toco/tooling_util.h b/tensorflow/contrib/lite/toco/tooling_util.h index a7e77a02eb1f78bc2d4841d1bccdc9eac3e31a95..05360e3b0a8510f57724d98dab0e2097b1062ced 100644 --- a/tensorflow/contrib/lite/toco/tooling_util.h +++ b/tensorflow/contrib/lite/toco/tooling_util.h @@ -28,6 +28,7 @@ limitations under the License. #if TOCO_SUPPORT_PORTABLE_PROTOS #include "third_party/protobuf/src/google/protobuf/text_format.h" #endif // TOCO_SUPPORT_PORTABLE_PROTOS +#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/model_flags.pb.h" #include "tensorflow/contrib/lite/toco/runtime/types.h" @@ -54,12 +55,19 @@ absl::string_view FindLongestCommonPrefix(absl::string_view a, absl::string_view b); string LogName(const Operator& op); +string ArrayDataTypeName(ArrayDataType data_type); + bool IsInputArray(const Model& model, const string& name); bool IsArrayConsumed(const Model& model, const string& name); 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); + +// Deletes the op and any of its input and output arrays if they are unused +// after the op has been deleted. +void DeleteOpAndArraysIfUnused(Model* model, Operator* op); std::vector>::const_iterator FindOpWithOutput( const Model& model, const string& array_name); @@ -67,10 +75,13 @@ 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); @@ -136,75 +147,14 @@ void FixNoOrphanedArray(Model* model); void ResolveModelFlags(const ModelFlags& model_flags, Model* model); template -void GetQuantizationParamsFromMinMax(const ModelFlags& model_flags, - const MinMax& minmax, +void GetQuantizationParamsFromMinMax(const MinMax& minmax, QuantizationParams* quantization_params) { using Integer = DataType; - const Integer qmin = std::numeric_limits::min(); - const Integer qmax = std::numeric_limits::max(); - const double qmin_double = qmin; - const double qmax_double = qmax; const double rmin = minmax.min; const double rmax = minmax.max; - // 0 should always be a representable value. Let's assume that the initial - // min,max range contains 0. - CHECK_LE(rmin, 0.); - CHECK_GE(rmax, 0.); - if (rmin == rmax) { - // Special case where the min,max range is a point. Should be {0}. - CHECK_EQ(rmin, 0.); - CHECK_EQ(rmax, 0.); - quantization_params->zero_point = 0; - quantization_params->scale = 0.; - return; - } - // General case. - // - // First determine the scale. - const double scale = (rmax - rmin) / (qmax_double - qmin_double); - - // Zero-point computation. - // First the initial floating-point computation. The zero-point can be - // determined from solving an affine equation for any known pair - // (real value, corresponding quantized value). - // We know two such pairs: (rmin, qmin) and (rmax, qmax). - // The arithmetic error on the zero point computed from either pair - // will be roughly machine_epsilon * (sum of absolute values of terms) - // so we want to use the variant that adds the smaller terms. - const double zero_point_from_min = qmin_double - rmin / scale; - const double zero_point_from_max = qmax_double - rmax / scale; - const double zero_point_from_min_error = - std::abs(qmin_double) + std::abs(rmin / scale); - const double zero_point_from_max_error = - std::abs(qmax_double) + std::abs(rmax / scale); - - const double zero_point_double = - zero_point_from_min_error < zero_point_from_max_error - ? zero_point_from_min - : zero_point_from_max; - - // Now we need to nudge the zero point to be an integer - // (our zero points are integer, and this is motivated by the requirement - // to be able to represent the real value "0" exactly as a quantized value, - // which is required in multiple places, for example in Im2col with SAME - // padding). - Integer nudged_zero_point = 0; - if (zero_point_double < qmin_double) { - nudged_zero_point = qmin; - } else if (zero_point_double > qmax_double) { - nudged_zero_point = qmax; - } else { - nudged_zero_point = static_cast(std::round(zero_point_double)); - } - // The zero point should always be in the range of quantized value, - // [qmin, qmax]. - CHECK_GE(nudged_zero_point, qmin); - CHECK_LE(nudged_zero_point, qmax); - - // Finally, store the result nudged quantization params. - quantization_params->zero_point = nudged_zero_point; - quantization_params->scale = scale; + *quantization_params = + ::tflite::ChooseQuantizationParams(rmin, rmax); } void CheckIsReadyForQuantization(const Model& model); @@ -255,6 +205,11 @@ 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); @@ -264,6 +219,11 @@ int AxesCount(AxesOrder 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, @@ -283,6 +243,23 @@ 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 diff --git a/tensorflow/contrib/lite/toco/types.proto b/tensorflow/contrib/lite/toco/types.proto index 318fd4b7b2c2df093562e73c3fe707675ee98876..03bd6150bc86bb27221814cd191b17f1a09585fa 100644 --- a/tensorflow/contrib/lite/toco/types.proto +++ b/tensorflow/contrib/lite/toco/types.proto @@ -34,4 +34,7 @@ enum IODataType { // String, not quantized STRING = 5; + + // Int16, quantized + QUANTIZED_INT16 = 6; } diff --git a/tensorflow/contrib/lite/tools/BUILD b/tensorflow/contrib/lite/tools/BUILD index 1bffcfb987330c5d067d7f986a486fcf93e57ee7..b5abbc0712599814e078d19bc015bc7bf1812f95 100644 --- a/tensorflow/contrib/lite/tools/BUILD +++ b/tensorflow/contrib/lite/tools/BUILD @@ -4,6 +4,7 @@ package(default_visibility = [ licenses(["notice"]) # Apache 2.0 +load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") load("//tensorflow:tensorflow.bzl", "tf_cc_binary") py_binary( @@ -45,7 +46,15 @@ tf_cc_binary( "//tensorflow/contrib/lite:framework", "//tensorflow/contrib/lite:string_util", "//tensorflow/contrib/lite/kernels:builtin_ops", - ], + ] + select({ + "//tensorflow:android": [ + "//tensorflow/core:android_tensorflow_lib", + ], + "//conditions:default": [ + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + ], + }), ) cc_library( @@ -99,8 +108,11 @@ cc_library( 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", ], ) @@ -108,13 +120,21 @@ cc_test( name = "verifier_test", size = "small", srcs = ["verifier_test.cc"], + tags = [ + "tflite_not_portable", + ], 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", ], ) + +tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/tools/benchmark_model.cc b/tensorflow/contrib/lite/tools/benchmark_model.cc index 6ae3ab57294a92162b15f326630ac202a9ba2a82..93c80e0f5e021f76bff6858b0ea3370724393d6d 100644 --- a/tensorflow/contrib/lite/tools/benchmark_model.cc +++ b/tensorflow/contrib/lite/tools/benchmark_model.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. @@ -25,36 +25,89 @@ limitations under the License. #include "tensorflow/contrib/lite/model.h" #include "tensorflow/contrib/lite/string_util.h" #include "tensorflow/contrib/lite/tools/mutable_op_resolver.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/util/command_line_flags.h" #ifdef TFLITE_CUSTOM_OPS_HEADER void RegisterSelectedOps(::tflite::MutableOpResolver* resolver); #endif -#define LOG(x) std::cerr +namespace tflite { -#define CHECK(x) \ - if (!(x)) { \ - LOG(ERROR) << #x << "failed"; \ - exit(1); \ +using ::tensorflow::Env; +using ::tensorflow::str_util::Split; +using ::tensorflow::str_util::SplitAndParseAsFloats; +using ::tensorflow::str_util::SplitAndParseAsInts; + +struct InputLayerInfo { + string name; + TfLiteType data_type; + std::vector shape; + // Note that initialization_values is currently unused. + std::vector initialization_values; +}; + +template +void FillRandomValue(T* ptr, const std::vector& sizes, + const std::function& random_func) { + int num_elements = 1; + for (int dim : sizes) { + num_elements *= dim; + } + for (int i = 0; i < num_elements; ++i) { + *ptr++ = random_func(); } +} -namespace tensorflow { -namespace benchmark_tflite_model { +void FillRandomString(tflite::DynamicBuffer* buffer, + const std::vector& sizes, + const std::function& random_func) { + int num_elements = 1; + for (int dim : sizes) { + num_elements *= dim; + } + for (int i = 0; i < num_elements; ++i) { + auto str = random_func(); + buffer->AddString(str.data(), str.length()); + } +} -std::unique_ptr model; -std::unique_ptr interpreter; +TfLiteType TfLiteTypeFromString(const string& input_layer_type) { + if (input_layer_type == "string") + return kTfLiteString; + else if (input_layer_type == "float") + return kTfLiteFloat32; + else if (input_layer_type == "uint8") + return kTfLiteUInt8; + else if (input_layer_type == "int32") + return kTfLiteInt32; + else if (input_layer_type == "int64") + return kTfLiteInt64; + else + return kTfLiteNoType; +} -void InitImpl(const std::string& graph, const std::vector& sizes, - const std::string& input_layer_type, int num_threads) { - CHECK(graph.c_str()); +std::vector ShapeFromTfLiteTensor(TfLiteTensor* t) { + std::vector result; + result.reserve(t->dims->size); + for (int i = 0; i < t->dims->size; ++i) { + result.push_back(t->dims->data[i]); + } + CHECK(!result.empty()) << "Found no shapes in model"; + return result; +} - model = tflite::FlatBufferModel::BuildFromFile(graph.c_str()); +bool CreateInterpreter(const string& graph, + std::unique_ptr* model, + std::unique_ptr* interpreter) { + *model = tflite::FlatBufferModel::BuildFromFile(graph.c_str()); if (!model) { - LOG(FATAL) << "Failed to mmap model " << graph; + std::cerr << "Failed to load model " << graph << std::endl; + return false; } - LOG(INFO) << "Loaded model " << graph; - model->error_reporter(); - LOG(INFO) << "resolved reporter"; #ifdef TFLITE_CUSTOM_OPS_HEADER tflite::MutableOpResolver resolver; @@ -63,34 +116,360 @@ void InitImpl(const std::string& graph, const std::vector& sizes, tflite::ops::builtin::BuiltinOpResolver resolver; #endif - tflite::InterpreterBuilder(*model, resolver)(&interpreter); - if (!interpreter) { - LOG(FATAL) << "Failed to construct interpreter"; + tflite::InterpreterBuilder(*(model->get()), resolver)(interpreter); + if (!(*interpreter)) { + std::cerr << "Failed to construct interpreter" << std::endl; + return false; } + return true; +} + +bool PrepareInterpreter(const std::vector inputs, + int num_threads, bool use_nnapi, + Interpreter* interpreter) { if (num_threads != -1) { interpreter->SetNumThreads(num_threads); } - int input = interpreter->inputs()[0]; + interpreter->UseNNAPI(use_nnapi); - if (input_layer_type != "string") { - interpreter->ResizeInputTensor(input, sizes); + // Check that all names and types match + for (const InputLayerInfo& input : inputs) { + for (int i : interpreter->inputs()) { + TfLiteTensor* t = interpreter->tensor(i); + CHECK_EQ(t->name, input.name) + << "Tensor # " << i << " is named " << t->name + << " but flags call it " << input.name; + CHECK_EQ(t->type, input.data_type) + << "Could not match the type of input tensor " << t->name; + } + } + + // Resize all non-string tensors. + for (const InputLayerInfo& input : inputs) { + for (int i : interpreter->inputs()) { + TfLiteTensor* t = interpreter->tensor(i); + if (t->type != kTfLiteString) { + interpreter->ResizeInputTensor(i, input.shape); + } + } } if (interpreter->AllocateTensors() != kTfLiteOk) { - LOG(FATAL) << "Failed to allocate tensors!"; + std::cerr << "Failed to allocate tensors!" << std::endl; + return false; + } + + // Set the values of the input tensors. + for (int i : interpreter->inputs()) { + TfLiteTensor* t = interpreter->tensor(i); + std::vector sizes = ShapeFromTfLiteTensor(t); + + // TODO(ahentz): below we ignore the O-th dimension (number of batches). + if (t->type == kTfLiteFloat32) { + FillRandomValue( + interpreter->typed_tensor(i), + std::vector(sizes.begin() + 1, sizes.end()), + []() { return static_cast(rand()) / RAND_MAX - 0.5f; }); + } else if (t->type == kTfLiteUInt8) { + FillRandomValue( + interpreter->typed_tensor(i), + std::vector(sizes.begin() + 1, sizes.end()), + []() { return static_cast(rand()) % 255; }); + } else if (t->type == kTfLiteString) { + tflite::DynamicBuffer buffer; + FillRandomString(&buffer, sizes, []() { + return "we're have some friends over saturday to hang out in the yard"; + }); + buffer.WriteToTensor(interpreter->tensor(i)); + } else { + std::cerr << "Don't know how to populate tensor " << t->name + << " of type " << t->type << std::endl; + return false; + } + } + return true; +} + +bool PopulateInputLayerInfo(const string& names_string, + const string& shapes_string, + const string& types_string, + const string& values_string, + std::vector* info) { + std::vector names = Split(names_string, ','); + std::vector shapes = Split(shapes_string, ':'); + std::vector types = Split(types_string, ','); + std::vector values = Split(values_string, ':'); + + if (names.size() != shapes.size()) { + LOG(ERROR) << "The number of items in" + << " --input_layer_shape (" << shapes_string << ", with " + << shapes.size() << " items)" + << " must match the number of items in" + << " --input_layer (" << names_string << ", with " + << names.size() << " items)." + << " For example --input_layer=input1,input2" + << " --input_layer_shape=1,224,224,4:1,20"; + return false; + } + if (names.size() != types.size()) { + LOG(ERROR) << "The number of items in" + << " --input_layer_type (" << types_string << ", with " + << types.size() << " items)" + << " must match the number of items in" + << " --input_layer (" << names_string << ", with " + << names.size() << " items)." + << " For example --input_layer=input1,input2" + << " --input_layer_type=float,int"; + return false; + } + + for (int i = 0; i < names.size(); ++i) { + info->push_back(InputLayerInfo()); + InputLayerInfo& input = info->back(); + + input.name = names[i]; + + input.data_type = TfLiteTypeFromString(types[i]); + CHECK(input.data_type != kTfLiteNoType) + << types[i] << " was an invalid type"; + + CHECK(SplitAndParseAsInts(shapes[i], ',', &input.shape)) + << "Incorrect size string specified: " << shapes[i]; + for (int dim : input.shape) { + if (dim == -1) { + LOG(ERROR) << "Any unknown sizes in the shapes (-1's) must be replaced" + << " with the size you want to benchmark with."; + return false; + } + } + + if (i < values.size()) { + CHECK(SplitAndParseAsFloats(values[i], ',', &input.initialization_values)) + << "Incorrect initialization values string specified: " << values[i]; + } + } + + return true; +} + +bool RunBenchmark(Interpreter* interpreter, int64_t* inference_time_us) { + const int64_t start_time = Env::Default()->NowMicros(); + + if (interpreter->Invoke() != kTfLiteOk) { + std::cerr << "Failed to invoke!"; + return false; } + + const int64_t end_time = Env::Default()->NowMicros(); + *inference_time_us = end_time - start_time; + return true; +} + +class Latencies { + public: + void AddMeasurement(int64_t time_us) { + max_ = std::max(time_us, max_); + min_ = std::min(time_us, min_); + ++count_; + sum_ += time_us; + squared_sum_ += static_cast(time_us) * time_us; + } + + double avg() const { + if (count_ == 0) return std::numeric_limits::quiet_NaN(); + return static_cast(sum_) / count_; + } + + int64_t std_deviation() const { + if (count_ == 0 || min_ == max_) return 0; + return sqrt(squared_sum_ / count_ - avg() * avg()); + } + + void OutputToStream(std::ostream* stream) const { + *stream << "count=" << count_; + if (count_ == 0) return; + *stream << " min=" << min_ << " max=" << max_; + *stream << " avg=" << avg() << " std=" << std_deviation(); + } + + private: + int64_t count_ = 0; + int64_t min_ = std::numeric_limits::max(); + int64_t max_ = std::numeric_limits::min(); + int64_t sum_ = 0; + double squared_sum_ = 0; +}; + +bool TimeMultipleRuns(Interpreter* interpreter, double sleep_seconds, + int num_runs, int64* total_time_us) { + // Convert the run_delay string into a timespec. + timespec req; + req.tv_sec = static_cast(sleep_seconds); + req.tv_nsec = (sleep_seconds - req.tv_sec) * 1000000000; + + *total_time_us = 0; + + std::cout << "Running benchmark for " << num_runs + << " iterations: " << std::endl; + + Latencies latencies; + for (int i = 0; i < num_runs; ++i) { + int64_t time_us; + bool run_status = RunBenchmark(interpreter, &time_us); + latencies.AddMeasurement(time_us); + *total_time_us += time_us; + if (!run_status) { + std::cout << "Failed on run " << i << std::endl; + return false; + } + + // If requested, sleep between runs for an arbitrary amount of time. + // This can be helpful to determine the effect of mobile processor + // scaling and thermal throttling. + if (sleep_seconds > 0.0) { +#ifdef PLATFORM_WINDOWS + Sleep(sleep_seconds * 1000); +#else + nanosleep(&req, nullptr); +#endif + } + } + latencies.OutputToStream(&std::cout); + std::cout << std::endl; + + return true; } int Main(int argc, char** argv) { - InitImpl("", {}, "", 1); + using tensorflow::Flag; + using tensorflow::Flags; + + string graph; // e.g.: /data/local/tmp/tfl_inception-v1_model.fb + string input_layer_string; // e.g.: input + string input_layer_shape_string; // e.g.: 1,224,224,3 + string input_layer_type_string; // e.g.: float + string input_layer_values_string; + string output_layer_string; // e.g.: output + int num_runs = 50; + string run_delay = "-1.0"; + int num_threads = -1; + string benchmark_name = ""; + string output_prefix = ""; + int warmup_runs = 1; + bool use_nnapi = false; + + std::vector flag_list = { + Flag("graph", &graph, "graph file name"), + // All the following flags are optional, but can be used in order + // to benchmark different input shapes. + Flag("input_layer", &input_layer_string, "input layer names"), + Flag("input_layer_shape", &input_layer_shape_string, "input layer shape"), + Flag("input_layer_type", &input_layer_type_string, "input layer type"), + Flag("input_layer_values", &input_layer_values_string, + "values to initialize the inputs with"), + Flag("output_layer", &output_layer_string, "output layer name"), + Flag("num_runs", &num_runs, "number of runs"), + Flag("run_delay", &run_delay, "delay between runs in seconds"), + Flag("num_threads", &num_threads, "number of threads"), + Flag("benchmark_name", &benchmark_name, "benchmark name"), + Flag("output_prefix", &output_prefix, "benchmark output prefix"), + Flag("warmup_runs", &warmup_runs, "how many runs to initialize model"), + Flag("use_nnapi", &use_nnapi, "use nnapi api"), + }; + string usage = Flags::Usage(argv[0], flag_list); + const bool parse_result = Flags::Parse(&argc, argv, flag_list); + tensorflow::port::InitMain(argv[0], &argc, &argv); + + if (!parse_result) { + std::cerr << usage << std::endl; + return -1; + } + + std::cout << "Graph: [" << graph << "]" << std::endl; + if (!input_layer_string.empty()) { + std::cout << "Input layers: [" << input_layer_string << "]" << std::endl; + std::cout << "Input shapes: [" << input_layer_shape_string << "]" + << std::endl; + std::cout << "Input types: [" << input_layer_type_string << "]" + << std::endl; + } + if (!output_layer_string.empty()) { + std::cout << "Output layers: [" << output_layer_string << "]" << std::endl; + } + std::cout << "Num runs: [" << num_runs << "]" << std::endl; + std::cout << "Inter-run delay (seconds): [" << run_delay << "]" << std::endl; + std::cout << "Num threads: [" << num_threads << "]" << std::endl; + if (!benchmark_name.empty()) { + std::cout << "Benchmark name: [" << benchmark_name << "]" << std::endl; + std::cout << "Output prefix: [" << output_prefix << "]" << std::endl; + } + std::cout << "Warmup runs: [" << warmup_runs << "]" << std::endl; + std::cout << "Use nnapi : [" << use_nnapi << "]" << std::endl; + + if (graph.empty()) { + std::cout + << "Please specify the name of your TF Lite input file with --graph" + << std::endl; + return -1; + } + + std::vector inputs; + if (!PopulateInputLayerInfo(input_layer_string, input_layer_shape_string, + input_layer_type_string, + input_layer_values_string, &inputs)) { + return -1; + } + + int64 initialization_start_us = Env::Default()->NowMicros(); + + std::unique_ptr model; + std::unique_ptr interpreter; + if (!CreateInterpreter(graph, &model, &interpreter)) { + return -1; + } + if (!PrepareInterpreter(inputs, num_threads, use_nnapi, interpreter.get())) { + return -1; + } + + int64 initialization_end_us = Env::Default()->NowMicros(); + + const double initialization_time_s = + (initialization_end_us - initialization_start_us) / 1000000.0f; + std::cout << "Initialized session in " << initialization_time_s << "s" + << std::endl; + + const double sleep_seconds = std::strtod(run_delay.c_str(), nullptr); + + // If requested, run through the graph first to preinitialize everything + // before the benchmarking runs. + int64 warmup_time_us = 0; + if (warmup_runs > 0) { + if (!TimeMultipleRuns(interpreter.get(), sleep_seconds, warmup_runs, + &warmup_time_us)) { + std::cerr << "Warmup failed" << std::endl; + return -1; + } + } + + // Capture overall inference time without stat logging overhead. This is the + // timing data that can be compared to other libaries. + int64 no_stat_time_us = 0; + if (!TimeMultipleRuns(interpreter.get(), sleep_seconds, num_runs, + &no_stat_time_us)) { + std::cerr << "Timing failed." << std::endl; + return -1; + } + + std::cout << "Average inference timings in us: " << no_stat_time_us / num_runs + << " , Warmup: " + << (warmup_runs > 0 ? warmup_time_us / warmup_runs : 0) << ", " + << std::endl; + return 0; } -} // namespace benchmark_tflite_model -} // namespace tensorflow +} // namespace tflite -int main(int argc, char** argv) { - return tensorflow::benchmark_tflite_model::Main(argc, argv); -} +int main(int argc, char** argv) { return ::tflite::Main(argc, argv); } diff --git a/tensorflow/contrib/lite/tools/verifier.cc b/tensorflow/contrib/lite/tools/verifier.cc index 95a0895379845d8887939e0217270b30ea5584ca..59c74205f0a311ec12ff87f46622041605fb493b 100644 --- a/tensorflow/contrib/lite/tools/verifier.cc +++ b/tensorflow/contrib/lite/tools/verifier.cc @@ -14,13 +14,32 @@ 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)) { @@ -30,14 +49,186 @@ const Model* VerifyFlatbufferAndGetModel(const void* buf, size_t len) { } } +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) { +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; } - - return model->version() == TFLITE_SCHEMA_VERSION; + 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 index 03e1f22b7e87baf6d1586dde5812fc854d9e2c4c..c2ee11215c861ed7b27696a8d786bb6e2a48e930 100644 --- a/tensorflow/contrib/lite/tools/verifier.h +++ b/tensorflow/contrib/lite/tools/verifier.h @@ -18,13 +18,18 @@ limitations under the License. #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. -bool Verify(const void* buf, size_t len); +// * 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 diff --git a/tensorflow/contrib/lite/tools/verifier_test.cc b/tensorflow/contrib/lite/tools/verifier_test.cc index 0481a55a78e5e1cd0821df13ffaf84bbe28a1b8e..b3e611f999b2837efbf8876bd989db44c408b8c7 100644 --- a/tensorflow/contrib/lite/tools/verifier_test.cc +++ b/tensorflow/contrib/lite/tools/verifier_test.cc @@ -12,7 +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/contrib/lite/tools/verifier.h" +#include +#include + #include "flatbuffers/flatbuffers.h" #include "flatbuffers/util.h" #include @@ -20,7 +22,10 @@ limitations under the License. #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 { @@ -28,31 +33,77 @@ using flatbuffers::FlatBufferBuilder; using flatbuffers::Offset; using flatbuffers::Vector; -// Class that abstracts the list of buffers at the end of the TF Lite structure -class DeferredBufferWriter { +// Build single subgraph model. +class TfLiteFlatbufferModelBuilder { public: - DeferredBufferWriter() { - data_.push_back({}); // sentinel empty buffer. + TfLiteFlatbufferModelBuilder() { + buffers_.push_back( + CreateBuffer(builder_, builder_.CreateVector(std::vector{}))); } - Offset>> BuildBuffers(FlatBufferBuilder *builder) { - std::vector> buffer_vector; - for (const auto &vec : data_) { - auto data_buffer = builder->CreateVector(vec.data(), vec.size()); - buffer_vector.push_back(tflite::CreateBuffer(*builder, data_buffer)); + 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_); } - return builder->CreateVector(buffer_vector); } - // Registers a buffer index and takes ownership of the data to write to it. - int Record(std::vector data) { - int buffer_index = data_.size(); - data_.emplace_back(std::move(data)); - return buffer_index; + 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: - std::vector> data_; + FlatBufferBuilder builder_; + MutableOpResolver resolver_; + TfLiteRegistration fake_op_; + std::vector> operators_; + std::vector> operator_codes_; + std::vector> tensors_; + std::vector> buffers_; }; TEST(VerifyModel, TestEmptyModel) { @@ -62,43 +113,27 @@ TEST(VerifyModel, TestEmptyModel) { /*description=*/0, /*buffers=*/0); ::tflite::FinishModelBuffer(builder, model); - ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize())); + ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize(), + MutableOpResolver{}, DefaultErrorReporter())); } TEST(VerifyModel, TestSimpleModel) { - FlatBufferBuilder builder; - auto inputs = builder.CreateVector({0}); - auto outputs = builder.CreateVector({1}); - auto operator_codes = builder.CreateVector(std::vector>{ - CreateOperatorCodeDirect(builder, BuiltinOperator_CUSTOM, "test")}); - auto operators = - builder.CreateVector(std::vector>{CreateOperator( - builder, /*opcode_index=*/0, - /*inputs=*/builder.CreateVector({0}), - /*outputs=*/builder.CreateVector({1}), BuiltinOptions_NONE, - /*builtin_options=*/0, - /*custom_options=*/0, ::tflite::CustomOptionsFormat_FLEXBUFFERS)}); - std::vector shape; - auto tensors = builder.CreateVector(std::vector>{ - CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/0, - "input", /*quantization=*/0), - CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/0, - "output", /*quantization=*/0)}); - auto subgraph = std::vector>( - {CreateSubGraph(builder, tensors, inputs, outputs, operators, - builder.CreateString("Main"))}); - - auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, operator_codes, - builder.CreateVector(subgraph), - builder.CreateString("SmartReply"), /*buffers=*/0); - - ::tflite::FinishModelBuffer(builder, model); - ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize())); + 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) { - string model = "123"; - ASSERT_FALSE(Verify(model.data(), model.size())); + std::string model = "123"; + ASSERT_FALSE(Verify(model.data(), model.size(), MutableOpResolver{}, + /*error_reporter=*/nullptr)); } TEST(VerifyModel, TestUnsupportedVersion) { @@ -106,7 +141,8 @@ TEST(VerifyModel, TestUnsupportedVersion) { 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())); + ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), + MutableOpResolver{}, DefaultErrorReporter())); } TEST(VerifyModel, TestRandomModificationIsNotAllowed) { @@ -116,20 +152,136 @@ TEST(VerifyModel, TestRandomModificationIsNotAllowed) { /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); ::tflite::FinishModelBuffer(builder, model); - string model_content(reinterpret_cast(builder.GetBufferPointer()), - builder.GetSize()); + 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())) + 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) { +int main(int argc, char** argv) { ::tflite::LogToStderr(); ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); diff --git a/tensorflow/contrib/lite/util.cc b/tensorflow/contrib/lite/util.cc new file mode 100644 index 0000000000000000000000000000000000000000..fb4af07d060cac3a6a4e01c7d625b6db5241f10d --- /dev/null +++ b/tensorflow/contrib/lite/util.cc @@ -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. +==============================================================================*/ +#include "tensorflow/contrib/lite/util.h" + +namespace tflite { + +TfLiteIntArray* ConvertVectorToTfLiteIntArray(const std::vector& input) { + return ConvertArrayToTfLiteIntArray(input.size(), input.data()); +} + +TfLiteIntArray* ConvertArrayToTfLiteIntArray(const int rank, const int* dims) { + TfLiteIntArray* output = TfLiteIntArrayCreate(rank); + for (size_t i = 0; i < rank; i++) { + output->data[i] = dims[i]; + } + return output; +} + +bool EqualArrayAndTfLiteIntArray(const TfLiteIntArray* a, const int b_size, + const int* b) { + if (!a) return false; + if (a->size != b_size) return false; + for (int i = 0; i < a->size; ++i) { + if (a->data[i] != b[i]) return false; + } + return true; +} + +} // namespace tflite diff --git a/tensorflow/contrib/lite/util.h b/tensorflow/contrib/lite/util.h new file mode 100644 index 0000000000000000000000000000000000000000..a34db35823104414cce028b9119397da085d05b1 --- /dev/null +++ b/tensorflow/contrib/lite/util.h @@ -0,0 +1,40 @@ +/* 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 file provides general C++ utility functions in TFLite. +// For example: Converting between `TfLiteIntArray`, `std::vector` and +// Flatbuffer vectors. These functions can't live in `context.h` since it's pure +// C. + +#ifndef TENSORFLOW_CONTRIB_LITE_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_UTIL_H_ + +#include +#include "tensorflow/contrib/lite/context.h" + +namespace tflite { + +// Converts a `std::vector` to a `TfLiteIntArray`. +TfLiteIntArray* ConvertVectorToTfLiteIntArray(const std::vector& input); + +TfLiteIntArray* ConvertArrayToTfLiteIntArray(const int rank, const int* dims); + +// Checks whether a `TfLiteIntArray` and an int array have matching elements. +bool EqualArrayAndTfLiteIntArray(const TfLiteIntArray* a, const int b_size, + const int* b); + +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_UTIL_H_ diff --git a/tensorflow/contrib/lite/util_test.cc b/tensorflow/contrib/lite/util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..04579c53aa4835c47d812c89a1554a0d2f2f30b8 --- /dev/null +++ b/tensorflow/contrib/lite/util_test.cc @@ -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. +==============================================================================*/ + +#include +#include +#include + +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/util.h" + +namespace tflite { +namespace { + +TEST(ConvertVectorToTfLiteIntArray, TestWithVector) { + std::vector input = {1, 2}; + TfLiteIntArray* output = ConvertVectorToTfLiteIntArray(input); + ASSERT_NE(output, nullptr); + EXPECT_EQ(output->size, 2); + EXPECT_EQ(output->data[0], 1); + EXPECT_EQ(output->data[1], 2); + TfLiteIntArrayFree(output); +} + +TEST(ConvertVectorToTfLiteIntArray, TestWithEmptyVector) { + std::vector input; + TfLiteIntArray* output = ConvertVectorToTfLiteIntArray(input); + ASSERT_NE(output, nullptr); + EXPECT_EQ(output->size, 0); + TfLiteIntArrayFree(output); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lookup/BUILD b/tensorflow/contrib/lookup/BUILD index 8ca03f4193f260ce32f942ccaf76a8260b282156..0a6edc33c57435e9466b27cb1d33570d7a98cb12 100644 --- a/tensorflow/contrib/lookup/BUILD +++ b/tensorflow/contrib/lookup/BUILD @@ -46,6 +46,7 @@ tf_py_test( "//tensorflow/python:variables", ], grpc_enabled = True, + tags = ["no_windows"], # TODO: needs investigation on Windows ) filegroup( diff --git a/tensorflow/contrib/lookup/lookup_ops.py b/tensorflow/contrib/lookup/lookup_ops.py index a430dac4ec43ce31f0b5aaae5e7b0b51d25c9632..62f1c810fc72ba7e27c4553006c947f8fa0ef629 100644 --- a/tensorflow/contrib/lookup/lookup_ops.py +++ b/tensorflow/contrib/lookup/lookup_ops.py @@ -341,23 +341,21 @@ class MutableHashTable(LookupInterface): # training to work correctly. Use the node name if no shared_name has been # explicitly specified. use_node_name_sharing = checkpoint and shared_name is None - # pylint: disable=protected-access if self._default_value.get_shape().ndims == 0: - self._table_ref = gen_lookup_ops._mutable_hash_table_v2( + self._table_ref = gen_lookup_ops.mutable_hash_table_v2( shared_name=shared_name, use_node_name_sharing=use_node_name_sharing, key_dtype=key_dtype, value_dtype=value_dtype, name=name) else: - self._table_ref = gen_lookup_ops._mutable_hash_table_of_tensors_v2( + self._table_ref = gen_lookup_ops.mutable_hash_table_of_tensors_v2( shared_name=shared_name, use_node_name_sharing=use_node_name_sharing, key_dtype=key_dtype, value_dtype=value_dtype, value_shape=self._default_value.get_shape(), name=name) - # pylint: enable=protected-access super(MutableHashTable, self).__init__(key_dtype, value_dtype, self._table_ref.op.name.split( "/")[-1]) @@ -378,9 +376,7 @@ class MutableHashTable(LookupInterface): with ops.name_scope(name, "%s_Size" % self._name, [self._table_ref]) as name: with ops.colocate_with(self._table_ref): - - # pylint: disable=protected-access - return gen_lookup_ops._lookup_table_size_v2(self._table_ref, name=name) + return gen_lookup_ops.lookup_table_size_v2(self._table_ref, name=name) def lookup(self, keys, name=None): """Looks up `keys` in a table, outputs the corresponding values. @@ -406,8 +402,7 @@ class MutableHashTable(LookupInterface): with ops.name_scope(name, "%s_lookup_table_find" % self._name, (self._table_ref, keys, self._default_value)) as name: with ops.colocate_with(self._table_ref): - # pylint: disable=protected-access - values = gen_lookup_ops._lookup_table_find_v2( + values = gen_lookup_ops.lookup_table_find_v2( self._table_ref, keys, self._default_value, name=name) values.set_shape(keys.get_shape().concatenate(self._value_shape)) @@ -437,7 +432,7 @@ class MutableHashTable(LookupInterface): [self._table_ref, keys, values]) as name: with ops.colocate_with(self._table_ref): # pylint: disable=protected-access - op = gen_lookup_ops._lookup_table_insert_v2( + op = gen_lookup_ops.lookup_table_insert_v2( self._table_ref, keys, values, name=name) return op @@ -454,8 +449,7 @@ class MutableHashTable(LookupInterface): with ops.name_scope(name, "%s_lookup_table_export_values" % self._name, [self._table_ref]) as name: with ops.colocate_with(self._table_ref): - # pylint: disable=protected-access - exported_keys, exported_values = gen_lookup_ops._lookup_table_export_v2( + exported_keys, exported_values = gen_lookup_ops.lookup_table_export_v2( self._table_ref, self._key_dtype, self._value_dtype, name=name) exported_values.set_shape(exported_keys.get_shape().concatenate( @@ -477,7 +471,7 @@ class MutableHashTable(LookupInterface): def restore(self, restored_tensors, unused_restored_shapes): # pylint: disable=protected-access with ops.colocate_with(self.op._table_ref): - return gen_lookup_ops._lookup_table_import_v2( + return gen_lookup_ops.lookup_table_import_v2( self.op._table_ref, restored_tensors[0], restored_tensors[1]) @@ -551,8 +545,7 @@ class MutableDenseHashTable(LookupInterface): # explicitly specified. use_node_name_sharing = checkpoint and shared_name is None empty_key = ops.convert_to_tensor(empty_key, dtype=key_dtype) - # pylint: disable=protected-access - self._table_ref = gen_lookup_ops._mutable_dense_hash_table_v2( + self._table_ref = gen_lookup_ops.mutable_dense_hash_table_v2( empty_key=empty_key, shared_name=shared_name, use_node_name_sharing=use_node_name_sharing, @@ -560,7 +553,6 @@ class MutableDenseHashTable(LookupInterface): value_shape=self._value_shape, initial_num_buckets=initial_num_buckets, name=name) - # pylint: enable=protected-access super(MutableDenseHashTable, self).__init__( key_dtype, value_dtype, self._table_ref.op.name.split("/")[-1]) @@ -580,8 +572,7 @@ class MutableDenseHashTable(LookupInterface): with ops.name_scope(name, "%s_Size" % self._name, [self._table_ref]) as name: with ops.colocate_with(self._table_ref): - # pylint: disable=protected-access - return gen_lookup_ops._lookup_table_size_v2(self._table_ref, name=name) + return gen_lookup_ops.lookup_table_size_v2(self._table_ref, name=name) def lookup(self, keys, name=None): """Looks up `keys` in a table, outputs the corresponding values. @@ -607,8 +598,7 @@ class MutableDenseHashTable(LookupInterface): with ops.name_scope(name, "%s_lookup_table_find" % self._name, [self._table_ref, keys]) as name: with ops.colocate_with(self._table_ref): - # pylint: disable=protected-access - values = gen_lookup_ops._lookup_table_find_v2( + values = gen_lookup_ops.lookup_table_find_v2( self._table_ref, keys, self._default_value, name=name) if keys.get_shape().ndims is not None and keys.get_shape().ndims > 0: @@ -640,8 +630,7 @@ class MutableDenseHashTable(LookupInterface): with ops.name_scope(name, "%s_lookup_table_insert" % self._name, [self._table_ref, keys, values]) as name: with ops.colocate_with(self._table_ref): - # pylint: disable=protected-access - op = gen_lookup_ops._lookup_table_insert_v2( + op = gen_lookup_ops.lookup_table_insert_v2( self._table_ref, keys, values, name=name) return op @@ -658,8 +647,7 @@ class MutableDenseHashTable(LookupInterface): with ops.name_scope(name, "%s_lookup_table_export_values" % self._name, [self._table_ref]) as name: with ops.colocate_with(self._table_ref): - # pylint: disable=protected-access - exported_keys, exported_values = gen_lookup_ops._lookup_table_export_v2( + exported_keys, exported_values = gen_lookup_ops.lookup_table_export_v2( self._table_ref, self._key_dtype, self._value_dtype, name=name) exported_values.set_shape(exported_keys.get_shape().concatenate( @@ -681,5 +669,5 @@ class MutableDenseHashTable(LookupInterface): def restore(self, restored_tensors, unused_restored_shapes): # pylint: disable=protected-access with ops.colocate_with(self.op._table_ref): - return gen_lookup_ops._lookup_table_import_v2( + return gen_lookup_ops.lookup_table_import_v2( self.op._table_ref, restored_tensors[0], restored_tensors[1]) 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..2b9eee4ef7b418e2b90d388d2f165537b8660a9a 100644 --- a/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py +++ b/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py @@ -50,16 +50,12 @@ def pairwise_distance(feature, squared=False): pairwise_distances: 2-D Tensor of size [number of data, number of data]. """ pairwise_distances_squared = math_ops.add( + math_ops.reduce_sum(math_ops.square(feature), axis=[1], keepdims=True), math_ops.reduce_sum( - math_ops.square(feature), - axis=[1], - keep_dims=True), - math_ops.reduce_sum( - math_ops.square( - array_ops.transpose(feature)), + math_ops.square(array_ops.transpose(feature)), axis=[0], - keep_dims=True)) - 2.0 * math_ops.matmul( - feature, array_ops.transpose(feature)) + keepdims=True)) - 2.0 * math_ops.matmul(feature, + array_ops.transpose(feature)) # Deal with numerical inaccuracies. Set small negatives to zero. pairwise_distances_squared = math_ops.maximum(pairwise_distances_squared, 0.0) @@ -132,10 +128,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 + math_ops.multiply(data - axis_minimums, mask), dim, + keepdims=True) + axis_minimums return masked_maximums @@ -151,10 +147,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 + math_ops.multiply(data - axis_maximums, mask), dim, + keepdims=True) + axis_maximums return masked_minimums @@ -202,8 +198,7 @@ def triplet_semihard_loss(labels, embeddings, margin=1.0): mask_final = array_ops.reshape( math_ops.greater( math_ops.reduce_sum( - math_ops.cast( - mask, dtype=dtypes.float32), 1, keep_dims=True), + math_ops.cast(mask, dtype=dtypes.float32), 1, keepdims=True), 0.0), [batch_size, batch_size]) mask_final = array_ops.transpose(mask_final) @@ -290,7 +285,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 +390,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 +443,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 + math_ops.multiply(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) @@ -467,10 +462,11 @@ def lifted_struct_loss(labels, embeddings, margin=1.0): array_ops.transpose(max_elements), [-1, 1]) loss_exp_left = array_ops.reshape( - 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]) + math_ops.reduce_sum( + math_ops.multiply( + math_ops.exp(diff_tiled - max_elements_vect), 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 +682,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 c573cf15da6aa756bf6840206af1663769d0181d..05e8d9064bea748c935859f5f9b4c7e646f504cf 100644 --- a/tensorflow/contrib/makefile/Makefile +++ b/tensorflow/contrib/makefile/Makefile @@ -407,7 +407,7 @@ $(MARCH_OPTION) \ -I$(JETPACK)/cuda/extras/CUPTI/include - LIBS += \ + CUDA_LIBS := \ -ltfcuda \ -lcudart_static \ -lcudnn \ @@ -420,10 +420,10 @@ $(MARCH_OPTION) \ -lculibos \ -lcurand_static - OBJDIR := $(OBJDIR)Tegra/ - LIBDIR := $(LIBDIR)Tegra/ - BINDIR := $(BINDIR)Tegra/ - DEPDIR := $(DEPDIR)Tegra/ + 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 \ @@ -677,6 +677,7 @@ endif # TEGRA 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 TF_CC_SRCS += tensorflow/contrib/makefile/downloads/fft2d/fftsg.c +TF_CC_SRCS += tensorflow/core/common_runtime/gpu/gpu_id_manager.cc # Also include the op and kernel definitions. TF_CC_SRCS += $(shell cat $(MAKEFILE_DIR)/tf_op_files.txt) PBT_CC_SRCS := $(shell cat $(MAKEFILE_DIR)/tf_pb_text_files.txt) @@ -729,7 +730,7 @@ $(BENCHMARK_NAME): $(BENCHMARK_OBJS) $(LIB_PATH) $(CUDA_LIB_DEPS) @mkdir -p $(dir $@) $(CXX) $(CXXFLAGS) $(INCLUDES) \ -o $(BENCHMARK_NAME) $(BENCHMARK_OBJS) \ - $(LIBFLAGS) $(TEGRA_LIBS) $(LIB_PATH) $(LDFLAGS) $(LIBS) + $(LIBFLAGS) $(TEGRA_LIBS) $(LIB_PATH) $(LDFLAGS) $(LIBS) $(CUDA_LIBS) # NVCC compilation rules for Tegra ifeq ($(BUILD_FOR_TEGRA),1) diff --git a/tensorflow/contrib/makefile/README.md b/tensorflow/contrib/makefile/README.md index 0613de2cabe2065f1e4a816f2295d41b69159c10..6c3b02e12b3082be8bfcc316c4c6122931eb5f76 100644 --- a/tensorflow/contrib/makefile/README.md +++ b/tensorflow/contrib/makefile/README.md @@ -130,6 +130,107 @@ 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 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"]), +``` + +If you are building for Android TV (Shield TV devices), replace "portrait" with "landscape" for android:screenOrientation in all four activities in tensorflow/examples/android/AndroidManifest.xml + +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 @@ -268,7 +369,7 @@ 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 aggresive optimization of the operators and the resulting library may not work with other graphs but will reduce the size of the final library. +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 diff --git a/tensorflow/contrib/makefile/build_all_android.sh b/tensorflow/contrib/makefile/build_all_android.sh index 281c4653c627661ae39592e2ea982d04104c30dd..fc88f59e0948e1d3ed7cce9b809bf30ba280af12 100755 --- a/tensorflow/contrib/makefile/build_all_android.sh +++ b/tensorflow/contrib/makefile/build_all_android.sh @@ -37,7 +37,7 @@ fi ARCH=armeabi-v7a -while getopts "Es:t:Tx:a" opt_name; do +while getopts "Es:t:Tx:a:" opt_name; do case "$opt_name" in E) ENABLE_EXPERIMENTAL_HEXNN_OPS="true";; s) SUB_MAKEFILES="${OPTARG}";; @@ -52,7 +52,7 @@ shift $((OPTIND - 1)) if [ "$ARCH" == "tegra" ]; then if [[ -z "${JETPACK}" ]]; then - export JETPACK="$HOME/JetPack_Android_3.0" + export JETPACK="$HOME/JetPack_Android_3.2" fi if [ ! -d ${JETPACK} ]; then echo "Can't find Jetpack at ${JETPACK}" diff --git a/tensorflow/contrib/makefile/build_all_ios.sh b/tensorflow/contrib/makefile/build_all_ios.sh index a18df256f976c3c0ac4cefe1c884d951e63ef823..0a458a27b3ac9b1a24b0f42de2f0166d515e8cd9 100755 --- a/tensorflow/contrib/makefile/build_all_ios.sh +++ b/tensorflow/contrib/makefile/build_all_ios.sh @@ -80,10 +80,9 @@ if [[ ! -z "${OPTIMIZE_FOR_GRAPH}" ]]; then 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 + ${PRNT_SLCTV_BIN} --graphs=${OPTIMIZE_FOR_GRAPH} > ${TOP_SRCDIR}/tensorflow/core/framework/ops_to_register.h fi if [[ "${ONLY_MAKE_TENSORFLOW}" != "true" ]]; then @@ -96,7 +95,7 @@ if [[ "${ONLY_MAKE_TENSORFLOW}" != "true" ]]; then if [[ -z "${BUILD_ARCH}" ]]; then # Compile protobuf for the target iOS device architectures. - tensorflow/contrib/makefile/compile_ios_protobuf.sh -a ${DEFAULT_ARCH} + 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} @@ -111,7 +110,7 @@ if [[ -z "${BUILD_ARCH}" ]]; then TARGET_NSYNC_LIB=`tensorflow/contrib/makefile/compile_nsync.sh -t ios` else # arch specified so build just that - TARGET_NSYNC_LIB=`tensorflow/contrib/makefile/compile_nsync.sh -t ios -a ${BUILD_ARCH}` + TARGET_NSYNC_LIB=`tensorflow/contrib/makefile/compile_nsync.sh -t ios -a "${BUILD_ARCH}"` fi export HOST_NSYNC_LIB TARGET_NSYNC_LIB diff --git a/tensorflow/contrib/makefile/proto_text_cc_files.txt b/tensorflow/contrib/makefile/proto_text_cc_files.txt index d56e388477db6239cfb577f7e2754321ff33bd82..77c936d8c5b99033ff5c5e149a6ce6613b603132 100644 --- a/tensorflow/contrib/makefile/proto_text_cc_files.txt +++ b/tensorflow/contrib/makefile/proto_text_cc_files.txt @@ -17,6 +17,7 @@ tensorflow/core/platform/env_time.cc tensorflow/core/platform/setround.cc tensorflow/core/platform/denormal.cc tensorflow/core/platform/default/tracing.cc +tensorflow/core/platform/default/mutex.cc tensorflow/core/platform/default/logging.cc tensorflow/core/platform/cpu_info.cc tensorflow/core/lib/wav/wav_io.cc 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 203ff4f890a3b0ed32caa1406508b100dd47bcad..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 @@ -84,7 +84,7 @@ 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" diff --git a/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in b/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in index d9277ed60cb456208572ca1ad8df530648faef82..3081084ee76e41de801f49a67c1fec07f4ff03b9 100644 --- a/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in +++ b/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in @@ -54,7 +54,7 @@ $(INFERENCE_SO_PATH): $(LIB_OBJS) $(INFERENCE_OBJS) $(CUDA_LIB_DEPS) -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 9a1ab503178f8c5fb4ba77917bd2cdad0ac72cdf..5a812af4e95fe7a05b9c2634b0cc1d860fb7f619 100644 --- a/tensorflow/contrib/makefile/tf_op_files.txt +++ b/tensorflow/contrib/makefile/tf_op_files.txt @@ -293,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/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 d3dce46bfb6e9c77cc7ae107b323a9bc7074c47e..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 @@ -83,6 +84,7 @@ 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 diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops.py b/tensorflow/contrib/metrics/python/ops/metric_ops.py index 55946c128b1a46b8368aedd9f857c1902c4c4586..81f05e7ce587ed1da67a17efbbeb809dbe7fc0b3 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops.py @@ -38,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. @@ -739,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) @@ -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 @@ -3356,8 +3646,8 @@ def cohen_kappa(labels, `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' + if context.executing_eagerly(): + 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.' @@ -3430,6 +3720,7 @@ def cohen_kappa(labels, __all__ = [ + 'auc_with_confidence_intervals', 'aggregate_metric_map', 'aggregate_metrics', 'cohen_kappa', diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py index e067f08babd9a900e876545d427c91e5ff808f04..33eb655fb660f0ecdfe1c5ab870d7f17690ae3ff 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py @@ -2128,6 +2128,205 @@ class StreamingDynamicAUCTest(test.TestCase): self.assertAlmostEqual(0.90277, auc.eval(), delta=1e-5) +class AucWithConfidenceIntervalsTest(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 AucWithConfidenceIntervalData 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. + + 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): diff --git a/tensorflow/contrib/model_pruning/README.md b/tensorflow/contrib/model_pruning/README.md index d286750c257e9a78a82c95c1fc872b3ca6972203..52b659c69fdfc507e6259e928d79c65471f2f025 100644 --- a/tensorflow/contrib/model_pruning/README.md +++ b/tensorflow/contrib/model_pruning/README.md @@ -134,7 +134,7 @@ $ bazel-bin/$examples_dir/cifar10/cifar10_eval --run_once ### Block Sparsity -For some hardware architectures, it may be beneficial to induce spatially correlated sparsity. To train models in which the weight tensors have block sparse structure, set *block_height* and *block_width* hyperparameters to the desired block configuration (2x2, 4x4, 4x1, 1x8, etc). Currently, block sparsity is supported for weight tensors with rank 2 only. The matrix is partitioned into non-overlapping blocks of size *[block_height, block_dim]* and the either the average or max absolute value in this block is taken as a proxy for the entire block (set by *block_pooling_function* hyperparameter). +For some hardware architectures, it may be beneficial to induce spatially correlated sparsity. To train models in which the weight tensors have block sparse structure, set *block_height* and *block_width* hyperparameters to the desired block configuration (2x2, 4x4, 4x1, 1x8, etc). Currently, block sparsity is only supported for weight tensors which can be squeezed to rank 2. The matrix is partitioned into non-overlapping blocks of size *[block_height, block_dim]* and the either the average or max absolute value in this block is taken as a proxy for the entire block (set by *block_pooling_function* hyperparameter). The convolution layer tensors are always pruned used block dimensions of [1,1]. ## References diff --git a/tensorflow/contrib/model_pruning/python/layers/layers.py b/tensorflow/contrib/model_pruning/python/layers/layers.py index dfebb9a6794056dd43b0699ccbcc5797f2f172f7..466daf204a1ae86a7f37107342046305ea7249fc 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 @@ -215,7 +214,7 @@ def masked_convolution(inputs, elif data_format == 'NCHW': df = 'channels_first' else: - raise ValueError('Unsupported data fromat', data_format) + raise ValueError('Unsupported data format', data_format) layer = layer_class( filters=num_outputs, diff --git a/tensorflow/contrib/model_pruning/python/pruning.py b/tensorflow/contrib/model_pruning/python/pruning.py index d16af9da19816211ee22f6ea48a347f0b9a4e612..5146a4a2de7806041991c04958de378b2d3dc810 100644 --- a/tensorflow/contrib/model_pruning/python/pruning.py +++ b/tensorflow/contrib/model_pruning/python/pruning.py @@ -216,7 +216,7 @@ def _partitioned_variable_assign(partitioned_var, new_value): """Assign op for partitioned variables. Args: - partitioned_var: A partitioned tensotflow variable + partitioned_var: A partitioned tensorflow variable new_value: Value to be assigned to the variable var Returns: @@ -523,7 +523,8 @@ class Pruning(object): """Performs block-granular masking of the weights. Block pruning occurs only if the block_height or block_width is > 1 and - if the weight tensor has ndims = 2. Otherwise, elementwise pruning occurs. + if the weight tensor, when squeezed, has ndims = 2. Otherwise, elementwise + pruning occurs. Args: weights: The weight tensor that needs to be masked. threshold: The current threshold value. The function will compute a new @@ -540,7 +541,8 @@ class Pruning(object): Raises: ValueError: if block pooling function is not AVG or MAX """ - if weights.get_shape().ndims != 2 or self._block_dim == [1, 1]: + squeezed_weights = array_ops.squeeze(weights) + if squeezed_weights.get_shape().ndims != 2 or self._block_dim == [1, 1]: return self._update_mask(weights, threshold) if self._block_pooling_function not in ['AVG', 'MAX']: @@ -549,9 +551,11 @@ class Pruning(object): with ops.name_scope(weights.op.name + '_pruning_ops'): abs_weights = math_ops.abs( - array_ops.reshape( - weights, [1, weights.get_shape()[0], - weights.get_shape()[1], 1])) + array_ops.reshape(weights, [ + 1, + squeezed_weights.get_shape()[0], + squeezed_weights.get_shape()[1], 1 + ])) pool_window = [self._block_dim[0], self._block_dim[1]] pooled_weights = nn_ops.pool( abs_weights, @@ -572,9 +576,10 @@ class Pruning(object): array_ops.ones(self._block_dim)) sliced_mask = array_ops.slice( updated_mask, [0, 0], - [weights.get_shape()[0], - weights.get_shape()[1]]) - return smoothed_threshold, sliced_mask + [squeezed_weights.get_shape()[0], + squeezed_weights.get_shape()[1]]) + return smoothed_threshold, array_ops.reshape(sliced_mask, + array_ops.shape(weights)) def _get_mask_assign_ops(self): # Make sure the assignment ops have not already been added to the list diff --git a/tensorflow/contrib/model_pruning/python/pruning_test.py b/tensorflow/contrib/model_pruning/python/pruning_test.py index 1767b4bb94a9bb56bc6a4933423ad27d8cf3ed35..89e65713197afc6ed37346cb67a6e9be3fa9290f 100644 --- a/tensorflow/contrib/model_pruning/python/pruning_test.py +++ b/tensorflow/contrib/model_pruning/python/pruning_test.py @@ -140,6 +140,23 @@ class PruningTest(test.TestCase): [0.0, -0.3, 0.0, -0.4]]) expected_mask = [[0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1], [1, 1, 1, 1]] + self._blockMasking(param_list + ["block_pooling_function=MAX"], weights_max, + expected_mask) + self._blockMasking(param_list + ["block_pooling_function=AVG"], weights_avg, + expected_mask) + + def testBlockMaskingWithHigherDimensions(self): + param_list = ["block_height=2", "block_width=2", "threshold_decay=0"] + + # Weights as in testBlockMasking, but with one extra dimension. + weights_avg = constant_op.constant( + [[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4], + [0.3, 0.3, 0.4, 0.4]]]) + weights_max = constant_op.constant( + [[[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2], [0.3, 0.0, 0.4, 0.0], + [0.0, -0.3, 0.0, -0.4]]]) + expected_mask = [[[0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1], [1, 1, 1, 1]]] + self._blockMasking(param_list + ["block_pooling_function=MAX"], weights_max, expected_mask) self._blockMasking(param_list + ["block_pooling_function=AVG"], diff --git a/tensorflow/contrib/mpi/mpi_utils.h b/tensorflow/contrib/mpi/mpi_utils.h index fa297c28cb47d43ba927ab941854bd472d90b465..df055ff56731140b3bd09704c70e65f81362f763 100644 --- a/tensorflow/contrib/mpi/mpi_utils.h +++ b/tensorflow/contrib/mpi/mpi_utils.h @@ -24,6 +24,8 @@ limitations under the License. #include "tensorflow/core/lib/strings/str_util.h" +// Skip MPI C++ bindings support, this matches the usage in other places +#define OMPI_SKIP_MPICXX #include "third_party/mpi/mpi.h" #define MPI_CHECK(cmd) \ do { \ diff --git a/tensorflow/contrib/nccl/BUILD b/tensorflow/contrib/nccl/BUILD index 5ac96007df7ee08b1e32aacd28f83768859810a9..94d01efee1546feca89a7e88acedf915b1dfb3a4 100644 --- a/tensorflow/contrib/nccl/BUILD +++ b/tensorflow/contrib/nccl/BUILD @@ -52,6 +52,7 @@ tf_cuda_cc_test( "manual", "multi_gpu", "no_oss", + "noguitar", "notap", ], deps = @@ -136,6 +137,7 @@ cuda_py_test( "manual", "multi_gpu", "no_oss", + "noguitar", "notap", ], ) diff --git a/tensorflow/contrib/nccl/python/ops/nccl_ops.py b/tensorflow/contrib/nccl/python/ops/nccl_ops.py index 8dc038b9ac992de7db8b762e3697c6693099e192..794372a1f4b0dcc41bcf0da611f5bc2ec9301973 100644 --- a/tensorflow/contrib/nccl/python/ops/nccl_ops.py +++ b/tensorflow/contrib/nccl/python/ops/nccl_ops.py @@ -267,5 +267,5 @@ def _check_device(tensor, expected=None): def _check_graph_mode(): - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError('Nccl ops are not supported in eager mode') 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 b24e332e4aea7f0ef981909558dcd6d730ca08a7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/lstm1d.py +++ /dev/null @@ -1,184 +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 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]): - lstm_cell = rnn_cell.BasicLSTMCell(noutput) - if reverse: - inputs = array_ops.reverse_v2(inputs, [0]) - outputs, _ = rnn.dynamic_rnn( - lstm_cell, inputs, time_major=True, dtype=inputs.dtype) - 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 32925569a82c43be75a0b6e93d7d781cda3d53f4..a2dbb8052bfa1634d27c8b38a9bb6ca27fae42a2 100644 --- a/tensorflow/contrib/nearest_neighbor/kernels/heap.h +++ b/tensorflow/contrib/nearest_neighbor/kernels/heap.h @@ -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_)) { 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/scaled_softplus.py b/tensorflow/contrib/nn/python/ops/scaled_softplus.py index fcbfbc239ca5b8a1d4b17b403f99b7eb05db47b0..7184ef2b66ec4662af3a37def070ab151d6e7c15 100644 --- a/tensorflow/contrib/nn/python/ops/scaled_softplus.py +++ b/tensorflow/contrib/nn/python/ops/scaled_softplus.py @@ -30,9 +30,7 @@ def _reduce_and_reshape_grad(g, t): """Returns the gradient, sum-reduced and reshaped to `t`'s shape.""" shape = array_ops.shape(t) g_shape = array_ops.shape(g) - # pylint: disable=protected-access - bcast_dims, _ = gen_array_ops._broadcast_gradient_args(shape, g_shape) - # pylint: enable=protected-access + bcast_dims, _ = gen_array_ops.broadcast_gradient_args(shape, g_shape) return array_ops.reshape(math_ops.reduce_sum(g, bcast_dims), shape) diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD index 827279bd476f9666a972f43ad557fde6d0b6c59a..bacf15bbd6140caf647552f0dca02209634ae56b 100644 --- a/tensorflow/contrib/opt/BUILD +++ b/tensorflow/contrib/opt/BUILD @@ -52,6 +52,9 @@ py_test( name = "external_optimizer_test", srcs = ["python/training/external_optimizer_test.py"], srcs_version = "PY2AND3", + tags = [ + "no-internal-py3", + ], deps = [ ":opt_py", "//tensorflow/python:array_ops", diff --git a/tensorflow/contrib/opt/python/training/addsign_test.py b/tensorflow/contrib/opt/python/training/addsign_test.py index bd19ee3e7ac514448c6d79272abb86a154f55e9a..08d45ed73f3ae4b580d7078272e79fef22ef67c5 100644 --- a/tensorflow/contrib/opt/python/training/addsign_test.py +++ b/tensorflow/contrib/opt/python/training/addsign_test.py @@ -97,7 +97,7 @@ class AddSignTest(test.TestCase): global_step=global_step) neg_update = opt.apply_gradients(zip([-grads0, -grads1], [var0, var1]), global_step=global_step) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(variables.global_variables_initializer()) # Fetch params to validate initial values self.assertAllClose([1.0, 2.0], self.evaluate(var0)) @@ -108,13 +108,13 @@ class AddSignTest(test.TestCase): # last 3 steps with negative gradient (sign(gm) should be -1) for t in range(1, 8): if t < 5: - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(update) elif t > 1: opt.apply_gradients(zip([grads0, grads1], [var0, var1]), global_step=global_step) else: - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(neg_update) elif t > 1: opt.apply_gradients(zip([-grads0, -grads1], [var0, var1]), diff --git a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py index 716ee9cdf704a14a6e433c7f92ccb91739f70655..5763593b81497f5d6945ff1e5d000042d295c093 100644 --- a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py +++ b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py @@ -150,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: 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_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/opt/python/training/powersign_test.py b/tensorflow/contrib/opt/python/training/powersign_test.py index ff7b1a72d47d8ef54980905323bcaf358c988a82..5214082dd66f00eadadad71d50f7e00b178b8c10 100644 --- a/tensorflow/contrib/opt/python/training/powersign_test.py +++ b/tensorflow/contrib/opt/python/training/powersign_test.py @@ -99,7 +99,7 @@ class PowerSignTest(test.TestCase): neg_update = opt.apply_gradients(zip([-grads0, -grads1], [var0, var1]), global_step=global_step) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(variables.global_variables_initializer()) # Fetch params to validate initial values self.assertAllClose([1.0, 2.0], self.evaluate(var0)) @@ -110,13 +110,13 @@ class PowerSignTest(test.TestCase): # last 3 steps with negative gradient (sign(gm) should be -1) for t in range(1, 8): if t < 5: - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(update) elif t > 1: opt.apply_gradients(zip([grads0, grads1], [var0, var1]), global_step=global_step) else: - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(neg_update) elif t > 1: opt.apply_gradients(zip([-grads0, -grads1], [var0, var1]), diff --git a/tensorflow/contrib/opt/python/training/variable_clipping_optimizer.py b/tensorflow/contrib/opt/python/training/variable_clipping_optimizer.py index 74036082f0ca2bae23b30deb1b1986befd6601d8..3c0b8394be51e8744b5461a00a99ead5e45d90b2 100644 --- a/tensorflow/contrib/opt/python/training/variable_clipping_optimizer.py +++ b/tensorflow/contrib/opt/python/training/variable_clipping_optimizer.py @@ -109,7 +109,7 @@ class VariableClippingOptimizer(optimizer.Optimizer): def _clip_dense(self, var): with self._maybe_colocate_with(var): - updated_var_value = var._ref() # pylint: disable=protected-access + updated_var_value = var.read_value() normalized_var = clip_ops.clip_by_norm( updated_var_value, self._max_norm, self._vars_to_clip_dims[var]) delta = updated_var_value - normalized_var 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/predictor/predictor_factories.py b/tensorflow/contrib/predictor/predictor_factories.py index 04b5d5bdf158dc6a478d7a24b538c75d1dca8d45..6e77e934fe19851eea9ed0b74eb7aecc76f6237a 100644 --- a/tensorflow/contrib/predictor/predictor_factories.py +++ b/tensorflow/contrib/predictor/predictor_factories.py @@ -53,7 +53,7 @@ def from_contrib_estimator(estimator, `Estimator`. """ if isinstance(estimator, core_estimator.Estimator): - raise TypeError('Espected estimator to be of type ' + raise TypeError('Expected estimator to be of type ' 'tf.contrib.learn.Estimator, but got type ' 'tf.python.estimator.Estimator. You likely want to call ' 'from_estimator.') @@ -88,7 +88,7 @@ def from_estimator(estimator, `Estimator`. """ if isinstance(estimator, contrib_estimator.Estimator): - raise TypeError('Espected estimator to be of type ' + raise TypeError('Expected estimator to be of type ' 'tf.python.estimator.Estimator, but got type ' 'tf.contrib.learn.Estimator. You likely want to call ' 'from_contrib_estimator.') diff --git a/tensorflow/contrib/py2tf/BUILD b/tensorflow/contrib/py2tf/BUILD index 3e846aefeb30e29de8b00f76c0b8d7c6053e8099..d91220f6ddb859ff52d4e5853948cb667981009b 100644 --- a/tensorflow/contrib/py2tf/BUILD +++ b/tensorflow/contrib/py2tf/BUILD @@ -18,69 +18,14 @@ py_library( name = "py2tf", srcs = [ "__init__.py", - "api.py", - "config.py", - "conversion.py", - "naming.py", ], srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/py2tf/converters", + "//tensorflow/contrib/py2tf/impl", "//tensorflow/contrib/py2tf/pyct", - "//tensorflow/contrib/py2tf/pyct/static_analysis", + "//tensorflow/contrib/py2tf/utils", "@gast_archive//:gast", "@six_archive//:six", ], ) - -# Separate target that allows access to internal symbols for testing. -py_library( - name = "py2tf_internal", - 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", - "@gast_archive//:gast", - "@six_archive//:six", - ], -) - -py_test( - name = "api_test", - srcs = ["api_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":py2tf_internal", - "//tensorflow/python:client_testlib", - ], -) - -py_test( - name = "conversion_test", - srcs = ["conversion_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":py2tf_internal", - "//tensorflow/python:client_testlib", - "@gast_archive//:gast", - ], -) - -py_test( - name = "naming_test", - srcs = ["naming_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":py2tf_internal", - "//tensorflow/python:client_testlib", - ], -) diff --git a/tensorflow/contrib/py2tf/__init__.py b/tensorflow/contrib/py2tf/__init__.py index d187da99e065cb2d31ae4e45a9570378f9d1bf27..6531183cb59af774299eb767cce111d2ec6f32b4 100644 --- a/tensorflow/contrib/py2tf/__init__.py +++ b/tensorflow/contrib/py2tf/__init__.py @@ -21,11 +21,18 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.py2tf.api import to_code -from tensorflow.contrib.py2tf.api import to_graph +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.impl.api import convert +from tensorflow.contrib.py2tf.impl.api import converted_call +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 = ['to_graph', 'to_code'] +_allowed_symbols = [ + 'to_graph', 'to_code', 'convert', 'graph_ready', 'converted_call', 'utils', + 'PyFlowParseError' +] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/py2tf/converters/BUILD b/tensorflow/contrib/py2tf/converters/BUILD index b61fda3e9122ad35ae9f84bcd94b7df2a53e9584..4bb6f76019739fc3b5bf4bf52e302a698693db5a 100644 --- a/tensorflow/contrib/py2tf/converters/BUILD +++ b/tensorflow/contrib/py2tf/converters/BUILD @@ -18,16 +18,20 @@ py_library( name = "converters", srcs = [ "asserts.py", - "break_canonicalization.py", + "break_statements.py", "builtin_functions.py", "call_trees.py", - "continue_canonicalization.py", + "continue_statements.py", "control_flow.py", "decorators.py", - "for_canonicalization.py", + "for_loops.py", + "ifexp.py", + "list_comprehension.py", + "lists.py", "logical_expressions.py", - "print_functions.py", + "name_scopes.py", "side_effect_guards.py", + "single_return.py", ], srcs_version = "PY2AND3", visibility = ["//tensorflow:__subpackages__"], @@ -45,8 +49,11 @@ py_library( visibility = ["//tensorflow:__subpackages__"], deps = [ ":converters", + "//tensorflow/contrib/py2tf/pyct", "//tensorflow/contrib/py2tf/pyct/static_analysis", + "//tensorflow/contrib/py2tf/utils", "@gast_archive//:gast", + "@six_archive//:six", ], ) @@ -56,18 +63,27 @@ py_test( srcs_version = "PY2AND3", deps = [ ":test_lib", - "//tensorflow/contrib/py2tf/pyct", "//tensorflow/python:client_testlib", ], ) py_test( - name = "break_canonicalization_test", - srcs = ["break_canonicalization_test.py"], + 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", + tags = ["no_windows"], # TODO: needs investigation on Windows + deps = [ + ":test_lib", "//tensorflow/python:client_testlib", ], ) @@ -76,9 +92,30 @@ py_test( name = "call_trees_test", srcs = ["call_trees_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/impl", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "continue_statements_test", + srcs = ["continue_statements_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//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", ], ) @@ -89,26 +126,23 @@ py_test( srcs_version = "PY2AND3", deps = [ ":test_lib", - "//tensorflow/contrib/py2tf/pyct", "//tensorflow/python:client_testlib", ], ) py_test( - name = "continue_canonicalization_test", - srcs = ["continue_canonicalization_test.py"], + name = "for_loops_test", + srcs = ["for_loops_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", + name = "name_scopes_test", + srcs = ["name_scopes_test.py"], deps = [ ":test_lib", "//tensorflow/contrib/py2tf/pyct", @@ -117,22 +151,21 @@ py_test( ) py_test( - name = "builtin_functions_test", - srcs = ["builtin_functions_test.py"], + 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 = "for_canonicalization_test", - srcs = ["for_canonicalization_test.py"], + name = "lists_test", + srcs = ["lists_test.py"], + srcs_version = "PY2AND3", deps = [ ":test_lib", - "//tensorflow/contrib/py2tf/pyct", "//tensorflow/python:client_testlib", ], ) @@ -143,26 +176,39 @@ py_test( srcs_version = "PY2AND3", deps = [ ":test_lib", - "//tensorflow/contrib/py2tf/pyct", "//tensorflow/python:client_testlib", ], ) py_test( - name = "print_functions_test", - srcs = ["print_functions_test.py"], + name = "side_effect_guards_test", + srcs = ["side_effect_guards_test.py"], + srcs_version = "PY2AND3", + tags = [ + # TODO(mdan): Fix. + "flaky", + "notap", + ], + deps = [ + ":test_lib", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "single_return_test", + srcs = ["single_return_test.py"], srcs_version = "PY2AND3", deps = [ ":test_lib", "//tensorflow/contrib/py2tf/pyct", "//tensorflow/python:client_testlib", - "@gast_archive//:gast", ], ) py_test( - name = "side_effect_guards_test", - srcs = ["side_effect_guards_test.py"], + name = "ifexp_test", + srcs = ["ifexp_test.py"], srcs_version = "PY2AND3", deps = [ ":test_lib", diff --git a/tensorflow/contrib/py2tf/converters/asserts.py b/tensorflow/contrib/py2tf/converters/asserts.py index 2d6ee1d09829b538815dbb9794868c13f51578fc..5b9b8e772bed82df2429fd6cb94dbf7b565e22b3 100644 --- a/tensorflow/contrib/py2tf/converters/asserts.py +++ b/tensorflow/contrib/py2tf/converters/asserts.py @@ -35,7 +35,7 @@ class AssertsTransformer(transformer.Base): # Note: The lone tf.Assert call will be wrapped with control_dependencies # by side_effect_guards. template = """ - tf.Assert(test, [tf.constant(msg)]) + tf.Assert(test, [msg]) """ if node.msg is None: diff --git a/tensorflow/contrib/py2tf/converters/break_canonicalization.py b/tensorflow/contrib/py2tf/converters/break_statements.py similarity index 81% rename from tensorflow/contrib/py2tf/converters/break_canonicalization.py rename to tensorflow/contrib/py2tf/converters/break_statements.py index 2ae65e3007466409433e9b4ea0081898907e19ac..bfb709c5e32c6f19dc0fd109df61ece925d701a3 100644 --- a/tensorflow/contrib/py2tf/converters/break_canonicalization.py +++ b/tensorflow/contrib/py2tf/converters/break_statements.py @@ -22,13 +22,15 @@ 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(gast.NodeTransformer): +class BreakCanonicalizationTransformer(transformer.Base): """Canonicalizes continue statements into additional conditionals.""" - def __init__(self, namer): - self.namer = namer + def __init__(self, context): + super(BreakCanonicalizationTransformer, self).__init__(context) # This is a stack structure, to correctly process nested loops. self.break_uses = [] @@ -67,9 +69,10 @@ class BreakCanonicalizationTransformer(gast.NodeTransformer): def visit_While(self, node): self.generic_visit(node.test) - scope = anno.getanno(node, 'body_scope') + scope = anno.getanno(node, NodeAnno.BODY_SCOPE) - break_var = self.namer.new_symbol('break_requested', scope.referenced) + 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]: @@ -89,9 +92,10 @@ class BreakCanonicalizationTransformer(gast.NodeTransformer): def visit_For(self, node): self.generic_visit(node.target) self.generic_visit(node.iter) - scope = anno.getanno(node, 'body_scope') + scope = anno.getanno(node, NodeAnno.BODY_SCOPE) - break_var = self.namer.new_symbol('break_requested', scope.referenced) + 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]: @@ -112,7 +116,5 @@ class BreakCanonicalizationTransformer(gast.NodeTransformer): return self._create_break_trigger() -def transform(node, namer): - transformer = BreakCanonicalizationTransformer(namer) - node = transformer.visit(node) - return node +def transform(node, context): + return BreakCanonicalizationTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/break_canonicalization_test.py b/tensorflow/contrib/py2tf/converters/break_statements_test.py similarity index 52% rename from tensorflow/contrib/py2tf/converters/break_canonicalization_test.py rename to tensorflow/contrib/py2tf/converters/break_statements_test.py index b5ba2ad923dfeb73b38169494f6c7ea16ee815f1..095fcdff07d44ecc6b9bb7f8d3e2c7c43df72a02 100644 --- a/tensorflow/contrib/py2tf/converters/break_canonicalization_test.py +++ b/tensorflow/contrib/py2tf/converters/break_statements_test.py @@ -12,25 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for break_canonicalization module.""" +"""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_canonicalization -from tensorflow.contrib.py2tf.converters import control_flow +from tensorflow.contrib.py2tf.converters import break_statements from tensorflow.contrib.py2tf.converters import converter_test_base -from tensorflow.contrib.py2tf.pyct import compiler from tensorflow.python.platform import test -class TestNamer(control_flow.SymbolNamer): - - def new_symbol(self, name_root, _): - return name_root - - class BreakCanonicalizationTest(converter_test_base.TestCase): def test_basic_break(self): @@ -44,15 +36,15 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v - node = self.parse_and_analyze(test_fn, {}, include_type_analysis=False) - node = break_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) - 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)) + 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): @@ -76,16 +68,17 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v - node = self.parse_and_analyze(test_fn, {}, include_type_analysis=False) - node = break_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) - # 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])) + 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): @@ -104,15 +97,15 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v, u, w - node = self.parse_and_analyze(test_fn, {}, include_type_analysis=False) - node = break_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) - 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)) + 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__': diff --git a/tensorflow/contrib/py2tf/converters/builtin_functions.py b/tensorflow/contrib/py2tf/converters/builtin_functions.py index 7f6b64a34c1b95f0dd6b92dbc587da672e6c9c28..f1129ef153e6be6cbcbbf4bab63c4fe32ec77147 100644 --- a/tensorflow/contrib/py2tf/converters/builtin_functions.py +++ b/tensorflow/contrib/py2tf/converters/builtin_functions.py @@ -21,33 +21,57 @@ from __future__ import print_function import gast from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer -class BuiltinFunctionTransformer(gast.NodeTransformer): - """Transforms Print nodes to Call so they can be handled as functions.""" +class BuiltinFunctionTransformer(transformer.Base): + """Handles builtin functions. - # TODO(mdan): Bring print_functions in here. + This transformer only covers functions that are translated into a + TF equivalent, like `len`. + """ - def _convert_len(self, node): + def __init__(self, context): + super(BuiltinFunctionTransformer, self).__init__(context) + + # pylint:disable=invalid-name + + def _convert_builtin(self, node): template = """ - tf.shape(args)[0] + py2tf_utils.dynamic_builtin(func, args) """ - new_call = templates.replace(template, args=node.args)[0].value - return new_call + return templates.replace(template, func=node.func, args=node.args)[0].value - # pylint:disable=invalid-name + def _convert_print(self, node): + template = """ + py2tf_utils.dynamic_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 in ('len', 'range'): + return self._convert_builtin(node) + # Print needs to be handled separately because it can be read as statement. + 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): - transformer = BuiltinFunctionTransformer() - node = transformer.visit(node) - return node +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 index b5358da6bc0be06ec1f59d0ef58d926289b5b78f..eb60a1d8ae2b56907df8f3ffafe7604883cfc2a9 100644 --- a/tensorflow/contrib/py2tf/converters/builtin_functions_test.py +++ b/tensorflow/contrib/py2tf/converters/builtin_functions_test.py @@ -18,11 +18,16 @@ 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.contrib.py2tf.pyct import compiler 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 @@ -34,14 +39,78 @@ class BuiltinFunctionsTest(converter_test_base.TestCase): return len(a) node = self.parse_and_analyze(test_fn, {'len': len}) - node = builtin_functions.transform(node) - result = compiler.ast_to_object(node) - setattr(result, 'tf', array_ops) - - with self.test_session() as sess: - self.assertEqual(3, - sess.run( - result.test_fn(constant_op.constant([0, 0, 0])))) + 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])))) + + self.assertEqual(3, result.test_fn([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__': diff --git a/tensorflow/contrib/py2tf/converters/call_trees.py b/tensorflow/contrib/py2tf/converters/call_trees.py index 4c238b7fb9e7bdb25b0ded9a4c4a6b13508f19c1..ca8726f9160d106ebd82e01e399e65fb77b02aab 100644 --- a/tensorflow/contrib/py2tf/converters/call_trees.py +++ b/tensorflow/contrib/py2tf/converters/call_trees.py @@ -27,6 +27,7 @@ import types import gast from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import inspect_utils from tensorflow.contrib.py2tf.pyct import parser from tensorflow.contrib.py2tf.pyct import templates from tensorflow.contrib.py2tf.pyct import transformer @@ -72,9 +73,8 @@ class CallTreeTransformer(transformer.Base): self.uncompiled_modules = uncompiled_modules self.nocompile_decorators = nocompile_decorators - # pylint:disable=invalid-name - def _resolve_name(self, node): + """Used to resolve decorator info.""" if isinstance(node, gast.Call): return self._resolve_name(node.func) if isinstance(node, gast.Name): @@ -99,7 +99,13 @@ class CallTreeTransformer(transformer.Base): (owner_type, node.attr)) return None + def _function_is_compilable(self, target_entity): + """Determines whether an entity can be compiled at all.""" + # TODO(mdan): This is just a placeholder. Implement. + return not isinstance(target_entity, types.BuiltinFunctionType) + def _should_compile(self, node, fqn): + """Determines whether an entity should be compiled in the context.""" for i in range(1, len(fqn)): if fqn[:i] in self.uncompiled_modules: return False @@ -141,33 +147,6 @@ class CallTreeTransformer(transformer.Base): 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') @@ -182,7 +161,11 @@ class CallTreeTransformer(transformer.Base): target_fqn, live_entity=target_entity) do_rename = True else: - owner_type = self._determine_function_owner(target_entity) + if anno.hasanno(node.func, 'parent_type'): + owner_type = anno.getanno(node.func, 'parent_type') + else: + # Fallback - not reliable. + owner_type = inspect_utils.getmethodclass(target_entity) new_name, do_rename = self.context.namer.compiled_function_name( target_fqn, live_entity=target_entity, owner_type=owner_type) @@ -192,32 +175,45 @@ class CallTreeTransformer(transformer.Base): # 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 = gast.Name(new_name, gast.Load(), None) + node.func = templates.replace('func_name', func_name=new_name)[0] return node def _wrap_to_py_func_no_return(self, node): - args_scope = anno.getanno(node, 'args_scope') # TODO(mdan): Properly handle varargs, kwargs, etc. template = """ - def wrapper(args): - call(args) - return 1 - tf.py_func(wrapper, [args], [tf.int64]) + py2tf_utils.wrap_py_func(func, None, (original_args,), True) """ - wrapper_def, call_expr = templates.replace( - template, - call=node.func, - wrapper=self.context.namer.compiled_function_name(node.func.id)[0], - args=tuple(gast.Name(n, gast.Load(), None) for n in args_scope.used)) - anno.setanno(call_expr.value, 'args_scope', args_scope) - # TODO(mdan): Rename this annotation to 'graph_ready' - anno.setanno(wrapper_def, 'skip_processing', True) - - return (wrapper_def, call_expr) + return templates.replace(template, func=node.func, original_args=node.args) + + def _insert_dynamic_conversion(self, node): + """Inlines a dynamic conversion for a dynamic function.""" + # TODO(mdan): Pass information on the statically compiled functions. + # Having access to the statically compiled functions can help avoid + # unnecessary compilation. + # For example, this would lead to function `a` being compiled twice: + # + # def a(): + # v = b + # b() + # def b(): + # a() + # + # This is really a problem with recursive calls, which currently can + # only be gated by a static condition, and should be rare. + # TODO(mdan): It probably makes sense to use dynamic conversion every time. + # Before we could convert all the time though, we'd need a reasonable + # caching mechanism. + template = """ + py2tf_api.converted_call(func, True, False, {}, original_args) + """ + call_expr = templates.replace( + template, func=node.func, original_args=node.args) + new_call = call_expr[0].value + # TODO(mdan): Improve the template mechanism to better support this. + new_call.keywords = node.keywords + return new_call - def _function_is_compilable(self, target_entity): - # TODO(mdan): This is just a placeholder. Implement. - return not isinstance(target_entity, types.BuiltinFunctionType) + # pylint:disable=invalid-name def visit_Expr(self, node): if isinstance(node.value, gast.Call): @@ -258,9 +254,9 @@ class CallTreeTransformer(transformer.Base): raise NotImplementedError('py_func with return values') else: if self.context.recursive: - raise NotImplementedError('Could not resolve target function.') + node = self._insert_dynamic_conversion(node) else: - # TODO(mdan): Double check. Is this reachable code? + # Unresolved functions are allowed in non-recursive mode. pass return node diff --git a/tensorflow/contrib/py2tf/converters/call_trees_test.py b/tensorflow/contrib/py2tf/converters/call_trees_test.py index e63c10de0fed72333a6d571f9b9a4f1cb50b5f1d..d482a9ef7897388839bbf8f9e4bfc5839d42b2d7 100644 --- a/tensorflow/contrib/py2tf/converters/call_trees_test.py +++ b/tensorflow/contrib/py2tf/converters/call_trees_test.py @@ -20,23 +20,11 @@ from __future__ import print_function from tensorflow.contrib.py2tf.converters import call_trees from tensorflow.contrib.py2tf.converters import converter_test_base -from tensorflow.contrib.py2tf.pyct import compiler from tensorflow.python.framework import constant_op from tensorflow.python.ops import math_ops from tensorflow.python.platform import test -class TestNamer(call_trees.FunctionNamer): - - def compiled_function_name(self, - original_fqn, - live_entity=None, - owner_type=None): - if owner_type is not None: - return None, False - return ('renamed_%s' % '_'.join(original_fqn)), True - - class CallTreesTest(converter_test_base.TestCase): def test_basic(self): @@ -50,14 +38,29 @@ class CallTreesTest(converter_test_base.TestCase): 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}, namer=TestNamer()) + 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_dynamic_function(self): + + def test_fn_1(): + raise ValueError('This should be masked by the mock.') + + def test_fn_2(f): + return f() + 3 + + node = self.parse_and_analyze(test_fn_2, {}) node = call_trees.transform(node, self.ctx, (), ()) - result = compiler.ast_to_object(node) - # Only test_fn_2 is transformed, so we'll insert renamed_test_fn_1 manually. - setattr(result, 'renamed_test_fn_1', renamed_test_fn_1) - self.assertEquals(3, result.test_fn_2(1)) + with self.compiled(node) as result: + # 10 = 7 (from the mock) + 3 (from test_fn_2) + self.assertEquals(10, result.test_fn_2(test_fn_1)) def test_simple_methods(self): @@ -71,13 +74,36 @@ class CallTreesTest(converter_test_base.TestCase): node = self.parse_and_analyze( TestClass.test_fn_2, {'TestClass': TestClass}, - namer=TestNamer(), + namer=converter_test_base.FakeNoRenameNamer(), arg_types={'self': (TestClass.__name__, TestClass)}) node = call_trees.transform(node, self.ctx, (), ()) - result = compiler.ast_to_object(node) - tc = TestClass() - self.assertEquals(3, result.test_fn_2(tc, 1)) + 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): @@ -86,26 +112,22 @@ class CallTreesTest(converter_test_base.TestCase): 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 - }, - namer=TestNamer()) + 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__,))), ()) - result = compiler.ast_to_object(node) - setattr(result, 'math_ops', math_ops) - setattr(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)) - result_val = sess.run(result_tensor) - self.assertEquals(3, result_val) + 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__': diff --git a/tensorflow/contrib/py2tf/converters/continue_canonicalization.py b/tensorflow/contrib/py2tf/converters/continue_statements.py similarity index 84% rename from tensorflow/contrib/py2tf/converters/continue_canonicalization.py rename to tensorflow/contrib/py2tf/converters/continue_statements.py index 486f0f6509d67d9d981e43ea6e5c77d14e6b23fc..4069a678b118b56b59d2e5491bb80cf52efd8143 100644 --- a/tensorflow/contrib/py2tf/converters/continue_canonicalization.py +++ b/tensorflow/contrib/py2tf/converters/continue_statements.py @@ -18,17 +18,17 @@ 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 ContinueCanonicalizationTransformer(gast.NodeTransformer): +class ContinueCanonicalizationTransformer(transformer.Base): """Canonicalizes continue statements into additional conditionals.""" - def __init__(self, namer): - self.namer = namer + def __init__(self, context): + super(ContinueCanonicalizationTransformer, self).__init__(context) # This is a stack structure, to correctly process nested loops. self.continuation_uses = [] @@ -76,7 +76,7 @@ class ContinueCanonicalizationTransformer(gast.NodeTransformer): return reorganized_nodes def _process_loop_block(self, block, scope): - cont_var = self.namer.new_symbol('cont_requested', scope.referenced) + 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]: @@ -87,7 +87,8 @@ class ContinueCanonicalizationTransformer(gast.NodeTransformer): def visit_While(self, node): self.generic_visit(node.test) node.body = self._process_loop_block(node.body, - anno.getanno(node, 'body_scope')) + anno.getanno(node, + NodeAnno.BODY_SCOPE)) for n in node.orelse: self.generic_visit(n) return node @@ -96,7 +97,8 @@ class ContinueCanonicalizationTransformer(gast.NodeTransformer): self.generic_visit(node.target) self.generic_visit(node.iter) node.body = self._process_loop_block(node.body, - anno.getanno(node, 'body_scope')) + anno.getanno(node, + NodeAnno.BODY_SCOPE)) for n in node.orelse: self.generic_visit(n) return node @@ -122,6 +124,4 @@ class ContinueCanonicalizationTransformer(gast.NodeTransformer): def transform(node, namer): - transformer = ContinueCanonicalizationTransformer(namer) - node = transformer.visit(node) - return node + return ContinueCanonicalizationTransformer(namer).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/continue_canonicalization_test.py b/tensorflow/contrib/py2tf/converters/continue_statements_test.py similarity index 50% rename from tensorflow/contrib/py2tf/converters/continue_canonicalization_test.py rename to tensorflow/contrib/py2tf/converters/continue_statements_test.py index c1fe903a2dd332626c8e64826652723c30ac412a..a598dcd1aed29478b7e3fe27e3c1b20010247dd9 100644 --- a/tensorflow/contrib/py2tf/converters/continue_canonicalization_test.py +++ b/tensorflow/contrib/py2tf/converters/continue_statements_test.py @@ -12,25 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for continue_canonicalization module.""" +"""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_canonicalization -from tensorflow.contrib.py2tf.converters import control_flow +from tensorflow.contrib.py2tf.converters import continue_statements from tensorflow.contrib.py2tf.converters import converter_test_base -from tensorflow.contrib.py2tf.pyct import compiler from tensorflow.python.platform import test -class TestNamer(control_flow.SymbolNamer): - - def new_symbol(self, name_root, _): - return name_root - - class ContinueCanonicalizationTest(converter_test_base.TestCase): def test_basic_continue(self): @@ -44,15 +36,15 @@ class ContinueCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v - node = self.parse_and_analyze(test_fn, {}, include_type_analysis=False) - node = continue_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node = continue_statements.transform(node, self.ctx) - 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)) + 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): @@ -65,14 +57,14 @@ class ContinueCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v - node = self.parse_and_analyze(test_fn, {}, include_type_analysis=False) - node = continue_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node = continue_statements.transform(node, self.ctx) - 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])) + 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): @@ -91,15 +83,15 @@ class ContinueCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v, u, w - node = self.parse_and_analyze(test_fn, {}, include_type_analysis=False) - node = continue_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node = continue_statements.transform(node, self.ctx) - 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)) + 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__': diff --git a/tensorflow/contrib/py2tf/converters/control_flow.py b/tensorflow/contrib/py2tf/converters/control_flow.py index a40c7b28f7bc3b8483b0b18cf11dbf99456df645..762c26f0c77e13c077761ceec41cb29db9149a35 100644 --- a/tensorflow/contrib/py2tf/converters/control_flow.py +++ b/tensorflow/contrib/py2tf/converters/control_flow.py @@ -21,7 +21,10 @@ 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): @@ -40,33 +43,65 @@ class SymbolNamer(object): raise NotImplementedError() -class SymbolRenamer(gast.NodeTransformer): - - def __init__(self, name_map): - self.name_map = name_map - - def visit_Name(self, node): - if node.id in self.name_map: - node.id = self.name_map[node.id] - return node - - -class ControlFlowTransformer(gast.NodeTransformer): +class ControlFlowTransformer(transformer.Base): """Transforms control flow structures like loops an conditionals.""" - def __init__(self, namer): - self.namer = namer + 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, 'body_scope') - orelse_scope = anno.getanno(node, 'orelse_scope') + 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( @@ -75,7 +110,7 @@ class ControlFlowTransformer(gast.NodeTransformer): raise ValueError( 'The else branch creates new symbols that the if branch does not.') - all_modified = tuple(body_scope.modified | orelse_scope.modified) + modified = tuple(body_scope.modified | orelse_scope.modified) all_referenced = body_scope.referenced | orelse_scope.referenced # Alias the closure variables inside the conditional functions @@ -86,80 +121,109 @@ class ControlFlowTransformer(gast.NodeTransformer): (body_scope.created | orelse_scope.created)) aliased_orig_names = tuple(need_alias) aliased_new_names = tuple( - self.namer.new_symbol(s, all_referenced) for s in aliased_orig_names) + 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 = node.body - node_body = [SymbolRenamer(alias_map).visit(n) for n in node_body] - node_orelse = node.orelse - node_orelse = [SymbolRenamer(alias_map).visit(n) for n in node_orelse] + 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) - if len(all_modified) == 1: - results = gast.Name(all_modified[0], None, 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: - results = gast.Tuple( - tuple(gast.Name(s, None, None) for s in all_modified), None) + body_returns = templates.replace('tf.ones(())')[0].value - template = """ - def body_name(): - aliased_new_names, = aliased_orig_names, - body - return (all_results,) - def orelse_name(): - aliased_new_names, = aliased_orig_names, - orelse - return (all_results,) - results = tf.cond(test, body_name, orelse_name) - """ - body_name = self.namer.new_symbol('if_true', all_referenced) - return templates.replace( - template, - test=node.test, - body_name=body_name, + 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, - orelse_name=self.namer.new_symbol('if_false', all_referenced), - orelse=node_orelse, + 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), - all_results=tuple(alias_map[s] if s in aliased_orig_names else s - for s in all_modified), - results=results) + 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, 'body_scope') - body_closure = tuple(body_scope.modified - body_scope.created) - - if len(body_closure) == 1: - state = body_closure[0] + 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) + if not state: + # TODO(mdan): Implement this properly. + # To complete this statement, we need to check whether any variable + # created inside the body scope is used before being modified outside the + # scope. This should be done during activity analysis, and in general + # should cover the case where variables may not be initialized. + raise ValueError('cannot convert while loop: no outputs') + + 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 = tuple(body_closure) - state_ast_tuple = gast.Tuple( - tuple(gast.Name(n, None, None) for n in state), None) + 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): + def test_name(state_ssf): return test - def body_name(state): + def body_name(state_ssf): body - return state, - state_ast_tuple = tf.while_loop(test_name, body_name, [state]) + 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.namer.new_symbol('loop_test', body_scope.referenced), - test=node.test, - body_name=self.namer.new_symbol('loop_body', body_scope.referenced), - body=node.body) + 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, namer): - transformer = ControlFlowTransformer(namer) - node = transformer.visit(node) +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 index 054e33750dbae86559a9575dfecde64132b9a2cd..b785b284a7fb7a0257551326c88b44a341b295ba 100644 --- a/tensorflow/contrib/py2tf/converters/control_flow_test.py +++ b/tensorflow/contrib/py2tf/converters/control_flow_test.py @@ -20,23 +20,11 @@ from __future__ import print_function from tensorflow.contrib.py2tf.converters import control_flow from tensorflow.contrib.py2tf.converters import converter_test_base -from tensorflow.contrib.py2tf.pyct import compiler from tensorflow.python.framework import constant_op from tensorflow.python.ops import control_flow_ops from tensorflow.python.platform import test -class TestNamer(control_flow.SymbolNamer): - - 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 - - class ControlFlowTest(converter_test_base.TestCase): def test_simple_while(self): @@ -50,13 +38,12 @@ class ControlFlowTest(converter_test_base.TestCase): return s, i, n node = self.parse_and_analyze(test_fn, {}) - node = control_flow.transform(node, TestNamer()) - result = compiler.ast_to_object(node) - setattr(result, 'tf', control_flow_ops) + node = control_flow.transform(node, self.ctx) - with self.test_session() as sess: - self.assertEqual((10, 5, 5), - sess.run(result.test_fn(constant_op.constant(5)))) + 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): @@ -66,12 +53,11 @@ class ControlFlowTest(converter_test_base.TestCase): return n node = self.parse_and_analyze(test_fn, {}) - node = control_flow.transform(node, TestNamer()) - result = compiler.ast_to_object(node) - setattr(result, 'tf', control_flow_ops) + node = control_flow.transform(node, self.ctx) - with self.test_session() as sess: - self.assertEqual(0, sess.run(result.test_fn(constant_op.constant(5)))) + 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): @@ -85,15 +71,14 @@ class ControlFlowTest(converter_test_base.TestCase): return a, b node = self.parse_and_analyze(test_fn, {}) - node = control_flow.transform(node, TestNamer()) - result = compiler.ast_to_object(node) - setattr(result, 'tf', control_flow_ops) + node = control_flow.transform(node, self.ctx) - 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)))) + 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): @@ -103,12 +88,11 @@ class ControlFlowTest(converter_test_base.TestCase): return n node = self.parse_and_analyze(test_fn, {}) - node = control_flow.transform(node, TestNamer()) - result = compiler.ast_to_object(node) - setattr(result, 'tf', control_flow_ops) + node = control_flow.transform(node, self.ctx) - with self.test_session() as sess: - self.assertEqual(-1, sess.run(result.test_fn(constant_op.constant(1)))) + 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__': diff --git a/tensorflow/contrib/py2tf/converters/converter_test_base.py b/tensorflow/contrib/py2tf/converters/converter_test_base.py index 6bfa55443c4fb09bdac4872d2fd2a9aefda74298..8c08c5492a4b10d4abb0ec3b19b39d5b17e41a0a 100644 --- a/tensorflow/contrib/py2tf/converters/converter_test_base.py +++ b/tensorflow/contrib/py2tf/converters/converter_test_base.py @@ -18,15 +18,88 @@ 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.static_analysis import access +from tensorflow.contrib.py2tf.pyct import pretty_printer +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 FakeNoRenameNamer(FakeNamer): + + def compiled_function_name(self, original_fqn, **_): + return str(original_fqn), False + + class TestCase(test.TestCase): + """Base class for unit tests in this module. Contains relevant utilities.""" + + @contextlib.contextmanager + def compiled(self, node, *symbols): + source = None + + self.dynamic_calls = [] + def converted_call(*args): + """Mock version of api.converted_call.""" + self.dynamic_calls.append(args) + return 7 + + try: + result, source = compiler.ast_to_object(node) + result.tf = self.make_fake_mod('fake_tf', *symbols) + result.py2tf_utils = utils + result.py2tf_api = self.make_fake_mod('fake_api', converted_call) + yield result + except Exception: # pylint:disable=broad-except + if source is None: + print('Offending AST:\n%s' % pretty_printer.fmt(node, color=False)) + else: + print('Offending compiled code:\n%s' % source) + raise + + def make_fake_mod(self, name, *symbols): + fake_mod = imp.new_module(name) + for s in symbols: + if hasattr(s, '__name__'): + setattr(fake_mod, s.__name__, s) + elif hasattr(s, 'name'): + # This is a bit of a hack, but works for things like tf.int32 + setattr(fake_mod, s.name, s) + else: + raise ValueError('can not attach %s - what should be its name?' % s) + return fake_mod + + def attach_namespace(self, module, **ns): + for k, v in ns.items(): + setattr(module, k, v) def parse_and_analyze(self, test_fn, @@ -34,17 +107,21 @@ class TestCase(test.TestCase): namer=None, arg_types=None, include_type_analysis=True, + owner_type=None, recursive=True): node, source = parser.parse_entity(test_fn) ctx = context.EntityContext( - namer=namer, + namer=namer or FakeNamer(), source_code=source, source_file=None, namespace=namespace, arg_values=None, arg_types=arg_types, - recursive=recursive) - node = access.resolve(node, ctx) + owner_type=owner_type, + recursive=recursive, + type_annotation_func=utils.set_element_type) + 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) diff --git a/tensorflow/contrib/py2tf/converters/decorators.py b/tensorflow/contrib/py2tf/converters/decorators.py index 3f620c1cd2d9b75f82410754a7e812e13eabe3ae..68bf241ef33292f0581ccb3c44f313f853c92ba7 100644 --- a/tensorflow/contrib/py2tf/converters/decorators.py +++ b/tensorflow/contrib/py2tf/converters/decorators.py @@ -33,6 +33,7 @@ class DecoratorsTransformer(gast.NodeTransformer): def __init__(self, remove_decorators): self.remove_decorators = remove_decorators + self.additional_dependencies = set() # pylint:disable=invalid-name @@ -44,13 +45,38 @@ class DecoratorsTransformer(gast.NodeTransformer): dec_func = dec.func else: dec_func = dec + + # Special cases. + # TODO(mdan): Is there any way we can treat these more generically? + # We may want to forego using decorators altogether if we can't + # properly support them. + if isinstance(dec_func, gast.Name) and dec_func.id in ('classmethod',): + # Assumption: decorators are only visible in the AST when converting + # a function inline (via another decorator). + # In that case, the converted function is no longer part of the + # original object that it was declared into. + # This is currently verified by tests. + continue + 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 + kept_decorators.append((dec, dec_value)) + + for _, dec_value in kept_decorators: + if dec_value.__module__ == '__main__': + raise ValueError( + 'decorator "%s" was not allowed because it is declared ' + 'in the module "%s". To fix this, declare it in a separate ' + 'module that we can import it from.' % (dec_value, + dec_value.__module__)) + else: + self.additional_dependencies.add(dec_value) + + node.decorator_list = [dec for dec, _ in kept_decorators] return node # pylint:enable=invalid-name @@ -59,4 +85,4 @@ class DecoratorsTransformer(gast.NodeTransformer): def transform(node, remove_decorators): transformer = DecoratorsTransformer(remove_decorators) node = transformer.visit(node) - return node + return node, transformer.additional_dependencies diff --git a/tensorflow/contrib/py2tf/converters/decorators_test.py b/tensorflow/contrib/py2tf/converters/decorators_test.py index f50d593043aeb76d63beb3cb6c301122c9ed8948..c75e5461746f27d14a54b7ac06e7f77d868372c8 100644 --- a/tensorflow/contrib/py2tf/converters/decorators_test.py +++ b/tensorflow/contrib/py2tf/converters/decorators_test.py @@ -18,78 +18,121 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import textwrap +from functools import wraps 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): +# The Python parser only briefly captures decorators into the AST. +# The interpreter desugars them on load, and the decorated function loses any +# trace of the decorator (which is notmally what you would expect, since +# they are meant to be transparent). +# However, decorators are still visible when you analyze the function +# from inside a decorator, before it was applied - as is the case +# with our conversion decorators. - def test_function_decorator(self): - def function_decorator(): +def simple_decorator(f): + return lambda a: f(a) + 1 - def decorator(f): - return lambda a: f(a) + 1 - return decorator +def self_removing_decorator(removing_wrapper): + def decorator(f): + @wraps(f) + def wrapper(*args): + # This removing wrapper is defined in the test below. This setup is so + # intricate just to simulate how we use the transformer in practice. + transformed_f = removing_wrapper(f, (self_removing_decorator,)) + return transformed_f(*args) + 1 + return wrapper + 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 +class DecoratorsTest(converter_test_base.TestCase): + + def _remover_wrapper(self, f, remove_decorators): + namespace = { + 'self_removing_decorator': self_removing_decorator, + 'simple_decorator': simple_decorator + } + node = self.parse_and_analyze(f, namespace) + node, _ = decorators.transform(node, remove_decorators=remove_decorators) + result, _ = compiler.ast_to_object(node) + return getattr(result, f.__name__) - node = self.parse_and_analyze(static_wrapper, - {'function_decorator': function_decorator}) - node = node.body[0].body[0] + def test_noop(self): - node = decorators.transform(node, remove_decorators=()) - result = compiler.ast_to_object( - node, - source_prefix=textwrap.dedent(tf_inspect.getsource(function_decorator))) - self.assertEqual(2, result.test_fn(1)) + def test_fn(a): + return a - node = decorators.transform(node, remove_decorators=(function_decorator,)) - result = compiler.ast_to_object(node) + node = self.parse_and_analyze(test_fn, {}) + node, deps = decorators.transform(node, remove_decorators=()) + result, _ = compiler.ast_to_object(node) + + self.assertFalse(deps) self.assertEqual(1, result.test_fn(1)) - def test_simple_decorator(self): + def test_function(self): - def simple_decorator(f): - return lambda a: f(a) + 1 + @self_removing_decorator(self._remover_wrapper) + def test_fn(a): + return a - # 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(): + # 2 = 1 (a) + 1 (decorator applied exactly once) + self.assertEqual(2, test_fn(1)) - @simple_decorator - def test_fn(a): # pylint:disable=unused-variable + def test_method(self): + + class TestClass(object): + + @self_removing_decorator(self._remover_wrapper) + def test_fn(self, a): return a - node = self.parse_and_analyze(static_wrapper, - {'simple_decorator': simple_decorator}) - node = node.body[0].body[0] + # 2 = 1 (a) + 1 (decorator applied exactly once) + self.assertEqual(2, TestClass().test_fn(1)) - node = decorators.transform(node, remove_decorators=()) - result = compiler.ast_to_object( - node, - source_prefix=textwrap.dedent(tf_inspect.getsource(simple_decorator))) - self.assertEqual(2, result.test_fn(1)) + def test_multiple_decorators(self): - node = decorators.transform(node, remove_decorators=(simple_decorator,)) - result = compiler.ast_to_object(node) - self.assertEqual(1, result.test_fn(1)) + class TestClass(object): + + # Note that reversing the order of this two doesn't work. + @classmethod + @self_removing_decorator(self._remover_wrapper) + def test_fn(cls, a): + return a + + # 2 = 1 (a) + 1 (decorator applied exactly once) + self.assertEqual(2, TestClass.test_fn(1)) + + def test_nested_decorators(self): + + @self_removing_decorator(self._remover_wrapper) + def test_fn(a): + @simple_decorator + def inner_fn(b): + return b + 11 + return inner_fn(a) + + with self.assertRaises(ValueError): + test_fn(1) + + # TODO(mdan): Uncomment this test once converter_test_base is updated. + # (can't do it now because it has unrelated pending changes) + # def test_nested_decorators(self): + # + # @self_removing_decorator(self._remover_wrapper) + # def test_fn(a): + # @imported_decorator + # def inner_fn(b): + # return b + 11 + # return inner_fn(a) + # + # # 14 = 1 (a) + 1 (simple_decorator) + 11 (inner_fn) + # self.assertEqual(14, test_fn(1)) if __name__ == '__main__': diff --git a/tensorflow/contrib/py2tf/converters/for_canonicalization.py b/tensorflow/contrib/py2tf/converters/for_loops.py similarity index 62% rename from tensorflow/contrib/py2tf/converters/for_canonicalization.py rename to tensorflow/contrib/py2tf/converters/for_loops.py index c284689b904c6f372f30e83c259416a51babe4a6..4297c1cf2a3632e097973280cc985fc48da64475 100644 --- a/tensorflow/contrib/py2tf/converters/for_canonicalization.py +++ b/tensorflow/contrib/py2tf/converters/for_loops.py @@ -22,32 +22,33 @@ 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 ForLoopCanonicalizationTransformer(gast.NodeTransformer): +class ForLoopCanonicalizationTransformer(transformer.Base): """Canonicalizes for loops (e.g. into while loops).""" - def __init__(self, namer): - self.namer = namer + def __init__(self, context): + super(ForLoopCanonicalizationTransformer, self).__init__(context) def visit_For(self, node): self.generic_visit(node) - body_scope = anno.getanno(node, 'body_scope') - - # TODO(mdan): Distinguish between `for i in n` and `for i in range(n)` - # Or maybe we should replace range with tf.range? - + body_scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + i_var = self.context.namer.new_symbol('i', body_scope.referenced) + n_var = self.context.namer.new_symbol('n', body_scope.referenced) + iterated_var = self.context.namer.new_symbol('iterated', + body_scope.referenced) + # TODO(mdan): Use TensorListFromTensor(loop_iter) here. if anno.hasanno(node, 'extra_cond'): template = """ i = 0 - n = len(loop_iter) + iterated = loop_iter + n = len(iterated) while i < n and extra_cond: - # TODO(mdan): Use TensorListFromTensor(loop_iter) here. - target = loop_iter[i] + target = iterated[i] body i += 1 """ @@ -56,26 +57,29 @@ class ForLoopCanonicalizationTransformer(gast.NodeTransformer): loop_iter=node.iter, target=node.target, body=node.body, - i=self.namer.new_symbol('i', body_scope.referenced), - n=self.namer.new_symbol('n', body_scope.referenced), + i=i_var, + n=n_var, + iterated=iterated_var, extra_cond=anno.getanno(node, 'extra_cond')) else: template = """ i = 0 - n = len(loop_iter) + iterated = loop_iter + n = len(iterated) while i < n: - # TODO(mdan): Use TensorListFromTensor(loop_iter) here. - target = loop_iter[i] - body # pylint:disable=pointless-statement + target = iterated[i] + body i += 1 """ - return templates.replace( + repl = templates.replace( template, loop_iter=node.iter, target=node.target, body=node.body, - i=self.namer.new_symbol('i', body_scope.referenced), - n=self.namer.new_symbol('n', body_scope.referenced)) + i=i_var, + n=n_var, + iterated=iterated_var) + return repl def visit_Continue(self, node): assert False, 'continue statement should be desugared at this point' @@ -84,7 +88,5 @@ class ForLoopCanonicalizationTransformer(gast.NodeTransformer): assert False, 'break statement should be desugared at this point' -def transform(node, namer): - transformer = ForLoopCanonicalizationTransformer(namer) - node = transformer.visit(node) - return node +def transform(node, context): + return ForLoopCanonicalizationTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/for_canonicalization_test.py b/tensorflow/contrib/py2tf/converters/for_loops_test.py similarity index 55% rename from tensorflow/contrib/py2tf/converters/for_canonicalization_test.py rename to tensorflow/contrib/py2tf/converters/for_loops_test.py index a6e6350fd45e9c9575af9c12d3d0c4e9b89bee41..b6e3e8c8d8d4960977e2b72b56a3fab8329ad2a7 100644 --- a/tensorflow/contrib/py2tf/converters/for_canonicalization_test.py +++ b/tensorflow/contrib/py2tf/converters/for_loops_test.py @@ -12,25 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for for_canonicalization module.""" +"""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 control_flow from tensorflow.contrib.py2tf.converters import converter_test_base -from tensorflow.contrib.py2tf.converters import for_canonicalization -from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.contrib.py2tf.converters import for_loops from tensorflow.python.platform import test -class TestNamer(control_flow.SymbolNamer): - - def new_symbol(self, name_root, _): - return name_root - - class ControlFlowTest(converter_test_base.TestCase): def test_basic_for(self): @@ -42,13 +34,36 @@ class ControlFlowTest(converter_test_base.TestCase): return s node = self.parse_and_analyze(test_fn, {}) - node = for_canonicalization.transform(node, TestNamer()) - result = compiler.ast_to_object(node) + 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)) + + def test_for_with_iterated_expression(self): + + eval_count = [0] + + def count_evals(x): + eval_count[0] += 1 + return x + + def test_fn(n): + s = 0 + for e in count_evals(range(n)): + s += e + return s + + node = self.parse_and_analyze(test_fn, {'count_evals': count_evals}) + node = for_loops.transform(node, self.ctx) - l = [1, 2, 3] - self.assertEqual(test_fn(l), result.test_fn(l)) - l = [] - self.assertEqual(test_fn(l), result.test_fn(l)) + with self.compiled(node) as result: + result.count_evals = count_evals + self.assertEqual(test_fn(5), result.test_fn(5)) + # count_evals ran twice, once for test_fn and another for result.test_fn + self.assertEqual(eval_count[0], 2) if __name__ == '__main__': diff --git a/tensorflow/contrib/py2tf/converters/ifexp.py b/tensorflow/contrib/py2tf/converters/ifexp.py new file mode 100644 index 0000000000000000000000000000000000000000..5fd6f348af0df81a6ff35745da603bd431130e20 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/ifexp.py @@ -0,0 +1,49 @@ +# 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 the ternary conditional operator.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer + + +class IfExp(transformer.Base): + """Canonicalizes all IfExp nodes into plain conditionals.""" + + def visit_IfExp(self, node): + template = """ + py2tf_utils.run_cond(test, lambda: body, lambda: orelse) + """ + desugared_ifexp = templates.replace_as_expression( + template, test=node.test, body=node.body, orelse=node.orelse) + return desugared_ifexp + + +def transform(node, context): + """Desugar IfExp nodes into plain conditionals. + + Args: + node: an AST node to transform + context: a context object + + Returns: + new_node: an AST with no IfExp nodes, only conditionals. + """ + + node = IfExp(context).visit(node) + return node diff --git a/tensorflow/contrib/py2tf/converters/ifexp_test.py b/tensorflow/contrib/py2tf/converters/ifexp_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9c357ef35b550833bcb79d39f0bdbc6d758d31a5 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/ifexp_test.py @@ -0,0 +1,106 @@ +# 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 ifexp module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import ifexp +from tensorflow.python.platform import test + + +class IfExpTest(converter_test_base.TestCase): + + def compiled_fn(self, test_fn, *args): + node = self.parse_and_analyze(test_fn, {}) + node = ifexp.transform(node, self.ctx) + module = self.compiled(node, *args) + return module + + def test_simple(self): + + def test_fn(x): + return 1 if x else 0 + + with self.compiled_fn(test_fn) as result: + result.py2tf_util = utils + for x in [0, 1]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_fn(self): + + def f(x): + return 3 * x + + def test_fn(x): + y = f(x * x if x > 0 else x) + return y + + with self.compiled_fn(test_fn) as result: + result.py2tf_util = utils + result.f = f + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_exp(self): + + def test_fn(x): + return x * x if x > 0 else x + + with self.compiled_fn(test_fn) as result: + result.py2tf_util = utils + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_nested(self): + + def test_fn(x): + return x * x if x > 0 else x if x else 1 + + with self.compiled_fn(test_fn) as result: + result.py2tf_util = utils + for x in [-2, 0, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_in_cond(self): + + def test_fn(x): + if x > 0: + return x * x if x < 5 else x * x * x + return -x + + with self.compiled_fn(test_fn) as result: + result.py2tf_util = utils + for x in [-2, 2, 5]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_assign_in_cond(self): + + def test_fn(x): + if x > 0: + x = -x if x < 5 else x + return x + + with self.compiled_fn(test_fn) as result: + result.py2tf_util = utils + for x in [-2, 2, 5]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + +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/lists.py b/tensorflow/contrib/py2tf/converters/lists.py new file mode 100644 index 0000000000000000000000000000000000000000..06e1dad8f4d652da78ed39309f5b40598e368ea6 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/lists.py @@ -0,0 +1,103 @@ +# 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 list operations. + +This includes converting Python lists to TensorArray/TensorList. +""" + +# TODO(mdan): Elaborate the logic here. +# TODO(mdan): Does it even make sense to attempt to try to use TAs? +# The current rule (always convert to TensorArray) is naive and insufficient. +# In general, a better mechanism could look like: +# * convert to TensorList by default +# * leave as Python list if the user explicitly forbids it +# * convert to TensorArray only when complete write once behavior can be +# guaranteed (e.g. list comprehensions) + +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.python.framework import dtypes + + +class ListTransformer(transformer.Base): + """Converts lists and related operations to their TF counterpart.""" + + def _empty_list(self, node): + if not anno.hasanno(node, 'element_type'): + raise NotImplementedError( + 'type inference for empty lists is not yet supported; ' + 'use utils.set_element_type(, ) to continue') + dtype = anno.getanno(node, 'element_type') + if not isinstance(dtype, dtypes.DType): + # TODO(mdan): Allow non-TF dtypes? + # That would be consistent with the dynamic dispatch pattern, but + # we must make sure that doesn't become confusing. + raise NotImplementedError('element type "%s" not yet supported' % dtype) + + dtype_name = dtype.name + # TODO(mdan): Does it ever make sense not to use tensor lists? + template = """ + tf.TensorArray(tf.dtype_name, size=0, dynamic_size=True) + """ + return templates.replace_as_expression(template, dtype_name=dtype_name) + + def _pre_populated_list(self, node): + raise NotImplementedError() + + def visit_Expr(self, node): + node = self.generic_visit(node) + if isinstance(node.value, gast.Call): + call_node = node.value + qn = anno.getanno(call_node.func, anno.Basic.QN) + + if qn.qn[-1] == 'append' and (len(call_node.args) == 1): + template = """ + target = py2tf_utils.dynamic_list_append(target, element) + """ + node = templates.replace( + template, + target=qn.parent.ast(), + element=call_node.args[0]) + return node + + def visit_Assign(self, node): + node = self.generic_visit(node) + + # Only convert lists when they are assigned to a variable, e.g.: + # l = [] + # TODO(mdan): This rule should be improved. + if len(node.targets) != 1: + return node + if not isinstance(node.value, gast.List): + return node + if not isinstance(node.value.ctx, gast.Load): + return node + + if node.value.elts: + node.value = self._pre_populated_list(node.value) + else: + node.value = self._empty_list(node.value) + return node + + +def transform(node, context): + return ListTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/lists_test.py b/tensorflow/contrib/py2tf/converters/lists_test.py new file mode 100644 index 0000000000000000000000000000000000000000..671a1cc7b1225061a00731596c536c4403e0bdff --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/lists_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 lists module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import lists +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.platform import test + + +class ListTest(converter_test_base.TestCase): + + def test_empty_annotated_list(self): + + def test_fn(): + l = [] + utils.set_element_type(l, dtypes.int32) + l.append(1) + return l + + node = self.parse_and_analyze(test_fn, {'dtypes': dtypes, 'utils': utils}) + node = lists.transform(node, self.ctx) + + with self.compiled(node, tensor_array_ops.TensorArray, + dtypes.int32) as result: + # TODO(mdan): Attach these additional modules automatically. + result.utils = utils + result.dtypes = dtypes + with self.test_session() as sess: + self.assertEqual(test_fn(), sess.run(result.test_fn().stack())) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/logical_expressions.py b/tensorflow/contrib/py2tf/converters/logical_expressions.py index df980d41c9c57e325bee9a1fa870d9c95f46ea41..10192e6a036c4a44aa1e6f1b4a390579bd703373 100644 --- a/tensorflow/contrib/py2tf/converters/logical_expressions.py +++ b/tensorflow/contrib/py2tf/converters/logical_expressions.py @@ -23,52 +23,116 @@ from __future__ import print_function import gast -from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer -class LogicalExpressionTransformer(gast.NodeTransformer): +# TODO(mdan): Properly extrack boolean ops according to lazy eval rules. +# Note that this isn't completely safe either, because tensors may have control +# dependencies. +# Note that for loops that should be done after the loop was converted to +# tf.while_loop so that the expanded conditionals are properly scoped. + +# Used to signal that an operand is safe for non-lazy evaluation. +SAFE_BOOLEAN_OPERAND = 'SAFE_BOOLEAN_OPERAND' + + +class LogicalExpressionTransformer(transformer.Base): """Converts logical expressions to corresponding TF calls.""" - def __init__(self): + def __init__(self, context): + super(LogicalExpressionTransformer, self).__init__(context) # 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', + gast.And: 'logical_and', + gast.Eq: 'equal', + gast.Gt: 'greater', + gast.GtE: 'greater_equal', + gast.Lt: 'less', + gast.LtE: 'less_equal', + gast.Not: 'logical_not', + gast.NotEq: 'not_equal', + gast.Or: 'logical_or', + gast.USub: 'negative', + gast.Is: 'py2tf_utils.dynamic_is', + gast.IsNot: 'py2tf_utils.dynamic_is_not' } + def _expect_simple_symbol(self, operand): + if isinstance(operand, gast.Name): + return + if anno.hasanno(operand, SAFE_BOOLEAN_OPERAND): + return + raise NotImplementedError( + 'only simple local variables are supported in logical and compound ' + 'comparison expressions; for example, we support "a or b" but not ' + '"a.x or b"; for a workaround, assign the expression to a local ' + 'variable and use that instead, for example "tmp = a.x", "tmp or b"') + + def _matching_tf_op(self, operator): + op_type = type(operator) + mapped_op = self.op_mapping.get(op_type) + if not mapped_op: + raise NotImplementedError('operator %s is not yet supported' % op_type) + return mapped_op + + def _inline_tf_op(self, op_name, args): + if 'py2tf_utils' in op_name: + # TODO(alexbw): explicitly spelling out the attribute function name + # until fix for issue highlighted in cl/188931581 lands. + template = """ + py2tf_utils.op_name(args) + """ + op_name = op_name.replace('py2tf_utils.', '') + else: + template = """ + tf.op_name(args) + """ + replacement = templates.replace_as_expression( + template, op_name=op_name, args=args) + anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True) + return replacement + 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 + ops_and_comps = list(zip(node.ops, node.comparators)) + left = node.left + op_tree = None + + # Repeated comparisons are converted to conjunctions: + # a < b < c -> a < b and b < c + while ops_and_comps: + op, right = ops_and_comps.pop(0) + binary_comparison = self._inline_tf_op(self._matching_tf_op(op), + (left, right)) + if isinstance(left, gast.Name) and isinstance(right, gast.Name): + anno.setanno(binary_comparison, SAFE_BOOLEAN_OPERAND, True) + if op_tree: + self._expect_simple_symbol(right) + op_tree = self._inline_tf_op('logical_and', + (binary_comparison, op_tree)) + else: + op_tree = binary_comparison + left = right + assert op_tree is not None + return op_tree 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 + return self._inline_tf_op(self._matching_tf_op(node.op), node.operand) 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 + node_values = node.values + right = node.values.pop() + self._expect_simple_symbol(right) + while node_values: + left = node_values.pop() + self._expect_simple_symbol(left) + right = self._inline_tf_op(self._matching_tf_op(node.op), (left, right)) + return right + + +def transform(node, context): + return LogicalExpressionTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/logical_expressions_test.py b/tensorflow/contrib/py2tf/converters/logical_expressions_test.py index d711065099b24ad814104e6460e6ca551b31b3e6..eb28c309a429f2267cc1ae1f6f65a8cde0ad91b8 100644 --- a/tensorflow/contrib/py2tf/converters/logical_expressions_test.py +++ b/tensorflow/contrib/py2tf/converters/logical_expressions_test.py @@ -20,7 +20,6 @@ from __future__ import print_function from tensorflow.contrib.py2tf.converters import converter_test_base from tensorflow.contrib.py2tf.converters import logical_expressions -from tensorflow.contrib.py2tf.pyct import compiler from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -33,13 +32,12 @@ class GradientsFunctionTest(converter_test_base.TestCase): return a == b node = self.parse_and_analyze(test_fn, {}) - node = logical_expressions.transform(node) - result = compiler.ast_to_object(node) - setattr(result, 'tf', math_ops) + node = logical_expressions.transform(node, self.ctx) - with self.test_session() as sess: - self.assertTrue(sess.run(result.test_fn(1, 1))) - self.assertFalse(sess.run(result.test_fn(1, 2))) + 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): @@ -47,12 +45,12 @@ class GradientsFunctionTest(converter_test_base.TestCase): return (a or b) and (a or b or c) node = self.parse_and_analyze(test_fn, {}) - node = logical_expressions.transform(node) - result = compiler.ast_to_object(node) - setattr(result, 'tf', math_ops) + node = logical_expressions.transform(node, self.ctx) - with self.test_session() as sess: - self.assertTrue(sess.run(result.test_fn(True, False, True))) + 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__': diff --git a/tensorflow/contrib/py2tf/converters/name_scopes.py b/tensorflow/contrib/py2tf/converters/name_scopes.py new file mode 100644 index 0000000000000000000000000000000000000000..c702823fcf047fcad3254318bd323d2b8fddd700 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/name_scopes.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. +# ============================================================================== +"""Wraps a function body with a `name_scope` of the function name. +""" + +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 FunctionNameScopeTransformer(transformer.Base): + """Wrap a function body with a `name_scope` of the function name.""" + + def __init__(self, context): + super(FunctionNameScopeTransformer, self).__init__(context) + self._function_level = 0 + + def visit_FunctionDef(self, node): + self._function_level += 1 + try: + self.generic_visit(node) + finally: + self._function_level -= 1 + scope_name = node.name + if self._function_level == 0 and self.context.owner_type is not None: + scope_name = '{}/{}'.format(self.context.owner_type.__name__, scope_name) + node.body = templates.replace( + 'with tf.name_scope(scope_name): body', + scope_name=gast.Str(scope_name), + body=node.body) + return node + + +def transform(node, context): + return FunctionNameScopeTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/name_scopes_test.py b/tensorflow/contrib/py2tf/converters/name_scopes_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a8ca341602ee5f06dbb812643a58794339d98afe --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/name_scopes_test.py @@ -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. +# ============================================================================== +"""Tests for for_canonicalization 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 name_scopes +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.platform import test + + +class FunctionNameScopeTransformer(converter_test_base.TestCase): + + def test_basic_name(self): + + def test_fn(l): + a = 5 + l += a + return l + + node = self.parse_and_analyze(test_fn, {}) + node = name_scopes.transform(node, self.ctx) + + with self.compiled(node, ops.name_scope) as result: + result_op = result.test_fn(constant_op.constant([1, 2, 3])) + self.assertIn('test_fn/', result_op.op.name) + + def test_nested_name(self): + + def test_fn(l): + + def body(i): + return i**2 + + l += [4] + return body(l) + + node = self.parse_and_analyze(test_fn, {}) + node = name_scopes.transform(node, self.ctx) + + with self.compiled(node, ops.name_scope) as result: + result_op = result.test_fn(constant_op.constant([1, 2, 3])) + first_result_input_name = result_op.op.inputs[0].name + second_result_input_name = result_op.op.inputs[1].name + self.assertIn('test_fn/', first_result_input_name) + self.assertNotIn('body/', first_result_input_name) + self.assertIn('test_fn/body/', second_result_input_name) + + def test_class_name(self): + + class TestClass(object): + + def test_fn(self, l): + + def body(i): + return i**2 + + l += [4] + return body(l) + + # Note that 'TestClass' was needed in the namespace here. + node = self.parse_and_analyze( + TestClass, {'TestClass': TestClass}, owner_type=TestClass) + node = name_scopes.transform(node, self.ctx) + + with self.compiled(node, ops.name_scope) as result: + result_op = result.TestClass().test_fn(constant_op.constant([1, 2, 3])) + first_result_input_name = result_op.op.inputs[0].name + second_result_input_name = result_op.op.inputs[1].name + self.assertIn('TestClass/test_fn/', first_result_input_name) + self.assertNotIn('body/', first_result_input_name) + self.assertIn('TestClass/test_fn/body/', second_result_input_name) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/print_functions.py b/tensorflow/contrib/py2tf/converters/print_functions.py deleted file mode 100644 index 5da738c4954fb628212562b73641e1fc27032168..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/py2tf/converters/print_functions.py +++ /dev/null @@ -1,51 +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. -# ============================================================================== -"""Compatibility support. Converts Print nodes to function calls.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import gast - -from tensorflow.contrib.py2tf.pyct import anno - - -class PrintFunctionTransformer(gast.NodeTransformer): - """Transforms Print nodes to Call so they can be handled as functions.""" - - # pylint:disable=invalid-name - - def visit_Print(self, node): - self.generic_visit(node) - for n in node.values: - n.ctx = gast.Param() - call_node = gast.Call( - func=gast.Name('print', gast.Load(), None), - args=node.values, - keywords=[]) - anno.setanno(call_node.func, 'live_val', print) - anno.setanno(call_node.func, 'fqn', 'print') - anno.setanno(call_node, 'args_scope', anno.getanno(node, 'args_scope')) - node = gast.Expr(call_node) - return node - - # pylint:enable=invalid-name - - -def transform(node): - transformer = PrintFunctionTransformer() - node = transformer.visit(node) - return node diff --git a/tensorflow/contrib/py2tf/converters/side_effect_guards.py b/tensorflow/contrib/py2tf/converters/side_effect_guards.py index ae9632396673b5d32c2cd16a51d3701387ac6d5f..30976b3ec6db5a6607023ac804d9d54cfb296190 100644 --- a/tensorflow/contrib/py2tf/converters/side_effect_guards.py +++ b/tensorflow/contrib/py2tf/converters/side_effect_guards.py @@ -34,12 +34,14 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from contextlib import contextmanager - 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): @@ -57,52 +59,61 @@ class SymbolNamer(object): raise NotImplementedError() -class SideEffectGuardTransformer(gast.NodeTransformer): +class SideEffectGuardTransformer(transformer.Base): """Adds control dependencies to functions with side effects.""" - def __init__(self, namer): - self.namer = namer - self.indent_next = False - self.next_indent_owner = None + 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 self.indent_next: - assert self.next_indent_owner is not None - current_dest.append(self.next_indent_owner) - current_dest = self.next_indent_owner.body - self.next_indent_owner = None - self.indent_next = False - if not current_dest: + 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 that can we can ' - 'condition on the statement.') + 'it is not followed by any computation which ' + 'the statement could gate.') return new_nodes def visit_FunctionDef(self, node): - if anno.hasanno(node, 'skip_processing'): - return node node.body = self._visit_and_reindent(node.body) return node - def _gate_symbols(self, guard_statement, guarded_args): - # TODO(mdan): This won't work for variables. - template = """ - (args,) = (tf.identity(a) for a in (args,)) - """ - guards = templates.replace(template, args=tuple(guarded_args)) - guard_statement.body.extend(guards) - return guard_statement + 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) @@ -111,46 +122,69 @@ class SideEffectGuardTransformer(gast.NodeTransformer): # opt.minimize(loss) # or: # tf.py_func(...) - args_scope = anno.getanno(node.value, 'args_scope') - temp_name = self.namer.new_symbol('temp', args_scope.parent.referenced) - # TODO(mdan): Unsafe reference modification! - args_scope.mark_write(temp_name) - template = """ - temp_result = call - if temp_result is not None: - if not isinstance(temp_result, (list, tuple)): - temp_result = (temp_result,) - ctx = tf.control_dependencies(temp_result) - else: - ctx = contextmanager(lambda: (yield))() - with ctx: - # TODO(mdan): Also insert ops to re-fetch if variables are involved. - pass # Will be removed below. - """ - # TODO(mdan): This is brittle. Reorganize the mechanism. - statements = templates.replace( - template, call=node.value, temp_result=temp_name) - control_deps_guard = statements[-1] - control_deps_guard.body = [] # 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. - guarded_args = tuple(args_scope.used & (args_scope.parent.modified - | args_scope.parent.returned)) + 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: - node = tuple(statements[:-1]) + ( - self._gate_symbols(control_deps_guard, 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: - node = tuple(statements[:-1]) - # The mechanism will insert the guard statement later. - self.indent_next = True - self.next_indent_owner = control_deps_guard + 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, namer): - transformer = SideEffectGuardTransformer(namer) - return transformer.visit(node) +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 index 5c56973dc2ae5d1976a68f040772e856cdaeabf5..463db2e770213ba9636d2537b095a77dece5d8f6 100644 --- a/tensorflow/contrib/py2tf/converters/side_effect_guards_test.py +++ b/tensorflow/contrib/py2tf/converters/side_effect_guards_test.py @@ -20,41 +20,145 @@ 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.contrib.py2tf.pyct import compiler +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 array_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 TestNamer(side_effect_guards.SymbolNamer): +class SideEffectGuardsTest(converter_test_base.TestCase): - def new_symbol(self, name_root, _): - return name_root + def test_side_effect_on_return_only_variable(self): + tf = None -class SideEffectGuardsTest(converter_test_base.TestCase): + def test_fn(a): + tf.assign(a, a + 1) + return a - def test_transform(self): + 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): - state_ops.assign(a, a + 1) + 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, {'state_ops': state_ops}) - node = side_effect_guards.transform(node, TestNamer()) - result = compiler.ast_to_object(node) - setattr(result, 'state_ops', state_ops) + 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 - # TODO(mdan): Configure the namespaces instead of doing these hacks. - ops.identity = array_ops.identity - setattr(result, 'tf', ops) + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) - with self.test_session() as sess: - v = variables.Variable(2) - sess.run(v.initializer) - self.assertEqual(3, sess.run(result.test_fn(v))) + 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__': diff --git a/tensorflow/contrib/py2tf/converters/single_return.py b/tensorflow/contrib/py2tf/converters/single_return.py new file mode 100644 index 0000000000000000000000000000000000000000..1194b98f5ebeffa79a41fc3b32aa79ffd8cc407b --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/single_return.py @@ -0,0 +1,317 @@ +# 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 functions with multiple returns to use just one.""" + +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 + + +# TODO(mdan): Move this logic into transformer_base. +class BodyVisitor(transformer.Base): + """Walks breadth- or depth-first the list-of-nodes bodies of AST nodes.""" + + def __init__(self, context, depth_first=False): + self.depth_first = depth_first + self.changes_made = False + super(BodyVisitor, self).__init__(context) + + def visit_nodelist(self, nodelist): + for node in nodelist: + if isinstance(node, list): + node = self.visit_nodelist(node) + else: + node = self.generic_visit(node) + return nodelist + + def visit_If(self, node): + if self.depth_first: + node = self.generic_visit(node) + node.body = self.visit_nodelist(node.body) + node.orelse = self.visit_nodelist(node.orelse) + if not self.depth_first: + node = self.generic_visit(node) + return node + + def visit_For(self, node): + if self.depth_first: + node = self.generic_visit(node) + node.body = self.visit_nodelist(node.body) + node.orelse = self.visit_nodelist(node.orelse) + if not self.depth_first: + node = self.generic_visit(node) + return node + + def visit_While(self, node): + if self.depth_first: + node = self.generic_visit(node) + node.body = self.visit_nodelist(node.body) + node.orelse = self.visit_nodelist(node.orelse) + if not self.depth_first: + node = self.generic_visit(node) + return node + + def visit_Try(self, node): + if self.depth_first: + node = self.generic_visit(node) + node.body = self.visit_nodelist(node.body) + node.orelse = self.visit_nodelist(node.orelse) + node.finalbody = self.visit_nodelist(node.finalbody) + for i in range(len(node.handlers)): + node.handlers[i].body = self.visit_nodelist(node.handlers[i].body) + if not self.depth_first: + node = self.generic_visit(node) + return node + + def visit_With(self, node): + if self.depth_first: + node = self.generic_visit(node) + node.body = self.visit_nodelist(node.body) + if not self.depth_first: + node = self.generic_visit(node) + return node + + def visit_FunctionDef(self, node): + if self.depth_first: + node = self.generic_visit(node) + node.body = self.visit_nodelist(node.body) + self.generic_visit(node) + if not self.depth_first: + node = self.generic_visit(node) + return node + + +class FoldElse(BodyVisitor): + + def visit_nodelist(self, nodelist): + for i in range(len(nodelist)): + node = nodelist[i] + if isinstance(node, gast.If): + true_branch_returns = isinstance(node.body[-1], gast.Return) + false_branch_returns = len(node.orelse) and isinstance( + node.orelse[-1], gast.Return) + # If the last node in the if body is a return, + # then every line after this if statement effectively + # belongs in the else. + if true_branch_returns and not false_branch_returns: + for j in range(i + 1, len(nodelist)): + nodelist[i].orelse.append(ast_util.copy_clean(nodelist[j])) + if nodelist[i + 1:]: + self.changes_made = True + return nodelist[:i + 1] + elif not true_branch_returns and false_branch_returns: + for j in range(i + 1, len(nodelist)): + nodelist[i].body.append(ast_util.copy_clean(nodelist[j])) + if nodelist[i + 1:]: + self.changes_made = True + return nodelist[:i + 1] + elif true_branch_returns and false_branch_returns: + if nodelist[i + 1:]: + raise ValueError( + 'Unreachable code after conditional where both branches return.' + ) + return nodelist + elif isinstance(node, gast.Return) and nodelist[i + 1:]: + raise ValueError( + 'Cannot have statements after a return in the same basic block') + return nodelist + + +def contains_return(node): + for n in gast.walk(node): + if isinstance(n, gast.Return): + return True + return False + + +class LiftReturn(transformer.Base): + """Move return statements out of If and With blocks.""" + + def __init__(self, context): + self.changes_made = False + self.common_return_name = None + super(LiftReturn, self).__init__(context) + + def visit_If(self, node): + # Depth-first traversal of if statements + node = self.generic_visit(node) + + # We check if both branches return, and if so, lift the return out of the + # conditional. We don't enforce that the true and false branches either + # both return or both do not, because FoldElse might move a return + # into a branch after this transform completes. FoldElse and LiftReturn + # are alternately run until the code reaches a fixed point. + true_branch_returns = isinstance(node.body[-1], gast.Return) + false_branch_returns = len(node.orelse) and isinstance( + node.orelse[-1], gast.Return) + if true_branch_returns and false_branch_returns: + node.body[-1] = templates.replace( + 'a = b', a=self.common_return_name, b=node.body[-1].value)[0] + node.orelse[-1] = templates.replace( + 'a = b', a=self.common_return_name, b=node.orelse[-1].value)[0] + return_node = templates.replace('return a', a=self.common_return_name)[0] + self.changes_made = True + return [node, return_node] + else: + return node + + def visit_With(self, node): + # Depth-first traversal of syntax + node = self.generic_visit(node) + + # If the with statement returns, lift the return + if isinstance(node.body[-1], gast.Return): + node.body[-1] = templates.replace( + 'a = b', a=self.common_return_name, b=node.body[-1].value)[0] + return_node = templates.replace('return a', a=self.common_return_name)[0] + node = self.generic_visit(node) + self.changes_made = True + return [node, return_node] + else: + return node + + def visit_FunctionDef(self, node): + # Ensure we're doing depth-first traversal + last_return_name = self.common_return_name + body_scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + referenced_names = body_scope.referenced + self.common_return_name = self.context.namer.new_symbol( + 'return_', referenced_names) + node = self.generic_visit(node) + self.common_return_name = last_return_name + return node + + +class DetectReturnInUnsupportedControlFlow(gast.NodeVisitor): + """Throws an error if code returns inside loops or try/except.""" + + # First, throw an error if we detect a return statement in a loop. + # TODO(alexbw): we need to learn to handle returns inside a loop, + # but don't currently have the TF constructs to do so (need something + # that looks vaguely like a goto). + + def __init__(self): + self.cant_return = False + super(DetectReturnInUnsupportedControlFlow, self).__init__() + + def visit_While(self, node): + self.cant_return = True + self.generic_visit(node) + self.cant_return = False + + def visit_For(self, node): + self.cant_return = True + self.generic_visit(node) + self.cant_return = False + + def visit_Try(self, node): + self.cant_return = True + self.generic_visit(node) + self.cant_return = False + + def visit_Return(self, node): + if self.cant_return: + raise ValueError( + 'Pyflow currently does not support `return` statements in loops. ' + 'Try assigning to a variable in the while loop, and returning ' + 'outside of the loop') + + +class DetectReturnInConditional(gast.NodeVisitor): + """Assert that no return statements are present in conditionals.""" + + def __init__(self): + self.cant_return = False + super(DetectReturnInConditional, self).__init__() + + def visit_If(self, node): + self.cant_return = True + self.generic_visit(node) + self.cant_return = False + + def visit_Return(self, node): + if self.cant_return: + raise ValueError( + 'After transforms, a conditional contained a `return `statement, ' + 'which is not allowed. This is a bug, and should not happen.') + + +class DetectReturnInFunctionDef(gast.NodeVisitor): + + def visit_FunctionDef(self, node): + self.generic_visit(node) + if not contains_return(node): + raise ValueError( + 'Each function definition should contain at least one return.') + + +def transform(node, context): + """Ensure a function has only a single return. + + This transforms an AST node with multiple returns successively into containing + only a single return node. + There are a few restrictions on what we can handle: + - An AST being transformed must contain at least one return. + - No returns allowed in loops. We have to know the type of the return value, + and we currently don't have either a type inference system to discover it, + nor do we have a mechanism for late type binding in TensorFlow. + - After all transformations are finished, a Return node is not allowed inside + control flow. If we were unable to move a return outside of control flow, + this is an error. + + Args: + node: an AST node to transform + context: a context object + + Returns: + new_node: an AST with a single return value + + Raises: + ValueError: if the AST is structured so that we can't perform the + transform. + """ + # Make sure that the function has at least one return statement + # TODO(alexbw): turning off this assertion for now -- + # we need to not require this in e.g. class constructors. + # DetectReturnInFunctionDef().visit(node) + + # Make sure there's no returns in unsupported locations (loops, try/except) + DetectReturnInUnsupportedControlFlow().visit(node) + + while True: + + # Try to lift all returns out of if statements and with blocks + lr = LiftReturn(context) + node = lr.visit(node) + changes_made = lr.changes_made + fe = FoldElse(context) + node = fe.visit(node) + changes_made = changes_made or fe.changes_made + + if not changes_made: + break + + # Make sure we've scrubbed all returns from conditionals + DetectReturnInConditional().visit(node) + + return node diff --git a/tensorflow/contrib/py2tf/converters/single_return_test.py b/tensorflow/contrib/py2tf/converters/single_return_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2ea7a9d6d3e25c8dafd8f211994c8fe99bd0e781 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/single_return_test.py @@ -0,0 +1,189 @@ +# 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 single_return 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 single_return +from tensorflow.python.framework.ops import name_scope +from tensorflow.python.platform import test + + +class SingleReturnTest(converter_test_base.TestCase): + + def compiled_fn(self, test_fn, *args): + node = self.parse_and_analyze(test_fn, {}) + node = single_return.transform(node, self.ctx) + module = self.compiled(node, *args) + return module + + def test_noop(self): + # Noop + def test_fn(x): + return x + + with self.compiled_fn(test_fn) as result: + self.assertEqual(test_fn(2.0), result.test_fn(2.0)) + + def test_return_expression(self): + # ANF + def test_fn(x): + return x * x + + with self.compiled_fn(test_fn) as result: + x = 2 + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_merge(self): + # Simple merge + def test_fn(x): + if x > 0: + return x + else: + return x * x + + with self.compiled_fn(test_fn) as result: + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_orphan_branch(self): + + def test_fn(x): + if x > 0: + return x + + with self.assertRaises(ValueError): + self.compiled_fn(test_fn) + + def test_lift_body_into_false_branch(self): + + def test_fn(x): + if x > 0: + return x + return x * x + + with self.compiled_fn(test_fn) as result: + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_lift_body_into_true_branch(self): + + def test_fn(x): + if x < 0: + x *= x + else: + # TODO(alexbw): linter bug here that requires us suppress this warning. + return x # pylint: disable=undefined-loop-variable + return x + + with self.compiled_fn(test_fn) as result: + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_nested_if(self): + + def test_fn(x): + if x > 0: + if x < 5: + return x + else: + return x * x + else: + return x * x * x + + with self.compiled_fn(test_fn) as result: + for x in [-2, 2, 5]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_context_manager(self): + + def test_fn(x): + + with name_scope(''): + return x * x + + with self.compiled_fn(test_fn) as result: + result.name_scope = name_scope + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_context_manager_in_conditional(self): + + def test_fn(x): + if x > 0: + with name_scope(''): + return x * x + else: + return x + + with self.compiled_fn(test_fn, name_scope) as result: + result.name_scope = name_scope + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def text_conditional_in_context_manager(self): + + def test_fn(x): + with name_scope(''): + if x > 0: + return x * x + else: + return x + + with self.compiled_fn(test_fn) as result: + result.name_scope = name_scope + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_no_return(self): + + def test_fn(x): + x *= x + + with self.compiled_fn(test_fn) as result: + self.assertEqual(test_fn(2), result.test_fn(2)) + + def test_nested_functiondefs(self): + + def test_fn(x): + + def inner_fn(y): + if y > 0: + return y * y + else: + return y + + return inner_fn(x) + + with self.compiled_fn(test_fn) as result: + for x in [-2, 2]: + self.assertEqual(test_fn(x), result.test_fn(x)) + + def test_loop(self): + + def test_fn(x): + for _ in range(10): + return x + return x + + with self.assertRaises(ValueError): + self.compiled_fn(test_fn) + + +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/api.py b/tensorflow/contrib/py2tf/impl/api.py similarity index 64% rename from tensorflow/contrib/py2tf/api.py rename to tensorflow/contrib/py2tf/impl/api.py index ca1f4e2645ee20fd78c0d837885823d2e199537a..883b304089024363f41cabde2cb74c49f01ae836 100644 --- a/tensorflow/contrib/py2tf/api.py +++ b/tensorflow/contrib/py2tf/impl/api.py @@ -23,10 +23,13 @@ from functools import wraps import gast import six -from tensorflow.contrib.py2tf import config -from tensorflow.contrib.py2tf import conversion +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 inspect_utils from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.utils import builtins +from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import tf_inspect # TODO(mdan): Properly document the type hints. @@ -83,15 +86,16 @@ def convert_inline(f, *args, **kwargs): return convert(arg_value_hints)(f)(*args, **kwargs) -def convert(recursive=False, arg_types=None): +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. + 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: @@ -108,27 +112,7 @@ def convert(recursive=False, arg_types=None): @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, - arg_values=arg_values, - arg_types=arg_types, - partial_types=partial_types) - return wrapped(*args, **kwargs) + return converted_call(f, recursive, verbose, arg_types, *args, **kwargs) # Sometimes the decorator is just desugared, making it impossible to detect. # This attribute makes detection easier. @@ -138,8 +122,81 @@ def convert(recursive=False, arg_types=None): return decorator +def converted_call(f, recursive, verbose, arg_types, *args, **kwargs): + """Compiles a function call inline.""" + # TODO(mdan): This needs cleanup. + # In particular, we may want to avoid renaming functions altogether. + + if conversion.is_whitelisted_for_graph(f): + return f(*args, **kwargs) + + unknown_arg_value = object() # Sentinel for arguments of unknown value + + if tf_inspect.isbuiltin(f): + return builtins.dynamic_builtin(f, *args, **kwargs) + + if tf_inspect.isfunction(f) or tf_inspect.ismethod(f): + # Regular functions + target_entity = f + arg_map_target = f + effective_args = args + f_class = inspect_utils.getmethodclass(f) + + if f_class is not None: + partial_types = (f_class,) + else: + partial_types = () + + elif tf_inspect.isclass(f): + # Constructors + target_entity = f + arg_map_target = f.__init__ + effective_args = (unknown_arg_value,) + args + partial_types = () + + elif hasattr(f, '__call__') and hasattr(f, '__class__'): + # Callable objects + target_entity = f.__call__ + arg_map_target = f.__call__ + effective_args = (f,) + args + partial_types = (f.__class__,) + + else: + NotImplementedError('unknown callable type "%s"' % type(f)) + + arg_values = tf_inspect.getcallargs(arg_map_target, *args, **kwargs) + for name, arg in arg_values.items(): + if arg is unknown_arg_value: + continue + 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) + + # When called from within a decorator, this is the only indication that + # the function is a method - it appears that the decorator is applied + # before the method is bound. + if not partial_types: + if 'self' in arg_values: + if tf_inspect.isclass(arg_values['self'].__class__): + partial_types = (arg_values['self'].__class__,) + elif 'cls' in arg_values: + if tf_inspect.isclass(arg_values['cls']): + partial_types = (arg_values['cls'],) + + converted_f = to_graph( + target_entity, + recursive=recursive, + verbose=verbose, + arg_values=arg_values, + arg_types=arg_types, + partial_types=partial_types) + return converted_f(*effective_args, **kwargs) + + def to_graph(e, recursive=True, + verbose=False, arg_values=None, arg_types=None, partial_types=None): @@ -155,6 +212,7 @@ def to_graph(e, 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 @@ -170,22 +228,26 @@ def to_graph(e, conversion_map = conversion.ConversionMap( recursive=recursive, nocompile_decorators=(convert, graph_ready, convert_inline), - partial_types=partial_types) + partial_types=partial_types, + api_module=tf_inspect.getmodule(to_graph)) _, 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)) + module.body.extend(parser.parse_str(import_line).body) for dep in conversion_map.dependency_cache.values(): module.body.append(dep) - compiled_node = compiler.ast_to_object(module) + 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_node.__dict__.update(inspect_utils.getnamespace(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 @@ -213,7 +275,8 @@ def to_code(e, conversion_map = conversion.ConversionMap( recursive=recursive, nocompile_decorators=(convert, graph_ready, convert_inline), - partial_types=partial_types) + partial_types=partial_types, + api_module=tf_inspect.getmodule(to_graph)) conversion.entity_to_graph(e, conversion_map, arg_values, arg_types) imports = '\n'.join(config.COMPILED_IMPORT_STATEMENTS) diff --git a/tensorflow/contrib/py2tf/api_test.py b/tensorflow/contrib/py2tf/impl/api_test.py similarity index 86% rename from tensorflow/contrib/py2tf/api_test.py rename to tensorflow/contrib/py2tf/impl/api_test.py index 2384447708d7e0ab5dbfbeb592a47353f1909f50..13f8e66018920a5b13f8bd3f00c67d3bbdd519aa 100644 --- a/tensorflow/contrib/py2tf/api_test.py +++ b/tensorflow/contrib/py2tf/impl/api_test.py @@ -18,21 +18,26 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.py2tf import api -from tensorflow.contrib.py2tf import config +from tensorflow.contrib.py2tf import utils +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 +tf = utils.fake_tf() + + 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 __future__ import print_function', + 'from tensorflow.contrib.py2tf import utils as ' + 'py2tf_utils', + 'tf = py2tf_utils.fake_tf()' + ) def test_decorator_recurses(self): @@ -45,7 +50,7 @@ class ApiTest(test.TestCase): @api.convert(recursive=True) def test_method(self, x, s, a): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= self.called_member(a) return x @@ -61,11 +66,11 @@ class ApiTest(test.TestCase): class TestClass(object): def called_member(self, a): - return math_ops.negative(a) + return tf.negative(a) @api.convert(recursive=False) def test_method(self, x, s, a): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= self.called_member(a) return x @@ -82,11 +87,11 @@ class ApiTest(test.TestCase): @api.graph_ready def called_member(self, a): - return math_ops.negative(a) + return tf.negative(a) @api.convert(recursive=True) def test_method(self, x, s, a): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= self.called_member(a) return x @@ -109,7 +114,7 @@ class ApiTest(test.TestCase): @api.convert(recursive=True) def test_method(self, x, s, a): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= self.called_member(a) return x @@ -131,7 +136,7 @@ class ApiTest(test.TestCase): @api.convert(recursive=True) def test_method(self, x, s, a): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= api.convert_inline(self.called_member, a) return x @@ -147,11 +152,11 @@ class ApiTest(test.TestCase): class TestClass(object): def called_member(self, a): - return math_ops.negative(a) + return tf.negative(a) @api.convert(recursive=True) def test_method(self, x, s, a): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= api.graph_ready(self.called_member(a)) return x @@ -164,7 +169,7 @@ class ApiTest(test.TestCase): def test_to_graph_basic(self): def test_fn(x, s): - while math_ops.reduce_sum(x) > s: + while tf.reduce_sum(x) > s: x //= 2 return x @@ -176,14 +181,14 @@ class ApiTest(test.TestCase): def test_to_code_basic(self): def test_fn(x, s): - while math_ops.reduce_sum(x) > s: + while tf.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, 'tf\\.while_loop') + self.assertRegexpMatches(compiled_code, 'py2tf_utils\\.run_while') self.assertIsNotNone(parser.parse_str(compiled_code)) diff --git a/tensorflow/contrib/py2tf/config.py b/tensorflow/contrib/py2tf/impl/config.py similarity index 70% rename from tensorflow/contrib/py2tf/config.py rename to tensorflow/contrib/py2tf/impl/config.py index 8c502a7a9e546dd9b9b40d7cf6d3c9821038afb3..bdbc6663dd65ed66c55ad2d2e52428084bbea219 100644 --- a/tensorflow/contrib/py2tf/config.py +++ b/tensorflow/contrib/py2tf/impl/config.py @@ -18,6 +18,9 @@ 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, @@ -27,12 +30,21 @@ PYTHON_LITERALS = { DEFAULT_UNCOMPILED_MODULES = set(( ('tensorflow',), + (utils.__name__,), + + # All of tensorflow's subpackages. Unlike the root tf module, they don't + # have well-known names. Not refering to the module directly to avoid + # circular imports. + (utils.__name__[:-len('.contrib.py2tf.utils')],), )) NO_SIDE_EFFECT_CONSTRUCTORS = set(('tensorflow',)) # TODO(mdan): Also allow controlling the generated names (for testability). COMPILED_IMPORT_STATEMENTS = ( - 'from contextlib import contextmanager', + 'from __future__ import print_function', 'import tensorflow as tf', -) + 'from tensorflow.contrib.py2tf.impl import api as ' + 'py2tf_api', + 'from tensorflow.contrib.py2tf import utils as ' + 'py2tf_utils') diff --git a/tensorflow/contrib/py2tf/conversion.py b/tensorflow/contrib/py2tf/impl/conversion.py similarity index 61% rename from tensorflow/contrib/py2tf/conversion.py rename to tensorflow/contrib/py2tf/impl/conversion.py index 67ca52d1944c8155d16e7e17411c1afb210d606b..37b24ab55fdd1b03e12e9afe06530e3c26218b61 100644 --- a/tensorflow/contrib/py2tf/conversion.py +++ b/tensorflow/contrib/py2tf/impl/conversion.py @@ -19,26 +19,32 @@ from __future__ import division from __future__ import print_function import gast -import six -from tensorflow.contrib.py2tf import config -from tensorflow.contrib.py2tf import naming +from tensorflow.contrib.py2tf import utils from tensorflow.contrib.py2tf.converters import asserts -from tensorflow.contrib.py2tf.converters import break_canonicalization +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_canonicalization +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_canonicalization +from tensorflow.contrib.py2tf.converters import for_loops +from tensorflow.contrib.py2tf.converters import ifexp +from tensorflow.contrib.py2tf.converters import lists from tensorflow.contrib.py2tf.converters import logical_expressions -from tensorflow.contrib.py2tf.converters import print_functions +from tensorflow.contrib.py2tf.converters import name_scopes from tensorflow.contrib.py2tf.converters import side_effect_guards +from tensorflow.contrib.py2tf.converters import single_return +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 inspect_utils from tensorflow.contrib.py2tf.pyct import parser -from tensorflow.contrib.py2tf.pyct.static_analysis import access +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.contrib.py2tf.utils import type_hints from tensorflow.python.util import tf_inspect @@ -46,7 +52,9 @@ from tensorflow.python.util import tf_inspect class ConversionMap(object): - """ConversionMaps keep track of converting function hierarchies. + """ConversionMap keeps track of converting function hierarchies. + + This object is mutable, and is updated as functions are converted. Attributes: recursive: Whether to recusrively convert any functions that the decorator @@ -55,18 +63,26 @@ class ConversionMap(object): off. dependency_cache: dict[object]: ast; maps original entities to their converted AST + additional_imports: set(object); additional entities which for any reason + cannot be attached after loading and need to be explicitly imported + in the generated code name_map: dict[string]: string; maps original entities to the name of their converted counterparts + api_module: A reference to the api module. The reference needs to be passed + to avoid circular dependencies. """ # TODO(mdan): Rename to ConversionContext, and pull in additional flags. - def __init__(self, recursive, nocompile_decorators, partial_types): + def __init__(self, recursive, nocompile_decorators, partial_types, + api_module): self.recursive = recursive self.nocompile_decorators = nocompile_decorators self.partial_types = partial_types if partial_types else () self.dependency_cache = {} + self.additional_imports = set() self.name_map = {} + self.api_module = api_module def new_namer(self, namespace): return naming.Namer(namespace, self.recursive, self.name_map, @@ -87,6 +103,24 @@ class ConversionMap(object): self.dependency_cache[original_entity] = converted_ast +def is_whitelisted_for_graph(o): + """Check whether an entity is whitelisted for use in graph mode. + + Examples of whitelisted entities include all members of the tensorflow + package. + + Args: + o: A Python entity. + Returns: + Boolean + """ + m = tf_inspect.getmodule(o) + for prefix, in config.DEFAULT_UNCOMPILED_MODULES: + if m.__name__.startswith(prefix): + return True + return False + + def entity_to_graph(o, conversion_map, arg_values, arg_types): """Compile a Python entity into equivalent TensorFlow. @@ -126,14 +160,20 @@ def entity_to_graph(o, conversion_map, arg_values, arg_types): 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, {}, {}) + while True: + candidate = None + for obj in conversion_map.name_map.keys(): + if obj not in conversion_map.dependency_cache: + candidate = obj + break + if candidate is None: + break + if (hasattr(candidate, 'im_class') and + getattr(candidate, '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(candidate, conversion_map, {}, {}) return node, new_name @@ -141,11 +181,12 @@ def entity_to_graph(o, conversion_map, arg_values, arg_types): 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) + method_filter = lambda m: tf_inspect.isfunction(m) or tf_inspect.ismethod(m) + members = tf_inspect.getmembers(c, predicate=method_filter) if not members: - raise ValueError('Cannot convert %s: it has no member methods.') + raise ValueError('Cannot convert %s: it has no member methods.' % c) - class_globals = None + class_namespace = None for _, m in members: node, _ = function_to_graph( m, @@ -154,37 +195,49 @@ def class_to_graph(c, conversion_map): 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) + if class_namespace is None: + class_namespace = inspect_utils.getnamespace(m) converted_members[m] = node - namer = conversion_map.new_namer(class_globals) + namer = conversion_map.new_namer(class_namespace) class_name = namer.compiled_class_name(c.__name__, c) node = gast.ClassDef( class_name, bases=[], keywords=[], - body=converted_members.values(), + body=list(converted_members.values()), decorator_list=[]) return node, class_name +def _add_self_references(namespace, api_module): + """Self refs are only required for analysis and are not used directly.""" + # 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.') + + # We also make reference to the api module for dynamic conversion, but + # to avoid circular references we don't import it here. + if 'py2tf_api' not in namespace: + namespace['py2tf_api'] = api_module + elif namespace['py2tf_api'] != api_module: + raise ValueError( + 'The module name "py2tf_api" is reserved and may not be used.') + + 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 + namespace = inspect_utils.getnamespace(f) + _add_self_references(namespace, conversion_map.api_module) namer = conversion_map.new_namer(namespace) + ctx = context.EntityContext( namer=namer, source_code=source, @@ -192,8 +245,10 @@ def function_to_graph(f, conversion_map, arg_values, arg_types, namespace=namespace, arg_values=arg_values, arg_types=arg_types, - recursive=conversion_map.recursive) - node = node_to_graph(node, ctx, conversion_map.nocompile_decorators) + owner_type=owner_type, + recursive=conversion_map.recursive, + type_annotation_func=type_hints.set_element_type) + node, deps = 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) @@ -204,11 +259,15 @@ def function_to_graph(f, conversion_map, arg_values, arg_types, node.name = new_name conversion_map.update_name_map(namer) + # TODO(mdan): Use this at compilation. + conversion_map.additional_imports.update(deps) + return node, new_name def _static_analysis_pass(node, ctx): - node = access.resolve(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 @@ -233,10 +292,7 @@ def node_to_graph(node, ctx, nocompile_decorators): # TODO(mdan): Factor out common elements. # These include: - # * keeping track of symbols that have been created - # * marking nodes (e.g. py_func wrappers) to suppress further processing # * code move between blocks - # * insertion of new global references # * visiting blocks in transformers # Certain steps, especially canonicalization, insert new symbols into the @@ -244,29 +300,43 @@ def node_to_graph(node, ctx, nocompile_decorators): # to re-run the analysis. node = _static_analysis_pass(node, ctx) - node = decorators.transform(node, nocompile_decorators) - node = break_canonicalization.transform(node, ctx.namer) + + # TODO(mdan): Clean this up. + # Some intermediate analyses are not required, and some comments got orphaned. + + # 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 = ifexp.transform(node, ctx) + node, deps = 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_canonicalization.transform(node, ctx.namer) + node = continue_statements.transform(node, ctx) ctx.namespace['len'] = len node = _static_analysis_pass(node, ctx) - node = for_canonicalization.transform(node, ctx.namer) - # for_canonicalization may insert new global references. - node = builtin_functions.transform(node) - # builtin_functions may insert new global references. - ctx.namespace['print'] = print + node = single_return.transform(node, ctx) + + node = _static_analysis_pass(node, ctx) + node = lists.transform(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 = print_functions.transform(node) node = call_trees.transform(node, ctx, config.DEFAULT_UNCOMPILED_MODULES, nocompile_decorators) - node = control_flow.transform(node, ctx.namer) - node = logical_expressions.transform(node) - node = side_effect_guards.transform(node, ctx.namer) + node = control_flow.transform(node, ctx) - return node + # control_flow may create new symbols and change scopes. + node = _static_analysis_pass(node, ctx) + node = logical_expressions.transform(node, ctx) + node = side_effect_guards.transform(node, ctx) + node = name_scopes.transform(node, ctx) + + return node, deps diff --git a/tensorflow/contrib/py2tf/conversion_test.py b/tensorflow/contrib/py2tf/impl/conversion_test.py similarity index 68% rename from tensorflow/contrib/py2tf/conversion_test.py rename to tensorflow/contrib/py2tf/impl/conversion_test.py index 26f915f4f46e54c9648ae6b35415c4e2639af774..9ff256aace7a0e7ac5e7ac07e580b8bed7d8df6f 100644 --- a/tensorflow/contrib/py2tf/conversion_test.py +++ b/tensorflow/contrib/py2tf/impl/conversion_test.py @@ -20,15 +20,26 @@ from __future__ import print_function import gast -from tensorflow.contrib.py2tf import conversion +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.impl import conversion +from tensorflow.python.framework import constant_op from tensorflow.python.platform import test class ConversionTest(test.TestCase): + def test_is_whitelisted_for_graph(self): + + def test_fn(): + return constant_op.constant(1) + + self.assertFalse(conversion.is_whitelisted_for_graph(test_fn)) + self.assertTrue(conversion.is_whitelisted_for_graph(utils)) + self.assertTrue(conversion.is_whitelisted_for_graph(constant_op.constant)) + def test_entity_to_graph_unsupported_types(self): with self.assertRaises(ValueError): - conversion_map = conversion.ConversionMap(True, (), ()) + conversion_map = conversion.ConversionMap(True, (), (), None) conversion.entity_to_graph('dummy', conversion_map, None, None) def test_entity_to_graph_callable(self): @@ -36,7 +47,7 @@ class ConversionTest(test.TestCase): def f(a): return a - conversion_map = conversion.ConversionMap(True, (), ()) + conversion_map = conversion.ConversionMap(True, (), (), None) 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) @@ -49,14 +60,17 @@ class ConversionTest(test.TestCase): def f(a): return g(a) - conversion_map = conversion.ConversionMap(True, (), ()) + conversion_map = conversion.ConversionMap(True, (), (), None) 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) + # need the extra .body[0] in order to step past the with tf.name_scope('f') + # that is added automatically self.assertEqual( - 'tf__g', conversion_map.dependency_cache[f].body[0].value.func.id) + 'tf__g', + conversion_map.dependency_cache[f].body[0].body[0].value.func.id) self.assertEqual('tf__g', conversion_map.dependency_cache[g].name) diff --git a/tensorflow/contrib/py2tf/naming.py b/tensorflow/contrib/py2tf/impl/naming.py similarity index 84% rename from tensorflow/contrib/py2tf/naming.py rename to tensorflow/contrib/py2tf/impl/naming.py index 5c7e4c5f95ac5a23061e82bd74fb992bab55af11..51326091de13715c32d0a79279f1d3274e48ad10 100644 --- a/tensorflow/contrib/py2tf/naming.py +++ b/tensorflow/contrib/py2tf/impl/naming.py @@ -18,6 +18,8 @@ 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. @@ -103,11 +105,26 @@ class Namer(object): 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 - n = 0 - while (new_name in self.global_namespace - or new_name in reserved_locals - or new_name in self.generated_names): + + 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) diff --git a/tensorflow/contrib/py2tf/naming_test.py b/tensorflow/contrib/py2tf/impl/naming_test.py similarity index 98% rename from tensorflow/contrib/py2tf/naming_test.py rename to tensorflow/contrib/py2tf/impl/naming_test.py index 5cf0a3da2c970d36051601f3ef19162ba241eab0..beb4e54937bbb91b19157c9b9e3c528353206c62 100644 --- a/tensorflow/contrib/py2tf/naming_test.py +++ b/tensorflow/contrib/py2tf/impl/naming_test.py @@ -18,7 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.py2tf import naming +from tensorflow.contrib.py2tf.impl import naming from tensorflow.python.platform import test diff --git a/tensorflow/contrib/py2tf/pyct/BUILD b/tensorflow/contrib/py2tf/pyct/BUILD index 1b2408ba0ea37e596b70d0f824bc04167fc9f412..edec5f7712d08247437c9e95d743e59dafffcd7b 100644 --- a/tensorflow/contrib/py2tf/pyct/BUILD +++ b/tensorflow/contrib/py2tf/pyct/BUILD @@ -1,5 +1,7 @@ licenses(["notice"]) # Apache 2.0 +exports_files(["LICENSE"]) + load("//tensorflow:tensorflow.bzl", "py_test") filegroup( @@ -19,10 +21,13 @@ py_library( srcs = [ "__init__.py", "anno.py", + "ast_util.py", "compiler.py", "context.py", + "inspect_utils.py", "parser.py", "pretty_printer.py", + "qual_names.py", "templates.py", "transformer.py", ], @@ -46,6 +51,17 @@ py_test( ], ) +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"], @@ -57,6 +73,17 @@ py_test( ], ) +py_test( + name = "inspect_utils_test", + srcs = ["inspect_utils_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) + py_test( name = "parser_test", srcs = ["parser_test.py"], @@ -77,6 +104,16 @@ py_test( ], ) +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"], diff --git a/tensorflow/contrib/py2tf/pyct/anno.py b/tensorflow/contrib/py2tf/pyct/anno.py index 889e4ba4ffaed887faffb8736e4a59502da99e81..cc4a7edf02ed7556c9a552d8730e4c7875038c83 100644 --- a/tensorflow/contrib/py2tf/pyct/anno.py +++ b/tensorflow/contrib/py2tf/pyct/anno.py @@ -21,6 +21,30 @@ 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] @@ -38,3 +62,16 @@ def setanno(node, key, value, field_name='___pyct_anno'): # 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) + + +def copyanno(from_node, to_node, key, field_name='___pyct_anno'): + if hasanno(from_node, key, field_name): + setanno(to_node, key, getanno(from_node, key, field_name), field_name) diff --git a/tensorflow/contrib/py2tf/pyct/anno_test.py b/tensorflow/contrib/py2tf/pyct/anno_test.py index 19e3b4576210c3715620fc7002c91c5130b46ed0..6c29918fdfaaa0224f20a2c3cb2ea8088f3eb52b 100644 --- a/tensorflow/contrib/py2tf/pyct/anno_test.py +++ b/tensorflow/contrib/py2tf/pyct/anno_test.py @@ -24,6 +24,9 @@ from tensorflow.contrib.py2tf.pyct import anno from tensorflow.python.platform import test +# TODO(mdan): Consider strong types instead of primitives. + + class AnnoTest(test.TestCase): def test_basic(self): @@ -37,6 +40,22 @@ class AnnoTest(test.TestCase): 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') + + def test_copyanno(self): + node_1 = ast.Name() + anno.setanno(node_1, 'foo', 3) + + node_2 = ast.Name() + anno.copyanno(node_1, node_2, 'foo') + anno.copyanno(node_1, node_2, 'bar') + + self.assertTrue(anno.hasanno(node_2, 'foo')) + self.assertFalse(anno.hasanno(node_2, 'bar')) + 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..a871ccad6fc7ea1487e41fd6da3ce6120bdcbcbd --- /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(qual_names.QN('b'), attr='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 index fc71469d1eaeb92352e3b50cb743621d7e5eb1d5..507dbc7ed3de9c0b8874164e97a3d1d149e42423 100644 --- a/tensorflow/contrib/py2tf/pyct/compiler.py +++ b/tensorflow/contrib/py2tf/pyct/compiler.py @@ -22,6 +22,7 @@ from __future__ import division from __future__ import print_function # TODO(mdan): Use six for compatibility here. +import atexit import imp import os import tempfile @@ -38,10 +39,14 @@ def ast_to_source(node, indentation): astor.string_repr.pretty_string) generator.visit(node) generator.result.append('\n') - return astor.source_repr.pretty_source(generator.result).lstrip() + # In some versions of Python, literals may appear as actual values. This + # ensures everything is string. + code = map(str, generator.result) + return astor.source_repr.pretty_source(code).lstrip() -def ast_to_object(node, indentation=' ', source_prefix=None): +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 @@ -51,6 +56,8 @@ def ast_to_object(node, indentation=' ', source_prefix=None): 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. @@ -63,4 +70,6 @@ def ast_to_object(node, indentation=' ', source_prefix=None): f.write(source_prefix) f.write('\n') f.write(source) - return imp.load_source(module_name, f.name) + 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 index e0cde43566310b99bac5035285154fde906fa127..243f4c81538f5853a01ff444f2ff16ccf7cd5d62 100644 --- a/tensorflow/contrib/py2tf/pyct/compiler_test.py +++ b/tensorflow/contrib/py2tf/pyct/compiler_test.py @@ -23,11 +23,28 @@ import textwrap import gast from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.contrib.py2tf.pyct import parser from tensorflow.python.platform import test +from tensorflow.python.util import tf_inspect class CompilerTest(test.TestCase): + def test_parser_compile_idempotent(self): + + def test_fn(x): + a = True + b = '' + if a: + b = x + 1 + return b + + self.assertEqual( + textwrap.dedent(tf_inspect.getsource(test_fn)), + tf_inspect.getsource( + compiler.ast_to_object( + parser.parse_entity(test_fn)[0].body[0])[0].test_fn)) + def test_ast_to_source(self): node = gast.If( test=gast.Num(1), @@ -41,6 +58,7 @@ class CompilerTest(test.TestCase): targets=[gast.Name('a', gast.Store(), None)], value=gast.Str('c')) ]) + self.assertEqual( textwrap.dedent(""" if 1: @@ -70,15 +88,19 @@ class CompilerTest(test.TestCase): decorator_list=[], returns=None) - mod = compiler.ast_to_object(node) + module, source = compiler.ast_to_object(node) - self.assertEqual(2, mod.f(1)) - with open(mod.__file__, 'r') as temp_output: + 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(""" - def f(a): - return a + 1 - """).strip(), + textwrap.dedent(expected_source).strip(), temp_output.read().strip()) diff --git a/tensorflow/contrib/py2tf/pyct/context.py b/tensorflow/contrib/py2tf/pyct/context.py index fef74ebefa290369c7310af6d7e4faeef44d9aee..b34015cfd2888f0dbeb6492b9e7335d561bf4763 100644 --- a/tensorflow/contrib/py2tf/pyct/context.py +++ b/tensorflow/contrib/py2tf/pyct/context.py @@ -22,6 +22,8 @@ from __future__ import print_function class EntityContext(object): """Contains information about an entity, like source code. + In general, objects of this class should be considered immutable. + Attributes: namer: Namer that matches the contract of all converters. source_code: The entity's source code. @@ -30,14 +32,18 @@ class EntityContext(object): (excluding parameters). arg_values: Dict[str->*], containing parameter values, if known. arg_types: Dict[str->*], containing parameter types, if known. + owner_type: The surrounding class type of the function, if present. """ + # TODO(mdan): Remove the default and update tests. def __init__(self, namer, source_code, source_file, namespace, arg_values, - arg_types, recursive): + arg_types, owner_type, recursive, type_annotation_func=None): 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.owner_type = owner_type self.recursive = recursive + self.type_annotation_func = type_annotation_func diff --git a/tensorflow/contrib/py2tf/pyct/inspect_utils.py b/tensorflow/contrib/py2tf/pyct/inspect_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d19c6ed75e0f0651781d6e1ed80f7be11fb8a5a4 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/inspect_utils.py @@ -0,0 +1,119 @@ +# 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 entity inspection utilities. + +This module contains whatever inspect doesn't offer out of the box. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import itertools + +import six + +from tensorflow.python.util import tf_inspect + + +def getnamespace(f): + """Returns the complete namespace of a function. + + Namespace is defined here as the mapping of all non-local variables to values. + This includes the globals and the closure variables. Note that this captures + the entire globals collection of the function, and may contain extra symbols + that it does not actually use. + + Args: + f: User defined function. + Returns: + A dict mapping symbol names to values. + """ + namespace = dict(six.get_function_globals(f)) + closure = six.get_function_closure(f) + freevars = six.get_function_code(f).co_freevars + if freevars and closure: + for name, cell in zip(freevars, closure): + namespace[name] = cell.cell_contents + return namespace + + +def getmethodclass(m): + """Resolves a function's owner, e.g. a method's class. + + Note that this returns the object that the function was retrieved from, not + necessarily the class where it was defined. + + This function relies on Python stack frame support in the interpreter, and + has the same limitations that inspect.currentframe. + + Limitations. This function will only work correctly if the owned class is + visible in the caller's global or local variables. + + Args: + m: A user defined function + + Returns: + The class that this function was retrieved from, or None if the function + is not an object or class method, or the class that owns the object or + method is not visible to m. + + Raises: + ValueError: if the class could not be resolved for any unexpected reason. + """ + + # Instance method and class methods: should be bound to a non-null "self". + # If self is a class, then it's a class method. + if hasattr(m, '__self__'): + if m.__self__: + if tf_inspect.isclass(m.__self__): + return m.__self__ + return type(m.__self__) + + # Class, static and unbound methods: search all defined classes in any + # namespace. This is inefficient but more robust method. + owners = [] + caller_frame = tf_inspect.currentframe().f_back + try: + # TODO(mdan): This doesn't consider cell variables. + # TODO(mdan): This won't work if the owner is hidden inside a container. + # Cell variables may be pulled using co_freevars and the closure. + for v in itertools.chain(caller_frame.f_locals.values(), + caller_frame.f_globals.values()): + if hasattr(v, m.__name__): + candidate = getattr(v, m.__name__) + # Py2 methods may be bound or unbound, extract im_func to get the + # underlying function. + if hasattr(candidate, 'im_func'): + candidate = candidate.im_func + if hasattr(m, 'im_func'): + m = m.im_func + if candidate is m: + owners.append(v) + finally: + del caller_frame + + if owners: + if len(owners) == 1: + return owners[0] + + # If multiple owners are found, and are not subclasses, raise an error. + owner_types = tuple(o if tf_inspect.isclass(o) else type(o) for o in owners) + for o in owner_types: + if tf_inspect.isclass(o) and issubclass(o, tuple(owner_types)): + return o + raise ValueError('Found too many owners of %s: %s' % (m, owners)) + + return None diff --git a/tensorflow/contrib/py2tf/pyct/inspect_utils_test.py b/tensorflow/contrib/py2tf/pyct/inspect_utils_test.py new file mode 100644 index 0000000000000000000000000000000000000000..5528ac851f74bd7b7dacdbe7b930945afa8c9783 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/inspect_utils_test.py @@ -0,0 +1,230 @@ +# 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 unspect_utils module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from functools import wraps + +import six + +from tensorflow.contrib.py2tf.pyct import inspect_utils +from tensorflow.python.platform import test + + +def decorator(f): + return f + + +def function_decorator(): + def dec(f): + return f + return dec + + +def wrapping_decorator(): + def dec(f): + def replacement(*_): + return None + + @wraps(f) + def wrapper(*args, **kwargs): + return replacement(*args, **kwargs) + return wrapper + return dec + + +class TestClass(object): + + def member_function(self): + pass + + @decorator + def decorated_member(self): + pass + + @function_decorator() + def fn_decorated_member(self): + pass + + @wrapping_decorator() + def wrap_decorated_member(self): + pass + + @staticmethod + def static_method(): + pass + + @classmethod + def class_method(cls): + pass + + +def free_function(): + pass + + +def factory(): + return free_function + + +def free_factory(): + def local_function(): + pass + return local_function + + +class InspectUtilsTest(test.TestCase): + + def test_getnamespace_globals(self): + ns = inspect_utils.getnamespace(factory) + self.assertEqual(ns['free_function'], free_function) + + def test_getnamespace_hermetic(self): + + # Intentionally hiding the global function to make sure we don't overwrite + # it in the global namespace. + free_function = object() # pylint:disable=redefined-outer-name + + def test_fn(): + return free_function + + ns = inspect_utils.getnamespace(test_fn) + globs = six.get_function_globals(test_fn) + self.assertTrue(ns['free_function'] is free_function) + self.assertFalse(globs['free_function'] is free_function) + + def test_getnamespace_locals(self): + + def called_fn(): + return 0 + + closed_over_list = [] + closed_over_primitive = 1 + + def local_fn(): + closed_over_list.append(1) + local_var = 1 + return called_fn() + local_var + closed_over_primitive + + ns = inspect_utils.getnamespace(local_fn) + self.assertEqual(ns['called_fn'], called_fn) + self.assertEqual(ns['closed_over_list'], closed_over_list) + self.assertEqual(ns['closed_over_primitive'], closed_over_primitive) + self.assertTrue('local_var' not in ns) + + def test_getmethodclass(self): + + self.assertEqual( + inspect_utils.getmethodclass(free_function), None) + self.assertEqual( + inspect_utils.getmethodclass(free_factory()), None) + + self.assertEqual( + inspect_utils.getmethodclass(TestClass.member_function), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(TestClass.decorated_member), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(TestClass.fn_decorated_member), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(TestClass.wrap_decorated_member), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(TestClass.static_method), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(TestClass.class_method), + TestClass) + + test_obj = TestClass() + self.assertEqual( + inspect_utils.getmethodclass(test_obj.member_function), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.decorated_member), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.fn_decorated_member), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.wrap_decorated_member), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.static_method), + TestClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.class_method), + TestClass) + + def test_getmethodclass_locals(self): + + def local_function(): + pass + + class LocalClass(object): + + def member_function(self): + pass + + @decorator + def decorated_member(self): + pass + + @function_decorator() + def fn_decorated_member(self): + pass + + @wrapping_decorator() + def wrap_decorated_member(self): + pass + + self.assertEqual( + inspect_utils.getmethodclass(local_function), None) + + self.assertEqual( + inspect_utils.getmethodclass(LocalClass.member_function), + LocalClass) + self.assertEqual( + inspect_utils.getmethodclass(LocalClass.decorated_member), + LocalClass) + self.assertEqual( + inspect_utils.getmethodclass(LocalClass.fn_decorated_member), + LocalClass) + self.assertEqual( + inspect_utils.getmethodclass(LocalClass.wrap_decorated_member), + LocalClass) + + test_obj = LocalClass() + self.assertEqual( + inspect_utils.getmethodclass(test_obj.member_function), + LocalClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.decorated_member), + LocalClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.fn_decorated_member), + LocalClass) + self.assertEqual( + inspect_utils.getmethodclass(test_obj.wrap_decorated_member), + LocalClass) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/pretty_printer.py b/tensorflow/contrib/py2tf/pyct/pretty_printer.py index 5e70c0ed833c10012e6a5b4cb26e9e4198162693..bacc1e4a7774ec5b84495255042392fe089150d5 100644 --- a/tensorflow/contrib/py2tf/pyct/pretty_printer.py +++ b/tensorflow/contrib/py2tf/pyct/pretty_printer.py @@ -25,24 +25,30 @@ import termcolor class PrettyPrinter(gast.NodeVisitor): """Print AST nodes.""" - def __init__(self): + 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 termcolor.colored(node.__class__.__name__, None, attrs=['bold']) + return self._color(node.__class__.__name__, None, ['bold']) def _field(self, name): - return termcolor.colored(name, 'blue') + return self._color(name, 'blue') def _value(self, name): - return termcolor.colored(name, 'magenta') + return self._color(name, 'magenta') def _warning(self, name): - return termcolor.colored(name, 'red') + return self._color(name, 'red') def _indent(self): - return termcolor.colored('| ' * self.indent_lvl, None, attrs=['dark']) + return self._color('| ' * self.indent_lvl, None, ['dark']) def _print(self, s): self.result += s @@ -76,6 +82,16 @@ class PrettyPrinter(gast.NodeVisitor): 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): @@ -87,8 +103,8 @@ class PrettyPrinter(gast.NodeVisitor): self.indent_lvl -= 1 -def fmt(node): - printer = PrettyPrinter() +def fmt(node, color=True): + printer = PrettyPrinter(color) if isinstance(node, (list, tuple)): for n in node: printer.visit(n) diff --git a/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py b/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py index 65e5b1d9191749a0caeeda48df37690564a8fc1e..81e3f47b80b6cb3bb7ba9f4a1787d03df4151a99 100644 --- a/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py +++ b/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py @@ -24,10 +24,6 @@ from tensorflow.contrib.py2tf.pyct import pretty_printer from tensorflow.python.platform import test -def f(x): - return x + 1 - - class PrettyPrinterTest(test.TestCase): def test_format(self): diff --git a/tensorflow/contrib/py2tf/pyct/qual_names.py b/tensorflow/contrib/py2tf/pyct/qual_names.py new file mode 100644 index 0000000000000000000000000000000000000000..6bcbaeb2aeb3043919e84bc6599edf5aee583c6d --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/qual_names.py @@ -0,0 +1,211 @@ +# 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 collections + +import gast + +from tensorflow.contrib.py2tf.pyct import anno + + +class Symbol(collections.namedtuple('Symbol', ['name'])): + """Represents a Python symbol.""" + + +class StringLiteral(collections.namedtuple('StringLiteral', ['value'])): + """Represents a Python string literal.""" + + def __str__(self): + return '\'%s\'' % self.value + + def __repr__(self): + return str(self) + + +class NumberLiteral(collections.namedtuple('NumberLiteral', ['value'])): + """Represents a Python numeric literal.""" + + def __str__(self): + return '%s' % self.value + + def __repr__(self): + return str(self) + + +# TODO(mdan): Use subclasses to remove the has_attr has_subscript booleans. +class QN(object): + """Represents a qualified name.""" + + def __init__(self, base, attr=None, subscript=None): + if attr is not None and subscript is not None: + raise ValueError('A QN can only be either an attr or a subscript, not ' + 'both: attr={}, subscript={}.'.format(attr, subscript)) + self._has_attr = False + self._has_subscript = False + + if attr is not None: + if not isinstance(base, QN): + raise ValueError( + 'for attribute QNs, base must be a QN; got instead "%s"' % base) + if not isinstance(attr, str): + raise ValueError('attr may only be a string; got instead "%s"' % attr) + self._parent = base + # TODO(mdan): Get rid of the tuple - it can only have 1 or 2 elements now. + self.qn = (base, attr) + self._has_attr = True + + elif subscript is not None: + if not isinstance(base, QN): + raise ValueError('For subscript QNs, base must be a QN.') + self._parent = base + self.qn = (base, subscript) + self._has_subscript = True + + else: + if not isinstance(base, (str, StringLiteral, NumberLiteral)): + # TODO(mdan): Require Symbol instead of string. + raise ValueError( + 'For simple QNs, base must be a string or a Literal object.') + assert '.' not in base and '[' not in base and ']' not in base + self._parent = None + self.qn = (base,) + + def is_symbol(self): + return isinstance(self.qn[0], str) + + def is_composite(self): + return len(self.qn) > 1 + + def has_subscript(self): + return self._has_subscript + + def has_attr(self): + return self._has_attr + + @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 + (self._has_attr, self._has_subscript)) + + def __eq__(self, other): + return (isinstance(other, QN) and self.qn == other.qn and + self.has_subscript() == other.has_subscript() and + self.has_attr() == other.has_attr()) + + def __str__(self): + if self.has_subscript(): + return str(self.qn[0]) + '[' + str(self.qn[1]) + ']' + if self.has_attr(): + return '.'.join(map(str, self.qn)) + else: + return str(self.qn[0]) + + def __repr__(self): + return str(self) + + def ssf(self): + """Simple symbol form.""" + ssfs = [n.ssf() if isinstance(n, QN) else n for n in self.qn] + ssf_string = '' + for i in range(0, len(self.qn) - 1): + if self.has_subscript(): + delimiter = '_sub_' + else: + delimiter = '_' + ssf_string += ssfs[i] + delimiter + return ssf_string + ssfs[-1] + + def ast(self): + # The caller must adjust the context appropriately. + if self.has_subscript(): + return gast.Subscript(self.parent.ast(), gast.Index(self.qn[-1].ast()), + None) + if self.has_attr(): + return gast.Attribute(self.parent.ast(), self.qn[-1], None) + + base = self.qn[0] + if isinstance(base, str): + return gast.Name(base, None, None) + elif isinstance(base, StringLiteral): + return gast.Str(base.value) + elif isinstance(base, NumberLiteral): + return gast.Num(base.value) + else: + assert False, ('the constructor should prevent types other than ' + 'str, StringLiteral and NumberLiteral') + + +class QnResolver(gast.NodeTransformer): + """Annotates nodes with QN information. + + Note: Not using NodeAnnos to avoid circular dependencies. + """ + + def visit_Call(self, node): + node = self.generic_visit(node) + # This helps treat the following cases uniformly: + # a = b[i] + # a = b()[i] + anno.copyanno(node.func, node, anno.Basic.QN) + return node + + def visit_Name(self, node): + node = self.generic_visit(node) + anno.setanno(node, anno.Basic.QN, QN(node.id)) + return node + + def visit_Attribute(self, node): + node = self.generic_visit(node) + anno.setanno(node, anno.Basic.QN, + QN(anno.getanno(node.value, anno.Basic.QN), attr=node.attr)) + return node + + def visit_Subscript(self, node): + node = self.generic_visit(node) + s = node.slice + if not isinstance(s, gast.Index): + # TODO(mdan): Support range and multi-dimensional indices. + # Continuing silently because some demos use these. + return node + if isinstance(s.value, gast.Num): + subscript = QN(NumberLiteral(s.value.n)) + elif isinstance(s.value, gast.Str): + subscript = QN(StringLiteral(s.value.s)) + else: + subscript = anno.getanno(node.slice.value, anno.Basic.QN) + anno.setanno(node, anno.Basic.QN, + QN(anno.getanno(node.value, anno.Basic.QN), + subscript=subscript)) + 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..f2cd8e98f02213c9035fdb5b20e0862f0a8fd3f6 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/qual_names_test.py @@ -0,0 +1,213 @@ +# 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.contrib.py2tf.pyct.qual_names import QN +from tensorflow.contrib.py2tf.pyct.qual_names import resolve +from tensorflow.python.platform import test + + +class QNTest(test.TestCase): + + def test_basic(self): + a = 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 = QN(a, attr='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',)) + + def test_subscripts(self): + a = QN('a') + b = QN('b') + a_sub_b = QN(a, subscript=b) + self.assertEqual(a_sub_b.qn, (a, b)) + self.assertEqual(str(a_sub_b), 'a[b]') + self.assertEqual(a_sub_b.ssf(), 'a_sub_b') + self.assertEqual(a_sub_b.ast().value.id, 'a') + self.assertEqual(a_sub_b.ast().slice.value.id, 'b') + self.assertTrue(a_sub_b.is_composite()) + self.assertTrue(a_sub_b.has_subscript()) + self.assertEqual(a_sub_b.parent.qn, ('a',)) + + c = QN('c') + b_sub_c = QN(b, subscript=c) + a_sub_b_sub_c = QN(a, subscript=b_sub_c) + self.assertEqual(a_sub_b_sub_c.qn, (a, b_sub_c)) + self.assertTrue(a_sub_b.is_composite()) + self.assertTrue(a_sub_b_sub_c.is_composite()) + self.assertTrue(a_sub_b.has_subscript()) + self.assertTrue(a_sub_b_sub_c.has_subscript()) + self.assertEqual(b_sub_c.qn, (b, c)) + self.assertEqual(str(a_sub_b_sub_c), 'a[b[c]]') + self.assertEqual(a_sub_b_sub_c.ssf(), 'a_sub_b_sub_c') + self.assertEqual(a_sub_b_sub_c.ast().value.id, 'a') + self.assertEqual(a_sub_b_sub_c.ast().slice.value.value.id, 'b') + self.assertEqual(a_sub_b_sub_c.ast().slice.value.slice.value.id, 'c') + self.assertEqual(b_sub_c.ast().slice.value.id, 'c') + self.assertEqual(a_sub_b_sub_c.parent.qn, ('a',)) + with self.assertRaises(ValueError): + QN('a', 'b') + + def test_equality(self): + a = QN('a') + a2 = QN('a') + a_b = QN(a, attr='b') + self.assertEqual(a2.qn, ('a',)) + with self.assertRaises(ValueError): + _ = a.parent + + a_b2 = QN(a, attr='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) + a_sub_b = QN(a, subscript='b') + a_sub_b2 = QN(a, subscript='b') + self.assertTrue(a_sub_b == a_sub_b2) + self.assertFalse(a_sub_b == a_b) + + def test_nested_attrs_subscripts(self): + a = QN('a') + b = QN('b') + c = QN('c') + b_sub_c = QN(b, subscript=c) + a_sub_b_sub_c = QN(a, subscript=b_sub_c) + + b_dot_c = QN(b, attr='c') + a_sub__b_dot_c = QN(a, subscript=b_dot_c) + + a_sub_b = QN(a, subscript=b) + a_sub_b__dot_c = QN(a_sub_b, attr='c') + + a_dot_b = QN(a, attr='b') + a_dot_b_sub_c = QN(a_dot_b, subscript=c) + + self.assertEqual(str(a_sub_b_sub_c), 'a[b[c]]') + self.assertEqual(str(a_sub__b_dot_c), 'a[b.c]') + self.assertEqual(str(a_sub_b__dot_c), 'a[b].c') + self.assertEqual(str(a_dot_b_sub_c), 'a.b[c]') + + self.assertNotEqual(a_sub_b_sub_c, a_sub__b_dot_c) + self.assertNotEqual(a_sub_b_sub_c, a_sub_b__dot_c) + self.assertNotEqual(a_sub_b_sub_c, a_dot_b_sub_c) + + self.assertNotEqual(a_sub__b_dot_c, a_sub_b__dot_c) + self.assertNotEqual(a_sub__b_dot_c, a_dot_b_sub_c) + + self.assertNotEqual(a_sub_b__dot_c, a_dot_b_sub_c) + + def test_hashable(self): + d = {QN('a'): 'a', QN('b'): 'b'} + self.assertEqual(d[QN('a')], 'a') + self.assertEqual(d[QN('b')], 'b') + self.assertTrue(QN('c') not in d) + + def test_literals(self): + a = QN('a') + a_sub_str_b = QN(a, subscript=QN(qual_names.StringLiteral('b'))) + a_sub_b = QN(a, subscript=QN('b')) + + self.assertNotEqual(a_sub_str_b, a_sub_b) + self.assertNotEqual(hash(a_sub_str_b), hash(a_sub_b)) + + a_sub_three = QN(a, subscript=QN(qual_names.NumberLiteral(3))) + self.assertEqual(a_sub_three.ast().slice.value.n, 3) + + +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 = 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') + + def test_subscript_resolve(self): + samples = """ + x[i] + x[i.b] + a.b[c] + a.b[x.y] + a[z[c]] + a[b[c[d]]] + a[b].c + a.b.c[d].e.f + a.b[c[d]].e.f + a.b[c[d.e.f].g].h + """ + nodes = resolve(parser.parse_str(textwrap.dedent(samples))) + nodes = tuple(n.value for n in nodes.body) + + self.assertQNStringIs(nodes[0], 'x[i]') + self.assertQNStringIs(nodes[1], 'x[i.b]') + self.assertQNStringIs(nodes[2], 'a.b[c]') + self.assertQNStringIs(nodes[3], 'a.b[x.y]') + self.assertQNStringIs(nodes[4], 'a[z[c]]') + self.assertQNStringIs(nodes[5], 'a[b[c[d]]]') + self.assertQNStringIs(nodes[6], 'a[b].c') + self.assertQNStringIs(nodes[7], 'a.b.c[d].e.f') + self.assertQNStringIs(nodes[8], 'a.b[c[d]].e.f') + self.assertQNStringIs(nodes[9], 'a.b[c[d.e.f].g].h') + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD b/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD index 32e2954fffca3b9f512116648117904b85a60e25..2799b56a0042e99b8f8b38100d07c5afaef9f424 100644 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD @@ -17,7 +17,8 @@ filegroup( py_library( name = "static_analysis", srcs = [ - "access.py", + "activity.py", + "annos.py", "live_values.py", "type_info.py", ], @@ -30,8 +31,8 @@ py_library( ) py_test( - name = "access_test", - srcs = ["access_test.py"], + name = "activity_test", + srcs = ["activity_test.py"], srcs_version = "PY2AND3", deps = [ ":static_analysis", @@ -59,6 +60,7 @@ py_test( deps = [ ":static_analysis", "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/contrib/py2tf/utils", "//tensorflow/python:client_testlib", ], ) diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/access_test.py b/tensorflow/contrib/py2tf/pyct/static_analysis/access_test.py deleted file mode 100644 index df0283b54d534443c0d4f044894903a4cb375f88..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/access_test.py +++ /dev/null @@ -1,236 +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 access 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.static_analysis import access -from tensorflow.python.platform import test - - -class ScopeTest(test.TestCase): - - def test_basic(self): - scope = access.Scope(None) - self.assertFalse(scope.has('foo')) - - scope.mark_read('foo') - self.assertFalse(scope.has('foo')) - - scope.mark_write('foo') - self.assertTrue(scope.has('foo')) - - scope.mark_read('bar') - self.assertFalse(scope.has('bar')) - - def test_copy(self): - scope = access.Scope(None) - scope.mark_write('foo') - - other = access.Scope(None) - other.copy_from(scope) - - self.assertTrue('foo' in other.created) - - scope.mark_write('bar') - scope.copy_from(other) - - self.assertFalse('bar' in scope.created) - - scope.mark_write('bar') - scope.merge_from(other) - - self.assertTrue('bar' in scope.created) - self.assertFalse('bar' in other.created) - - def test_nesting(self): - scope = access.Scope(None) - scope.mark_write('foo') - scope.mark_read('bar') - - child = access.Scope(scope) - self.assertTrue(child.has('foo')) - self.assertTrue(scope.has('foo')) - - child.mark_write('bar') - self.assertTrue(child.has('bar')) - self.assertFalse(scope.has('bar')) - - def test_referenced(self): - scope = access.Scope(None) - scope.mark_read('a') - - child = access.Scope(scope) - child.mark_read('b') - - child2 = access.Scope(child, isolated=False) - child2.mark_read('c') - - self.assertTrue('c' in child2.referenced) - self.assertTrue('b' in child2.referenced) - self.assertFalse('a' in child2.referenced) - - self.assertTrue('c' in child.referenced) - self.assertTrue('b' in child.referenced) - self.assertFalse('a' in child.referenced) - - -class AccessResolverTest(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 = access.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, - 'is_local')) # c in b = c - self.assertTrue(anno.getanno(node.body[0].body[1].test.left, - 'is_local')) # b in b > 0 - self.assertTrue(anno.getanno(node.body[0].body[2].value, - 'is_local')) # b in return b - - def assertScopeIs(self, scope, used, modified, created): - self.assertItemsEqual(used, scope.used) - self.assertItemsEqual(modified, scope.modified) - self.assertItemsEqual(created, 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, '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, '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, '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, 'body_scope'), ('b',), ('b', 'c'), ('c',)) - self.assertScopeIs( - anno.getanno(while_node, 'body_parent_scope'), ('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, 'body_scope'), ('b',), ('b', 'c'), ('c',)) - self.assertScopeIs( - anno.getanno(for_node, 'body_parent_scope'), ('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, '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, 'body_parent_scope'), ('x', 'z', 'u'), - ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u')) - self.assertScopeIs( - anno.getanno(if_node, 'orelse_scope'), ('x', 'y'), ('x', 'y', 'u'), - ('y', 'u')) - self.assertScopeIs( - anno.getanno(if_node, 'body_parent_scope'), ('x', 'z', 'u'), - ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u')) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/access.py b/tensorflow/contrib/py2tf/pyct/static_analysis/activity.py similarity index 63% rename from tensorflow/contrib/py2tf/pyct/static_analysis/access.py rename to tensorflow/contrib/py2tf/pyct/static_analysis/activity.py index 33629f87d1d9fe32fac5dd05ba5feaecf4be7230..87fc8c979c4e3310fb3aa82b0f23d909b0170cda 100644 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/access.py +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/activity.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Access information (reads, writes) resolution.""" +"""Activity analysis.""" from __future__ import absolute_import from __future__ import division @@ -24,6 +24,8 @@ import gast from tensorflow.contrib.py2tf.pyct import anno from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.qual_names import QN +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno # TODO(mdan): Add support for PY3 (e.g. Param vs arg). @@ -69,13 +71,33 @@ class Scope(object): tuple(self.modified)) def copy_from(self, other): + """Recursively copies the contents of this scope from another scope.""" + if (self.parent is None) != (other.parent is None): + raise ValueError('cannot copy scopes of different structures') + if other.parent is not None: + self.parent.copy_from(other.parent) + self.isolated = other.isolated 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) + @classmethod + def copy_of(cls, other): + if other.parent is not None: + parent = cls.copy_of(other.parent) + else: + parent = None + new_copy = cls(parent) + new_copy.copy_from(other) + return new_copy + def merge_from(self, other): + if (self.parent is None) != (other.parent is None): + raise ValueError('cannot merge scopes of different structures') + if other.parent is not None: + self.parent.merge_from(other.parent) self.modified |= other.modified self.created |= other.created self.used |= other.used @@ -112,6 +134,14 @@ class Scope(object): 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): @@ -132,39 +162,48 @@ class Scope(object): self.parent.mark_returned(name) -class AccessResolver(transformer.Base): +class ActivityAnalizer(transformer.Base): """Annotates nodes with local scope information. See Scope.""" - def __init__(self, context): - super(AccessResolver, self).__init__(context) - self.scope = Scope(None) + def __init__(self, context, parent_scope): + super(ActivityAnalizer, self).__init__(context) + self.scope = Scope(parent_scope) self._in_return_statement = False - def visit_Name(self, node): - # TODO(mdan): This is insufficient for object fields, e.g. hp.learning_rate. - self.generic_visit(node) + def _track_symbol(self, node): + qn = anno.getanno(node, anno.Basic.QN) + if isinstance(node.ctx, gast.Store): - self.scope.mark_write(node.id) + self.scope.mark_write(qn) elif isinstance(node.ctx, gast.Load): - anno.setanno(node, 'is_local', self.scope.has(node.id)) - self.scope.mark_read(node.id) + 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(node.id) - self.scope.mark_param(node.id) + self.scope.mark_creation(qn) + self.scope.mark_param(qn) else: - raise ValueError('Unknown context %s for node %s.' % (type(node.ctx), - node.id)) - anno.setanno(node, 'is_modified_since_entry', - self.scope.is_modified_since_entry(node.id)) - anno.setanno(node, 'is_param', self.scope.is_param(node.id)) + 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(node.id) + 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): @@ -173,7 +212,7 @@ class AccessResolver(transformer.Base): self.scope = args_scope for n in node.values: self.visit(n) - anno.setanno(node, 'args_scope', args_scope) + anno.setanno(node, NodeAnno.ARGS_SCOPE, args_scope) self.scope = current_scope return node @@ -186,7 +225,7 @@ class AccessResolver(transformer.Base): # TODO(mdan): Account starargs, kwargs for n in node.keywords: self.visit(n) - anno.setanno(node, 'args_scope', args_scope) + anno.setanno(node, NodeAnno.ARGS_SCOPE, args_scope) self.scope = current_scope self.visit(node.func) return node @@ -197,7 +236,7 @@ class AccessResolver(transformer.Base): self.scope = block_scope for n in block: self.visit(n) - anno.setanno(node, '%s_scope' % scope_name, block_scope) + anno.setanno(node, scope_name, block_scope) self.scope = current_scope return node @@ -206,39 +245,58 @@ class AccessResolver(transformer.Base): # 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) + before_parent = Scope.copy_of(self.scope) after_children = [] - for child, name in children: + for child, scope_name in children: self.scope.copy_from(before_parent) - parent = self._process_block_node(parent, child, name) - after_child = Scope(None) - after_child.copy_from(self.scope) + parent = self._process_block_node(parent, child, scope_name) + after_child = Scope.copy_of(self.scope) after_children.append(after_child) for after_child in after_children: self.scope.merge_from(after_child) - for child, name in children: - # TODO(mdan): We don't need this - we have the parent link from scope. - anno.setanno(parent, '%s_parent_scope' % name, self.scope) return parent + def visit_FunctionDef(self, node): + if self.scope: + qn = QN(node.name) + self.scope.mark_write(qn) + current_scope = self.scope + fndef_scope = Scope(current_scope, isolated=True) + self.scope = fndef_scope + self.generic_visit(node) + anno.setanno(node, NodeAnno.BODY_SCOPE, fndef_scope) + self.scope = current_scope + return node + + def visit_With(self, node): + current_scope = self.scope + with_scope = Scope(current_scope, isolated=False) + self.scope = with_scope + self.generic_visit(node) + anno.setanno(node, NodeAnno.BODY_SCOPE, with_scope) + self.scope = current_scope + return node + def visit_If(self, node): self.visit(node.test) - node = self._process_parallel_blocks( - node, ((node.body, 'body'), (node.orelse, 'orelse'))) + 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, 'body'), (node.orelse, 'orelse'))) + 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, 'body'), (node.orelse, 'orelse'))) + node = self._process_parallel_blocks(node, + ((node.body, NodeAnno.BODY_SCOPE), + (node.orelse, NodeAnno.ORELSE_SCOPE))) return node def visit_Return(self, node): @@ -248,5 +306,5 @@ class AccessResolver(transformer.Base): return node -def resolve(node, context): - return AccessResolver(context).visit(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..b16d15b39d8eb4c444cbc50ae62baa3a8fcc7841 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/activity_test.py @@ -0,0 +1,329 @@ +# 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_from(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_copy_of(self): + scope = activity.Scope(None) + scope.mark_read(QN('foo')) + + self.assertTrue(QN('foo') in activity.Scope.copy_of(scope).used) + + child_scope = activity.Scope(scope) + child_scope.mark_read(QN('bar')) + + self.assertTrue(QN('bar') in activity.Scope.copy_of(child_scope).used) + + 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, + owner_type=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 assertScopeIsRmc(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.assertScopeIsRmc(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.assertScopeIsRmc( + 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.assertScopeIsRmc( + anno.getanno(while_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'), + ('c',)) + self.assertScopeIsRmc( + 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.assertScopeIsRmc( + anno.getanno(for_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'), ('c',)) + self.assertScopeIsRmc( + 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.assertScopeIsRmc( + 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.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent, ('x', 'z', 'u'), + ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u')) + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), ('x', 'y'), + ('x', 'y', 'u'), ('y', 'u')) + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent, ('x', 'z', 'u'), + ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u')) + + def test_nested_if_else_creation(self): + + def test_fn(b): + if b > 0: + if b < 5: + a = b + else: + a = b * b + return a + + node = self._parse_and_analyze(test_fn) + inner_if_node = node.body[0].body[0].body[0] + self.assertScopeIsRmc( + anno.getanno(inner_if_node, NodeAnno.BODY_SCOPE), ('b',), ('a',), + ('a',)) + self.assertScopeIsRmc( + anno.getanno(inner_if_node, NodeAnno.ORELSE_SCOPE), ('b',), ('a',), + ('a',)) + + def test_function_def(self): + + def test_fn(a): + + def f(x): + y = x * x + return y + + b = a + for i in a: + c = b + b -= f(i) + return b, c + + node = self._parse_and_analyze(test_fn) + fndef_node = node.body[0].body[0] + + self.assertScopeIsRmc( + anno.getanno(fndef_node, + NodeAnno.BODY_SCOPE).parent, ('b', 'i', 'f', 'c', 'a'), + ('f', 'b', 'c', 'i'), ('f', 'a', 'b', 'c', 'i')) + self.assertScopeIsRmc( + anno.getanno(fndef_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('y',), ( + 'x', + 'y', + )) + + 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.assertScopeIsRmc( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'a.b', 'a.c'), (), + ()) + if_node = node.body[0].body[1] + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('a',), ('a.b',), ()) + self.assertScopeIsRmc( + 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..5254b83ca7c775867fc2ad5ef0a0ad93ac483ba0 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/annos.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. +# ============================================================================== +"""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 + # These flags are boolean. + 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 + # Scopes are represented by objects of type activity.Scope. + 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).') + + # Type and Value annotations + # Type annotations are represented by objects of type type_info.Type. + STATIC_INFO = ( + 'The type or value information that should be asserted about the entity ' + 'referenced by the symbol holding this annotation, irrespective of the ' + 'execution context.') diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py index 5a2903e6b5909e64e7bf2c9ea733a391da9292a1..0388be5d252389f2f3516c8b27828905d6475589 100644 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py @@ -16,7 +16,7 @@ Live values are extracted from the known execution context. -Requires annotations generated by AccessResolver. +Requires activity analysis annotations. """ from __future__ import absolute_import @@ -27,6 +27,7 @@ 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): @@ -44,12 +45,12 @@ class LiveValueResolver(transformer.Base): def visit_Name(self, node): self.generic_visit(node) if isinstance(node.ctx, gast.Load): - assert anno.hasanno(node, 'is_local'), node - symbol_is_local = anno.getanno(node, 'is_local') - assert anno.hasanno(node, 'is_modified_since_entry'), node - symbol_is_modified = anno.getanno(node, 'is_modified_since_entry') - assert anno.hasanno(node, 'is_param'), node - symbol_is_param = anno.getanno(node, 'is_param') + 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: @@ -60,7 +61,11 @@ class LiveValueResolver(transformer.Base): anno.setanno(node, 'live_val', obj) anno.setanno(node, 'fqn', (obj.__name__,)) else: - raise ValueError('Could not resolve symbol "%s".' % node.id) + 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. @@ -81,6 +86,7 @@ class LiveValueResolver(transformer.Base): if not hasattr(parent_object, node.attr): raise AttributeError('%s has no attribute %s' % (parent_object, node.attr)) + anno.setanno(node, 'parent_type', type(parent_object)) 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. @@ -91,13 +97,14 @@ class LiveValueResolver(transformer.Base): # 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, 'parent_type', parent_type) 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, 'is_local'), stem_name + 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 diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py index f3057b346677ac7fc8d36c1c14663c614f7e4ea9..c133a455b3dd328689102634c6076f366212ac25 100644 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py @@ -21,7 +21,8 @@ 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.static_analysis import access +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 @@ -45,8 +46,10 @@ class LiveValuesResolverTest(test.TestCase): namespace=namespace, arg_values=None, arg_types=arg_types, + owner_type=None, recursive=True) - node = access.resolve(node, ctx) + 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) @@ -100,6 +103,7 @@ class LiveValuesResolverTest(test.TestCase): 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, anno.getanno(func_node, 'parent_type')) self.assertEquals(('TestClass', 'member'), anno.getanno(func_node, 'fqn')) diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py index cf74142cbe9b1396585ac3e3c6d9858013f05ac2..5556a58c025da695bcef10352c597c7c8dd612d9 100644 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py @@ -14,9 +14,29 @@ # ============================================================================== """Type resolution. +This analyzer uses known live values to further infer object types. This +may include for instance constructed objects and object member functions. + +In addition, the analyzer will also process annotations for TF (staged) type +annotations. + Requires annotations generated by LiveValuesResolver. """ +# TODO(mdan): This would be more robust with a CFG. +# Situations with multiple reaching modifications (e.g. modified inside and +# outside a control flow statement) should be more robustly detected and +# analyzed. + +# TODO(mdan): Look into using Python AST's type annotation fields instead. +# It would be desirable to use that mechanism if we can. +# Some caveats to consider: We may need to annotate other nodes like +# Attribute. It may also not be feasible for us to faithfully to replicate +# PY3's type annotations where it isn't available. It would also require us +# to design rigorous type definitions that can accommodate Python types +# as well as TensorFLow dtypes and shapes. + + from __future__ import absolute_import from __future__ import division from __future__ import print_function @@ -29,7 +49,7 @@ from tensorflow.python.util import tf_inspect class Scope(object): - """Encloses symbol value references. + """Tracks symbol value references. Attributes: values: A dict mapping string to gast.Node, containing the value that was @@ -116,31 +136,36 @@ class TypeInfoResolver(transformer.Base): return node def _process_function_arg(self, arg_name): - if self.function_level == 1 and arg_name in self.context.arg_types: + 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 = gast.Name(arg_name, gast.Load(), None) - type_string, type_obj = self.context.arg_types[arg_name] + 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(node.arg) + 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(node.id) - elif isinstance(node.ctx, gast.Load) and self.scope.hasval(node.id): + 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(node.id) - 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')) + # then for future references to `a` we should have definition = `b` + definition = self.scope.getval(qn) + if anno.hasanno(definition, 'type'): + anno.setanno(node, 'type', anno.getanno(definition, 'type')) + anno.setanno(node, 'type_fqn', anno.getanno(definition, 'type_fqn')) + if anno.hasanno(definition, 'element_type'): + anno.setanno(node, 'element_type', + anno.getanno(definition, 'element_type')) return node def _process_variable_assignment(self, source, targets): @@ -159,16 +184,11 @@ class TypeInfoResolver(transformer.Base): for t in targets: if isinstance(t, gast.Tuple): for i, e in enumerate(t.elts): - self.scope.setval(e.id, - gast.Subscript( - source, gast.Index(i), ctx=gast.Store())) - elif isinstance(t, gast.Name): - self.scope.setval(t.id, source) - elif isinstance(t, gast.Attribute): - if not (isinstance(t.value, gast.Name) and t.value.id == 'self'): - raise ValueError( - 'Dont know how to handle assignment to attributes of objects' - ' other than "self": [%s].%s' % (t.value, t.attr)) + 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) @@ -184,6 +204,34 @@ class TypeInfoResolver(transformer.Base): self._process_variable_assignment(node.value, node.targets) return node + def visit_Call(self, node): + if anno.hasanno(node.func, 'live_val'): + # Symbols targeted by the "set_type" marker function are assigned the data + # type that it specified. + if (anno.getanno(node.func, 'live_val') is + self.context.type_annotation_func): + # Expecting the actual type to be the second argument. + if len(node.args) != 2: + raise ValueError('"%s" must have exactly two parameters' + % self.context.type_annotation_func) + if not anno.hasanno(node.args[0], anno.Basic.QN): + raise ValueError('the first argument of "%s" must by a symbol' + % self.context.type_annotation_func) + if not anno.hasanno(node.args[1], 'live_val'): + raise ValueError( + 'the second argument of "%s" must be statically resolvable' % + self.context.type_annotation_func) + target_symbol = anno.getanno(node.args[0], anno.Basic.QN) + element_type = anno.getanno(node.args[1], 'live_val') + # Find the definition of this symbol and annotate it with the given + # data type. That in turn will cause future uses of the symbol + # to receive the same type annotation. + definition = self.scope.getval(target_symbol) + anno.setanno(node, 'element_type', element_type) + anno.setanno(definition, 'element_type', element_type) + # TODO(mdan): Should we update references between definition and here? + return self.generic_visit(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 index 68fa1ee92a72b2c3d015b84724e63a17069c196c..0d9d5a85f055b170ea6e493e8ac185f1298ebf3c 100644 --- a/tensorflow/contrib/py2tf/pyct/static_analysis/type_info_test.py +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info_test.py @@ -18,10 +18,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.py2tf import utils 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.static_analysis import access +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 @@ -55,7 +57,10 @@ class ScopeTest(test.TestCase): class TypeInfoResolverTest(test.TestCase): - def _parse_and_analyze(self, test_fn, namespace, arg_types=None): + def _parse_and_analyze(self, + test_fn, + namespace, + arg_types=None): node, source = parser.parse_entity(test_fn) ctx = context.EntityContext( namer=None, @@ -64,8 +69,11 @@ class TypeInfoResolverTest(test.TestCase): namespace=namespace, arg_values=None, arg_types=arg_types, - recursive=True) - node = access.resolve(node, ctx) + owner_type=None, + recursive=True, + type_annotation_func=utils.set_element_type) + 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, {}) @@ -172,6 +180,22 @@ class TypeInfoResolverTest(test.TestCase): method_call = node.body[0].body[1].value.func self.assertFalse(anno.hasanno(method_call, 'live_val')) + def test_type_annotation(self): + + class Foo(object): + pass + + def test_fn(): + f = [] + f = utils.set_element_type(f, Foo) + return f + + node = self._parse_and_analyze(test_fn, {'Foo': Foo, 'utils': utils}) + f_def = node.body[0].body[0].value + self.assertEqual(anno.getanno(f_def, 'element_type'), Foo) + f_ref = node.body[0].body[1].value + self.assertEqual(anno.getanno(f_ref, 'element_type'), Foo) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/py2tf/pyct/templates.py b/tensorflow/contrib/py2tf/pyct/templates.py index 6be526f20ddfd5c2e847287bfbf11e640e2026e6..cdd71dc56de33cde46d6115085350a321093d792 100644 --- a/tensorflow/contrib/py2tf/pyct/templates.py +++ b/tensorflow/contrib/py2tf/pyct/templates.py @@ -22,12 +22,13 @@ from __future__ import division from __future__ import print_function import ast -import copy 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): @@ -41,6 +42,7 @@ class ReplaceTransformer(gast.NodeTransformer): 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. @@ -57,39 +59,75 @@ class ReplaceTransformer(gast.NodeTransformer): 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', + 'A function name can only be replaced by a Name node. Found: %s' % repl) node.name = repl.id return node - def visit_Name(self, node): - if node.id in self.replacements: - # TODO(mdan): Sanitize the nodes by erasing scope-dependent annotations. - new_nodes = copy.copy(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: - e.ctx = node.ctx - n.ctx = node.ctx - if len(new_nodes) == 1: - new_nodes, = new_nodes - return new_nodes + 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_Attribute(self, node): + node = self.generic_visit(node) + if node.attr not in self.replacements: + return node + repl = self.replacements[node.attr] + if not isinstance(repl, gast.Name): + raise ValueError( + 'An attribute can only be replaced by a Name node. Found: %s' % repl) + node.attr = repl.id + return 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 _strings_to_names(n): +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 [_strings_to_names(e) for e in n] + return [_convert_to_ast(e) for e in n] if isinstance(n, tuple): - return tuple(_strings_to_names(e) for e in n) + return tuple(_convert_to_ast(e) for e in n) return n @@ -122,5 +160,22 @@ def replace(template, **replacements): raise ValueError('Expected string template, got %s' % type(template)) tree = parser.parse_str(textwrap.dedent(template)) for k in replacements: - replacements[k] = _strings_to_names(replacements[k]) - return ReplaceTransformer(replacements).visit(tree).body + 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) + + +def replace_as_expression(template, **replacements): + """Variant of replace that generates expressions, instead of code blocks.""" + replacement = replace(template, **replacements) + if len(replacement) != 1: + raise ValueError( + 'single expression expected; for more general templates use replace') + node = replacement[0] + if not isinstance(node, gast.Expr): + raise ValueError( + 'the template is expected to generate an expression node; instead ' + 'found %s' % node) + return node.value diff --git a/tensorflow/contrib/py2tf/pyct/templates_test.py b/tensorflow/contrib/py2tf/pyct/templates_test.py index 1143131283cd92c42abfc73d5728fac96cc31c23..d7835b80a7f53c3ba012d01cac34b68c57bfe348 100644 --- a/tensorflow/contrib/py2tf/pyct/templates_test.py +++ b/tensorflow/contrib/py2tf/pyct/templates_test.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import imp + import gast from tensorflow.contrib.py2tf.pyct import compiler @@ -27,6 +29,17 @@ 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): @@ -36,7 +49,7 @@ class TemplatesTest(test.TestCase): """ node = templates.replace(template, a='b')[0] - result = compiler.ast_to_object(node) + result, _ = compiler.ast_to_object(node) self.assertEquals(7, result.test_fn(2)) def test_replace_function_name(self): @@ -48,10 +61,10 @@ class TemplatesTest(test.TestCase): """ node = templates.replace(template, fname='test_fn')[0] - result = compiler.ast_to_object(node) + result, _ = compiler.ast_to_object(node) self.assertEquals(7, result.test_fn(2)) - def test_code_block(self): + def test_replace_code_block(self): template = """ def test_fn(a): block @@ -65,9 +78,46 @@ class TemplatesTest(test.TestCase): 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) + result, _ = compiler.ast_to_object(node) self.assertEquals(3, result.test_fn(1)) + def test_replace_attribute(self): + template = """ + def test_fn(a): + return a.foo + """ + + node = templates.replace(template, foo='b')[0] + result, _ = compiler.ast_to_object(node) + mod = imp.new_module('test') + mod.b = 3 + self.assertEquals(3, result.test_fn(mod)) + + with self.assertRaises(ValueError): + templates.replace(template, foo=1) + + def replace_as_expression(self): + template = """ + foo(a) + """ + + node = templates.replace(template, foo='bar', a='baz') + self.assertTrue(node is gast.Call) + self.assertEqual(node.func.id, 'bar') + self.assertEqual(node.func.args[0].id, 'baz') + + def replace_as_expression_restrictions(self): + template = """ + foo(a) + bar(b) + """ + with self.assertRaises(ValueError): + templates.replace_as_expression(template) + with self.assertRaises(ValueError): + templates.replace('') + with self.assertRaises(ValueError): + templates.replace('a = b') + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/py2tf/pyct/transformer.py b/tensorflow/contrib/py2tf/pyct/transformer.py index 8a836b7c1bf1dfdcfda84588be8fd904d970579e..57016bb4ce84776dfc8dfbe380322a03eb4b37b8 100644 --- a/tensorflow/contrib/py2tf/pyct/transformer.py +++ b/tensorflow/contrib/py2tf/pyct/transformer.py @@ -23,6 +23,7 @@ import sys import gast import six +from tensorflow.contrib.py2tf.pyct import anno from tensorflow.contrib.py2tf.pyct import pretty_printer @@ -43,17 +44,26 @@ class Base(gast.NodeTransformer): self._col_offset = 0 self.context = context + def debug_print(self, node): + """Helper method useful for debugging.""" + if __debug__: + print(pretty_printer.fmt(node)) + return node + def visit(self, node): + source_code = self.context.source_code + source_file = self.context.source_file try: - source_code = self.context.source_code - source_file = self.context.source_file 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, NotImplementedError) as e: - msg = '%s: %s\nOccurred at node:\n%s' % (e.__class__.__name__, str(e), - pretty_printer.fmt(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: diff --git a/tensorflow/contrib/py2tf/utils/BUILD b/tensorflow/contrib/py2tf/utils/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..8bc338e801aa283967f4f6e6a659df9683cbc154 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/BUILD @@ -0,0 +1,112 @@ +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", + "builtins.py", + "context_managers.py", + "misc.py", + "multiple_dispatch.py", + "py_func.py", + "tensor_list.py", + "testing.py", + "type_check.py", + "type_hints.py", + ], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:__subpackages__"], + deps = [ + "//tensorflow/python:list_ops", + "//tensorflow/python:script_ops", + "@six_archive//:six", + ], +) + +py_test( + name = "builtins_test", + srcs = ["builtins_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +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", + tags = ["no_windows"], # TODO: needs investigation on Windows + 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..d9d8e3468966bc9da31c3fc756a9660f5ff7d115 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/__init__.py @@ -0,0 +1,34 @@ +# 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.builtins import dynamic_builtin +from tensorflow.contrib.py2tf.utils.builtins import dynamic_print +from tensorflow.contrib.py2tf.utils.builtins import dynamic_range +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 dynamic_is +from tensorflow.contrib.py2tf.utils.multiple_dispatch import dynamic_is_not +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.py_func import wrap_py_func +from tensorflow.contrib.py2tf.utils.tensor_list import dynamic_list_append +from tensorflow.contrib.py2tf.utils.testing import fake_tf +from tensorflow.contrib.py2tf.utils.type_check import is_tensor +from tensorflow.contrib.py2tf.utils.type_hints import set_element_type diff --git a/tensorflow/contrib/py2tf/utils/builtins.py b/tensorflow/contrib/py2tf/utils/builtins.py new file mode 100644 index 0000000000000000000000000000000000000000..3cb62b55d4d23545af4d641ecab1663ee7f7b876 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/builtins.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. +# ============================================================================== +"""Builtin conversion utilities.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import six + +from tensorflow.contrib.py2tf.utils import py_func +from tensorflow.contrib.py2tf.utils import type_check +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import logging_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.util import tf_inspect + + +def dynamic_builtin(f, *args, **kwargs): + """Converts a builtin function call inline.""" + # Some built-ins may be objects. + if not tf_inspect.isbuiltin(f) and f not in (range,): + return f(*args, **kwargs) + + if f is len: + return dynamic_len(*args, **kwargs) + if six.PY2 and f is xrange: + return dynamic_range(*args, **kwargs) + if f is range: + return dynamic_range(*args, **kwargs) + + raise NotImplementedError( + 'The "%s" builtin is not yet supported.' % f.__name__) + + +def dynamic_len(list_or_tensor): + """Implementation of len using dynamic dispatch.""" + if tensor_util.is_tensor(list_or_tensor): + shape = list_or_tensor.shape + if not shape: + raise ValueError( + 'len requires non-zero rank for tensor "%s"' % list_or_tensor) + return array_ops.shape(list_or_tensor)[0] + + return len(list_or_tensor) + + +def dynamic_range(start_or_stop, stop=None, step=None): + """Implementation of range using dynamic dispatch.""" + if type_check.is_tensor(start_or_stop, stop, step): + if step is not None: + return math_ops.range(start_or_stop, stop, step) + if stop is not None: + return math_ops.range(start_or_stop, stop) + return math_ops.range(start_or_stop) + + if step is not None: + return range(start_or_stop, stop, step) + elif stop is not None: + return range(start_or_stop, stop) + return range(start_or_stop) + + +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 dynamic_print(*values): + """Implementartion of print using dynamic dispatch. + + 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/builtins_test.py b/tensorflow/contrib/py2tf/utils/builtins_test.py new file mode 100644 index 0000000000000000000000000000000000000000..59b3573d38c5bd98f416c7b77d1bc772cb8069dd --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/builtins_test.py @@ -0,0 +1,111 @@ +# 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 builtins 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 builtins +from tensorflow.python.framework import constant_op +from tensorflow.python.platform import test + + +class BuiltinsTest(test.TestCase): + + def test_dynamic_len_tf_scalar(self): + a = constant_op.constant(1) + + with self.assertRaises(ValueError): + with self.test_session() as sess: + sess.run(builtins.dynamic_builtin(len, a)) + + def test_dynamic_len_tf_array(self): + a = constant_op.constant([1, 2, 3]) + + with self.test_session() as sess: + self.assertEqual(3, sess.run(builtins.dynamic_builtin(len, a))) + + def test_dynamic_len_tf_matrix(self): + a = constant_op.constant([[1, 2], [3, 4]]) + + with self.test_session() as sess: + self.assertEqual(2, sess.run(builtins.dynamic_builtin(len, a))) + + def test_dynamic_len_py_list(self): + a = [3] * 5 + + self.assertEqual(5, builtins.dynamic_builtin(len, a)) + + def test_dynamic_range_all_python(self): + self.assertListEqual(list(builtins.dynamic_builtin(range, 3)), [0, 1, 2]) + self.assertListEqual(list(builtins.dynamic_builtin(range, 1, 3)), [1, 2]) + self.assertListEqual( + list(builtins.dynamic_builtin(range, 2, 0, -1)), [2, 1]) + + def test_dynamic_range_tf(self): + with self.test_session() as sess: + self.assertAllEqual( + sess.run(builtins.dynamic_builtin(range, constant_op.constant(3))), + [0, 1, 2]) + self.assertAllEqual( + sess.run(builtins.dynamic_builtin(range, 1, constant_op.constant(3))), + [1, 2]) + self.assertAllEqual( + sess.run( + builtins.dynamic_builtin(range, 2, 0, constant_op.constant(-1))), + [2, 1]) + + def test_dynamic_range_detection(self): + def range(x): # pylint:disable=redefined-builtin + return x + + # Functions that just have the names of builtins are ignored. + self.assertEqual(builtins.dynamic_builtin(range, 1), 1) + if six.PY2: + self.assertListEqual( + list(builtins.dynamic_builtin(xrange, 3)), [0, 1, 2]) + self.assertListEqual( + list(builtins.dynamic_builtin(six.moves.range, 3)), [0, 1, 2]) + self.assertListEqual( + list(builtins.dynamic_builtin(six.moves.xrange, 3)), [0, 1, 2]) + + def test_dynamic_print_tf(self): + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + with self.test_session() as sess: + sess.run(builtins.dynamic_print('test message', 1)) + self.assertEqual(out_capturer.getvalue(), 'test message 1\n') + finally: + sys.stdout = sys.__stdout__ + + def test_dynamic_print_complex(self): + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + with self.test_session() as sess: + sess.run(builtins.dynamic_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/context_managers.py b/tensorflow/contrib/py2tf/utils/context_managers.py new file mode 100644 index 0000000000000000000000000000000000000000..3d150a95817b83c4d7aaa78dc250092dcc4c5a9b --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/context_managers.py @@ -0,0 +1,49 @@ +# 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 +from tensorflow.python.ops import tensor_array_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. + """ + def control_dependency_handle(t): + if isinstance(t, tensor_array_ops.TensorArray): + return t.flow + return t + + 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_value = tuple(control_dependency_handle(t) for t in 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..404f6e44e59d8bd6131367e3234843f03b351910 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/context_managers_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 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.framework import dtypes +from tensorflow.python.ops import tensor_array_ops +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( + tensor_array_ops.TensorArray(dtypes.int32, size=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..8aedd4cd64798660cc07364c45487399986c9be6 --- /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.misc import alias_tensors +from tensorflow.python.framework.constant_op import constant +from tensorflow.python.ops.variables import Variable +from tensorflow.python.platform import test + + +class MiscTest(test.TestCase): + + def test_alias_single_tensor(self): + a = constant(1) + + new_a = 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(1) + v = Variable(2) + s = 'a' + l = [1, 2, 3] + + new_a, new_v, new_s, new_l = 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..da7a942703d83b55edbd1607cb49ad4137daeb04 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/multiple_dispatch.py @@ -0,0 +1,105 @@ +# 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 +from tensorflow.python.ops import math_ops + + +def dynamic_is(left, right): + if is_tensor(left, right): + return math_ops.equal(left.name, right.name) + else: + return left is right + + +def dynamic_is_not(left, right): + if is_tensor(left, right): + return math_ops.not_equal(left.name, right.name) + else: + return left is not right + + +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..8d89b6898a366fe90ee1d43a55d0a7f10690224b --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/multiple_dispatch_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 multiple_dispatch.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +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_dynamic_is_python(self): + a = np.eye(3) + also_a = a + not_actually_a = np.eye(3) + should_be_true1 = multiple_dispatch.dynamic_is(a, also_a) + should_be_false1 = multiple_dispatch.dynamic_is_not(a, also_a) + should_be_true2 = multiple_dispatch.dynamic_is_not(a, not_actually_a) + should_be_false2 = multiple_dispatch.dynamic_is(a, not_actually_a) + self.assertTrue(should_be_true1) + self.assertTrue(should_be_true2) + self.assertFalse(should_be_false1) + self.assertFalse(should_be_false2) + + def test_dynamic_is_tf(self): + with Session().as_default(): + a = constant([2.0]) + also_a = a + not_actually_a = constant([2.0]) + should_be_true1 = multiple_dispatch.dynamic_is(a, also_a) + should_be_false1 = multiple_dispatch.dynamic_is_not(a, also_a) + should_be_true2 = multiple_dispatch.dynamic_is_not(a, not_actually_a) + should_be_false2 = multiple_dispatch.dynamic_is(a, not_actually_a) + self.assertTrue(should_be_true1.eval()) + self.assertTrue(should_be_true2.eval()) + self.assertFalse(should_be_false1.eval()) + self.assertFalse(should_be_false2.eval()) + + 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/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..2556f412891b4f0b954af5a6f0193341a6a5020a --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/tensor_list.py @@ -0,0 +1,68 @@ +# 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.framework import ops +from tensorflow.python.ops import list_ops +from tensorflow.python.ops import tensor_array_ops + + +def dynamic_list_append(target, element): + """Converts a list append call inline.""" + if isinstance(target, tensor_array_ops.TensorArray): + return target.write(target.size(), element) + # TODO(mdan): What's the right way to check this? + # TODO(mdan): We may not need this branch. + # It may be possible to use TensorList alone if the loop body will not + # require wrapping it, although we'd have to think about an autoboxing + # mechanism for lists received as parameter. + if isinstance(target, ops.Tensor): + return list_ops.tensor_list_push_back(target, element) + + # Python targets (including TensorList): fallback to their original append. + target.append(element) + return target + + +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..110e4d105e934d9d752afc2ccc0c53c99b70d41d --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/tensor_list_test.py @@ -0,0 +1,117 @@ +# 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 dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework.constant_op import constant +from tensorflow.python.ops import list_ops +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.platform import test + + +class TensorListTest(test.TestCase): + + def _shape(self, shape_tuple): + return constant(shape_tuple, dtypes.int32) + + def test_dynamic_list_append(self): + l = [] + l = tl.dynamic_list_append(l, 1) + self.assertListEqual(l, [1]) + + l = list_ops.empty_tensor_list(self._shape(()), dtypes.int32) + l = tl.dynamic_list_append(l, 1) + s = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32) + with self.test_session() as sess: + self.assertAllEqual(sess.run(s), [1]) + + l = tensor_array_ops.TensorArray(dtypes.int32, size=0, dynamic_size=True) + l = tl.dynamic_list_append(l, 1) + s = l.stack() + with self.test_session() as sess: + self.assertAllEqual(sess.run(s), [1]) + + l = tl.TensorList(self._shape(()), dtypes.int32) + l = tl.dynamic_list_append(l, 1) + with self.test_session() as sess: + self.assertAllEqual(sess.run(l[0]), 1) + + 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/quantize/python/copy_graph.py b/tensorflow/contrib/py2tf/utils/testing.py similarity index 68% rename from tensorflow/contrib/quantize/python/copy_graph.py rename to tensorflow/contrib/py2tf/utils/testing.py index 0376fcba82b99feabdba3b683f9db9a32db51efb..cb4785d0dc0f4674b3560418daeb6733364b21e7 100644 --- a/tensorflow/contrib/quantize/python/copy_graph.py +++ b/tensorflow/contrib/py2tf/utils/testing.py @@ -12,21 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utility to copy a tf.Graph.""" +"""Testing utilities.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function +import imp + from tensorflow.python.framework import ops -from tensorflow.python.training import saver as saver_lib +from tensorflow.python.ops import math_ops -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 fake_tf(): + """Creates a fake module that looks like TensorFlow, for testing.""" + mod = imp.new_module('tensorflow') + mod_contents = dict() + mod_contents.update(math_ops.__dict__) + mod_contents.update(ops.__dict__) + mod_contents.update(mod.__dict__) + mod.__dict__.update(mod_contents) + return mod diff --git a/tensorflow/contrib/bayesflow/python/ops/halton_sequence.py b/tensorflow/contrib/py2tf/utils/type_check.py similarity index 63% rename from tensorflow/contrib/bayesflow/python/ops/halton_sequence.py rename to tensorflow/contrib/py2tf/utils/type_check.py index 49d747d538f5a4aa3134d28ba00a651cb509fa41..b9b2b451a4e22684a19f0d10fbf5e4fae5d6152b 100644 --- a/tensorflow/contrib/bayesflow/python/ops/halton_sequence.py +++ b/tensorflow/contrib/py2tf/utils/type_check.py @@ -1,4 +1,4 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -12,22 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Support for low discrepancy Halton sequences. - -""" +"""Utilities used in py2tf-generated code.""" 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.halton_sequence_impl import * -# pylint: enable=wildcard-import -from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.framework import tensor_util + + +def is_tensor(*args): + """Check if any arguments are tensors. -_allowed_symbols = [ - 'sample', -] + Args: + *args: Python objects that may or may not be tensors. -remove_undocumented(__name__, _allowed_symbols) + Returns: + True if any *args are TensorFlow types, False if none are. + """ + return any([tensor_util.is_tensor(a) for a in args]) diff --git a/tensorflow/contrib/py2tf/converters/print_functions_test.py b/tensorflow/contrib/py2tf/utils/type_check_test.py similarity index 57% rename from tensorflow/contrib/py2tf/converters/print_functions_test.py rename to tensorflow/contrib/py2tf/utils/type_check_test.py index 475196ce102955b350acf9bf94255997f875f62c..7d0428e9cccecdc67511e236bc00655a055aea29 100644 --- a/tensorflow/contrib/py2tf/converters/print_functions_test.py +++ b/tensorflow/contrib/py2tf/utils/type_check_test.py @@ -12,33 +12,31 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for print_functions module.""" +"""Tests for type_check.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import gast +import numpy -from tensorflow.contrib.py2tf.converters import converter_test_base -from tensorflow.contrib.py2tf.converters import print_functions -from tensorflow.contrib.py2tf.pyct import compiler +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 PrintFunctionsTest(converter_test_base.TestCase): +class TypeCheckTest(test.TestCase): - def test_transform(self): - - def test_fn(a): - print(a) - - node = self.parse_and_analyze(test_fn, {'print': print}) - node = print_functions.transform(node) - result = compiler.ast_to_object(node) - - result.test_fn('a') - self.assertTrue(isinstance(node.body[0].body[0].value, gast.Call)) + 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__': diff --git a/tensorflow/contrib/py2tf/utils/type_hints.py b/tensorflow/contrib/py2tf/utils/type_hints.py new file mode 100644 index 0000000000000000000000000000000000000000..aeb9e545610460afbe364dfcfc7a54b9aede29fe --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/type_hints.py @@ -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. +# ============================================================================== +"""No-op utilities that provide static type hints. + +These are used when the data type is not known at creation, for instance in the +case of empty lists. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +def set_element_type(entity, dtype, shape=None): + """Indicates that the entity is expected hold items of specified type. + + This function is a no-op. Its presence merely marks the data type of its + argument. The staged TensorFlow ops will reflect and assert this data type. + + Args: + entity: A Tensor or TensorArray. + dtype: TensorFlow dtype value to assert for entity. + shape: Optional shape to assert for entity. + Returns: + The value of entity, unchanged. + """ + del dtype + del shape + return entity diff --git a/tensorflow/contrib/quantization/README.md b/tensorflow/contrib/quantization/README.md new file mode 100644 index 0000000000000000000000000000000000000000..359950aaf3d89c1f3e8fda21cbd27fb89217d918 --- /dev/null +++ b/tensorflow/contrib/quantization/README.md @@ -0,0 +1,7 @@ +The contrib/quantization package exposes a few TensorFlow quantization operations. + +If you are looking for quantized training rewrites that allow for training +quantized models that work with +[TensorFlow Lite](https://www.tensorflow.org/mobile/tflite/), you should look at +the [contrib/quantize](https://www.tensorflow.org/api_docs/python/tf/contrib/quantize) +package. diff --git a/tensorflow/contrib/quantize/BUILD b/tensorflow/contrib/quantize/BUILD index 3c5b34a0a6adb2f4e340a8e378c1eb51a2e2b534..0b7629620418340d803753be0df1f04c342dc490 100644 --- a/tensorflow/contrib/quantize/BUILD +++ b/tensorflow/contrib/quantize/BUILD @@ -13,6 +13,21 @@ 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", + "//tensorflow/python:variable_scope", + ], +) + py_library( name = "graph_matcher", srcs = [ @@ -75,11 +90,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 +110,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 +124,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 +155,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 +166,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 +213,6 @@ py_test( "//tensorflow/python:math_ops", "//tensorflow/python:nn_ops", "//tensorflow/python:platform_test", - "//tensorflow/python:training", ], ) @@ -229,12 +224,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 +239,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/README.md b/tensorflow/contrib/quantize/README.md index 40541729da5fd9d0ae75579e11f20999337de124..348c824a4072c3329ac4a3441c19c71598bc9c03 100644 --- a/tensorflow/contrib/quantize/README.md +++ b/tensorflow/contrib/quantize/README.md @@ -1,9 +1,9 @@ +# Quantized Training Rewrites + tf.contrib.quantize provides tools for transforming graphs to include ops to model quantization of weights, biases and activations during both training and inference. This is done using the -[fake quantization op] -(https://www.tensorflow.org/versions/r0.12/api_docs/python/array_ops/fake_quantization), -which is described below: +[fake quantization op](https://www.tensorflow.org/versions/r0.12/api_docs/python/array_ops/fake_quantization). Recent literature has shown that fixed point networks provide comparable performance to floating point networks [1]. This is achieved by modeling the @@ -14,56 +14,52 @@ updated at high precision as this is needed to ensure sufficient precision in accumulating tiny adjustments to the parameters. However, for the forward pass, the parameters and activations are quantized to the desired lower precision. -![drawing](g3doc/drawings/Fake_Quantization.jpg) - -###Forward pass - - - - -\begin{equation*} -f_Q(x) = \Delta\text{ }round\left(\frac{sat\left(x\right)-x_{min}}{\Delta}\right) -\end{equation*} - - -where - -$$ -\begin{equation*} -sat(x) = -\left\{ - \begin{array}{ll} - x_{min} & \mbox{if } x \le x_{min} \\ - x & \mbox{if } x_{min} \leq x \leq x_{max} \\ - x_{max} & \mbox{if } x_{max} \le x - \end{array} -\right. -\end{equation*} -$$ - - -where $$\Delta$$ is the Quantizer Step size, given by -$$\Delta =\frac{x_{max} - x_{min} }{255} $$ and $$x_{min} $$ and $$x_{max}$$ are -the minimum and maximum values of the variable under consideration. Note that -the rounding performed is deterministic and corresponds to asymmetric rounding, -which is supported in almost all hardware platforms. - -###Backward pass -For the backward pass, we model the quantizer as a piecewise linear block, with -derivatives that are non-zero only in the linear region. - - - -\begin{equation*} -\frac{df_Q(x)}{dx}=1, x_{min} \leq x \leq x_{max},\text{ 0 elsewhere } -\end{equation*} - -Therefore, the backward pass through the quantizer reduces to passing through -the gradients as long as the inputs to the quantizer are in the linear region. -Otherwise, the gradients are set to zero. - -Note that the quantizer is fully specified by the min and max values of the -variables being quantized. +## How to use the Rewrites + +tf.contrib.quantize provides two rewrites, one to train for quantization and +one to create a [TensorFlow Lite](https://www.tensorflow.org/mobile/tflite/) +compatible eval graph. + +``` +# Build forward pass of model. +… +loss = tf.losses.get_total_loss() + +# Call the training rewrite which rewrites the graph in-place with FakeQuantization nodes +# and folds batchnorm for training. +# It is often needed to finetune a floating point model for quantization with this training tool. +# When training from scratch, quant_delay can be used to activate quantization after +# training to convergence with the float graph, effectively finetuning the model. +tf.contrib.quantize.create_training_graph(quant_delay=2000000) + +# Call backward pass optimizer as usual. +optimizer = tf.train.GradientDescentOptimizer(learning_rate) +optimizer.minimize(loss) +``` + +Additionally, the rewritten eval graph is non-trivially different from the +training graph due the effects of quantization on batch normalization. Thus, +we offer a separate rewrite for the eval_graph. + +``` +# Build eval model +… +logits = tf.nn.softmax_cross_entropy_with_logits(...) + +# Call the eval rewrite which rewrites the graph in-place with FakeQuantization nodes +# and fold batchnorm for eval. +tf.contrib.quantize.create_eval_graph() + +# Save the checkpoint and eval graph proto to disk for freezing and providing to TFLite. +with open(eval_graph_file, ‘w’) as f: + f.write(str(g.as_graph_def())) +saver = tf.train.Saver() +saver.save(sess, checkpoint_name) +``` + +These rewrites are an active area of research and experimentation, so the +rewrites and quantized training will likely not work across all models, though +we hope to work towards generalizing these techniques. [1] P.Gysel, "HARDWARE-ORIENTED APPROXIMATION OF CONVOLUTIONAL diff --git a/tensorflow/contrib/quantize/g3doc/drawings/Fake_Quantization.jpg b/tensorflow/contrib/quantize/g3doc/drawings/Fake_Quantization.jpg deleted file mode 100644 index fdc7ae40cec757cc0a93d50eca6c8698a4697d07..0000000000000000000000000000000000000000 Binary files a/tensorflow/contrib/quantize/g3doc/drawings/Fake_Quantization.jpg and /dev/null differ diff --git a/tensorflow/contrib/quantize/python/common.py b/tensorflow/contrib/quantize/python/common.py index d0b0674c31239ee903f5ab7ef9ae0262bb20d189..3138149468c87b649f120a81b28c60947dcf2c4e 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,13 @@ 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 array_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 +93,33 @@ 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 + '/Identity: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. We return an identity to force resource variables and + # normal variables to return a tensor of the same name. + return array_ops.identity( + 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..06c62f2d265503bf42d46fb682a398ce1f4d15fb --- /dev/null +++ b/tensorflow/contrib/quantize/python/common_test.py @@ -0,0 +1,67 @@ +# 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 variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.platform import googletest + + +class CommonTest(test_util.TensorFlowTestCase): + + def testCreateOrGetQuantizationStep(self): + self._TestCreateOrGetQuantizationStep(False) + + def testCreateOrGetQuantizationStepResourceVar(self): + self._TestCreateOrGetQuantizationStep(True) + + def _TestCreateOrGetQuantizationStep(self, use_resource): + g = ops.Graph() + with session.Session(graph=g) as sess: + variable_scope.get_variable_scope().set_use_resource(use_resource) + 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..1afcbb850496fa1f654b7340946e087d403d3549 100644 --- a/tensorflow/contrib/quantize/python/fold_batch_norms.py +++ b/tensorflow/contrib/quantize/python/fold_batch_norms.py @@ -26,14 +26,16 @@ 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 from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import variable_scope 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 +43,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 +66,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 +79,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 +90,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 +110,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 +139,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 +155,64 @@ 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( - 'FusedBatchNorm', - inputs=[ - conv_pattern, gamma_pattern, beta_pattern, mean_pattern, - variance_pattern - ]) - matmul_batch_norm_pattern = graph_matcher.OpTypePattern( + 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])) + + 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]) - conv_matcher = graph_matcher.GraphMatcher(conv_batch_norm_pattern) - matmul_matcher = graph_matcher.GraphMatcher(matmul_bn_output_reshape_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 = 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] + + # Ensure that the output tensor has consumers, otherwise this is a dangling + # node and not a match. + if not output_tensor.consumers(): + continue - 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) @@ -220,50 +237,32 @@ def _FindFusedBatchNorms(graph): # The batch variance used during forward and backward prop is biased, # i.e it is calculated as: V=sum(x(k)-mu)^2/N. For the moving average # calculation, the variance is corrected by the term N/N-1 (Bessel's - # correction). The variance tensor read from FuseBatchNorm has 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 +271,155 @@ def _FindFusedBatchNorms(graph): gamma_tensor=gamma_tensor, beta_tensor=beta_tensor, mean_tensor=mean_tensor, - 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=batch_epsilon) + + +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 -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 - - @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 containing 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) + recip_sigma = math_ops.rsqrt(match.variance_tensor + match.batch_epsilon) + 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 +427,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. @@ -345,8 +439,16 @@ def _FoldUnfusedBatchNorms(graph): for bn in common.BatchNormGroups(graph): has_scaling = _HasScaling(graph, input_to_ops_map, bn) + if not _IsValidUnfusedBatchNorm(graph, bn): + continue + # 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 +470,101 @@ 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: +def _IsValidUnfusedBatchNorm(graph, context): + """Checks that the output of the unfused batch norm has consumers.""" + add_shift = graph.get_operation_by_name( + context + '/BatchNorm/batchnorm/add_1') + # Ensure that the output tensor of batch norm has consumers, otherwise this + # is a dangling node and not a match. + return bool(add_shift.outputs[0].consumers()) - 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 = 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_mean = base_context + '/BatchNorm/moments/Squeeze' + op_suffix_variance = base_context + '/BatchNorm/moments/Squeeze_1' + 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' + + if variable_scope.get_variable_scope().use_resource: + op_suffix_gamma = base_context + '/BatchNorm/gamma/Read/ReadVariableOp' + op_suffix_moving_variance = ( + base_context + '/BatchNorm/moving_variance/Read/ReadVariableOp') + op_suffix_moving_mean = ( + base_context + '/BatchNorm/moving_mean/Read/ReadVariableOp') + else: + op_suffix_gamma = base_context + '/BatchNorm/gamma' + op_suffix_moving_variance = base_context + '/BatchNorm/moving_variance/read' + op_suffix_moving_mean = base_context + '/BatchNorm/moving_mean/read' + + # 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 = 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=batch_epsilon) + + +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 +574,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 +595,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 +640,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 +662,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 +794,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): + 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 = batch_epsilon + + @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(self): + return self._batch_epsilon + + @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..a4f7b1b22139588be29171126d43b872d6658168 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 = [] @@ -282,8 +282,8 @@ def _FakeQuantWithMinMaxVars(inputs, min_var, max_var, per_channel, num_bits, Args: inputs: a tensor containing values to be quantized. min_var: a variable containing quantization range lower end(s). - max_var: a variable containing quantization range lupper end(s). - per_channel: a boolean specifying whether to use per-channel quantizatioh. + max_var: a variable containing quantization range upper end(s). + per_channel: a boolean specifying whether to use per-channel quantization. num_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^num_bits - 1] or wide range [0; 2^num_bits - 1]. diff --git a/tensorflow/contrib/quantize/python/quantize.py b/tensorflow/contrib/quantize/python/quantize.py index 50a2b4c91c9e7a2681f6041646a023a4225fb0c5..f1c8e678f71342b1241753275c78e09cac5460e9 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,36 @@ 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. +_WEIGHT_TYPES = {'Variable', 'VariableV2', 'VarHandleOp'} 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.GLOBAL_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 +60,292 @@ 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|ReadVariableOp', 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) + + # Match the final layer, where there will not be an activation and instead + # the output of the final BiasAdd must be quantized, so we treat it as the + # 'activation_op' in the _LayerMatch. + # TODO(suharshs): Figure out how to quantize this final layer across many + # models. + final_layer_matcher = graph_matcher.GraphMatcher(bias_add_pattern) + for match_result in final_layer_matcher.match_graph(graph): + layer_op = match_result.get_op(layer_pattern) + weight_tensor = match_result.get_tensor(weight_pattern) + activation_op = match_result.get_op(bias_add_pattern) + yield _LayerMatch(layer_op, weight_tensor, activation_op, None, None) + + +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.GLOBAL_VARIABLES, + 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. + 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 bbd9743d8014ce495a4967e7484981f7e60ae4a3..92557ac00003256df9b6224a35d0632fd4b5901e 100644 --- a/tensorflow/contrib/quantize/python/quantize_graph.py +++ b/tensorflow/contrib/quantize/python/quantize_graph.py @@ -18,177 +18,206 @@ 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. + + Variables added by the rewrite get added to the global variables collection. + + 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. - - 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. + 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. + tf.Operation. """ - return _create_graph( + # 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, - elements=elements, - device_name_or_function=device_name_or_function) + quant_delay=quant_delay, + freeze_bn_delay=freeze_bn_delay) -def create_eval_graph(input_graph, elements=None, device_name_or_function=None): - """Returns a transformed eval input_graph for simulated quantization. +def create_eval_graph(input_graph=None): + """Rewrites an eval input_graph in place for simulated quantization. - The forward pass has fake quantization ops inserted to simulate the error - introduced by quantization. + Variables added by the rewrite get added to the global variables collection. - 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. + 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. - 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. + Args: + input_graph: The tf.Graph to be transformed, if None then defaults to the + default graph. Raises: ValueError: If elements contains an element that isn't a tf.Tensor or - tf.Operation. + tf.Operation. """ - return _create_graph( - input_graph=input_graph, - is_training=False, - elements=elements, - device_name_or_function=device_name_or_function) + _create_graph(input_graph=input_graph, is_training=False) -def experimental_create_training_graph(input_graph, - elements=None, - device_name_or_function=None): - """Returns a transformed training input_graph for simulated quantization. +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. + + Variables added by the rewrite get added to the global variables collection. This function has additional experimental options not (yet) available to create_training_graph. The resulting behavior may be undefined. - The forward pass has fake quantization ops inserted to simulate the error - introduced by quantization. - 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. + 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. - 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. + 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, - elements=None, - device_name_or_function=None): - """Returns a transformed eval input_graph for simulated quantization. +def experimental_create_eval_graph(input_graph=None, + weight_bits=8, + activation_bits=8): + """Rewrites an eval input_graph in place for simulated quantization. + + Variables added by the rewrite get added to the global variables collection. 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 514862a0ab5b796718a04aa65a46e7a7e3b86330..b9d03c1bc059fe7bcce75978f503cbbf76090dbd 100644 --- a/tensorflow/contrib/quantize/python/quantize_graph_test.py +++ b/tensorflow/contrib/quantize/python/quantize_graph_test.py @@ -20,13 +20,11 @@ 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 @@ -34,7 +32,7 @@ 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 _RunTestOverParameters(self, test_fn): + def _RunTestOverAllRewrites(self, test_fn): rewrite_fns = [ quantize_graph.create_training_graph, quantize_graph.create_eval_graph, @@ -44,85 +42,202 @@ class QuantizeGraphTest(test_util.TensorFlowTestCase): for fn in rewrite_fns: test_fn(fn) - def testReturnedElements(self): - self._RunTestOverParameters(self._TestReturnElements) + 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 _TestReturnElements(self, fn): - 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] - q_graph, returned_elements = fn(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 testNoReturnElements(self): - self._RunTestOverParameters(self._TestNoReturnElements) - - def _TestNoReturnElements(self, fn): - graph = ops.Graph() - with graph.as_default(): - a = constant_op.constant(1.0) - b = variables.Variable(2.0) - _ = a + b - q_graph = fn(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 testDeviceName(self): - self._RunTestOverParameters(self._TestDeviceName) - - def _TestDeviceName(self, 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' - q_graph = fn(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. + rewrite_fn(graph) - Function is defined purely to shorten the name so that it stops wrapping. - - 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 testIdempotent(self): + self._RunTestOverAllRewrites(self._TestIdempotent) + + def _TestIdempotent(self, rewrite_fn): + with ops.Graph().as_default() as g: + self._ConvLayer() + rewrite_fn() + graph_def_before = str(g.as_graph_def()) + # Ensuring that calling the rewrite again doesn't add more nodes. + rewrite_fn() + graph_def_after = str(g.as_graph_def()) + self.assertEqual(graph_def_before, graph_def_after) + + 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..db745aa56212af6a9c20e06ee9e4e5d6e27cf3c3 100644 --- a/tensorflow/contrib/quantize/python/quantize_parameterized_test.py +++ b/tensorflow/contrib/quantize/python/quantize_parameterized_test.py @@ -28,8 +28,8 @@ from tensorflow.python.ops import control_flow_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 variable_scope from tensorflow.python.platform import googletest -from tensorflow.python.training import training batch_norm = layers.batch_norm conv2d = layers.conv2d @@ -57,10 +57,95 @@ class QuantizeTest(test_util.TensorFlowTestCase): (array_ops.identity, 'Identity', True, 5000), ] for params in parameters_list: - test_fn(params[0], params[1], params[2], params[3]) + # Test everything with resource variables and normal variables. + test_fn(params[0], params[1], params[2], params[3], False) + test_fn(params[0], params[1], params[2], params[3], True) + + def _AssertCorrectQuantizedGraphWithoutBatchNorm( + self, graph, scope, layer, activation_op_name, with_bypass, delay, + use_resource): + 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) + + # Assemble the expected inputs. + if use_resource: + expected_inputs = [ + scope + '/weights_quant/FakeQuantWithMinMaxVars/ReadVariableOp', + scope + '/weights_quant/FakeQuantWithMinMaxVars/ReadVariableOp_1', + ] + if layer == 'DepthwiseConv2dNative': + expected_inputs.append(scope + '/depthwise/ReadVariableOp') + else: + expected_inputs.append(scope + '/' + layer + '/ReadVariableOp') + else: + expected_inputs = [ + scope + '/weights_quant/AssignMinLast', + scope + '/weights_quant/AssignMaxLast', + ] + if layer == 'DepthwiseConv2dNative': + expected_inputs.append(scope + '/depthwise_weights/read') + else: + expected_inputs.append(scope + '/weights/read') + + self._AssertInputOpsAre(weights_quant, expected_inputs) + if delay and delay > 0: + output_op_name = scope + '/weights_quant/delayed_quant/Switch_1' + else: + if layer == 'DepthwiseConv2dNative': + output_op_name = scope + '/depthwise' + else: + output_op_name = scope + '/' + layer + + self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) + + if with_bypass: + conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + + quantization_node_name) + self.assertEqual(conv_quant.type, quantization_node_name) + if use_resource: + expected_inputs = [ + scope + '/conv_quant/FakeQuantWithMinMaxVars/ReadVariableOp', + scope + '/conv_quant/FakeQuantWithMinMaxVars/ReadVariableOp_1', + scope + '/BiasAdd', + ] + else: + expected_inputs = [ + scope + '/conv_quant/AssignMinEma', + scope + '/conv_quant/AssignMaxEma', scope + '/BiasAdd' + ] + self._AssertInputOpsAre(conv_quant, expected_inputs) + output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' + if delay else 'test/Add') + self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) + + act_quant = graph.get_operation_by_name('test/act_quant/' + + quantization_node_name) + self.assertEqual(act_quant.type, quantization_node_name) + if use_resource: + expected_inputs = [ + 'test/act_quant/FakeQuantWithMinMaxVars/ReadVariableOp', + 'test/act_quant/FakeQuantWithMinMaxVars/ReadVariableOp_1', + 'test/' + activation_op_name, + ] + else: + expected_inputs = [ + 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', + 'test/' + activation_op_name + ] + self._AssertInputOpsAre(act_quant, expected_inputs) + output_op_name = ('test/act_quant/delayed_quant/Switch_1' + if delay else 'control_dependency') + self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) + self._AssertIdempotent(graph) + + def testQuantize_Conv2dWithoutBatchNorm(self): + self._RunWithoutBatchNormTestOverParameters( + self._TestQuantize_Conv2dWithoutBatchNorm) def _TestQuantize_Conv2dWithoutBatchNorm(self, activation, activation_op_name, - with_bypass, delay): + with_bypass, delay, use_resource): """Tests quantization: inputs -> Conv2d no batch norm -> Activation. Args: @@ -70,20 +155,25 @@ class QuantizeTest(test_util.TensorFlowTestCase): 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. + use_resource: Bool, when true uses resource variables. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - + variable_scope.get_variable_scope().set_use_resource(use_resource) 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 out_depth = 3 if with_bypass else 32 activation_fn = None if with_bypass else activation scope = 'test/test2' if with_bypass else 'test' - node = conv2d(inputs, out_depth, [5, 5], stride=stride, padding='SAME', - weights_initializer=self._WeightInit(0.09), - activation_fn=activation_fn, scope=scope) + node = conv2d( + inputs, + out_depth, [5, 5], + stride=stride, + padding='SAME', + weights_initializer=self._WeightInit(0.09), + activation_fn=activation_fn, + scope=scope) if with_bypass: node = math_ops.add(inputs, node, name='test/Add') node = activation(node, name='test/' + activation_op_name) @@ -91,51 +181,18 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - quantize.Quantize(graph, 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/AssignMinLast', - scope + '/weights_quant/AssignMaxLast', scope + '/weights/read' - ] - self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + '/Conv2D' - self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) + quantize.Quantize(graph, True, quant_delay=delay) - if with_bypass: - conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + - quantization_node_name) - self.assertEqual(conv_quant.type, quantization_node_name) - expected_inputs = [ - scope + '/conv_quant/AssignMinEma', - scope + '/conv_quant/AssignMaxEma', scope + '/BiasAdd' - ] - self._AssertInputOpsAre(conv_quant, expected_inputs) - output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' - if delay else 'test/Add') - self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) + self._AssertCorrectQuantizedGraphWithoutBatchNorm( + graph, scope, 'Conv2D', activation_op_name, with_bypass, delay, + use_resource) - act_quant = graph.get_operation_by_name('test/act_quant/' + - quantization_node_name) - self.assertEqual(act_quant.type, quantization_node_name) - - expected_inputs = [ - 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', - 'test/' + activation_op_name - ] - self._AssertInputOpsAre(act_quant, expected_inputs) - output_op_name = ('test/act_quant/delayed_quant/Switch_1' - if delay else 'control_dependency') - self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) - - def testQuantize_Conv2dWithoutBatchNorm(self): + def testQuantize_FCWithoutBatchNorm(self): self._RunWithoutBatchNormTestOverParameters( - self._TestQuantize_Conv2dWithoutBatchNorm) + self._TestQuantize_FCWithoutBatchNorm) def _TestQuantize_FCWithoutBatchNorm(self, activation, activation_op_name, - with_bypass, delay): + with_bypass, delay, use_resource): """Tests quantization: inputs -> FC no batch norm -> Activation. Args: @@ -145,71 +202,40 @@ class QuantizeTest(test_util.TensorFlowTestCase): 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. + use_resource: Bool, when true uses resource variables. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - + variable_scope.get_variable_scope().set_use_resource(use_resource) batch_size, depth = 5, 256 inputs = array_ops.zeros((batch_size, depth)) out_depth = 256 if with_bypass else 128 activation_fn = None if with_bypass else activation scope = 'test/test2' if with_bypass else 'test' - node = fully_connected(inputs, out_depth, - weights_initializer=self._WeightInit(0.03), - activation_fn=activation_fn, scope=scope) + node = fully_connected( + inputs, + out_depth, + weights_initializer=self._WeightInit(0.03), + activation_fn=activation_fn, + scope=scope) if with_bypass: node = math_ops.add(inputs, node, name='test/Add') node = activation(node, name='test/' + activation_op_name) update_barrier = control_flow_ops.no_op(name='update_barrier') with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') + quantize.Quantize(graph, True, quant_delay=delay) - quantize.Quantize(graph, 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/AssignMinLast', - scope + '/weights_quant/AssignMaxLast', scope + '/weights/read' - ] - self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + '/MatMul' - self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) + self._AssertCorrectQuantizedGraphWithoutBatchNorm( + graph, scope, 'MatMul', activation_op_name, with_bypass, delay, + use_resource) - if with_bypass: - conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + - quantization_node_name) - self.assertEqual(conv_quant.type, quantization_node_name) - expected_inputs = [ - scope + '/conv_quant/AssignMinEma', - scope + '/conv_quant/AssignMaxEma', scope + '/BiasAdd' - ] - self._AssertInputOpsAre(conv_quant, expected_inputs) - output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' - if delay else 'test/Add') - self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) - - act_quant = graph.get_operation_by_name('test/act_quant/' + - quantization_node_name) - self.assertEqual(act_quant.type, quantization_node_name) - expected_inputs = [ - 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', - 'test/' + activation_op_name - ] - self._AssertInputOpsAre(act_quant, expected_inputs) - output_op_name = ('test/act_quant/delayed_quant/Switch_1' - if delay else 'control_dependency') - self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) - - def testQuantize_FCWithoutBatchNorm(self): + def testQuantize_DepthwiseConv2dWithoutBatchNorm(self): self._RunWithoutBatchNormTestOverParameters( - self._TestQuantize_FCWithoutBatchNorm) + self._TestQuantize_DepthwiseConv2dWithoutBatchNorm) def _TestQuantize_DepthwiseConv2dWithoutBatchNorm( - self, activation, activation_op_name, with_bypass, delay): + self, activation, activation_op_name, with_bypass, delay, use_resource): """Tests quantization: inputs -> DWConv2d no batch norm -> Activation. Args: @@ -219,70 +245,36 @@ class QuantizeTest(test_util.TensorFlowTestCase): 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. + use_resource: Bool, when true uses resource variables. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - + variable_scope.get_variable_scope().set_use_resource(use_resource) 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 activation_fn = None if with_bypass else activation scope = 'test/test2' if with_bypass else 'test' - node = separable_conv2d(inputs, None, [5, 5], stride=stride, - depth_multiplier=1.0, padding='SAME', - weights_initializer=self._WeightInit(0.09), - activation_fn=activation_fn, scope=scope) + node = separable_conv2d( + inputs, + None, [5, 5], + stride=stride, + depth_multiplier=1.0, + padding='SAME', + weights_initializer=self._WeightInit(0.09), + activation_fn=activation_fn, + scope=scope) if with_bypass: node = math_ops.add(inputs, node, name='test/Add') node = activation(node, name='test/' + activation_op_name) update_barrier = control_flow_ops.no_op(name='update_barrier') with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') + quantize.Quantize(graph, True, quant_delay=delay) - quantize.Quantize(graph, 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/AssignMinLast', - scope + '/weights_quant/AssignMaxLast', - scope + '/depthwise_weights/read' - ] - self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + '/depthwise' - self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) - - if with_bypass: - conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + - quantization_node_name) - self.assertEqual(conv_quant.type, quantization_node_name) - expected_inputs = [ - scope + '/conv_quant/AssignMinEma', - scope + '/conv_quant/AssignMaxEma', scope + '/BiasAdd' - ] - self._AssertInputOpsAre(conv_quant, expected_inputs) - output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' - if delay else 'test/Add') - self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) - - act_quant = graph.get_operation_by_name('test/act_quant/' + - quantization_node_name) - self.assertEqual(act_quant.type, quantization_node_name) - expected_inputs = [ - 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', - 'test/' + activation_op_name - ] - self._AssertInputOpsAre(act_quant, expected_inputs) - output_op_name = ('test/act_quant/delayed_quant/Switch_1' - if delay else 'control_dependency') - self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) - - def testQuantize_DepthwiseConv2dWithoutBatchNorm(self): - self._RunWithoutBatchNormTestOverParameters( - self._TestQuantize_DepthwiseConv2dWithoutBatchNorm) + self._AssertCorrectQuantizedGraphWithoutBatchNorm( + graph, scope, 'DepthwiseConv2dNative', activation_op_name, with_bypass, + delay, use_resource) def _RunBatchNormTestOverParameters(self, test_fn): # TODO(suharshs): Use parameterized test once OSS TF supports it. @@ -314,42 +306,88 @@ class QuantizeTest(test_util.TensorFlowTestCase): (array_ops.identity, 'Identity', True, 5000, True) ] for params in parameters_list: - test_fn(params[0], params[1], params[2], params[3], params[4]) + # Test everything with resource variables and normal variables. + test_fn(params[0], params[1], params[2], params[3], params[4], False) + test_fn(params[0], params[1], params[2], params[3], params[4], True) - def _TestQuantize_Conv2dWithBatchNorm(self, activation, activation_op_name, - with_bypass, delay, fused_batch_norm): - """Tests quantization: inputs -> Conv2d with batch norm -> Activation. + def _AssertCorrectQuantizedGraphWithBatchNorm(self, graph, scope, layer, + activation_op_name, with_bypass, + delay, use_resource): + 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) + if use_resource: + expected_inputs = [ + scope + '/weights_quant/FakeQuantWithMinMaxVars/ReadVariableOp', + scope + '/weights_quant/FakeQuantWithMinMaxVars/ReadVariableOp_1', + ] + else: + expected_inputs = [ + scope + '/weights_quant/' + 'AssignMinLast', + scope + '/weights_quant/' + 'AssignMaxLast' + ] + expected_inputs.append(scope + '/mul_fold') - 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) + self._AssertInputOpsAre(weights_quant, expected_inputs) + if layer == 'DepthwiseConv2dNative': + output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' + if delay else '/depthwise_Fold') + else: + output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' + if delay else '/' + layer + '_Fold') + self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) + + if with_bypass: + conv_quant = graph.get_operation_by_name( + scope + '/conv_quant/' + quantization_node_name) + self.assertEqual(conv_quant.type, quantization_node_name) + + if use_resource: + expected_inputs = [ + scope + '/conv_quant/FakeQuantWithMinMaxVars/ReadVariableOp', + scope + '/conv_quant/FakeQuantWithMinMaxVars/ReadVariableOp_1', + ] + else: + expected_inputs = [ + scope + '/conv_quant/AssignMinEma', + scope + '/conv_quant/AssignMaxEma', + ] + expected_inputs.append(scope + '/add_fold') + + self._AssertInputOpsAre(conv_quant, expected_inputs) + output_op_name = ( + scope + '/conv_quant/delayed_quant/Switch_1' if delay else 'test/Add') + self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) + + act_quant = graph.get_operation_by_name( + 'test/act_quant/' + quantization_node_name) + self.assertEqual(act_quant.type, quantization_node_name) + + if use_resource: + expected_inputs = [ + 'test/act_quant/FakeQuantWithMinMaxVars/ReadVariableOp', + 'test/act_quant/FakeQuantWithMinMaxVars/ReadVariableOp_1', + ] + else: + expected_inputs = [ + 'test/act_quant/AssignMinEma', + 'test/act_quant/AssignMaxEma', + ] + expected_inputs.append('test/' + activation_op_name) + + self._AssertInputOpsAre(act_quant, expected_inputs) + output_op_name = ('test/act_quant/delayed_quant/Switch_1' + if delay else 'control_dependency') + self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) + self._AssertIdempotent(graph) def testQuantize_Conv2dWithBatchNorm(self): self._RunBatchNormTestOverParameters(self._TestQuantize_Conv2dWithBatchNorm) - def _testQuantize_Conv2dWithBatchNorm(self, activation, activation_op_name, + def _TestQuantize_Conv2dWithBatchNorm(self, activation, activation_op_name, with_bypass, delay, fused_batch_norm, - use_ema): + use_resource): """Tests quantization: inputs -> Conv2d with batch norm -> Activation. Args: @@ -360,12 +398,11 @@ 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. + use_resource: Bool, when true uses resource variables. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - + variable_scope.get_variable_scope().set_use_resource(use_resource) 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 @@ -382,7 +419,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): normalizer_params=self._BatchNormParams(fused_batch_norm), scope=scope) - # Manually add a bypass (optionaly) and an activation. + # Manually add a bypass (optional) and an activation. if with_bypass: node = math_ops.add(inputs, node, name='test/Add') @@ -392,86 +429,19 @@ 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, True, quant_delay=delay) - quantize.Quantize( - graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema) - - 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' - ] - self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' - if (delay and use_ema) else '/Conv2D_Fold') - self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) - - if with_bypass: - conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + - quantization_node_name) - self.assertEqual(conv_quant.type, quantization_node_name) - expected_inputs = [ - scope + '/conv_quant/AssignMinEma', - scope + '/conv_quant/AssignMaxEma', scope + '/add_fold' - ] - self._AssertInputOpsAre(conv_quant, expected_inputs) - output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' - if delay else 'test/Add') - self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) - - act_quant = graph.get_operation_by_name('test/act_quant/' + - quantization_node_name) - self.assertEqual(act_quant.type, quantization_node_name) - expected_inputs = [ - 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', - 'test/' + activation_op_name - ] - self._AssertInputOpsAre(act_quant, expected_inputs) - output_op_name = ('test/act_quant/delayed_quant/Switch_1' - 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) + self._AssertCorrectQuantizedGraphWithBatchNorm( + graph, scope, 'Conv2D', activation_op_name, with_bypass, delay, + use_resource) def testQuantize_FCWithBatchNorm(self): self._RunBatchNormTestOverParameters(self._TestQuantize_FCWithBatchNorm) - def _testQuantize_FCWithBatchNorm(self, activation, activation_op_name, + def _TestQuantize_FCWithBatchNorm(self, activation, activation_op_name, with_bypass, delay, fused_batch_norm, - use_ema): + use_resource): """Tests quantization: inputs -> FC with batch norm -> Activation. Args: @@ -482,12 +452,11 @@ 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. + use_resource: Bool, when true uses resource variables. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - + variable_scope.get_variable_scope().set_use_resource(use_resource) batch_size, depth = 5, 256 inputs = array_ops.zeros((batch_size, depth)) out_depth = 256 if with_bypass else 128 @@ -501,7 +470,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): normalizer_params=self._BatchNormParams(fused_batch_norm), scope=scope) - # Manually add a bypass (optionaly) and an activation. + # Manually add a bypass (optional) and an activation. if with_bypass: node = math_ops.add(inputs, node, name='test/Add') @@ -511,88 +480,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' - ] - self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' - if delay and use_ema else '/MatMul_Fold') - self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) - - if with_bypass: - conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + - quantization_node_name) - self.assertEqual(conv_quant.type, quantization_node_name) - expected_inputs = [ - scope + '/conv_quant/AssignMinEma', - scope + '/conv_quant/AssignMaxEma', scope + '/add_fold' - ] - self._AssertInputOpsAre(conv_quant, expected_inputs) - output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' - if delay else 'test/Add') - self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) - - act_quant = graph.get_operation_by_name('test/act_quant/' + - quantization_node_name) - self.assertEqual(act_quant.type, quantization_node_name) - expected_inputs = [ - 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', - 'test/' + activation_op_name - ] - self._AssertInputOpsAre(act_quant, expected_inputs) - output_op_name = ('test/act_quant/delayed_quant/Switch_1' - 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) + self._AssertCorrectQuantizedGraphWithBatchNorm( + graph, scope, 'MatMul', activation_op_name, with_bypass, delay, + use_resource) 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, use_resource): """Tests quantization: inputs -> DWConv2d with batch norm -> Activation. Args: @@ -603,12 +505,11 @@ 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. + use_resource: Bool, when true uses resource variables. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - + variable_scope.get_variable_scope().set_use_resource(use_resource) 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 @@ -625,7 +526,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): normalizer_params=self._BatchNormParams(fused_batch_norm), scope=scope) - # Manually add a bypass (optionaly) and an activation. + # Manually add a bypass (optional) and an activation. if with_bypass: node = math_ops.add(inputs, node, name='test/Add') @@ -635,50 +536,22 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - fold_batch_norms.FoldBatchNorms(graph) - - quantize.Quantize( - graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema) - 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' - ] - self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' - if delay and use_ema else '/depthwise_Fold') - self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) + fold_batch_norms.FoldBatchNorms(graph, is_training=True) + quantize.Quantize(graph, True, quant_delay=delay) - if with_bypass: - conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' + - quantization_node_name) - self.assertEqual(conv_quant.type, quantization_node_name) - expected_inputs = [ - scope + '/conv_quant/AssignMinEma', - scope + '/conv_quant/AssignMaxEma', scope + '/add_fold' - ] - self._AssertInputOpsAre(conv_quant, expected_inputs) - output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1' - if delay else 'test/Add') - self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name]) + self._AssertCorrectQuantizedGraphWithBatchNorm( + graph, scope, 'DepthwiseConv2dNative', activation_op_name, + with_bypass, delay, use_resource) - act_quant = graph.get_operation_by_name('test/act_quant/' + - quantization_node_name) - self.assertEqual(act_quant.type, quantization_node_name) - expected_inputs = [ - 'test/act_quant/AssignMinEma', 'test/act_quant/AssignMaxEma', - 'test/' + activation_op_name - ] - self._AssertInputOpsAre(act_quant, expected_inputs) - output_op_name = ('test/act_quant/delayed_quant/Switch_1' - if delay else 'control_dependency') - self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) + def _AssertIdempotent(self, graph): + # Ensure that calling the rewrite again doesn't change the graph. + graph_def_before = str(graph.as_graph_def()) + with graph.as_default(): + # Ensuring that calling the rewrite again doesn't add more nodes. + fold_batch_norms.FoldBatchNorms(graph, is_training=True) + quantize.Quantize(graph, True) + graph_def_after = str(graph.as_graph_def()) + self.assertEqual(graph_def_before, graph_def_after) def _BatchNormParams(self, fused=False): return {'center': True, 'scale': True, 'decay': 1.0 - 0.003, 'fused': fused} @@ -692,7 +565,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): stddev: Standard deviation of normal variable. Returns: - An initialized that initialzes with a truncated normal variable. + An initialized that initializes with a truncated normal variable. """ return init_ops.truncated_normal_initializer(stddev=stddev) diff --git a/tensorflow/contrib/quantize/python/quantize_test.py b/tensorflow/contrib/quantize/python/quantize_test.py index 1e4dd7cf67dbfbd16386fd740c7dcc83e05ad82a..b2e5707a6d5bccd90ae9663ad7e8bf1a0482a297 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,14 +106,35 @@ 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/' + quantization_node_name) self.assertEqual(add_quant.type, quantization_node_name) + def testFinalLayerQuantized(self): + self._RunTestOverParameters(self._TestFinalLayerQuantized) + + def _TestFinalLayerQuantized(self, is_training): + graph = ops.Graph() + with graph.as_default(): + batch_size, height, width, depth = 5, 128, 128, 3 + input1 = array_ops.zeros((batch_size, height, width, depth)) + _ = conv2d( + input1, + 32, [5, 5], + stride=2, + padding='SAME', + weights_initializer=self._WeightInit(0.09), + activation_fn=None, + scope='test') + # Ensure that the a FakeQuant operation is in the outputs of the BiasAdd. + bias_add_op = graph.get_operation_by_name('test/BiasAdd') + quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8) + self.assertTrue('FakeQuantWithMinMaxVars' in + [op.type for op in bias_add_op.outputs[0].consumers()]) + def _WeightInit(self, stddev): """Returns truncated normal variable initializer. @@ -111,7 +144,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): stddev: Standard deviation of normal variable. Returns: - An initialized that initialzes with a truncated normal variable. + An initialized that initializes with a truncated normal variable. """ return init_ops.truncated_normal_initializer(stddev=stddev) 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/remote_fused_graph/pylib/BUILD b/tensorflow/contrib/remote_fused_graph/pylib/BUILD index 27f0a7f58f78135f1d73ae04bd1e76ef496fa549..54c66271cd43c85bc7be139dbdf70b9178ad1622 100644 --- a/tensorflow/contrib/remote_fused_graph/pylib/BUILD +++ b/tensorflow/contrib/remote_fused_graph/pylib/BUILD @@ -38,7 +38,6 @@ py_test( size = "small", srcs = ["python/ops/remote_fused_graph_ops_test.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], deps = [ ":remote_fused_graph_ops_py", "//tensorflow/core:protos_all_py", 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/ops/gru_ops.cc b/tensorflow/contrib/rnn/ops/gru_ops.cc index e91d1e8a80ed252e5f89e116fb0a325be67e3941..9c8e40851a0cc5bd7f37f94a62ecdef7248660c1 100644 --- a/tensorflow/contrib/rnn/ops/gru_ops.cc +++ b/tensorflow/contrib/rnn/ops/gru_ops.cc @@ -69,7 +69,7 @@ Element-wise dot product of a and b is represented by ab Element-wise dot product is represented by \circ Matrix multiplication is represented by * -Baises are initialized with : +Biases are initialized with : `b_ru` - constant_initializer(1.0) `b_c` - constant_initializer(0.0) 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 59d7b2466ca70fd9a4e2863b56dc920cee656079..d41fc0b3ac1cee4eacc88cb0f41df1f9ee59e7c3 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,8 +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 @@ -157,7 +155,7 @@ class RNNCellTest(test.TestCase): m.name: np.array([[0.1, 0.1]]) }) # Smoke test - self.assertAllClose(res[0], [[0.509682, 0.509682]]) + self.assertAllClose(res[0], [[0.509682, 0.509682]]) def testSRUCellWithDiffSize(self): with self.test_session() as sess: @@ -167,9 +165,10 @@ 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., 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.55255556, 0.55255556]]) @@ -188,6 +187,8 @@ class RNNCellTest(test.TestCase): ], state_is_tuple=False) self.assertEqual(cell.dtype, None) + self.assertEqual("cell-0", cell._checkpoint_dependencies[0].name) + self.assertEqual("cell-1", cell._checkpoint_dependencies[1].name) g, out_m = cell(x, m) # Layer infers the input type. self.assertEqual(cell.dtype, dtype.name) 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 57521c6a9ba0b2d66639017b09c541e270276323..de5df912921932056526e1e6dc5dbb905735f775 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py @@ -869,7 +869,7 @@ class LSTMTest(test.TestCase): num_proj = 4 max_length = 8 sequence_length = [4, 6] - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() with self.test_session(graph=ops_lib.Graph()) as sess: initializer = init_ops.random_uniform_initializer( -0.01, 0.01, seed=self._seed) @@ -934,8 +934,7 @@ class LSTMTest(test.TestCase): if in_graph_mode: self.assertAllEqual(outputs_static, outputs_dynamic) else: - self.assertAllEqual( - array_ops.stack(outputs_static).numpy(), outputs_dynamic.numpy()) + self.assertAllEqual(array_ops.stack(outputs_static), outputs_dynamic) self.assertAllEqual(np.hstack(state_static), np.hstack(state_dynamic)) @test_util.run_in_graph_and_eager_modes() @@ -946,7 +945,7 @@ class LSTMTest(test.TestCase): num_proj = 4 max_length = 8 sequence_length = [4, 6] - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() with self.test_session(graph=ops_lib.Graph()) as sess: initializer = init_ops.random_uniform_initializer( -0.01, 0.01, seed=self._seed) @@ -1022,10 +1021,9 @@ class LSTMTest(test.TestCase): if in_graph_mode: self.assertAllEqual(outputs_static, outputs_dynamic) else: - self.assertAllEqual( - array_ops.stack(outputs_static).numpy(), outputs_dynamic.numpy()) - state_static = [s.numpy() for s in nest.flatten(state_static)] - state_dynamic = [s.numpy() for s in nest.flatten(state_dynamic)] + self.assertAllEqual(array_ops.stack(outputs_static), outputs_dynamic) + state_static = nest.flatten(state_static) + state_dynamic = nest.flatten(state_dynamic) self.assertAllEqual(np.hstack(state_static), np.hstack(state_dynamic)) def _testDynamicEquivalentToStaticRNN(self, use_sequence_length): @@ -1043,7 +1041,7 @@ class LSTMTest(test.TestCase): else: sequence_length = None - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() # TODO(b/68017812): Eager ignores operation seeds, so we need to create a # single cell and reuse it across the static and dynamic RNNs. Remove this diff --git a/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py b/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py index 7957edf68cc8a1461fccfc2de93ad5250dc9fdb5..ffd24218944e150a32b1b915288ab1df90afb45c 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py @@ -54,7 +54,7 @@ def blocks_match(sess, use_peephole): initializer = init_ops.random_uniform_initializer(-0.01, 0.01, seed=19890212) with variable_scope.variable_scope("test", initializer=initializer): - # magic naming so that the cells pick up these variables and resuse them + # magic naming so that the cells pick up these variables and reuse them if use_peephole: wci = variable_scope.get_variable( "rnn/lstm_cell/w_i_diag", shape=[cell_size], dtype=dtypes.float32) 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 8a3894ef9d7042e66b52edefdf08b278dcc6c4f4..63fdd91d368d97007280871f3886e5649e6b2e86 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py @@ -455,8 +455,8 @@ class RNNCellTest(test.TestCase): self.assertAllClose(np.concatenate(res[1], axis=1), expected_state) def testAttentionCellWrapperFailures(self): - with self.assertRaisesRegexp(TypeError, - "The parameter cell is not RNNCell."): + with self.assertRaisesRegexp( + TypeError, rnn_cell_impl.ASSERT_LIKE_RNNCELL_ERROR_REGEXP): contrib_rnn_cell.AttentionCellWrapper(None, 0) num_units = 8 @@ -878,7 +878,6 @@ class RNNCellTest(test.TestCase): 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]]], dtype=np.float32) @@ -912,7 +911,6 @@ class RNNCellTest(test.TestCase): 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]] @@ -954,7 +952,6 @@ class RNNCellTest(test.TestCase): 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], @@ -1031,57 +1028,92 @@ class RNNCellTest(test.TestCase): num_units = 4 number_of_groups = 1 - with self.test_session() as sess: - with variable_scope.variable_scope( - "root1", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.ones([batch_size, num_units]) - # When number_of_groups = 1, G-LSTM is equivalent to regular LSTM - gcell = contrib_rnn_cell.GLSTMCell( - 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) - gh, gs = gcell(x, zero_state) - h, g = cell(x, zero_state) + # Try with input dimension equal to num_units or not. + for num_inputs in [num_units, num_units + number_of_groups]: + with self.test_session() as sess: + with variable_scope.variable_scope( + "root1_%d" % num_inputs, + initializer=init_ops.constant_initializer(0.5)): + x = array_ops.ones([batch_size, num_inputs]) + # When number_of_groups = 1, G-LSTM is equivalent to regular LSTM + gcell = contrib_rnn_cell.GLSTMCell( + 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) + gh, gs = gcell(x, zero_state) + h, g = cell(x, zero_state) - sess.run([variables.global_variables_initializer()]) - glstm_result = sess.run([gh, gs]) - lstm_result = sess.run([h, g]) + sess.run([variables.global_variables_initializer()]) + glstm_result = sess.run([gh, gs]) + lstm_result = sess.run([h, g]) - self.assertAllClose(glstm_result[0], lstm_result[0], 1e-5) - self.assertAllClose(glstm_result[1], lstm_result[1], 1e-5) + self.assertAllClose(glstm_result[0], lstm_result[0], 1e-5) + self.assertAllClose(glstm_result[1], lstm_result[1], 1e-5) # Test that G-LSTM subgroup act like corresponding sub-LSTMs batch_size = 2 num_units = 4 number_of_groups = 2 - with self.test_session() as sess: + # Try with num_inputs equal to or not equal to num_units. + for num_inputs in [num_units, num_units + number_of_groups]: + with self.test_session() as sess: + with variable_scope.variable_scope( + "root2_%d" % num_inputs, + initializer=init_ops.constant_initializer(0.5)): + # input for G-LSTM with 2 groups + glstm_input = array_ops.ones([batch_size, num_inputs]) + 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) + 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_inputs / 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) + h, g = cell(lstm_input, cell_zero_state) + + sess.run([variables.global_variables_initializer()]) + [gh_res, h_res] = sess.run([gh, h]) + self.assertAllClose(gh_res[:, 0:int(num_units / number_of_groups)], + h_res, 1e-5) + self.assertAllClose(gh_res[:, int(num_units / number_of_groups):], + h_res, 1e-5) + + def testGLSTMCellFailure(self): + batch_size = 2 + num_units = 4 + number_of_groups = 2 + with self.test_session(): with variable_scope.variable_scope( - "root2", initializer=init_ops.constant_initializer(0.5)): - # input for G-LSTM with 2 groups - glstm_input = array_ops.ones([batch_size, num_units]) + "glstm_failure", initializer=init_ops.constant_initializer(0.5)): 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) - 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) - h, g = cell(lstm_input, cell_zero_state) + # Try an input with statically-unknown innermost dimension. + glstm_input = array_ops.placeholder( + dtypes.float32, shape=[batch_size, None]) + with self.assertRaisesRegexp(ValueError, + "input size must be statically known"): + gcell(glstm_input, gcell_zero_state) - sess.run([variables.global_variables_initializer()]) - [gh_res, h_res] = sess.run([gh, h]) - self.assertAllClose(gh_res[:, 0:int(num_units / number_of_groups)], - h_res, 1e-5) - self.assertAllClose(gh_res[:, int(num_units / number_of_groups):], - h_res, 1e-5) + # Try an input whose innermost dimension isn't divisible into groups. + glstm_input = array_ops.placeholder( + dtypes.float32, shape=[batch_size, 3]) + with self.assertRaisesRegexp( + ValueError, + r"input size \(3\) must be divisible by number_of_groups \(2\)"): + gcell(glstm_input, gcell_zero_state) class LayerNormBasicLSTMCellTest(test.TestCase): @@ -1168,7 +1200,7 @@ class LayerNormBasicLSTMCellTest(test.TestCase): h1 = array_ops.zeros([1, 2]) state1 = rnn_cell.LSTMStateTuple(c1, h1) state = (state0, state1) - single_cell = lambda: contrib_rnn_cell.LayerNormBasicLSTMCell(2, layer_norm=False) + single_cell = lambda: contrib_rnn_cell.LayerNormBasicLSTMCell(2, layer_norm=False) # pylint: disable=line-too-long cell = rnn_cell.MultiRNNCell([single_cell() for _ in range(2)]) g, out_m = cell(x, state) sess.run([variables.global_variables_initializer()]) @@ -1200,7 +1232,7 @@ class LayerNormBasicLSTMCellTest(test.TestCase): self.assertAllClose(expected_state1_h, actual_state1_h, 1e-5) with variable_scope.variable_scope( - "other", initializer=init_ops.constant_initializer(0.5)) as vs: + "other", initializer=init_ops.constant_initializer(0.5)): x = array_ops.zeros( [1, 3]) # Test BasicLSTMCell with input_size != num_units. c = array_ops.zeros([1, 2]) @@ -1549,98 +1581,7 @@ 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) - - -class WeightNormLSTMCellTest(test.TestCase): - """Compared cell output with pre-calculated values.""" - - def _cell_output(self, cell): - """Calculate cell output""" + """Calculates cell output.""" with self.test_session() as sess: init = init_ops.constant_initializer(0.5) @@ -1667,7 +1608,7 @@ class WeightNormLSTMCellTest(test.TestCase): return actual_state_c, actual_state_h def testBasicCell(self): - """Tests cell w/o peepholes and w/o normalisation""" + """Tests cell w/o peepholes and w/o normalisation.""" def cell(): return contrib_rnn_cell.WeightNormLSTMCell(2, @@ -1683,7 +1624,7 @@ class WeightNormLSTMCellTest(test.TestCase): self.assertAllClose(expected_h, actual_h, 1e-5) def testNonbasicCell(self): - """Tests cell with peepholes and w/o normalisation""" + """Tests cell with peepholes and w/o normalisation.""" def cell(): return contrib_rnn_cell.WeightNormLSTMCell(2, @@ -1698,9 +1639,8 @@ class WeightNormLSTMCellTest(test.TestCase): 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""" + """Tests cell w/o peepholes and with normalisation.""" def cell(): return contrib_rnn_cell.WeightNormLSTMCell(2, @@ -1716,7 +1656,7 @@ class WeightNormLSTMCellTest(test.TestCase): self.assertAllClose(expected_h, actual_h, 1e-5) def testNonBasicCellWithNorm(self): - """Tests cell with peepholes and with normalisation""" + """Tests cell with peepholes and with normalisation.""" def cell(): return contrib_rnn_cell.WeightNormLSTMCell(2, diff --git a/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py b/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py index 8109ebc718353300f94536c5d7ae3332da584a1d..645f82624bf67b96ffc8520289b293b45f0e69e2 100644 --- a/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py @@ -40,7 +40,6 @@ from tensorflow.python.util import nest # pylint: disable=protected-access,invalid-name RNNCell = rnn_cell_impl.RNNCell -_like_rnncell = rnn_cell_impl._like_rnncell _WEIGHTS_VARIABLE_NAME = rnn_cell_impl._WEIGHTS_VARIABLE_NAME _BIAS_VARIABLE_NAME = rnn_cell_impl._BIAS_VARIABLE_NAME # pylint: enable=protected-access,invalid-name @@ -221,8 +220,7 @@ class EmbeddingWrapper(RNNCell): ValueError: if embedding_classes is not positive. """ super(EmbeddingWrapper, self).__init__(_reuse=reuse) - if not _like_rnncell(cell): - raise TypeError("The parameter cell is not RNNCell.") + rnn_cell_impl.assert_like_rnncell("cell", cell) if embedding_classes <= 0 or embedding_size <= 0: raise ValueError("Both embedding_classes and embedding_size must be > 0: " "%d, %d." % (embedding_classes, embedding_size)) @@ -301,8 +299,7 @@ class InputProjectionWrapper(RNNCell): super(InputProjectionWrapper, self).__init__(_reuse=reuse) if input_size is not None: logging.warn("%s: The input_size parameter is deprecated.", self) - if not _like_rnncell(cell): - raise TypeError("The parameter cell is not RNNCell.") + rnn_cell_impl.assert_like_rnncell("cell", cell) self._cell = cell self._num_proj = num_proj self._activation = activation @@ -356,8 +353,7 @@ class OutputProjectionWrapper(RNNCell): ValueError: if output_size is not positive. """ super(OutputProjectionWrapper, self).__init__(_reuse=reuse) - if not _like_rnncell(cell): - raise TypeError("The parameter cell is not RNNCell.") + rnn_cell_impl.assert_like_rnncell("cell", cell) if output_size < 1: raise ValueError("Parameter output_size must be > 0: %d." % output_size) self._cell = cell 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..9e61fc54d10c1b75786450060e428c73974760a7 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 @@ -480,8 +480,7 @@ class LSTMBlockWrapper(base_layer.Layer): """Run this LSTM on inputs, starting from the given state. Args: - inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]` - or a list of `time_len` tensors of shape `[batch_size, input_size]`. + inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]`. initial_state: a tuple `(initial_cell_state, initial_output)` with tensors of shape `[batch_size, self._num_units]`. If this is not provided, the cell is expected to create a zero initial state of type `dtype`. @@ -572,9 +571,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 5fee2e93e4e575f647b3e6f132a57c57056726ad..2f6ae9f3678e58dae67bf777991641b10e42ef94 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/rnn_cell.py @@ -32,12 +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.ops import nn_impl from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import nest @@ -424,8 +424,9 @@ 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], [-1, self._num_units]) @@ -533,7 +534,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): initializer: (optional) The initializer to use for the weight and projection matrices, default None. num_unit_shards: (optional) int, default 1, How to split the weight - matrix. If > 1,the weight matrix is stored across num_unit_shards. + matrix. If > 1, the weight matrix is stored across num_unit_shards. forget_bias: (optional) float, default 1.0, The initial bias of the forget gates, used to reduce the scale of forgetting at the beginning of the training. @@ -992,7 +993,7 @@ class BidirectionalGridLSTMCell(GridLSTMCell): initializer: (optional) The initializer to use for the weight and projection matrices, default None. num_unit_shards: (optional) int, default 1, How to split the weight - matrix. If > 1,the weight matrix is stored across num_unit_shards. + matrix. If > 1, the weight matrix is stored across num_unit_shards. forget_bias: (optional) float, default 1.0, The initial bias of the forget gates, used to reduce the scale of forgetting at the beginning of the training. @@ -1142,8 +1143,7 @@ class AttentionCellWrapper(rnn_cell_impl.RNNCell): `state_is_tuple` is `False` or if attn_length is zero or less. """ super(AttentionCellWrapper, self).__init__(_reuse=reuse) - if not rnn_cell_impl._like_rnncell(cell): # pylint: disable=protected-access - raise TypeError("The parameter cell is not RNNCell.") + rnn_cell_impl.assert_like_rnncell("cell", cell) if nest.is_sequence(cell.state_size) and not state_is_tuple: raise ValueError( "Cell returns tuple of states, but the flag " @@ -2058,16 +2058,19 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): initializers=None, name="conv_lstm_cell"): """Construct ConvLSTMCell. + Args: conv_ndims: Convolution dimensionality (1, 2 or 3). input_shape: Shape of the input as int tuple, excluding the batch size. output_channels: int, number of output channels of the conv LSTM. kernel_shape: Shape of kernel as in tuple (of size 1,2 or 3). - use_bias: Use bias in convolutions. + use_bias: (bool) Use bias in convolutions. skip_connection: If set to `True`, concatenate the input to the - output of the conv LSTM. Default: `False`. + output of the conv LSTM. Default: `False`. forget_bias: Forget bias. + initializers: Unused. name: Name of the module. + Raises: ValueError: If `skip_connection` is `True` and stride is different from 1 or if `input_shape` is incompatible with `conv_ndims`. @@ -2130,7 +2133,7 @@ class Conv1DLSTMCell(ConvLSTMCell): def __init__(self, name="conv_1d_lstm_cell", **kwargs): """Construct Conv1DLSTM. See `ConvLSTMCell` for more details.""" - super(Conv1DLSTMCell, self).__init__(conv_ndims=1, **kwargs) + super(Conv1DLSTMCell, self).__init__(conv_ndims=1, name=name, **kwargs) class Conv2DLSTMCell(ConvLSTMCell): @@ -2141,7 +2144,7 @@ class Conv2DLSTMCell(ConvLSTMCell): def __init__(self, name="conv_2d_lstm_cell", **kwargs): """Construct Conv2DLSTM. See `ConvLSTMCell` for more details.""" - super(Conv2DLSTMCell, self).__init__(conv_ndims=2, **kwargs) + super(Conv2DLSTMCell, self).__init__(conv_ndims=2, name=name, **kwargs) class Conv3DLSTMCell(ConvLSTMCell): @@ -2152,19 +2155,23 @@ class Conv3DLSTMCell(ConvLSTMCell): def __init__(self, name="conv_3d_lstm_cell", **kwargs): """Construct Conv3DLSTM. See `ConvLSTMCell` for more details.""" - super(Conv3DLSTMCell, self).__init__(conv_ndims=3, **kwargs) + super(Conv3DLSTMCell, self).__init__(conv_ndims=3, name=name, **kwargs) def _conv(args, filter_size, num_features, bias, bias_start=0.0): - """convolution: + """Convolution. + Args: args: a Tensor or a list of Tensors of dimension 3D, 4D or 5D, batch x n, Tensors. filter_size: int tuple of filter height and width. num_features: int, number of features. + bias: Whether to use biases in the convolution layer. bias_start: starting value to initialize the bias; 0 by default. + Returns: A 3D, 4D, or 5D Tensor with shape [batch ... num_features] + Raises: ValueError: if some of the arguments has unspecified or wrong shape. """ @@ -2224,6 +2231,13 @@ class GLSTMCell(rnn_cell_impl.RNNCell): O. Kuchaiev and B. Ginsburg "Factorization Tricks for LSTM Networks", ICLR 2017 workshop. + + In brief, a G-LSTM cell consists of one LSTM sub-cell per group, where each + sub-cell operates on an evenly-sized sub-vector of the input and produces an + evenly-sized sub-vector of the output. For example, a G-LSTM cell with 128 + units and 4 groups consists of 4 LSTMs sub-cells with 32 units each. If that + G-LSTM cell is fed a 200-dim input, then each sub-cell receives a 50-dim part + of the input and produces a 32-dim part of the output. """ def __init__(self, @@ -2285,7 +2299,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 @@ -2297,7 +2311,7 @@ class GLSTMCell(rnn_cell_impl.RNNCell): return self._output_size def _get_input_for_group(self, inputs, group_id, group_size): - """Slices inputs into groups to prepare for processing by cell's groups + """Slices inputs into groups to prepare for processing by cell's groups. Args: inputs: cell input or it's previous state, @@ -2319,9 +2333,12 @@ class GLSTMCell(rnn_cell_impl.RNNCell): """Run one step of G-LSTM. Args: - inputs: input Tensor, 2D, [batch x num_units]. - state: this must be a tuple of state Tensors, both `2-D`, - with column sizes `c_state` and `m_state`. + inputs: input Tensor, 2D, [batch x num_inputs]. num_inputs must be + statically-known and evenly divisible into groups. The innermost + vectors of the inputs are split into evenly-sized sub-vectors and fed + into the per-group LSTM sub-cells. + state: this must be a tuple of state Tensors, both `2-D`, with column + sizes `c_state` and `m_state`. Returns: A tuple containing: @@ -2336,11 +2353,24 @@ class GLSTMCell(rnn_cell_impl.RNNCell): Raises: ValueError: If input size cannot be inferred from inputs via - static shape inference. + static shape inference, or if the input shape is incompatible + with the number of groups. """ (c_prev, m_prev) = state self._batch_size = inputs.shape[0].value or array_ops.shape(inputs)[0] + + # If the input size is statically-known, calculate and validate its group + # size. Otherwise, use the output group size. + input_size = inputs.shape[1].value + if input_size is None: + raise ValueError("input size must be statically known") + if input_size % self._number_of_groups != 0: + raise ValueError( + "input size (%d) must be divisible by number_of_groups (%d)" % + (input_size, self._number_of_groups)) + input_group_size = int(input_size / self._number_of_groups) + dtype = inputs.dtype scope = vs.get_variable_scope() with vs.variable_scope(scope, initializer=self._initializer): @@ -2353,15 +2383,16 @@ class GLSTMCell(rnn_cell_impl.RNNCell): 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(inputs, group_id, input_group_size), 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 + 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) @@ -2680,8 +2711,8 @@ class LayerNormLSTMCell(rnn_cell_impl.RNNCell): return m, new_state -class SRUCell(rnn_cell_impl._LayerRNNCell): - """SRU, Simple Recurrent Unit +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). @@ -2729,12 +2760,13 @@ class SRUCell(rnn_cell_impl._LayerRNNCell): input_depth = inputs_shape[1].value + # pylint: disable=protected-access self._kernel = self.add_variable( rnn_cell_impl._WEIGHTS_VARIABLE_NAME, shape=[input_depth, 4 * self._num_units]) - + # pylint: enable=protected-access self._bias = self.add_variable( - rnn_cell_impl._BIAS_VARIABLE_NAME, + rnn_cell_impl._BIAS_VARIABLE_NAME, # pylint: disable=protected-access shape=[2 * self._num_units], initializer=init_ops.constant_initializer(0.0, dtype=self.dtype)) @@ -2743,7 +2775,7 @@ class SRUCell(rnn_cell_impl._LayerRNNCell): def call(self, inputs, state): """Simple recurrent unit (SRU) with num_units cells.""" - U = math_ops.matmul(inputs, self._kernel) + U = math_ops.matmul(inputs, self._kernel) # pylint: disable=invalid-name x_bar, f_intermediate, r_intermediate, x_tx = array_ops.split( value=U, num_or_size_splits=4, axis=1) @@ -2873,6 +2905,7 @@ class WeightNormLSTMCell(rnn_cell_impl.RNNCell): Args: args: a 2D Tensor or a list of 2D, batch x n, Tensors. output_size: int, second dimension of W[i]. + norm: bool, whether to normalize the weights. bias: boolean, whether to add a bias term or not. bias_initializer: starting value to initialize the bias (default is all zeros). diff --git a/tensorflow/contrib/saved_model/BUILD b/tensorflow/contrib/saved_model/BUILD index 245fe07f2bcdaddb2bc47c0e1234dc1f19bd85e3..b10757df47a4c38b98692c3d1995b3f040f0b478 100644 --- a/tensorflow/contrib/saved_model/BUILD +++ b/tensorflow/contrib/saved_model/BUILD @@ -53,6 +53,7 @@ py_test( size = "small", srcs = ["python/saved_model/reader_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows visibility = ["//visibility:private"], deps = [ ":saved_model_py", 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/python/kernel_tests/attention_wrapper_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py index b427dff88b2d586ccf8c512bb498cdaf879ac781..c4139dde492a3c88ec3d5f19973314ac3fd81fb3 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py @@ -222,6 +222,9 @@ class AttentionWrapperTest(test.TestCase): self.assertEqual( (None, batch_size, None), tuple(state_alignment_history.get_shape().as_list())) + nest.assert_same_structure( + cell.state_size, + cell.zero_state(batch_size, dtypes.float32)) # Remove the history from final_state for purposes of the # remainder of the tests. final_state = final_state._replace(alignment_history=()) # pylint: disable=protected-access 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 926554031775202d7f7d9018cf6ae4efb34fe96b..178328619f087789df040489cd150ba018cc8d14 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 @@ -27,6 +27,7 @@ from tensorflow.contrib.seq2seq.python.ops import beam_search_ops from tensorflow.contrib.seq2seq.python.ops import decoder 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.layers import core as layers_core from tensorflow.python.ops import array_ops @@ -70,6 +71,98 @@ class TestGatherTree(test.TestCase): self.assertAllEqual(expected_result, res_) + def _test_gather_tree_from_array(self, + depth_ndims=0, + merged_batch_beam=False): + array = np.array( + [[[1, 2, 3], [4, 5, 6], [7, 8, 9], [0, 0, 0]], + [[2, 3, 4], [5, 6, 7], [8, 9, 10], [11, 12, 0]]]).transpose([1, 0, 2]) + parent_ids = np.array( + [[[0, 0, 0], [0, 1, 1], [2, 1, 2], [-1, -1, -1]], + [[0, 0, 0], [1, 1, 0], [2, 0, 1], [0, 1, 0]]]).transpose([1, 0, 2]) + expected_array = np.array( + [[[2, 2, 2], [6, 5, 6], [7, 8, 9], [0, 0, 0]], + [[2, 3, 2], [7, 5, 7], [8, 9, 8], [11, 12, 0]]]).transpose([1, 0, 2]) + sequence_length = [[3, 3, 3], [4, 4, 3]] + + array = ops.convert_to_tensor( + array, dtype=dtypes.float32) + parent_ids = ops.convert_to_tensor( + parent_ids, dtype=dtypes.int32) + expected_array = ops.convert_to_tensor( + expected_array, dtype=dtypes.float32) + + max_time = array_ops.shape(array)[0] + batch_size = array_ops.shape(array)[1] + beam_width = array_ops.shape(array)[2] + + def _tile_in_depth(tensor): + # Generate higher rank tensors by concatenating tensor and tensor + 1. + for _ in range(depth_ndims): + tensor = array_ops.stack([tensor, tensor + 1], -1) + return tensor + + if merged_batch_beam: + array = array_ops.reshape( + array, [max_time, batch_size * beam_width]) + expected_array = array_ops.reshape( + expected_array, [max_time, batch_size * beam_width]) + + if depth_ndims > 0: + array = _tile_in_depth(array) + expected_array = _tile_in_depth(expected_array) + + sorted_array = beam_search_decoder.gather_tree_from_array( + array, parent_ids, sequence_length) + + with self.test_session() as sess: + sorted_array = sess.run(sorted_array) + expected_array = sess.run(expected_array) + self.assertAllEqual(expected_array, sorted_array) + + def test_gather_tree_from_array_scalar(self): + self._test_gather_tree_from_array() + + def test_gather_tree_from_array_1d(self): + self._test_gather_tree_from_array(depth_ndims=1) + + def test_gather_tree_from_array_1d_with_merged_batch_beam(self): + self._test_gather_tree_from_array(depth_ndims=1, merged_batch_beam=True) + + def test_gather_tree_from_array_2d(self): + self._test_gather_tree_from_array(depth_ndims=2) + + +class TestArrayShapeChecks(test.TestCase): + + def _test_array_shape_dynamic_checks(self, static_shape, dynamic_shape, + batch_size, beam_width, is_valid=True): + t = array_ops.placeholder_with_default( + np.random.randn(*static_shape).astype(np.float32), + shape=dynamic_shape) + + batch_size = array_ops.constant(batch_size) + check_op = beam_search_decoder._check_batch_beam(t, batch_size, beam_width) # pylint: disable=protected-access + + with self.test_session() as sess: + if is_valid: + sess.run(check_op) + else: + with self.assertRaises(errors.InvalidArgumentError): + sess.run(check_op) + + def test_array_shape_dynamic_checks(self): + self._test_array_shape_dynamic_checks( + (8, 4, 5, 10), (None, None, 5, 10), 4, 5, is_valid=True) + self._test_array_shape_dynamic_checks( + (8, 20, 10), (None, None, 10), 4, 5, is_valid=True) + self._test_array_shape_dynamic_checks( + (8, 21, 10), (None, None, 10), 4, 5, is_valid=False) + self._test_array_shape_dynamic_checks( + (8, 4, 6, 10), (None, None, None, 10), 4, 5, is_valid=False) + self._test_array_shape_dynamic_checks( + (8, 4), (None, None), 4, 5, is_valid=False) + class TestEosMasking(test.TestCase): """Tests EOS masking used in beam search.""" @@ -319,7 +412,8 @@ class TestLargeBeamStep(test.TestCase): class BeamSearchDecoderTest(test.TestCase): - def _testDynamicDecodeRNN(self, time_major, has_attention): + def _testDynamicDecodeRNN(self, time_major, has_attention, + with_alignment_history=False): encoder_sequence_length = np.array([3, 2, 3, 1, 1]) decoder_sequence_length = np.array([2, 0, 1, 2, 3]) batch_size = 5 @@ -359,7 +453,7 @@ class BeamSearchDecoderTest(test.TestCase): cell=cell, attention_mechanism=attention_mechanism, attention_layer_size=attention_depth, - alignment_history=False) + alignment_history=with_alignment_history) cell_state = cell.zero_state( dtype=dtypes.float32, batch_size=batch_size_tensor * beam_width) if has_attention: @@ -420,6 +514,12 @@ class BeamSearchDecoderTest(test.TestCase): def testDynamicDecodeRNNBatchMajorYesAttention(self): self._testDynamicDecodeRNN(time_major=False, has_attention=True) + def testDynamicDecodeRNNBatchMajorYesAttentionWithAlignmentHistory(self): + self._testDynamicDecodeRNN( + time_major=False, + has_attention=True, + with_alignment_history=True) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py index 95dea312f3a4e77176a4bc4af290ad48c078deda..9ff8a343f124193b56e3bf46624ab1db927301a0 100644 --- a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py +++ b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py @@ -331,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]) @@ -924,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, @@ -1153,9 +1152,7 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): is a list, and its length does not match that of `attention_layer_size`. """ super(AttentionWrapper, self).__init__(name=name) - if not rnn_cell_impl._like_rnncell(cell): # pylint: disable=protected-access - raise TypeError( - "cell must be an RNNCell, saw type: %s" % type(cell).__name__) + rnn_cell_impl.assert_like_rnncell("cell", cell) if isinstance(attention_mechanism, (list, tuple)): self._is_multi = True attention_mechanisms = attention_mechanism @@ -1281,7 +1278,8 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): attention_state=self._item_or_tuple( a.state_size for a in self._attention_mechanisms), alignment_history=self._item_or_tuple( - () for _ in self._attention_mechanisms)) # sometimes a TensorArray + a.alignments_size if self._alignment_history else () + for a in self._attention_mechanisms)) # sometimes a TensorArray def zero_state(self, batch_size, dtype): """Return an initial (zero) state tuple for this `AttentionWrapper`. @@ -1321,22 +1319,26 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): cell_state = nest.map_structure( lambda s: array_ops.identity(s, name="checked_cell_state"), cell_state) + initial_alignments = [ + attention_mechanism.initial_alignments(batch_size, dtype) + for attention_mechanism in self._attention_mechanisms] return AttentionWrapperState( cell_state=cell_state, time=array_ops.zeros([], dtype=dtypes.int32), attention=_zero_state_tensors(self._attention_layer_size, batch_size, dtype), - alignments=self._item_or_tuple( - attention_mechanism.initial_alignments(batch_size, dtype) - for attention_mechanism in self._attention_mechanisms), + alignments=self._item_or_tuple(initial_alignments), attention_state=self._item_or_tuple( attention_mechanism.initial_state(batch_size, dtype) for attention_mechanism in self._attention_mechanisms), alignment_history=self._item_or_tuple( - tensor_array_ops.TensorArray(dtype=dtype, size=0, - dynamic_size=True) + tensor_array_ops.TensorArray( + dtype, + size=0, + dynamic_size=True, + element_shape=alignment.shape) if self._alignment_history else () - for _ in self._attention_mechanisms)) + for alignment in initial_alignments)) def call(self, inputs, state): """Perform a step of attention-wrapped RNN. diff --git a/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py b/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py index ed226239b860e2250072a28a5538b816642ec54b..7eb95e5a70de985dca0d4b565ba03bdf454b6161 100644 --- a/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py @@ -59,8 +59,7 @@ class BasicDecoder(decoder.Decoder): Raises: TypeError: if `cell`, `helper` or `output_layer` have an incorrect type. """ - if not rnn_cell_impl._like_rnncell(cell): # pylint: disable=protected-access - raise TypeError("cell must be an RNNCell, received: %s" % type(cell)) + rnn_cell_impl.assert_like_rnncell("cell", cell) if not isinstance(helper, helper_py.Helper): raise TypeError("helper must be a Helper, received: %s" % type(helper)) if (output_layer is not None diff --git a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py index d6184d61095f727f9dcab56fe59e2601868c1624..a26107b0d71497e2d21d735f02e2556873709cda 100644 --- a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py @@ -35,6 +35,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.platform import tf_logging from tensorflow.python.util import nest __all__ = [ @@ -121,14 +122,114 @@ def tile_batch(t, multiplier, name=None): return nest.map_structure(lambda t_: _tile_batch(t_, multiplier), t) +def gather_tree_from_array(t, parent_ids, sequence_length): + """Calculates the full beams for `TensorArray`s. + + Args: + t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of + shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]` + where `s` is the depth shape. + parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`. + sequence_length: The sequence length of shape `[batch_size, beam_width]`. + + Returns: + A `Tensor` which is a stacked `TensorArray` of the same size and type as + `t` and where beams are sorted in each `Tensor` according to `parent_ids`. + """ + max_time = parent_ids.shape[0].value or array_ops.shape(parent_ids)[0] + batch_size = parent_ids.shape[1].value or array_ops.shape(parent_ids)[1] + beam_width = parent_ids.shape[2].value or array_ops.shape(parent_ids)[2] + + # Generate beam ids that will be reordered by gather_tree. + beam_ids = array_ops.expand_dims( + array_ops.expand_dims(math_ops.range(beam_width), 0), 0) + beam_ids = array_ops.tile(beam_ids, [max_time, batch_size, 1]) + + mask = array_ops.sequence_mask( + sequence_length, maxlen=max_time, dtype=dtypes.int32) + mask = array_ops.transpose(mask, perm=[2, 0, 1]) + + # Use beam_width + 1 to mark the end of beam. + masked_beam_ids = (beam_ids * mask) + (1 - mask) * (beam_width + 1) + + max_sequence_lengths = math_ops.to_int32( + math_ops.reduce_max(sequence_length, axis=1)) + sorted_beam_ids = beam_search_ops.gather_tree( + step_ids=masked_beam_ids, + parent_ids=parent_ids, + max_sequence_lengths=max_sequence_lengths, + end_token=beam_width + 1) + + # For out of range steps, simply copy the same beam. + sorted_beam_ids = array_ops.where( + math_ops.cast(mask, dtypes.bool), x=sorted_beam_ids, y=beam_ids) + + # Generate indices for gather_nd. + time_ind = array_ops.tile(array_ops.reshape( + math_ops.range(max_time), [-1, 1, 1]), [1, batch_size, beam_width]) + batch_ind = array_ops.tile(array_ops.reshape( + math_ops.range(batch_size), [-1, 1, 1]), [1, max_time, beam_width]) + batch_ind = array_ops.transpose(batch_ind, perm=[1, 0, 2]) + indices = array_ops.stack([time_ind, batch_ind, sorted_beam_ids], -1) + + # Gather from a tensor with collapsed additional dimensions. + gather_from = t + final_shape = array_ops.shape(gather_from) + gather_from = array_ops.reshape( + gather_from, [max_time, batch_size, beam_width, -1]) + ordered = array_ops.gather_nd(gather_from, indices) + ordered = array_ops.reshape(ordered, final_shape) + + return ordered + + def _check_maybe(t): - if isinstance(t, tensor_array_ops.TensorArray): - raise TypeError( - "TensorArray state is not supported by BeamSearchDecoder: %s" % t.name) if t.shape.ndims is None: raise ValueError( "Expected tensor (%s) to have known rank, but ndims == None." % t) +def _check_static_batch_beam_maybe(shape, batch_size, beam_width): + """Raises an exception if dimensions are known statically and can not be + reshaped to [batch_size, beam_size, -1]. + """ + reshaped_shape = tensor_shape.TensorShape([batch_size, beam_width, None]) + if (batch_size is not None and shape[0].value is not None + and (shape[0] != batch_size * beam_width + or (shape.ndims >= 2 and shape[1].value is not None + and (shape[0] != batch_size or shape[1] != beam_width)))): + tf_logging.warn("TensorArray reordering expects elements to be " + "reshapable to %s which is incompatible with the " + "current shape %s. Consider setting " + "reorder_tensor_arrays to False to disable TensorArray " + "reordering during the beam search." + % (reshaped_shape, shape)) + return False + return True + +def _check_batch_beam(t, batch_size, beam_width): + """Returns an Assert operation checking that the elements of the stacked + TensorArray can be reshaped to [batch_size, beam_size, -1]. At this point, + the TensorArray elements have a known rank of at least 1. + """ + error_message = ("TensorArray reordering expects elements to be " + "reshapable to [batch_size, beam_size, -1] which is " + "incompatible with the dynamic shape of %s elements. " + "Consider setting reorder_tensor_arrays to False to disable " + "TensorArray reordering during the beam search." + % (t.name)) + rank = t.shape.ndims + shape = array_ops.shape(t) + if rank == 2: + condition = math_ops.equal(shape[1], batch_size * beam_width) + else: + condition = math_ops.logical_or( + math_ops.equal(shape[1], batch_size * beam_width), + math_ops.logical_and( + math_ops.equal(shape[1], batch_size), + math_ops.equal(shape[2], beam_width))) + return control_flow_ops.Assert(condition, [error_message]) + + class BeamSearchDecoder(decoder.Decoder): """BeamSearch sampling decoder. @@ -173,7 +274,8 @@ class BeamSearchDecoder(decoder.Decoder): initial_state, beam_width, output_layer=None, - length_penalty_weight=0.0): + length_penalty_weight=0.0, + reorder_tensor_arrays=True): """Initialize the BeamSearchDecoder. Args: @@ -188,6 +290,12 @@ class BeamSearchDecoder(decoder.Decoder): `tf.layers.Dense`. Optional layer to apply to the RNN output prior to storing the result or sampling. length_penalty_weight: Float weight to penalize length. Disabled with 0.0. + reorder_tensor_arrays: If `True`, `TensorArray`s' elements within the cell + state will be reordered according to the beam search path. If the + `TensorArray` can be reordered, the stacked form will be returned. + Otherwise, the `TensorArray` will be returned as is. Set this flag to + `False` if the cell state contains `TensorArray`s that are not amenable + to reordering. Raises: TypeError: if `cell` is not an instance of `RNNCell`, @@ -195,14 +303,14 @@ class BeamSearchDecoder(decoder.Decoder): ValueError: If `start_tokens` is not a vector or `end_token` is not a scalar. """ - if not rnn_cell_impl._like_rnncell(cell): # pylint: disable=protected-access - raise TypeError("cell must be an RNNCell, received: %s" % type(cell)) + rnn_cell_impl.assert_like_rnncell("cell", cell) # pylint: disable=protected-access 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 self._output_layer = output_layer + self._reorder_tensor_arrays = reorder_tensor_arrays if callable(embedding): self._embedding_fn = embedding @@ -300,12 +408,13 @@ class BeamSearchDecoder(decoder.Decoder): """ finished, start_inputs = self._finished, self._start_inputs + dtype = nest.flatten(self._initial_cell_state)[0].dtype 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) + on_value=ops.convert_to_tensor(0.0, dtype=dtype), + off_value=ops.convert_to_tensor(-np.Inf, dtype=dtype), + dtype=dtype) initial_state = BeamSearchDecoderState( cell_state=self._initial_cell_state, @@ -342,6 +451,11 @@ class BeamSearchDecoder(decoder.Decoder): outputs.parent_ids, max_sequence_lengths=max_sequence_lengths, end_token=self._end_token) + if self._reorder_tensor_arrays: + final_state = final_state._replace(cell_state=nest.map_structure( + lambda t: self._maybe_sort_array_beams( + t, outputs.parent_ids, final_state.lengths), + final_state.cell_state)) outputs = FinalBeamSearchDecoderOutput( beam_search_decoder_output=outputs, predicted_ids=predicted_ids) return outputs, final_state @@ -432,9 +546,10 @@ class BeamSearchDecoder(decoder.Decoder): returned unchanged. Raises: - TypeError: If `t` is an instance of `TensorArray`. ValueError: If the rank of `t` is not statically known. """ + if isinstance(t, tensor_array_ops.TensorArray): + return t _check_maybe(t) if t.shape.ndims >= 1: return self._split_batch_beams(t, s) @@ -455,15 +570,55 @@ class BeamSearchDecoder(decoder.Decoder): A reshaped version of t with shape `[batch_size, beam_width] + s`. Raises: - TypeError: If `t` is an instance of `TensorArray`. ValueError: If the rank of `t` is not statically known. """ + if isinstance(t, tensor_array_ops.TensorArray): + return t _check_maybe(t) if t.shape.ndims >= 2: return self._merge_batch_beams(t, s) else: return t + def _maybe_sort_array_beams(self, t, parent_ids, sequence_length): + """Maybe sorts beams within a `TensorArray`. + + Args: + t: A `TensorArray` of size `max_time` that contains `Tensor`s of shape + `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]` where + `s` is the depth shape. + parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`. + sequence_length: The sequence length of shape `[batch_size, beam_width]`. + + Returns: + A `TensorArray` where beams are sorted in each `Tensor` or `t` itself if + it is not a `TensorArray` or does not meet shape requirements. + """ + if not isinstance(t, tensor_array_ops.TensorArray): + return t + # pylint: disable=protected-access + if (not t._infer_shape or not t._element_shape + or t._element_shape[0].ndims is None + or t._element_shape[0].ndims < 1): + shape = ( + t._element_shape[0] if t._infer_shape and t._element_shape + else tensor_shape.TensorShape(None)) + tf_logging.warn("The TensorArray %s in the cell state is not amenable to " + "sorting based on the beam search result. For a " + "TensorArray to be sorted, its elements shape must be " + "defined and have at least a rank of 1, but saw shape: %s" + % (t.handle.name, shape)) + return t + shape = t._element_shape[0] + # pylint: enable=protected-access + if not _check_static_batch_beam_maybe( + shape, tensor_util.constant_value(self._batch_size), self._beam_width): + return t + t = t.stack() + with ops.control_dependencies( + [_check_batch_beam(t, self._batch_size, self._beam_width)]): + return gather_tree_from_array(t, parent_ids, sequence_length) + def step(self, time, inputs, state, name=None): """Perform a decoding step. @@ -570,7 +725,6 @@ 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 = array_ops.reshape(scores, [batch_size, -1]) # Pick the next beams according to the specified successors function @@ -724,7 +878,7 @@ def _mask_probs(probs, eos_token, finished): eos_token, vocab_size, dtype=probs.dtype, - on_value=0., + on_value=ops.convert_to_tensor(0., dtype=probs.dtype), off_value=probs.dtype.min) finished_probs = array_ops.tile( array_ops.reshape(finished_row, [1, 1, -1]), @@ -759,6 +913,8 @@ def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size, output: Gathered tensor of shape tf.shape(gather_from)[:1+len(gather_shape)] or the original tensor if its dimensions are too small. """ + if isinstance(gather_from, tensor_array_ops.TensorArray): + return gather_from _check_maybe(gather_from) if gather_from.shape.ndims >= len(gather_shape): return _tensor_gather_helper( diff --git a/tensorflow/contrib/session_bundle/BUILD b/tensorflow/contrib/session_bundle/BUILD index 67011c8fef6c4f54db2626ffe7ae1299bddbb352..3ad88a8a22966a072e0c43d6c41369333aba5cbe 100644 --- a/tensorflow/contrib/session_bundle/BUILD +++ b/tensorflow/contrib/session_bundle/BUILD @@ -165,6 +165,7 @@ py_test( name = "gc_test", srcs = ["gc_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows visibility = ["//visibility:private"], deps = [ ":gc", diff --git a/tensorflow/contrib/session_bundle/bundle_shim.py b/tensorflow/contrib/session_bundle/bundle_shim.py index 3149875e41f6f77b3bcbc0ab1a150cfdc59ad2ba..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") == legacy_constants.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") == legacy_constants.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") != legacy_constants.GENERIC_SIGNATURE or - output_signature.WhichOneof("type") != legacy_constants.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/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/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/mfcc_ops.py b/tensorflow/contrib/signal/python/ops/mfcc_ops.py index 6cef95f742515709f0f41632358c2d8663daed2c..4e842f7f10ae07448cc07e5f636ae80a820e656f 100644 --- a/tensorflow/contrib/signal/python/ops/mfcc_ops.py +++ b/tensorflow/contrib/signal/python/ops/mfcc_ops.py @@ -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 c7a54cb9a2e9535efbdc179f1463cef379ebb1f9..40f484fd78302163ba36142dec057478fe899189 100644 --- a/tensorflow/contrib/slim/README.md +++ b/tensorflow/contrib/slim/README.md @@ -94,7 +94,7 @@ of thin wrapper functions in [variables.py](https://www.tensorflow.org/code/tensorflow/contrib/framework/python/ops/variables.py) which allow callers to easily define variables. -For example, to create a `weight` variable, initialize it using a truncated +For example, to create a `weights` variable, initialize it using a truncated normal distribution, regularize it with an `l2_loss` and place it on the `CPU`, one need only declare the following: @@ -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: diff --git a/tensorflow/contrib/slim/python/slim/data/BUILD b/tensorflow/contrib/slim/python/slim/data/BUILD index 5daabbd62e7e63608a7a86a8b7fb0bc0d570b28b..7aa1684839184900547e3304358315246ac2b140 100644 --- a/tensorflow/contrib/slim/python/slim/data/BUILD +++ b/tensorflow/contrib/slim/python/slim/data/BUILD @@ -61,6 +61,7 @@ py_test( name = "dataset_data_provider_test", srcs = ["dataset_data_provider_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":dataset", ":dataset_data_provider", diff --git a/tensorflow/contrib/slim/python/slim/data/parallel_reader.py b/tensorflow/contrib/slim/python/slim/data/parallel_reader.py index ad5e985487190e72b9eb2809da964f3d7b34ef94..b3343aef47d9f352c3bcbef4afbe8f9bf2560e6d 100644 --- a/tensorflow/contrib/slim/python/slim/data/parallel_reader.py +++ b/tensorflow/contrib/slim/python/slim/data/parallel_reader.py @@ -221,7 +221,7 @@ def parallel_read(data_sources, the data will be cycled through indefinitely. num_readers: a integer, number of Readers to create. reader_kwargs: an optional dict, of kwargs for the reader. - shuffle: boolean, wether should shuffle the files and the records by using + shuffle: boolean, whether should shuffle the files and the records by using RandomShuffleQueue as common_queue. dtypes: A list of types. The length of dtypes must equal the number of elements in each record. If it is None it will default to 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 f5a9299d263450ba89617f38bf7a4c5cbc359cb1..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 @@ -236,7 +236,7 @@ class SingleEvaluationTest(test.TestCase): def _prepareCheckpoint(self, checkpoint_path): init_op = control_flow_ops.group(variables.global_variables_initializer(), variables.local_variables_initializer()) - saver = saver_lib.Saver() + saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V1) with self.test_session() as sess: sess.run(init_op) saver.save(sess, checkpoint_path) diff --git a/tensorflow/contrib/slim/python/slim/learning.py b/tensorflow/contrib/slim/python/slim/learning.py index 54362c87b561595697ee64b9d5e565fdc3f0bbe0..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: @@ -770,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/solvers/python/kernel_tests/linear_equations_test.py b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py index 7b609ae96b20a5c3d078777cc8fbb475e5eebb1b..a1282847bef981717d7fdb1474adbbaaae4621c0 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py @@ -47,8 +47,8 @@ def _get_linear_equations_tests(dtype_, use_static_shape_, shape_): 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()) + 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) @@ -66,18 +66,30 @@ def _get_linear_equations_tests(dtype_, use_static_shape_, shape_): x = array_ops.placeholder(dtype_) jacobi = array_ops.placeholder(dtype_) operator = util.create_operator(a) - preconditioners = [None, util.identity_operator(a), - util.create_operator(jacobi)] + 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) + 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}) + 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) diff --git a/tensorflow/contrib/solvers/python/kernel_tests/util_test.py b/tensorflow/contrib/solvers/python/kernel_tests/util_test.py index 12e94369cbae462c21867657119cd2dd9ee29651..5d7534657bff27f7169e6a97bf4b03d4f6a35bc9 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/util_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/util_test.py @@ -85,9 +85,11 @@ class UtilTest(test.TestCase): 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}) + [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) diff --git a/tensorflow/contrib/solvers/python/ops/least_squares.py b/tensorflow/contrib/solvers/python/ops/least_squares.py index fb7c0eb649c5216736b239d1a423cdaf7079f582..6e164f53420675d149ded6c1f42ca87bd89b158c 100644 --- a/tensorflow/contrib/solvers/python/ops/least_squares.py +++ b/tensorflow/contrib/solvers/python/ops/least_squares.py @@ -33,7 +33,7 @@ def cgls(operator, rhs, tol=1e-6, max_iter=20, name="cgls"): r"""Conjugate gradient least squares solver. Solves a linear least squares problem \\(||A x - rhs||_2\\) for a single - righ-hand side, using an iterative, matrix-free algorithm where the action of + right-hand side, using an iterative, matrix-free algorithm where the action of the matrix A is represented by `operator`. The CGLS algorithm implicitly applies the symmetric conjugate gradient algorithm to the normal equations \\(A^* A x = A^* rhs\\). The iteration terminates when either diff --git a/tensorflow/contrib/solvers/python/ops/linear_equations.py b/tensorflow/contrib/solvers/python/ops/linear_equations.py index 4dfaa97ac9834ca3c13a9f8e8d721ddaba33bf7d..9305c6a11c4ec898c82553773e8e7277a54ab82e 100644 --- a/tensorflow/contrib/solvers/python/ops/linear_equations.py +++ b/tensorflow/contrib/solvers/python/ops/linear_equations.py @@ -26,6 +26,7 @@ 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 @@ -40,7 +41,7 @@ def conjugate_gradient(operator, r"""Conjugate gradient solver. Solves a linear system of equations `A*x = rhs` for selfadjoint, positive - definite matrix `A` and righ-hand side vector `rhs`, using an iterative, + definite matrix `A` and right-hand side vector `rhs`, using an iterative, matrix-free algorithm where the action of the matrix A is represented by `operator`. The iteration terminates when either the number of iterations exceeds `max_iter` or when the residual norm has been reduced to `tol` @@ -84,10 +85,9 @@ def conjugate_gradient(operator, cg_state = collections.namedtuple("CGState", ["i", "x", "r", "p", "gamma"]) def stopping_criterion(i, state): - return math_ops.logical_and(i < max_iter, - linalg_ops.norm(state.r) > tol) + return math_ops.logical_and(i < max_iter, linalg_ops.norm(state.r) > tol) - 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 @@ -108,8 +108,7 @@ def conjugate_gradient(operator, rhs = array_ops.expand_dims(rhs, -1) if x is None: x = array_ops.expand_dims( - array_ops.zeros( - n, dtype=rhs.dtype.base_dtype), -1) + array_ops.zeros(n, dtype=rhs.dtype.base_dtype), -1) r0 = rhs else: x = array_ops.expand_dims(x, -1) @@ -119,7 +118,7 @@ def conjugate_gradient(operator, else: p0 = preconditioner.apply(r0) gamma0 = util.dot(r0, p0) - tol = tol * linalg_ops.norm(r0) + tol *= linalg_ops.norm(r0) i = constant_op.constant(0, dtype=dtypes.int32) state = cg_state(i=i, x=x, r=r0, p=p0, gamma=gamma0) _, state = control_flow_ops.while_loop(stopping_criterion, cg_step, 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/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..c1724c6e43c0365ceac18bb1d9dfea7351cca060 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.executing_eagerly() 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.executing_eagerly(): + 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, @@ -267,7 +269,7 @@ def _make_summary_writer(name, factory, **kwargs): resource = gen_summary_ops.summary_writer(shared_name=name) # TODO(apassos): Consider doing this instead. # node = factory(resource, **kwargs) - # if not context.in_eager_mode(): + # if not context.executing_eagerly(): # ops.get_default_session().run(node) ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, factory(resource, **kwargs)) @@ -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.") + if context.executing_eagerly(): + return None return ops.get_collection(ops.GraphKeys._SUMMARY_COLLECTION) # pylint: disable=protected-access @@ -311,7 +309,7 @@ def summary_writer_initializer_op(): Raises: RuntimeError: If in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "tf.contrib.summary.summary_writer_initializer_op is only " "supported in graph mode.") @@ -330,8 +328,12 @@ def summary_writer_function(name, tensor, function, family=None): Returns: The result of writing the summary. """ + name_scope = ops.get_name_scope() + if name_scope: + # Add a slash to allow reentering the name scope. + name_scope += "/" def record(): - with summary_op_util.summary_scope( + with ops.name_scope(name_scope), summary_op_util.summary_scope( name, family, values=[tensor]) as (tag, scope): with ops.control_dependencies([function(tag, scope)]): return constant_op.constant(True) @@ -479,7 +481,7 @@ def graph(param, step=None, name=None): Raises: TypeError: If `param` isn't already a @{tf.Tensor} in graph mode. """ - if not context.in_eager_mode() and not isinstance(param, ops.Tensor): + if not context.executing_eagerly() and not isinstance(param, ops.Tensor): raise TypeError("graph() needs a tf.Tensor (e.g. tf.placeholder) in graph " "mode, but was: %s" % type(param)) writer = context.context().summary_writer_resource diff --git a/tensorflow/contrib/summary/summary_ops_graph_test.py b/tensorflow/contrib/summary/summary_ops_graph_test.py index 2b7806f80d020e0064b0f5cf32fd765a9ee993d1..3aba04540eba12092d884cca10e23546eb91c91d 100644 --- a/tensorflow/contrib/summary/summary_ops_graph_test.py +++ b/tensorflow/contrib/summary/summary_ops_graph_test.py @@ -85,6 +85,38 @@ class DbTest(summary_test_util.SummaryDbTest): self.assertEqual(len(events), 2) self.assertEqual(events[1].summary.value[0].tag, 'my_scalar') + def testScalarSummaryNameScope(self): + """Test record_summaries_every_n_global_steps and all_summaries().""" + with ops.Graph().as_default(), self.test_session() as sess: + global_step = training_util.get_or_create_global_step() + global_step.initializer.run() + with ops.device('/cpu:0'): + step_increment = state_ops.assign_add(global_step, 1) + sess.run(step_increment) # Increment global step from 0 to 1 + + logdir = tempfile.mkdtemp() + with summary_ops.create_file_writer(logdir, max_queue=0, + name='t2').as_default(): + with summary_ops.record_summaries_every_n_global_steps(2): + summary_ops.initialize() + with ops.name_scope('scope'): + summary_op = summary_ops.scalar('my_scalar', 2.0) + + # Neither of these should produce a summary because + # global_step is 1 and "1 % 2 != 0" + sess.run(summary_ops.all_summary_ops()) + sess.run(summary_op) + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(len(events), 1) + + # Increment global step from 1 to 2 and check that the summary + # is now written + sess.run(step_increment) + sess.run(summary_ops.all_summary_ops()) + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(len(events), 2) + self.assertEqual(events[1].summary.value[0].tag, 'scope/my_scalar') + def testSummaryGraphModeCond(self): with ops.Graph().as_default(), self.test_session(): training_util.get_or_create_global_step() diff --git a/tensorflow/contrib/summary/summary_ops_test.py b/tensorflow/contrib/summary/summary_ops_test.py index dfaa4182bb867cc03480320eaf1804da36206655..c756f8b27055f9cf86a311e485d97745a3c7a95b 100644 --- a/tensorflow/contrib/summary/summary_ops_test.py +++ b/tensorflow/contrib/summary/summary_ops_test.py @@ -29,7 +29,7 @@ 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 ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import state_ops @@ -59,12 +59,6 @@ _NUMPY_NUMERIC_TYPES = { 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') - def testShouldRecordSummary(self): self.assertFalse(summary_ops.should_record_summaries()) with summary_ops.always_record_summaries(): @@ -114,6 +108,20 @@ class TargetTest(test_util.TensorFlowTestCase): self.assertEqual(len(events), 2) self.assertEqual(events[1].summary.value[0].tag, 'scalar') + def testSummaryNameScope(self): + training_util.get_or_create_global_step() + logdir = tempfile.mkdtemp() + with summary_ops.create_file_writer( + logdir, max_queue=0, + name='t2').as_default(), summary_ops.always_record_summaries(): + + with ops.name_scope('scope'): + summary_ops.scalar('scalar', 2.0) + + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(len(events), 2) + self.assertEqual(events[1].summary.value[0].tag, 'scope/scalar') + def testSummaryGlobalStep(self): step = training_util.get_or_create_global_step() logdir = tempfile.mkdtemp() 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 58a7fa095d8356229fdb5879bea99d316113c828..07b6b1f142ba4d113c0d45db5a8a2ce3f66f4154 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", @@ -552,7 +553,6 @@ py_test( srcs = ["client/random_forest_test.py"], srcs_version = "PY2AND3", tags = [ - "no_windows", "nomac", # b/63258195 "notsan", ], diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h index 04e6b0a735320dd024e326a94ef910593a326245..dc3e9fe79d32a19930d500b62b520eddb4b41aa8 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h @@ -468,7 +468,7 @@ class FixedSizeSparseClassificationGrowStats : 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 { diff --git a/tensorflow/contrib/tensorboard/BUILD b/tensorflow/contrib/tensorboard/BUILD index 2e0a46ffe432341a423ac159deb7745d9ef15374..db2e000ef875c354edcdf0fa50fd3ac6d4907c17 100644 --- a/tensorflow/contrib/tensorboard/BUILD +++ b/tensorflow/contrib/tensorboard/BUILD @@ -9,11 +9,11 @@ exports_files(["LICENSE"]) # For platform specific build config load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library") +load("//tensorflow:tensorflow.bzl", "py_test") tf_proto_library( name = "protos_all", srcs = glob(["**/*.proto"]), - go_api_version = 2, visibility = ["//visibility:public"], ) diff --git a/tensorflow/contrib/tensorboard/db/summary_file_writer.cc b/tensorflow/contrib/tensorboard/db/summary_file_writer.cc index d891e86e53f4d760bfaea0e67601cfda037a4564..85b3e7231bcb433e9510522597c03c5f764f06cf 100644 --- a/tensorflow/contrib/tensorboard/db/summary_file_writer.cc +++ b/tensorflow/contrib/tensorboard/db/summary_file_writer.cc @@ -42,14 +42,14 @@ class SummaryFileWriter : public SummaryWriterInterface { if (is_dir.code() != tensorflow::error::NOT_FOUND) { return is_dir; } - TF_RETURN_IF_ERROR(env_->CreateDir(logdir)); + TF_RETURN_IF_ERROR(env_->RecursivelyCreateDir(logdir)); } mutex_lock ml(mu_); events_writer_ = tensorflow::MakeUnique(io::JoinPath(logdir, "events")); - if (!events_writer_->InitWithSuffix(filename_suffix)) { - return errors::Unknown("Could not initialize events writer."); - } + 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(); @@ -151,9 +151,8 @@ class SummaryFileWriter : public SummaryWriterInterface { events_writer_->WriteEvent(*e); } queue_.clear(); - if (!events_writer_->Flush()) { - return errors::InvalidArgument("Could not flush events file."); - } + TF_RETURN_WITH_CONTEXT_IF_ERROR(events_writer_->Flush(), + "Could not flush events file."); last_flush_ = env_->NowMicros(); return Status::OK(); } diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD index 28f571e1f015435b33406d0b802bb8e059c6e5ee..906cc3f0344e7cb641589bd522e33d658150d3b5 100644 --- a/tensorflow/contrib/tensorrt/BUILD +++ b/tensorflow/contrib/tensorrt/BUILD @@ -1,5 +1,6 @@ # Description: -# Wrap NVIDIA TensorRT (http://developer.nvidia.com/tensorrt) with tensorflow. +# 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__"]) @@ -8,7 +9,19 @@ 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", @@ -32,6 +45,234 @@ tf_cuda_cc_test( ]), ) +tf_custom_op_library( + name = "python/ops/_trt_engine_op.so", + srcs = [ + "ops/trt_calib_op.cc", + "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_calib_op.cc", + "kernels/trt_engine_op.cc", + ], + hdrs = [ + "kernels/trt_calib_op.h", + "kernels/trt_engine_op.h", + ], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = [ + ":trt_logging", + ":trt_resources", + "//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(), + # TODO(laigd) + alwayslink = 1, # buildozer: disable=alwayslink-with-hdrs +) + +tf_gen_op_libs( + op_lib_names = [ + "trt_engine_op", + "trt_calib_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", + gen_locally = True, + deps = [ + ":trt_calib_op_op_lib", + ":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", + "//tensorflow/python:errors", + ], +) + +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", + ], +) + +tf_cuda_library( + name = "trt_resources", + srcs = [ + "resources/trt_int8_calibrator.cc", + "resources/trt_resource_manager.cc", + ], + hdrs = [ + "resources/trt_int8_calibrator.h", + "resources/trt_resource_manager.h", + "resources/trt_resources.h", + ], + deps = [ + ":trt_logging", + "//tensorflow/core:framework_headers_lib", + "//tensorflow/core:framework_lite", + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]), +) + +# Library for the node-level conversion portion of TensorRT operation creation +tf_cuda_library( + name = "trt_conversion", + srcs = [ + "convert/convert_graph.cc", + "convert/convert_nodes.cc", + ], + hdrs = [ + "convert/convert_graph.h", + "convert/convert_nodes.h", + ], + deps = [ + ":segment", + ":trt_logging", + ":trt_resources", + "//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( diff --git a/tensorflow/contrib/tensorrt/README.md b/tensorflow/contrib/tensorrt/README.md new file mode 100644 index 0000000000000000000000000000000000000000..461e627e99c38fe492dc8a8a1c9fb06e6cfda3a2 --- /dev/null +++ b/tensorflow/contrib/tensorrt/README.md @@ -0,0 +1,25 @@ +Using TensorRT in TensorFlow +============================ + +This module provides necessary bindings and introduces TRT_engine_op +operator that wraps a subgraph in TensorRT. This is still a work in progress +but should be useable with most common graphs. + +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. + +```shell +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 can be found in test/test_tftrt.py directory diff --git a/tensorflow/contrib/tensorrt/__init__.py b/tensorflow/contrib/tensorrt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..140ad4828208ae4844a49bf664955b50cd9e51cd --- /dev/null +++ b/tensorflow/contrib/tensorrt/__init__.py @@ -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. +# ============================================================================= +"""Exposes the python wrapper for TensorRT graph transforms.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import errors + +# pylint: disable=unused-import,wildcard-import,g-import-not-at-top +try: + from tensorflow.contrib.tensorrt.python import * +except errors.NotFoundError as e: + no_trt_message = ( + '**** Failed to initialize TensorRT. This is either because the TensorRT' + ' installation path is not in LD_LIBRARY_PATH, or because you do not have' + ' it installed. If not installed, please go to' + ' https://developer.nvidia.com/tensorrt to download and install' + ' TensorRT ****') + print(no_trt_message) + raise e +# pylint: enable=unused-import,wildcard-import,g-import-not-at-top diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc new file mode 100644 index 0000000000000000000000000000000000000000..eea8c8efa28c1b7f9a1979311e920ca5c89f830a --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc @@ -0,0 +1,427 @@ +/* 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 + +#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" // NOLINT + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +namespace convert { +namespace { + +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", + "Mean", + "AvgPool", + "ConcatV2", + "DepthwiseConv2dNative", + "FusedBatchNorm", + "FusedBatchNormV2", + // TODO(ben,jie): ... + }; + // 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); + } else { + VLOG(2) << edge->src()->name() << " N, "; + } + } + } +} + +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()) { + VLOG(2) << edge->dst()->name() << " Y, "; + outgoing_edges->insert(edge); + } else { + VLOG(2) << edge->dst()->name() << " N, "; + } + } + } +} + +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; +} +// TODO(sami): convert references to pointers +struct ConvertGraphParams { + ConvertGraphParams( + tensorflow::Graph& inp_graph, + const std::vector& output_node_names, + const std::set& subgraph_node_id_numbers, + size_t max_supported_batch_size, size_t max_consumed_workspace_size_bytes, + const tensorflow::grappler::GraphProperties& current_graph_properties, + std::unordered_map>* output_edges, + int engine_precision_mode) + : graph(inp_graph), + output_names(output_node_names), + subgraph_node_ids(subgraph_node_id_numbers), + max_batch_size(max_supported_batch_size), + max_workspace_size_bytes(max_consumed_workspace_size_bytes), + graph_properties(current_graph_properties), + output_edge_map(output_edges), + precision_mode(engine_precision_mode) {} + tensorflow::Graph& graph; + const std::vector& output_names; + const std::set& subgraph_node_ids; + size_t max_batch_size; + size_t max_workspace_size_bytes; + const tensorflow::grappler::GraphProperties& graph_properties; + std::unordered_map>* output_edge_map; + int precision_mode; + std::vector> subgraph_inputs; + std::vector> subgraph_outputs; + tensorflow::EdgeSet subgraph_incoming_edges; + tensorflow::EdgeSet subgraph_outgoing_edges; +}; + +static tensorflow::Status FillSubGraphEdgeSets(ConvertGraphParams* p) { + GetSubGraphIncomingEdges(p->graph, p->subgraph_node_ids, + &p->subgraph_incoming_edges); + for (const tensorflow::Edge* edge : p->subgraph_incoming_edges) { + p->subgraph_inputs.push_back({edge->src()->id(), edge->src_output()}); + } + auto output_name_to_index_map = BuildTensorNameMap(p->output_names); + std::set> subgraph_outputs_set; + // Collect outputs referenced from output_names + for (int node_id : p->subgraph_node_ids) { + tensorflow::Node* node = p->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}); + } + } + } + GetSubGraphOutgoingEdges(p->graph, p->subgraph_node_ids, + &p->subgraph_outgoing_edges); + for (const tensorflow::Edge* edge : p->subgraph_outgoing_edges) { + subgraph_outputs_set.insert({edge->src()->id(), edge->src_output()}); + } + p->subgraph_outputs.reserve(subgraph_outputs_set.size()); + p->subgraph_outputs.insert(p->subgraph_outputs.begin(), + subgraph_outputs_set.begin(), + subgraph_outputs_set.end()); + return tensorflow::Status::OK(); +}; + +tensorflow::Status GetCalibNode(ConvertGraphParams* params) { + TF_RETURN_IF_ERROR(FillSubGraphEdgeSets(params)); + tensorflow::NodeDef trt_node_def; + SubGraphParams s(params->graph, params->subgraph_node_ids, + params->subgraph_inputs, params->subgraph_outputs, + params->max_batch_size, params->max_workspace_size_bytes, + params->graph_properties, params->output_edge_map, + &trt_node_def, params->precision_mode); + TF_RETURN_IF_ERROR(InjectCalibrationNode(s)); + tensorflow::Status status; + tensorflow::Node* trt_node = params->graph.AddNode(trt_node_def, &status); + + TF_RETURN_IF_ERROR(status); + + for (auto in_edge : + params->subgraph_incoming_edges) { // loop over incoming edges and + // attach them to calib node + // tensorflow::Node* src_node = in_edge->src(); + auto src_output = in_edge->src_output(); + auto dst_node = in_edge->dst(); + auto dst_input = in_edge->dst_input(); + VLOG(1) << " update edge " << trt_node->name() << ":" << src_output + << " -> " << dst_node->name() << ":" << dst_input; + TF_RETURN_IF_ERROR( + params->graph.UpdateEdge(trt_node, src_output, dst_node, dst_input)); + } + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertSubGraphToTensorRT(ConvertGraphParams* params) { + TF_RETURN_IF_ERROR(FillSubGraphEdgeSets(params)); + tensorflow::NodeDef trt_node_def; + + SubGraphParams s(params->graph, params->subgraph_node_ids, + params->subgraph_inputs, params->subgraph_outputs, + params->max_batch_size, params->max_workspace_size_bytes, + params->graph_properties, params->output_edge_map, + &trt_node_def, params->precision_mode); + TF_RETURN_IF_ERROR(ConvertSubGraphToTensorRTNodeDef(s)); + tensorflow::Status status; + tensorflow::Node* trt_node = params->graph.AddNode(trt_node_def, &status); + + // AddNode does not wire edges. + // Re-map incoming edges to use the new TRT node instead of the orig subgraph + std::map, int> subgraph_edge_to_input_map; + for (size_t i = 0; i < params->subgraph_inputs.size(); ++i) { + subgraph_edge_to_input_map.insert({params->subgraph_inputs.at(i), i}); + } + for (const tensorflow::Edge* edge : params->subgraph_incoming_edges) { + std::pair old_src = {edge->src()->id(), edge->src_output()}; + int new_src_output = subgraph_edge_to_input_map.at(old_src); + params->graph.AddEdge(edge->src(), edge->src_output(), trt_node, + new_src_output); + params->graph.RemoveEdge(edge); + } + + VLOG(2) << "new wiring edges: " << trt_node->in_edges().size(); + for (const tensorflow::Edge* edge : trt_node->in_edges()) { + VLOG(2) << edge->src()->name() << " port: " << edge->src_output(); + } + + 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 < params->subgraph_outputs.size(); ++i) { + subgraph_edge_to_output_map.insert({params->subgraph_outputs.at(i), i}); + } + TF_RETURN_IF_ERROR(status); + for (const tensorflow::Edge* edge : params->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(params->graph.UpdateEdge( + trt_node, new_src_output, edge->dst(), edge->dst_input())); + } + // Remove the original subgraph + for (int node_id : params->subgraph_node_ids) { + tensorflow::Node* node = params->graph.FindNodeId(node_id); + // Don't remove the input placeholders + if (node->type_string() == "Placeholder") { + continue; + } + params->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 ConvertCalibGraphToInferGraph( + const tensorflow::GraphDef& graph_def, tensorflow::GraphDef* infer_graph) { + VLOG(0) << "Starting Calib Conversion"; + tensorflow::Graph graph(tensorflow::OpRegistry::Global()); + TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToGraph( + tensorflow::GraphConstructorOptions(), graph_def, &graph)); + // get calib nodes + std::vector calib_nodes; + for (auto node : graph.op_nodes()) { + if (node->type_string() == "TRTCalibOp") { + VLOG(1) << "Found Calib Node"; + calib_nodes.push_back(node); + } + } + VLOG(0) << "Num Calib nodes in graph= " << calib_nodes.size(); + if (calib_nodes.size() == 0) + return tensorflow::errors::FailedPrecondition( + "Graph doesn't contain any calibration nodes!." + " Please generate calibration graph and run calibration first"); + for (auto n : calib_nodes) { + TF_RETURN_IF_ERROR( + tensorrt::convert::ConvertCalibrationNodeToEngineNode(graph, n)); + } + graph.ToGraphDef(infer_graph); + return tensorflow::Status::OK(); +} + +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, + int precision_mode = FP32MODE, int minimum_segment_size = 3) { + // 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}}); + + // single machine + int num_cpu_cores = tensorflow::grappler::GetNumAvailableLogicalCPUCores(); + int num_gpus = tensorflow::grappler::GetNumAvailableGPUs(); + VLOG(2) << "cpu_cores: " << num_cpu_cores; + VLOG(2) << "gpus: " << num_gpus; + + 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 = minimum_segment_size; + 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)); + std::unordered_map> output_edge_map; + int count = 0; + float total_num_nodes_in_segments = 0.; + for (auto s : segments) { + total_num_nodes_in_segments += s.size(); + } + for (const std::set& subgraph_node_names : segments) { + std::set subgraph_node_ids; + size_t max_mem_per_engine = + max_workspace_size_bytes * + ((float)subgraph_node_names.size() / total_num_nodes_in_segments); + std::stringstream oss; + for (const string& node_name : subgraph_node_names) { + oss << " " << node_name; + subgraph_node_ids.insert(node_map.at(node_name)->id()); + } + VLOG(2) << "Subgraph nodes" << oss.str(); + ConvertGraphParams p(graph, output_names, subgraph_node_ids, max_batch_size, + max_mem_per_engine, static_graph_properties, + &output_edge_map, precision_mode); + if (precision_mode == INT8MODE) { + TF_RETURN_IF_ERROR(GetCalibNode(&p)); + } else { + tensorflow::Status status = ConvertSubGraphToTensorRT(&p); + if (status != tensorflow::Status::OK()) { + LOG(WARNING) << "subgraph conversion error for subgraph_index:" << count + << " due to: \n" + << status.ToString() << " SKIPPING......"; + } + count++; + } + } + 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..e01e4a5328061ad527b2dac6e2e4ef1559bd914d --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_graph.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_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 { + +// This method converts an already generated calibration graph which was used in +// calibration runs to an inference graph +tensorflow::Status ConvertCalibGraphToInferGraph( + const tensorflow::GraphDef& graph_def, tensorflow::GraphDef* new_graph_def); + +// 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 allowance 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, + int precision_mode, int minimum_segment_size); + +} // 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..92a692baa743ecb957cf1ffc156e962335b99b3f --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc @@ -0,0 +1,2662 @@ +/* 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/contrib/tensorrt/log/trt_logger.h" +#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" +#include "tensorflow/contrib/tensorrt/resources/trt_resources.h" +#include "tensorflow/core/framework/node_def.pb.h" // NOLINT +#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/str_util.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 "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 { +using ::tensorflow::strings::StrCat; + +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 " + tensorflow::DataTypeString(tf_dtype)); + } + 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; +} + +string GetCommonNameScope(const string& op_name_a, const string& op_name_b) { + size_t last_scope_separator = 0; + for (size_t i = 0; i < std::min(op_name_a.size(), op_name_b.size()); ++i) { + if (op_name_a[i] != op_name_b[i]) { + break; + } else if (op_name_a[i] == '/') { + last_scope_separator = i + 1; + } + } + return op_name_a.substr(0, last_scope_separator); +} + +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 <> +std::vector TFAttrs::get>(string key) const { + auto attr = this->at(key)->list().s(); + 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 <> +float TFAttrs::get(string key) const { + return this->at(key)->f(); +} + +template <> +bool TFAttrs::get(string key) const { + return this->at(key)->b(); +} + +// TODO(jie): reorder4 & reorder2 should be merged? +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()]; + } + } + } + } +} + +template +void Reorder2(nvinfer1::DimsHW shape, const T* idata, nvinfer1::DimsHW istrides, + T* odata, nvinfer1::DimsHW ostrides) { + for (int h = 0; h < shape.h(); ++h) { + for (int w = 0; w < shape.w(); ++w) { + odata[h * ostrides.h() + w * ostrides.w()] = + idata[h * ostrides.h() + w * ostrides.w()]; + } + } +} + +// TODO(jie): fallback to tensorflow!! +void ReorderCKtoKC(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights) { + int c = iweights.shape_.d[0]; + int k = iweights.shape_.d[1]; + oweights->shape_.d[0] = k; + oweights->shape_.d[1] = c; + nvinfer1::DimsHW istrides = {1, k}; + nvinfer1::DimsHW ostrides = {c, 1}; + switch (iweights.type_) { + case tensorflow::DataType::DT_FLOAT: { + Reorder2({k, c}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + case tensorflow::DataType::DT_HALF: { + Reorder2( + {k, c}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + default: + LOG(FATAL) << "Unsupported type in reorder expected fp32 or fp16 but got " + << DataTypeString(iweights.type_); + } +} + +void ReorderRSCKToKCRS(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights, int num_groups) { + 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]; + // TRT requires GKcRS, while TF depthwise has RSCK + // where c=1, C=G + VLOG(2) << "num_groups: " << num_groups; + int c = iweights.shape_.d[2] / num_groups; + VLOG(2) << "c" << iweights.shape_.d[2] << " then " << c; + int k = iweights.shape_.d[3] * num_groups; + VLOG(2) << "k" << iweights.shape_.d[3] << " then " << k; + oweights->shape_.d[0] = k / num_groups; + oweights->shape_.d[1] = c * num_groups; + 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; + } + case tensorflow::DataType::DT_HALF: { + Reorder4( + {k, c, r, s}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + + default: + LOG(FATAL) << "Unsupported type, expected fp32 or fp16 but got " + << DataTypeString(iweights.type_); + } +} + +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()); +} + +class Converter; + +using OpConverter = + std::function&, + std::vector*)>; + +class Converter { + std::unordered_map trt_tensors_; + std::unordered_map op_registry_; + nvinfer1::INetworkDefinition* trt_network_; + std::list> temp_bufs_; + tensorflow::tensorrt::TRTWeightStore* weight_store_; + bool fp16_; + void register_op_converters(); + std::vector get_inputs( + const tensorflow::NodeDef& node_def) { + std::vector inputs; + for (auto const& input_name : node_def.input()) { + /************************************************************************* + * TODO(jie) handle case 1) here + * Normalizes the inputs and extracts associated metadata: + * 1) Inputs can contain a colon followed by a suffix of characters. + * That suffix may be a single number (e.g. inputName:1) or several + * word characters separated from a number by a colon + * (e.g. inputName:foo:1). The + * latter case is used to denote inputs and outputs of functions. + * 2) Control dependency inputs contain caret at the beginning and we + * remove this and annotate the edge as a control dependency. + ************************************************************************/ + string name = input_name[0] == '^' ? input_name.substr(1) : input_name; + auto first = name.find_first_of(':'); + if (first != string::npos && first + 2 == name.size() && + name[first + 1] == '0') + name.erase(first); + + VLOG(2) << "retrieve input: " << name; + if (trt_tensors_.count(name)) { + inputs.push_back(trt_tensors_.at(name)); + } else { + LOG(FATAL) << "input: " << name << " not available for node at, " + << node_def.name(); + } + } + return inputs; + } + + public: + explicit Converter(nvinfer1::INetworkDefinition* trt_network, + tensorflow::tensorrt::TRTWeightStore* ws, bool fp16) + : trt_network_(trt_network), weight_store_(ws), fp16_(fp16) { + this->register_op_converters(); + } + tensorflow::tensorrt::TRTWeightStore* weight_store() { return weight_store_; } + 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 + weight_store_->store_.push_back(std::vector(weights.size_bytes())); + weights.SetValues(weight_store_->store_.back().data()); + return weights; + } + bool isFP16() { return fp16_; }; + 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 = StrCat(output_name, ":", 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 / 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; + }; + } +}; + +template <> +std::function LambdaFactory::unary() { + switch (op) { + case OP_CATEGORY::RSQRT: { + VLOG(2) << "RSQRT GETS DONE"; + return [](Eigen::half t) -> Eigen::half { + return Eigen::half(1.0 / sqrt(float(t))); + }; + } + case OP_CATEGORY::NEG: + return [](Eigen::half t) -> Eigen::half { return -t; }; + default: + VLOG(2) << "Not supported op for unary: " << static_cast(op); + return nullptr; + } +} +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; + } + case tensorflow::DataType::DT_HALF: { + 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; + } + case tensorflow::DataType::DT_HALF: { + 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, + const std::vector& 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")); + + 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 back to TF +// approach for constant folding +tensorflow::Status ConstantFoldBinary( + Converter& ctx, const tensorflow::NodeDef& node_def, + const std::vector& 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 num_dims = weights_input_l.shape_.nbDims; + nvinfer1::Dims output_shape; + output_shape.nbDims = num_dims; + VLOG(2) << "nb_dims: " << num_dims + << ", the other: " << weights_input_r.shape_.nbDims; + for (int i = 0; i < num_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); + + 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 + nvinfer1::DataType ttype; + TF_CHECK_OK(ConvertDType(weights.type_, &ttype)); + + // Check scale mode + auto dims_w = weights.shape_; + auto dims_t = tensor->getDimensions(); + + // default to element-wise + auto scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; + + // TODO(jie): maybe use a permutation instead to support more cases; + bool permutation_flag = false; + + if (weights.count() == 1) { + VLOG(2) << "UNIFORM"; + scale_mode = nvinfer1::ScaleMode::kUNIFORM; + } else { + // no broadcasting on Batch dimension; + VLOG(2) << "WEIGHTS DIM: " << dims_w.nbDims + << " tensor DIM: " << dims_t.nbDims; + if (dims_w.nbDims == dims_t.nbDims + 1) { + if (dims_w.d[0] == 1) { + for (int i = 1; i < dims_w.nbDims; i++) { + dims_w.d[i - 1] = dims_w.d[i]; + } + dims_w.nbDims--; + } else { + return tensorflow::errors::InvalidArgument( + "Binary op cannot operate on batch, " + node_def.name()); + } + } + + if (dims_w.nbDims == dims_t.nbDims && dims_w.d[0] == dims_t.d[0]) { + scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; + // default is element; + for (int i = 1; i < dims_w.nbDims; i++) { + if (dims_w.d[i] != dims_t.d[i]) { + // if dimension does not match, switch back to channel; + VLOG(2) << "channel"; + scale_mode = nvinfer1::ScaleMode::kCHANNEL; + break; + } + } + // if channel as candidate, validate it + if (scale_mode == nvinfer1::ScaleMode::kCHANNEL) { + for (int i = 1; i < dims_w.nbDims; i++) { + if (dims_w.d[i] != 1) + return tensorflow::errors::InvalidArgument( + "Weight shape not compatible at, " + node_def.name()); + } + } else { + VLOG(2) << "elementwise"; + } + } else if (dims_w.nbDims == 1 && + dims_w.d[0] == dims_t.d[dims_t.nbDims - 1]) { + // channel wise and broadcast required; + permutation_flag = true; + scale_mode = nvinfer1::ScaleMode::kCHANNEL; + } else { + return tensorflow::errors::InvalidArgument( + "Weight shape not compatible at, " + node_def.name()); + } + } + + // transpose last dimension + std::vector permutation(dims_t.nbDims + 1); + if (permutation_flag) { + if (scale_mode == nvinfer1::ScaleMode::kCHANNEL && dims_t.nbDims > 1) { + // we swap the last dimension into channel for trt. + // because of tensorflow default broadcasting rules. + for (int i = 0; i < static_cast(permutation.size()); i++) { + permutation[i] = i; + } + permutation[1] = dims_t.nbDims; + permutation[dims_t.nbDims] = 1; + tensor = ctx.TransposeTensor(const_cast(tensor), + permutation); + } else { + return tensorflow::errors::InvalidArgument( + "Transpose cannot be applied, " + 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); + // transpose back dimension + if (permutation_flag) { + output_tensor = ctx.TransposeTensor(output_tensor, permutation); + } + + // Pass the output + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +enum class ConvolutionType { DEFAULT, DEPTHWISE_CONV }; + +tensorflow::Status ConvertConv2DHelper( + Converter& ctx, const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs, + int group // group ==0 specifies depthwise conv +) { + const nvinfer1::ITensor* 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") { + tensor = ctx.TransposeTensor(const_cast(tensor), + {0, 3, 1, 2}); + h_index = 1; + w_index = 2; + // TODO(jie): transpose it + } + + // tensor after transpose (NCHW) + auto tensor_dim = tensor->getDimensions(); + + int num_groups = group; + if (num_groups == 0) // depthwise convolution + num_groups = tensor_dim.d[0]; + VLOG(2) << "groups count: " << num_groups; + + TRT_ShapedWeights weights_rsck = inputs.at(1).weights(); + TRT_ShapedWeights weights = ctx.get_temp_weights_like(weights_rsck); + ReorderRSCKToKCRS(weights_rsck, &weights, num_groups); + TRT_ShapedWeights biases(weights.type_); + int noutput = weights.shape_.d[0] * num_groups; + nvinfer1::DimsHW kernel_size; + kernel_size.h() = weights.shape_.d[2]; + kernel_size.w() = weights.shape_.d[3]; + VLOG(2) << "kernel size: " << kernel_size.h() << ", " << kernel_size.w(); + + // TODO(jie): stride. (NHWC/NCHW) + auto tf_stride = attrs.get>("strides"); + VLOG(2) << "h_INDEX" << h_index << ", w_index " << w_index; + VLOG(2) << "stride!!!: " << tf_stride[0] << tf_stride[1] << tf_stride[2] + << tf_stride[3]; + nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); + + 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()); + layer->setNbGroups(num_groups); + 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 ConvertConv2DHelper( + Converter& ctx, const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs, ConvolutionType type) { + switch (type) { + case ConvolutionType::DEFAULT: + return ConvertConv2DHelper(ctx, node_def, inputs, outputs, 1); + case ConvolutionType::DEPTHWISE_CONV: + return ConvertConv2DHelper(ctx, node_def, inputs, outputs, 0); + } + return tensorflow::errors::Unimplemented("unsupported convolution type at, " + + node_def.name()); +} + +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}, + {"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, + const std::vector& 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) { + return ConvertConv2DHelper(ctx, node_def, inputs, outputs, + ConvolutionType::DEFAULT); +} + +tensorflow::Status ConvertConv2DDepthwise( + Converter& ctx, const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs) { + return ConvertConv2DHelper(ctx, node_def, inputs, outputs, + ConvolutionType::DEPTHWISE_CONV); +} + +tensorflow::Status ConvertPool(Converter& ctx, + const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs) { + const nvinfer1::ITensor* 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 if (node_def.op() == "AvgPool") + type = nvinfer1::PoolingType::kAVERAGE; + 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, + const std::vector& inputs, + std::vector* outputs) { + const nvinfer1::ITensor* 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, + const std::vector& 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; + const nvinfer1::ITensor* 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 !!!!"; + } + + auto dims = tensor->getDimensions(); + VLOG(2) << "tensor dimensions: " << dims.nbDims; + for (int i = 0; i < dims.nbDims; i++) { + VLOG(2) << "i: " << dims.d[i]; + } + dims = weights.shape_; + VLOG(2) << "tensor dimensions: " << dims.nbDims; + for (int i = 0; i < dims.nbDims; i++) { + VLOG(2) << "i: " << dims.d[i]; + } + + nvinfer1::ScaleMode mode = nvinfer1::ScaleMode::kCHANNEL; + if (weights.shape_.d[0] == 1) { + mode = nvinfer1::ScaleMode::kUNIFORM; + } + + nvinfer1::IScaleLayer* layer = + ctx.network()->addScale(*const_cast(tensor), mode, + 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, + const std::vector& 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(); + VLOG(2) << "size: " << weights_tensor.float_val_size(); + scalar_shape = GetTensorShape(tensor); + for (int i = 0; i < scalar_shape.nbDims; i++) + VLOG(2) << scalar_shape.d[i]; + if (GetShapeSize(scalar_shape) != weights_tensor.float_val_size()) { + if (weights_tensor.float_val_size() == 1 || + scalar_shape.d[0] == weights_tensor.float_val_size()) { + scalar_shape.nbDims = 1; + // no dimension provided. flatten it + scalar_shape.d[0] = weights_tensor.float_val_size(); + scalar_shape.type[0] = nvinfer1::DimensionType::kSPATIAL; + } else { + LOG(FATAL) << "Broadcast on weights only supports kCHANNEL and" + << " kUNIFORM, at: " << node_def.name(); + } + } + } else { + VLOG(2) << "Dimensions: " << tensor.dims(); + scalar_shape.nbDims = 1; + // no dimension provided. flatten it + scalar_shape.d[0] = weights_tensor.float_val_size(); + 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; + } + } + if (ctx.isFP16()) { + auto dtype_new = tensorflow::DataType::DT_HALF; + size_t len_data = tensorflow::DataTypeSize(dtype_new); + for (int i = 0; i < scalar_shape.nbDims; i++) + len_data *= scalar_shape.d[i]; + ctx.weight_store()->store_.push_back(std::vector(len_data)); + void* dst = static_cast(&(ctx.weight_store()->store_.back()[0])); + tensorflow::Tensor temp_tensor(tensorflow::DT_HALF, tensor.shape()); + auto half_tensor = temp_tensor.flat(); + Eigen::DefaultDevice defd; + half_tensor.device(defd) = + tensor.flat().template cast(); + memcpy(dst, half_tensor.data(), len_data); // store into weight store + weights = TRT_ShapedWeights(dtype_new, dst, scalar_shape); + } else { + size_t len_data = tensorflow::DataTypeSize(dtype); + for (int i = 0; i < scalar_shape.nbDims; i++) + len_data *= scalar_shape.d[i]; + ctx.weight_store()->store_.push_back(std::vector(len_data)); + void* dst = static_cast(&(ctx.weight_store()->store_.back()[0])); + std::vector tensor_data( + weights_tensor.float_val().begin(), + weights_tensor.float_val() + .end()); // make a local copy first to flatten + memcpy(dst, tensor_data.data(), len_data); // store into weight store + weights = TRT_ShapedWeights(dtype, dst, scalar_shape); + } + } else if (!weights_tensor.int_val().empty()) { + VLOG(2) << "int!!!" << node_def.name(); + nvinfer1::Dims scalar_shape; + if (tensor.dims() > 0) { + VLOG(2) << "dimensions: " << tensor.dims(); + scalar_shape = GetTensorShape(tensor); + if (GetShapeSize(scalar_shape) != weights_tensor.int_val_size()) { + if (weights_tensor.int_val_size() == 1 || + scalar_shape.d[0] == weights_tensor.int_val_size()) { + scalar_shape.nbDims = 1; + // no dimension provided. flatten it + scalar_shape.d[0] = weights_tensor.int_val_size(); + scalar_shape.type[0] = nvinfer1::DimensionType::kSPATIAL; + } else { + LOG(FATAL) << "Broadcast on weights only supports kCHANNEL and" + << " kUNIFORM, at: " << node_def.name(); + } + } + } else { + VLOG(2) << "dimensions: " << tensor.dims(); + scalar_shape.nbDims = 1; + // no dimension provided. flatten it + scalar_shape.d[0] = weights_tensor.int_val_size(); + 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; + } + } + if (ctx.isFP16()) { + auto dtype_new = tensorflow::DataType::DT_HALF; + size_t len_data = tensorflow::DataTypeSize(dtype_new); + for (int i = 0; i < scalar_shape.nbDims; i++) + len_data *= scalar_shape.d[i]; + ctx.weight_store()->store_.push_back(std::vector(len_data)); + void* dst = static_cast(&(ctx.weight_store()->store_.back()[0])); + tensorflow::Tensor temp_tensor(tensorflow::DT_HALF, tensor.shape()); + TTypes::Flat half_tensor = temp_tensor.flat(); + Eigen::DefaultDevice defd; + switch (dtype) { + case (tensorflow::DT_INT32): { + half_tensor.device(defd) = + tensor.flat().template cast(); + break; + } + case (tensorflow::DT_INT16): { + half_tensor.device(defd) = + tensor.flat().template cast(); + break; + } + case (tensorflow::DT_INT8): { + half_tensor.device(defd) = + tensor.flat().template cast(); + break; + } + case (tensorflow::DT_UINT8): { + half_tensor.device(defd) = + tensor.flat().template cast(); + break; + } + default: + return tensorflow::errors::InvalidArgument( + "Datatype " + tensorflow::DataTypeString(dtype) + + " for FP16 conversion"); + break; + }; + memcpy(dst, half_tensor.data(), len_data); // store into weight store + weights = TRT_ShapedWeights(dtype_new, dst, scalar_shape); + } else { + size_t len_data = tensorflow::DataTypeSize(dtype); + for (int i = 0; i < scalar_shape.nbDims; i++) + len_data *= scalar_shape.d[i]; + size_t len_tensor = weights_tensor.int_val_size() * sizeof(int32); + len_data = std::max(len_data, len_tensor); + ctx.weight_store()->store_.push_back(std::vector(len_data)); + void* dst = static_cast(&(ctx.weight_store()->store_.back()[0])); + std::vector tensor_data( + weights_tensor.int_val().begin(), + weights_tensor.int_val() + .end()); // make a local copy first to flatten + // doesn't have to be contiguous + memcpy(dst, tensor_data.data(), len_tensor); // store into weight store + weights = TRT_ShapedWeights(dtype, dst, 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, + const std::vector& inputs, + std::vector* outputs) { + outputs->push_back(inputs.at(0)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertBinary(Converter& ctx, + const tensorflow::NodeDef& node_def, + const std::vector& 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, + const std::vector& 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, + const std::vector& 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; + const nvinfer1::ITensor* 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) == 0) { + 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); + } + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertPad(Converter& ctx, + const tensorflow::NodeDef& node_def, + const std::vector& 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; + const nvinfer1::ITensor* 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(); +} + +tensorflow::Status ConvertConcat(Converter& ctx, + const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs) { + // not including the last input (axis) here + int input_size = static_cast(inputs.size()) - 1; + + if (!inputs.at(0).is_tensor()) + return tensorflow::errors::InvalidArgument( + "Concat in TRT support only Tensor input, at " + node_def.name()); + + // We are retrieving the axis + TRT_ShapedWeights axis = inputs.at(input_size).weights(); + + TFAttrs attrs(node_def); + // auto attr_size = attrs.at("N")->i(); + // auto data_type = attrs.get("T"); + auto index_type = attrs.get("Tidx"); + + // TODO(jie): handle data type + // Only expect to handle INT32 as index attributes for now + if (index_type != tensorflow::DataType::DT_INT32) + return tensorflow::errors::Unimplemented( + "Tidx supports only DT_INT32, at " + node_def.name()); + + int index = *(static_cast(const_cast(axis.GetValues()))); + + // TODO(jie): early termination with no-op (attr_size==1) + + auto dim = inputs.at(0).tensor()->getDimensions(); + // dimension check + if (index > dim.nbDims + 1) + return tensorflow::errors::InvalidArgument( + "Concatenate on axis out of dimension range, at " + node_def.name()); + + if (index == 0) + return tensorflow::errors::InvalidArgument( + "Concatenate on batch dimension not supported, at " + node_def.name()); + + // incase we need permutation; + std::vector permutation_order(dim.nbDims + 1); + + for (int i = 0; i < dim.nbDims + 1; i++) permutation_order[i] = i; + + if (index != 1) { + permutation_order[1] = index - 1; + permutation_order[index - 1] = 1; + } + + std::vector inputs_vec; + // Shap chack (all input tensor should have same shape) + // starting from 0 since we are probably also doing transpose here; + for (int i = 0; i < input_size; i++) { + auto tensor_i = inputs.at(i).tensor(); + auto dim_i = tensor_i->getDimensions(); + if (dim_i.nbDims != dim.nbDims) + return tensorflow::errors::InvalidArgument( + "Concatenate receives inputs with inconsistent dimensions, at " + + node_def.name()); + + for (int j = 0; j < dim.nbDims; j++) { + // check dimension consistency on non-concatenate axis + if (j != index - 1 && dim_i.d[j] != dim.d[j]) + return tensorflow::errors::InvalidArgument( + "Concatenate receives inputs with inconsistent shape, at" + + node_def.name()); + } + + // TRT does concatenation only on channel! + if (index != 1) + tensor_i = ctx.TransposeTensor(const_cast(tensor_i), + permutation_order); + + inputs_vec.push_back(tensor_i); + } + + // nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + nvinfer1::IConcatenationLayer* layer = ctx.network()->addConcatenation( + const_cast(inputs_vec.data()), + inputs_vec.size()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + if (index != 1) { + output_tensor = ctx.TransposeTensor(output_tensor, permutation_order); + } + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertFusedBatchNorm( + Converter& ctx, const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs) { + TFAttrs attrs(node_def); + float epsilon = attrs.get("epsilon"); + auto data_format = attrs.get("data_format"); + if (data_format != "NCHW") { + return tensorflow::errors::Unimplemented( + "only data_format=NCHW is supported, at " + node_def.name()); + } + bool is_training = attrs.get("is_training"); + if (is_training) { + return tensorflow::errors::Unimplemented( + "only is_training=false is supported, at " + node_def.name()); + } + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + TRT_ShapedWeights scale_weights = inputs.at(1).weights(); + TRT_ShapedWeights offset_weights = inputs.at(2).weights(); + TRT_ShapedWeights mean_weights = inputs.at(3).weights(); + TRT_ShapedWeights variance_weights = inputs.at(4).weights(); + TRT_ShapedWeights dummy_power_weights(scale_weights.type_); + TRT_ShapedWeights combined_scale_weights = + ctx.get_temp_weights_like(scale_weights); + TRT_ShapedWeights combined_offset_weights = + ctx.get_temp_weights_like(offset_weights); + size_t nweight = scale_weights.count(); + if ((scale_weights.type_ == offset_weights.type_) && + (mean_weights.type_ == variance_weights.type_) && + (scale_weights.type_ == variance_weights.type_)) { + if ((scale_weights.type_ != tensorflow::DataType::DT_FLOAT) && + (scale_weights.type_ != tensorflow::DataType::DT_HALF)) { + return tensorflow::errors::Unimplemented( + "only float32 or float16 weight data type is supported, for node " + + node_def.name() + " got " + + tensorflow::DataTypeString(scale_weights.type_)); + } + if (scale_weights.type_ == tensorflow::DT_FLOAT) { + for (size_t i = 0; i < nweight; ++i) { + float scale = (static_cast(scale_weights.GetValues()))[i]; + float offset = + (static_cast(offset_weights.GetValues()))[i]; + float mean = (static_cast(mean_weights.GetValues()))[i]; + float variance = + (static_cast(variance_weights.GetValues()))[i]; + float& combined_scale_ref = const_cast( + static_cast(combined_scale_weights.GetValues()))[i]; + float& combined_offset_ref = const_cast( + static_cast(combined_offset_weights.GetValues()))[i]; + combined_scale_ref = scale / sqrtf(variance + epsilon); + combined_offset_ref = offset - mean * combined_scale_ref; + } + } else { + const Eigen::half* scale_vals = + (static_cast(scale_weights.GetValues())); + const Eigen::half* off_vals = + (static_cast(offset_weights.GetValues())); + const Eigen::half* mean_vals = + (static_cast(mean_weights.GetValues())); + const Eigen::half* variance_vals = + (static_cast(variance_weights.GetValues())); + Eigen::half* comb_scale_vals = const_cast( + static_cast(combined_scale_weights.GetValues())); + Eigen::half* comb_off_vals = const_cast( + static_cast(combined_offset_weights.GetValues())); + for (size_t i = 0; i < nweight; ++i) { + float scale(scale_vals[i]); + float offset(off_vals[i]); + float mean(mean_vals[i]); + float variance(variance_vals[i]); + float combined_scale_ref = scale / sqrtf(variance + epsilon); + comb_scale_vals[i] = Eigen::half(combined_scale_ref); + float combined_offset_ref = offset - mean * combined_scale_ref; + comb_off_vals[i] = Eigen::half(combined_offset_ref); + } + } + } + nvinfer1::IScaleLayer* layer = ctx.network()->addScale( + *const_cast(tensor), nvinfer1::ScaleMode::kCHANNEL, + combined_offset_weights.GetWeightsForTRT(), + combined_scale_weights.GetWeightsForTRT(), + dummy_power_weights.GetWeightsForTRT()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertMatMul(Converter& ctx, + const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs) { + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + + // TODO(jie): transpose! + TFAttrs attrs(node_def); + + TRT_ShapedWeights weights_ck = inputs.at(1).weights(); + TRT_ShapedWeights weights = ctx.get_temp_weights_like(weights_ck); + ReorderCKtoKC(weights_ck, &weights); + TRT_ShapedWeights biases(weights.type_); + + int noutput = weights.shape_.d[0]; + + nvinfer1::IFullyConnectedLayer* layer = ctx.network()->addFullyConnected( + *const_cast(tensor), noutput, weights, biases); + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertReshape( + Converter& ctx, const tensorflow::NodeDef& node_def, + const std::vector& 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; + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + auto dims = tensor->getDimensions(); + // restore implicit batch dimension + + TRT_ShapedWeights shape = inputs.at(1).weights(); + + TFAttrs attrs(node_def); + + auto padding_type = attrs.get("Tshape"); + + if (shape.shape_.nbDims != 1) + return tensorflow::errors::InvalidArgument( + "reshape new shape is not 1 dimensional, at " + node_def.name()); + + // Only expect to handle INT32 as attributes for now + if (padding_type != tensorflow::DataType::DT_INT32) + return tensorflow::errors::Unimplemented( + "reshape new shape supports only DT_INT32, at " + node_def.name()); + + auto shape_data = static_cast(const_cast(shape.GetValues())); + + if (shape_data[0] != -1) + return tensorflow::errors::InvalidArgument( + "reshape new shape first dimension is not -1, at " + node_def.name()); + + auto shape_num_dims = shape.shape_.d[0]; + VLOG(2) << "shape dimensions: " << shape_num_dims; + int volume_w = 1; + for (int i = 1; i < shape.shape_.d[0]; i++) volume_w *= shape_data[i]; + + int volume_t = 1; + for (int i = 0; i < dims.nbDims; i++) volume_t *= dims.d[i]; + + VLOG(2) << "volume: " << volume_t << " volume weights: " << volume_w; + if (volume_w != volume_t) + return tensorflow::errors::InvalidArgument( + "volume does not agree between tensor and new shape, at " + + node_def.name()); + + nvinfer1::IShuffleLayer* layer = + ctx.network()->addShuffle(*const_cast(tensor)); + + nvinfer1::Dims reshape_dims; + VLOG(2) << "new dimension: " << shape_num_dims - 1; + reshape_dims.nbDims = shape_num_dims - 1; + for (int32_t i = 0; i < reshape_dims.nbDims; ++i) { + reshape_dims.d[i] = shape_data[i + 1]; + } + layer->setReshapeDimensions(reshape_dims); + VLOG(2) << "new dimension: " << shape_num_dims - 1; + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + auto dims_output = output_tensor->getDimensions(); + VLOG(2) << "output tensor dimension:" << dims_output.nbDims; + 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_["DepthwiseConv2dNative"] = ConvertConv2DDepthwise; + op_registry_["Relu"] = ConvertActivation; + op_registry_["MaxPool"] = ConvertPool; + op_registry_["AvgPool"] = ConvertPool; + // This could be really handled as ConvertBinary + op_registry_["BiasAdd"] = ConvertScale; + op_registry_["Const"] = ConvertConst; + // TODO(ben,jie): this is a temp hack. + op_registry_["Identity"] = ConvertIdentity; // Identity should be removed + + // 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 + + op_registry_["ConcatV2"] = ConvertConcat; + op_registry_["MatMul"] = ConvertMatMul; + op_registry_["Reshape"] = ConvertReshape; + op_registry_["FusedBatchNorm"] = ConvertFusedBatchNorm; + op_registry_["FusedBatchNormV2"] = ConvertFusedBatchNorm; +} + +} // namespace +tensorflow::Status GetTensorRTGraph(tensorrt::convert::SubGraphParams& s) { + return tensorflow::errors::Unimplemented("Not implemented yet"); +} +tensorflow::Status ConvertCalibrationNodeToEngineNode( + tensorflow::Graph& graph, tensorflow::Node* c_node) { + const auto ndef = c_node->def(); + + TFAttrs attrs(ndef); + std::vector segment_nodes( + attrs.get>("segment_nodes")); + std::vector output_nodes( + attrs.get>("segment_output_names")); + std::vector input_names( + attrs.get>("input_names")); + string res_name = attrs.get("resource_name"); + VLOG(1) << "Node name " << c_node->name() << " res_name " << res_name; + string engine_name = "my_trt_op"; + { + const auto node_id = tensorflow::str_util::Split(res_name, "_"); + engine_name += node_id.back(); + } + std::map node_maps; + + for (auto n : graph.op_nodes()) { + node_maps.insert({n->name(), n}); + } + VLOG(1) << "Output Nodes:"; + std::vector out_types; + std::vector out_edges; + for (auto& i : output_nodes) { + auto node_port = tensorflow::str_util::Split(i, ":"); + VLOG(1) << " " << i << " in graph " << node_maps.count(i); + auto out_node_name = node_port.at(0); + if (node_port.size() > 1) { + VLOG(1) << "Multi port output" << node_port.at(0) << " " + << node_port.at(1) << " size=" << node_port.size(); + } + auto node_it = node_maps.find(out_node_name); + if (node_it != node_maps.end()) { + tensorflow::Node* out_node = node_it->second; + int port = 0; + if (node_port.size() == 2) { + port = std::strtoul(node_port.at(1).c_str(), nullptr, 10); + out_types.push_back(out_node->output_type(port)); + } else { + out_types.push_back(out_node->output_type(0)); + } + for (auto out_edge : out_node->out_edges()) { + if (out_edge->src_output() == port) { + out_edges.push_back(out_edge); + break; + } + } + } else { + LOG(WARNING) << " couldn't find output node " << out_node_name; + } + } + VLOG(1) << "Input Nodes:"; + for (auto& i : input_names) { + VLOG(1) << " " << i << " in graph " << node_maps.count(i); + } + auto trt_rm = tensorflow::tensorrt::TRTResourceManager::instance(); + auto resmgr = trt_rm->getManager("TRTCalibOps"); + tensorflow::tensorrt::TRTCalibrationResource* calib_res = nullptr; + auto status = resmgr->Lookup(res_name, res_name, &calib_res); + if (!status.ok() || !calib_res->calibrator_) { + return tensorflow::errors::FailedPrecondition( + "You must run calibration" + " and inference conversion in the same proces"); + } + + calib_res->calibrator_->setDone(); + calib_res->thr_->join(); + delete calib_res->thr_; + if (!calib_res->engine_) { + LOG(FATAL) << "Calibration failed!, engine is nullptr. Did you run " + "calibration graph?"; + } + auto weight_rmgr = trt_rm->getManager("WeightStore"); + TF_CHECK_OK(weight_rmgr->Delete( + res_name, res_name)); + auto engine_plan = calib_res->engine_->serialize(); + calib_res->engine_->destroy(); + calib_res->network_->destroy(); + calib_res->builder_->destroy(); + calib_res->thr_ = nullptr; + calib_res->engine_ = nullptr; + calib_res->builder_ = nullptr; + tensorflow::NodeDefBuilder op_builder(engine_name, "TRTEngineOp"); + std::vector income_edges; + for (const auto in_edge : c_node->in_edges()) { + auto src = in_edge->src(); + int dest_port = in_edge->dst_input(); + income_edges.emplace_back(src->name(), in_edge->src_output(), + c_node->input_type(dest_port)); + } + tensorflow::gtl::ArraySlice input_list( + income_edges); + op_builder.Input(input_list); + tensorflow::NodeDef engine_node; + const char* engine_plan_data = static_cast(engine_plan->data()); + string engine_plan_string(engine_plan_data, + engine_plan_data + engine_plan->size()); + status = op_builder.Attr("serialized_engine", engine_plan_string) + .Attr("input_nodes", input_names) + .Attr("output_nodes", output_nodes) + .Attr("OutT", out_types) + .Finalize(&engine_node); + if (!status.ok()) { + LOG(ERROR) << "Engine Node creation failed"; + return status; + } + auto trt_engine_node = graph.AddNode(engine_node, &status); + TF_CHECK_OK(status); + for (size_t i = 0; i < out_edges.size(); i++) { + VLOG(1) << "Connecting trt_engine_node output " << i << " with " + << out_edges.at(i)->dst()->name() << " port " + << out_edges.at(i)->dst_input(); + TF_RETURN_IF_ERROR(graph.UpdateEdge(trt_engine_node, i, + out_edges.at(i)->dst(), + out_edges.at(i)->dst_input())); + } + VLOG(1) << "Segment nodes:"; + for (auto& i : segment_nodes) { + VLOG(1) << " " << i << " in graph " << node_maps.count(i); + auto it = node_maps.find(i); + if (it != node_maps.end()) { + graph.RemoveNode(it->second); + } + } + graph.RemoveNode(c_node); + return tensorflow::Status::OK(); +} + +tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) { + // Visit nodes in reverse topological order and construct the TRT network. + + // Toposort + std::vector order_vec; + tensorflow::GetPostOrder(s.graph, &order_vec); + // Select just the subgraph + std::list order; + for (tensorflow::Node* node : order_vec) { + if (s.subgraph_node_ids.count(node->id())) { + order.push_front(node); // we want topological order to construct the + // network layer by layer + } + } + // topological order is needed to build TRT network + static int static_id = 0; + string subgraph_name_scope; + if (!order.empty()) { + subgraph_name_scope = order.front()->name(); + } + for (const tensorflow::Node* node : order) { + subgraph_name_scope = GetCommonNameScope(subgraph_name_scope, node->name()); + } + // TODO(sami,ben,jie): proper naming! + string calib_op_name = + StrCat(subgraph_name_scope, "my_trt_calib_op_", static_id); + string engine_name = StrCat(subgraph_name_scope, "my_trt_op", static_id); + static_id++; + auto trt_rmgr = tensorflow::tensorrt::TRTResourceManager::instance(); + auto op_rmgr = trt_rmgr->getManager("TRTCalibOps"); + auto op_res = new tensorflow::tensorrt::TRTCalibrationResource(); + TF_CHECK_OK(op_rmgr->Create(calib_op_name, calib_op_name, op_res)); + op_res->logger_ = new tensorflow::tensorrt::Logger(); + op_res->builder_ = nvinfer1::createInferBuilder(*(op_res->logger_)); + + if (!op_res->builder_) { + return tensorflow::errors::Internal( + "failed to create TensorRT builder object"); + } + + op_res->network_ = op_res->builder_->createNetwork(); + if (!op_res->network_) { + return tensorflow::errors::Internal( + "failed to create TensorRT network object"); + } + + // Build the network + auto weight_rmgr = trt_rmgr->getManager("WeightStore"); + auto ws = new tensorflow::tensorrt::TRTWeightStore(); + TF_CHECK_OK(weight_rmgr->Create(calib_op_name, calib_op_name, ws)); + Converter converter(op_res->network_, ws, s.precision_mode == FP16MODE); + std::vector input_names; + std::vector input_dtypes; + for (const std::pair& input : s.input_inds) { + VLOG(2) << "parsing input. Node id= " << input.first; + int node_id = input.first; + int output_idx = input.second; + tensorflow::Node* node = s.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 (!s.graph_properties.HasOutputProperties(node_name)) + return tensorflow::errors::Internal("failed to find input node: " + + node_name); + + auto op_info_vec = s.graph_properties.GetOutputProperties(node_name); + if (static_cast(op_info_vec.size()) < output_idx) + return tensorflow::errors::Internal( + "accessing output index of: ", output_idx, ", at node: ", node_name, + "with output entry from shape_map: ", 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: " << output_idx + << ", at node: " << node_name + << "with output entry from shape_map: " << 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 = StrCat(node_name, ":", 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; + int trt_engine_op_output_idx = 0; + for (const std::pair& output : s.output_inds) { + int node_id = output.first; + int output_idx = output.second; + tensorflow::Node* node = s.graph.FindNodeId(node_id); + string op_name = node->name(); + string tensor_name = op_name; + + s.output_edge_map->insert( + {trt_engine_op_output_idx == 0 + ? engine_name + : StrCat(engine_name, ":", trt_engine_op_output_idx), + {output_idx, tensor_name}}); + trt_engine_op_output_idx++; + if (output_idx != 0) { + tensor_name = StrCat(tensor_name, ":", output_idx); + } + VLOG(1) << "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"; + + // Build the engine + op_res->builder_->setMaxBatchSize(s.max_batch_size); + op_res->builder_->setMaxWorkspaceSize(s.max_workspace_size_bytes); + + // Build the TRT op + // TODO(sami,ben,jie): proper naming! + tensorflow::NodeDefBuilder op_builder(calib_op_name, "TRTCalibOp"); + std::vector income_edges; + for (size_t i = 0; i < input_names.size(); ++i) { + int output_idx = s.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)); + VLOG(1) << calib_op_name << " input " << i << " = " << input_names.at(i) + << ":" << output_idx + << " dType= " << tensorflow::DataTypeString(input_dtypes.at(i)); + income_edges.push_back(incoming_edge); + } + tensorflow::gtl::ArraySlice input_list( + income_edges); + op_builder.Input(input_list); + std::vector segment_names; + segment_names.reserve(s.subgraph_node_ids.size()); + for (int i : s.subgraph_node_ids) { + auto node = s.graph.FindNodeId(i); + segment_names.push_back(node->name()); + } + LOG(INFO) << "finished op preparation"; + + auto status = op_builder.Attr("segment_nodes", segment_names) + .Attr("input_names", input_names) + .Attr("segment_output_names", output_names) + .Attr("resource_name", calib_op_name) + .Finalize(s.trt_node); + + LOG(INFO) << status.ToString(); + LOG(INFO) << "finished op building"; + + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertSubGraphToTensorRTNodeDef( + tensorrt::convert::SubGraphParams& s) { + // Visit nodes in reverse topological order and construct the TRT network. + + // Toposort + std::vector order_vec; + tensorflow::GetPostOrder(s.graph, &order_vec); + // Select just the subgraph + std::list order; + for (tensorflow::Node* node : order_vec) { + if (s.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"); + } + + string subgraph_name_scope; + if (!order.empty()) { + subgraph_name_scope = order.front()->name(); + } + for (const tensorflow::Node* node : order) { + subgraph_name_scope = GetCommonNameScope(subgraph_name_scope, node->name()); + } + static int static_id = 0; + // TODO(sami,ben,jie): proper naming! + string engine_name = StrCat(subgraph_name_scope, "my_trt_op"); + engine_name = StrCat(engine_name, static_id++); + auto trt_rmgr = tensorflow::tensorrt::TRTResourceManager::instance(); + auto weight_rmgr = trt_rmgr->getManager("WeightStore"); + auto ws = new tensorflow::tensorrt::TRTWeightStore(); + TF_CHECK_OK(weight_rmgr->Create(engine_name, engine_name, ws)); + + // Build the network + Converter converter(trt_network.get(), ws, s.precision_mode == FP16MODE); + + std::vector input_names; + std::vector input_dtypes; + for (const std::pair& input : s.input_inds) { + VLOG(2) << "parsing input!!!!!"; + int node_id = input.first; + int output_idx = input.second; + tensorflow::Node* node = s.graph.FindNodeId(node_id); + auto node_name = node->name(); + // input_names should use the node name in the graph + // here it should be the input tensor name -> matching the binding + // insert original node name without port + auto tensor_name = node_name; + if (output_idx != 0) { + tensor_name = StrCat(tensor_name, ":", output_idx); + } + + VLOG(2) << "input name: " << node_name << " tensor_name: " << tensor_name + << " idx: " << output_idx; + + auto shape_inference_node_name = node_name; + auto shape_inference_output_idx = output_idx; + // rewire the shape inference to original node in the graph + if (s.output_edge_map->count(tensor_name)) { + shape_inference_node_name = s.output_edge_map->at(tensor_name).second; + shape_inference_output_idx = s.output_edge_map->at(tensor_name).first; + } + if (shape_inference_output_idx < 0) continue; + VLOG(2) << "shapeinference name: " << shape_inference_node_name + << " idx: " << shape_inference_output_idx; + + if (!s.graph_properties.HasOutputProperties(shape_inference_node_name)) + return tensorflow::errors::Internal("failed to find input node: " + + shape_inference_node_name); + + auto op_info_vec = + s.graph_properties.GetOutputProperties(shape_inference_node_name); + if (static_cast(op_info_vec.size()) <= shape_inference_output_idx) + return tensorflow::errors::Internal( + "accessing output index of: ", shape_inference_output_idx, + ", at node: ", shape_inference_node_name, + " with output entry from shape_map: ", op_info_vec.size()); + + auto op_info = op_info_vec.at(shape_inference_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: " << output_idx + << ", at node: " << node_name + << " with output entry from shape_map: " << 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; + + // TODO(jie): TRT 3.x only support 4 dimensional input tensor. + // update the code once TRT 4.0 comes out. + if (op_info.shape().dim_size() != 4) + return tensorflow::errors::Unimplemented("require 4 dimensional input"); + + 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 = StrCat(node_name, ":", output_idx); + } + + input_names.push_back(input_tensor_name); + 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; + int trt_engine_op_output_idx = 0; + for (const std::pair& output : s.output_inds) { + int node_id = output.first; + int output_idx = output.second; + tensorflow::Node* node = s.graph.FindNodeId(node_id); + string op_name = node->name(); + string tensor_name = op_name; + + s.output_edge_map->insert( + {trt_engine_op_output_idx == 0 + ? engine_name + : StrCat(engine_name, ":", trt_engine_op_output_idx), + {output_idx, tensor_name}}); + trt_engine_op_output_idx++; + if (output_idx != 0) + tensorflow::strings::StrAppend(&tensor_name, ":", 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"; + + // Build the engine + trt_builder->setMaxBatchSize(s.max_batch_size); + trt_builder->setMaxWorkspaceSize(s.max_workspace_size_bytes); + VLOG(0) << "Max batch size= " << s.max_batch_size + << " max workspace size= " << s.max_workspace_size_bytes; + if (s.precision_mode == FP16MODE) { + trt_builder->setHalf2Mode(true); + VLOG(0) << "Using FP16 precision mode"; + } + LOG(INFO) << "starting build engine"; + string engine_plan_string; + { + auto trt_engine = + infer_object(trt_builder->buildCudaEngine(*converter.network())); + VLOG(0) << "Built network"; + if (trt_engine.get() == nullptr) { + return tensorflow::errors::Internal("Engine building failure"); + } + 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()); + } + TF_RETURN_IF_ERROR(weight_rmgr->Delete( + engine_name, engine_name)); + LOG(INFO) << "finished engine " << engine_name; + + // Build the TRT op + tensorflow::NodeDefBuilder op_builder(engine_name, "TRTEngineOp"); + std::vector income_edges; + VLOG(2) << "input edge size: " << input_names.size(); + for (size_t i = 0; i < input_names.size(); ++i) { + VLOG(2) << "input edges: " << i << " " << input_names.at(i); + int output_idx = s.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(s.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..954a1e72f8604371fc00e088a67b4d411314dda6 --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.h @@ -0,0 +1,87 @@ +/* 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 +#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 { + +const int FP32MODE = 0; +const int FP16MODE = 1; +const int INT8MODE = 2; + +struct SubGraphParams { + SubGraphParams( + tensorflow::Graph& inp_graph, + const std::set& subgraph_node_id_numbers, + const std::vector>& input_indices, + const std::vector>& output_indices, + size_t max_supported_batch_size, size_t max_consumed_workspace_size_bytes, + const tensorflow::grappler::GraphProperties& current_graph_properties, + std::unordered_map>* output_edges, + tensorflow::NodeDef* constructed_trt_node, + int engine_precision_mode = FP32MODE) + : graph(inp_graph), + subgraph_node_ids(subgraph_node_id_numbers), + input_inds(input_indices), + output_inds(output_indices), + max_batch_size(max_supported_batch_size), + max_workspace_size_bytes(max_consumed_workspace_size_bytes), + graph_properties(current_graph_properties), + output_edge_map(output_edges), + trt_node(constructed_trt_node), + precision_mode(engine_precision_mode) {} + + 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_properties; + std::unordered_map>* output_edge_map; + tensorflow::NodeDef* trt_node; + const int precision_mode; +}; + +// TODO(sami): Replace references with const reference or pointers +tensorflow::Status ConvertSubGraphToTensorRTNodeDef(SubGraphParams& params); +tensorflow::Status InjectCalibrationNode(SubGraphParams& params); +tensorflow::Status ConvertCalibrationNodeToEngineNode(tensorflow::Graph& graph, + tensorflow::Node* c_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_calib_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..aea44fd8a2fcc4c359a6cb0c98ae34711708326e --- /dev/null +++ b/tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc @@ -0,0 +1,136 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/contrib/tensorrt/kernels/trt_calib_op.h" +#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h" +#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" +#include "tensorflow/contrib/tensorrt/resources/trt_resources.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/stream_executor.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda_runtime_api.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { + +TRTCalibOp::TRTCalibOp(OpKernelConstruction* context) : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("segment_nodes", &segment_nodes_)); + OP_REQUIRES_OK(context, context->GetAttr("input_names", &input_names_)); + OP_REQUIRES_OK(context, context->GetAttr("resource_name", &resource_name_)); +}; + +#define TYPECASE(dt, X, Y) \ + case dt: { \ + return (void*)X->flat::Type>().data(); \ + } + +void* GetTensorAddress(const Tensor* tensor_ptr) { + auto tensor_type = tensor_ptr->dtype(); + switch (tensor_type) { + TYPECASE(tensorflow::DT_FLOAT, tensor_ptr, dest_ptr); + TYPECASE(tensorflow::DT_HALF, tensor_ptr, dest_ptr); + TYPECASE(tensorflow::DT_INT8, tensor_ptr, dest_ptr); + default: { + LOG(FATAL) << "Unsupported Data type " + << tensorflow::DataTypeString(tensor_type); + return nullptr; + } + } +} + +void TRTCalibOp::Compute(tensorflow::OpKernelContext* ctx) { + // TODO(aaroey): make sure ctx->resource_mgr() is used in future PR. + auto trt_rm = tensorflow::tensorrt::TRTResourceManager::instance(); + auto res_mgr = trt_rm->getManager("TRTCalibOps"); + tensorflow::tensorrt::TRTCalibrationResource* calib_res = nullptr; + auto status = res_mgr->Lookup(resource_name_, resource_name_, &calib_res); + + if (!status.ok()) { + ctx->SetStatus(status); + return; + } + int num_inputs = ctx->num_inputs(); + // first run instantiate calibrator + if (calib_res->calibrator_ == nullptr) { + dev_tensors_.resize(num_inputs); + int batch_size = ctx->input(0).dim_size(0); + VLOG(1) << " Constructing calibrator"; + for (int i = 0; i < num_inputs; i++) { + // allocate workspace on device for inputs + const tensorflow::Tensor& t = ctx->input(i); + OP_REQUIRES_OK(ctx, + ctx->allocate_persistent(t.dtype(), t.shape(), + &dev_tensors_.at(i), nullptr)); + const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx); + CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes()); + void* device_address = GetTensorAddress(device_tensor); + device_buffers_.emplace(input_names_.at(i), + std::pair( + device_address, device_tensor->TotalBytes())); + } + + calib_res->calibrator_ = + new TRTInt8Calibrator(device_buffers_, batch_size, resource_name_); + string label(resource_name_); + calib_res->thr_ = new std::thread([calib_res, label]() { + VLOG(1) << "Starting calibration thread, Calibration Resource @ " + << calib_res; + calib_res->builder_->setInt8Calibrator(calib_res->calibrator_); + calib_res->builder_->setInt8Mode(true); + calib_res->engine_ = calib_res->builder_->buildCudaEngine( + *calib_res->network_); // will loop until we terminate calibrator + VLOG(1) << "Calibration loop terminated " << label; + }); + VLOG(1) << "initialized calibrator resource"; + } // calibrator initialized + + // Pass input data to calibrator + std::unordered_map input_data; + for (int i = 0; i < num_inputs; i++) { + const Tensor& t = ctx->input(i); + void* data_address = GetTensorAddress(&t); + const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx); + CHECK_EQ(t.TotalBytes(), + device_tensor->TotalBytes()); // use the tensor so FW keeps it + input_data.emplace(input_names_.at(i), data_address); + ctx->set_output(i, t); + } + VLOG(2) << "Filled map for sending"; + // copied from cuda_kernel_helper since it seems only valid in *.cu.cc files + const cudaStream_t* stream = CHECK_NOTNULL( + reinterpret_cast(ctx->op_device_context() + ->stream() + ->implementation() + ->CudaStreamMemberHack())); + calib_res->calibrator_->setBatch(input_data, *stream); + VLOG(2) << "Passed calibration data"; + // TODO(aaroey): make sure we wait for the completion of calibration on the + // last batch in future PR. +}; + +#undef TYPECASE + +REGISTER_KERNEL_BUILDER(Name("TRTCalibOp").Device(DEVICE_GPU), TRTCalibOp); + +} // namespace tensorrt +} // namespace tensorflow +#endif +#endif diff --git a/tensorflow/contrib/tensorrt/kernels/trt_calib_op.h b/tensorflow/contrib/tensorrt/kernels/trt_calib_op.h new file mode 100644 index 0000000000000000000000000000000000000000..23df9db32f077a080eaff7479fcbe90d6a504c42 --- /dev/null +++ b/tensorflow/contrib/tensorrt/kernels/trt_calib_op.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_KERNELS_TRT_CALIB_OP_H +#define TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_CALIB_OP_H + +#include +#include +#include +#include +#include +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +namespace tensorflow { +namespace tensorrt { +// TODO(sami): Convert this to async kernel! +class TRTCalibOp : public OpKernel { + public: + explicit TRTCalibOp(OpKernelConstruction* context); + + void Compute(OpKernelContext* context) override; + + private: + string resource_name_; + std::vector segment_nodes_; + std::vector input_names_; + std::vector shapes_; + std::unordered_map> device_buffers_; + std::vector dev_tensors_; +}; +} // namespace tensorrt +} // namespace tensorflow +#endif +#endif +#endif // TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_CALIB_OP_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..b32371b642f38b0851955a4a3beab97b86e1f6a0 --- /dev/null +++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc @@ -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. +==============================================================================*/ +#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 { +static ::tensorflow::tensorrt::Logger logger; +namespace gpu = ::perftools::gputools; +using IRuntime = nvinfer1::IRuntime; +using Dims = nvinfer1::Dims; + +namespace tensorrt { + +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 + // TODO(jie): cudaSetDevice make sure trt engine is allocated on the same + // gpu where the input/output is also located. + int gpu_id = context->device()->tensorflow_gpu_device_info()->gpu_id; + cudaSetDevice(gpu_id); + int device; + cudaGetDevice(&device); + if (gpu_id != device) LOG(FATAL) << "set device failed!"; + + // 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 + + 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; + 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); + if (num_batch > trt_engine_ptr_->getMaxBatchSize()) { + LOG(FATAL) << "input tensor batch larger than max_batch_size: " + << trt_engine_ptr_->getMaxBatchSize(); + } + } else if (num_batch != input_shape.dim_size(0)) { + LOG(FATAL) << "input data inconsistent batch size"; + 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; + } + } + + 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())); + + // TODO(jie): trt enqueue does not return error + auto ret = trt_execution_context_ptr_->enqueue(num_batch, &buffers[0], + *stream, nullptr); + VLOG(2) << "enqueue returns: " << ret; + // sync should be done by TF. +} + +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..dda0dc9e712eb726800abfb6084f4f708d04825b --- /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) << name_ << " " << msg; + break; + } + case Severity::kWARNING: { + LOG(WARNING) << name_ << " " << msg; + break; + } + case Severity::kERROR: { + LOG(ERROR) << name_ << " " << msg; + break; + } + case Severity::kINTERNAL_ERROR: { + LOG(FATAL) << name_ << " " << 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..7f3544f8cfda8dce13881e1f8f4388b640e315f4 --- /dev/null +++ b/tensorflow/contrib/tensorrt/log/trt_logger.h @@ -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. +==============================================================================*/ + +#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 { + public: + Logger(string name = "DefaultLogger") : name_(name){}; + void log(nvinfer1::ILogger::Severity severity, const char* msg) override; + + private: + 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_calib_op.cc b/tensorflow/contrib/tensorrt/ops/trt_calib_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..4835e5065068ec7a59995eb7f6126b31aecf6704 --- /dev/null +++ b/tensorflow/contrib/tensorrt/ops/trt_calib_op.cc @@ -0,0 +1,37 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" +namespace tensorflow { + +REGISTER_OP("TRTCalibOp") + .Attr("segment_nodes: list(string)") // names of the ops in segment + .Attr("segment_output_names: list(string)") // names of the output ops in + // segment + .Attr("input_names: list(string)") // names of the inputs for + // passing into tensorrt + .Attr("resource_name: string") + .Attr("InT: list({int8, float16, float32})") + .Input("in_tensor: InT") + .Output("out_tensor: InT") + .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) { + for (int i = 0; i < c->num_inputs(); i++) { + c->set_output(i, c->input(i)); + } + return Status::OK(); + }); + +} // namespace tensorflow 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/bayesflow/python/ops/custom_grad.py b/tensorflow/contrib/tensorrt/python/__init__.py similarity index 59% rename from tensorflow/contrib/bayesflow/python/ops/custom_grad.py rename to tensorflow/contrib/tensorrt/python/__init__.py index ca1ecb9c40204c3c723fa3423cfe148e823adc28..0b2321b5fc7bcbd53c01d1c97cafcfcb229a83ef 100644 --- a/tensorflow/contrib/bayesflow/python/ops/custom_grad.py +++ b/tensorflow/contrib/tensorrt/python/__init__.py @@ -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. @@ -11,24 +11,15 @@ # WITHOUT 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 for specifying custom gradients. - -See ${python/contrib.bayesflow.custom_gradient}. -""" +# ============================================================================= +"""Exposes the python wrapper for TensorRT graph transforms.""" 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.custom_grad_impl import * -# pylint: enable=wildcard-import -from tensorflow.python.util.all_util import remove_undocumented - -_allowed_symbols = [ - 'custom_gradient', -] - -remove_undocumented(__name__, _allowed_symbols) +# 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 calib_graph_to_infer_graph +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..338475d90ea55ab2c1bb8df77f27a71a4a36a5dd --- /dev/null +++ b/tensorflow/contrib/tensorrt/python/trt_convert.py @@ -0,0 +1,163 @@ +# 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 calib_convert +from tensorflow.contrib.tensorrt.wrap_conversion import trt_convert +from tensorflow.core.framework import graph_pb2 +from tensorflow.core.protobuf import rewriter_config_pb2 +from tensorflow.python.framework import errors +from tensorflow.python.framework import errors_impl as _impl +from tensorflow.python.framework import meta_graph +from tensorflow.python.framework import ops +from tensorflow.python.grappler import tf_optimizer +from tensorflow.python.util import compat +# pylint: enable=unused-import,line-too-long + + +# 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, + precision_mode="FP32", + minimum_segment_size=3): + """Python wrapper for the TRT transformation. + + 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) + precision_mode: one of 'FP32', 'FP16' and 'INT8' + minimum_segment_size: the minimum number of nodes required for a subgraph to + be replaced by TRTEngineOp. + + Returns: + New GraphDef with TRTEngineOps placed in graph replacing subgraphs. + + Raises: + ValueError: if the provided precision mode is invalid. + RuntimeError: if the returned status message is malformed. + """ + supported_precision_modes = {"FP32": 0, "FP16": 1, "INT8": 2} + if precision_mode.upper() not in supported_precision_modes: + raise ValueError(("precision mode '{}' is not supported." + "It should be one of {}").format( + precision_mode, "{'FP32', 'FP16', 'INT8'}")) + mode = supported_precision_modes[precision_mode.upper()] + + 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, mode, minimum_segment_size) + 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 + + +def calib_graph_to_infer_graph(calibration_graph_def): + """Convert an existing calibration graph to inference graph. + + Args: + calibration_graph_def: the calibration GraphDef object with calibration data + Returns: + New GraphDef with TRTEngineOps placed in graph replacing calibration nodes. + Raises: + RuntimeError: if the returned status message is malformed. + """ + + def py2string(inp): + return inp + + def py3string(inp): + return inp.decode("utf-8") + + if _six.PY2: + to_string = py2string + else: + to_string = py3string + + graph_str = calibration_graph_def.SerializeToString() + out = calib_convert(graph_str) + status = to_string(out[0]) + output_graph_def_string = out[1] + del graph_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/resources/trt_int8_calibrator.cc b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc new file mode 100644 index 0000000000000000000000000000000000000000..dc7c93f869f5ef7c8eaa2a87eed26cfe69597fdb --- /dev/null +++ b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc @@ -0,0 +1,129 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h" + +#include +#include +#include + +#include "tensorflow/core/platform/logging.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda_runtime_api.h" + +namespace tensorflow { +namespace tensorrt { + +// set the batch size before constructing the thread to execute engine +int TRTInt8Calibrator::getBatchSize() const { return batch_size_; } + +TRTInt8Calibrator::TRTInt8Calibrator( + const std::unordered_map>& dev_buffers, + int batch_size, string engine_name) + : batch_size_(batch_size), + done_(false), + dev_buffers_(dev_buffers), + calib_running_(false), + batch_is_set_(false), + engine_name_(engine_name) {} + +bool TRTInt8Calibrator::setBatch(const std::unordered_map& data, + const cudaStream_t stream) { + tensorflow::mutex_lock lock(cond_mtx_); + while ((calib_running_ || batch_is_set_) && + !done_) { // wait while calibration is running + cond_.wait(lock); + } + if (done_) return false; + CHECK(!calib_running_ && !batch_is_set_); + VLOG(1) << "Set Batch Waiting finished"; + for (const auto it : data) { + auto devptr = dev_buffers_.find(it.first); + if (devptr == dev_buffers_.end()) { + LOG(FATAL) << "FATAL " << engine_name_ << " input name '" << it.first + << "' does not match with the buffer names"; + } + const auto& d = devptr->second; + + // TODO(aaroey): we should not use sync copy on default stream. Make sure + // stream->ThenMemcpy() is used in future PRs. + // TODO(sami,aaroey): Need to figure out a way to ensure synchronization + // between stream, perhaps using a tensor? + auto status = cudaMemcpyAsync(d.first, it.second, d.second, + cudaMemcpyDeviceToDevice, stream); + if (status != cudaSuccess) { + LOG(FATAL) << "cudaMemcpy " << engine_name_ << " for '" << it.first + << "' failed with " << status; + } + } + + // TODO(Sami, aaorey): Find an alternative way! + cudaStreamSynchronize( + stream); // we have to wait for the stream before returning! + batch_is_set_ = true; + cond_.notify_all(); + return true; +} + +bool TRTInt8Calibrator::getBatch(void** bindings, const char** names, + int num_bindings) { + tensorflow::mutex_lock lock(cond_mtx_); + calib_running_ = false; + cond_.notify_all(); + while ((!batch_is_set_ && !done_)) { // wait until new batch arrives + cond_.wait(lock); + + } + if (done_) { + return false; + } + + for (int i = 0; i < num_bindings; i++) { + auto it = dev_buffers_.find(names[i]); + if (it == dev_buffers_.end()) { + LOG(FATAL) << "Calibration engine asked for unknown tensor name '" + << names[i] << "' at position " << i; + } + + bindings[i] = it->second.first; + } + batch_is_set_ = false; + calib_running_ = true; + return true; +} + +const void* TRTInt8Calibrator::readCalibrationCache(std::size_t& length) { + return nullptr; +} + +void TRTInt8Calibrator::setDone() { + tensorflow::mutex_lock lock(cond_mtx_); + done_ = true; + cond_.notify_all(); +} + +void TRTInt8Calibrator::writeCalibrationCache(const void* ptr, + std::size_t length) {} +TRTInt8Calibrator::~TRTInt8Calibrator() { + VLOG(1) << "Destroying calibrator for " << engine_name_; +} + +} // namespace tensorrt +} // namespace tensorflow + +#endif +#endif diff --git a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h new file mode 100644 index 0000000000000000000000000000000000000000..d77aa2c5ab184756adaee38f88180b3c128ebe03 --- /dev/null +++ b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.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_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_ + +#include +#include +#include +#include +#include "tensorflow/core/platform/mutex.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +#include "cuda/include/cuda_runtime_api.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +// This class provides a 1 element queue to match TFs push model to +// TRTs pull model for calibration. When TRT implements a means for +// a push calibration This class should be updated accordingly + +struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator { + public: + TRTInt8Calibrator( + const std::unordered_map>& dev_buffers, + int batch_size, string engine_name); + int getBatchSize() const override; + bool getBatch(void* bindings[], const char* names[], + int num_bindings) override; + bool setBatch(const std::unordered_map& data, + const cudaStream_t stream); + void setDone(); + const void* readCalibrationCache(std::size_t& length) override; + void writeCalibrationCache(const void* ptr, std::size_t length) override; + ~TRTInt8Calibrator(); + + private: + const int batch_size_; + tensorflow::mutex cond_mtx_; // mutex for condition_variable + tensorflow::condition_variable cond_; // condition variable to implement + // producer-consumer queue for + // calibration + bool done_; + const std::unordered_map> + dev_buffers_; // map to keep tensorrt input buffers and sizes keyed with + // buffer names + bool calib_running_; + bool batch_is_set_; + string engine_name_; +}; + +} // namespace tensorrt +} // namespace tensorflow + +#endif +#endif +#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_ diff --git a/tensorflow/contrib/tensorrt/resources/trt_resource_manager.cc b/tensorflow/contrib/tensorrt/resources/trt_resource_manager.cc new file mode 100644 index 0000000000000000000000000000000000000000..e663eed4dd6704e2f41bde1dfabd411e86669ecd --- /dev/null +++ b/tensorflow/contrib/tensorrt/resources/trt_resource_manager.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" +#include "tensorflow/core/platform/logging.h" + +namespace tensorflow { +namespace tensorrt { + +std::shared_ptr +tensorflow::tensorrt::TRTResourceManager::getManager(const string& op_name) { + // mutex is held for lookup only. Most instantiations where mutex will be held + // longer will be during op creation and should be ok. + tensorflow::mutex_lock lock(map_mutex_); + auto s = managers_.find(op_name); + if (s == managers_.end()) { + auto it = managers_.emplace( + op_name, std::make_shared(op_name)); + VLOG(1) << "Returning a new manager " << op_name; + return it.first->second; + } + VLOG(1) << "Returning old manager " << op_name; + return s->second; +} + +} // namespace tensorrt +} // namespace tensorflow diff --git a/tensorflow/contrib/tensorrt/resources/trt_resource_manager.h b/tensorflow/contrib/tensorrt/resources/trt_resource_manager.h new file mode 100644 index 0000000000000000000000000000000000000000..5f8ad491d3c13e8911b0b95c3e95e19afe4d59c0 --- /dev/null +++ b/tensorflow/contrib/tensorrt/resources/trt_resource_manager.h @@ -0,0 +1,49 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_ +#include + +#include +#include +#include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/platform/mutex.h" + +namespace tensorflow { +namespace tensorrt { + +class TRTResourceManager { + TRTResourceManager() = default; + + public: + static std::shared_ptr instance() { + static std::shared_ptr instance_( + new TRTResourceManager); + return instance_; + } + // returns a manager for given op, if it doesn't exists it creates one + std::shared_ptr getManager(const string& op_name); + + private: + std::unordered_map> + managers_; + tensorflow::mutex map_mutex_; +}; + +} // namespace tensorrt +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCE_TRT_RESOURCE_MANAGER_H_ diff --git a/tensorflow/contrib/tensorrt/resources/trt_resources.h b/tensorflow/contrib/tensorrt/resources/trt_resources.h new file mode 100644 index 0000000000000000000000000000000000000000..3c85968ae7acf5c5fc567be6805a5d226b1094c7 --- /dev/null +++ b/tensorflow/contrib/tensorrt/resources/trt_resources.h @@ -0,0 +1,95 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRTRESOURCES_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRTRESOURCES_H_ + +#include +#include +#include +#include +#include +#include "tensorflow/contrib/tensorrt/log/trt_logger.h" +#include "tensorflow/core/framework/resource_mgr.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +class TRTCalibrationResource : public tensorflow::ResourceBase { + public: + TRTCalibrationResource() + : calibrator_(nullptr), + builder_(nullptr), + network_(nullptr), + engine_(nullptr), + logger_(nullptr), + thr_(nullptr) {} + string DebugString() override { + std::stringstream oss; + oss << " Calibrator = " << std::hex << calibrator_ << std::dec << std::endl + << " Builder = " << std::hex << builder_ << std::dec << std::endl + << " Network = " << std::hex << network_ << std::dec << std::endl + << " Engine = " << std::hex << engine_ << std::dec << std::endl + << " Logger = " << std::hex << logger_ << std::dec << std::endl + << " Thread = " << std::hex << thr_ << std::dec << std::endl; + return oss.str(); + } + ~TRTCalibrationResource() { + VLOG(0) << "Destroying Calibration Resource " << std::endl << DebugString(); + } + TRTInt8Calibrator* calibrator_; + nvinfer1::IBuilder* builder_; + nvinfer1::INetworkDefinition* network_; + nvinfer1::ICudaEngine* engine_; + tensorflow::tensorrt::Logger* logger_; + // TODO(sami): Use threadpool threads! + std::thread* thr_; +}; + +class TRTWeightStore : public tensorflow::ResourceBase { + public: + TRTWeightStore() {} + std::list> store_; + string DebugString() override { + std::stringstream oss; + size_t lenBytes = 0; + for (const auto& v : store_) { + lenBytes += v.size() * sizeof(uint8_t); + } + oss << " Number of entries = " << store_.size() << std::endl + << " Total number of bytes = " + << store_.size() * sizeof(std::vector) + lenBytes << std::endl; + return oss.str(); + } + virtual ~TRTWeightStore() { VLOG(1) << "Destroying store" << DebugString(); } +}; + +class TRTEngineResource : public tensorflow::ResourceBase { + public: + TRTEngineResource() : runtime_(nullptr), ctx_(nullptr){}; + string DebugString() override { return string(""); } + nvinfer1::IRuntime* runtime_; + nvinfer1::IExecutionContext* ctx_; +}; + +} // namespace tensorrt +} // namespace tensorflow +#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCEMGR_TRTRESOURCES_H_ +#endif +#endif 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/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py new file mode 100644 index 0000000000000000000000000000000000000000..ad01bedd8fa066e914b05b20dbc47d9aabe790d9 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/test_tftrt.py @@ -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. +# ============================================================================== +"""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): + """Run given graphdef once.""" + 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 + + +# Use real data that is representative of the inference dataset +# for calibration. For this test script it is random data. +def run_calibration(gdef, dumm_inp): + """Run given calibration graph multiple times.""" + 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: + # run over real calibration data here, we are mimicking a calibration set of + # 30 different batches. Use as much calibration data as you want + for _ in range(30): + 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) + orig_graph = get_simple_graph_def() # use a frozen graph for inference + # Get optimized graph + trt_graph = trt.create_inference_graph( + input_graph_def=orig_graph, + outputs=["output"], + max_batch_size=inp_dims[0], + max_workspace_size_bytes=1 << 25, + precision_mode="FP32", # TRT Engine precision "FP32","FP16" or "INT8" + minimum_segment_size=2 # minimum number of nodes in an engine + ) + o1 = run_graph(orig_graph, 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 + fp16_graph = trt.create_inference_graph( + input_graph_def=orig_graph, + outputs=["output"], + max_batch_size=inp_dims[0], + max_workspace_size_bytes=1 << 25, + precision_mode="FP16", # TRT Engine precision "FP32","FP16" or "INT8" + minimum_segment_size=2 # minimum number of nodes in an engine + ) + int8_calib_gdef = trt.create_inference_graph( + input_graph_def=orig_graph, + outputs=["output"], + max_batch_size=inp_dims[0], + max_workspace_size_bytes=1 << 25, + precision_mode="INT8", # TRT Engine precision "FP32","FP16" or "INT8" + minimum_segment_size=2 # minimum number of nodes in an engine + ) + o4 = run_graph(fp16_graph, dummy_input) + _ = run_calibration(int8_calib_gdef, dummy_input) + int8_graph = trt.calib_graph_to_infer_graph(int8_calib_gdef) + o5 = run_graph(int8_graph, dummy_input) + assert np.allclose(o1, o4) + assert np.allclose(o1, o5) + print("Pass") diff --git a/tensorflow/contrib/tensorrt/trt_conversion.i b/tensorflow/contrib/tensorrt/trt_conversion.i new file mode 100644 index 0000000000000000000000000000000000000000..46480e99a113afb34702b0ecd71468d4bdc83f98 --- /dev/null +++ b/tensorflow/contrib/tensorrt/trt_conversion.i @@ -0,0 +1,186 @@ +/* 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; +%unignore calib_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, + int precision_mode, + int minimum_segment_size + // 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(precision_mode < 0 || precision_mode > 2){ + out_status = "InvalidArgument;Invalid precision_mode"; + 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, ""}; + } + tensorflow::GraphDef outGraph; + tensorflow::Status conversion_status = + tensorflow::tensorrt::convert::ConvertGraphDefToTensorRT( + graph_def, output_names, max_batch_size, max_workspace_size_bytes, + &outGraph, precision_mode, minimum_segment_size); + 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 calib_convert(string graph_def_string // const tensorflow::GraphDef& + // 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, ""}; + } + + tensorflow::GraphDef outGraph; + tensorflow::Status conversion_status = + tensorflow::tensorrt::convert::ConvertCalibGraphToInferGraph(graph_def, + &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 calib_convert(string graph_def_string); + +std::pair trt_convert(string graph_def_string, + std::vector output_names, + size_t max_batch_size, + size_t max_workspace_size_bytes, + int precision_mode, int minimum_segment_size); + + +%unignoreall diff --git a/tensorflow/contrib/testing/python/framework/fake_summary_writer.py b/tensorflow/contrib/testing/python/framework/fake_summary_writer.py index f2065c666255984c8ab770fc10f682b1eabad095..15a415df303df5be44e89c00005cb253ae2af286 100644 --- a/tensorflow/contrib/testing/python/framework/fake_summary_writer.py +++ b/tensorflow/contrib/testing/python/framework/fake_summary_writer.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function from tensorflow.core.framework import summary_pb2 +from tensorflow.python.framework import test_util from tensorflow.python.summary.writer import writer from tensorflow.python.summary.writer import writer_cache @@ -85,7 +86,11 @@ class FakeSummaryWriter(object): if expected_added_graphs is not None: test_case.assertEqual(expected_added_graphs, self._added_graphs) if expected_added_meta_graphs is not None: - test_case.assertEqual(expected_added_meta_graphs, self._added_meta_graphs) + test_case.assertEqual(len(expected_added_meta_graphs), + len(self._added_meta_graphs)) + for expected, actual in zip(expected_added_meta_graphs, + self._added_meta_graphs): + test_util.assert_meta_graph_protos_equal(test_case, expected, actual) if expected_session_logs is not None: test_case.assertEqual(expected_session_logs, self._added_session_logs) diff --git a/tensorflow/contrib/timeseries/examples/BUILD b/tensorflow/contrib/timeseries/examples/BUILD index bb86ecb2209f9bed3ad6c37f4b23bc7b361e1bd6..70bf67c7793ba96a8be0ced58421f778a5776d6c 100644 --- a/tensorflow/contrib/timeseries/examples/BUILD +++ b/tensorflow/contrib/timeseries/examples/BUILD @@ -25,7 +25,10 @@ py_test( srcs = ["predict_test.py"], data = ["data/period_trend.csv"], srcs_version = "PY2AND3", - tags = ["notsan"], # b/67513579 + tags = [ + "no_windows", # TODO: needs investigation on Windows + "notsan", # b/67513579 + ], deps = [ ":predict", "//tensorflow/python:client_testlib", diff --git a/tensorflow/contrib/timeseries/examples/data/multivariate_periods.csv b/tensorflow/contrib/timeseries/examples/data/multivariate_periods.csv index 02a60d1cf61765c7c916803fe918d8b7b186405e..9b15b4f0b26f11ac3281ca4206654872984628b6 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.,strkeya +1,0.108010001864,1.41645361423,2.1686839775,2.94963962176,4.1263503303,1.,0.,strkeyb +2,-0.800567600028,1.0172132907,1.96434754116,2.99885333086,4.04300485864,1.,0.,strkey +3,0.0607042871898,0.719540073421,1.9765012584,2.89265588817,4.0951014426,1.,0.,strkey +4,0.933712200629,0.28052120776,1.41018552514,2.69232603996,4.06481164223,1.,0.,strkey +5,-0.171730652974,0.260054421028,1.48770816369,2.62199129293,4.44572807842,1.,0.,strkey +6,-1.00180162933,0.333045158863,1.50006392277,2.88888309683,4.24755865606,1.,0.,strkey +7,0.0580061875336,0.688929398826,1.56543458772,2.99840358953,4.52726873347,1.,0.,strkey +8,0.764139447412,1.24704875327,1.77649279698,3.13578593851,4.63238922951,1.,0.,strkey +9,-0.230331874785,1.47903998963,2.03547545751,3.20624030377,4.77980005228,1.,0.,strkey +10,-1.03846045211,2.01133000781,2.31977503972,3.67951536251,5.09716775897,1.,0.,strkeyc +11,0.188643592253,2.23285349038,2.68338482249,3.49817168611,5.24928239634,1.,0.,strkey +12,0.91207302309,2.24244446841,2.71362604985,3.96332587625,5.37802271594,1.,0.,strkey +13,-0.296588665881,2.02594634141,3.07733910479,3.99698324956,5.56365901394,1.,0.,strkey +14,-0.959961476551,1.45078629833,3.18996420137,4.3763059609,5.65356015609,1.,0.,strkey +15,0.46313530679,1.01141441548,3.4980215948,4.20224896882,5.88842247449,1.,0.,strkey +16,0.929354125798,0.626635305936,3.70508262244,4.51791573544,5.73945973251,1.,0.,strkey +17,-0.519110731957,0.269249223148,3.39866823332,4.46802003061,5.82768174382,1.,0.,strkey +18,-0.924330981367,0.349602834684,3.21762413294,4.72803587499,5.94918925767,1.,0.,strkey +19,0.253239387885,0.345158023497,3.11071425333,4.79311566935,5.9489259713,1.,0.,strkey +20,0.637408390225,0.698996675371,3.25232492145,4.73814732384,5.9612010251,1.,0.,strkey +21,-0.407396859412,1.17456342803,2.49526823723,4.59323415742,5.82501686811,1.,0.,strkey +22,-0.967485452118,1.66655933642,2.47284606244,4.58316034754,5.88721406681,1.,0.,strkey +23,0.474480867904,1.95018556323,2.0228950072,4.48651142819,5.8255943735,1.,0.,strkey +24,1.04309652155,2.23519892356,1.91924131572,4.19094661783,5.87457348436,1.,0.,strkey +25,-0.517861513772,2.12501967336,1.70266619979,4.05280882887,5.72160912899,1.,0.,strkey +26,-0.945301585146,1.65464653549,1.81567174251,3.92309850635,5.58270493814,1.,0.,strkey +27,0.501153868974,1.40600764889,1.53991387719,3.72853247942,5.60169001727,1.,0.,strkey +28,0.972859524418,1.00344321868,1.5175642828,3.64092376655,5.10567722582,1.,0.,strkey +29,-0.70553406135,0.465306263885,1.7038540803,3.33236870312,5.09182481555,1.,0.,strkey +30,-0.946093634916,0.294539309453,1.88052827037,2.93011492669,4.97354922696,1.,0.,strkey +31,0.47922123231,0.308465865031,2.03445883031,2.90772899045,4.86241793548,1.,0.,strkey +32,0.754030014252,0.549752241167,2.46115815089,2.95063349534,4.71834614627,1.,0.,strkey +33,-0.64875949826,0.894615488148,2.5922463381,2.81269864022,4.43480095104,1.,0.,strkey +34,-0.757829951086,1.39123914261,2.69258079904,2.61834837315,4.36580046156,1.,0.,strkey +35,0.565653301088,1.72360022693,2.97794913834,2.80403840334,4.27327248459,1.,0.,strkey +36,0.867440092372,2.21100730052,3.38648090792,2.84057515729,4.12210169576,1.,0.,strkey +37,-0.894567758095,2.17549105818,3.45532493329,2.90446025717,4.00251740584,1.,0.,strkeyd +38,-0.715442356893,2.15105389965,3.52041791902,3.03650393392,4.12809249577,1.,0.,strkey +39,0.80671703672,1.81504564517,3.60463324866,3.00747789871,3.98440762467,1.,0.,strkey +40,0.527014790142,1.31803513865,3.43842186337,3.3332594663,4.03232406566,1.,0.,strkey +41,-0.795936862129,0.847809114454,3.09875133548,3.52863155938,3.94883924909,1.,0.,strkey +42,-0.610245806946,0.425530441018,2.92581949152,3.77238736123,4.27287245021,1.,0.,strkey +43,0.611662279431,0.178432049837,2.48128214822,3.73212087883,4.17319013831,1.,0.,strkey +44,0.650866553108,0.220341648392,2.41694642022,4.2609098519,4.27271645905,1.,0.,strkey +45,-0.774156982023,0.632667602331,2.05474356052,4.32889204886,4.18029723271,1.,0.,strkey +46,-0.714058448409,0.924562377599,1.75706135146,4.52492718422,4.3972678094,1.,0.,strkey +47,0.889627293379,1.46207968841,1.78299357672,4.64466731095,4.56317887554,1.,0.,strkey +48,0.520140662861,1.8996333843,1.41377633823,4.48899091177,4.78805049769,1.,0.,strkey +49,-1.03816935616,2.08997002059,1.51218375351,4.84167764204,4.93026048606,1.,0.,strkey +50,-0.40772951362,2.30878972136,1.44144415128,4.76854460997,5.01538444629,1.,0.,strkey +51,0.792730684781,1.91367048509,1.58887384677,4.71739397335,5.25690012199,1.,0.,strkey +52,0.371311881576,1.67565079528,1.81688563053,4.60353107555,5.44265822961,1.,0.,strkey +53,-0.814398070371,1.13374634126,1.80328814859,4.72264252878,5.52674761122,1.,0.,strkey +54,-0.469017949323,0.601244136627,2.29690896736,4.49859178859,5.54126153454,1.,0.,strkey +55,0.871044371426,0.407597593794,2.7499112487,4.19060637761,5.57693767301,1.,0.,strkey +56,0.523764933017,0.247705192709,3.09002071379,4.02095509006,5.80510362182,1.,0.,strkey +57,-0.881326403531,0.31513103164,3.11358205718,3.96079100808,5.81000652365,1.,0.,strkey +58,-0.357928025339,0.486163915865,3.17884556771,3.72634990659,5.85693642011,1.,0.,strkey +59,0.853038779822,1.04218094475,3.45835384454,3.36703969978,5.9585988449,1.,0.,strkey +60,0.435311516013,1.59715085283,3.63313338588,3.11276729421,5.93643818229,1.,0.,strkey +61,-1.02703719138,1.92205832542,3.47606111735,3.06247155999,6.02106646259,1.,0.,strkey +62,-0.246661325557,2.14653802542,3.29446326567,2.89936259181,5.67531541272,1.,0.,strkey +63,1.02554736569,2.25943737733,3.07031591528,2.78176218013,5.78206328989,1.,0.,strkey +64,0.337814475969,2.07589147224,2.80356226089,2.55888206331,5.7094075496,1.,0.,strkey +65,-1.12023369929,1.25333011618,2.56497288445,2.77361359194,5.50799418376,1.,0.,strkey +66,-0.178980246554,1.11937139901,2.51598681313,2.91438309151,5.47469577206,1.,0.,strkey +67,0.97550951531,0.60553823137,2.11657741073,2.88081098981,5.37034999502,1.,0.,strkey +68,0.136653357206,0.365828836075,1.97386033165,3.13217903204,5.07254490219,1.,0.,strkey +69,-1.05607596951,0.153152115069,1.52110743825,3.01308794192,5.08902539125,1.,0.,strkey +70,-0.13095280331,0.337113974483,1.52703079853,3.16687131599,4.86649398514,1.,0.,strkey +71,1.07081057754,0.714247566736,1.53761382634,3.45151989484,4.75892309166,1.,0.,strkey +72,0.0153410376082,1.24631231847,1.61690939161,3.85481994498,4.35683752832,1.,0.,strkey +73,-0.912801257303,1.60791309476,1.8729264524,4.03037260012,4.36072588913,1.,0.,strkey +74,-0.0894895640338,2.02535207407,1.93484909619,4.09557485132,4.35327025188,1.,0.,strkey +75,0.978646999652,2.20085086625,2.09003440427,4.27542353033,4.1805058388,1.,0.,strkey +76,-0.113312642876,2.2444100761,2.50789248839,4.4151861502,4.03267168136,1.,0.,strkey +77,-1.00215099149,1.84305628445,2.61691237246,4.45425147595,3.81203553766,1.,0.,strkey +78,-0.0183234614205,1.49573923116,2.99308471214,4.71134960112,4.0273804959,1.,0.,strkey +79,1.0823738177,1.12211589848,3.27079386925,4.94288270502,4.01851068083,1.,0.,strkey +80,0.124370187893,0.616474412808,3.4284236674,4.76942168327,3.9749536483,1.,0.,strkey +81,-0.929423379352,0.290977090976,3.34131726136,4.78590392707,4.10190661656,1.,0.,strkey +82,0.23766302648,0.155302052254,3.49779513794,4.64605656795,4.15571321107,1.,0.,strkey +83,1.03531486192,0.359702776204,3.4880725919,4.48167586667,4.21134561991,1.,0.,strkey +84,-0.261234571382,0.713877760378,3.42756426614,4.426443869,4.25208300527,1.,0.,strkey +85,-1.03572442277,1.25001113691,2.96908341113,4.25500915322,4.25723010649,1.,0.,strkey +86,0.380034261243,1.70543355622,2.73605932518,4.16703432307,4.63700400788,1.,0.,strkey +87,1.03734873488,1.97544410562,2.55586572141,3.84976673263,4.55282864289,1.,0.,strkey +88,-0.177344253372,2.22614526325,2.09565864891,3.77378097953,4.82577400298,1.,0.,strkey +89,-0.976821526892,2.18385079177,1.78522284118,3.67768223554,5.06302440873,1.,0.,strkey +90,0.264820472091,1.86981946157,1.50048403865,3.43619796921,5.05651761669,1.,0.,strkey +91,1.05642344868,1.47568646076,1.51347671977,3.20898518885,5.50149047462,1.,0.,strkey +92,-0.311607433358,1.04226467636,1.52089650905,3.02291865417,5.4889046232,1.,0.,strkey +93,-0.724285777937,0.553052311957,1.48573560173,2.7365973598,5.72549174225,1.,0.,strkey +94,0.519859192905,0.226520626591,1.61543723167,2.84102086852,5.69330622288,1.,0.,strkey +95,1.0323195039,0.260873217055,1.81913034804,2.83951143848,5.90325028086,1.,0.,strkey +96,-0.53285682538,0.387695521405,1.70935609313,2.57977050631,5.79579213161,1.,0.,strkey +97,-0.975127997215,0.920948771589,2.51292643636,2.71004616612,5.87016469227,1.,0.,strkey +98,0.540246804099,1.36445470181,2.61949412896,2.98482553485,6.02447664937,1.,0.,strkey +99,0.987764008058,1.85581989607,2.84685706149,2.94760204892,6.0212151724,1.,0.,strkey diff --git a/tensorflow/contrib/timeseries/examples/known_anomaly.py b/tensorflow/contrib/timeseries/examples/known_anomaly.py index 7659dd308a7ee1b70d6688b85e4f6157ddee0540..c08c0b0acb917f527d7efa91874d6405b9220083 100644 --- a/tensorflow/contrib/timeseries/examples/known_anomaly.py +++ b/tensorflow/contrib/timeseries/examples/known_anomaly.py @@ -46,12 +46,12 @@ def train_and_evaluate_exogenous(csv_file_name=_DATA_FILE, train_steps=300): # Indicate the format of our exogenous feature, in this case a string # representing a boolean value. - string_feature = tf.contrib.layers.sparse_column_with_keys( - column_name="is_changepoint", keys=["no", "yes"]) + string_feature = tf.feature_column.categorical_column_with_vocabulary_list( + key="is_changepoint", vocabulary_list=["no", "yes"]) # Specify the way this feature is presented to the model, here using a one-hot # encoding. - one_hot_feature = tf.contrib.layers.one_hot_column( - sparse_id_column=string_feature) + one_hot_feature = tf.feature_column.indicator_column( + categorical_column=string_feature) estimator = tf.contrib.timeseries.StructuralEnsembleRegressor( periodicities=12, diff --git a/tensorflow/contrib/timeseries/examples/lstm.py b/tensorflow/contrib/timeseries/examples/lstm.py index c834430b95d37f7c313c54c086b9c323cf3e921f..2eee878196bb64b523c491ca808ca8d6ff5dd36c 100644 --- a/tensorflow/contrib/timeseries/examples/lstm.py +++ b/tensorflow/contrib/timeseries/examples/lstm.py @@ -20,12 +20,14 @@ 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 @@ -46,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 @@ -56,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.feature_column`s 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__( @@ -63,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() @@ -70,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 @@ -102,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) @@ -129,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 @@ -141,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): @@ -162,36 +178,85 @@ 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. + categorical_column = tf.feature_column.categorical_column_with_hash_bucket( + key="categorical_exogenous_feature", hash_bucket_size=16) + 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.feature_column.numeric_column( + "2d_exogenous_feature", shape=(2,)), + tf.feature_column.embedding_column( + categorical_column=categorical_column, dimension=10)] 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 + + ("categorical_exogenous_feature",)), + # Data types other than for `times` need to be specified if they aren't + # float32. In this case one of our exogenous features has string dtype. + column_dtypes=((tf.int64,) + (tf.float32,) * 7 + (tf.string,))) 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), + "categorical_exogenous_feature": numpy.array( + ["strkey"] * 100)[None, :, None]} (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 fff972c1f3277ad5d83673a202a50d1e6f7df210..64f5cd83575e1eeaa35360c4dc8c5a0d1f384066 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/BUILD @@ -140,11 +140,13 @@ py_library( "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", "//tensorflow/python:state_ops", + "//tensorflow/python:summary", "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python/estimator:estimator_py", "//tensorflow/python/estimator:export", "//tensorflow/python/estimator:head", + "//tensorflow/python/estimator:metric_keys", ], ) @@ -154,9 +156,7 @@ py_test( "head_test.py", ], srcs_version = "PY2AND3", - tags = [ - "no_pip_gpu", # b/63391119 - ], + tags = ["no_pip_gpu"], # b/63391119 deps = [ ":feature_keys", ":head", @@ -425,6 +425,7 @@ py_test( srcs_version = "PY2AND3", tags = [ "no_pip_gpu", # b/63391119 + "no_windows", # TODO: needs investigation on Windows ], deps = [ ":feature_keys", diff --git a/tensorflow/contrib/timeseries/python/timeseries/estimators.py b/tensorflow/contrib/timeseries/python/timeseries/estimators.py index 3738dfa154d4f39b9562446972443ed88f3fbe8b..8d13343e82340dae11b0be54e3bc3152060dca36 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/estimators.py +++ b/tensorflow/contrib/timeseries/python/timeseries/estimators.py @@ -29,10 +29,12 @@ from tensorflow.contrib.timeseries.python.timeseries.state_space_models.filterin from tensorflow.python.estimator import estimator_lib from tensorflow.python.estimator.export import export_lib +from tensorflow.python.feature_column import feature_column 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 parsing_ops from tensorflow.python.training import training as train @@ -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,29 @@ 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 - placeholders[feature_key] = array_ops.placeholder( - dtype=value_tensor.dtype, name=feature_key, shape=feature_shape) + if self._model.exogenous_feature_columns: + 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. + parsed_features = ( + feature_column.make_parse_example_spec( + self._model.exogenous_feature_columns)) + placeholder_features = parsing_ops.parse_example( + serialized=array_ops.placeholder( + shape=[None], dtype=dtypes.string), + features=parsed_features) + exogenous_feature_shapes = { + key: (value.get_shape(), value.dtype) for key, value + in placeholder_features.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_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. @@ -327,11 +340,11 @@ class StructuralEnsembleRegressor(StateSpaceRegressor): determine the model size. Learning autoregressive coefficients typically requires more steps and a smaller step size than other components. - exogenous_feature_columns: A list of tf.contrib.layers.FeatureColumn - objects (for example tf.contrib.layers.embedding_column) corresponding - to exogenous features which provide extra information to the model but - are not part of the series to be predicted. Passed to - tf.contrib.layers.input_from_feature_columns. + exogenous_feature_columns: A list of `tf.feature_column`s (for example + `tf.feature_column.embedding_column`) corresponding to exogenous + features which provide extra information to the model but are not part + of the series to be predicted. Passed to + `tf.feature_column.input_layer`. exogenous_update_condition: A function taking two Tensor arguments, `times` (shape [batch size]) and `features` (a dictionary mapping exogenous feature keys to Tensors with shapes [batch size, ...]), and diff --git a/tensorflow/contrib/timeseries/python/timeseries/head.py b/tensorflow/contrib/timeseries/python/timeseries/head.py index f0330bfbbd6e8067e5d085376acdf2e6bcaccb6a..f4d9351432ef32f22334667766528907eac77f19 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/head.py +++ b/tensorflow/contrib/timeseries/python/timeseries/head.py @@ -26,6 +26,7 @@ from tensorflow.contrib.timeseries.python.timeseries import feature_keys from tensorflow.python.estimator import estimator_lib from tensorflow.python.estimator.canned import head as head_lib +from tensorflow.python.estimator.canned import metric_keys from tensorflow.python.estimator.export import export_lib from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -35,6 +36,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.util import nest +from tensorflow.python.summary import summary def time_series_regression_head(model, @@ -71,11 +73,34 @@ class _TimeSeriesRegressionHead(head_lib._Head): # pylint:disable=protected-acc self.input_statistics_generator = input_statistics_generator self._name = name + @property + def name(self): + return self._name + + # TODO(terrytangyuan): consolidate `model_outputs` and `_Head.LossSpec` + # once `_Head.create_loss` becomes extendable + def create_loss(self, features, mode, logits=None, labels=None): + """See `_Head`.""" + model_outputs = self.state_manager.define_loss( + self.model, features, mode) + summary.scalar( + head_lib._summary_key(self._name, metric_keys.MetricKeys.LOSS), + model_outputs.loss) + return model_outputs + + @property + def logits_dimension(self): + """See `_Head`.""" + return 1 + def _train_ops(self, features): """Add training ops to the graph.""" - with variable_scope.variable_scope("model"): - model_outputs = self.state_manager.define_loss( - self.model, features, estimator_lib.ModeKeys.TRAIN) + mode = estimator_lib.ModeKeys.TRAIN + with variable_scope.variable_scope( + "model", + # Use ResourceVariables to avoid race conditions. + use_resource=True): + model_outputs = self.create_loss(features, mode) train_op = optimizers.optimize_loss( model_outputs.loss, @@ -85,31 +110,14 @@ class _TimeSeriesRegressionHead(head_lib._Head): # pylint:disable=protected-acc learning_rate=None) return estimator_lib.EstimatorSpec( loss=model_outputs.loss, - mode=estimator_lib.ModeKeys.TRAIN, + mode=mode, train_op=train_op) - # TODO(terrytangyuan): suffix summary and metrics keys by `"/" + name` - @property - def name(self): - return self._name - - # TODO(terrytangyuan): unused for now. Need to decouple - # `state_manager.define_loss` to satisfy the extendable return signature of - # `_Head.create_loss`. - def create_loss(self, features, mode, logits, labels): - """See `_Head`.""" - return None - - # TODO(terrytangyuan): check label dimension - @property - def logits_dimension(self): - return None - def _evaluate_ops(self, features): """Add ops for evaluation (aka filtering) to the graph.""" - with variable_scope.variable_scope("model"): - model_outputs = self.state_manager.define_loss( - self.model, features, estimator_lib.ModeKeys.EVAL) + mode = estimator_lib.ModeKeys.EVAL + with variable_scope.variable_scope("model", use_resource=True): + model_outputs = self.create_loss(features, mode) metrics = {} # Just output in-sample predictions for the last chunk seen for prediction_key, prediction_value in model_outputs.predictions.items(): @@ -122,13 +130,13 @@ class _TimeSeriesRegressionHead(head_lib._Head): # pylint:disable=protected-acc model_outputs.end_state)) return estimator_lib.EstimatorSpec( loss=model_outputs.loss, - mode=estimator_lib.ModeKeys.EVAL, + mode=mode, eval_metric_ops=metrics, predictions={}) def _predict_ops(self, features): """Add ops for prediction to the graph.""" - with variable_scope.variable_scope("model"): + with variable_scope.variable_scope("model", use_resource=True): prediction = self.model.predict(features=features) prediction[feature_keys.PredictionResults.TIMES] = features[ feature_keys.PredictionFeatures.TIMES] @@ -137,12 +145,11 @@ class _TimeSeriesRegressionHead(head_lib._Head): # pylint:disable=protected-acc def _serving_ops(self, features): """Add ops for serving to the graph.""" - with variable_scope.variable_scope("model"): + with variable_scope.variable_scope("model", use_resource=True): prediction_outputs = self.model.predict(features=features) with variable_scope.variable_scope("model", reuse=True): - filtering_outputs = self.state_manager.define_loss( - self.model, features, estimator_lib.ModeKeys.EVAL) - + filtering_outputs = self.create_loss( + features, estimator_lib.ModeKeys.EVAL) return estimator_lib.EstimatorSpec( mode=estimator_lib.ModeKeys.PREDICT, export_outputs={ @@ -191,7 +198,7 @@ class _TimeSeriesRegressionHead(head_lib._Head): # pylint:disable=protected-acc def create_estimator_spec(self, features, mode, labels=None): """Performs basic error checking and returns an EstimatorSpec.""" - with ops.name_scope("head"): + with ops.name_scope(self._name, "head"): if labels: raise ValueError( "The model received a `labels` dictionary, which is " diff --git a/tensorflow/contrib/timeseries/python/timeseries/input_pipeline.py b/tensorflow/contrib/timeseries/python/timeseries/input_pipeline.py index d4ee59036624cffb216709e096981d362670e416..04225333b9377447f46d32663df76aece97a51e7 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/input_pipeline.py +++ b/tensorflow/contrib/timeseries/python/timeseries/input_pipeline.py @@ -500,6 +500,41 @@ class CSVReader(ReaderBaseTimeSeriesParser): return features +class TFExampleReader(ReaderBaseTimeSeriesParser): + """Reads and parses `tf.Example`s from a TFRecords file.""" + + def __init__(self, + filenames, + features): + """Configure `tf.Example` parsing. + + Args: + filenames: A filename or list of filenames to read the time series + from. Each line must have columns corresponding to `column_names`. + features: A dictionary mapping from feature keys to `tf.FixedLenFeature` + objects. Must include `TrainEvalFeatures.TIMES` (scalar integer) and + `TrainEvalFeatures.VALUES` (floating point vector) features. + Raises: + ValueError: If required times/values features are not present. + """ + if feature_keys.TrainEvalFeatures.TIMES not in features: + raise ValueError("'{}' is a required column.".format( + feature_keys.TrainEvalFeatures.TIMES)) + if feature_keys.TrainEvalFeatures.VALUES not in features: + raise ValueError("'{}' is a required column.".format( + feature_keys.TrainEvalFeatures.VALUES)) + self._features = features + super(TFExampleReader, self).__init__(filenames=filenames) + + def _get_reader(self): + return io_ops.TFRecordReader() + + def _process_records(self, examples): + """Parse `tf.Example`s into `Tensors`.""" + return parsing_ops.parse_example( + serialized=examples, features=self._features) + + class TimeSeriesInputFn(object): """Base for classes which create batches of windows from a time series.""" diff --git a/tensorflow/contrib/timeseries/python/timeseries/input_pipeline_test.py b/tensorflow/contrib/timeseries/python/timeseries/input_pipeline_test.py index ed78a835a4d451e9e7d18bb833d8ebed6c05a195..703537abf0fe3985aaf0434cc633cb410dd6bd4c 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/input_pipeline_test.py +++ b/tensorflow/contrib/timeseries/python/timeseries/input_pipeline_test.py @@ -27,7 +27,11 @@ from tensorflow.contrib.timeseries.python.timeseries import input_pipeline from tensorflow.contrib.timeseries.python.timeseries import test_utils from tensorflow.contrib.timeseries.python.timeseries.feature_keys import TrainEvalFeatures +from tensorflow.core.example import example_pb2 +from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors +from tensorflow.python.lib.io import tf_record +from tensorflow.python.ops import parsing_ops from tensorflow.python.ops import variables from tensorflow.python.platform import test from tensorflow.python.training import coordinator as coordinator_lib @@ -52,6 +56,21 @@ def _make_csv_time_series(num_features, num_samples, test_tmpdir): return filename +def _make_tfexample_series(num_features, num_samples, test_tmpdir): + _, data_file = tempfile.mkstemp(dir=test_tmpdir) + with tf_record.TFRecordWriter(data_file) as writer: + for i in range(num_samples): + example = example_pb2.Example() + times = example.features.feature[TrainEvalFeatures.TIMES] + times.int64_list.value.append(i) + values = example.features.feature[TrainEvalFeatures.VALUES] + values.float_list.value.extend( + [float(i) * 2. + feature_number + for feature_number in range(num_features)]) + writer.write(example.SerializeToString()) + return data_file + + def _make_numpy_time_series(num_features, num_samples): times = numpy.arange(num_samples) values = times[:, None] * 2. + numpy.arange(num_features)[None, :] @@ -107,6 +126,19 @@ class RandomWindowInputFnTests(test.TestCase): time_series_reader = input_pipeline.CSVReader([filename]) self._test_out_of_order(time_series_reader, discard_out_of_order=False) + def test_tfexample_sort_out_of_order(self): + filename = _make_tfexample_series( + num_features=1, num_samples=50, + test_tmpdir=self.get_temp_dir()) + time_series_reader = input_pipeline.TFExampleReader( + [filename], + features={ + TrainEvalFeatures.TIMES: parsing_ops.FixedLenFeature( + shape=[], dtype=dtypes.int64), + TrainEvalFeatures.VALUES: parsing_ops.FixedLenFeature( + shape=[1], dtype=dtypes.float32)}) + self._test_out_of_order(time_series_reader, discard_out_of_order=False) + def test_numpy_sort_out_of_order(self): data = _make_numpy_time_series(num_features=1, num_samples=50) time_series_reader = input_pipeline.NumpyReader(data) @@ -183,6 +215,20 @@ class RandomWindowInputFnTests(test.TestCase): self._test_multivariate(time_series_reader=time_series_reader, num_features=2) + def test_tfexample_multivariate(self): + filename = _make_tfexample_series( + num_features=2, num_samples=50, + test_tmpdir=self.get_temp_dir()) + time_series_reader = input_pipeline.TFExampleReader( + [filename], + features={ + TrainEvalFeatures.TIMES: parsing_ops.FixedLenFeature( + shape=[], dtype=dtypes.int64), + TrainEvalFeatures.VALUES: parsing_ops.FixedLenFeature( + shape=[2], dtype=dtypes.float32)}) + self._test_multivariate(time_series_reader=time_series_reader, + num_features=2) + def test_numpy_multivariate(self): data = _make_numpy_time_series(num_features=3, num_samples=50) time_series_reader = input_pipeline.NumpyReader(data) @@ -248,6 +294,20 @@ class WholeDatasetInputFnTests(test.TestCase): self._whole_dataset_input_fn_test_template( time_series_reader=time_series_reader, num_features=1, num_samples=50) + def test_tfexample(self): + filename = _make_tfexample_series( + num_features=4, num_samples=100, + test_tmpdir=self.get_temp_dir()) + time_series_reader = input_pipeline.TFExampleReader( + [filename], + features={ + TrainEvalFeatures.TIMES: parsing_ops.FixedLenFeature( + shape=[], dtype=dtypes.int64), + TrainEvalFeatures.VALUES: parsing_ops.FixedLenFeature( + shape=[4], dtype=dtypes.float32)}) + self._whole_dataset_input_fn_test_template( + time_series_reader=time_series_reader, num_features=4, num_samples=100) + def test_numpy(self): data = _make_numpy_time_series(num_features=4, num_samples=100) time_series_reader = input_pipeline.NumpyReader(data) diff --git a/tensorflow/contrib/timeseries/python/timeseries/model.py b/tensorflow/contrib/timeseries/python/timeseries/model.py index b32b5c5494ae14187954b900119678a5b53a3602..7644764a7459db3951fe9a2790389713dd412a8f 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/model.py +++ b/tensorflow/contrib/timeseries/python/timeseries/model.py @@ -21,17 +21,17 @@ from __future__ import print_function import abc import collections -from tensorflow.contrib import layers - from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.contrib.timeseries.python.timeseries.feature_keys import PredictionFeatures from tensorflow.contrib.timeseries.python.timeseries.feature_keys import TrainEvalFeatures +from tensorflow.python.feature_column import feature_column 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 math_ops +from tensorflow.python.ops import parsing_ops from tensorflow.python.ops import tensor_array_ops from tensorflow.python.ops import variable_scope @@ -65,11 +65,11 @@ class TimeSeriesModel(object): Args: num_features: Number of features for the time series - exogenous_feature_columns: A list of tf.contrib.layers.FeatureColumn - objects (for example tf.contrib.layers.embedding_column) corresponding - to exogenous features which provide extra information to the model but - are not part of the series to be predicted. Passed to - tf.contrib.layers.input_from_feature_columns. + exogenous_feature_columns: A list of `tf.feature_column`s (for example + `tf.feature_column.embedding_column`) corresponding to exogenous + features which provide extra information to the model but are not + part of the series to be predicted. Passed to + `tf.feature_column.input_layer`. dtype: The floating point datatype to use. """ if exogenous_feature_columns: @@ -83,6 +83,11 @@ class TimeSeriesModel(object): self._stats_means = None self._stats_sigmas = None + @property + def exogenous_feature_columns(self): + """`tf.feature_colum`s 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 +255,26 @@ 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(): + parsed_features = ( + feature_column.make_parse_example_spec( + self._exogenous_feature_columns)) + placeholder_features = parsing_ops.parse_example( + serialized=array_ops.placeholder(shape=[None], dtype=dtypes.string), + features=parsed_features) + embedded = feature_column.input_layer( + features=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. @@ -285,13 +310,13 @@ class TimeSeriesModel(object): # Avoid shape warnings when embedding "scalar" exogenous features (those # with only batch and window dimensions); input_from_feature_columns # expects input ranks to match the embedded rank. - if tensor.get_shape().ndims == 1: + if tensor.get_shape().ndims == 1 and tensor.dtype != dtypes.string: exogenous_features_single_batch_dimension[name] = tensor[:, None] else: exogenous_features_single_batch_dimension[name] = tensor embedded_exogenous_features_single_batch_dimension = ( - layers.input_from_feature_columns( - columns_to_tensors=exogenous_features_single_batch_dimension, + feature_column.input_layer( + features=exogenous_features_single_batch_dimension, feature_columns=self._exogenous_feature_columns, trainable=True)) exogenous_regressors = array_ops.reshape( @@ -358,8 +383,8 @@ class SequentialTimeSeriesModel(TimeSeriesModel): may use _scale_back_data or _scale_back_variance to return predictions to the input scale. dtype: The floating point datatype to use. - exogenous_feature_columns: A list of tf.contrib.layers.FeatureColumn - objects. See `TimeSeriesModel`. + exogenous_feature_columns: A list of `tf.feature_column`s objects. See + `TimeSeriesModel`. exogenous_update_condition: A function taking two Tensor arguments `times` (shape [batch size]) and `features` (a dictionary mapping exogenous feature keys to Tensors with shapes [batch size, ...]) and returning a diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD index c86d06e9236962cbabbc56afa1cfe213e0c78bc0..07df7bc9a5cd45e0b54e78aaa6ea02fedc584cc3 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD @@ -40,6 +40,7 @@ py_test( timeout = "long", # Moderate but for asan srcs = ["state_space_model_test.py"], srcs_version = "PY2AND3", + tags = ["no_windows"], # TODO: needs investigation on Windows deps = [ ":state_space_model", "//tensorflow/contrib/layers:layers_py", diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py index 6257002647ed53bbde3ace11a6b45e4e2cdeb57d..951c6546d5fed77e0cfa98a4e774b804639d7dad 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py @@ -112,11 +112,11 @@ class StateSpaceModelConfiguration( exogenous_noise_decreases: If True, exogenous regressors can "set" model state, decreasing uncertainty. If both this parameter and exogenous_noise_increases are False, exogenous regressors are ignored. - exogenous_feature_columns: A list of tf.contrib.layers.FeatureColumn - objects (for example tf.contrib.layers.embedding_column) corresponding - to exogenous features which provide extra information to the model but - are not part of the series to be predicted. Passed to - tf.contrib.layers.input_from_feature_columns. + exogenous_feature_columns: A list of `tf.feature_column`s (for example + `tf.feature_column.embedding_column`) corresponding to exogenous + features which provide extra information to the model but are not part + of the series to be predicted. Passed to + `tf.feature_column.input_layer`. exogenous_update_condition: A function taking two Tensor arguments `times` (shape [batch size]) and `features` (a dictionary mapping exogenous feature keys to Tensors with shapes [batch size, ...]) and returning a 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 5980fc5d5deccc151b01c72fa19b734a7c485bdc..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 diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index a7d54d8a0cc4598c26d1c7bd62f5b0aa1070701b..f9d433a45b12630ef46048dae4be4c922e84053e 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -24,6 +24,7 @@ cc_library( name = "all_ops", deps = [ ":cross_replica_ops_op_lib", + ":host_compute_ops_op_lib", ":infeed_ops_op_lib", ":outfeed_ops_op_lib", ":replication_ops_op_lib", @@ -36,7 +37,9 @@ 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", @@ -67,6 +70,7 @@ py_library( tf_gen_op_libs( op_lib_names = [ "cross_replica_ops", + "host_compute_ops", "infeed_ops", "outfeed_ops", "replication_ops", @@ -76,6 +80,7 @@ tf_gen_op_libs( deps = [ "//tensorflow/contrib/tpu/proto:tpu_embedding_config_proto_cc", "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core:protos_all_cc", ], ) @@ -83,6 +88,7 @@ tf_custom_op_library( name = "python/ops/_tpu_ops.so", srcs = [ "ops/cross_replica_ops.cc", + "ops/host_compute_ops.cc", "ops/infeed_ops.cc", "ops/outfeed_ops.cc", "ops/replication_ops.cc", @@ -99,6 +105,7 @@ tf_gen_op_wrapper_py( name = "tpu_ops", deps = [ ":cross_replica_ops_op_lib", + ":host_compute_ops_op_lib", ":infeed_ops_op_lib", ":outfeed_ops_op_lib", ":replication_ops_op_lib", @@ -161,6 +168,7 @@ py_library( ], srcs_version = "PY2AND3", deps = [ + ":datasets", ":profiler", ":tpu_py", "//tensorflow/contrib/tpu/proto:topology_proto_py", @@ -179,6 +187,33 @@ py_library( ], ) +py_library( + name = "datasets", + srcs = [ + "python/tpu/datasets.py", + ], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:function", + "//tensorflow/python:functional_ops", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/ops:iterator_ops", + "//tensorflow/python/data/ops:readers", + ], +) + +tf_py_test( + name = "datasets_test", + srcs = ["python/tpu/datasets_test.py"], + additional_deps = [ + "//tensorflow/python:client_testlib", + ":datasets", + ], + grpc_enabled = True, +) + tf_py_test( name = "tpu_test", size = "small", @@ -190,6 +225,7 @@ tf_py_test( "//tensorflow/python:framework", "//tensorflow/python:layers", ], + tags = ["no_windows"], # TODO: needs investigation on Windows ) tf_py_test( diff --git a/tensorflow/contrib/tpu/ops/host_compute_ops.cc b/tensorflow/contrib/tpu/ops/host_compute_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..48aeb81ac1311d3acd4972810f0a27a382f8b136 --- /dev/null +++ b/tensorflow/contrib/tpu/ops/host_compute_ops.cc @@ -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. +==============================================================================*/ + +#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("_XlaSendFromHost") + .Input("inputs: Tinputs") + .Input("dynamic_key: string") + .Attr("Tinputs: list(type) >= 0") + .Attr("key: string") + .Attr("device_ordinal: int") + .SetIsStateful() + .SetShapeFn(::tensorflow::shape_inference::NoOutputs) + .Doc(R"doc( +A placeholder op for multiple values that will be sent from TensorFlow to a +running XLA computation. + +inputs: A list of tensors that will be sent to the XLA computation. +dynamic_key: The key sent at runtime by the compile node to identify which +execution the transfer corresponds to. +Tinputs: The element types of each element in `inputs`. +key: A key that is unique in the computation and associates the send with the consumer in +the XLA computation. +device_ordinal: The device to use. +)doc"); + +REGISTER_OP("_XlaRecvAtHost") + .Input("dynamic_key: string") + .Output("outputs: Toutputs") + .Attr("Toutputs: list(type) >= 0") + .Attr("key: string") + .Attr("device_ordinal: int") + .SetIsStateful() + .SetShapeFn(::tensorflow::shape_inference::UnknownShape) + .Doc(R"doc( +A placeholder op for multiple values that will be sent to TensorFlow from a +running XLA computation. + +dynamic_key: The key sent at runtime by the compile node to identify which +execution the transfer corresponds to. +outputs: A list of tensors that will be received from the XLA computation. +Toutputs: The element types of each element in `outputs`. +key: A key that is unique in the computation and associates the send with the consumer in +the XLA computation. +device_ordinal: The device to use. +)doc"); + +} // namespace tensorflow 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 28417b89e0d4e0c5b2ca4f4794d29ab8a31049d7..7bf5c21d0b526ee5e32448f75d39eca8add6d877 100644 --- a/tensorflow/contrib/tpu/ops/tpu_configuration_ops.cc +++ b/tensorflow/contrib/tpu/ops/tpu_configuration_ops.cc @@ -191,6 +191,7 @@ REGISTER_OP("ConfigureDistributedTPU") .Output("topology: string") .Attr("embedding_config: string = ''") .Attr("tpu_embedding_config: string = ''") + .Attr("is_global_init: bool = false") .SetIsStateful() .SetShapeFn(shape_inference::UnknownShape) .Doc(R"doc( @@ -202,6 +203,7 @@ topology. tpu_embedding_config: Serialized tensorflow.tpu.TPUEmbeddingConfiguration that describes the embedding lookups of the program. embedding_config: Reserved. Do not use. +is_global_init: Reserved. Do not use. )doc"); REGISTER_OP("ShutdownDistributedTPU") @@ -212,4 +214,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 index cc32a265286951a1e4d59228da6b3ac83a75c5e9..72d37f774cc518c559b5953561957a799a7da568 100644 --- a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc +++ b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc @@ -50,7 +50,7 @@ namespace tensorflow { // 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 +// saving a checkpoint, the model must Retrieve the parameters back into the // host CPU memory. REGISTER_OP("TPUEmbeddingLoadGradientDescentParameters") @@ -263,7 +263,7 @@ REGISTER_OP("TPUEmbeddingReceiveActivations") .SetIsStateful() .SetShapeFn(tpu_embedding_config_util::ActivationShapes) .Doc(R"doc( -An op that receives embeddng activations on the TPU. +An op that receives embedding activations on the TPU. The TPU system performs the embedding lookups and aggregations specified by the arguments to TPUEmbeddingEnqueueSparseBatch. The results of these @@ -293,7 +293,7 @@ REGISTER_OP("TPUEmbeddingActivations") 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 +from the Tensors returned 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. diff --git a/tensorflow/contrib/tpu/profiler/BUILD b/tensorflow/contrib/tpu/profiler/BUILD index 198da0203a7d17249c4f50110713121b74d5ca4f..0a52d0b13b7c8749ad44377659714d297ffec3ee 100644 --- a/tensorflow/contrib/tpu/profiler/BUILD +++ b/tensorflow/contrib/tpu/profiler/BUILD @@ -18,7 +18,7 @@ filegroup( visibility = ["//tensorflow:__subpackages__"], ) -tf_proto_library_cc( +tf_proto_library( name = "tpu_profiler_proto", srcs = ["tpu_profiler.proto"], has_services = 1, @@ -98,16 +98,36 @@ tf_cc_test( ], ) -tf_proto_library_cc( +tf_proto_library( name = "op_profile_proto", srcs = ["op_profile.proto"], cc_api_version = 2, visibility = ["//visibility:public"], ) -tf_proto_library_cc( +tf_proto_library( name = "tf_op_stats_proto", srcs = ["tf_op_stats.proto"], cc_api_version = 2, visibility = ["//visibility:public"], ) + +tf_proto_library( + name = "tpu_profiler_analysis_proto", + srcs = ["tpu_profiler_analysis.proto"], + has_services = 1, + cc_api_version = 2, + cc_grpc_version = 1, + protodeps = [":tpu_profiler_proto"] + tf_additional_all_protos(), + visibility = ["//visibility:public"], +) + +py_library( + name = "tpu_profiler_analysis_pb2_grpc", + srcs = ["tpu_profiler_analysis_pb2_grpc.py"], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + ":tpu_profiler_analysis_proto_py", + ], +) diff --git a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc index 6a05a2abf6fa57938af2b9b02ce394a9f6b7fe6e..b1ef9fde37fe0647965f0818895be37d2d56d207 100644 --- a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc @@ -29,6 +29,7 @@ limitations under the License. #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" @@ -47,6 +48,19 @@ string GetCurrentTimeStampAsString() { return s; } +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; @@ -60,11 +74,14 @@ ProfileResponse Profile(const string& service_addr, int duration_ms, ::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; @@ -101,7 +118,14 @@ int main(int argc, char** argv) { 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); @@ -130,6 +154,8 @@ int main(int argc, char** argv) { << 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. diff --git a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc index 64e4e6275dd0619796b5088f3b5f25b134a89b22..ebd6185faad28ae7a22eb33f6b358eb2344c9c22 100644 --- a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc @@ -151,8 +151,7 @@ Status WriteTensorboardTPUProfile(const string& logdir, const string& run, TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(profile_run_dir)); // Ignore computation_graph for now. - const bool empty_trace = response.encoded_trace().empty(); - if (!empty_trace) { + 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,11 +162,9 @@ Status WriteTensorboardTPUProfile(const string& logdir, const string& run, TF_RETURN_IF_ERROR(DumpOpProfileToLogDirectory(profile_run_dir, response.op_profile(), os)); } - if (!empty_trace && !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)); } return Status::OK(); diff --git a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h index 2f8656a37be031fbe72fdda355040f88605fde78..29ef977bacfd61e163be49558c5b94277ed479c1 100644 --- a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h +++ b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h @@ -29,6 +29,8 @@ namespace tpu { // - Op profile // - 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); 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 885466e5d17b726f4334674fe15f61ed2a529788..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 @@ -25,17 +25,34 @@ import sys import tensorflow as tf -flags.DEFINE_string('service_addr', None, - 'Address of TPU profiler service e.g. localhost:8466') +# 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.') + +# 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") + '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.") + '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.") + 'Set to false to profile longer TPU ' + 'device traces.') FLAGS = flags.FLAGS EXECUTABLE = 'data/capture_tpu_profile' @@ -46,16 +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='+logdir) - cmd.append('--service_addr='+FLAGS.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()) + 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 3dffebe66801bf5f6956ac4297e8dbededf5a4f7..8d99835b64152629c66607e6792495eb36319eb8 100644 --- a/tensorflow/contrib/tpu/profiler/pip_package/setup.py +++ b/tensorflow/contrib/tpu/profiler/pip_package/setup.py @@ -20,7 +20,7 @@ from __future__ import print_function from setuptools import setup -_VERSION = '1.5.0-rc1' +_VERSION = '1.6.0-rc1' CONSOLE_SCRIPTS = [ 'capture_tpu_profile=cloud_tpu_profiler.main:run_main', @@ -47,20 +47,16 @@ setup( # 4 - Beta # 5 - Production/Stable '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', @@ -69,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 2094294baad63ae73712c8648b588accd4551ef8..e5c798aa2f463a1a2d7cb041ba9b51569958f4fd 100644 --- a/tensorflow/contrib/tpu/profiler/tf_op_stats.proto +++ b/tensorflow/contrib/tpu/profiler/tf_op_stats.proto @@ -77,6 +77,8 @@ message StepInfoResult { // The infeed duration in picoseconds. // Can turn into a map if we want a variable number of ops. optional uint64 infeed_duration_ps = 3; + // The start time of this step in picoseconds. + optional uint64 begin_ps = 4; } // Result proto for a sequence of steps. @@ -155,6 +157,54 @@ message RunEnvironmentResult { repeated HostDependentJobInfoResult host_dependent_job_info = 6; } +// The types of host operations that are tracked. +enum HostOp { + // Invalid host op. + kINVALIDHostOp = 0; + // Each of host op type has two parts: + // (1) the stage where the op happens and (2) the op name. + // stage = Input Data Producer, op = Get Next Batch. + kInputDataProducerGetNextBatch = 1; + // stage = Input Data Producer, op = Session Run. + kInputDataProducerSessionRun = 2; + // stage = Input Data Producer, op = Forward Batch. + kInputDataProducerForwardBatch = 3; + // stage = Infeed Thread, op = Get Next Batch. + kInfeedThreadGetNextBatch = 4; + // stage = Infeed Thread, op = Session Run. + kInfeedThreadSessionRun = 5; + // stage = Infeed Thread, op = Forward Batch. + kInfeedThreadForwardBatch = 6; + // stage = Outfeed Thread, op = Get Next Batch. + kOutfeedThreadGetNextBatch = 7; + // stage = Outfeed Thread, op = Session Run. + kOutfeedThreadSessionRun = 8; + // stage = Outfeed Thread, op = Forward Batch. + kOutfeedThreadForwardBatch = 9; +} + +// Result proto for the host ops per TPU step. +message HostOpsPerTpuStep { + // Whether the data in this message is valid. + optional bool valid = 1 [default = false]; + // The current TPU step number. + optional uint32 tpu_step_num = 2; + // The beginning time of the current TPU step on the device in picoseconds. + optional uint64 tpu_step_begin_ps = 3; + // The ending time of the current TPU step on the device in picoseconds. + optional uint64 tpu_step_end_ps = 4; + // For each possible host operation, maps to the difference between the TPU + // step number that the host op targets and the current TPU step number. + // The key is HostOp, value is the step difference. + map step_diffs = 5; +} + +// Result proto for the host ops for all TPU steps. +message HostOpsResult { + // A sequence of HostOpsPerTpuStep (one for each TPU step) + repeated HostOpsPerTpuStep host_op_sequence = 1; +} + // Result proto for TfStatsHelper. message TfOpStats { // The result for the TF-metric database. @@ -171,4 +221,6 @@ message TfOpStats { optional double matrix_unit_utilization_percent = 6; // The run environment of this profiling session. optional RunEnvironmentResult run_environment = 7; + // The result for the host operations. + optional HostOpsResult host_ops = 8; } diff --git a/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto b/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto new file mode 100644 index 0000000000000000000000000000000000000000..a4fc8d4e879eb85522f35663c9c628ecd5ef562c --- /dev/null +++ b/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto @@ -0,0 +1,73 @@ +syntax = "proto3"; +package tensorflow; + +import "tensorflow/contrib/tpu/profiler/tpu_profiler.proto"; + +message NewProfileSessionRequest { + ProfileRequest request = 1; + string repository_root = 2; + repeated string hosts = 3; +} + +message NewProfileSessionResponse { + // Auxiliary error_message. + string error_message = 1; + // If success, return session identifier for future reference. + string session_id = 2; +} + +message EnumProfileSessionsAndToolsRequest { + string repository_root = 1; +} + +message ProfileSessionInfo { + string session_id = 1; + // Which tool data is available for consumption. + repeated string available_tools = 2; +} + +message EnumProfileSessionsAndToolsResponse { + // Auxiliary error_message. + string error_message = 1; + // If success, the returned sessions information are stored here. + repeated ProfileSessionInfo sessions = 2; +} + +message ProfileSessionDataRequest { + string repository_root = 1; + string session_id = 2; + // Which tool + string tool_name = 3; + // Tool's specific parameters. e.g. TraceViewer's viewport etc + map parameters = 4; +} + +message ProfileSessionDataResponse { + // Auxiliary error_message. + string error_message = 1; + + // Output format. e.g. "json" or "proto" or "blob" + string output_format = 2; + + // TODO(jiesun): figure out whether to put bytes or oneof tool specific proto. + bytes output = 3; +} +//////////////////////////////////////////////////////////////////////////////// +// TPUProfileAnalysis service provide entry point for profiling TPU and for +// serving profiled data to Tensorboard through GRPC +//////////////////////////////////////////////////////////////////////////////// +service TPUProfileAnalysis { + // Starts a profiling session, blocks until it completes. + // TPUProfileAnalysis service delegate this to TPUProfiler service. + // Populate the profiled data in repository, then return status to caller. + rpc NewSession(NewProfileSessionRequest) returns (NewProfileSessionResponse) { + } + // Enumerate existing sessions and return available profile tools. + rpc EnumSessions(EnumProfileSessionsAndToolsRequest) + returns (EnumProfileSessionsAndToolsResponse) { + } + // Retrieve specific tool's data for specific session. + rpc GetSessionToolData(ProfileSessionDataRequest) + returns (ProfileSessionDataResponse) { + } +} diff --git a/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis_pb2_grpc.py b/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis_pb2_grpc.py new file mode 100644 index 0000000000000000000000000000000000000000..c28fef22a9d3736748b1b56135302d5ec7845720 --- /dev/null +++ b/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis_pb2_grpc.py @@ -0,0 +1,138 @@ +# 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. +# ============================================================================== +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +# +# Do not use pylint on generated code. +# pylint: disable=missing-docstring,g-short-docstring-punctuation,g-no-space-after-docstring-summary,invalid-name,line-too-long,unused-argument,g-doc-args +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import grpc + +from third_party.tensorflow.contrib.tpu.profiler import tpu_profiler_analysis_pb2 as third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2 + + +class TPUProfileAnalysisStub(object): + """////////////////////////////////////////////////////////////////////////////// + + TPUProfileAnalysis service provide entry point for profiling TPU and for + serving profiled data to Tensorboard through GRPC + ////////////////////////////////////////////////////////////////////////////// + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.NewSession = channel.unary_unary( + '/tensorflow.TPUProfileAnalysis/NewSession', + request_serializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + NewProfileSessionRequest.SerializeToString, + response_deserializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + NewProfileSessionResponse.FromString, + ) + self.EnumSessions = channel.unary_unary( + '/tensorflow.TPUProfileAnalysis/EnumSessions', + request_serializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + EnumProfileSessionsAndToolsRequest.SerializeToString, + response_deserializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + EnumProfileSessionsAndToolsResponse.FromString, + ) + self.GetSessionToolData = channel.unary_unary( + '/tensorflow.TPUProfileAnalysis/GetSessionToolData', + request_serializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + ProfileSessionDataRequest.SerializeToString, + response_deserializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + ProfileSessionDataResponse.FromString, + ) + + +class TPUProfileAnalysisServicer(object): + """////////////////////////////////////////////////////////////////////////////// + + TPUProfileAnalysis service provide entry point for profiling TPU and for + serving profiled data to Tensorboard through GRPC + ////////////////////////////////////////////////////////////////////////////// + """ + + def NewSession(self, request, context): + """Starts a profiling session, blocks until it completes. + TPUProfileAnalysis service delegate this to TPUProfiler service. + Populate the profiled data in repository, then return status to caller. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def EnumSessions(self, request, context): + """Enumerate existing sessions and return available profile tools. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetSessionToolData(self, request, context): + """Retrieve specific tool's data for specific session. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_TPUProfileAnalysisServicer_to_server(servicer, server): + rpc_method_handlers = { + 'NewSession': + grpc.unary_unary_rpc_method_handler( + servicer.NewSession, + request_deserializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + NewProfileSessionRequest.FromString, + response_serializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + NewProfileSessionResponse.SerializeToString, + ), + 'EnumSessions': + grpc.unary_unary_rpc_method_handler( + servicer.EnumSessions, + request_deserializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + EnumProfileSessionsAndToolsRequest.FromString, + response_serializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + EnumProfileSessionsAndToolsResponse.SerializeToString, + ), + 'GetSessionToolData': + grpc.unary_unary_rpc_method_handler( + servicer.GetSessionToolData, + request_deserializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + ProfileSessionDataRequest.FromString, + response_serializer= + third__party_dot_tensorflow_dot_contrib_dot_tpu_dot_profiler_dot_tpu__profiler__analysis__pb2. + ProfileSessionDataResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'tensorflow.TPUProfileAnalysis', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ops.py b/tensorflow/contrib/tpu/python/ops/tpu_ops.py index 1c970655d0e464912d9b0a514fca0523bc604726..14c63a79763300dcfe8d6c8e09b90f8e9c772358 100644 --- a/tensorflow/contrib/tpu/python/ops/tpu_ops.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ops.py @@ -47,7 +47,8 @@ if platform.system() != "Windows": # types are supported. _SUPPORTED_INFEED_DTYPES = set([ - dtypes.bool, dtypes.int32, dtypes.bfloat16, dtypes.float32 + dtypes.bool, dtypes.int32, dtypes.int64, dtypes.bfloat16, dtypes.float32, + dtypes.complex64 ]) def infeed_dequeue(dtype, shape, name=None): diff --git a/tensorflow/contrib/tpu/python/tpu/datasets.py b/tensorflow/contrib/tpu/python/tpu/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..465c668fd8b42f150892f8e4b52de76c6fe13fa9 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/datasets.py @@ -0,0 +1,184 @@ +# 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. +# ====================================== +"""Library of Cloud TPU helper functions for data loading.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.data.python.ops import batching +from tensorflow.contrib.data.python.ops import interleave_ops +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.ops import iterator_ops +from tensorflow.python.data.ops import readers +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import function +from tensorflow.python.framework import ops +from tensorflow.python.ops import functional_ops + + +def _TextLineDataset(filename): + buffer_size = 8 * 1024 * 1024 # 8 MiB per file + dataset = readers.TextLineDataset(filename, buffer_size=buffer_size) + return dataset + + +def _TFRecordDataset(filename): + buffer_size = 8 * 1024 * 1024 # 8 MiB per file + dataset = readers.TFRecordDataset(filename, buffer_size=buffer_size) + return dataset + + +_FILETYPE_MAP = { + 'tfrecord': _TFRecordDataset, + 'textline': _TextLineDataset, + 'text': _TextLineDataset, +} + + +def StreamingFilesDataset(files, + filetype=None, + file_reader_job=None, + worker_job=None, + num_epochs=None, + filename_shuffle_buffer_size=None, + num_parallel_reads=None, + batch_transfer_size=None, + sloppy=None): + """StreamingFilesDataset constructs a dataset to stream from workers (GCE VM). + + Because Cloud TPUs are allocated over the network, a Cloud TPU cannot read + files local to your GCE VM. In order to train using files stored on your local + VM (e.g. on local SSD for extreme performance), use the StreamingFilesDataset + helper to generate a dataset to feed your Cloud TPU with files from your GCE + VM. + + The resulting dataset may return an OutOfRangeError if there are no files + found as a result of the fileglob expansion. + + Note: StreamingFilesDataset assumes that the session is using a + TPUClusterResolver and has therefore a worker and a coordinator job. File + loading will be done on the coordinator job. + + Args: + files: A string glob to match files, or a `tf.data.Dataset` generating file + names. + filetype: A string (one of 'tfrecord', or 'textline') or a single-argument + TensorFlow function that when given a filename returns a dataset. + file_reader_job: An optional string that corresponds to the job that should + perform the file reads. + worker_job: An optional string that corresponds to the job that should + process the tensors (i.e. your GPU or TPU worker). + num_epochs: The number of epochs through the training set that should be + generated. By default, it will repeat infinitely. + filename_shuffle_buffer_size: An optional integer whose value controls the + shuffling of the file names. If you would like to read from the files in + the same order, set to 0 or False. + num_parallel_reads: An optional integer controlling the number of files to + read from concurrently. (Set to 1 for no parallelism.) + batch_transfer_size: An optional integer controlling the batching used to + amortize the remote function invocation overhead. Set to a very large + number to increase throughput. Set to a very small number to reduce memory + consumption. Set to False to skip batching. + sloppy: (Optional.) If `False`, read input data while maintaining a + deterministic order. (This may have significant performance impacts.) + sloppy defaults to: True. + Returns: + A `tf.data.Dataset` with an infinite stream of elements generated by a + parallel interleaving of the set of files matched (or generated) by `files` + with a type is the output of the dataset specified by `filetype`. + + Raises: + ValueError: if any argument is not of the expected type. + """ + if filetype is None: + filetype = 'tfrecord' + + if isinstance(filetype, str): + if filetype not in _FILETYPE_MAP: + raise ValueError('Unexpected filetype: %s' % filetype) + reader_fn = _FILETYPE_MAP[filetype] + elif callable(filetype): + reader_fn = filetype + else: + raise ValueError('filetype should be a string or a callable') + + file_reader_job = file_reader_job or 'coordinator' + + worker_job = worker_job or 'worker' + + if filename_shuffle_buffer_size is None: + filename_shuffle_buffer_size = 4096 + + num_parallel_reads = num_parallel_reads or 8 + + if batch_transfer_size is None: + batch_transfer_size = 256 + + if sloppy is None: + sloppy = True + + with ops.device('/job:%s' % file_reader_job): + if isinstance(files, str): + source_dataset = dataset_ops.Dataset.list_files(files) + elif isinstance(files, dataset_ops.Dataset): + source_dataset = files + else: + raise ValueError('files was not a string or a dataset: %s' % files) + + if filename_shuffle_buffer_size: + source_dataset = source_dataset.shuffle( + buffer_size=filename_shuffle_buffer_size) + + # NOTE: We perform the `repeat` on the source dataset, because the output + # dataset does not currently have enough information to recreate an iterator + # over the source dataset when it reaches the end. + source_dataset = source_dataset.repeat(num_epochs) + + source_dataset = source_dataset.apply( + interleave_ops.parallel_interleave( + reader_fn, cycle_length=num_parallel_reads, sloppy=sloppy)) + + if batch_transfer_size: + source_dataset = source_dataset.batch(batch_transfer_size) + + source_dataset = source_dataset.prefetch(1) + + source_iterator = source_dataset.make_one_shot_iterator() + source_handle = source_iterator.string_handle() + + @function.Defun(dtypes.string) + def LoadingFunc(h): + remote_iterator = iterator_ops.Iterator.from_string_handle( + h, source_dataset.output_types, source_dataset.output_shapes) + return remote_iterator.get_next() + + def MapFn(unused_input): + return functional_ops.remote_call( + args=[source_handle], + Tout=[dtypes.string], + f=LoadingFunc, + target='/job:%s/replica:0/task:0/cpu:0' % file_reader_job) + + with ops.device('/job:%s' % worker_job): + output_dataset = dataset_ops.Dataset.range(2).repeat().map( + MapFn, num_parallel_calls=4 if sloppy else None) + output_dataset = output_dataset.prefetch(1) + + if batch_transfer_size: + # Undo the batching used during the transfer. + output_dataset = output_dataset.apply(batching.unbatch()).prefetch(1) + + return output_dataset diff --git a/tensorflow/contrib/tpu/python/tpu/datasets_test.py b/tensorflow/contrib/tpu/python/tpu/datasets_test.py new file mode 100644 index 0000000000000000000000000000000000000000..918cf0ed8e513de0d4207f7d2aac61ad886c8288 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/datasets_test.py @@ -0,0 +1,181 @@ +# 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. +# ============================================================================== +"""TPU datasets tests.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +from tensorflow.contrib.tpu.python.tpu import datasets +from tensorflow.core.protobuf import cluster_pb2 +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.ops import readers +from tensorflow.python.lib.io import python_io +from tensorflow.python.platform import test +from tensorflow.python.training import server_lib +from tensorflow.python.util import compat + +_NUM_FILES = 10 +_NUM_ENTRIES = 20 + + +class DatasetsTest(test.TestCase): + + def setUp(self): + super(DatasetsTest, self).setUp() + self._coord = server_lib.Server.create_local_server() + self._worker = server_lib.Server.create_local_server() + + self._cluster_def = cluster_pb2.ClusterDef() + worker_job = self._cluster_def.job.add() + worker_job.name = 'worker' + worker_job.tasks[0] = self._worker.target[len('grpc://'):] + coord_job = self._cluster_def.job.add() + coord_job.name = 'coordinator' + coord_job.tasks[0] = self._coord.target[len('grpc://'):] + + session_config = config_pb2.ConfigProto(cluster_def=self._cluster_def) + + self._sess = session.Session(self._worker.target, config=session_config) + + def testTextLineDataset(self): + all_contents = [] + for i in range(_NUM_FILES): + filename = os.path.join(self.get_temp_dir(), 'text_line.%d.txt' % i) + contents = [] + for j in range(_NUM_ENTRIES): + contents.append(compat.as_bytes('%d: %d' % (i, j))) + with open(filename, 'wb') as f: + f.write(b'\n'.join(contents)) + all_contents.extend(contents) + + dataset = datasets.StreamingFilesDataset( + os.path.join(self.get_temp_dir(), 'text_line.*.txt'), filetype='text') + + iterator = dataset.make_initializable_iterator() + self._sess.run(iterator.initializer) + get_next = iterator.get_next() + + retrieved_values = [] + for _ in range(4 * len(all_contents)): + retrieved_values.append(compat.as_bytes(self._sess.run(get_next))) + + self.assertEqual(set(all_contents), set(retrieved_values)) + + def testTFRecordDataset(self): + all_contents = [] + for i in range(_NUM_FILES): + filename = os.path.join(self.get_temp_dir(), 'tf_record.%d' % i) + writer = python_io.TFRecordWriter(filename) + for j in range(_NUM_ENTRIES): + record = compat.as_bytes('Record %d of file %d' % (j, i)) + writer.write(record) + all_contents.append(record) + writer.close() + + dataset = datasets.StreamingFilesDataset( + os.path.join(self.get_temp_dir(), 'tf_record*'), filetype='tfrecord') + + iterator = dataset.make_initializable_iterator() + self._sess.run(iterator.initializer) + get_next = iterator.get_next() + + retrieved_values = [] + for _ in range(4 * len(all_contents)): + retrieved_values.append(compat.as_bytes(self._sess.run(get_next))) + + self.assertEqual(set(all_contents), set(retrieved_values)) + + def testTFRecordDatasetFromDataset(self): + filenames = [] + all_contents = [] + for i in range(_NUM_FILES): + filename = os.path.join(self.get_temp_dir(), 'tf_record.%d' % i) + filenames.append(filename) + writer = python_io.TFRecordWriter(filename) + for j in range(_NUM_ENTRIES): + record = compat.as_bytes('Record %d of file %d' % (j, i)) + writer.write(record) + all_contents.append(record) + writer.close() + + filenames = dataset_ops.Dataset.from_tensor_slices(filenames) + + dataset = datasets.StreamingFilesDataset(filenames, filetype='tfrecord') + + iterator = dataset.make_initializable_iterator() + self._sess.run(iterator.initializer) + get_next = iterator.get_next() + + retrieved_values = [] + for _ in range(4 * len(all_contents)): + retrieved_values.append(compat.as_bytes(self._sess.run(get_next))) + + self.assertEqual(set(all_contents), set(retrieved_values)) + + def testArbitraryReaderFunc(self): + + def MakeRecord(i, j): + return compat.as_bytes('%04d-%04d' % (i, j)) + + record_bytes = len(MakeRecord(10, 200)) + + all_contents = [] + for i in range(_NUM_FILES): + filename = os.path.join(self.get_temp_dir(), 'fixed_length.%d' % i) + with open(filename, 'wb') as f: + for j in range(_NUM_ENTRIES): + record = MakeRecord(i, j) + f.write(record) + all_contents.append(record) + + def FixedLengthFile(filename): + return readers.FixedLengthRecordDataset(filename, record_bytes) + + dataset = datasets.StreamingFilesDataset( + os.path.join(self.get_temp_dir(), 'fixed_length*'), + filetype=FixedLengthFile) + + iterator = dataset.make_initializable_iterator() + self._sess.run(iterator.initializer) + get_next = iterator.get_next() + + retrieved_values = [] + for _ in range(4 * len(all_contents)): + retrieved_values.append(compat.as_bytes(self._sess.run(get_next))) + + self.assertEqual(set(all_contents), set(retrieved_values)) + + def testUnexpectedFiletypeString(self): + with self.assertRaises(ValueError): + datasets.StreamingFilesDataset( + os.path.join(self.get_temp_dir(), '*'), filetype='foo') + + def testUnexpectedFiletypeType(self): + with self.assertRaises(ValueError): + datasets.StreamingFilesDataset( + os.path.join(self.get_temp_dir(), '*'), filetype=3) + + def testUnexpectedFilesType(self): + with self.assertRaises(ValueError): + datasets.StreamingFilesDataset(123, filetype='tfrecord') + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/tpu/python/tpu/device_assignment.py b/tensorflow/contrib/tpu/python/tpu/device_assignment.py index ee202610a8a8a1406363b3010771e7806d5d84bf..726b2d248e3086e1882004827076ed3e563d960d 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 + logical_core. + Raises: + ValueError: If no replica exists 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_config.py b/tensorflow/contrib/tpu/python/tpu/tpu_config.py index 0c2580211ab7674d841ca1953c9327df9488bb8e..38b5ea23103730630ae8e1cdd7b9180a501013c5 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config.py @@ -23,7 +23,10 @@ import collections import json import os +import numpy as np + from tensorflow.contrib.tpu.python.tpu import util as util_lib +from tensorflow.core.protobuf import config_pb2 from tensorflow.python.estimator import run_config as run_config_lib from tensorflow.python.platform import tf_logging as logging @@ -31,29 +34,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.Topology} 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 +82,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 +100,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 +128,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) @@ -103,6 +141,7 @@ class RunConfig(run_config_lib.RunConfig): tpu_config=None, evaluation_master=None, master=None, + cluster=None, **kwargs): """Constructs a RunConfig. @@ -111,16 +150,26 @@ 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. + cluster: a ClusterResolver + **kwargs: keyword config parameters. + + Raises: + ValueError: if cluster is not None and the provided session_config has a + cluster_def already. """ super(RunConfig, self).__init__(**kwargs) self._tpu_config = tpu_config or TPUConfig() + self._cluster = cluster - # If user sets master and/or evaluation_master explicilty, including empty + # If user sets master and/or evaluation_master explicitly, including empty # string '', take it. Otherwise, take the values set by parent class. if master is not None: + if cluster is not None: + raise ValueError('Both master and cluster are set.') self._master = master + else: + if cluster: + self._master = cluster.master() if evaluation_master is not None: self._evaluation_master = evaluation_master @@ -134,6 +183,20 @@ class RunConfig(run_config_lib.RunConfig): # evaluation_master to master, unless user overwrites it. self._evaluation_master = self._master + # Set the ClusterSpec to use + if cluster: + self._cluster_spec = cluster.cluster_spec() + + # Merge the cluster_def into the ConfigProto. + if self._session_config is None: # pylint: disable=access-member-before-definition + self._session_config = config_pb2.ConfigProto(allow_soft_placement=True) + if self._session_config.HasField('cluster_def'): + raise ValueError( + 'You cannot provide a ClusterResolver and ' + 'session_config.cluster_def.') + self._session_config.cluster_def.CopyFrom( + self._cluster_spec.as_cluster_def()) + @property def evaluation_master(self): return self._evaluation_master @@ -146,6 +209,10 @@ class RunConfig(run_config_lib.RunConfig): def tpu_config(self): return self._tpu_config + @property + def cluster(self): + return self._cluster + def replace(self, **kwargs): if 'tpu_config' not in kwargs: return super(RunConfig, self).replace(**kwargs) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py b/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py index 60884aa32f932413b49ea2193a145828489ea04c..37ef3dbe1e66efe18b13ab9153ee346c08b9774a 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py @@ -43,6 +43,16 @@ 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): 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..3bac2db77e95520a6c9c4c17658267a9a6588d94 --- /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, based 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 distinguish 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 9d715bb23683ae4b35cef70ffa7c373dd613b909..52a3ee8283dc374ab14b07e0265f5a7c15a5589f 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import collections -from contextlib import contextmanager import copy import signal import threading @@ -34,18 +33,19 @@ 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 @@ -61,6 +61,7 @@ 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. @@ -68,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 @@ -126,7 +133,7 @@ def _increase_eval_step_op(iterations_per_loop): """Returns an op to increase the eval step for TPU evaluation. Args: - iterations_per_loop: Tensor. The number of eval steps runnining in TPU + iterations_per_loop: Tensor. The number of eval steps running in TPU system before returning to CPU host for each `Session.run`. Returns: @@ -140,234 +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 global_batch_size(self): - mode = self._assert_mode() - if mode == model_fn_lib.ModeKeys.EVAL and self._eval_batch_size is None: - raise RuntimeError('Internal error, EVAL on TPU is not enabled, but ' - '`global_batch_size` is called.') - return (self._train_batch_size - if mode == model_fn_lib.ModeKeys.TRAIN else self._eval_batch_size) - - @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. @@ -415,14 +194,13 @@ class TPUEstimatorSpec( 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. `host_call` currently works for train() and - evaluate(). The function's graph 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}. Note that `host_call` does not - currently work if `use_tpu` is set to False. + 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, @@ -454,10 +232,18 @@ class TPUEstimatorSpec( def as_estimator_spec(self): """Creates an equivalent `EstimatorSpec` used by CPU train/eval.""" + 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 = _OutfeedHostCall.create_cpu_hostcall( - {'eval_metrics': self.eval_metrics})['eval_metrics'] + 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, @@ -466,7 +252,10 @@ class TPUEstimatorSpec( train_op=self.train_op, eval_metric_ops=eval_metric_ops, export_outputs=self.export_outputs, - scaffold=scaffold) + scaffold=scaffold, + training_hooks=hooks, + evaluation_hooks=hooks, + prediction_hooks=hooks) class _OpQueueContext(object): @@ -501,6 +290,22 @@ class _OpQueueContext(object): self._thread.join() +class _OpSignalOnceQueueContext(_OpQueueContext): + """Manages work queue and thread for a infeed/outfeed thread. + + This subclass only signals once. + """ + + def __init__(self, name, target, args): + super(_OpSignalOnceQueueContext, self).__init__(name, target, args) + self._has_signaled = False + + def send_next_batch_signal(self, iterations): + if not self._has_signaled: + self._queue.put(iterations) + self._has_signaled = True + + class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): """A Session hook setting up the TPU initialization, infeed, and outfeed. @@ -510,12 +315,19 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): dequeue. """ - def __init__(self, ctx, enqueue_ops, dequeue_ops): + 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 + + self._run_infeed_loop_on_coordinator = run_infeed_loop_on_coordinator self._initial_infeed_sleep_secs = ( ctx.config.tpu_config.initial_infeed_sleep_secs) + self._session_cancel_timer = None self._feed_error = None @@ -567,7 +379,7 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): # 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 @@ -580,7 +392,7 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): # If the main session is still running, the infeed/outfeed errors are # legitimate, and should be logged. - if not self._finished: + 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() @@ -598,15 +410,15 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): logging.info('%s thread starting after sleep', self._name) try: - if _WRAP_INPUT_FN_INTO_WHILE_LOOP: - for _ in queue_ctx.read_iteration_counts(): - session.run(self._enqueue_ops) - else: + 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) - logging.debug('Infeed thread finished, shutting down.') + 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) @@ -617,16 +429,20 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): 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) + 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_ops, options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000)) logging.info('Start infeed thread controller') - self._infeed_controller = _OpQueueContext( + self._infeed_controller = self._create_infeed_controller( name='InfeedController', target=self._run_infeed, args=(session,)) logging.info('Start outfeed thread controller') @@ -634,17 +450,18 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): name='OutfeedController', target=self._run_outfeed, args=(session,)) def before_run(self, run_context): - if self._feed_error: - logging.warning('Feed error occurred, terminating session.') - run_context.request_stop() - return + 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) - # TODO(xiejw): Refactor the outfeed dequeue into tf.while_loop. logging.info('Dequeue next (%d) batch(es) of data from outfeed.', iterations) self._outfeed_controller.send_next_batch_signal(iterations) @@ -665,6 +482,16 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): 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): """Hook that requests stop at a specified step. @@ -752,6 +579,47 @@ class _SetEvalIterationsHook(session_run_hook.SessionRunHook): self._iterations_per_loop_var.load(self._num_steps, session=session) +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 provided) + # batch 1: images, labels, stop = 0 (user provided) + # ... + # batch 99: images, labels, stop = 0 (user provided) + # 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" prediction, 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.""" @@ -763,11 +631,14 @@ def generate_per_core_enqueue_ops_fn_for_host(ctx, input_fn, 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) @@ -790,36 +661,76 @@ def generate_per_core_enqueue_ops_fn_for_host(ctx, input_fn, 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 - inputs_structure_recorder.validate_and_record_structure(features, labels) + 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, 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 + unsharded_tensor_list, + 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 + return enqueue_ops_fn, captured_infeed_queue, hooks, is_dataset class _InputPipeline(object): @@ -840,7 +751,7 @@ class _InputPipeline(object): 2. (features, labels) Internally, form 1 is reformed to `(features, None)` as features and labels - are passed separatedly to underlying methods. For TPU training, TPUEstimator + are passed separately to underlying methods. For TPU training, TPUEstimator may expect multiple `features` and `labels` tuples one for each core. TPUEstimator allows various different structures for inputs (namely `features` @@ -859,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 @@ -866,19 +778,24 @@ 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 [] - return tensor_or_dict.keys() if isinstance(tensor_or_dict, dict) else [] + return sorted(tensor_or_dict.keys()) if isinstance( + tensor_or_dict, dict) else [] # Extract structure. has_labels = labels is not None 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' @@ -891,7 +808,7 @@ 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: @@ -907,6 +824,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): @@ -932,7 +852,11 @@ class _InputPipeline(object): 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( @@ -949,13 +873,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. @@ -983,25 +914,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 @@ -1015,6 +955,7 @@ class _InputPipeline(object): self._ctx, self._input_fn, self._inputs_structure_recorder)) if _WRAP_INPUT_FN_INTO_WHILE_LOOP: + run_infeed_loop_on_coordinator = False enqueue_ops.append( _wrap_computation_in_while_loop( device=host_device, op_fn=enqueue_ops_fn)) @@ -1028,15 +969,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: + 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_computation_in_while_loop( - device=host_device, op_fn=enqueue_ops_fn)) + wrap_fn(device=host_device, op_fn=enqueue_ops_fn)) else: enqueue_ops.append(enqueue_ops_fn()) infeed_queues.append(captured_infeed_queue.get()) @@ -1044,7 +1002,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 @@ -1101,7 +1059,8 @@ 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) @@ -1110,23 +1069,28 @@ class _ModelFnWrapper(object): 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)) loss, train_op = estimator_spec.loss, estimator_spec.train_op - host_call_outfeed_ops = [] if isinstance(estimator_spec, TPUEstimatorSpec): captured_scaffold_fn.capture(estimator_spec.scaffold_fn) - if 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() else: captured_scaffold_fn.capture(None) - with ops.control_dependencies([train_op] + host_call_outfeed_ops): - return array_ops.identity(loss) + # We must run train_op to update the variables prior to running the + # outfeed. + with ops.control_dependencies([train_op]): + 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 @@ -1153,15 +1117,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 `_OutfeedHostCall` instance. + A tuple of eval_fn, host_calls, and captured scaffold_fn. The eval_fn + representing the eval step for TPU. """ 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): @@ -1174,6 +1139,8 @@ class _ModelFnWrapper(object): 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) @@ -1182,7 +1149,55 @@ class _ModelFnWrapper(object): return eval_step, host_calls, captured_scaffold_fn - def _call_model_fn(self, features, labels): + 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 = {} @@ -1213,7 +1228,7 @@ class _ModelFnWrapper(object): 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`. @@ -1268,11 +1283,14 @@ class _OutfeedHostCall(object): 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]) - if len(host_call[1]) != len(fn_args): + # 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 metric_fn.'.format(name)) + '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 create_cpu_hostcall(host_calls): @@ -1331,7 +1349,9 @@ class _OutfeedHostCall(object): # TODO(jhseu): Consider deduping tensors. for name in self._names: tensors.extend(self._tensors[name]) - return [tpu_ops.outfeed_enqueue_tuple(tensors)] + + with ops.device(tpu.core(0)): + return [tpu_ops.outfeed_enqueue_tuple(tensors)] def create_tpu_hostcall(self): """Sends the tensors through outfeed and runs the host_fn on CPU. @@ -1364,10 +1384,11 @@ class _OutfeedHostCall(object): for shape in self._tensor_shapes[name]: tensor_shapes.append(shape) - # Outfeed ops execute on each JF node. Note: we must constraint it such that - # we have at most one outfeed dequeue and enqueue. + # 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(self._ctx.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=tensor_dtypes, shapes=tensor_shapes) @@ -1414,6 +1435,34 @@ class _OutfeedHostCall(object): 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) + + 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.""" @@ -1464,23 +1513,23 @@ 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. + 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, 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 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` (See docstring for TPUConfig for details). - 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. + - For evaluation and prediction, `model_fn` gets per-core batch size and + `input_fn` get per-host batch size. + + Evaluation + ========== + + `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 @@ -1488,13 +1537,22 @@ class TPUEstimator(estimator_lib.Estimator): 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. - 2. `input_fn` for evaluation should not throw OutOfRange error for all - evaluation steps and all batches should have the same size. + 1. TPU evaluation only works on a single host (one TPU worker). + + 2. `input_fn` for evaluation should **NOT** raise an end-of-input exception + (`OutOfRangeError` or `StopIteration`). And all evaluation steps and all + batches should have the same size. Example (MNIST): + ---------------- + ``` # The metric Fn which runs on CPU. def metric_fn(labels, logits): @@ -1530,8 +1588,120 @@ class TPUEstimator(estimator_lib.Estimator): })) ``` - Predict support on TPU is not yet implemented. So, `predict` and - `export_savedmodel` are executed on CPU, even if `use_tpu` is true. + Prediction + ========== + + Prediction on TPU is an experimental feature to support large batch inference. + It is not designed for latency-critical system. In addition, due to some + usability issues, for prediction with small dataset, CPU `.predict`, i.e., + creating a new `TPUEstimator` instance with `use_tpu=False`, might be more + convenient. + + Note: In contrast to TPU training/evaluation, the `input_fn` for prediction + *should* raise an end-of-input exception (`OutOfRangeError` or + `StopIteration`), which serves as the stopping signal to `TPUEstimator`. To be + precise, the ops created by `input_fn` produce one batch of the data. + The `predict()` API processes one batch at a time. When reaching the end of + the data source, an end-of-input exception should be raised by one of these + operations. The user usually does not need to do this manually. As long as the + dataset is not repeated forever, the `tf.data` API will raise an end-of-input + exception automatically after the last batch has been produced. + + Note: Estimator.predict returns a Python generator. Please consume all the + data from the generator so that TPUEstimator can shutdown the TPU system + properly for user. + + Current limitations: + -------------------- + 1. TPU prediction only works on a single host (one TPU worker). + + 2. `input_fn` must return a `Dataset` instance rather than `features`. In + fact, .train() and .evaluate() also support Dataset as return value. + + 3. Each batch returned by `Dataset`'s iterator must have the *same static* + shape. This means two things: + - batch_size cannot be `None` + - the final batch must be padded by user to a full batch. + + Example (MNIST): + ---------------- + ``` + height = 32 + width = 32 + total_examples = 100 + + def predict_input_fn(params): + batch_size = params['batch_size'] + + images = tf.random_uniform( + [total_examples, height, width, 3], minval=-1, maxval=1) + + dataset = tf.data.Dataset.from_tensor_slices(images) + dataset = dataset.batch(batch_size) + dataset = dataset.map(lambda images: {'image': images}) + + def pad(tensor, missing_count): + # Pads out the batch dimension to the complete batch_size. + rank = len(tensor.shape) + assert rank > 0 + padding = tf.stack([[0, missing_count]] + [[0, 0]] * (rank - 1)) + padded_shape = (batch_size,) + tuple(tensor.shape[1:]) + padded_tensor = tf.pad(tensor, padding) + padded_tensor.set_shape(padded_shape) + return padded_tensor + + def pad_batch_if_incomplete(batch_features): + # Pads out the batch dimension for all features. + real_batch_size = tf.shape(batch_features["image"])[0] + + missing_count = tf.constant(batch_size, tf.int32) - real_batch_size + + padded_features = { + key: pad(tensor, missing_count) + for key, tensor in batch_features.iteritems() + } + padding_mask = tf.concat( + [ + tf.zeros((real_batch_size, 1), dtype=tf.int32), + tf.ones((missing_count, 1), dtype=tf.int32) + ], + axis=0) + padding_mask.set_shape((batch_size, 1)) + padded_features["is_padding"] = padding_mask + return padded_features + + dataset = dataset.map(pad_batch_if_incomplete) + + return dataset + + def model_fn(features, labels, params, mode): + # Generate predictions, called 'output', from features['image'] + + if mode == tf.estimator.ModeKeys.PREDICT: + return tf.contrib.tpu.TPUEstimatorSpec( + mode=mode, + predictions={ + 'predictions': output, + 'is_padding': features['is_padding'] + }) + + tpu_est = TPUEstimator( + model_fn=model_fn, + ..., + predict_batch_size=16) + + # Fully consume the generator so that TPUEstimator can shutdown the TPU + # system. + for item in tpu_est.predict(input_fn=input_fn): + # Filter out item if the `is_padding` is 1. + # Process the 'predictions' + ``` + + Exporting + ========= + + Exporting `SavedModel` support on TPU is not yet implemented. So, + `export_savedmodel` is executed on CPU, even if `use_tpu` is true. """ def __init__(self, @@ -1542,6 +1712,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. @@ -1560,18 +1731,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) @@ -1595,30 +1765,24 @@ class TPUEstimator(estimator_lib.Estimator): _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 @@ -1638,8 +1802,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. @@ -1720,6 +1887,9 @@ class TPUEstimator(estimator_lib.Estimator): if 'config' in input_fn_args: kwargs['config'] = config + if 'mode' in input_fn_args: + kwargs['mode'] = mode + with self._ctx.with_mode(mode) as ctx: # Setting the batch size in params first. This helps user to have same # input_fn for use_tpu=True/False. @@ -1727,7 +1897,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) @@ -1746,6 +1918,17 @@ class TPUEstimator(estimator_lib.Estimator): return _input_fn + def _validate_features_in_predict_input(self, result): + """Skip the validation. + + For TPUEstimator, we do not need to check the result type. `_InputPipeline` + has stronger check. Parent class's check generates confusing warning msg. + + Args: + result: `features` returned by input_fn. + """ + pass + def _augment_model_fn(self, model_fn, batch_axis): """Returns a new model_fn, which wraps the TPU support.""" @@ -1754,8 +1937,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) @@ -1765,7 +1953,7 @@ 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: @@ -1775,7 +1963,12 @@ class TPUEstimator(estimator_lib.Estimator): if host_ops is None: host_ops = [] hooks = [ - TPUInfeedOutfeedSessionHook(ctx, enqueue_ops, host_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( @@ -1784,7 +1977,7 @@ class TPUEstimator(estimator_lib.Estimator): '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() @@ -1799,50 +1992,121 @@ class TPUEstimator(estimator_lib.Estimator): train_op=control_flow_ops.group(*update_ops), scaffold=scaffold) - # Now eval. - total_loss, host_calls, 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() - 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'] + + 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) + hooks = [ - TPUInfeedOutfeedSessionHook(ctx, enqueue_ops, - eval_update_ops + host_ops), - ] + _StoppingPredictHook(scalar_stopping_signal), + TPUInfeedOutfeedSessionHookForPrediction(ctx, enqueue_ops, + host_ops), + ] + input_hooks + + # TODO(b/73813593): Delete this logging once the bug is resolved. + logging.info( + 'If the Tensors in TPUEstimatorSpec.predictions dict are large, ' + 'you might observe the TPU program getting stuck (b/73813593). ' + 'Consider using small Tensors in the predictions dict to verify ' + 'the issue and report on the bug.') 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 @@ -1850,7 +2114,6 @@ class TPUEstimator(estimator_lib.Estimator): 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, host_calls, captured_scaffold_fn = ( @@ -1865,8 +2128,9 @@ def _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): (loss,) = tpu.shard( multi_tpu_eval_steps_on_single_shard, inputs=[], - num_shards=num_cores, - outputs_from_all_shards=False) + num_shards=ctx.num_replicas, + outputs_from_all_shards=False, + device_assignment=ctx.device_assignment) scaffold = _get_scaffold(captured_scaffold_fn) return loss, host_calls, scaffold @@ -1874,7 +2138,6 @@ def _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): 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, host_call, captured_scaffold_fn = ( @@ -1889,11 +2152,40 @@ def _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): (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, host_call, scaffold + return dummy_predict_op, host_calls, scaffold def _wrap_computation_in_while_loop(device, op_fn): @@ -1914,6 +2206,29 @@ def _wrap_computation_in_while_loop(device, op_fn): 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(): """Validate graph before running distributed training. @@ -2004,3 +2319,199 @@ class _CapturingContext(control_flow_ops.ControlFlowContext): def __exit__(self, _, __, ___): # pylint: disable=invalid-name 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): + """The map fn to insert signals.""" + if len(args) == 1: + # Unpack the single Tensor/dict argument as features. This is required + # for the input_fn returns no labels. + args = args[0] + 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/training_loop.py b/tensorflow/contrib/tpu/python/tpu/training_loop.py index 3d7896127a99653167f164873331a2cc95f656e8..82a75d02552b7b013452945a76b16c2c2fb9fa82 100644 --- a/tensorflow/contrib/tpu/python/tpu/training_loop.py +++ b/tensorflow/contrib/tpu/python/tpu/training_loop.py @@ -170,7 +170,7 @@ def while_loop(condition, body, inputs=None, infeed_queue=None, name=None): def repeat(n, body, inputs=None, infeed_queue=None, name=None): - """Builds a training loop that executes a fixed number of interations. + """Builds a training loop that executes a fixed number of iterations. The set of loop-carried tensors correspond to `inputs`. `body` must be a function that takes and returns the values of the diff --git a/tensorflow/contrib/tpu/tpu_estimator.md b/tensorflow/contrib/tpu/tpu_estimator.md index ca1255b16b1575d291df51dfde696b36c38359ae..4ef8f9eebdb165e5fe221be8670276bf943159b3 100644 --- a/tensorflow/contrib/tpu/tpu_estimator.md +++ b/tensorflow/contrib/tpu/tpu_estimator.md @@ -231,7 +231,7 @@ Refer to this link for all [Cloud TPU documentation](https://cloud.google.com/tp ### Profiling -You can profile the `worker` by using instructions as spcified in the [Cloud TPU Tools](https://cloud.google.com/tpu/docs/cloud-tpu-tools). +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? diff --git a/tensorflow/contrib/training/BUILD b/tensorflow/contrib/training/BUILD index cccaa2b833ee764921508a5b6d6affe0b8822ede..6ae2f382528c37ae647b73ea01a7f88c07580c78 100644 --- a/tensorflow/contrib/training/BUILD +++ b/tensorflow/contrib/training/BUILD @@ -26,6 +26,7 @@ py_library( "python/training/resample.py", "python/training/sampling_ops.py", "python/training/sequence_queueing_state_saver.py", + "python/training/tensor_queue_dataset.py", "python/training/training.py", "python/training/tuner.py", ], @@ -285,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( @@ -301,7 +324,6 @@ tf_proto_library( name = "protos_all", srcs = glob(["**/*.proto"]), cc_api_version = 2, - go_api_version = 2, java_api_version = 2, visibility = ["//visibility:public"], ) diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py index fdfd27d6a414933b0bec824bae512c45dac24d3c..95e051e3b5bb9f8075e66891a45c64a27bca68d1 100644 --- a/tensorflow/contrib/training/python/training/hparam.py +++ b/tensorflow/contrib/training/python/training/hparam.py @@ -358,6 +358,8 @@ class HParams(object): ``` """ + _HAS_DYNAMIC_ATTRIBUTES = True # Required for pytype checks. + def __init__(self, hparam_def=None, model_structure=None, **kwargs): """Create an instance of `HParams` from keyword arguments. 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/loader.py b/tensorflow/contrib/util/loader.py index f4283cd9ed6eb185ec56bdfef45f6ddc5253d128..dca01d26f4c6291074d855322c8215027cfd2ace 100644 --- a/tensorflow/contrib/util/loader.py +++ b/tensorflow/contrib/util/loader.py @@ -42,9 +42,10 @@ def load_op_library(path): plugin. """ if os.name == 'nt': - # To avoid makeing every user_ops aware of windows, re-write - # the file extension from .so to .dll. - path = re.sub(r'\.so$', '.dll', path) + # To avoid making every user_ops aware of windows, re-write + # the file extension from .so to .dll if .so file doesn't exist. + if not os.path.exists(path): + path = re.sub(r'\.so$', '.dll', path) # Currently we have only some user_ops as dlls on windows - don't try # to load them if the dll is not found. diff --git a/tensorflow/contrib/verbs/README.md b/tensorflow/contrib/verbs/README.md index 1b99f4ce4f645d0c59b2552cf26f47495cbbba73..4b6104a8b4d542b1d8a9cb3e48eeed4950d791cd 100644 --- a/tensorflow/contrib/verbs/README.md +++ b/tensorflow/contrib/verbs/README.md @@ -25,9 +25,9 @@ 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. Few remedies are possible: -1. The memory is pinned, transfered, 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. +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 alloactors, 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 registred address, which will allow us to do direct RDMA writes to it. +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. For best performance, we will adopt HKUST 0 copies approach in our solution. This means: @@ -77,7 +77,7 @@ When the receiver receives the **RDMA_MESSAGE_META_DATA_RESPONSE**, it will loca 1. Update the local meta-data cache. 2. Reallocate the result/proxy tensors. -3. Re-send the tensor request. For tracability, the new message has a different name: **RDMA_MESSAGE_TENSOR_RE_REQUEST**. +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. @@ -93,7 +93,7 @@ When the receiver receives the RDMA write, it will locate the relevant **RdmaTen 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 immediatly. + * If the tensor is ready before the request arrives, than the tensor is put in a local table. When the request arrives, it will invoke the callback immediately. In code it is done by calling **RecvLocalAsync()**, which receives the tensor's key, step-id, and the callback. 2. When the callback is invoked, the relevant tensor is removed from the tag matching table. In the case where we need to send the tensor's meta-data, the **RdmaTensorResponse** will store a copy of the tensor until the re-request arrives. 3. The sending of protocol messages (**RDMA_MESSAGE_TENSOR_REQUEST**, **RDMA_MESSAGE_META_DATA_RESPONSE** and **RDMA_MESSAGE_TENSOR_RE_REQUEST**) is done by the class **RdmaMessageBuffer**. All messages are sent using RDMA writes from/to fixed messages buffers. This implies that we cannot send on a specific channel more than one message at a time. In order to synchronize the messages, the **RdmaMessageBuffer** holds the a local and remote buffer statuses which can be either busy or idle. When a write is issued, both statuses will be changed to busy. When the write-complete event is received, the local status is changed to idle. When the write is received on the remote side, the remote side will parse the message, and return an ACK back to the sending side on which the sending side will update the remote status to idle. When both the local and remote statuses are idle, the next message can be sent. @@ -115,7 +115,7 @@ When the receiver receives the RDMA write, it will locate the relevant **RdmaTen * 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 avaliable). + * 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. @@ -153,7 +153,7 @@ When the receiver receives the RDMA write, it will locate the relevant **RdmaTen * 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-requset after meta-data update and reallocation of result/proxy tensors. +* **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. diff --git a/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md b/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md index 956b8f2147cf8154b6f1ade006d7bff194864c9b..da6fdd48e19e9d1503d1537926b1c464a0e77589 100644 --- a/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md +++ b/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md @@ -64,7 +64,7 @@ The protocol messages themselves will remain mostly unchanged at the first stage * type - The message type. * request_index - Request index. * is_dead/data_type/tensor_shape/tensor_bytes - The up-to-date meta-data. -* **RDMA_MESSAGE_BUFFER_RESPONSE** - (receiver ==> sender) Tensor re-requset after meta-data update and reallocation of result/proxy tensors. +* **RDMA_MESSAGE_BUFFER_RESPONSE** - (receiver ==> sender) Tensor re-request after meta-data update and reallocation of result/proxy tensors. * type - The message type. * name (name_size) - Name of the requested tensor. * step_id - Step ID. diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 063a4e9d306221132931ebcf8ece19746794d6c6..c7d12920754da150d34d3069093a798cee38008f 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -193,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", ] @@ -211,7 +212,6 @@ ADDITIONAL_CORE_PROTO_SRCS = [ "protobuf/named_tensor.proto", "protobuf/saved_model.proto", "protobuf/tensorflow_server.proto", - "util/event.proto", "util/test_log.proto", ] @@ -220,7 +220,6 @@ tf_proto_library( srcs = CORE_PROTO_SRCS + ADDITIONAL_CORE_PROTO_SRCS, cc_api_version = 2, default_header = True, - go_api_version = 2, j2objc_api_version = 1, java_api_version = 2, js_api_version = 2, @@ -314,6 +313,7 @@ cc_library( "lib/gtl/optional.h", "lib/gtl/priority_queue_util.h", "lib/hash/crc32c.h", + "lib/hash/hash.h", "lib/histogram/histogram.h", "lib/io/buffered_inputstream.h", "lib/io/compression.h", @@ -339,6 +339,7 @@ cc_library( "lib/strings/strcat.h", "lib/strings/stringprintf.h", "platform/abi.h", + "platform/context.h", "platform/cpu_feature_guard.h", "platform/cpu_info.h", "platform/dynamic_annotations.h", @@ -376,6 +377,17 @@ cc_library( 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"], @@ -433,6 +445,7 @@ tf_cuda_library( "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", @@ -453,6 +466,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", @@ -467,6 +481,7 @@ tf_cuda_library( "framework/type_index.h", "framework/type_traits.h", "framework/types.h", + "framework/visitable_allocator.h", "public/version.h", "util/activation_mode.h", "util/bcast.h", @@ -785,6 +800,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", @@ -824,6 +840,7 @@ 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", @@ -973,22 +990,15 @@ filegroup( # Core sources for Android builds. filegroup( - name = "mobile_srcs", + name = "mobile_srcs_no_runtime", srcs = [ ":proto_text_srcs_all", - "//tensorflow/core/kernels:android_srcs", "//tensorflow/core/platform/default/build_config:android_srcs", - "//tensorflow/core/util/ctc:android_srcs", - "//tensorflow/core/util/tensor_bundle:android_srcs", ] + glob( [ "client/**/*.cc", - "common_runtime/**/*.h", - "common_runtime/**/*.cc", "framework/**/*.h", "framework/**/*.cc", - "graph/**/*.h", - "graph/**/*.cc", "lib/**/*.h", "lib/**/*.cc", "platform/**/*.h", @@ -1004,7 +1014,6 @@ filegroup( "**/*main.cc", "debug/**/*", "framework/op_gen_*", - "graph/dot.*", "lib/jpeg/**/*", "lib/png/**/*", "lib/gif/**/*", @@ -1021,13 +1030,52 @@ filegroup( "platform/stream_executor.*", "platform/windows/**/*", "user_ops/**/*.cu.cc", + "util/ctc/*.h", + "util/ctc/*.cc", + "util/tensor_bundle/*.h", + "util/tensor_bundle/*.cc", + "common_runtime/gpu/**/*", + "common_runtime/gpu_device_factory.*", + ], + ), + visibility = ["//visibility:public"], +) + +filegroup( + name = "mobile_srcs_only_runtime", + srcs = [ + "//tensorflow/core/kernels:android_srcs", + "//tensorflow/core/util/ctc:android_srcs", + "//tensorflow/core/util/tensor_bundle:android_srcs", + ] + glob( + [ + "common_runtime/**/*.h", + "common_runtime/**/*.cc", + "graph/**/*.h", + "graph/**/*.cc", + ], + exclude = [ + "**/*test.*", + "**/*testutil*", + "**/*testlib*", + "**/*main.cc", "common_runtime/gpu/**/*", "common_runtime/gpu_device_factory.*", + "graph/dot.*", ], ), visibility = ["//visibility:public"], ) +filegroup( + name = "mobile_srcs", + srcs = [ + ":mobile_srcs_no_runtime", + ":mobile_srcs_only_runtime", + ], + visibility = ["//visibility:public"], +) + # Native library support for Android applications. Does not contain # operators, use :android_tensorflow_lib if you want full operator # support. @@ -1320,6 +1368,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", @@ -1355,6 +1410,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 @@ -1457,6 +1519,7 @@ LIB_INTERNAL_PUBLIC_HEADERS = tf_additional_lib_hdrs() + [ "lib/strings/base64.h", "lib/strings/ordered_code.h", "lib/strings/proto_text_util.h", + "lib/strings/proto_serialization.h", "lib/strings/scanner.h", "lib/wav/wav_io.h", "platform/demangle.h", @@ -1602,6 +1665,25 @@ cc_library( ], ) +cc_library( + name = "tflite_portable_logging", + srcs = [], + hdrs = [ + "lib/bfloat16/bfloat16.h", + "platform/default/integral_types.h", + "platform/default/logging.h", + "platform/logging.h", + "platform/macros.h", + "platform/platform.h", + "platform/types.h", + ], + copts = tf_copts(), + linkopts = ["-ldl"], + deps = [ + "//tensorflow/core/platform/default/build_config:logging", + ], +) + cc_library( name = "android_jpeg_internal", srcs = if_android([ @@ -1714,6 +1796,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", @@ -1731,6 +1816,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({ @@ -1748,6 +1834,7 @@ FRAMEWORK_INTERNAL_PUBLIC_HEADERS = [ "framework/tracking_allocator.h", # only needed for tests "framework/unique_tensor_references.h", "framework/variant.h", + "framework/visitable_allocator.h", "platform/variant_coding.h", "util/command_line_flags.h", "util/env_var.h", @@ -1788,6 +1875,13 @@ cc_header_only_library( ], ) +cc_header_only_library( + name = "core_cpu_headers_lib", + deps = [ + ":core_cpu_lib", + ], +) + tf_cuda_library( name = "framework_internal_impl", srcs = FRAMEWORK_INTERNAL_PRIVATE_HEADERS + [ @@ -1800,6 +1894,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", ], @@ -1814,6 +1911,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({ @@ -1849,7 +1947,7 @@ tf_cuda_library( ) + if_mkl( [ "//third_party/mkl:intel_binary_blob", - "@mkl_dnn//:mkl_dnn", + "@mkl_dnn", ], ), alwayslink = 1, @@ -1861,7 +1959,6 @@ cc_header_only_library( deps = [ ":framework", ":reader_base", - "@nsync//:nsync_headers", ], ) @@ -1902,7 +1999,6 @@ cc_library( tf_cuda_library( name = "cuda_device_functions", hdrs = ["util/cuda_device_functions.h"], - cuda_deps = ["//third_party_gpus/cuda:cuda_headers"], visibility = ["//visibility:public"], deps = [":framework_lite"], ) @@ -1929,6 +2025,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", @@ -1949,12 +2046,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, @@ -1971,18 +2065,24 @@ tf_cuda_library( CORE_CPU_BASE_HDRS = GRAPH_HDRS + [ "common_runtime/device.h", + "common_runtime/device_mgr.h", + "common_runtime/eval_const_tensor.h", "common_runtime/graph_runner.h", "common_runtime/shape_refiner.h", "framework/versions.h", + "common_runtime/process_function_library_runtime.h", + "common_runtime/function.h", ] tf_cuda_library( name = "core_cpu_base", srcs = [ + "common_runtime/eval_const_tensor.cc", "common_runtime/shape_refiner.cc", "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", @@ -2018,30 +2118,28 @@ CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [ "common_runtime/costmodel_manager.h", "common_runtime/debugger_state_interface.h", "common_runtime/device_factory.h", - "common_runtime/device_mgr.h", "common_runtime/device_set.h", "common_runtime/dma_helper.h", "common_runtime/eigen_thread_pool.h", "common_runtime/executor.h", - "common_runtime/function.h", "common_runtime/graph_optimizer.h", "common_runtime/local_device.h", "common_runtime/memory_types.h", "common_runtime/mkl_cpu_allocator.h", "common_runtime/optimization_registry.h", "common_runtime/pending_counts.h", - "common_runtime/process_function_library_runtime.h", "common_runtime/process_util.h", "common_runtime/profile_handler.h", "common_runtime/renamed_device.h", "common_runtime/rendezvous_mgr.h", "common_runtime/rendezvous_util.h", + "common_runtime/scoped_allocator.h", + "common_runtime/scoped_allocator_mgr.h", "common_runtime/session_factory.h", "common_runtime/placer.h", "common_runtime/stats_publisher_interface.h", "common_runtime/step_stats_collector.h", "common_runtime/threadpool_device.h", - "common_runtime/visitable_allocator.h", "graph/gradients.h", "graph/quantize_training.h", ] + if_mkl(["graph/mkl_graph_util.h"]) @@ -2067,6 +2165,7 @@ tf_cuda_library( "common_runtime/graph_runner.cc", "common_runtime/local_device.cc", "common_runtime/memory_types.cc", + "common_runtime/mkl_cpu_allocator.cc", "common_runtime/optimization_registry.cc", "common_runtime/parallel_concat_optimizer.cc", "common_runtime/placer.cc", @@ -2075,6 +2174,8 @@ tf_cuda_library( "common_runtime/renamed_device.cc", "common_runtime/rendezvous_mgr.cc", "common_runtime/rendezvous_util.cc", + "common_runtime/scoped_allocator.cc", + "common_runtime/scoped_allocator_mgr.cc", "common_runtime/session.cc", "common_runtime/session_factory.cc", "common_runtime/session_options.cc", @@ -2106,6 +2207,7 @@ tf_cuda_library( ] + if_mkl( [ "//third_party/mkl:intel_binary_blob", + "@mkl_dnn", ], ), alwayslink = 1, @@ -2150,14 +2252,12 @@ tf_cuda_library( ] + if_mkl( [ "//third_party/mkl:intel_binary_blob", - "@mkl_dnn//:mkl_dnn", + "@mkl_dnn", ], ) + tf_additional_core_deps() + if_static([":core_cpu_impl"]), alwayslink = 1, ) -# This library is deprecated and no longer publicly available. -# Do not add more uses of it. cc_library( name = "regexp_internal", hdrs = [ @@ -2227,12 +2327,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", @@ -2251,7 +2364,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", @@ -2266,6 +2378,7 @@ tf_cuda_library( ":core_cpu_lib", ":framework", ":framework_internal", + ":gpu_id", ":gpu_init_impl", ":gpu_lib", ":graph", @@ -2855,6 +2968,7 @@ tf_cc_tests_gpu( linkstatic = tf_kernel_tests_linkstatic(), deps = [ ":gpu_headers_lib", + ":gpu_id", ":gpu_runtime", ":test", ], @@ -2866,7 +2980,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", ], @@ -2878,6 +2992,7 @@ tf_cc_tests_gpu( ":direct_session", ":framework", ":framework_internal", + ":gpu_id", ":gpu_runtime", ":lib", ":lib_internal", @@ -3045,6 +3160,7 @@ tf_cc_test( ":core_cpu", ":core_cpu_internal", ":framework", + ":lib", ":test", ":test_main", ":testlib", @@ -3093,6 +3209,7 @@ tf_cc_test( "//tensorflow/core/kernels:dense_update_ops", "//tensorflow/core/kernels:fifo_queue_op", "//tensorflow/core/kernels:function_ops", + "//tensorflow/core/kernels:identity_n_op", "//tensorflow/core/kernels:identity_op", "//tensorflow/core/kernels:matmul_op", "//tensorflow/core/kernels:ops_util", @@ -3135,6 +3252,7 @@ tf_cc_test( "//tensorflow/core/kernels:fifo_queue_op", "//tensorflow/core/kernels:function_ops", "//tensorflow/core/kernels:identity_op", + "//tensorflow/core/kernels:identity_n_op", "//tensorflow/core/kernels:matmul_op", "//tensorflow/core/kernels:ops_util", "//tensorflow/core/kernels:queue_ops", @@ -3209,6 +3327,43 @@ tf_cc_test( size = "small", srcs = ["common_runtime/function_test.cc"], linkstatic = tf_kernel_tests_linkstatic(), + tags = [ + "manual", + "no_oss", + ], + deps = [ + ":core", + ":core_cpu", + ":core_cpu_internal", + ":direct_session_internal", + ":framework", + ":framework_internal", + ":lib", + ":lib_internal", + ":ops", + ":protos_all_cc", + ":test", + ":test_main", + ":testlib", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:cc_ops_internal", + "//tensorflow/cc:function_ops", + "//tensorflow/cc:functional_ops", + "//tensorflow/core/kernels:cast_op", + "//tensorflow/core/kernels:cwise_op", + "//tensorflow/core/kernels:function_ops", + "//tensorflow/core/kernels:matmul_op", + "//tensorflow/core/kernels:random_ops", + "//tensorflow/core/kernels:shape_ops", + "//third_party/eigen3", + ], +) + +tf_cc_test( + name = "common_runtime_function_threadpool_test", + size = "small", + srcs = ["common_runtime/function_threadpool_test.cc"], + linkstatic = tf_kernel_tests_linkstatic(), deps = [ ":core", ":core_cpu", @@ -3237,6 +3392,21 @@ tf_cc_test( ], ) +tf_cc_test( + name = "common_runtime_scoped_allocator_mgr_test", + size = "small", + srcs = ["common_runtime/scoped_allocator_mgr_test.cc"], + linkstatic = tf_kernel_tests_linkstatic(), + deps = [ + ":core_cpu", + ":core_cpu_internal", + ":framework", + ":lib", + ":test", + ":test_main", + ], +) + tf_cc_test_gpu( name = "gpu_allocator_retry_test", size = "medium", @@ -3273,6 +3443,7 @@ tf_cc_test_gpu( ":direct_session", ":framework", ":framework_internal", + ":gpu_id", ":gpu_runtime", ":lib", ":lib_internal", @@ -3433,6 +3604,7 @@ tf_cc_tests( "ops/parsing_ops_test.cc", "ops/random_ops_test.cc", "ops/set_ops_test.cc", + "ops/shape_function_test.cc", "ops/sparse_ops_test.cc", "ops/spectral_ops_test.cc", "ops/state_ops_test.cc", @@ -3539,6 +3711,13 @@ filegroup( "lib/gif/testdata/optimized.gif", # BMP data "lib/bmp/testdata/lena.bmp", + # SSIM, PSNR data + "lib/ssim/testdata/checkerboard1.png", + "lib/ssim/testdata/checkerboard2.png", + "lib/ssim/testdata/checkerboard3.png", + "lib/psnr/testdata/cat_q20.jpg", + "lib/psnr/testdata/cat_q72.jpg", + "lib/psnr/testdata/cat_q95.jpg", ], visibility = ["//visibility:public"], ) @@ -3591,6 +3770,18 @@ filegroup( visibility = ["//tensorflow:__subpackages__"], ) +alias( + name = "android_srcs_no_runtime", + actual = ":mobile_srcs_no_runtime", + visibility = ["//visibility:public"], +) + +alias( + name = "android_srcs_only_runtime", + actual = ":mobile_srcs_only_runtime", + visibility = ["//visibility:public"], +) + alias( name = "android_srcs", actual = ":mobile_srcs", diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD index 81187ff6b772633105e0962d9da8f87d6cfd9558..58dbac4e8edac7079d315fbfcdafbd136793df0b 100644 --- a/tensorflow/core/api_def/BUILD +++ b/tensorflow/core/api_def/BUILD @@ -96,6 +96,7 @@ tf_cc_test( srcs = ["api_test.cc"], data = [ ":base_api_def", + ":python_api_def", ], deps = [ ":excluded_ops_lib", diff --git a/tensorflow/core/api_def/api_test.cc b/tensorflow/core/api_def/api_test.cc index 112c55ccc3ba1262b48c1b6c0890b3ae22744383..477a0b670e49f8aa4ee8c250d4957886eb865ed5 100644 --- a/tensorflow/core/api_def/api_test.cc +++ b/tensorflow/core/api_def/api_test.cc @@ -41,8 +41,9 @@ namespace tensorflow { namespace { constexpr char kDefaultApiDefDir[] = "tensorflow/core/api_def/base_api"; +constexpr char kPythonApiDefDir[] = + "tensorflow/core/api_def/python_api"; constexpr char kApiDefFilePattern[] = "api_def_*.pbtxt"; -} // namespace // Reads golden ApiDef files and returns a map from file name to ApiDef file // contents. @@ -66,9 +67,93 @@ void GetGoldenApiDefs(Env* env, const string& api_files_dir, } } -class ApiTest : public ::testing::Test { +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?"; + } +} + +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(); + } + } +} + +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; + } + 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; + } + } + 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(); + } + } +} + +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(); + } + } +} +} // namespace + +class BaseApiTest : public ::testing::Test { protected: - ApiTest() { + BaseApiTest() { OpRegistry::Global()->Export(false, &ops_); const std::vector multi_line_fields = {"description"}; @@ -80,7 +165,7 @@ class ApiTest : public ::testing::Test { }; // Check that all ops have an ApiDef. -TEST_F(ApiTest, AllOpsAreInApiDef) { +TEST_F(BaseApiTest, AllOpsAreInApiDef) { auto* excluded_ops = GetExcludedOps(); for (const auto& op : ops_.op()) { if (excluded_ops->find(op.name()) != excluded_ops->end()) { @@ -94,16 +179,8 @@ TEST_F(ApiTest, AllOpsAreInApiDef) { } // Check that ApiDefs have a corresponding op. -TEST_F(ApiTest, AllApiDefsHaveCorrespondingOp) { - 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?"; - } +TEST_F(BaseApiTest, AllApiDefsHaveCorrespondingOp) { + TestAllApiDefsHaveCorrespondingOp(ops_, api_defs_map_); } string GetOpDefHasDocStringError(const string& op_name) { @@ -117,7 +194,7 @@ string GetOpDefHasDocStringError(const string& op_name) { // Check that OpDef's do not have descriptions and summaries. // Descriptions and summaries must be in corresponding ApiDefs. -TEST_F(ApiTest, OpDefsShouldNotHaveDocs) { +TEST_F(BaseApiTest, OpDefsShouldNotHaveDocs) { auto* excluded_ops = GetExcludedOps(); for (const auto& op : ops_.op()) { if (excluded_ops->find(op.name()) != excluded_ops->end()) { @@ -143,62 +220,56 @@ TEST_F(ApiTest, OpDefsShouldNotHaveDocs) { // Checks that input arg names in an ApiDef match input // arg names in corresponding OpDef. -TEST_F(ApiTest, AllApiDefInputArgsAreValid) { - for (const auto& op : ops_.op()) { - const auto& api_def = api_defs_map_[op.name()]; - 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(); - } - } +TEST_F(BaseApiTest, AllApiDefInputArgsAreValid) { + TestAllApiDefInputArgsAreValid(ops_, api_defs_map_); } // Checks that output arg names in an ApiDef match output // arg names in corresponding OpDef. -TEST_F(ApiTest, AllApiDefOutputArgsAreValid) { - for (const auto& op : ops_.op()) { - const auto& api_def = api_defs_map_[op.name()]; - 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; - } - } - 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(); - } - } +TEST_F(BaseApiTest, AllApiDefOutputArgsAreValid) { + TestAllApiDefOutputArgsAreValid(ops_, api_defs_map_); } // Checks that attribute names in an ApiDef match attribute // names in corresponding OpDef. -TEST_F(ApiTest, AllApiDefAttributeNamesAreValid) { - for (const auto& op : ops_.op()) { - const auto& api_def = api_defs_map_[op.name()]; - 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(); - } +TEST_F(BaseApiTest, AllApiDefAttributeNamesAreValid) { + TestAllApiDefAttributeNamesAreValid(ops_, api_defs_map_); +} + +class PythonApiTest : public ::testing::Test { + protected: + PythonApiTest() { + OpRegistry::Global()->Export(false, &ops_); + const std::vector multi_line_fields = {"description"}; + + Env* env = Env::Default(); + GetGoldenApiDefs(env, kPythonApiDefDir, &api_defs_map_); } + OpList ops_; + std::unordered_map api_defs_map_; +}; + +// Check that ApiDefs have a corresponding op. +TEST_F(PythonApiTest, AllApiDefsHaveCorrespondingOp) { + TestAllApiDefsHaveCorrespondingOp(ops_, api_defs_map_); } + +// Checks that input arg names in an ApiDef match input +// arg names in corresponding OpDef. +TEST_F(PythonApiTest, AllApiDefInputArgsAreValid) { + TestAllApiDefInputArgsAreValid(ops_, api_defs_map_); +} + +// Checks that output arg names in an ApiDef match output +// arg names in corresponding OpDef. +TEST_F(PythonApiTest, AllApiDefOutputArgsAreValid) { + TestAllApiDefOutputArgsAreValid(ops_, api_defs_map_); +} + +// 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: < [1, 2, 4, 7, 8] +idx ==> [0, 0, 1, 2, 2, 2, 3, 4, 4] +count ==> [2, 1, 3, 1, 2] +``` + +For an `2-D` tensor `x` with `axis = 0`: + +``` +# tensor 'x' is [[1, 0, 0], +# [1, 0, 0], +# [2, 0, 0]] +y, idx, count = unique_with_counts(x, axis=0) +y ==> [[1, 0, 0], + [2, 0, 0]] +idx ==> [0, 0, 1] +count ==> [2, 1] +``` + +For an `2-D` tensor `x` with `axis = 1`: + +``` +# tensor 'x' is [[1, 0, 0], +# [1, 0, 0], +# [2, 0, 0]] +y, idx, count = unique_with_counts(x, axis=1) +y ==> [[1, 0], + [1, 0], + [2, 0]] +idx ==> [0, 1, 1] +count ==> [1, 2] +``` +END +} diff --git a/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMax.pbtxt b/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMax.pbtxt index 8298d62f253160847ee34bcdea5a81c7370e5124..4ca6780c95629de06db319db228f440219989793 100644 --- a/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMax.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMax.pbtxt @@ -14,20 +14,21 @@ Has same shape as data, except for dimension 0 which has size `num_segments`. END } - summary: "Computes the Max along segments of a tensor." + summary: "Computes the maximum along segments of a tensor." description: <::min()`. +If the maximum is empty for a given segment ID `i`, it outputs the smallest +possible value for the specific numeric type, +`output[i] = numeric_limits::lowest()`.

diff --git a/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMin.pbtxt b/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMin.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..55ea69b5dd5f7fda5c877ca5771ec2cbb86e3a9a --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentMin.pbtxt @@ -0,0 +1,33 @@ +op { + graph_op_name: "UnsortedSegmentMin" + in_arg { + name: "segment_ids" + description: <::max()`. +END +} diff --git a/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentProd.pbtxt b/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentProd.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..577ff53d60c5a174b4ba43a667885a6983b2dfb9 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_UnsortedSegmentProd.pbtxt @@ -0,0 +1,32 @@ +op { + graph_op_name: "UnsortedSegmentProd" + in_arg { + name: "segment_ids" + description: <
  • **Organization**: Google
    • diff --git a/tensorflow/docs_src/api_guides/python/contrib.distributions.md b/tensorflow/docs_src/api_guides/python/contrib.distributions.md index 7a3d509b75198461430195aa70a336f94b7f8cfa..533d7dac1373f61ca92dba288a7d29e07e0f37d3 100644 --- a/tensorflow/docs_src/api_guides/python/contrib.distributions.md +++ b/tensorflow/docs_src/api_guides/python/contrib.distributions.md @@ -17,7 +17,6 @@ initialized with parameters that define the distributions. * @{tf.contrib.distributions.Binomial} * @{tf.contrib.distributions.Bernoulli} -* @{tf.contrib.distributions.BernoulliWithSigmoidProbs} * @{tf.contrib.distributions.Beta} * @{tf.contrib.distributions.Categorical} * @{tf.contrib.distributions.Chi2} diff --git a/tensorflow/docs_src/api_guides/python/regression_examples.md b/tensorflow/docs_src/api_guides/python/regression_examples.md index dae50a8f032bae9421bc01d1ac4043fdaae30080..7de2be05521d9293e33664cdbbd7bf16b9ad7c52 100644 --- a/tensorflow/docs_src/api_guides/python/regression_examples.md +++ b/tensorflow/docs_src/api_guides/python/regression_examples.md @@ -38,7 +38,7 @@ The preceding examples rely on the following data set utility: Utility Description -     imports85.py + imports85.py This program provides utility functions that load the imports85 data set into formats that other TensorFlow programs (for example, linear_regression.py and diff --git a/tensorflow/docs_src/community/index.md b/tensorflow/docs_src/community/index.md index 8e67022648d4c7161b02072446371e6d7e7168e2..b706d9b2047a4ff9707772edb30bfd036bbffc24 100644 --- a/tensorflow/docs_src/community/index.md +++ b/tensorflow/docs_src/community/index.md @@ -5,6 +5,7 @@ This section contains the following documents: * @{$welcome$Welcome to the TensorFlow Community}, which explains how you can get involved, where to report issues, and where to join like-minded TensorFlow enthusiasts online. + * @{$roadmap$Roadmap}, which summarizes upcoming additions to TensorFlow. * @{$documentation$Writing TensorFlow Documentation}, which explains TensorFlow's documentation conventions. If you are modifying TensorFlow source code or documentation, please read this guide. diff --git a/tensorflow/docs_src/community/leftnav_files b/tensorflow/docs_src/community/leftnav_files index c1595d3c955bb87120fe6a6c9185c58e9db1097e..fab35024ad63e09adba1298eab52f7904eca1007 100644 --- a/tensorflow/docs_src/community/leftnav_files +++ b/tensorflow/docs_src/community/leftnav_files @@ -1,5 +1,6 @@ index.md welcome.md +roadmap.md documentation.md style_guide.md benchmarks.md diff --git a/tensorflow/docs_src/community/roadmap.md b/tensorflow/docs_src/community/roadmap.md new file mode 100644 index 0000000000000000000000000000000000000000..a3170a10f2d12ed272ee1d32da679f25916994c6 --- /dev/null +++ b/tensorflow/docs_src/community/roadmap.md @@ -0,0 +1,85 @@ +# Roadmap +**Last updated: Feb 15, 2018** + +TensorFlow is a rapidly moving, community supported project. This document is intended +to provide guidance about priorities and focus areas of the core set of TensorFlow +developers and about functionality that can be expected in the upcoming releases of +TensorFlow. Many of these areas are driven by community use cases, and we welcome +further +[contributions](https://github.com/tensorflow/tensorflow/blob/master/CONTRIBUTING.md) +to TensorFlow. + +The features below do not have concrete release dates. However, the majority can be +expected in the next one to two releases. + +### APIs +#### High Level APIs: +* Easy multi-GPU utilization with Estimators +* Easy-to-use high-level pre-made estimators for Gradient Boosted Trees, Time Series, and other models + +#### Eager Execution: +* Efficient utilization of multiple GPUs +* Distributed training (multi-machine) +* Performance improvements +* Simpler export to a GraphDef/SavedModel + +#### Keras API: +* Better integration with tf.data (ability to call `model.fit` with data tensors) +* Full support for Eager Execution (both Eager support for the regular Keras API, and ability +to create Keras models Eager- style via Model subclassing) +* Better distribution/multi-GPU support and TPU support (including a smoother model-to-estimator workflow) + +#### Official Models: +* A set of +[reference models](https://github.com/tensorflow/models/tree/master/official) +across image recognition, speech, object detection, and + translation that demonstrate best practices and serve as a starting point for + high-performance model development. + +#### Contrib: +* Deprecation notices added to parts of tf.contrib where preferred implementations exist outside of tf.contrib. +* As much as possible, large projects inside tf.contrib moved to separate repositories. +* The tf.contrib module will eventually be discontinued in its current form, experimental development will in future happen in other repositories. + + +#### Probabilistic Reasoning and Statistical Analysis: +* Rich set of tools for probabilistic and statistical analysis in tf.distributions + and tf.probability. These include new samplers, layers, optimizers, losses, and structured models +* Statistical tools for hypothesis testing, convergence diagnostics, and sample statistics +* Edward 2.0: High-level API for probabilistic programming + +### Platforms +#### TensorFlow Lite: +* Increased coverage of supported ops in TensorFlow Lite +* Easier conversion of a trained TensorFlow graph for use on TensorFlow Lite +* Support for GPU acceleration in TensorFlow Lite (iOS and Android) +* Support for hardware accelerators via Android NeuralNets API +* Improved CPU performance by quantization and other network optimizations (eg. pruning, distillation) +* Increased support for devices beyond Android and iOS (eg. RPi, Cortex-M) + +### Performance +#### Distributed TensorFlow: +* Multi-GPU support optimized for a variety of GPU topologies +* Improved mechanisms for distributing computations on several machines + +#### Optimizations: +* Mixed precision training support with initial example model and guide +* Native TensorRT support +* Int8 support for SkyLake via MKL +* Dynamic loading of SIMD-optimized kernels + +### Documentation and Usability: +* Updated documentation, tutorials and Getting Started guides +* Process to enable external contributions to tutorials, documentation, and blogs showcasing best practice use-cases of TensorFlow and high-impact applications + +### Community and Partner Engagement +#### Special Interest Groups: +* Mobilizing the community to work together in focused domains +* [tf-distribute](https://groups.google.com/a/tensorflow.org/forum/#!forum/tf-distribute): build and packaging of TensorFlow +* More to be identified and launched + +#### Community: +* Incorporate public feedback on significant design decisions via a Request-for-Comment (RFC) process +* Formalize process for external contributions to land in TensorFlow and associated projects +* Grow global TensorFlow communities and user groups +* Collaborate with partners to co-develop and publish research papers diff --git a/tensorflow/docs_src/community/welcome.md b/tensorflow/docs_src/community/welcome.md index d2d3f9edaed9fc3c921a98c95ae24ce168e00216..6d0458e678b5507fc722e2c3848e84ca2168e1e3 100644 --- a/tensorflow/docs_src/community/welcome.md +++ b/tensorflow/docs_src/community/welcome.md @@ -51,6 +51,8 @@ Europe: TensorFlow provides multiple communication paths. To pick the right path, please read the following list carefully: + * For new release announcements and security updates, subscribe to + [announce@tensorflow.org](https://groups.google.com/a/tensorflow.org/forum/#!forum/announce). * To ask or answer technical questions about TensorFlow, use [Stack Overflow](https://stackoverflow.com/questions/tagged/tensorflow). For example, ask or search Stack Overflow about a particular error message diff --git a/tensorflow/docs_src/deploy/index.md b/tensorflow/docs_src/deploy/index.md index 5831960b4f6e383a6babb0823893a5d9ec5017f0..07b1bc9257ff7b132d22ac186a2f462e9c784867 100644 --- a/tensorflow/docs_src/deploy/index.md +++ b/tensorflow/docs_src/deploy/index.md @@ -7,6 +7,8 @@ the following documents: a cluster of TensorFlow servers. * @{$hadoop$How to run TensorFlow on Hadoop}, which has a highly self-explanatory title. + * @{$s3$How to run TensorFlow with the S3 filesystem}, which explains how + to run TensorFlow with the S3 file system. * The entire document set for [TensorFlow serving](/serving), an open-source, flexible, high-performance serving system for machine-learned models designed for production environments. TensorFlow Serving provides diff --git a/tensorflow/docs_src/deploy/leftnav_files b/tensorflow/docs_src/deploy/leftnav_files index f8f8d578e602cac8dd814326e318ebe0e85ec700..c682e7add16c741279aedb40c1b12f4ca8f0286a 100644 --- a/tensorflow/docs_src/deploy/leftnav_files +++ b/tensorflow/docs_src/deploy/leftnav_files @@ -1,3 +1,4 @@ index.md distributed.md hadoop.md +s3.md diff --git a/tensorflow/docs_src/deploy/s3.md b/tensorflow/docs_src/deploy/s3.md new file mode 100644 index 0000000000000000000000000000000000000000..38f84286347622d1de0646cdc621d5fb1447e588 --- /dev/null +++ b/tensorflow/docs_src/deploy/s3.md @@ -0,0 +1,40 @@ +# How to run TensorFlow on S3 + +This document describes how to run TensorFlow on S3 file system. + +## S3 + +We assume that you are familiar with @{$reading_data$reading data}. + +To use S3 with TensorFlow, change the file paths you use to read and write +data to an S3 path. For example: + +```python +filenames = ["s3://bucketname/path/to/file1.tfrecord", + "s3://bucketname/path/to/file2.tfrecord"] +dataset = tf.data.TFRecordDataset(filenames) +``` + +When reading or writing data on S3 with your TensorFlow program, the behavior +could be controlled by various environmental variables: + +* **AWS_REGION**: By default, regional endpoint is used for S3, with region + controlled by `AWS_REGION`. If `AWS_REGION` is not specified, then + `us-east-1` is used. +* **S3_ENDPOINT**: The endpoint could be overridden explicitly with + `S3_ENDPOINT` specified. +* **S3_USE_HTTPS**: HTTPS is used to access S3 by default, unless + `S3_USE_HTTPS=0`. +* **S3_VERIFY_SSL**: If HTTPS is used, SSL verification could be disabled + with `S3_VERIFY_SSL=0`. + +To read or write objects in a bucket that is no publicly accessible, +AWS credentials must be provided through one of the following methods: + +* Set credentials in the AWS credentials profile file on the local system, + located at: `~/.aws/credentials` on Linux, macOS, or Unix, or + `C:\Users\USERNAME\.aws\credentials` on Windows. +* Set the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment + variables. +* If TensorFlow is deployed on an EC2 instance, specify an IAM role and then + give the EC2 instance access to that role. diff --git a/tensorflow/docs_src/extend/add_filesys.md b/tensorflow/docs_src/extend/add_filesys.md index f0591b7b7d8af478db067ecd3bdd949e75d813c9..06f11de4eb0ea7878b01cd37d994c5a40ec400be 100644 --- a/tensorflow/docs_src/extend/add_filesys.md +++ b/tensorflow/docs_src/extend/add_filesys.md @@ -81,6 +81,8 @@ filesystem implementations call their existing libraries. Examples include: plugin](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/platform/hadoop/hadoop_file_system.h) * [GCS plugin](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/platform/cloud/gcs_file_system.h) +* [S3 + plugin](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/platform/s3/s3_file_system.h) #### The File interfaces diff --git a/tensorflow/docs_src/get_started/checkpoints.md b/tensorflow/docs_src/get_started/checkpoints.md index dfa2110e691167f54e6ea8b7a832f0a88d0ec41a..4aa07c7f2a0b56aa6de6f42e30c364c348753a39 100644 --- a/tensorflow/docs_src/get_started/checkpoints.md +++ b/tensorflow/docs_src/get_started/checkpoints.md @@ -154,7 +154,7 @@ classifier = tf.estimator.DNNClassifier( The first time you call an Estimator's `train` method, TensorFlow saves a checkpoint to the `model_dir`. Each subsequent call to the Estimator's -`train`, `eval`, or `predict` method causes the following: +`train`, `evaluate`, or `predict` method causes the following: 1. The Estimator builds the model's [graph](https://developers.google.com/machine-learning/glossary/#graph) @@ -222,7 +222,7 @@ does not match the shape stored in checkpoint: [20] To run experiments in which you train and compare slightly different versions of a model, save a copy of the code that created each -`model-dir`, possibly by creating a separate git branch for each version. +`model_dir`, possibly by creating a separate git branch for each version. This separation will keep your checkpoints recoverable. ## Summary diff --git a/tensorflow/docs_src/get_started/custom_estimators.md b/tensorflow/docs_src/get_started/custom_estimators.md index 79c4ee75d01c745d9e492c5db9df11a93eca0477..941c3e16905a9062b3081ad0af6bcbc1621a146b 100644 --- a/tensorflow/docs_src/get_started/custom_estimators.md +++ b/tensorflow/docs_src/get_started/custom_estimators.md @@ -161,12 +161,12 @@ classifier = tf.estimator.Estimator( To implement a typical model function, you must do the following: -* (Define the model)[#define_the_model]. +* [Define the model](#define_the_model). * Specify additional calculations for each of the [three different modes](#modes): - * [Predict](#predict) - * [Evaluate](#evaluate) - * [Train](#train) + * [Predict](#predict) + * [Evaluate](#evaluate) + * [Train](#train) ## Define the model @@ -213,7 +213,7 @@ is connected to every node in the preceding layer. Here's the relevant code: ``` * The `units` parameter defines the number of output neurons in a given layer. -* The `activation` parameter defines the [activation function](https://developers.google.com/machine-learning/glossary/#a) — +* The `activation` parameter defines the [activation function](https://developers.google.com/machine-learning/glossary/#activation_function) — [Relu](https://developers.google.com/machine-learning/glossary/#ReLU) in this case. @@ -546,8 +546,8 @@ In brief, here's what the three graphs tell you: * accuracy: The accuracy is recorded by the following two lines: - * `eval_metric_ops={'my_accuracy': accuracy})`, during evaluation. - * `tf.summary.scalar('accuracy', accuracy[1])`, during training. + * `eval_metric_ops={'my_accuracy': accuracy})`, during evaluation. + * `tf.summary.scalar('accuracy', accuracy[1])`, during training. These tensorboard graphs are one of the main reasons it's important to pass a `global_step` to your optimizer's `minimize` method. The model can't record diff --git a/tensorflow/docs_src/get_started/datasets_quickstart.md b/tensorflow/docs_src/get_started/datasets_quickstart.md index a8a2ab6e56130c7805d48477301c63d88f87489c..c972e5e555eea1fab5a67fdecf13264897785519 100644 --- a/tensorflow/docs_src/get_started/datasets_quickstart.md +++ b/tensorflow/docs_src/get_started/datasets_quickstart.md @@ -28,8 +28,8 @@ def train_input_fn(features, labels, batch_size): # Shuffle, repeat, and batch the examples. dataset = dataset.shuffle(1000).repeat().batch(batch_size) - # Build the Iterator, and return the read end of the pipeline. - return dataset.make_one_shot_iterator().get_next() + # Return the dataset. + return dataset ``` Let's look at this more closely. @@ -40,7 +40,7 @@ This function expects three arguments. Arguments expecting an "array" can accept nearly anything that can be converted to an array with `numpy.array`. One exception is [`tuple`](https://docs.python.org/3/tutorial/datastructures.html#tuples-and-sequences) -which has special meaning for `Datasets`. +which, as we will see, has special meaning for `Datasets`. * `features`: A `{'feature_name':array}` dictionary (or [`DataFrame`](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.html)) @@ -73,11 +73,12 @@ Let's walk through the `train_input_fn()`. ### Slices -In the simplest cases, @{tf.data.Dataset.from_tensor_slices} function takes an -array and returns a @{tf.data.Dataset} representing slices of the array. For -example, an array containing the @{$tutorials/layers$mnist training data} -has a shape of `(60000, 28, 28)`. Passing this to `from_tensor_slices` returns -a `Dataset` object containing 60000 slices, each one a 28x28 image. +The function starts by using the @{tf.data.Dataset.from_tensor_slices} function +to create a @{tf.data.Dataset} representing slices of the array. The array is +sliced across the first dimension. For example, an array containing the +@{$tutorials/layers$mnist training data} has a shape of `(60000, 28, 28)`. +Passing this to `from_tensor_slices` returns a `Dataset` object containing +60000 slices, each one a 28x28 image. The code that returns this `Dataset` is as follows: @@ -89,18 +90,24 @@ mnist_ds = tf.data.Dataset.from_tensor_slices(mnist_x) print(mnist_ds) ``` -This will print the following line, showing the @{$programmers_guide/tensors#shapes$shapes} and @{$programmers_guide/tensors#data_types$types} of the items in -the dataset. Note that the dataset does not know how many items it contains. +This will print the following line, showing the +@{$programmers_guide/tensors#shapes$shapes} and +@{$programmers_guide/tensors#data_types$types} of the items in +the dataset. Note that a `Dataset` does not know how many items it contains. ``` None ``` -The dataset above represents a collection of simple arrays, but datasets are -much more powerful than this. Datasets transparently handle any nested -combination of dictionaries or tuples. For example, ensuring that `features` -is a standard dictionary, you can then convert the dictionary of arrays to -a `Dataset` of dictionaries as follows: +The `Dataset` above represents a simple collection of arrays, but datasets are +much more powerful than this. A `Dataset` can transparently handle any nested +combination of dictionaries or tuples (or +[`namedtuple`](https://docs.python.org/2/library/collections.html#collections.namedtuple) +). + +For example after converting the iris `features` +to a standard python dictionary, you can then convert the dictionary of arrays +to a `Dataset` of dictionaries as follows: ``` python dataset = tf.data.Dataset.from_tensor_slices(dict(features)) @@ -124,9 +131,9 @@ and `types` of the `Dataset` take on the same structure. This dataset contains dictionaries of @{$programmers_guide/tensors#rank$scalars}, all of type `tf.float64`. -The first line of `train_input_fn` uses the same functionality, but adds -another level of structure. It creates a dataset containing -`(features, labels)` pairs. +The first line of the iris `train_input_fn` uses the same functionality, but +adds another level of structure. It creates a dataset containing +`(features_dict, label)` pairs. The following code shows that the label is a scalar with type `int64`: @@ -164,14 +171,14 @@ dataset = dataset.shuffle(1000).repeat().batch(batch_size) ``` The @{tf.data.Dataset.shuffle$`shuffle`} method uses a fixed-size buffer to -shuffle the items as they pass through. Setting a `buffer_size` greater than -the number of examples in the `Dataset` ensures that the data is completely -shuffled. The Iris data set only contains 150 examples. +shuffle the items as they pass through. In this case the `buffer_size` is +greater than the number of examples in the `Dataset`, ensuring that the data is +completely shuffled (The Iris data set only contains 150 examples). -The @{tf.data.Dataset.repeat$`repeat`} method has the `Dataset` restart when +The @{tf.data.Dataset.repeat$`repeat`} method restarts the `Dataset` when it reaches the end. To limit the number of epochs, set the `count` argument. -The @{tf.data.Dataset.repeat$`batch`} method collects a number of examples and +The @{tf.data.Dataset.batch$`batch`} method collects a number of examples and stacks them, to create batches. This adds a dimension to their shape. The new dimension is added as the first dimension. The following code uses the `batch` method on the MNIST `Dataset`, from earlier. This results in a @@ -213,35 +220,16 @@ print(dataset) ### Return - +At this point the `Dataset` contains `(features_dict, labels)` pairs. +This is the format expected by the `train` and `evaluate` methods, so the +`input_fn` returns the dataset. -The `train`, `evaluate`, and `predict` methods of every Estimator require -input functions to return a `(features, label)` pair containing -@{$programmers_guide/tensors$tensorflow tensors}. The `train_input_fn` uses -the following line to convert the Dataset into the expected format: +The `labels` can/should be omitted when using the `predict` method. -```python -# Build the Iterator, and return the read end of the pipeline. -features_result, labels_result = dataset.make_one_shot_iterator().get_next() -``` + -The result is a structure of @{$programmers_guide/tensors$TensorFlow tensors}, -matching the layout of the items in the `Dataset`. -For an introduction to what these objects are and how to work with them, -see @{$programmers_guide/low_level_intro}. - -``` python -print((features_result, labels_result)) -``` - -```None -({ - 'SepalLength': , - 'PetalWidth': , - 'PetalLength': , - 'SepalWidth': }, -Tensor("IteratorGetNext_1:4", shape=(?,), dtype=int64)) -``` ## Reading a CSV File @@ -277,9 +265,6 @@ ds = tf.data.TextLineDataset(train_path).skip(1) ### Build a csv line parser -Ultimately we will need to parse each of the lines in the dataset, to -produce the necessary `(features, label)` pairs. - We will start by building a function to parse a single line. The following `iris_data.parse_line` function accomplishes this task using the diff --git a/tensorflow/docs_src/get_started/feature_columns.md b/tensorflow/docs_src/get_started/feature_columns.md index ad3e1fe3e3a4e3f5278e76bcaa0fc8eee2faf374..d8e4bec86357aabd2065be50d1197122c407c9d7 100644 --- a/tensorflow/docs_src/get_started/feature_columns.md +++ b/tensorflow/docs_src/get_started/feature_columns.md @@ -146,10 +146,10 @@ single input number into a four-element vector. Therefore, the model now can learn _four individual weights_ rather than just one; four weights creates a richer model than one weight. More importantly, bucketizing enables the model to clearly distinguish between different year categories since only one of the -elements is set (1) and the other three elements are cleared (0). When we just -use a single number (a year) as input, the model can only learn a linear -relationship. So, bucketing provides the model with additional flexibility that -the model can use to learn. +elements is set (1) and the other three elements are cleared (0). For example, +when we just use a single number (a year) as input, a linear model can only +learn a linear relationship. So, bucketing provides the model with additional +flexibility that the model can use to learn. The following code demonstrates how to create a bucketized feature: @@ -242,7 +242,7 @@ on an explicit vocabulary list. For example: # the elements in the vocabulary list. vocabulary_feature_column = tf.feature_column.categorical_column_with_vocabulary_list( - key="a feature returned by input_fn()", + key=feature_name_from_input_fn, vocabulary_list=["kitchenware", "electronics", "sports"]) ``` @@ -259,7 +259,7 @@ you place the vocabulary words in a separate file. For example: # the elements in the vocabulary file vocabulary_feature_column = tf.feature_column.categorical_column_with_vocabulary_file( - key="a feature returned by input_fn()", + key=feature_name_from_input_fn, vocabulary_file="product_class.txt", vocabulary_size=3) ``` diff --git a/tensorflow/docs_src/get_started/get_started_for_beginners.md b/tensorflow/docs_src/get_started/get_started_for_beginners.md index ea1c2fb3f473b9e39567c7607d3b3ad10d2de6b5..b88483be699630d2275850cbc7c461eeb90f5943 100644 --- a/tensorflow/docs_src/get_started/get_started_for_beginners.md +++ b/tensorflow/docs_src/get_started/get_started_for_beginners.md @@ -36,6 +36,7 @@ the following three: alt="Petal geometry compared for three iris species: Iris setosa, Iris virginica, and Iris versicolor" src="../images/iris_three_species.jpg">
+ **From left to right, [*Iris setosa*](https://commons.wikimedia.org/w/index.php?curid=170298) (by [Radomil](https://commons.wikimedia.org/wiki/User:Radomil), CC BY-SA 3.0), @@ -90,11 +91,10 @@ a number. Here's the representation scheme: A **model** is the relationship between features and the label. For the Iris problem, the model defines the relationship -between the sepal and petal measurements and the Iris species. -Some simple models can be described with a few lines of algebra; -more complex machine learning models -contain such a large number of interlacing mathematical functions and -parameters that they become hard to summarize mathematically. +between the sepal and petal measurements and the predicted Iris species. Some +simple models can be described with a few lines of algebra, but complex machine +learning models have a large number of parameters that are difficult to +summarize. Could you determine the relationship between the four features and the Iris species *without* using machine learning? That is, could you use @@ -188,6 +188,7 @@ provides a programming stack consisting of multiple API layers:
+ **The TensorFlow Programming Environment.**

 

@@ -331,7 +332,7 @@ interpret data is such a rich topic that we devote an entire From a code perspective, you build a list of `feature_column` objects by calling functions from the @{tf.feature_column} module. Each object describes an input to the model. To tell the model to interpret data as a floating-point value, -call @{tf.feature_column.numeric_column). In `premade_estimator.py`, all +call @{tf.feature_column.numeric_column}. In `premade_estimator.py`, all four features should be interpreted as literal floating-point values, so the code to create a feature column looks as follows: @@ -357,7 +358,7 @@ my_feature_columns = [ ### Select the type of model -We need the select the kind of model that will be trained. +We need to select the kind of model that will be trained. Lots of model types exist; picking the ideal type takes experience. We've selected a neural network to solve the Iris problem. [**Neural networks**](https://developers.google.com/machine-learning/glossary/#neural_network) @@ -380,6 +381,7 @@ fully connected neural network consisting of three hidden layers:
+ **A neural network with three hidden layers.**

 

@@ -568,6 +570,7 @@ of 0.5. The following suggests a more effective model: 5.5 2.5 4.0 1.3 1 1 + **A model that is 80% accurate.**

 

@@ -655,7 +658,9 @@ calls as follows: ```python predictions = classifier.predict( - input_fn=lambda:eval_input_fn(predict_x, batch_size=args.batch_size)) + input_fn=lambda:eval_input_fn(predict_x, + labels=None, + batch_size=args.batch_size)) ``` As with the `evaluate` method, our `predict` method also gathers examples @@ -700,7 +705,7 @@ for pred_dict, expec in zip(predictions, expected): class_id = pred_dict['class_ids'][0] probability = pred_dict['probabilities'][class_id] - print(template.format(SPECIES[class_id], 100 * probability, expec)) + print(template.format(iris_data.SPECIES[class_id], 100 * probability, expec)) ``` Running the program yields the following output: diff --git a/tensorflow/docs_src/get_started/premade_estimators.md b/tensorflow/docs_src/get_started/premade_estimators.md index 4f01f997c33c211e8cff81b6b268bb320aa794df..e50d2f542037c8537f79a2ae53a2cbb3448243c6 100644 --- a/tensorflow/docs_src/get_started/premade_estimators.md +++ b/tensorflow/docs_src/get_started/premade_estimators.md @@ -98,6 +98,7 @@ classifies Iris flowers into three different species based on the size of their alt="Petal geometry compared for three iris species: Iris setosa, Iris virginica, and Iris versicolor" src="../images/iris_three_species.jpg"> + **From left to right, [*Iris setosa*](https://commons.wikimedia.org/w/index.php?curid=170298) (by [Radomil](https://commons.wikimedia.org/wiki/User:Radomil), CC BY-SA 3.0), @@ -396,9 +397,9 @@ predictions and their probabilities: ``` python -for pred_dict, expec in zip(predictions, expected): - template = ('\nPrediction is "{}" ({:.1f}%), expected "{}"') +template = ('\nPrediction is "{}" ({:.1f}%), expected "{}"') +for pred_dict, expec in zip(predictions, expected): class_id = pred_dict['class_ids'][0] probability = pred_dict['probabilities'][class_id] diff --git a/tensorflow/docs_src/install/index.md b/tensorflow/docs_src/install/index.md index 3c8488643f071c147dfbc4e0b4b4760b0a817718..4f85383925bbb8a03372b020e448a0e604f3b999 100644 --- a/tensorflow/docs_src/install/index.md +++ b/tensorflow/docs_src/install/index.md @@ -3,7 +3,7 @@ We've built and tested TensorFlow on the following 64-bit laptop/desktop operating systems: - * MacOS X 10.11 (El Capitan) or later. + * macOS 10.12.6 (Sierra) or later. * Ubuntu 16.04 or later * Windows 7 or later. diff --git a/tensorflow/docs_src/install/install_c.md b/tensorflow/docs_src/install/install_c.md index 14add7c77e61534b011a95007da4a3918b26dc93..0481c97885df97675553eccb08989846a1275948 100644 --- a/tensorflow/docs_src/install/install_c.md +++ b/tensorflow/docs_src/install/install_c.md @@ -15,7 +15,7 @@ instructions might also work on other variants, we have only tested following requirements: * Linux, 64-bit, x86 - * macOS X, Version 10.11 (El Capitan) or higher + * macOS X, Version 10.12.6 (Sierra) or higher ## Installation @@ -38,7 +38,7 @@ enable TensorFlow for C: OS="linux" # Change to "darwin" for macOS TARGET_DIRECTORY="/usr/local" curl -L \ - "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.5.0.tar.gz" | + "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.6.0.tar.gz" | sudo tar -C $TARGET_DIRECTORY -xz The `tar` command extracts the TensorFlow C library into the `lib` diff --git a/tensorflow/docs_src/install/install_go.md b/tensorflow/docs_src/install/install_go.md index d2af9d98434a60ef0c82e56db602cdaad8f5fa37..8f89898c92d00eaf0d637e78b3ec897a34df6b89 100644 --- a/tensorflow/docs_src/install/install_go.md +++ b/tensorflow/docs_src/install/install_go.md @@ -17,7 +17,7 @@ instructions might also work on other variants, we have only tested following requirements: * Linux, 64-bit, x86 - * macOS X, 10.11 (El Capitan) or higher + * macOS X, 10.12.6 (Sierra) or higher ## Installation @@ -38,7 +38,7 @@ steps to install this library and enable TensorFlow for Go: TF_TYPE="cpu" # Change to "gpu" for GPU support TARGET_DIRECTORY='/usr/local' curl -L \ - "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.5.0.tar.gz" | + "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.6.0.tar.gz" | sudo tar -C $TARGET_DIRECTORY -xz The `tar` command extracts the TensorFlow C library into the `lib` diff --git a/tensorflow/docs_src/install/install_java.md b/tensorflow/docs_src/install/install_java.md index e5388c4b1e0202f468c0f6d4caf065c3c58a0312..0ee9c849e11448c4efc74263f6aacbb46056201d 100644 --- a/tensorflow/docs_src/install/install_java.md +++ b/tensorflow/docs_src/install/install_java.md @@ -18,7 +18,7 @@ instructions might also work on other variants, we have only tested following requirements: * Ubuntu 16.04 or higher; 64-bit, x86 - * macOS X 10.11 (El Capitan) or higher + * macOS 10.12.6 (Sierra) or higher * Windows 7 or higher; 64-bit, x86 The installation instructions for Android are in a separate @@ -36,7 +36,7 @@ following to the project's `pom.xml` to use the TensorFlow Java APIs: org.tensorflow tensorflow - 1.5.0 + 1.6.0 ``` @@ -65,7 +65,7 @@ As an example, these steps will create a Maven project that uses TensorFlow: org.tensorflow tensorflow - 1.5.0 + 1.6.0 @@ -123,12 +123,12 @@ instead: org.tensorflow libtensorflow - 1.5.0 + 1.6.0 org.tensorflow libtensorflow_jni_gpu - 1.5.0 + 1.6.0 ``` @@ -147,7 +147,7 @@ refer to the simpler instructions above instead. Take the following steps to install TensorFlow for Java on Linux or macOS: 1. Download - [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.5.0.jar), + [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.6.0.jar), which is the TensorFlow Java Archive (JAR). 2. Decide whether you will run TensorFlow for Java on CPU(s) only or with @@ -166,7 +166,7 @@ Take the following steps to install TensorFlow for Java on Linux or macOS: OS=$(uname -s | tr '[:upper:]' '[:lower:]') mkdir -p ./jni curl -L \ - "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.5.0.tar.gz" | + "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.6.0.tar.gz" | tar -xz -C ./jni ### Install on Windows @@ -174,10 +174,10 @@ Take the following steps to install TensorFlow for Java on Linux or macOS: Take the following steps to install TensorFlow for Java on Windows: 1. Download - [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.5.0.jar), + [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.6.0.jar), which is the TensorFlow Java Archive (JAR). 2. Download the following Java Native Interface (JNI) file appropriate for - [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.5.0.zip). + [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.6.0.zip). 3. Extract this .zip file. @@ -225,7 +225,7 @@ must be part of your `classpath`. For example, you can include the downloaded `.jar` in your `classpath` by using the `-cp` compilation flag as follows: -
javac -cp libtensorflow-1.5.0.jar HelloTF.java
+
javac -cp libtensorflow-1.6.0.jar HelloTF.java
### Running @@ -239,11 +239,11 @@ two files are available to the JVM: For example, the following command line executes the `HelloTF` program on Linux and macOS X: -
java -cp libtensorflow-1.5.0.jar:. -Djava.library.path=./jni HelloTF
+
java -cp libtensorflow-1.6.0.jar:. -Djava.library.path=./jni HelloTF
And the following command line executes the `HelloTF` program on Windows: -
java -cp libtensorflow-1.5.0.jar;. -Djava.library.path=jni HelloTF
+
java -cp libtensorflow-1.6.0.jar;. -Djava.library.path=jni HelloTF
If the program prints Hello from version, you've successfully installed TensorFlow for Java and are ready to use the API. If the program diff --git a/tensorflow/docs_src/install/install_linux.md b/tensorflow/docs_src/install/install_linux.md index cd8c14599fb587098d13782dc76e5a451807b56a..3e8744bf9d1e01d00c703d9153db8ac48bce1d50 100644 --- a/tensorflow/docs_src/install/install_linux.md +++ b/tensorflow/docs_src/install/install_linux.md @@ -41,7 +41,8 @@ must be installed on your system: [NVIDIA's documentation](https://developer.nvidia.com/cudnn). Ensure that you create the `CUDA_HOME` environment variable as described in the NVIDIA documentation. - * GPU card with CUDA Compute Capability 3.0 or higher. See + * GPU card with CUDA Compute Capability 3.0 or higher for building + from source and 3.5 or higher for our binaries. See [NVIDIA documentation](https://developer.nvidia.com/cuda-gpus) for a list of supported GPU cards. * The libcupti-dev library, which is the NVIDIA CUDA Profile Tools Interface. @@ -188,7 +189,7 @@ Take the following steps to install TensorFlow with Virtualenv: Virtualenv environment: 
(tensorflow)$ pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp34-cp34m-linux_x86_64.whl
+ https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp34-cp34m-linux_x86_64.whl If you encounter installation problems, see [Common Installation Problems](#common_installation_problems). @@ -293,7 +294,7 @@ take the following steps: 
      $ sudo pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp34-cp34m-linux_x86_64.whl
+     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp34-cp34m-linux_x86_64.whl
If this step fails, see @@ -356,24 +357,23 @@ where: to 6006. * TensorFlowCPUImage is required. It identifies the Docker container. Specify one of the following values: - * gcr.io/tensorflow/tensorflow, which is the TensorFlow CPU binary image. - * gcr.io/tensorflow/tensorflow:latest-devel, which is the latest + * tensorflow/tensorflow, which is the TensorFlow CPU binary image. + * tensorflow/tensorflow:latest-devel, which is the latest TensorFlow CPU Binary image plus source code. - * gcr.io/tensorflow/tensorflow:version, which is the + * tensorflow/tensorflow:version, which is the specified version (for example, 1.1.0rc1) of TensorFlow CPU binary image. - * gcr.io/tensorflow/tensorflow:version-devel, which is + * tensorflow/tensorflow:version-devel, which is the specified version (for example, 1.1.0rc1) of the TensorFlow GPU binary image plus source code. - gcr.io is the Google Container Registry. Note that some - TensorFlow images are also available at + TensorFlow images are available at [dockerhub](https://hub.docker.com/r/tensorflow/tensorflow/). For example, the following command launches the latest TensorFlow CPU binary image in a Docker container from which you can run TensorFlow programs in a shell: 
-$ docker run -it gcr.io/tensorflow/tensorflow bash
+$ docker run -it tensorflow/tensorflow bash
The following command also launches the latest TensorFlow CPU binary image in a @@ -381,7 +381,7 @@ Docker container. However, in this Docker container, you can run TensorFlow programs in a Jupyter notebook: 
-$ docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow
+$ docker run -it -p 8888:8888 tensorflow/tensorflow
Docker will download the TensorFlow binary image the first time you launch it. @@ -405,14 +405,14 @@ where: hostPort and containerPort to `8888`. * TensorFlowGPUImage specifies the Docker container. You must specify one of the following values: - * gcr.io/tensorflow/tensorflow:latest-gpu, which is the latest + * tensorflow/tensorflow:latest-gpu, which is the latest TensorFlow GPU binary image. - * gcr.io/tensorflow/tensorflow:latest-devel-gpu, which is + * tensorflow/tensorflow:latest-devel-gpu, which is the latest TensorFlow GPU Binary image plus source code. - * gcr.io/tensorflow/tensorflow:version-gpu, which is the + * tensorflow/tensorflow:version-gpu, which is the specified version (for example, 0.12.1) of the TensorFlow GPU binary image. - * gcr.io/tensorflow/tensorflow:version-devel-gpu, which is + * tensorflow/tensorflow:version-devel-gpu, which is the specified version (for example, 0.12.1) of the TensorFlow GPU binary image plus source code. @@ -421,7 +421,7 @@ following command launches the latest TensorFlow GPU binary image in a Docker container from which you can run TensorFlow programs in a shell: 
-$ nvidia-docker run -it gcr.io/tensorflow/tensorflow:latest-gpu bash
+$ nvidia-docker run -it tensorflow/tensorflow:latest-gpu bash
The following command also launches the latest TensorFlow GPU binary image @@ -429,13 +429,13 @@ in a Docker container. In this Docker container, you can run TensorFlow programs in a Jupyter notebook: 
-$ nvidia-docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow:latest-gpu
+$ nvidia-docker run -it -p 8888:8888 tensorflow/tensorflow:latest-gpu
The following command installs an older TensorFlow version (0.12.1): 
-$ nvidia-docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow:0.12.1-gpu
+$ nvidia-docker run -it -p 8888:8888 tensorflow/tensorflow:0.12.1-gpu
Docker will download the TensorFlow binary image the first time you launch it. @@ -480,8 +480,7 @@ Take the following steps to install TensorFlow in an Anaconda environment: 
      (tensorflow)$ pip install --ignore-installed --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp34-cp34m-linux_x86_64.whl
- + https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp34-cp34m-linux_x86_64.whl ## Validate your installation @@ -506,7 +505,7 @@ If you installed through Docker, start a Docker container from which you can run bash. For example: 
-$ docker run -it gcr.io/tensorflow/tensorflow bash
+$ docker run -it tensorflow/tensorflow bash
@@ -648,14 +647,14 @@ This section documents the relevant values for Linux installations. CPU only: 
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp27-none-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp27-none-linux_x86_64.whl
GPU support: 
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0-cp27-none-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.6.0-cp27-none-linux_x86_64.whl
Note that GPU support requires the NVIDIA hardware and software described in @@ -667,14 +666,14 @@ Note that GPU support requires the NVIDIA hardware and software described in CPU only: 
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp34-cp34m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp34-cp34m-linux_x86_64.whl
GPU support: 
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0-cp34-cp34m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.6.0-cp34-cp34m-linux_x86_64.whl
Note that GPU support requires the NVIDIA hardware and software described in @@ -686,14 +685,14 @@ Note that GPU support requires the NVIDIA hardware and software described in CPU only: 
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp35-cp35m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp35-cp35m-linux_x86_64.whl
GPU support: 
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0-cp35-cp35m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.6.0-cp35-cp35m-linux_x86_64.whl
@@ -705,14 +704,14 @@ Note that GPU support requires the NVIDIA hardware and software described in CPU only: 
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp36-cp36m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.6.0-cp36-cp36m-linux_x86_64.whl
GPU support: 
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0-cp36-cp36m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.6.0-cp36-cp36m-linux_x86_64.whl
diff --git a/tensorflow/docs_src/install/install_mac.md b/tensorflow/docs_src/install/install_mac.md index f49d3a2f08559780152bdb9016ece3d72a364ffb..205db8e6bdd438ee5502f6cd1d92c30ea9e33152 100644 --- a/tensorflow/docs_src/install/install_mac.md +++ b/tensorflow/docs_src/install/install_mac.md @@ -5,7 +5,11 @@ instructions might also work on other macOS variants, we have only tested (and we only support) these instructions on machines meeting the following requirements: - * macOS X 10.11 (El Capitan) or higher + * macOS 10.12.6 (Sierra) or higher + +Note: There are known, accuracy-affecting numerical issues before macOS 10.12.6 +(Sierra) that are described in +[GitHub#15933](https://github.com/tensorflow/tensorflow/issues/15933#issuecomment-366331383). Note: As of version 1.2, TensorFlow no longer provides GPU support on macOS. @@ -114,8 +118,8 @@ Take the following steps to install TensorFlow with Virtualenv: Python 2.7, the command to install TensorFlow in the active Virtualenv is as follows: - 
 $ pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0-py3-none-any.whl
+ 
 $ pip install --upgrade \
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.6.0-py2-none-any.whl
If you encounter installation problems, see [Common Installation Problems](#common-installation-problems). @@ -237,8 +241,8 @@ take the following steps: you are installing TensorFlow for Mac OS and Python 2.7 issue the following command: - 
 $ sudo pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0-py3-none-any.whl
+ 
 $ sudo pip install --upgrade \
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.6.0-py2-none-any.whl
If the preceding command fails, see [installation problems](#common-installation-problems). @@ -288,24 +292,23 @@ where: to 6006. * TensorFlowImage is required. It identifies the Docker container. You must specify one of the following values: - * gcr.io/tensorflow/tensorflow: TensorFlow binary image. - * gcr.io/tensorflow/tensorflow:latest-devel: TensorFlow + * tensorflow/tensorflow: TensorFlow binary image. + * tensorflow/tensorflow:latest-devel: TensorFlow Binary image plus source code. -gcr.io is the Google Container Registry. Note that some -TensorFlow images are also available at +The TensorFlow images are available at [dockerhub](https://hub.docker.com/r/tensorflow/tensorflow/). For example, the following command launches a TensorFlow CPU binary image in a Docker container from which you can run TensorFlow programs in a shell: -
$ docker run -it gcr.io/tensorflow/tensorflow bash
+
$ docker run -it tensorflow/tensorflow bash
The following command also launches a TensorFlow CPU binary image in a Docker container. However, in this Docker container, you can run TensorFlow programs in a Jupyter notebook: -
$ docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow
+
$ docker run -it -p 8888:8888 tensorflow/tensorflow
Docker will download the TensorFlow binary image the first time you launch it. @@ -347,7 +350,7 @@ Take the following steps to install TensorFlow in an Anaconda environment: TensorFlow for Python 2.7: 
 (targetDirectory)$ pip install --ignore-installed --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0-py2-none-any.whl
+ https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.6.0-py2-none-any.whl @@ -372,7 +375,7 @@ do the following: If you installed through Docker, start a Docker container that runs bash. For example: -
$ docker run -it gcr.io/tensorflow/tensorflow bash
+
$ docker run -it tensorflow/tensorflow bash
@@ -520,7 +523,7 @@ This section documents the relevant values for Mac OS installations. 
-https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0-py2-none-any.whl
+https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.6.0-py2-none-any.whl
@@ -528,5 +531,5 @@ https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0-py2-none-any. 
-https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0-py3-none-any.whl
+https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.6.0-py3-none-any.whl
diff --git a/tensorflow/docs_src/install/install_sources.md b/tensorflow/docs_src/install/install_sources.md index c1e4545d207f656673fa37bbeaee92c236b29bb4..c09c9c2c0c6eaabe9066a2630a3aec39f05fcfec 100644 --- a/tensorflow/docs_src/install/install_sources.md +++ b/tensorflow/docs_src/install/install_sources.md @@ -221,7 +221,7 @@ problem, do either of the following: * Download Xcode 7.2 and select it as your default by issuing the following command: - 
 $ sudo xcode-select -s /Application/Xcode-7.2/Xcode.app
+ 
 $ sudo xcode-select -s /Applications/Xcode-7.2/Xcode.app
**NOTE:** Your system must fulfill the NVIDIA software requirements described in one of the following documents: @@ -272,8 +272,6 @@ Found possible Python library paths: Please input the desired Python library path to use. Default is [/usr/lib/python2.7/dist-packages] Using python library path: /usr/local/lib/python2.7/dist-packages -Do you wish to build TensorFlow with MKL support? [y/N] -No MKL support will be enabled for TensorFlow Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native]: Do you wish to use jemalloc as the malloc implementation? [Y/n] jemalloc enabled @@ -361,10 +359,10 @@ Invoke `pip install` to install that pip package. The filename of the `.whl` file depends on your platform. For example, the following command will install the pip package -for TensorFlow 1.5.0 on Linux: +for TensorFlow 1.6.0 on Linux: 
-$ sudo pip install /tmp/tensorflow_pkg/tensorflow-1.5.0-py2-none-any.whl
+$ sudo pip install /tmp/tensorflow_pkg/tensorflow-1.6.0-py2-none-any.whl
## Validate your installation @@ -395,8 +393,7 @@ TensorFlow programs: 
Hello, TensorFlow!
-If you are new to TensorFlow, see @{$get_started/get_started$Getting Started with -TensorFlow}. +If you are new to TensorFlow, see @{$get_started/premade_estimators$Getting Started with TensorFlow}. If the system outputs an error message instead of a greeting, see [Common installation problems](#common_installation_problems). @@ -462,7 +459,8 @@ Stack Overflow and specify the `tensorflow` tag. **Linux** - + + @@ -480,6 +478,7 @@ Stack Overflow and specify the `tensorflow` tag. **Mac**
Version:CPU/GPU:Python Version:Compiler:Build Tools:cuDNN:CUDA:
tensorflow-1.6.0CPU2.7, 3.3-3.6GCC 4.8Bazel 0.9.0N/AN/A
tensorflow_gpu-1.6.0GPU2.7, 3.3-3.6GCC 4.8Bazel 0.9.079
tensorflow-1.5.0CPU2.7, 3.3-3.6GCC 4.8Bazel 0.8.0N/AN/A
tensorflow_gpu-1.5.0GPU2.7, 3.3-3.6GCC 4.8Bazel 0.8.079
tensorflow-1.4.0CPU2.7, 3.3-3.6GCC 4.8Bazel 0.5.4N/AN/A
+ @@ -493,6 +492,8 @@ Stack Overflow and specify the `tensorflow` tag. **Windows**
Version:CPU/GPU:Python Version:Compiler:Build Tools:cuDNN:CUDA:
tensorflow-1.6.0CPU2.7, 3.3-3.6Clang from xcodeBazel 0.8.1N/AN/A
tensorflow-1.5.0CPU2.7, 3.3-3.6Clang from xcodeBazel 0.8.1N/AN/A
tensorflow-1.4.0CPU2.7, 3.3-3.6Clang from xcodeBazel 0.5.4N/AN/A
tensorflow-1.3.0CPU2.7, 3.3-3.6Clang from xcodeBazel 0.4.5N/AN/A
+ + diff --git a/tensorflow/docs_src/install/install_windows.md b/tensorflow/docs_src/install/install_windows.md index 86a111c2ec1ee1e126b3e96c83eae68da7da7e41..2413bc9cfbbfd577ebd583be4da82994e8551c9e 100644 --- a/tensorflow/docs_src/install/install_windows.md +++ b/tensorflow/docs_src/install/install_windows.md @@ -17,7 +17,7 @@ You must choose one of the following types of TensorFlow to install: NVIDIA® GPU, you must install this version. Note that this version of TensorFlow is typically much easier to install (typically, in 5 or 10 minutes), so even if you have an NVIDIA GPU, we recommend - installing this version first. + installing this version first. Prebuilt binaries will use AVX instructions. * **TensorFlow with GPU support**. TensorFlow programs typically run significantly faster on a GPU than on a CPU. Therefore, if your system has a NVIDIA® GPU meeting the prerequisites shown below @@ -41,13 +41,14 @@ installed on your system: Note that cuDNN is typically installed in a different location from the other CUDA DLLs. Ensure that you add the directory where you installed the cuDNN DLL to your `%PATH%` environment variable. - * GPU card with CUDA Compute Capability 3.0 or higher. See + * GPU card with CUDA Compute Capability 3.0 or higher for building + from source and 3.5 or higher for our binaries. See [NVIDIA documentation](https://developer.nvidia.com/cuda-gpus) for a list of supported GPU cards. If you have a different version of one of the preceding packages, please change to the specified versions. In particular, the cuDNN version -must match exactly: TensorFlow will not load if it cannot find `cuDNN64_6.dll`. +must match exactly: TensorFlow will not load if it cannot find `cuDNN64_7.dll`. To use a different version of cuDNN, you must build from source. ## Determine how to install TensorFlow @@ -153,8 +154,7 @@ TensorFlow programs: 
Hello, TensorFlow!
-If you are new to TensorFlow, see @{$get_started/get_started$Getting Started with -TensorFlow}. +If you are new to TensorFlow, see @{$get_started/premade_estimators$Getting Started with TensorFlow}. If the system outputs an error message instead of a greeting, see [Common installation problems](#common_installation_problems). diff --git a/tensorflow/docs_src/install/leftnav_files b/tensorflow/docs_src/install/leftnav_files index 0e8b5ae7a17eb43cffc76d40692c4f0042de44af..e523e06f67aad508238ee0965f34ebe16c77bf90 100644 --- a/tensorflow/docs_src/install/leftnav_files +++ b/tensorflow/docs_src/install/leftnav_files @@ -1,16 +1,16 @@ index.md ### Python -install_linux.md -install_mac.md -install_windows.md -install_sources.md +install_linux.md: Ubuntu +install_mac.md: MacOS +install_windows.md: Windows +install_sources.md: From source >>> migration.md ### Other Languages -install_java.md -install_go.md -install_c.md +install_java.md: Java +install_go.md: Go +install_c.md: C diff --git a/tensorflow/docs_src/mobile/android_build.md b/tensorflow/docs_src/mobile/android_build.md index b5a1d5d7d1bf3b456ab24165e273969bdbd7bfca..08a5fbe41c87c88399682208c38bf7a892d8fc1a 100644 --- a/tensorflow/docs_src/mobile/android_build.md +++ b/tensorflow/docs_src/mobile/android_build.md @@ -90,8 +90,8 @@ using [ADB](https://developer.android.com/studio/command-line/adb.html). This requires some knowledge of build systems and Android developer tools, but we'll guide you through the basics here. -- First, follow our instructions for @{$install/install_sources$installing from - sources}. This will also guide you through installing Bazel and cloning the +- First, follow our instructions for @{$install/install_sources$installing from sources}. + This will also guide you through installing Bazel and cloning the TensorFlow code. - Download the Android [SDK](https://developer.android.com/studio/index.html) diff --git a/tensorflow/docs_src/mobile/leftnav_files b/tensorflow/docs_src/mobile/leftnav_files index ac50f528ba468d8a830c059539d3399f413f39c8..4cf134cc3c2c323405d769a5ced5d5a68f188203 100644 --- a/tensorflow/docs_src/mobile/leftnav_files +++ b/tensorflow/docs_src/mobile/leftnav_files @@ -2,6 +2,7 @@ index.md ### TensorFlow Lite tflite/index.md tflite/demo_android.md +tflite/demo_ios.md >>> ### TensorFlow Mobile mobile_intro.md diff --git a/tensorflow/docs_src/mobile/mobile_intro.md b/tensorflow/docs_src/mobile/mobile_intro.md index 17dbf1c3e6ad89768529864ba884274a51b3dfb2..69b63ae7d22ced9fd0299f17d1ae2d614c9a6be7 100644 --- a/tensorflow/docs_src/mobile/mobile_intro.md +++ b/tensorflow/docs_src/mobile/mobile_intro.md @@ -235,7 +235,7 @@ TensorFlow [on Github](https://github.com/tensorflow/models) that you can look through. Lean towards the simplest model you can find, and try to get started as soon as you have even a small amount of labelled data, since you’ll get the best results when you’re able to iterate quickly. The shorter the time it takes to -try training a model and running it in s real application, the better overall +try training a model and running it in its real application, the better overall results you’ll see. It’s common for an algorithm to get great training accuracy numbers but then fail to be useful within a real application because there’s a mismatch between the dataset and real usage. Prototype end-to-end usage as soon diff --git a/tensorflow/docs_src/mobile/optimizing.md b/tensorflow/docs_src/mobile/optimizing.md index 44cacff5dbbcb0685044c342184464b47a8ed090..ca9cb043e9282702044b0da26c27f5bf29141cae 100644 --- a/tensorflow/docs_src/mobile/optimizing.md +++ b/tensorflow/docs_src/mobile/optimizing.md @@ -290,8 +290,8 @@ run it on a 64-bit ARM device: You can interpret the results in exactly the same way as the desktop version above. If you have any trouble figuring out what the right input and output -names and types are, take a look at the @{$mobile/prepare_models$Preparing -models} page for details about detecting these for your model, and look at the +names and types are, take a look at the @{$mobile/prepare_models$Preparing models} +page for details about detecting these for your model, and look at the `summarize_graph` tool which may give you helpful information. diff --git a/tensorflow/docs_src/mobile/tflite/demo_android.md b/tensorflow/docs_src/mobile/tflite/demo_android.md index 79b567897cb8a38ed2e27e73aa7e8fee95f718b8..c94b5597a673a7e68aed517b325b9719b3b73bbd 100644 --- a/tensorflow/docs_src/mobile/tflite/demo_android.md +++ b/tensorflow/docs_src/mobile/tflite/demo_android.md @@ -8,6 +8,9 @@ You'll need an Android device running Android 5.0 or higher to run the demo. To get you started working with TensorFlow Lite on Android, we'll walk you through building and deploying our TensorFlow demo app in Android Studio. +Note: For a more detailed guide see the +[TFLite Codelab](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets-2-tflite/index.html#0) + It's also possible to build the demo app with Bazel, but we only recommend this for advanced users who are very familiar with the Bazel build environment. For more information on that, see our page [on Github](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite#building-tensorflow-lite-and-the-demo-app-from-source). diff --git a/tensorflow/docs_src/mobile/tflite/demo_ios.md b/tensorflow/docs_src/mobile/tflite/demo_ios.md new file mode 100644 index 0000000000000000000000000000000000000000..3ee9b1cbca6cfef98616bd33bbf91b756b4efa15 --- /dev/null +++ b/tensorflow/docs_src/mobile/tflite/demo_ios.md @@ -0,0 +1,68 @@ +# TensorFlow Lite Demo for iOS + +The TensorFlow Lite demo is a camera app that continuously classifies whatever +it sees from your device's back camera, using a quantized MobileNet model. These +instructions walk you through building and running the demo on an iOS device. + +## Prerequisites + +* You must have [Xcode](https://developer.apple.com/xcode/) installed and have a + valid Apple Developer ID, and have an iOS device set up and linked to your + developer account with all of the appropriate certificates. For these + instructions, we assume that you have already been able to build and deploy an + app to an iOS device with your current developer environment. + +* The demo app requires a camera and must be executed on a real iOS device. You + can build it and run with the iPhone Simulator but it won't have any camera + information to classify. + +* You don't need to build the entire TensorFlow library to run the demo, but you + will need to clone the TensorFlow repository if you haven't already: + + git clone https://github.com/tensorflow/tensorflow + +* You'll also need the Xcode command-line tools: + + xcode-select --install + + If this is a new install, you will need to run the Xcode application once to + agree to the license before continuing. + +## Building the iOS Demo App + +1. Install CocoaPods if you don't have it: + + sudo gem install cocoapods + +2. Download the model files used by the demo app (this is done from inside the + cloned directory): + + sh tensorflow/contrib/lite/examples/ios/download_models.sh + +3. Install the pod to generate the workspace file: + + cd tensorflow/contrib/lite/examples/ios/camera + pod install + + If you have installed this pod before and that command doesn't work, try + + pod update + + At the end of this step you should have a file called + `tflite_camera_example.xcworkspace`. + +4. Open the project in Xcode by typing this on the command line: + + open tflite_camera_example.xcworkspace + + This launches Xcode if it isn't open already and opens the + `tflite_camera_example` project. + +5. Build and run the app in Xcode. + + Note that as mentioned earlier, you must already have a device set up and + linked to your Apple Developer account in order to deploy the app on a + device. + +You'll have to grant permissions for the app to use the device's camera. Point +the camera at various objects and enjoy seeing how the model classifies things! diff --git a/tensorflow/docs_src/performance/datasets_performance.md b/tensorflow/docs_src/performance/datasets_performance.md index 4f95e17c3598c23645fad07441c267266e5ef34e..46b43b7673c561679e89fff0ae738b0e751fcff5 100644 --- a/tensorflow/docs_src/performance/datasets_performance.md +++ b/tensorflow/docs_src/performance/datasets_performance.md @@ -92,11 +92,11 @@ transform the data. Without pipelining, the CPU and the GPU/TPU sit idle much of the time: -![without pipelining](https://www.tensorflow.org/images/datasets_without_pipelining.png) +![without pipelining](/images/datasets_without_pipelining.png) With pipelining, idle time diminishes significantly: -![with pipelining](https://www.tensorflow.org/images/datasets_with_pipelining.png) +![with pipelining](/images/datasets_with_pipelining.png) The `tf.data` API provides a software pipelining mechanism through the @{tf.data.Dataset.prefetch} transformation, which can be used to decouple the @@ -139,7 +139,7 @@ multiple CPU cores. To make this possible, the `map` transformation provides the the following diagram illustrates the effect of setting `num_parallel_calls=2` to the `map` transformation: -![parallel map](https://www.tensorflow.org/images/datasets_parallel_map.png) +![parallel map](/images/datasets_parallel_map.png) Choosing the best value for the `num_parallel_calls` argument depends on your hardware, characteristics of your training data (such as its size and shape), @@ -213,7 +213,7 @@ number of datasets to overlap can be specified by the `cycle_length` argument. The following diagram illustrates the effect of supplying `cycle_length=2` to the `parallel_interleave` transformation: -![parallel io](https://www.tensorflow.org/images/datasets_parallel_io.png) +![parallel io](/images/datasets_parallel_io.png) To apply this change to our running example, change: diff --git a/tensorflow/docs_src/performance/leftnav_files b/tensorflow/docs_src/performance/leftnav_files index 316f023f43dcfe781c7819d1681335267ddd5f76..d11a7e5d07c3e6cfa092e7ac11189ce6c272c1ad 100644 --- a/tensorflow/docs_src/performance/leftnav_files +++ b/tensorflow/docs_src/performance/leftnav_files @@ -2,6 +2,7 @@ performance_guide.md datasets_performance.md performance_models.md benchmarks.md +quantization.md ### XLA xla/index.md @@ -11,6 +12,3 @@ xla/jit.md xla/operation_semantics.md xla/shapes.md xla/tfcompile.md - -### Quantization -quantization.md diff --git a/tensorflow/docs_src/performance/performance_guide.md b/tensorflow/docs_src/performance/performance_guide.md index 10e7ad7ada533c8da5e5b871b38809b90604685e..580a899ac4e4f5c3d97ce023f25083168fe00d01 100644 --- a/tensorflow/docs_src/performance/performance_guide.md +++ b/tensorflow/docs_src/performance/performance_guide.md @@ -78,7 +78,7 @@ training CIFAR-10 illustrates the use of the `tf.data` API along with The `tf.data` API utilizes C++ multi-threading and has a much lower overhead than the Python-based `queue_runner` that is limited by Python's multi-threading performance. A detailed performance guide for the `tf.data` API can be found -[here](#datasets_performance). +[here](@{$datasets_performance}). While feeding data using a `feed_dict` offers a high level of flexibility, in general `feed_dict` does not provide a scalable solution. If only a single GPU @@ -498,7 +498,7 @@ For TensorFlow source versions after 1.3.0: ```bash ./configure # Pick the desired options -bazel build --config=mkl -c opt //tensorflow/tools/pip_package:build_pip_package +bazel build --config=mkl --config=opt //tensorflow/tools/pip_package:build_pip_package ``` diff --git a/tensorflow/docs_src/performance/quantization.md b/tensorflow/docs_src/performance/quantization.md index 544274cab68934419e8601a4d9714d80335fca28..411889cb1c616130f809e6228cc692ba3f951d48 100644 --- a/tensorflow/docs_src/performance/quantization.md +++ b/tensorflow/docs_src/performance/quantization.md @@ -1,226 +1,253 @@ -# How to Quantize Neural Networks with TensorFlow - -When modern neural networks were being developed, the biggest challenge was -getting them to work at all! That meant that accuracy and speed during training -were the top priorities. Using floating point arithmetic was the easiest way to -preserve accuracy, and GPUs were well-equipped to accelerate those calculations, -so it's natural that not much attention was paid to other numerical formats. - -These days, we actually have a lot of models being deployed in commercial -applications. The computation demands of training grow with the number of -researchers, but the cycles needed for inference expand in proportion to users. -That means pure inference efficiency has become a burning issue for a lot of -teams. - -That is where quantization comes in. It's an umbrella term that covers a lot of -different techniques to store numbers and perform calculations on them in more -compact formats than 32-bit floating point. I am going to focus on eight-bit -fixed point, for reasons I'll go into more detail on later. - -[TOC] - -## Why does Quantization Work? - -Training neural networks is done by applying many tiny nudges to the weights, -and these small increments typically need floating point precision to work -(though there are research efforts to use quantized representations here too). - -Taking a pre-trained model and running inference is very different. One of the -magical qualities of deep networks is that they tend to cope very well with high -levels of noise in their inputs. If you think about recognizing an object in a -photo you've just taken, the network has to ignore all the CCD noise, lighting -changes, and other non-essential differences between it and the training -examples it's seen before, and focus on the important similarities instead. This -ability means that they seem to treat low-precision calculations as just another -source of noise, and still produce accurate results even with numerical formats -that hold less information. - -## Why Quantize? - -Neural network models can take up a lot of space on disk, with the original -AlexNet being over 200 MB in float format for example. Almost all of that size -is taken up with the weights for the neural connections, since there are often -many millions of these in a single model. Because they're all slightly different -floating point numbers, simple compression formats like zip don't compress them -well. They are arranged in large layers though, and within each layer the -weights tend to be normally distributed within a certain range, for example -3.0 -to 6.0. - -The simplest motivation for quantization is to shrink file sizes by storing the -min and max for each layer, and then compressing each float value to an -eight-bit integer representing the closest real number in a linear set of 256 -within the range. For example with the -3.0 to 6.0 range, a 0 byte would -represent -3.0, a 255 would stand for 6.0, and 128 would represent about 1.5. -I'll go into the exact calculations later, since there's some subtleties, but -this means you can get the benefit of a file on disk that's shrunk by 75%, and -then convert back to float after loading so that your existing floating-point -code can work without any changes. - -Another reason to quantize is to reduce the computational resources you need to -do the inference calculations, by running them entirely with eight-bit inputs -and outputs. This is a lot more difficult since it requires changes everywhere -you do calculations, but offers a lot of potential rewards. Fetching eight-bit -values only requires 25% of the memory bandwidth of floats, so you'll make much -better use of caches and avoid bottlenecking on RAM access. You can also -typically use SIMD operations that do many more operations per clock cycle. In -some case you'll have a DSP chip available that can accelerate eight-bit -calculations too, which can offer a lot of advantages. - -Moving calculations over to eight bit will help you run your models faster, and -use less power (which is especially important on mobile devices). It also opens -the door to a lot of embedded systems that can't run floating point code -efficiently, so it can enable a lot of applications in the IoT world. - -## Why Not Train in Lower Precision Directly? - -There have been some experiments training at lower bit depths, but the results -seem to indicate that you need higher than eight bit to handle the back -propagation and gradients. That makes implementing the training more -complicated, and so starting with inference made sense. We also already have a -lot of float models already that we use and know well, so being able to convert -them directly is very convenient. - -## How Can You Quantize Your Models? - -TensorFlow has production-grade support for eight-bit calculations built in. It -also has a process for converting many models trained in floating-point over to -equivalent graphs using quantized calculations for inference. For example, -here's how you can translate the latest GoogLeNet model into a version that uses -eight-bit computations: - -```sh -curl -L "https://storage.googleapis.com/download.tensorflow.org/models/inception_v3_2016_08_28_frozen.pb.tar.gz" | - tar -C tensorflow/examples/label_image/data -xz -bazel build tensorflow/tools/graph_transforms:transform_graph -bazel-bin/tensorflow/tools/graph_transforms/transform_graph \ - --in_graph=tensorflow/examples/label_image/data/inception_v3_2016_08_28_frozen.pb \ - --out_graph=/tmp/quantized_graph.pb \ - --inputs=input \ - --outputs=InceptionV3/Predictions/Reshape_1 \ - --transforms='add_default_attributes strip_unused_nodes(type=float, shape="1,299,299,3") - remove_nodes(op=Identity, op=CheckNumerics) fold_constants(ignore_errors=true) - fold_batch_norms fold_old_batch_norms quantize_weights quantize_nodes - strip_unused_nodes sort_by_execution_order' +# Fixed Point Quantization + +Quantization techniques store and calculate numbers in more compact formats. +[TensorFlow Lite](/mobile/tflite/) adds quantization that uses an 8-bit fixed +point representation. + +Since a challenge for modern neural networks is optimizing for high accuracy, the +priority has been improving accuracy and speed during training. Using floating +point arithmetic is an easy way to preserve accuracy and GPUs are designed to +accelerate these calculations. + +However, as more machine learning models are deployed to mobile devices, +inference efficiency has become a critical issue. Where the computational demand +for *training* grows with the amount of models trained on different +architectures, the computational demand for *inference* grows in proportion to +the amount of users. + +## Quantization benefits + + +Using 8-bit calculations help your models run faster and use less power. This is +especially important for mobile devices and embedded applications that can't run +floating point code efficiently, for example, Internet of Things (IoT) and +robotics devices. There are additional opportunities to extend this support to +more backends and research lower precision networks. + +### Smaller file sizes {: .hide-from-toc} + +Neural network models require a lot of space on disk. For example, the original +AlexNet requires over 200 MB for the float format—almost all of that for the +model's millions of weights. Because the weights are slightly different +floating point numbers, simple compression formats perform poorly (like zip). + +Weights fall in large layers of numerical values. For each layer, weights tend to +be normally distributed within a range. Quantization can shrink file sizes by +storing the minimum and maximum weight for each layer, then compress each +weight's float value to an 8-bit integer representing the closest real number in +a linear set of 256 within the range. + +### Faster inference {: .hide-from-toc} + +Since calculations are run entirely on 8-bit inputs and outputs, quantization +reduces the computational resources needed for inference calculations. This is +more involved, requiring changes to all floating point calculations, but results +in a large speed-up for inference time. + +### Memory efficiency {: .hide-from-toc} + +Since fetching 8-bit values only requires 25% of the memory bandwidth of floats, +more efficient caches avoid bottlenecks for RAM access. In many cases, the power +consumption for running a neural network is dominated by memory access. The +savings from using fixed-point 8-bit weights and activations are significant. + +Typically, SIMD operations are available that run more operations per clock +cycle. In some cases, a DSP chip is available that accelerates 8-bit calculations +resulting in a massive speedup. + +## Fixed point quantization techniques + +The goal is to use the same precision for weights and activations during both +training and inference. But an important difference is that training consists of +a forward pass and a backward pass, while inference only uses a forward pass. +When we train the model with quantization in the loop, we ensure that the forward +pass matches precision for both training and inference. + +To minimize the loss in accuracy for fully fixed point models (weights and +activations), train the model with quantization in the loop. This simulates +quantization in the forward pass of a model so weights tend towards values that +perform better during quantized inference. The backward pass uses quantized +weights and activations and models quantization as a straight through estimator. +(See Bengio et al., [2013](https://arxiv.org/abs/1308.3432)) + +Additionally, the minimum and maximum values for activations are determined +during training. This allows a model trained with quantization in the loop to be +converted to a fixed point inference model with little effort, eliminating the +need for a separate calibration step. + +## Quantization training with TensorFlow + +TensorFlow can train models with quantization in the loop. Because training +requires small gradient adjustments, floating point values are still used. To +keep models as floating point while adding the quantization error in the training +loop, @{$array_ops#Fake_quantization$fake quantization} nodes simulate the +effect of quantization in the forward and backward passes. + +Since it's difficult to add these fake quantization operations to all the +required locations in the model, there's a function available that rewrites the +training graph. To create a fake quantized training graph: + +``` +# Build forward pass of model. +loss = tf.losses.get_total_loss() + +# Call the training rewrite which rewrites the graph in-place with +# FakeQuantization nodes and folds batchnorm for training. It is +# often needed to fine tune a floating point model for quantization +# with this training tool. When training from scratch, quant_delay +# can be used to activate quantization after training to converge +# with the float graph, effectively fine-tuning the model. +tf.contrib.quantize.create_training_graph(quant_delay=2000000) + +# Call backward pass optimizer as usual. +optimizer = tf.train.GradientDescentOptimizer(learning_rate) +optimizer.minimize(loss) ``` -This will produce a new model that runs the same operations as the original, but -with eight bit calculations internally, and all weights quantized as well. If -you look at the file size, you'll see it's about a quarter of the original (23MB -versus 91MB). You can still run this model using exactly the same inputs and -outputs though, and you should get equivalent results. Here's an example: +The rewritten *eval graph* is non-trivially different from the *training graph* +since the quantization ops affect the batch normalization step. Because of this, +we've added a separate rewrite for the *eval graph*: -```sh -bazel build tensorflow/examples/label_image:label_image -bazel-bin/tensorflow/examples/label_image/label_image \ ---graph=/tmp/quantized_graph.pb \ +``` +# Build eval model +logits = tf.nn.softmax_cross_entropy_with_logits(...) + +# Call the eval rewrite which rewrites the graph in-place with +# FakeQuantization nodes and fold batchnorm for eval. +tf.contrib.quantize.create_eval_graph() + +# Save the checkpoint and eval graph proto to disk for freezing +# and providing to TFLite. +with open(eval_graph_file, ‘w’) as f: + f.write(str(g.as_graph_def())) +saver = tf.train.Saver() +saver.save(sess, checkpoint_name) +``` + +Methods to rewrite the training and eval graphs are an active area of research +and experimentation. Although rewrites and quantized training might not work or +improve performance for all models, we are working to generalize these +techniques. + +## Generating fully quantized models + +The previously demonstrated after-rewrite eval graph only *simulates* +quantization. To generate real fixed point computations from a trained +quantization model, convert it to a fixed point kernel. Tensorflow Lite supports +this conversion from the graph resulting from `create_eval_graph`. + +First, create a frozen graph that will be the input for the TensorFlow Lite +toolchain: + +``` +bazel build tensorflow/python/tools:freeze_graph && \ + bazel-bin/tensorflow/python/tools/freeze_graph \ + --input_graph=eval_graph_def.pb \ + --input_checkpoint=checkpoint \ + --output_graph=frozen_eval_graph.pb --output_node_names=outputs ``` -You'll see that this runs the newly-quantized graph, and outputs a very similar -answer to the original. - -You can run the same process on your own models saved out as GraphDefs, with the -input and output names adapted to those your network requires. I recommend that -you run them through the freeze_graph script first, to convert checkpoints into -constants stored in the file. - -## How Does the Quantization Process Work? - -We've implemented quantization by writing equivalent eight-bit versions of -operations that are commonly used during inference. These include convolution, -matrix multiplication, activation functions, pooling operations and -concatenation. The conversion script first replaces all the individual ops it -knows about with quantized equivalents. These are small sub-graphs that have -conversion functions before and after to move the data between float and -eight-bit. Below is an example of what they look like. First here's the original -Relu operation, with float inputs and outputs: - -![Relu Diagram](https://www.tensorflow.org/images/quantization0.png) - -Then, this is the equivalent converted subgraph, still with float inputs and -outputs, but with internal conversions so the calculations are done in eight -bit. - -![Converted Diagram](https://www.tensorflow.org/images/quantization1.png) - -The min and max operations actually look at the values in the input float -tensor, and then feeds them into the Dequantize operation that converts the -tensor into eight-bits. There are more details on how the quantized representation -works later on. - -Once the individual operations have been converted, the next stage is to remove -unnecessary conversions to and from float. If there are consecutive sequences of -operations that all have float equivalents, then there will be a lot of adjacent -Dequantize/Quantize ops. This stage spots that pattern, recognizes that they -cancel each other out, and removes them, like this: - -![Stripping Diagram](https://www.tensorflow.org/images/quantization2.png) - -Applied on a large scale to models where all of the operations have quantized -equivalents, this gives a graph where all of the tensor calculations are done in -eight bit, without having to convert to float. - -## What Representation is Used for Quantized Tensors? - -We approach converting floating-point arrays of numbers into eight-bit -representations as a compression problem. We know that the weights and -activation tensors in trained neural network models tend to have values that are -distributed across comparatively small ranges (for example you might have -15 to -+15 for weights, -500 to 1000 for activations on an image model, though the -exact numbers will vary). We also know from experiment that neural nets tend to -be very robust in the face of noise, and so the noise-like error produced by -quantizing down to a small set of values will not hurt the precision of the -overall results very much. We also want to pick a representation that's easy to -perform calculations on, especially the large matrix multiplications that form -the bulk of the work that's needed to run a model. - -These led us to pick a representation that has two floats to store the overall -minimum and maximum values that are represented by the lowest and highest -quantized value. Each entry in the quantized array represents a float value in -that range, distributed linearly between the minimum and maximum. For example, -if we have minimum = -10.0, and maximum = 30.0f, and an eight-bit array, here's -what the quantized values represent: +Provide this to the TensorFlow Lite Optimizing Converter (TOCO) to get a fully +quantized TensorFLow Lite model: ``` -Quantized | Float ---------- | ----- -0 | -10.0 -255 | 30.0 -128 | 10.0 +bazel build tensorflow/contrib/lite/toco:toco && \ + ./bazel-bin/third_party/tensorflow/contrib/lite/toco/toco \ + --input_file=frozen_eval_graph.pb \ + --output_file=tflite_model.tflite \ + --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \ + --inference_type=QUANTIZED_UINT8 \ + --input_shape="1,224, 224,3" \ + --input_array=input \ + --output_array=outputs \ + --std_value=127.5 --mean_value=127.5 ``` -The advantages of this format are that it can represent arbitrary magnitudes of -ranges, they don't have to be symmetrical, it can represent signed and unsigned -values, and the linear spread makes doing multiplications straightforward. There -are alternatives like [Song Han's code books](http://arxiv.org/pdf/1510.00149.pdf) -that can use lower bit depths by non-linearly distributing the float values -across the representation, but these tend to be more expensive to calculate on. - -The advantage of having a strong and clear definition of the quantized format is -that it's always possible to convert back and forth from float for operations -that aren't quantization-ready, or to inspect the tensors for debugging -purposes. One implementation detail in TensorFlow that we're hoping to improve -in the future is that the minimum and maximum float values need to be passed as -separate tensors to the one holding the quantized values, so graphs can get a -bit dense! - -The nice thing about the minimum and maximum ranges is that they can often be -pre-calculated. Weight parameters are constants known at load time, so their -ranges can also be stored as constants. We often know the ranges for inputs (for -examples images are usually RGB values in the range 0.0 to 255.0), and many -activation functions have known ranges too. This can avoid having to analyze the -outputs of an operation to determine the range, which we need to do for math ops -like convolution or matrix multiplication which produce 32-bit accumulated -results from 8-bit inputs. - -## What's Next? - -We've found that we can get extremely good performance on mobile and embedded -devices by using eight-bit arithmetic rather than floating-point. You can see -the framework we use to optimize matrix multiplications at -[gemmlowp](https://github.com/google/gemmlowp). We still need to apply all the -lessons we've learned to the TensorFlow ops to get maximum performance on -mobile, but we're actively working on that. Right now, this quantized -implementation is a reasonably fast and accurate reference implementation that -we're hoping will enable wider support for our eight-bit models on a wider -variety of devices. We also hope that this demonstration will encourage the -community to explore what's possible with low-precision neural networks. +See the documentation for @{tf.contrib.quantize} and +[TensorFlow Lite](/mobile/tflite/). + +## Quantized accuracy + +Fixed point [MobileNet](https://arxiv.org/abs/1704.0486) models are released with +8-bit weights and activations. Using the rewriters, these models achieve the +Top-1 accuracies listed in Table 1. For comparison, the floating point accuracies +are listed for the same models. The code used to generate these models +[is available](https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet_v1.md) +along with links to all of the pretrained mobilenet_v1 models. + +
+
Version:CPU/GPU:Python Version:Compiler:Build Tools:cuDNN:CUDA:
tensorflow-1.6.0CPU3.5-3.6MSVC 2015 update 3Cmake v3.6.3N/AN/A
tensorflow_gpu-1.6.0GPU3.5-3.6MSVC 2015 update 3Cmake v3.6.379
tensorflow-1.5.0CPU3.5-3.6MSVC 2015 update 3Cmake v3.6.3N/AN/A
tensorflow_gpu-1.5.0GPU3.5-3.6MSVC 2015 update 3Cmake v3.6.379
tensorflow-1.4.0CPU3.5-3.6MSVC 2015 update 3Cmake v3.6.3N/AN/A
+ + + + + + + + + + + + + + + + + + + + + + +
Image SizeDepthTop-1 Accuracy:
Floating point		Top-1 Accuracy:
Fixed point: 8 bit weights and activations
1280.250.4150.399
1280.50.5630.549
1280.750.6210.598
12810.6520.64
1600.250.4550.435
1600.50.5910.577
1600.750.6530.639
16010.680.673
1920.250.4770.458
1920.50.6170.604
1920.750.6720.662
19210.70.69
2240.250.4980.482
2240.50.6330.622
2240.750.6840.679
22410.7090.697
+
+ Table 1: MobileNet Top-1 accuracy on Imagenet Validation dataset. +
+ + +## Representation for quantized tensors + +TensorFlow approaches the conversion of floating-point arrays of numbers into +8-bit representations as a compression problem. Since the weights and activation +tensors in trained neural network models tend to have values that are distributed +across comparatively small ranges (for example, -15 to +15 for weights or -500 to +1000 for image model activations). And since neural nets tend to be robust +handling noise, the error introduced by quantizing to a small set of values +maintains the precision of the overall results within an acceptable threshold. A +chosen representation must perform fast calculations, especially the large matrix +multiplications that comprise the bulk of the computations while running a model. + +This is represented with two floats that store the overall minimum and maximum +values corresponding to the lowest and highest quantized value. Each entry in the +quantized array represents a float value in that range, distributed linearly +between the minimum and maximum. For example, with a minimum of -10.0 and maximum +of 30.0f, and an 8-bit array, the quantized values represent the following: + +
+ + + + + +
QuantizedFloat
0-10.0
25530.0
12810.0
+
+ Table 2: Example quantized value range +
+
+ +The advantages of this representation format are: + +* It efficiently represents an arbitrary magnitude of ranges. +* The values don't have to be symmetrical. +* The format represents both signed and unsigned values. +* The linear spread makes multiplications straightforward. + +Alternative techniques use lower bit depths by non-linearly distributing the +float values across the representation, but currently are more expensive in terms +of computation time. (See Han et al., +[2016](https://arxiv.org/abs/1510.00149).) + +The advantage of having a clear definition of the quantized format is that it's +always possible to convert back and forth from fixed-point to floating-point for +operations that aren't quantization-ready, or to inspect the tensors for +debugging. diff --git a/tensorflow/docs_src/performance/xla/jit.md b/tensorflow/docs_src/performance/xla/jit.md index d4dc3e57c8fb5ec2a979b6ba7ebe2a3b6c3a5f94..d9a979ccbd31773b9d227ff946486706844a8f81 100644 --- a/tensorflow/docs_src/performance/xla/jit.md +++ b/tensorflow/docs_src/performance/xla/jit.md @@ -157,7 +157,7 @@ to fuse Ops is visible by starting at `hlo_graph_0.dot` and viewing each diagram in succession. To Render the .dot file into a png, install -[GraphViz](http://www.graphviz.org/Download..php) and run: +[GraphViz](https://www.graphviz.org/download/) and run: ```shell dot -Tpng hlo_graph_80.dot -o hlo_graph_80.png diff --git a/tensorflow/docs_src/performance/xla/operation_semantics.md b/tensorflow/docs_src/performance/xla/operation_semantics.md index 1e9b8b35db65ef19a4bcb607b98af1e1de4e6d5b..5e39e710a0dba74dfd68a04367ce402362520590 100644 --- a/tensorflow/docs_src/performance/xla/operation_semantics.md +++ b/tensorflow/docs_src/performance/xla/operation_semantics.md @@ -45,27 +45,30 @@ feature dimension in `operand`), the operation calculates the gradients with respect to `operand`, `offset` and `scale` across all the other dimensions. The `feature_index` must be a valid index for the feature dimension in `operand`. -The three gradients are defined by the following formulas: +The three gradients are defined by the following formulas (Assuming a +4-dimensional tensor as `operand` and (l) is the index for feature dimension): -\\( \nabla x = \nabla y * \gamma * \sqrt{\sigma^2+\epsilon} \\) +\\( coef_l = \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h (\nabla y_{ijkl} * (x_{ijkl} - \mu_l) / (\sigma^2_{l}+\epsilon)) \\) -\\( \nabla \gamma = sum(\nabla y * (x - \mu) * \sqrt{\sigma^2 + \epsilon}) \\) +\\( \nabla x_{ijkl} = \gamma_{l} * (1/\sqrt{\sigma^2_{l}+\epsilon}) * [\nabla y_{ijkl} - mean(\nabla y) - (x_{ijkl} - \mu_{l}) * coef_l] \\) -\\( \nabla \beta = sum(\nabla y) \\) +\\( \nabla \beta_l = \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\) + +\\( \nabla \gamma_l = \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} * ((x_{ijkl} - \mu_l) / \sqrt{\sigma^2_{l}+\epsilon}) \\) The inputs `mean` and `variance` represents moments value across batch and spatial dimensions. The output type is a tuple of three handles: -|Outputs | Type | Semantics | -|------------- | ----------------------- | ------------------------------------| -|`grad_operand`| `ComputationDataHandle` | gradient with respect to input | -: : : `operand` : -|`grad_scale` | `ComputationDataHandle` | gradient with respect to input | -: : : `scale` : -|`grad_offset` | `ComputationDataHandle` | gradient with respect to input | -: : : `offset` : +|Outputs | Type | Semantics | +|------------- | ----------------------- | ------------------------------------ | +|`grad_operand`| `ComputationDataHandle` | gradient with respect to input | +: : : `operand` (\\( \nabla x\\)) : +|`grad_scale` | `ComputationDataHandle` | gradient with respect to input | +: : : `scale` (\\( \nabla \gamma\\)) : +|`grad_offset` | `ComputationDataHandle` | gradient with respect to input | +: : : `offset`(\\( \nabla \beta\\)) : ## BatchNormInference @@ -119,11 +122,11 @@ Normalizes an array across batch and spatial dimensions. | Arguments | Type | Semantics | | --------------- | ----------------------- | -------------------------------- | | `operand` | `ComputationDataHandle` | n dimensional array to be | -: : : normalized : +: : : normalized (x) : | `scale` | `ComputationDataHandle` | 1 dimensional array | : : : (\\(\gamma\\)) : | `offset` | `ComputationDataHandle` | 1 dimensional array | -: : : (\\(\beta\\ ) : +: : : (\\(\beta\\)) : | `epsilon` | `float` | Epsilon value (\\(\epsilon\\)) | | `feature_index` | `int64` | Index to feature dimension | : : : in `operand` : @@ -135,8 +138,8 @@ element in `operand`. The `feature_index` must be a valid index for the feature dimension in `operand`. The algorithm goes as follows for each batch in `operand` \\(x\\) that -contains `m` elements with `w` and `h` as the size of spatial dimensions ( -assuming `operand` is an 4 dimensional array): +contains `m` elements with `w` and `h` as the size of spatial dimensions +(assuming `operand` is an 4 dimensional array): - Calculates batch mean \\(\mu_l\\) for each feature `l` in feature dimension: \\(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\\) @@ -170,7 +173,7 @@ Similar to a `tf.bitcast` in TensorFlow, performs an element-wise bitcast operation from a data shape to a target shape. The dimensions must match, and the conversion is an element-wise one; e.g. `s32` elements become `f32` elements via bitcast routine. Bitcast is implemented as a low-level cast, so machines -with different floating point representations will give different results. +with different floating-point representations will give different results. `BitcastConvertType(operand, new_element_type)` @@ -252,9 +255,9 @@ Clamps an operand to within the range between a minimum and maximum value. Given an operand and minimum and maximum values, returns the operand if it is in the range between the minimum and maximum, else returns the minimum value if the operand is below this range or the maximum value if the operand is above this -range. That is, `clamp(a, x, b) = max(min(a, x), b)`. +range. That is, `clamp(a, x, b) = min(max(a, x), b)`. -All three arrays must be the same shape. Alternately, as a restricted form of +All three arrays must be the same shape. Alternatively, as a restricted form of [broadcasting](broadcasting.md), `min` and/or `max` can be a scalar of type `T`. Example with scalar `min` and `max`: @@ -351,7 +354,7 @@ each other) and contains the arguments in the order that they were specified. : : : concatenated between the `operands`. : With the exception of `dimension` all dimensions must be the same. This is -because XLA does not support "ragged" arrays Also note that rank-0 values +because XLA does not support "ragged" arrays. Also note that rank-0 values cannot be concatenated (as it's impossible to name the dimension along which the concatenation occurs). @@ -440,11 +443,13 @@ area and a computation is performed for each possible position of the window. | `lhs` | `ComputationDataHandle` | rank n+2 array of inputs | | `rhs` | `ComputationDataHandle` | rank n+2 array of kernel | : : : weights : -| `window_strides` | `ArraySlice` | n-d array of kernel strides | -| `padding` | `ArraySlice` | size n array of kernel strides| +| `padding` | `ArraySlice>` : padding : -| `lhs_dilation` | `ArraySlice` | n-d lhs dilation factor array | -| `rhs_dilation` | `ArraySlice` | n-d rhs dilation factor array | +| `lhs_dilation` | `ArraySlice` | size n lhs dilation factor | +: : : array | +| `rhs_dilation` | `ArraySlice` | size n rhs dilation factor +: : : array | Let n be the number of spatial dimensions. The `lhs` argument is a rank n+2 array describing the base area. This is called the input, even though of course @@ -468,7 +473,7 @@ filter/kernel/window. The dimensions are, in this order: window that moves across the base area. The `window_strides` argument specifies the stride of the convolutional window -in the spatial dimensions. For example, if the stride in a the first spatial +in the spatial dimensions. For example, if the stride in the first spatial dimension is 3, then the window can only be placed at coordinates where the first spatial index is divisible by 3. @@ -717,6 +722,7 @@ in 'dimension_numbers'. Associated contracting dimension numbers from the 'lhs' and 'rhs' do not need to be the same, but must be listed in the same order in both 'lhs/rhs_contracting_dimensions' arrays and have the same dimension sizes. +There must be exactly one contracting dimension on both 'lhs' and 'rhs'. Example with contracting dimension numbers: @@ -736,8 +742,9 @@ DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0}, ``` Associated batch dimension numbers from the 'lhs' and 'rhs' must have the same -dimension number, must be listed in the same order in both arrays, and must -have the same dimension sizes. +dimension number, must be listed in the same order in both arrays, must +have the same dimension sizes, and must be ordered before contracting and +non-contracting/non-batch dimension numbers. Example with batch dimension numbers (batch size 2, 2x2 matrices): @@ -769,6 +776,10 @@ DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0}, | [b0, m, k] `dot` [b0, k, n] | [b0, m, n] | batch matmul | | [b0, b1, m, k] `dot` [b0, b1, k, n] | [b0, b1, m, n] | batch matmul | +It follows that the resulting dimension number starts with the batch dimension, +then the 'lhs' non-contracting/non-batch dimension, and finally the 'rhs' +non-contracting/non-batch dimension. + ## DynamicSlice See also @@ -936,7 +947,7 @@ expand the rank of the lower-rank operand up to the rank of the higher-rank operand. `broadcast_dimensions` maps the dimensions of the lower-rank shape to the dimensions of the higher-rank shape. The unmapped dimensions of the expanded shape are filled with dimensions of size one. Degenerate-dimension broadcasting -then broadcasts the shapes along these degenerate dimension to equalize the +then broadcasts the shapes along these degenerate dimensions to equalize the shapes of both operands. The semantics are described in detail on the @{$broadcasting$broadcasting page}. @@ -1021,6 +1032,213 @@ Arguments | Type | Semantics The function is applied to each element in the `operand` array, resulting in an array with the same shape. It is allowed for `operand` to be a scalar (rank 0). +## Gather + +The XLA gather operation stitches together several slices (each slice at a +potentially different runtime offset) of an input tensor into an output tensor. + +### General Semantics + +See also +[`ComputationBuilder::Gather`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/computation_builder.h). +For a more intuitive description, see the "Informal Description" section below. + + `gather(operand, gather_indices, output_window_dims, elided_window_dims, window_bounds, gather_dims_to_operand_dims)` + +|Arguments | Type | Semantics | +|----------------- | ----------------------- | --------------------------------| +|`operand` | `ComputationDataHandle` | The tensor we’re gathering | +: : : from. : +|`gather_indices` | `ComputationDataHandle` | Tensor containing the starting | +: : : indices of the slices we're : +: : : we're stitching together into : +: : : the output tensor. : +|`index_vector_dim` | `int64` | The dimension in | +: : : `gather_indices` that contains : +: : : the starting indices. : +|`output_window_dims` | `ArraySlice` | The set of dimensions in the | +: : : output shape that are _window : +: : : dimensions_ (defined below). : +: : : Not all window dimensions may : +: : : be present in the output shape. : +|`elided_window_dims` | `ArraySlice` | The set of _window dimensions_ | +: : : that are not present in the output shape. : +: : : `window_bounds[i]` must be `1` for all `i` : +: : : in `elided_window_dims`. : +|`window_bounds` | `ArraySlice` | `window_bounds[i]` is the bounds | +: : : for window dimension `i`. This includes : +: : : both the window dimensions that are : +: : : explicitly part of the output shape (via : +: : : `output_window_dims`) and the window : +: : : dimensions that are elided (via : +: : : `elided_window_dims`). : +|`gather_dims_to_operand_dims` | `ArraySlice` | A dimension map (the | +: : : array is interpreted as mapping `i` to : +: : : `gather_dims_to_operand_dims[i]`) from : +: : : the gather indices in `gather_indices` to : +: : : the operand index space. It has to be : +: : : one-to-one and total. : + +For every index `Out` in the output tensor, we compute two things (more +precisely described later): + + - An index into `gather_indices.rank` - `1` dimensions of `gather_indices`, + which gives us a starting index of a slice, _operand slice_, in the operand + tensor. These `gather_indices.rank` - `1` dimensions are all the dimensions + in `gather_indices` except `index_vector_dim`. + + - A _window index_ that has the same rank as the operand. This index is + composed of the values in `Out` at dimensions `output_window_dims`, embedded + with zeroes according to `elided_window_dims`. + +The _window index_ is the relative index of the element in _operand slice_ that +should be present in the output at index `Out`. + +The output is a tensor of rank `output_window_dims.size` + `gather_indices.rank` +- `1`. Additionally, as a shorthand, we define `output_gather_dims` of type +`ArraySlice` as the set of dimensions in the output shape but not in +`output_window_dims`, in ascending order. E.g. if the output tensor has rank +`5`, `output_window_dims` is {`2`, `4`} then `output_gather_dims` is {`0`, `1`, +`3`} + +If `index_vector_dim` is equal to `gather_indices.rank` we implicitly +consider `gather_indices` to have a trailing `1` dimension (i.e. if +`gather_indices` was of shape `[6,7]` and `index_vector_dim` is `2` then +we implicitly consider the shape of `gather_indices` to be `[6,7,1]`). + +The bounds for the output tensor along dimension `i` is computed as follows: + + 1. If `i` is present in `output_gather_dims` (i.e. is equal to + `output_gather_dims[k]` for some `k`) then we pick the corresponding + dimension bounds out of `gather_indices.shape`, skipping + `index_vector_dim` (i.e. pick `gather_indices.shape.dims`[`k`] if `k` + < `index_vector_dim` and `gather_indices.shape.dims`[`k`+`1`] + otherwise). + 2. If `i` is present in `output_window_dims` (i.e. equal to + `output_window_dims`[`k`] for some `k`) then we pick the corresponding + bound out of `window_bounds` after accounting for `elided_window_dims` + (i.e. we pick `adjusted_window_bounds`[`k`] where `adjusted_window_bounds` + is `window_bounds` with the bounds at indices `elided_window_dims` + removed). + +The operand index `In` corresponding to an output index `Out` is computed as +follows: + + 1. Let `G` = { `Out`[`k`] for `k` in `output_gather_dims` }. Use `G` to slice + out vector `S` such that `S`[`i`] = `gather_indices`[Combine(`G`, `i`)] + where Combine(A, b) inserts b at position `index_vector_dim` into A. + Note that this is well defined even if `G` is empty -- if `G` is empty then + `S` = `gather_indices`. + 2. Create an index, `S``in`, into `operand` using `S` by + scattering `S` using the `gather_dims_to_operand_dims` map + (`S``in` is the starting indices for _operand slice_ mentioned + above). More precisely: + 1. `S``in`[`gather_dims_to_operand_dims`[`k`]] = `S`[`k`] if `k` < + `gather_dims_to_operand_dims.size`. + 2. `S``in`[`_`] = `0` otherwise. + 3. Create an index `W``in` into `operand` by scattering the indices + at the output window dimensions in `Out` according to + the `elided_window_dims` set (`W``in` is the _window index_ + mentioned above). More precisely: + 1. `W``in`[`window_dims_to_operand_dims`(`k`)] = `Out`[`k`] if + `k` < `output_window_dims.size` (`window_dims_to_operand_dims` is + defined below). + 2. `W``in`[`_`] = `0` otherwise. + 4. `In` is `W``in` + `S``in` where + is element-wise + addition. + +`window_dims_to_operand_dims` is the monotonic function with domain [`0`, +`output_window_dims.size`) and range [`0`, `operand.rank`) \ +`elided_window_dims`. So if, e.g., `output_window_dims.size` is `4`, +`operand.rank` is `6` and `elided_window_dims` is {`0`, `2`} then +`window_dims_to_operand_dims` is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}. + +### Informal Description and Examples + +`index_vector_dim` is set to `gather_indices.rank` - `1` in all of the +examples that follow. More interesting values for `index_vector_dim` +does not change the operation fundamentally, but makes the visual representation +more cumbersome. + +To get an intuition on how all of the above fits together, let's look at an +example that gathers 5 slices of shape `[8,6]` from a `[16,11]` tensor. The +position of a slice into the `[16,11]` tensor can be represented as an index +vector of shape `S64[2]`, so the set of 5 positions can be represented as a +`S64[5,2]` tensor. + +The behavior of the gather operation can then be depicted as an index +transformation that takes [`G`,`W``0`,`W``1`], an index in +the output shape, and maps it to an element in the input tensor in the following +way: + +
+ +
+ +We first select an (`X`,`Y`) vector from the gather indices tensor using `G`. +The element in the output tensor at index +[`G`,`W``0`,`W``1`] is then the element in the input +tensor at index [`X`+`W``0`,`Y`+`W``1`]. + +`window_bounds` is `[8,6]`, which decides the range of W`0` and +W`1`, and this in turn decides the bounds of the slice. + +This gather operation acts as a batch dynamic slice with `G` as the batch +dimension. + +The gather indices may be multidimensional. For instance, a more general +version of the example above using a "gather indices" tensor of shape `[4,5,2]` +would translate indices like this: + +
+ +
+ +Again, this acts as a batch dynamic slice `G``0` and +`G``1` as the batch dimensions. The window bounds are still `[8,6]`. + +The gather operation in XLA generalizes the informal semantics outlined above in +the following ways: + + 1. We can configure which dimensions in the output shape are the window + dimensions (dimensions containing `W``0`, `W``1` in + the last example). The output gather dimensions (dimensions containing + `G``0`, `G``1` in the last example) are defined to be + the output dimensions that are not window dimensions. + + 2. The number of output window dimensions explicitly present in the output + shape may be smaller than the input rank. These "missing" dimensions, which + are listed explicitly as `elided_window_dims`, must have a window bound of + `1`. Since they have a window bound of `1` the only valid index for them is + `0` and eliding them does not introduce ambiguity. + + 3. The slice extracted from the "Gather Indices" tensor ((`X`, `Y`) in the last + example) may have fewer elements than the input tensor rank, and an explicit + mapping dictates how the index should be expanded to have the same rank as + the input. + +As a final example, we use (2) and (3) to implement `tf.gather_nd`: + +
+ +
+ +`G``0` and `G``1` are used to slice out a starting index +from the gather indices tensor as usual, except the starting index has only one +element, `X`. Similarly, there is only one output window index with the value +`W``0`. However, before being used as indices into the input tensor, +these are expanded in accordance to "Gather Index Mapping" +(`gather_dims_to_operand_dims` in the formal description) and "Window Mapping" +(`window_dims_to_operand_dims` in the formal description) into +[`0`,`W``0`] and [`X`,`0`] respectively, adding up to +[`X`,`W``0`]. In other words, the output index +[`G``0`,`G``1`,`W``0`] maps to the input index +[`GatherIndices`[`G``0`,`G``1`,`0`],`X`] which gives us +the semantics for `tf.gather_nd`. + +`window_bounds` for this case is `[1,11]`. Intuitively this means that every +index `X` in the gather indices tensor picks an entire row and the result is the +concatenation of all these rows. ## GetTupleElement @@ -1075,7 +1293,7 @@ result2 = while (condition, init = result1) { ``` Nested tuple shapes are not supported. For an empty tuple shape, the Infeed -operation is effectively a nop and proceeds without reading any data from the +operation is effectively a no-op and proceeds without reading any data from the Infeed of the device. > Note: We plan to allow multiple Infeed operations without a total order, in @@ -1138,7 +1356,7 @@ dimension. `PaddingConfig` is a repeated field of `PaddingConfigDimension`, which contains three fields for each dimension: `edge_padding_low`, `edge_padding_high`, and -`interior_padding`. `edge_padding_low` and `edge_padding_high` specifies the +`interior_padding`. `edge_padding_low` and `edge_padding_high` specify the amount of padding added at the low-end (next to index 0) and the high-end (next to the highest index) of each dimension respectively. The amount of edge padding can be negative -- the absolute value of negative padding indicates the number @@ -1147,8 +1365,8 @@ the amount of padding added between any two elements in each dimension. Interior padding occurs logically before edge padding, so in the case of negative edge padding elements are removed from the interior-padded operand. This operation is a no-op if the edge padding pairs are all (0, 0) and the interior padding values -are all 0. Figure below shows examples of different `edge_padding` and -`interior_padding` values for a two dimensional array. +are all 0. The figure below shows examples of different `edge_padding` and +`interior_padding` values for a two-dimensional array.
diff --git a/tensorflow/docs_src/programmers_guide/datasets.md b/tensorflow/docs_src/programmers_guide/datasets.md index 9ede4ab83c1dcdb7370e83dfb9227fbb235d0689..9ccdbde627e6b2415835f7c0771eca1afa04f7f8 100644 --- a/tensorflow/docs_src/programmers_guide/datasets.md +++ b/tensorflow/docs_src/programmers_guide/datasets.md @@ -18,11 +18,11 @@ The `tf.data` API introduces two new abstractions to TensorFlow: tensors representing the image data and a label. There are two distinct ways to create a dataset: - * Creating a **source** (e.g. `Dataset.from_tensor_slices()`) constructs a + * Creating a **source** (e.g. `Dataset.from_tensor_slices()`) constructs a dataset from one or more `tf.Tensor` objects. - * Applying a **transformation** (e.g. `Dataset.batch()`) constructs a dataset + * Applying a **transformation** (e.g. `Dataset.batch()`) constructs a dataset from one or more `tf.data.Dataset` objects. * A `tf.data.Iterator` provides the main way to extract elements from a @@ -322,9 +322,39 @@ sess.run(iterator.initializer) next1, (next2, next3) = iterator.get_next() ``` -Note that evaluating *any* of `next1`, `next2`, or `next3` will advance the -iterator for all components. A typical consumer of an iterator will include all -components in a single expression. +Note that `next1`, `next2`, and `next3` are tensors produced by the +same op/node (created by `Iterator.get_next()`). Therefore, evaluating *any* of +these tensors will advance the iterator for all components. A typical consumer +of an iterator will include all components in a single expression. + +### Saving iterator state + +The @{tf.contrib.data.make_saveable_from_iterator} function creates a +`SaveableObject` from an iterator, which can be used to save and +restore the current state of the iterator (and, effectively, the whole input +pipeline). A saveable object thus created can be added to @{tf.train.Saver} +variables list or the `tf.GraphKeys.SAVEABLE_OBJECTS` collection for saving and +restoring in the same manner as a @{tf.Variable}. Refer to +@{$saved_model$Saving and Restoring} for details on how to save and restore +variables. + +```python +# Create saveable object from iterator. +saveable = tf.contrib.data.make_saveable_from_iterator(iterator) + +# Save the iterator state by adding it to the saveable objects collection. +tf.add_to_collection(tf.GraphKeys.SAVEABLE_OBJECTS, saveable) +saver = tf.train.Saver() + +with tf.Session() as sess: + + if should_checkpoint: + saver.save(path_to_checkpoint) + +# Restore the iterator state. +with tf.Session() as sess: + saver.restore(sess, path_to_checkpoint) +``` ## Reading input data diff --git a/tensorflow/docs_src/programmers_guide/debugger.md b/tensorflow/docs_src/programmers_guide/debugger.md index 9eaee2702829cbfd96cd56e832003724eba5bb1b..d1399814ee862f5f7ecc3f448d51fb3724fa3447 100644 --- a/tensorflow/docs_src/programmers_guide/debugger.md +++ b/tensorflow/docs_src/programmers_guide/debugger.md @@ -1,4 +1,4 @@ -# Debugging TensorFlow Programs +# TensorFlow Debugger @@ -23,8 +23,13 @@ debuggers such as Python's `pdb` due to TensorFlow's computation-graph paradigm. > installed using `pip install .whl`, however curses on Windows > may not work as reliably as curses on Linux or Mac. -This tutorial demonstrates how to use the **tfdbg** command-line interface -(CLI) to debug the appearance of [`nan`s](https://en.wikipedia.org/wiki/NaN) +> NOTE: This guide focuses on the command-line interface (CLI) of tfdbg. For +> guide on how to use the graphical user interface (GUI) of tfdbg, i.e., the +> **TensorBoard Debugger Plugin**, please visit +> [its README](https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/debugger/README.md). + +This tutorial demonstrates how to use the **tfdbg** CLI to debug the appearance +of [`nan`s](https://en.wikipedia.org/wiki/NaN) and [`inf`s](https://en.wikipedia.org/wiki/Infinity), a frequently-encountered type of bug in TensorFlow model development. The following example is for users who use the low-level @@ -214,7 +219,7 @@ navigate between these screens by clicking the `<--` and ### Other Features of the tfdbg CLI In addition to the commands listed above, the tfdbg CLI provides the following -addditional features: +additional features: * To navigate through previous tfdbg commands, type in a few characters followed by the Up or Down arrow keys. tfdbg will show you the history of @@ -454,7 +459,7 @@ accuracy_score = classifier.evaluate(x=test_set.data, [debug_tflearn_iris.py](https://www.tensorflow.org/code/tensorflow/python/debug/examples/debug_tflearn_iris.py), -based on {$tflearn$tf-learn's iris tutorial}, contains a full example of how to +based on [tf-learn's iris tutorial](https://www.tensorflow.org/versions/r1.2/get_started/tflearn), contains a full example of how to use the tfdbg with `Estimator`s. To run this example, do: ```none @@ -748,6 +753,7 @@ There are three possible workarounds or solutions: # For LocalCLIDebugHook hooks = [tf_debug.LocalCLIDebugHook(dump_root="/with/lots/of/space")] ``` + Make sure that the directory pointed to by dump_root is empty or nonexistent. tfdbg cleans up the dump directories before exiting. * Reduce the batch size used during the runs. @@ -806,3 +812,13 @@ sess.run(b) the constant-folding would not occur and `tfdbg` should show the intermediate tensor dumps. + +**Q**: Is there a GUI for tfdbg? + +**A**: Yes, the **TensorBoard Debugger Plugin** is the GUI of tfdbg. + It offers features such as inspection of the computation graph, + real-time visualization of tensor values, continuation to tensor + and conditional breakpoints, and tying tensors to their + graph-construction source code, all in the browser environment. + To get started, please visit + [its README](https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/debugger/README.md). diff --git a/tensorflow/docs_src/programmers_guide/faq.md b/tensorflow/docs_src/programmers_guide/faq.md index 70931f2862de98cb1e934f85919d558a3b36304a..392ac6f7f12532c3efce5bec1917691f55c7bee5 100644 --- a/tensorflow/docs_src/programmers_guide/faq.md +++ b/tensorflow/docs_src/programmers_guide/faq.md @@ -159,8 +159,7 @@ available. These operations allow you to build sophisticated @{$reading_data$input pipelines}, at the cost of making the TensorFlow computation somewhat more complicated. See the how-to documentation for -@{$reading_data#creating-threads-to-prefetch-using-queuerunner-objects$using -`QueueRunner` objects to drive queues and readers} +@{$reading_data#creating_threads_to_prefetch_using_queuerunner_objects$using `QueueRunner` objects to drive queues and readers} for more information on how to use them. ## Variables @@ -273,7 +272,7 @@ Prefer predefined TensorFlow operations such as @{tf.decode_raw}, If your data is not easily parsable with the built-in TensorFlow operations, consider converting it, offline, to a format that is easily parsable, such -as ${tf.python_io.TFRecordWriter$`TFRecord`} format. +as @{tf.python_io.TFRecordWriter$`TFRecord`} format. The more efficient method to customize the parsing behavior is to @{$adding_an_op$add a new op written in C++} that parses your diff --git a/tensorflow/docs_src/programmers_guide/graphs.md b/tensorflow/docs_src/programmers_guide/graphs.md index 2b4896c381052b5a3fb97385a18dbff82c2c0d89..e69b717432e6a8fab0085eb419dcbc0991cd9d28 100644 --- a/tensorflow/docs_src/programmers_guide/graphs.md +++ b/tensorflow/docs_src/programmers_guide/graphs.md @@ -125,14 +125,14 @@ an operation: @{tf.Tensor} accepts an optional `name` argument. For example, `tf.constant(42.0, name="answer")` creates a new @{tf.Operation} named `"answer"` and returns a @{tf.Tensor} named `"answer:0"`. If the default graph - already contained an operation named `"answer"`, the TensorFlow would append + already contains an operation named `"answer"`, then TensorFlow would append `"_1"`, `"_2"`, and so on to the name, in order to make it unique. * The @{tf.name_scope} function makes it possible to add a **name scope** prefix to all operations created in a particular context. The current name scope prefix is a `"/"`-delimited list of the names of all active @{tf.name_scope} context managers. If a name scope has already been used in the current - context, TensorFlow appens `"_1"`, `"_2"`, and so on. For example: + context, TensorFlow appends `"_1"`, `"_2"`, and so on. For example: ```python c_0 = tf.constant(0, name="c") # => operation named "c" @@ -210,9 +210,8 @@ with tf.device("/device:GPU:0"): # Operations created in this context will be pinned to the GPU. result = tf.matmul(weights, img) ``` - -If you are deploying TensorFlow in a @{$deploy/distributed$typical distributed -configuration}, you might specify the job name and task ID to place variables on +If you are deploying TensorFlow in a @{$deploy/distributed$typical distributed configuration}, +you might specify the job name and task ID to place variables on a task in the parameter server job (`"/job:ps"`), and the other operations on task in the worker job (`"/job:worker"`): @@ -336,20 +335,20 @@ described below. controls the behavior of the session. For example, some of the configuration options include: - * `allow_soft_placement`. Set this to `True` to enable a "soft" device + * `allow_soft_placement`. Set this to `True` to enable a "soft" device placement algorithm, which ignores @{tf.device} annotations that attempt to place CPU-only operations on a GPU device, and places them on the CPU instead. - * `cluster_def`. When using distributed TensorFlow, this option allows you + * `cluster_def`. When using distributed TensorFlow, this option allows you to specify what machines to use in the computation, and provide a mapping between job names, task indices, and network addresses. See @{tf.train.ClusterSpec.as_cluster_def} for details. - * `graph_options.optimizer_options`. Provides control over the optimizations + * `graph_options.optimizer_options`. Provides control over the optimizations that TensorFlow performs on your graph before executing it. - * `gpu_options.allow_growth`. Set this to `True` to change the GPU memory + * `gpu_options.allow_growth`. Set this to `True` to change the GPU memory allocator so that it gradually increases the amount of memory allocated, rather than allocating most of the memory at startup. diff --git a/tensorflow/docs_src/programmers_guide/index.md b/tensorflow/docs_src/programmers_guide/index.md index d45e666ce7b440bae20ba32d894526372af7e17b..e8c2fa6990c8ecfca1cfe76b3f813b4ae6917742 100644 --- a/tensorflow/docs_src/programmers_guide/index.md +++ b/tensorflow/docs_src/programmers_guide/index.md @@ -13,7 +13,7 @@ works. The units are as follows: ## Low Level APIs * @{$programmers_guide/low_level_intro}, which introduces the - basics of how you can to use TensorFlow outside of the high Level APIs. + basics of how you can use TensorFlow outside of the high Level APIs. * @{$programmers_guide/tensors}, which explains how to create, manipulate, and access Tensors--the fundamental object in TensorFlow. * @{$programmers_guide/variables}, which details how @@ -30,8 +30,12 @@ works. The units are as follows: can still be helpful. * @{$programmers_guide/saved_model}, which explains how to save and restore variables and models. + +## Accelerators + * @{$using_gpu} explains how TensorFlow assigns operations to devices and how you can change the arrangement manually. + * @{$using_tpu} explains how to modify `Estimator` programs to run on a TPU. ## ML Concepts diff --git a/tensorflow/docs_src/programmers_guide/leftnav_files b/tensorflow/docs_src/programmers_guide/leftnav_files index 38de3ccc3e474e6051976c810519212da8f5051e..3fe4cb2ddaee40d9d6c6470bee171dedb27ad890 100644 --- a/tensorflow/docs_src/programmers_guide/leftnav_files +++ b/tensorflow/docs_src/programmers_guide/leftnav_files @@ -10,7 +10,10 @@ tensors.md variables.md graphs.md saved_model.md + +### Accelerators using_gpu.md +using_tpu.md ### ML Concepts embedding.md @@ -19,9 +22,9 @@ embedding.md debugger.md ### TensorBoard -summaries_and_tensorboard.md -graph_viz.md -tensorboard_histograms.md +summaries_and_tensorboard.md: Visualizing Learning +graph_viz.md: Graphs +tensorboard_histograms.md: Histograms ### Misc version_compat.md diff --git a/tensorflow/docs_src/programmers_guide/low_level_intro.md b/tensorflow/docs_src/programmers_guide/low_level_intro.md index 8f6d3fbd46d8b76d6033d95fd51c1df45733f5a3..05709ad10a9275953d351e4a62cbf6d7fbffbbe3 100644 --- a/tensorflow/docs_src/programmers_guide/low_level_intro.md +++ b/tensorflow/docs_src/programmers_guide/low_level_intro.md @@ -286,6 +286,23 @@ while True: break ``` +If the `Dataset` depends on stateful operations you may need to +initialize the iterator before using it, as shown below: + +``` python +r = tf.random_normal([10,3]) +dataset = tf.data.Dataset.from_tensor_slices(r) +iterator = dataset.make_initializable_iterator() +next_row = iterator.get_next() + +sess.run(iterator.initializer) +while True: + try: + print(sess.run(next_row)) + except tf.errors.OutOfRangeError: + break +``` + For more details on Datasets and Iterators see: @{$programmers_guide/datasets}. ## Layers @@ -295,7 +312,7 @@ the same input. @{tf.layers$Layers} are the preferred way to add trainable parameters to a graph. Layers package together both the variables and the operations that act -on them, . For example a +on them. For example a [densely-connected layer](https://developers.google.com/machine-learning/glossary/#fully_connected_layer) performs a weighted sum across all inputs for each output and applies an optional @@ -478,7 +495,7 @@ good. Here's what we got; your own output will almost certainly differ: [ 0.10527515]] ``` -### loss +### Loss To optimize a model, you first need to define the loss. We'll use the mean square error, a standard loss for regression problems. @@ -504,7 +521,7 @@ TensorFlow provides [**optimizers**](https://developers.google.com/machine-learning/glossary/#optimizer) implementing standard optimization algorithms. These are implemented as sub-classes of @{tf.train.Optimizer}. They incrementally change each -variable in order to minimizethe loss. The simplest optimization algorithm is +variable in order to minimize the loss. The simplest optimization algorithm is [**gradient descent**](https://developers.google.com/machine-learning/glossary/#gradient_descent), implemented by @{tf.train.GradientDescentOptimizer}. It modifies each variable according to the magnitude of the derivative of loss with respect to diff --git a/tensorflow/docs_src/programmers_guide/saved_model.md b/tensorflow/docs_src/programmers_guide/saved_model.md index 9f50be5b31cd8b61b81426f50aa9ef9beb3138f2..55ee42dd6405db6bd34b064d71deaeb94839b0fa 100644 --- a/tensorflow/docs_src/programmers_guide/saved_model.md +++ b/tensorflow/docs_src/programmers_guide/saved_model.md @@ -1,35 +1,33 @@ -# Saving and Restoring +# Save and Restore -This document explains how to save and restore -@{$variables$variables} and models. +The @{tf.train.Saver} class provides methods to save and restore models. The +@{tf.saved_model.simple_save} function is an easy way to build a +@{tf.saved_model$saved model} suitable for serving. +[Estimators](@{$programmers_guide/estimators}) automatically save and restore +variables in the `model_dir`. +## Save and restore variables -## Saving and restoring variables +TensorFlow @{$variables} are the best way to represent shared, persistent state +manipulated by your program. The `tf.train.Saver` constructor adds `save` and +`restore` ops to the graph for all, or a specified list, of the variables in the +graph. The `Saver` object provides methods to run these ops, specifying paths +for the checkpoint files to write to or read from. -A TensorFlow variable provides the best way to represent shared, persistent -state manipulated by your program. (See @{$variables$Variables} for details.) -This section explains how to save and restore variables. -Note that Estimators automatically saves and restores variables -(in the `model_dir`). - -The `tf.train.Saver` class provides methods for saving and restoring models. -The `tf.train.Saver` constructor adds `save` and `restore` ops to the graph -for all, or a specified list, of the variables in the graph. The `Saver` -object provides methods to run these ops, specifying paths for the checkpoint -files to write to or read from. - -The saver will restore all variables already defined in your model. If you're +`Saver` restores all variables already defined in your model. If you're loading a model without knowing how to build its graph (for example, if you're writing a generic program to load models), then read the [Overview of saving and restoring models](#models) section later in this document. -TensorFlow saves variables in binary **checkpoint files** that, -roughly speaking, map variable names to tensor values. +TensorFlow saves variables in binary *checkpoint files* that map variable +names to tensor values. +Caution: TensorFlow model files are code. Be careful with untrusted code. +See [Using TensorFlow Securely](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) +for details. - -### Saving variables +### Save variables Create a `Saver` with `tf.train.Saver()` to manage all variables in the model. For example, the following snippet demonstrates how to call the @@ -61,9 +59,7 @@ with tf.Session() as sess: print("Model saved in path: %s" % save_path) ``` - - -### Restoring variables +### Restore variables The `tf.train.Saver` object not only saves variables to checkpoint files, it also restores variables. Note that when you restore variables you do not have @@ -92,14 +88,11 @@ with tf.Session() as sess: print("v2 : %s" % v2.eval()) ``` -Notes: - -* There is not a physical file called "/tmp/model.ckpt". It is the **prefix** - of filenames created for the checkpoint. Users only interact with the - prefix instead of physical checkpoint files. +Note: There is not a physical file called `/tmp/model.ckpt`. It is the *prefix* of +filenames created for the checkpoint. Users only interact with the prefix +instead of physical checkpoint files. - -### Choosing which variables to save and restore +### Choose variables to save and restore If you do not pass any arguments to `tf.train.Saver()`, the saver handles all variables in the graph. Each variable is saved under the name that was passed @@ -198,29 +191,42 @@ chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='v2', all_t -## Overview of saving and restoring models +## Save and restore models + +Use `SavedModel` to save and load your model—variables, the graph, and the +graph's metadata. This is a language-neutral, recoverable, hermetic +serialization format that enables higher-level systems and tools to produce, +consume, and transform TensorFlow models. TensorFlow provides several ways to +interact with `SavedModel`, including the @{tf.saved_model} APIs, +@{tf.estimator.Estimator}, and a command-line interface. + -When you want to save and load variables, the graph, and the -graph's metadata--basically, when you want to save or restore -your model--we recommend using SavedModel. -**SavedModel** is a language-neutral, recoverable, hermetic -serialization format. SavedModel enables higher-level systems -and tools to produce, consume, and transform TensorFlow models. -TensorFlow provides several mechanisms for interacting with -SavedModel, including tf.saved_model APIs, Estimator APIs and a CLI. +## Build and load a SavedModel +### Simple save -## APIs to build and load a SavedModel +The easiest way to create a `SavedModel` is to use the @{tf.saved_model.simple_save} +function: + +```python +simple_save(session, + export_dir, + inputs={"x": x, "y": y}, + outputs={"z": z}) +``` -This section focuses on the APIs for building and loading a SavedModel, -particularly when using lower-level TensorFlow APIs. +This configures the `SavedModel` so it can be loaded by +[TensorFlow serving](/serving/serving_basic) and supports the +[Predict API](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/predict.proto). +To access the classify, regress, or multi-inference APIs, use the manual +`SavedModel` builder APIs or an @{tf.estimator.Estimator}. +### Manually build a SavedModel -### Building a SavedModel +If your use case isn't covered by @{tf.saved_model.simple_save}, use the manual +@{tf.saved_model.builder$builder APIs} to create a `SavedModel`. -We provide a Python implementation of the SavedModel -@{tf.saved_model.builder$builder}. -The `SavedModelBuilder` class provides functionality to +The @{tf.saved_model.builder.SavedModelBuilder} class provides functionality to save multiple `MetaGraphDef`s. A **MetaGraph** is a dataflow graph, plus its associated variables, assets, and signatures. A **`MetaGraphDef`** is the protocol buffer representation of a MetaGraph. A **signature** is @@ -250,16 +256,51 @@ with tf.Session(graph=tf.Graph()) as sess: builder.add_meta_graph_and_variables(sess, [tag_constants.TRAINING], signature_def_map=foo_signatures, - assets_collection=foo_assets) + assets_collection=foo_assets, + strip_default_attrs=True) ... # Add a second MetaGraphDef for inference. with tf.Session(graph=tf.Graph()) as sess: ... - builder.add_meta_graph([tag_constants.SERVING]) + builder.add_meta_graph([tag_constants.SERVING], strip_default_attrs=True) ... builder.save() ``` + +#### Forward compatibility via `strip_default_attrs=True` + +Following the guidance below gives you forward compatibility only if the set of +Ops has not changed. + +The @{tf.saved_model.builder.SavedModelBuilder$`SavedModelBuilder`} class allows +users to control whether default-valued attributes must be stripped from the +@{$extend/tool_developers#nodes$`NodeDefs`} +while adding a meta graph to the SavedModel bundle. Both +@{tf.saved_model.builder.SavedModelBuilder.add_meta_graph_and_variables$`SavedModelBuilder.add_meta_graph_and_variables`} +and @{tf.saved_model.builder.SavedModelBuilder.add_meta_graph$`SavedModelBuilder.add_meta_graph`} +methods accept a Boolean flag `strip_default_attrs` that controls this behavior. + +If `strip_default_attrs` is `False`, the exported @{tf.MetaGraphDef} will have +the default valued attributes in all its @{tf.NodeDef} instances. +This can break forward compatibility with a sequence of events such as the +following: + +* An existing Op (`Foo`) is updated to include a new attribute (`T`) with a + default (`bool`) at version 101. +* A model producer such as a "trainer binary" picks up this change (version 101) + to the `OpDef` and re-exports an existing model that uses Op `Foo`. +* A model consumer (such as [Tensorflow Serving](/serving)) running an older + binary (version 100) doesn't have attribute `T` for Op `Foo`, but tries to + import this model. The model consumer doesn't recognize attribute `T` in a + `NodeDef` that uses Op `Foo` and therefore fails to load the model. +* By setting `strip_default_attrs` to True, the model producers can strip away + any default valued attributes in the `NodeDefs`. This helps ensure that newly + added attributes with defaults don't cause older model consumers to fail + loading models regenerated with newer training binaries. + +See [compatibility guidance](https://www.tensorflow.org/programmers_guide/version_compat) +for more information. ### Loading a SavedModel in Python @@ -285,7 +326,7 @@ with tf.Session(graph=tf.Graph()) as sess: ``` -### Loading a Savedmodel in C++ +### Load a SavedModel in C++ The C++ version of the SavedModel [loader](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/saved_model/loader.h) @@ -303,6 +344,30 @@ LoadSavedModel(session_options, run_options, export_dir, {kSavedModelTagTrain}, &bundle); ``` +### Load and serve a SavedModel in TensorFlow serving + +You can easily load and serve a SavedModel with the TensorFlow Serving Model +Server binary. See [instructions](https://www.tensorflow.org/serving/setup#installing_using_apt-get) +on how to install the server, or build it if you wish. + +Once you have the Model Server, run it with: +``` +tensorflow_model_server --port=port-numbers --model_name=your-model-name --model_base_path=your_model_base_path +``` +Set the port and model_name flags to values of your choosing. The +model_base_path flag expects to be to a base directory, with each version of +your model residing in a numerically named subdirectory. If you only have a +single version of your model, simply place it in a subdirectory like so: +* Place the model in /tmp/model/0001 +* Set model_base_path to /tmp/model + +Store different versions of your model in numerically named subdirectories of a +common base directory. For example, suppose the base directory is `/tmp/model`. +If you have only one version of your model, store it in `/tmp/model/0001`. If +you have two versions of your model, store the second version in +`/tmp/model/0002`, and so on. Set the `--model-base_path` flag to the base +directory (`/tmp/model`, in this example). TensorFlow Model Server will serve +the model in the highest numbered subdirectory of that base directory. ### Standard constants @@ -335,7 +400,7 @@ defined in: After training an `Estimator` model, you may want to create a service from that model that takes requests and returns a result. You can run such a -service locally on your machine or deploy it scalably in the cloud. +service locally on your machine or deploy it in the cloud. To prepare a trained Estimator for serving, you must export it in the standard SavedModel format. This section explains how to: @@ -347,7 +412,7 @@ SavedModel format. This section explains how to: * Serve the model from a local server and request predictions. -### Preparing serving inputs +### Prepare serving inputs During training, an @{$premade_estimators#input_fn$`input_fn()`} ingests data and prepares it for use by the model. At serving time, similarly, a @@ -421,14 +486,15 @@ to expect and how to map them to your model's expected inputs. By contrast, the *output* portion of the signature is determined by the model. -### Performing the export +### Perform the export To export your trained Estimator, call @{tf.estimator.Estimator.export_savedmodel} with the export base path and the `serving_input_receiver_fn`. ```py -estimator.export_savedmodel(export_dir_base, serving_input_receiver_fn) +estimator.export_savedmodel(export_dir_base, serving_input_receiver_fn, + strip_default_attrs=True) ``` This method builds a new graph by first calling the @@ -444,7 +510,7 @@ Session. > Note: It is your responsibility to garbage-collect old exports. > Otherwise, successive exports will accumulate under `export_dir_base`. -### Specifying the outputs of a custom model +### Specify the outputs of a custom model When writing a custom `model_fn`, you must populate the `export_outputs` element of the @{tf.estimator.EstimatorSpec} return value. This is a dict of @@ -476,7 +542,7 @@ indicating which `SignatureDef` will be served when an inference request does not specify one. -### Serving the exported model locally +### Serve the exported model locally For local deployment, you can serve your model using [TensorFlow Serving](https://github.com/tensorflow/serving), an open-source project that loads a @@ -495,7 +561,7 @@ bazel-bin/tensorflow_serving/model_servers/tensorflow_model_server --port=9000 - Now you have a server listening for inference requests via gRPC on port 9000! -### Requesting predictions from a local server +### Request predictions from a local server The server responds to gRPC requests according to the [PredictionService](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/prediction_service.proto#L15) @@ -588,7 +654,7 @@ passing in sample inputs in various formats (for example, Python expressions) and then fetching the output. -### Installing the SavedModel CLI +### Install the SavedModel CLI Broadly speaking, you can install TensorFlow in either of the following two ways: @@ -670,15 +736,15 @@ executing the computation graph later. For example: $ saved_model_cli show --dir \ /tmp/saved_model_dir --tag_set serve --signature_def serving_default The given SavedModel SignatureDef contains the following input(s): -inputs['x'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: x:0 + inputs['x'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: x:0 The given SavedModel SignatureDef contains the following output(s): -outputs['y'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y:0 + outputs['y'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y:0 Method name is: tensorflow/serving/predict ``` @@ -690,32 +756,32 @@ $ saved_model_cli show --dir /tmp/saved_model_dir --all MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs: signature_def['classify_x2_to_y3']: -The given SavedModel SignatureDef contains the following input(s): -inputs['inputs'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: x2:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['scores'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y3:0 -Method name is: tensorflow/serving/classify + The given SavedModel SignatureDef contains the following input(s): + inputs['inputs'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: x2:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['scores'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y3:0 + Method name is: tensorflow/serving/classify ... signature_def['serving_default']: -The given SavedModel SignatureDef contains the following input(s): -inputs['x'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: x:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['y'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y:0 -Method name is: tensorflow/serving/predict + The given SavedModel SignatureDef contains the following input(s): + inputs['x'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: x:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['y'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y:0 + Method name is: tensorflow/serving/predict ``` @@ -815,7 +881,7 @@ For example: `=[{"age":[22,24],"education":["BS","MS"]}]` ``` -#### Save Output +#### Save output By default, the SavedModel CLI writes output to stdout. If a directory is passed to `--outdir` option, the outputs will be saved as npy files named after @@ -824,7 +890,7 @@ output tensor keys under the given directory. Use `--overwrite` to overwrite existing output files. -#### TensorFlow Debugger (tfdbg) Integration +#### TensorFlow debugger (tfdbg) integration If `--tf_debug` option is set, the SavedModel CLI will use the TensorFlow Debugger (tfdbg) to watch the intermediate Tensors and runtime @@ -931,6 +997,3 @@ of checkpoints and assets: Each graph is associated with a specific set of tags, which enables identification during a load or restore operation. - - - diff --git a/tensorflow/docs_src/programmers_guide/summaries_and_tensorboard.md b/tensorflow/docs_src/programmers_guide/summaries_and_tensorboard.md index 05dfdfdc4d2257fc680e7fa99b666ef86e3bef09..fadfa03e78349801d69e0045991a8fa9a0a59df9 100644 --- a/tensorflow/docs_src/programmers_guide/summaries_and_tensorboard.md +++ b/tensorflow/docs_src/programmers_guide/summaries_and_tensorboard.md @@ -16,10 +16,17 @@ TensorBoard is fully configured, it looks like this:
-This tutorial is intended to get you started with simple TensorBoard usage. -There are other resources available as well! The [TensorBoard's GitHub](https://github.com/tensorflow/tensorboard) -has a lot more information on TensorBoard usage, including tips & tricks, and -debugging information. +This 30-minute tutorial is intended to get you started with simple TensorBoard +usage. It assumes a basic understanding of TensorFlow. + +There are other resources available as well! The [TensorBoard GitHub](https://github.com/tensorflow/tensorboard) +has a lot more information on using individual dashboards within TensorBoard +including tips & tricks and debugging information. + +## Setup + +[Install TensorFlow](https://www.tensorflow.org/install/). Installing TensorFlow +via pip should also automatically install TensorBoard. ## Serializing the data @@ -76,7 +83,7 @@ data than you need, though. Instead, consider running the merged summary op every `n` steps. The code example below is a modification of the -@{$layers$simple MNIST tutorial}, +[simple MNIST tutorial](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/mnist/mnist.py), in which we have added some summary ops, and run them every ten steps. If you run this and then launch `tensorboard --logdir=/tmp/tensorflow/mnist`, you'll be able to visualize statistics, such as how the weights or accuracy varied during @@ -214,4 +221,5 @@ corner. Each tab represents a set of serialized data that can be visualized. For in depth information on how to use the *graph* tab to visualize your graph, see @{$graph_viz$TensorBoard: Graph Visualization}. -For more usage information on TensorBoard in general, see the [TensorBoard's GitHub](https://github.com/tensorflow/tensorboard). +For more usage information on TensorBoard in general, see the +[TensorBoard GitHub](https://github.com/tensorflow/tensorboard). diff --git a/tensorflow/docs_src/programmers_guide/using_tpu.md b/tensorflow/docs_src/programmers_guide/using_tpu.md new file mode 100644 index 0000000000000000000000000000000000000000..a9c2cb3e33d4817b9a35400dcce9227ddd635ff4 --- /dev/null +++ b/tensorflow/docs_src/programmers_guide/using_tpu.md @@ -0,0 +1,395 @@ +# Using TPUs + +This document walks through the principal TensorFlow APIs necessary to make +effective use of a [Cloud TPU](https://cloud.google.com/tpu/), and highlights +the differences between regular TensorFlow usage, and usage on a TPU. + +This doc is aimed at users who: + +* Are familiar with TensorFlow's `Estimator` and `Dataset` APIs +* Have maybe [tried out a Cloud TPU](https://cloud.google.com/tpu/docs/quickstart) + using an existing model. +* Have, perhaps, skimmed the code of an example TPU model + [[1]](https://github.com/tensorflow/models/blob/master/official/mnist/mnist_tpu.py) + [[2]](https://github.com/tensorflow/tpu-demos/tree/master/cloud_tpu/models). +* Are interested in porting an existing `Estimator` model to + run on Cloud TPUs + +## TPUEstimator + +@{tf.estimator.Estimator$Estimators} are TensorFlow's model-level abstraction. +Standard `Estimators` can drive models on CPU and GPUs. You must use +@{tf.contrib.tpu.TPUEstimator} to drive a model on TPUs. + +Refer to TensorFlow's Getting Started section for an introduction to the basics +of using a @{$get_started/premade_estimators$pre-made `Estimator`}, and +@{$get_started/custom_estimators$custom `Estimator`s}. + +The `TPUEstimator` class differs somewhat from the `Estimator` class. + +The simplest way to maintain a model that can be run both on CPU/GPU or on a +Cloud TPU is to define the model's inference phase (from inputs to predictions) +outside of the `model_fn`. Then maintain separate implementations of the +`Estimator` setup and `model_fn`, both wrapping this inference step. For an +example of this pattern compare the `mnist.py` and `mnist_tpu.py` implementation in +[tensorflow/models](https://github.com/tensorflow/models/tree/master/official/mnist). + +### Running a `TPUEstimator` locally + +To create a standard `Estimator` you call the constructor, and pass it a +`model_fn`, for example: + +``` +my_estimator = tf.estimator.Estimator( + model_fn=my_model_fn) +``` + +The changes required to use a @{tf.contrib.tpu.TPUEstimator} on your local +machine are relatively minor. The constructor requires two additional arguments. +You should set the `use_tpu` argument to `False`, and pass a +@{tf.contrib.tpu.RunConfig} as the `config` argument, as shown below: + +``` python +my_tpu_estimator = tf.contrib.tpu.TPUEstimator( + model_fn=my_model_fn, + config=tf.contrib.tpu.RunConfig() + use_tpu=False) +``` + +Just this simple change will allow you to run a `TPUEstimator` locally. +The majority of example TPU models can be run in this local mode, +by setting the command line flags as follows: + + +``` +$> python mnist_tpu.py --use_tpu=false --master='' +``` + +Note: This `use_tpu=False` argument is useful for trying out the `TPUEstimator` +API. It is not meant to be a complete TPU compatibility test. Successfully +running a model locally in a `TPUEstimator` does not guarantee that it will +work on a TPU. + + +### Building a `tpu.RunConfig` + +While the default `RunConfig` is sufficient for local training, these settings +cannot be ignored in real usage. + +A more typical setup for a `RunConfig`, that can be switched to use a Cloud +TPU, might be as follows: + +``` python +import tempfile +import subprocess + +class FLAGS(object): + use_tpu=False + tpu_name=None + # Use a local temporary path for the `model_dir` + model_dir = tempfile.mkdtemp() + # Number of training steps to run on the Cloud TPU before returning control. + iterations = 50 + # A single Cloud TPU has 8 shards. + num_shards = 8 + +if FLAGS.use_tpu: + my_project_name = subprocess.check_output([ + 'gcloud','config','get-value','project']) + my_zone = subprocess.check_output([ + 'gcloud','config','get-value','compute/zone']) + cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver( + tpu_names=[FLAGS.tpu_name], + zone=my_zone, + project=my_project) + master = tpu_cluster_resolver.get_master() +else: + master = '' + +my_tpu_run_config = tf.contrib.tpu.RunConfig( + master=master, + evaluation_master=master, + model_dir=FLAGS.model_dir, + session_config=tf.ConfigProto( + allow_soft_placement=True, log_device_placement=True), + tpu_config=tf.contrib.tpu.TPUConfig(FLAGS.iterations, + FLAGS.num_shards), +) +``` + +Then you must pass the @{tf.contrib.tpu.RunConfig} to the constructor: + +``` python +my_tpu_estimator = tf.contrib.tpu.TPUEstimator( + model_fn=my_model_fn, + config = my_tpu_run_config, + use_tpu=FLAGS.use_tpu) +``` + +Typically the `FLAGS` would be set by command line arguments. To switch from +training locally to training on a cloud TPU you would need to: + +* Set `FLAGS.use_tpu` to `True` +* Set `FLAGS.tpu_name` so the `tf.contrib.cluster_resolver.TPUClusterResolver` can find it +* Set `FLAGS.model_dir` to a Google Cloud Storage bucket url (`gs://`). + + +## Optimizer + +When training on a cloud TPU you **must** wrap the optimizer in a +@{tf.contrib.tpu.CrossShardOptimizer}, which uses an `allreduce` to aggregate +gradients and broadcast the result to each shard (each TPU core). + +The `CrossShardOptimizer` is not compatible with local training. So, to have +the same code run both locally and on a Cloud TPU, add lines like the following: + +``` python +optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate) +if FLAGS.use_tpu: + optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer) +``` + +If you prefer to avoid a global `FLAGS` variable in your model code, one +approach is to set the optimizer as one of the `Estimator`'s params, +as follows: + +``` python +my_tpu_estimator = tf.contrib.tpu.TPUEstimator( + model_fn=my_model_fn, + config = my_tpu_run_config, + use_tpu=FLAGS.use_tpu, + params={'optimizer':optimizer}) +``` + +## Model Function + +This section details the changes you must make to the model function +(`model_fn()`) to make it `TPUEstimator` compatible. + +### Static shapes + +During regular usage TensorFlow attempts to determine the shapes of each +`tf.Tensor` during graph construction. During execution any unknown shape +dimensions are determined dynamically, +see @{$programmers_guide/tensors#shape$Tensor Shapes} for more details. + +To run on Cloud TPUs TensorFlow models are compiled using @{$xla$XLA}. +XLA uses a similar system for determining shapes at compile time. XLA requires +that all tensor dimensions be statically defined at compile time. All shapes +must evaluate to a constant, and not depend on external data, or stateful +operations like variables or a random number generator. + + +### Summaries + +Remove any use of `tf.summary` from your model. + +@{$summaries_and_tensorboard$TensorBoard summaries} are a great way see inside +your model. A minimal set of basic summaries are automatically recorded by the +`TPUEstimator`, to `event` files in the `model_dir`. Custom summaries, however, +are currently unsupported when training on a Cloud TPU. So while the +`TPUEstimator` will still run locally with summaries, it will fail if used on a +TPU. + +### Metrics + +Build your evaluation metrics dictionary in a stand-alone `metric_fn`. + + + +Evaluation metrics are an essential part of training a model. These are fully +supported on Cloud TPUs, but with a slightly different syntax. + +A standard @{tf.metrics} returns two tensors. The first returns the running +average of the metric value, while the second updates the running average and +returns the value for this batch: + +``` +running_average, current_batch = tf.metrics.accuracy(labels, predictions) +``` + +In a standard `Estimator` you create a dictionary of these pairs, and return it +as part of the `EstimatorSpec`. + +```python +my_metrics = {'accuracy': tf.metrics.accuracy(labels, predictions)} + +return tf.estimator.EstimatorSpec( + ... + eval_metric_ops=my_metrics +) +``` + +In a `TPUEstimator` you instead pass a function (which returns a metrics +dictionary) and a list of argument tensors, as shown below: + +```python +def my_metric_fn(labels, predictions): + return {'accuracy': tf.metrics.accuracy(labels, predictions)} + +return tf.contrib.tpu.TPUEstimatorSpec( + ... + eval_metrics=(my_metric_fn, [labels, predictions]) +) +``` + +### Use `TPUEstimatorSpec` + +`TPUEstimatorSpec` do not support hooks, and require function wrappers for +some fields. + +An `Estimator`'s `model_fn` must return an `EstimatorSpec`. An `EstimatorSpec` +is a simple structure of named fields containing all the `tf.Tensors` of the +model that the `Estimator` may need to interact with. + +`TPUEstimators` use a @{tf.contrib.tpu.TPUEstimatorSpec}. There are a few +differences between it and a standard @{tf.estimator.EstimatorSpec}: + + +* The `eval_metric_ops` must be wrapped into a `metrics_fn`, this field is + renamed `eval_metrics` ([see above](#metrics)). +* The @{tf.train.SessionRunHook$hooks} are unsupported, so these fields are + omitted. +* The @{tf.train.Scaffold$`scaffold`}, if used, must also be wrapped in a + function. This field is renamed to `scaffold_fn`. + +`Scaffold` and `Hooks` are for advanced usage, and can typically be omitted. + +## Input functions + +Input functions work mainly unchanged as they run on the host computer, not the +Cloud TPU itself. This section explains the two necessary adjustments. + +### Params argument + + + +The `input_fn` for a standard `Estimator` _can_ include a +`params` argument; the `input_fn` for a `TPUEstimator` *must* include a +`params` argument. This is necessary to allow the estimator to set the batch +size for each replica of the input stream. So the minimum signature for an +`input_fn` for a `TPUEstimator` is: + +``` +def my_input_fn(params): + pass +``` + +Where `params['batch-size']` will contain the batch size. + +### Static shapes and batch size + +The input pipeline generated by your `input_fn` is run on CPU. So it is mostly +free strict static shape requirements imposed by the XLA/TPU environment. The +one requirement is that the batches of data fed from your input pipeline to +the TPU have a static shape, as determined by the standard TensorFlow shape +inference algorithm. Intermediate tensors are free to have a dynamic shapes. +If shape inference has failed, but the shape is known it is possible to +impose the correct shape using `tf.set_shape()`. + +In the example below the shape +inference algorithm fails, but it is corrected using `set_shape`: + +``` +>>> x = tf.zeros(tf.constant([1,2,3])+1) +>>> x.shape + +TensorShape([Dimension(None), Dimension(None), Dimension(None)]) + +>>> x.set_shape([2,3,4]) +``` + +In many cases the batch size is the only unknown dimension. + +A typical input pipeline, using `tf.data`, will usually produce batches of a +fixed size. The last batch of a finite `Dataset`, however, is typically smaller, +containing just the remaining elements. Since a `Dataset` does not know its own +length or finiteness, the standard @{tf.data.Dataset.batch$`batch`} method +cannot determine if all batches will have a fixed size batch on its own: + +``` +>>> params = {'batch_size':32} +>>> ds = tf.data.Dataset.from_tensors([0, 1, 2]) +>>> ds = ds.repeat().batch(params['batch-size']) +>>> ds + + +``` + +The most straightforward fix is to +@{tf.data.Dataset.apply$apply} @{tf.contrib.data.batch_and_drop_remainder} +as follows: + +``` +>>> params = {'batch_size':32} +>>> ds = tf.data.Dataset.from_tensors([0, 1, 2]) +>>> ds = ds.repeat().apply( +... tf.contrib.data.batch_and_drop_remainder(params['batch-size'])) +>>> ds + + <_RestructuredDataset shapes: (32, 3), types: tf.int32> +``` + +The one downside to this approach is that, as the name implies, this batching +method throws out any fractional batch at the end of the dataset. This is fine +for an infinitely repeating dataset being used for training, but could be a +problem if you want to train for an exact number of epochs. + +To do an exact 1-epoch of _evaluation_ you can work around this by manually +padding the length of the batches, and setting the padding entries to have zero +weight when creating your `tf.metrics`. + +## Datasets + +Efficient use of the `tf.data.Dataset` API is critical when using a Cloud +TPU, as it is impossible to use the Cloud TPU's unless you can feed it data +quickly enough. See @{$datasets_performance} for details on dataset performance. + +For all but the simplest experimentation (using +@{tf.data.Dataset.from_tensor_slices} or other in-graph data) you will need to +store all data files read by the `TPUEstimator`'s `Dataset` in Google Cloud +Storage Buckets. + + + +For most use-cases, we recommend converting your data into `TFRecord` +format and using a @{tf.data.TFRecordDataset} to read it. This, however, is not +a hard requirement and you can use other dataset readers +(`FixedLengthRecordDataset` or `TextLineDataset`) if you prefer. + +Small datasets can be loaded entirely into memory using +@{tf.data.Dataset.cache}. + +Regardless of the data format used, it is strongly recommended that you +@{$performance_guide#use_large_files$use large files}, on the order of +100MB. This is especially important in this networked setting as the overhead +of opening a file is significantly higher. + +It is also important, regardless of the type of reader used, to enable buffering +using the `buffer_size` argument to the constructor. This argument is specified +in bytes. A minimum of a few MB (`buffer_size=8*1024*1024`) is recommended so +that data is available when needed. + +The TPU-demos repo includes +[a script](https://github.com/tensorflow/tpu-demos/blob/master/cloud_tpu/datasets/imagenet_to_gcs.py) +for downloading the imagenet dataset and converting it to an appropriate format. +This together with the imagenet +[models](https://github.com/tensorflow/tpu-demos/tree/master/cloud_tpu/models) +included in the repo demonstrate all of these best-practices. + + +## What Next + +For details on how to actually set up and run a Cloud TPU see: + + * [Google Cloud TPU Documentation](https://cloud.google.com/tpu/docs/) + +This document is by no means exhaustive. The best source of more detail on how +to make a Cloud TPU compatible model are the example models published in: + + * The [TPU Demos Repository.](https://github.com/tensorflow/tpu-demos/) + +For more information about tuning TensorFlow code for performance see: + + * The @{$performance$Performance Section.} + diff --git a/tensorflow/docs_src/programmers_guide/variables.md b/tensorflow/docs_src/programmers_guide/variables.md index 64250738056043e236b5eb236bcbf29375655260..e8cf7711552f4c83ed1e03e0753b580cc7505ddc 100644 --- a/tensorflow/docs_src/programmers_guide/variables.md +++ b/tensorflow/docs_src/programmers_guide/variables.md @@ -62,9 +62,10 @@ them. For this reason TensorFlow provides **collections**, which are named lists of tensors or other objects, such as `tf.Variable` instances. By default every `tf.Variable` gets placed in the following two collections: + * `tf.GraphKeys.GLOBAL_VARIABLES` --- variables that can be shared across -multiple devices, - * `tf.GraphKeys.TRAINABLE_VARIABLES`--- variables for which TensorFlow will + multiple devices, + * `tf.GraphKeys.TRAINABLE_VARIABLES` --- variables for which TensorFlow will calculate gradients. If you don't want a variable to be trainable, add it to the diff --git a/tensorflow/docs_src/programmers_guide/version_compat.md b/tensorflow/docs_src/programmers_guide/version_compat.md index a28f1385c87c7a083ee96977c5ab268c6977e17e..72e427c5f8f0f6581d528f4ead18699736eafd04 100644 --- a/tensorflow/docs_src/programmers_guide/version_compat.md +++ b/tensorflow/docs_src/programmers_guide/version_compat.md @@ -60,7 +60,8 @@ patch versions. The public APIs consist of * [`tensor_shape`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/tensor_shape.proto) * [`types`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/types.proto) -## What is *not* covered {not_covered} + +## What is *not* covered Some API functions are explicitly marked as "experimental" and can change in backward incompatible ways between minor releases. These include: @@ -182,7 +183,7 @@ Our versioning scheme has three requirements: * **Forward compatibility** to support scenarios where the producer of a graph or checkpoint is upgraded to a newer version of TensorFlow before the consumer. -* Enable evolving TensorFlow in incompatible ways. For example, removing Ops, +* Enable evolving TensorFlow in incompatible ways. For example, removing ops, adding attributes, and removing attributes. Note that while the `GraphDef` version mechanism is separate from the TensorFlow @@ -244,32 +245,51 @@ contains a main data version which is treated as either `producer` or `TF_CHECKPOINT_VERSION_MIN_CONSUMER`, and `TF_CHECKPOINT_VERSION_MIN_PRODUCER`. +### Add a new attribute with default to an existing op + +Following the guidance below gives you forward compatibility only if the set of +ops has not changed: + +1. If forward compatibility is desired, set `strip_default_attrs` to `True` + while exporting the model using either the + @{tf.saved_model.builder.SavedModelBuilder.add_meta_graph_and_variables$`add_meta_graph_and_variables`} + and @{tf.saved_model.builder.SavedModelBuilder.add_meta_graph$`add_meta_graph`} + methods of the `SavedModelBuilder` class, or + @{tf.estimator.Estimator.export_savedmodel$`Estimator.export_savedmodel`} +2. This strips off the default valued attributes at the time of + producing/exporting the models. This makes sure that the exported + @{tf.MetaGraphDef} does not contain the new op-attribute when the default + value is used. +3. Having this control could allow out-of-date consumers (for example, serving + binaries that lag behind training binaries) to continue loading the models + and prevent interruptions in model serving. + ### Evolving GraphDef versions This section explains how to use this versioning mechanism to make different types of changes to the `GraphDef` format. -#### Add an Op +#### Add an op -Add the new Op to both consumers and producers at the same time, and do not +Add the new op to both consumers and producers at the same time, and do not change any `GraphDef` versions. This type of change is automatically backward compatible, and does not impact forward compatibility plan since existing producer scripts will not suddenly use the new functionality. -#### Add an Op and switch existing Python wrappers to use it +#### Add an op and switch existing Python wrappers to use it 1. Implement new consumer functionality and increment the `GraphDef` version. 2. If it is possible to make the wrappers use the new functionality only in cases that did not work before, the wrappers can be updated now. 3. Change Python wrappers to use the new functionality. Do not increment - `min_consumer`, since models that do not use this Op should not break. + `min_consumer`, since models that do not use this op should not break. -#### Remove or restrict an Op's functionality +#### Remove or restrict an op's functionality -1. Fix all producer scripts (not TensorFlow itself) to not use the banned Op or +1. Fix all producer scripts (not TensorFlow itself) to not use the banned op or functionality. 2. Increment the `GraphDef` version and implement new consumer functionality - that bans the removed Op or functionality for GraphDefs at the new version + that bans the removed op or functionality for GraphDefs at the new version and above. If possible, make TensorFlow stop producing `GraphDefs` with the banned functionality. To do so, add the [`REGISTER_OP(...).Deprecated(deprecated_at_version, @@ -278,15 +298,15 @@ existing producer scripts will not suddenly use the new functionality. 4. Increase `min_producer` to the GraphDef version from (2) and remove the functionality entirely. -#### Change an Op's functionality +#### Change an op's functionality -1. Add a new similar Op named `SomethingV2` or similar and go through the +1. Add a new similar op named `SomethingV2` or similar and go through the process of adding it and switching existing Python wrappers to use it, which may take three weeks if forward compatibility is desired. -2. Remove the old Op (Can only take place with a major version change due to +2. Remove the old op (Can only take place with a major version change due to backward compatibility). -3. Increase `min_consumer` to rule out consumers with the old Op, add back the - old Op as an alias for `SomethingV2`, and go through the process to switch +3. Increase `min_consumer` to rule out consumers with the old op, add back the + old op as an alias for `SomethingV2`, and go through the process to switch existing Python wrappers to use it. 4. Go through the process to remove `SomethingV2`. @@ -294,6 +314,6 @@ existing producer scripts will not suddenly use the new functionality. 1. Bump the `GraphDef` version and add the bad version to `bad_consumers` for all new GraphDefs. If possible, add to `bad_consumers` only for GraphDefs - which contain a certain Op or similar. + which contain a certain op or similar. 2. If existing consumers have the bad version, push them out as soon as possible. diff --git a/tensorflow/docs_src/tutorials/deep_cnn.md b/tensorflow/docs_src/tutorials/deep_cnn.md index 679754020470dddfcffa76e62ca8f55a439ec4f5..6a4c9a9b0727208a158b1b57d13ca70290961ec2 100644 --- a/tensorflow/docs_src/tutorials/deep_cnn.md +++ b/tensorflow/docs_src/tutorials/deep_cnn.md @@ -268,7 +268,7 @@ in `cifar10_input.py`. `cifar10_train.py` periodically @{tf.train.Saver$saves} all model parameters in -@{$variables#saving-and-restoring$checkpoint files} +@{$programmers_guide/saved_model$checkpoint files} but it does *not* evaluate the model. The checkpoint file will be used by `cifar10_eval.py` to measure the predictive performance (see [Evaluating a Model](#evaluating-a-model) below). diff --git a/tensorflow/docs_src/tutorials/image_retraining.md b/tensorflow/docs_src/tutorials/image_retraining.md index df15bc0a9c3763aa51c2fc8cf36ce9fc3544ae68..93d7c86e42aa90d145d27b56edc0abfec7034686 100644 --- a/tensorflow/docs_src/tutorials/image_retraining.md +++ b/tensorflow/docs_src/tutorials/image_retraining.md @@ -115,7 +115,7 @@ process is progressing. The training's objective is to make the loss as small as possible, so you can tell if the learning is working by keeping an eye on whether the loss keeps trending downwards, ignoring the short-term noise. -By default this script will run 4,000 training steps. Each step chooses ten +By default this script will run 4,000 training steps. Each step chooses 100 images at random from the training set, finds their bottlenecks from the cache, and feeds them into the final layer to get predictions. Those predictions are then compared against the actual labels to update the final layer's weights @@ -349,31 +349,32 @@ results, but if you intend to deploy your model on mobile devices or other resource-constrained environments you may want to trade off a little accuracy for much smaller file sizes or faster speeds. To help with that, the [retrain.py script](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/image_retraining/retrain.py) -supports 32 different variations on the [Mobilenet architecture](https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html). +supports different variations on the [Mobilenet architecture](https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html). These are a little less precise than Inception v3, but can result in far -smaller file sizes (down to less than a megabyte) and can be many times faster +smaller file sizes (a few megabytes) and can be many times faster to run. To train with one of these models, pass in the `--architecture` flag, for example: ``` python tensorflow/examples/image_retraining/retrain.py \ - --image_dir ~/flower_photos --architecture mobilenet_0.25_128_quantized + --image_dir ~/flower_photos --architecture mobilenet_0.25_128 ``` -This will create a 941KB model file in `/tmp/output_graph.pb`, with 25% of the -parameters of the full Mobilenet, taking 128x128 sized input images, and with -its weights quantized down to eight bits on disk. You can choose '1.0', '0.75', -'0.50', or '0.25' to control the number of weight parameters, and so the file -size (and to some extent the speed), '224', '192', '160', or '128' for the input -image size, with smaller sizes giving faster speeds, and an optional -'_quantized' at the end to indicate whether the file should contain 8-bit or -32-bit float weights. +This will create a 1.9MB model file in `/tmp/output_graph.pb`, with only 25% of +the number of neurons of the full Mobilenet, and trained to take 128x128 sized +input images. + +You can choose '1.0', '0.75', '0.50', or '0.25' to control the number of +neurons (activations of hidden layers); the number of weights (and hence to +some extent the file size and speed) shrinks like the square of that fraction. +You can choose '224', '192', '160', or '128' for the input image size, +with smaller sizes giving faster speeds. The speed and size advantages come at a loss to accuracy of course, but for many purposes this isn't critical. They can also be somewhat offset with improved training data. For example, training with distortions allows me to get above 80% -accuracy on the flower data set even with the 0.25/128/quantized graph above. +accuracy on the flower data set even with the 0.25/128 graph above. If you're going to be using the Mobilenet models in label_image or your own programs, you'll need to feed in an image of the specified size converted to a @@ -395,3 +396,9 @@ python tensorflow/examples/label_image/label_image.py \ --input_mean=128 --input_std=128 \ --image=$HOME/flower_photos/daisy/21652746_cc379e0eea_m.jpg ``` + +For more information on deploying the retrained model to a mobile device, see +the [codelab version](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/#0) +of this tutorial, especially [part 2](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets-2-tflite/#0), which describes +[TensorFlow Lite](/mobile/tflite/) and the additional optimizations it offers +(including quantization of model weights). diff --git a/tensorflow/docs_src/tutorials/index.md b/tensorflow/docs_src/tutorials/index.md index 8c697e48e550c4e425db33bab7257532d209ac7a..af01d3eaa12157f82c981de005708509f6652cca 100644 --- a/tensorflow/docs_src/tutorials/index.md +++ b/tensorflow/docs_src/tutorials/index.md @@ -10,7 +10,7 @@ these tutorials. These tutorials cover different aspects of image recognition: - * @{$layers}, which introduces convolutional neural networks (CNNs) and + * @{$layers$MNIST}, which introduces convolutional neural networks (CNNs) and demonstrates how to build a CNN in TensorFlow. * @{$image_recognition}, which introduces the field of image recognition and uses a pre-trained model (Inception) for recognizing images. diff --git a/tensorflow/docs_src/tutorials/kernel_methods.md b/tensorflow/docs_src/tutorials/kernel_methods.md index 63f408c2ca304d6345ffff459b799b011f8d8035..b1f06ce0a3b35a79b60f18dfe8a40dda0c732f07 100644 --- a/tensorflow/docs_src/tutorials/kernel_methods.md +++ b/tensorflow/docs_src/tutorials/kernel_methods.md @@ -1,9 +1,9 @@ # Improving Linear Models Using Explicit Kernel Methods -Note: This document uses a deprecated version of ${tf.estimator}, -which has a ${tf.contrib.learn.estimator$different interface}. +Note: This document uses a deprecated version of @{tf.estimator}, +which has a different interface (see `tf.contrib.learn Estimator`). It also uses other `contrib` methods whose -${$version_compat#not_covered$API may not be stable}. +@{$version_compat#not_covered$API may not be stable}. In this tutorial, we demonstrate how combining (explicit) kernel methods with linear models can drastically increase the latters' quality of predictions diff --git a/tensorflow/docs_src/tutorials/layers.md b/tensorflow/docs_src/tutorials/layers.md index b898cbe29c2bac9ade341fe3b3566e42e133fc5b..9b17d0d4d52e6c5aa9b00b2e50840fa8cd4e8cb6 100644 --- a/tensorflow/docs_src/tutorials/layers.md +++ b/tensorflow/docs_src/tutorials/layers.md @@ -193,7 +193,7 @@ to calculate loss, configure the training op, and generate predictions. If you're already experienced with CNNs and @{$get_started/custom_estimators$TensorFlow `Estimator`s}, and find the above code intuitive, you may want to skim these sections or just skip ahead to ["Training and Evaluating the CNN MNIST -Classifier"](#training-and-evaluating-the-cnn-mnist-classifier). +Classifier"](#training_and_evaluating_the_cnn_mnist_classifier). ### Input Layer @@ -446,7 +446,7 @@ tf.nn.softmax(logits, name="softmax_tensor") > Note: We use the `name` argument to explicitly name this operation > `softmax_tensor`, so we can reference it later. (We'll set up logging for the -> softmax values in ["Set Up a Logging Hook"](#set-up-a-logging-hook). +> softmax values in ["Set Up a Logging Hook"](#set-up-a-logging-hook)). We compile our predictions in a dict, and return an `EstimatorSpec` object: @@ -534,9 +534,8 @@ if mode == tf.estimator.ModeKeys.TRAIN: ``` > Note: For a more in-depth look at configuring training ops for Estimator model -> functions, see @{$get_started/custom_estimators#defining-the-training-op-for-the-model$"Defining -> the training op for the model"} in the @{$get_started/custom_estimators$"Creating Estimations in -> tf.estimator"} tutorial. +> functions, see @{$get_started/custom_estimators#defining_the_training_op_for_the_model$"Defining the training op for the model"} +> in the @{$get_started/custom_estimators$"Creating Estimators in tf.estimator."} tutorial. ### Add evaluation metrics @@ -625,8 +624,8 @@ operation earlier when we generated the probabilities in `cnn_model_fn`. > Note: If you don't explicitly assign a name to an operation via the `name` > argument, TensorFlow will assign a default name. A couple easy ways to > discover the names applied to operations are to visualize your graph on -> @{$graph_viz$TensorBoard}) or to enable the @{$debugger$TensorFlow Debugger -> (tfdbg)}. +> @{$graph_viz$TensorBoard}) or to enable the +> @{$programmers_guide/debugger$TensorFlow Debugger (tfdbg)}. Next, we create the `LoggingTensorHook`, passing `tensors_to_log` to the `tensors` argument. We set `every_n_iter=50`, which specifies that probabilities @@ -635,7 +634,7 @@ should be logged after every 50 steps of training. ### Train the Model Now we're ready to train our model, which we can do by creating `train_input_fn` -ans calling `train()` on `mnist_classifier`. Add the following to `main()`: +and calling `train()` on `mnist_classifier`. Add the following to `main()`: ```python # Train the model diff --git a/tensorflow/docs_src/tutorials/leftnav_files b/tensorflow/docs_src/tutorials/leftnav_files index 41ffdc86010fb8407889df26eefa5fa59952c5da..888052428f951fa1a7cbd9c6d35497a056387097 100644 --- a/tensorflow/docs_src/tutorials/leftnav_files +++ b/tensorflow/docs_src/tutorials/leftnav_files @@ -1,22 +1,22 @@ index.md ### Images -layers.md -image_recognition.md -image_retraining.md +layers.md: MNIST +image_recognition.md: Image Recognition +image_retraining.md: Image Retraining deep_cnn.md ### Sequences recurrent.md -seq2seq.md -recurrent_quickdraw.md +seq2seq.md: Neural Machine Translation +recurrent_quickdraw.md: Drawing Classification audio_recognition.md ### Data Representation -wide.md -wide_and_deep.md +wide.md: Linear Models +wide_and_deep.md: Wide & Deep Learning word2vec.md -kernel_methods.md +kernel_methods.md: Kernel Methods ### Non-ML mandelbrot.md diff --git a/tensorflow/docs_src/tutorials/recurrent_quickdraw.md b/tensorflow/docs_src/tutorials/recurrent_quickdraw.md index e22536adb6f0b893602ff79612cfb01e10586a18..7584a76ba5d3a044377e7050993b6ad778633433 100644 --- a/tensorflow/docs_src/tutorials/recurrent_quickdraw.md +++ b/tensorflow/docs_src/tutorials/recurrent_quickdraw.md @@ -109,7 +109,8 @@ This download will take a while and download a bit more than 23GB of data. To convert the `ndjson` files to @{$python/python_io#tfrecords_format_details$TFRecord} files containing -${tf.train.Example} protos run the following command. +[`tf.train.Example`](https://www.tensorflow.org/code/tensorflow/core/example/example.proto) +protos run the following command. ```shell python create_dataset.py --ndjson_path rnn_tutorial_data \ diff --git a/tensorflow/docs_src/tutorials/wide.md b/tensorflow/docs_src/tutorials/wide.md index dba6f54c52ca5bf2569c66ad055329708de3991c..27ce75a30dd2acd5925702611042270e767b0c73 100644 --- a/tensorflow/docs_src/tutorials/wide.md +++ b/tensorflow/docs_src/tutorials/wide.md @@ -74,15 +74,15 @@ Here's a list of columns available in the Census Income dataset: | relationship | Categorical | Wife, Own-child, Husband, | : : : Not-in-family, Other-relative, : : : : Unmarried. : -| race | Categorical | White, Asian-Pac-Islander, | -: : : Amer-Indian-Eskimo, Other, Black. : +| race | Categorical | Amer-Indian-Eskimo, Asian-Pac- | +: : : Islander, Black, White, Other. : | gender | Categorical | Female, Male. | | capital_gain | Continuous | Capital gains recorded. | | capital_loss | Continuous | Capital Losses recorded. | | hours_per_week | Continuous | Hours worked per week. | | native_country | Categorical | Country of origin of the | : : : individual. : -| income | Categorical | ">50K" or "<=50K", meaning | +| income_bracket | Categorical | ">50K" or "<=50K", meaning | : : : whether the person makes more : : : : than $50,000 annually. : @@ -247,7 +247,7 @@ hours_per_week = tf.feature_column.numeric_column('hours_per_week') ### Making Continuous Features Categorical through Bucketization Sometimes the relationship between a continuous feature and the label is not -linear. As an hypothetical example, a person's income may grow with age in the +linear. As a hypothetical example, a person's income may grow with age in the early stage of one's career, then the growth may slow at some point, and finally the income decreases after retirement. In this scenario, using the raw `age` as a real-valued feature column might not be a good choice because the model can @@ -361,6 +361,16 @@ The first line of the final output should be something like `accuracy: 0.83557522`, which means the accuracy is 83.6%. Feel free to try more features and transformations and see if you can do even better! +After the model is evaluated, we can use the model to predict whether an individual has an annual income of over +50,000 dollars given an individual's information input. +```python + pred_iter = model.predict(input_fn=lambda: input_fn(FLAGS.test_data, 1, False, 1)) + for pred in pred_iter: + print(pred['classes']) +``` + +The model prediction output would be like `[b'1']` or `[b'0']` which means whether corresponding individual has an annual income of over 50,000 dollars or not. + If you'd like to see a working end-to-end example, you can download our [example code](https://github.com/tensorflow/models/tree/master/official/wide_deep/wide_deep.py) and set the `model_type` flag to `wide`. diff --git a/tensorflow/examples/android/AndroidManifest.xml b/tensorflow/examples/android/AndroidManifest.xml index bb75431a1f8bab2951299520903aa6e043f8415e..5c47ce6b673e4c9d635b867c1ccdc679f67c6ae5 100644 --- a/tensorflow/examples/android/AndroidManifest.xml +++ b/tensorflow/examples/android/AndroidManifest.xml @@ -40,6 +40,7 @@ + @@ -49,6 +50,7 @@ + @@ -58,6 +60,7 @@ + @@ -67,6 +70,7 @@ + diff --git a/tensorflow/examples/android/build.gradle b/tensorflow/examples/android/build.gradle index f7bdf8b816a8191770bc1ad59b890041b8e39912..0767726aa9a248fb073fbd4114f47d1b4ed6901b 100644 --- a/tensorflow/examples/android/build.gradle +++ b/tensorflow/examples/android/build.gradle @@ -56,10 +56,12 @@ def nativeOutDir = 'libs/' + cpuType def nativeBuildRule = 'buildNativeBazel' def demoLibPath = '../../../bazel-bin/tensorflow/examples/android/libtensorflow_demo.so' def inferenceLibPath = '../../../bazel-bin/tensorflow/contrib/android/libtensorflow_inference.so' + +// Override for Makefile builds. if (nativeBuildSystem == 'makefile') { nativeBuildRule = 'buildNativeMake' - demoLibPath = '../../../tensorflow/contrib/makefile/gen/lib/libtensorflow_demo.so' - inferenceLibPath = '../../../tensorflow/contrib/makefile/gen/lib/libtensorflow_inference.so' + demoLibPath = '../../../tensorflow/contrib/makefile/gen/lib/android_' + cpuType + '/libtensorflow_demo.so' + inferenceLibPath = '../../../tensorflow/contrib/makefile/gen/lib/android_' + cpuType + '/libtensorflow_inference.so' } // If building with Bazel, this is the location of the bazel binary. @@ -154,7 +156,8 @@ task buildNativeMake(type: Exec) { '-s', \ 'tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in', \ '-t', \ - 'libtensorflow_inference.so libtensorflow_demo.so' \ + 'libtensorflow_inference.so libtensorflow_demo.so all' \ + , '-a', cpuType \ //, '-T' // Uncomment to skip protobuf and speed up subsequent builds. } diff --git a/tensorflow/examples/android/res/animator/color_animation.xml b/tensorflow/examples/android/res/animator/color_animation.xml new file mode 100644 index 0000000000000000000000000000000000000000..891d8cc1d4f3e59d0371030fd763c5ad468e7887 --- /dev/null +++ b/tensorflow/examples/android/res/animator/color_animation.xml @@ -0,0 +1,30 @@ + + + + + diff --git a/tensorflow/examples/android/src/org/tensorflow/demo/CameraActivity.java b/tensorflow/examples/android/src/org/tensorflow/demo/CameraActivity.java index 8bd4abb154a8f8c74f2195d4acbb99d3d5d498ea..429138abe5338e63d602ef6005e7607d21e1e357 100644 --- a/tensorflow/examples/android/src/org/tensorflow/demo/CameraActivity.java +++ b/tensorflow/examples/android/src/org/tensorflow/demo/CameraActivity.java @@ -351,6 +351,10 @@ public abstract class CameraActivity extends Activity protected void setFragment() { String cameraId = chooseCamera(); + if (cameraId == null) { + Toast.makeText(this, "No Camera Detected", Toast.LENGTH_SHORT).show(); + finish(); + } Fragment fragment; if (useCamera2API) { @@ -416,7 +420,8 @@ public abstract class CameraActivity extends Activity @Override public boolean onKeyDown(final int keyCode, final KeyEvent event) { - if (keyCode == KeyEvent.KEYCODE_VOLUME_DOWN || keyCode == KeyEvent.KEYCODE_VOLUME_UP) { + if (keyCode == KeyEvent.KEYCODE_VOLUME_DOWN || keyCode == KeyEvent.KEYCODE_VOLUME_UP + || keyCode == KeyEvent.KEYCODE_BUTTON_L1 || keyCode == KeyEvent.KEYCODE_DPAD_CENTER) { debug = !debug; requestRender(); onSetDebug(debug); diff --git a/tensorflow/examples/android/src/org/tensorflow/demo/LegacyCameraConnectionFragment.java b/tensorflow/examples/android/src/org/tensorflow/demo/LegacyCameraConnectionFragment.java index bc0c738e53739a516bef268e6551cbb7741a6298..068c7b0d945669b8207097e81c03ade07bc7ca73 100644 --- a/tensorflow/examples/android/src/org/tensorflow/demo/LegacyCameraConnectionFragment.java +++ b/tensorflow/examples/android/src/org/tensorflow/demo/LegacyCameraConnectionFragment.java @@ -82,7 +82,7 @@ public class LegacyCameraConnectionFragment extends Fragment { try { Camera.Parameters parameters = camera.getParameters(); List focusModes = parameters.getSupportedFocusModes(); - if (focusModes != null + if (focusModes != null && focusModes.contains(Camera.Parameters.FOCUS_MODE_CONTINUOUS_PICTURE)) { parameters.setFocusMode(Camera.Parameters.FOCUS_MODE_CONTINUOUS_PICTURE); } diff --git a/tensorflow/examples/android/src/org/tensorflow/demo/SpeechActivity.java b/tensorflow/examples/android/src/org/tensorflow/demo/SpeechActivity.java index 184df1bdb42802bfe50b15429f09baeb5600e34f..1cddf3dc5568babb8c08c690fad143299f5ccca5 100644 --- a/tensorflow/examples/android/src/org/tensorflow/demo/SpeechActivity.java +++ b/tensorflow/examples/android/src/org/tensorflow/demo/SpeechActivity.java @@ -31,7 +31,8 @@ the RecognizeCommands helper class. package org.tensorflow.demo; -import android.animation.ValueAnimator; +import android.animation.AnimatorInflater; +import android.animation.AnimatorSet; import android.app.Activity; import android.content.pm.PackageManager; import android.media.AudioFormat; @@ -329,17 +330,13 @@ public class SpeechActivity extends Activity { labelIndex = i; } } - final View labelView = (View) labelsListView.getChildAt(labelIndex - 2); - ValueAnimator colorAnimation = - ValueAnimator.ofArgb(0x00b3ccff, 0xffb3ccff, 0x00b3ccff); - colorAnimation.setDuration(750); - colorAnimation.addUpdateListener( - new ValueAnimator.AnimatorUpdateListener() { - @Override - public void onAnimationUpdate(ValueAnimator animator) { - labelView.setBackgroundColor((int) animator.getAnimatedValue()); - } - }); + final View labelView = labelsListView.getChildAt(labelIndex - 2); + + AnimatorSet colorAnimation = + (AnimatorSet) + AnimatorInflater.loadAnimator( + SpeechActivity.this, R.animator.color_animation); + colorAnimation.setTarget(labelView); colorAnimation.start(); } } diff --git a/tensorflow/examples/android/src/org/tensorflow/demo/StylizeActivity.java b/tensorflow/examples/android/src/org/tensorflow/demo/StylizeActivity.java index 6a66ec3927be62f1f996eb18bb6c04ea66f43152..33ec65e9f73a1d04bcafdc09d1618b32e03b1dc0 100644 --- a/tensorflow/examples/android/src/org/tensorflow/demo/StylizeActivity.java +++ b/tensorflow/examples/android/src/org/tensorflow/demo/StylizeActivity.java @@ -16,8 +16,10 @@ package org.tensorflow.demo; +import android.app.UiModeManager; import android.content.Context; import android.content.res.AssetManager; +import android.content.res.Configuration; import android.graphics.Bitmap; import android.graphics.Bitmap.Config; import android.graphics.BitmapFactory; @@ -31,9 +33,11 @@ import android.graphics.Typeface; import android.media.ImageReader.OnImageAvailableListener; import android.os.Bundle; import android.os.SystemClock; +import android.util.DisplayMetrics; import android.util.Size; import android.util.TypedValue; import android.view.Display; +import android.view.KeyEvent; import android.view.MotionEvent; import android.view.View; import android.view.View.OnClickListener; @@ -43,6 +47,7 @@ import android.widget.BaseAdapter; import android.widget.Button; import android.widget.GridView; import android.widget.ImageView; +import android.widget.RelativeLayout; import android.widget.Toast; import java.io.IOException; import java.io.InputStream; @@ -381,6 +386,27 @@ public class StylizeActivity extends CameraActivity implements OnImageAvailableL grid = (GridView) findViewById(R.id.grid_layout); grid.setAdapter(adapter); grid.setOnTouchListener(gridTouchAdapter); + + // Change UI on Android TV + UiModeManager uiModeManager = (UiModeManager) getSystemService(UI_MODE_SERVICE); + if (uiModeManager.getCurrentModeType() == Configuration.UI_MODE_TYPE_TELEVISION) { + DisplayMetrics displayMetrics = new DisplayMetrics(); + getWindowManager().getDefaultDisplay().getMetrics(displayMetrics); + int styleSelectorHeight = displayMetrics.heightPixels; + int styleSelectorWidth = displayMetrics.widthPixels - styleSelectorHeight; + RelativeLayout.LayoutParams layoutParams = new RelativeLayout.LayoutParams(styleSelectorWidth, ViewGroup.LayoutParams.MATCH_PARENT); + + // Calculate number of style in a row, so all the style can show up without scrolling + int numOfStylePerRow = 3; + while (styleSelectorWidth / numOfStylePerRow * Math.ceil((float) (adapter.getCount() - 2) / numOfStylePerRow) > styleSelectorHeight) { + numOfStylePerRow++; + } + grid.setNumColumns(numOfStylePerRow); + layoutParams.addRule(RelativeLayout.ALIGN_PARENT_RIGHT); + grid.setLayoutParams(layoutParams); + adapter.buttons.clear(); + } + setStyle(adapter.items[0], 1.0f); } @@ -602,4 +628,38 @@ public class StylizeActivity extends CameraActivity implements OnImageAvailableL borderedText.drawLines(canvas, 10, canvas.getHeight() - 10, lines); } + + @Override + public boolean onKeyDown(int keyCode, KeyEvent event) { + int moveOffset = 0; + switch (keyCode) { + case KeyEvent.KEYCODE_DPAD_LEFT: + moveOffset = -1; + break; + case KeyEvent.KEYCODE_DPAD_RIGHT: + moveOffset = 1; + break; + case KeyEvent.KEYCODE_DPAD_UP: + moveOffset = -1 * grid.getNumColumns(); + break; + case KeyEvent.KEYCODE_DPAD_DOWN: + moveOffset = grid.getNumColumns(); + break; + default: + return super.onKeyDown(keyCode, event); + } + + // get the highest selected style + int currentSelect = 0; + float highestValue = 0; + for (int i = 0; i < adapter.getCount(); i++) { + if (adapter.items[i].value > highestValue) { + currentSelect = i; + highestValue = adapter.items[i].value; + } + } + setStyle(adapter.items[(currentSelect + moveOffset + adapter.getCount()) % adapter.getCount()], 1); + + return true; + } } diff --git a/tensorflow/examples/android/src/org/tensorflow/demo/tracking/MultiBoxTracker.java b/tensorflow/examples/android/src/org/tensorflow/demo/tracking/MultiBoxTracker.java index 2fe2ba539edc84e80baf36b6d1ac1e192bc92163..af6af2bc8f508a70aa7e44a7236f0e7ea5e3d71c 100644 --- a/tensorflow/examples/android/src/org/tensorflow/demo/tracking/MultiBoxTracker.java +++ b/tensorflow/examples/android/src/org/tensorflow/demo/tracking/MultiBoxTracker.java @@ -199,7 +199,7 @@ public class MultiBoxTracker { final int w, final int h, final int rowStride, - final int sensorOrienation, + final int sensorOrientation, final byte[] frame, final long timestamp) { if (objectTracker == null && !initialized) { @@ -209,7 +209,7 @@ public class MultiBoxTracker { objectTracker = ObjectTracker.getInstance(w, h, rowStride, true); frameWidth = w; frameHeight = h; - this.sensorOrientation = sensorOrienation; + this.sensorOrientation = sensorOrientation; initialized = true; if (objectTracker == null) { diff --git a/tensorflow/examples/get_started/regression/imports85.py b/tensorflow/examples/get_started/regression/imports85.py index 6bee556eb887a643b3a81691324736427ecc2707..4fdaceea9afee74550196031fe590c3a2abd20ed 100644 --- a/tensorflow/examples/get_started/regression/imports85.py +++ b/tensorflow/examples/get_started/regression/imports85.py @@ -131,11 +131,12 @@ def dataset(y_name="price", train_fraction=0.7): # booleans but we are dealing with symbolic tensors. return ~in_training_set(line) - base_dataset = (tf.contrib.data - # Get the lines from the file. - .TextLineDataset(path) - # drop lines with question marks. - .filter(has_no_question_marks)) + base_dataset = ( + tf.data + # Get the lines from the file. + .TextLineDataset(path) + # drop lines with question marks. + .filter(has_no_question_marks)) train = (base_dataset # Take only the training-set lines. diff --git a/tensorflow/examples/how_tos/reading_data/fully_connected_reader.py b/tensorflow/examples/how_tos/reading_data/fully_connected_reader.py index 461fb1c5173f66278eb585d30bd8749a58fb6245..307eede5c03780e9244b035f020fc7846290d4d9 100644 --- a/tensorflow/examples/how_tos/reading_data/fully_connected_reader.py +++ b/tensorflow/examples/how_tos/reading_data/fully_connected_reader.py @@ -1,4 +1,4 @@ -# Copyright 2015 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. @@ -45,6 +45,7 @@ VALIDATION_FILE = 'validation.tfrecords' def decode(serialized_example): + """Parses an image and label from the given `serialized_example`.""" features = tf.parse_single_example( serialized_example, # Defaults are not specified since both keys are required. @@ -66,6 +67,7 @@ def decode(serialized_example): def augment(image, label): + """Placeholder for data augmentation.""" # OPTIONAL: Could reshape into a 28x28 image and apply distortions # here. Since we are not applying any distortions in this # example, and the next step expects the image to be flattened @@ -74,9 +76,8 @@ def augment(image, label): def normalize(image, label): - # Convert from [0, 255] -> [-0.5, 0.5] floats. + """Convert `image` from [0, 255] -> [-0.5, 0.5] floats.""" image = tf.cast(image, tf.float32) * (1. / 255) - 0.5 - return image, label @@ -106,18 +107,23 @@ def inputs(train, batch_size, num_epochs): if train else VALIDATION_FILE) with tf.name_scope('input'): - # TFRecordDataset opens a protobuf and reads entries line by line - # could also be [list, of, filenames] + # TFRecordDataset opens a binary file and reads one record at a time. + # `filename` could also be a list of filenames, which will be read in order. dataset = tf.data.TFRecordDataset(filename) - dataset = dataset.repeat(num_epochs) - # map takes a python function and applies it to every sample + # The map transformation takes a function and applies it to every element + # of the dataset. dataset = dataset.map(decode) dataset = dataset.map(augment) dataset = dataset.map(normalize) - #the parameter is the queue size + # The shuffle transformation uses a finite-sized buffer to shuffle elements + # in memory. The parameter is the number of elements in the buffer. For + # completely uniform shuffling, set the parameter to be the same as the + # number of elements in the dataset. dataset = dataset.shuffle(1000 + 3 * batch_size) + + dataset = dataset.repeat(num_epochs) dataset = dataset.batch(batch_size) iterator = dataset.make_one_shot_iterator() @@ -153,7 +159,7 @@ def run_training(): sess.run(init_op) try: step = 0 - while True: #train until OutOfRangeError + while True: # Train until OutOfRangeError start_time = time.time() # Run one step of the model. The return values are diff --git a/tensorflow/examples/image_retraining/retrain.py b/tensorflow/examples/image_retraining/retrain.py index ec22684eaf63700c608c6ce45f22941555246b99..99a71206acbd533ec8bc5a9644435eacad564cd4 100644 --- a/tensorflow/examples/image_retraining/retrain.py +++ b/tensorflow/examples/image_retraining/retrain.py @@ -41,7 +41,6 @@ The subfolder names are important, since they define what label is applied to each image, but the filenames themselves don't matter. Once your images are prepared, you can run the training with a command like this: - ```bash bazel build tensorflow/examples/image_retraining:retrain && \ bazel-bin/tensorflow/examples/image_retraining/retrain \ @@ -70,17 +69,22 @@ on resource-limited platforms, you can try the `--architecture` flag with a Mobilenet model. For example: Run floating-point version of mobilenet: + ```bash python tensorflow/examples/image_retraining/retrain.py \ --image_dir ~/flower_photos --architecture mobilenet_1.0_224 ``` -Run quantized version of mobilenet: +Run mobilenet, instrumented for quantization: + ```bash python tensorflow/examples/image_retraining/retrain.py \ - --image_dir ~/flower_photos/ --architecture mobilenet_1.0_224_quantized + --image_dir ~/flower_photos/ --architecture mobilenet_1.0_224_quant ``` +These instrumented models can be converted to fully quantized mobile models via +TensorFlow Lite. + There are 32 different Mobilenet models to choose from, with a variety of file size and latency options. The first number can be '1.0', '0.75', '0.50', or '0.25' to control the size, and the second controls the input image size, either @@ -96,6 +100,12 @@ Visualize the summaries with this command: tensorboard --logdir /tmp/retrain_logs +To use with Tensorflow Serving: + +```bash +tensorflow_model_server --port=9000 --model_name=inception \ + --model_base_path=/tmp/saved_models/ +``` """ from __future__ import absolute_import from __future__ import division @@ -114,7 +124,6 @@ import numpy as np from six.moves import urllib import tensorflow as tf -from tensorflow.contrib.quantize.python import quant_ops from tensorflow.python.framework import graph_util from tensorflow.python.framework import tensor_shape from tensorflow.python.platform import gfile @@ -128,6 +137,9 @@ FLAGS = None # need to update these to reflect the values in the network you're using. MAX_NUM_IMAGES_PER_CLASS = 2 ** 27 - 1 # ~134M +# The location where variable checkpoints will be stored. +CHECKPOINT_NAME = '/tmp/_retrain_checkpoint' + def create_image_lists(image_dir, testing_percentage, validation_percentage): """Builds a list of training images from the file system. @@ -344,8 +356,8 @@ def maybe_download_and_extract(data_url): filepath, _ = urllib.request.urlretrieve(data_url, filepath, _progress) print() statinfo = os.stat(filepath) - tf.logging.info('Successfully downloaded', filename, statinfo.st_size, - 'bytes.') + tf.logging.info('Successfully downloaded %s %d bytes.', filename, + statinfo.st_size) print('Extracting file from ', filepath) tarfile.open(filepath, 'r:gz').extractall(dest_directory) else: @@ -738,9 +750,9 @@ def variable_summaries(var): tf.summary.histogram('histogram', var) -def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, - bottleneck_tensor_size, quantize_layer): - """Adds a new softmax and fully-connected layer for training. +def add_final_retrain_ops(class_count, final_tensor_name, bottleneck_tensor, + bottleneck_tensor_size, quantize_layer, is_training): + """Adds a new softmax and fully-connected layer for training and eval. We need to retrain the top layer to identify our new classes, so this function adds the right operations to the graph, along with some variables to hold the @@ -756,7 +768,9 @@ def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, bottleneck_tensor: The output of the main CNN graph. bottleneck_tensor_size: How many entries in the bottleneck vector. quantize_layer: Boolean, specifying whether the newly added layer should be - quantized. + instrumented for quantized. + is_training: Boolean, specifying whether the newly add layer is for training + or eval. Returns: The tensors for the training and cross entropy results, and tensors for the @@ -771,50 +785,41 @@ def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, ground_truth_input = tf.placeholder( tf.int64, [None], name='GroundTruthInput') - # Organizing the following ops as `final_training_ops` so they're easier - # to see in TensorBoard - layer_name = 'final_training_ops' + # Organizing the following ops so they are easier to see in TensorBoard. + layer_name = 'final_retrain_ops' with tf.name_scope(layer_name): with tf.name_scope('weights'): initial_value = tf.truncated_normal( [bottleneck_tensor_size, class_count], stddev=0.001) layer_weights = tf.Variable(initial_value, name='final_weights') - if quantize_layer: - quantized_layer_weights = quant_ops.MovingAvgQuantize( - layer_weights, is_training=True) - variable_summaries(quantized_layer_weights) - variable_summaries(layer_weights) + with tf.name_scope('biases'): layer_biases = tf.Variable(tf.zeros([class_count]), name='final_biases') - if quantize_layer: - quantized_layer_biases = quant_ops.MovingAvgQuantize( - layer_biases, is_training=True) - variable_summaries(quantized_layer_biases) - variable_summaries(layer_biases) with tf.name_scope('Wx_plus_b'): - if quantize_layer: - logits = tf.matmul(bottleneck_input, - quantized_layer_weights) + quantized_layer_biases - logits = quant_ops.MovingAvgQuantize( - logits, - init_min=-32.0, - init_max=32.0, - is_training=True, - num_bits=8, - narrow_range=False, - ema_decay=0.5) - tf.summary.histogram('pre_activations', logits) - else: - logits = tf.matmul(bottleneck_input, layer_weights) + layer_biases - tf.summary.histogram('pre_activations', logits) + logits = tf.matmul(bottleneck_input, layer_weights) + layer_biases + tf.summary.histogram('pre_activations', logits) final_tensor = tf.nn.softmax(logits, name=final_tensor_name) + # The tf.contrib.quantize functions rewrite the graph in place for + # quantization. The imported model graph has already been rewritten, so upon + # calling these rewrites, only the newly added final layer will be + # transformed. + if quantize_layer: + if is_training: + tf.contrib.quantize.create_training_graph() + else: + tf.contrib.quantize.create_eval_graph() + tf.summary.histogram('activations', final_tensor) + # If this is an eval graph, we don't need to add loss ops or an optimizer. + if not is_training: + return None, None, bottleneck_input, ground_truth_input, final_tensor + with tf.name_scope('cross_entropy'): cross_entropy_mean = tf.losses.sparse_softmax_cross_entropy( labels=ground_truth_input, logits=logits) @@ -850,13 +855,91 @@ def add_evaluation_step(result_tensor, ground_truth_tensor): return evaluation_step, prediction -def save_graph_to_file(sess, graph, graph_file_name): +def run_final_eval(sess, model_info, class_count, image_lists, jpeg_data_tensor, + decoded_image_tensor, resized_image_tensor, + bottleneck_tensor): + """Runs a final evaluation on an eval graph using the test data set. + + Args: + sess: Session for the train graph. + model_info: Model info dictionary from create_model_info() + class_count: Number of classes + image_lists: Dictionary of training images for each label. + jpeg_data_tensor: The layer to feed jpeg image data into. + decoded_image_tensor: The output of decoding and resizing the image. + resized_image_tensor: The input node of the recognition graph. + bottleneck_tensor: The bottleneck output layer of the CNN graph. + """ + (sess, bottleneck_input, ground_truth_input, evaluation_step, + prediction) = build_eval_session(model_info, class_count) + + test_bottlenecks, test_ground_truth, test_filenames = ( + get_random_cached_bottlenecks(sess, image_lists, FLAGS.test_batch_size, + 'testing', FLAGS.bottleneck_dir, + FLAGS.image_dir, jpeg_data_tensor, + decoded_image_tensor, resized_image_tensor, + bottleneck_tensor, FLAGS.architecture)) + test_accuracy, predictions = sess.run( + [evaluation_step, prediction], + feed_dict={ + bottleneck_input: test_bottlenecks, + ground_truth_input: test_ground_truth + }) + tf.logging.info('Final test accuracy = %.1f%% (N=%d)' % + (test_accuracy * 100, len(test_bottlenecks))) + + if FLAGS.print_misclassified_test_images: + tf.logging.info('=== MISCLASSIFIED TEST IMAGES ===') + for i, test_filename in enumerate(test_filenames): + if predictions[i] != test_ground_truth[i]: + tf.logging.info('%70s %s' % (test_filename, + list(image_lists.keys())[predictions[i]])) + + +def build_eval_session(model_info, class_count): + """Builds an restored eval session without train operations for exporting. + + Args: + model_info: Model info dictionary from create_model_info() + class_count: Number of classes + + Returns: + Eval session containing the restored eval graph. + The bottleneck input, ground truth, eval step, and prediction tensors. + """ + # If quantized, we need to create the correct eval graph for exporting. + eval_graph, bottleneck_tensor, _ = create_model_graph(model_info) + + eval_sess = tf.Session(graph=eval_graph) + with eval_graph.as_default(): + # Add the new layer for exporting. + (_, _, bottleneck_input, + ground_truth_input, final_tensor) = add_final_retrain_ops( + class_count, FLAGS.final_tensor_name, bottleneck_tensor, + model_info['bottleneck_tensor_size'], model_info['quantize_layer'], + False) + + # Now we need to restore the values from the training graph to the eval + # graph. + tf.train.Saver().restore(eval_sess, CHECKPOINT_NAME) + + evaluation_step, prediction = add_evaluation_step(final_tensor, + ground_truth_input) + + return (eval_sess, bottleneck_input, ground_truth_input, evaluation_step, + prediction) + + +def save_graph_to_file(graph, graph_file_name, model_info, class_count): + """Saves an graph to file, creating a valid quantized one if necessary.""" + sess, _, _, _, _ = build_eval_session(model_info, class_count) + graph = sess.graph + output_graph_def = graph_util.convert_variables_to_constants( sess, graph.as_graph_def(), [FLAGS.final_tensor_name]) with gfile.FastGFile(graph_file_name, 'wb') as f: f.write(output_graph_def.SerializeToString()) - return def prepare_file_system(): @@ -909,11 +992,10 @@ def create_model_info(architecture): return None version_string = parts[1] if (version_string != '1.0' and version_string != '0.75' and - version_string != '0.50' and version_string != '0.25'): + version_string != '0.5' and version_string != '0.25'): tf.logging.error( - """"The Mobilenet version should be '1.0', '0.75', '0.50', or '0.25', - but found '%s' for architecture '%s'""", - version_string, architecture) + """"The Mobilenet version should be '1.0', '0.75', '0.5', or '0.25', + but found '%s' for architecture '%s'""", version_string, architecture) return None size_string = parts[2] if (size_string != '224' and size_string != '192' and @@ -926,35 +1008,26 @@ def create_model_info(architecture): if len(parts) == 3: is_quantized = False else: - if parts[3] != 'quantized': + if parts[3] != 'quant': tf.logging.error( "Couldn't understand architecture suffix '%s' for '%s'", parts[3], architecture) return None is_quantized = True + data_url = 'http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/' + model_name = 'mobilenet_v1_' + version_string + '_' + size_string if is_quantized: - data_url = 'http://download.tensorflow.org/models/mobilenet_v1_' - data_url += version_string + '_' + size_string + '_quantized_frozen.tgz' - bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0' - resized_input_tensor_name = 'Placeholder:0' - model_dir_name = ('mobilenet_v1_' + version_string + '_' + size_string + - '_quantized_frozen') - model_base_name = 'quantized_frozen_graph.pb' - - else: - data_url = 'http://download.tensorflow.org/models/mobilenet_v1_' - data_url += version_string + '_' + size_string + '_frozen.tgz' - bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0' - resized_input_tensor_name = 'input:0' - model_dir_name = 'mobilenet_v1_' + version_string + '_' + size_string - model_base_name = 'frozen_graph.pb' + model_name += '_quant' + data_url += model_name + '.tgz' + bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0' + resized_input_tensor_name = 'input:0' + model_file_name = model_name + '_frozen.pb' bottleneck_tensor_size = 1001 input_width = int(size_string) input_height = int(size_string) input_depth = 3 - model_file_name = os.path.join(model_dir_name, model_base_name) input_mean = 127.5 input_std = 127.5 else: @@ -1004,6 +1077,47 @@ def add_jpeg_decoding(input_width, input_height, input_depth, input_mean, return jpeg_data, mul_image +def export_model(model_info, class_count, saved_model_dir): + """Exports model for serving. + + Args: + model_info: The modelinfo for the current model. + class_count: The number of classes. + saved_model_dir: Directory in which to save exported model and variables. + """ + # The SavedModel should hold the eval graph. + sess, _, _, _, _ = build_eval_session(model_info, class_count) + graph = sess.graph + with graph.as_default(): + input_tensor = model_info['resized_input_tensor_name'] + in_image = sess.graph.get_tensor_by_name(input_tensor) + inputs = {'image': tf.saved_model.utils.build_tensor_info(in_image)} + + out_classes = sess.graph.get_tensor_by_name('final_result:0') + outputs = { + 'prediction': tf.saved_model.utils.build_tensor_info(out_classes) + } + + signature = tf.saved_model.signature_def_utils.build_signature_def( + inputs=inputs, + outputs=outputs, + method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME) + + legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op') + + # Save out the SavedModel. + builder = tf.saved_model.builder.SavedModelBuilder(saved_model_dir) + builder.add_meta_graph_and_variables( + sess, [tf.saved_model.tag_constants.SERVING], + signature_def_map={ + tf.saved_model.signature_constants. + DEFAULT_SERVING_SIGNATURE_DEF_KEY: + signature + }, + legacy_init_op=legacy_init_op) + builder.save() + + def main(_): # Needed to make sure the logging output is visible. # See https://github.com/tensorflow/tensorflow/issues/3047 @@ -1018,11 +1132,6 @@ def main(_): tf.logging.error('Did not recognize architecture flag') return -1 - # Set up the pre-trained graph. - maybe_download_and_extract(model_info['data_url']) - graph, bottleneck_tensor, resized_image_tensor = ( - create_model_graph(model_info)) - # Look at the folder structure, and create lists of all the images. image_lists = create_image_lists(FLAGS.image_dir, FLAGS.testing_percentage, FLAGS.validation_percentage) @@ -1041,6 +1150,19 @@ def main(_): FLAGS.flip_left_right, FLAGS.random_crop, FLAGS.random_scale, FLAGS.random_brightness) + # Set up the pre-trained graph. + maybe_download_and_extract(model_info['data_url']) + graph, bottleneck_tensor, resized_image_tensor = ( + create_model_graph(model_info)) + + # Add the new layer that we'll be training. + with graph.as_default(): + (train_step, cross_entropy, bottleneck_input, + ground_truth_input, final_tensor) = add_final_retrain_ops( + class_count, FLAGS.final_tensor_name, bottleneck_tensor, + model_info['bottleneck_tensor_size'], model_info['quantize_layer'], + True) + with tf.Session(graph=graph) as sess: # Set up the image decoding sub-graph. jpeg_data_tensor, decoded_image_tensor = add_jpeg_decoding( @@ -1064,15 +1186,8 @@ def main(_): decoded_image_tensor, resized_image_tensor, bottleneck_tensor, FLAGS.architecture) - # Add the new layer that we'll be training. - (train_step, cross_entropy, bottleneck_input, ground_truth_input, - final_tensor) = add_final_training_ops( - len(image_lists.keys()), FLAGS.final_tensor_name, bottleneck_tensor, - model_info['bottleneck_tensor_size'], model_info['quantize_layer']) - # Create the operations we need to evaluate the accuracy of our new layer. - evaluation_step, prediction = add_evaluation_step( - final_tensor, ground_truth_input) + evaluation_step, _ = add_evaluation_step(final_tensor, ground_truth_input) # Merge all the summaries and write them out to the summaries_dir merged = tf.summary.merge_all() @@ -1082,6 +1197,10 @@ def main(_): validation_writer = tf.summary.FileWriter( FLAGS.summaries_dir + '/validation') + # Create a train saver that is used to restore values into an eval graph + # when exporting models. + train_saver = tf.train.Saver() + # Set up all our weights to their initial default values. init = tf.global_variables_initializer() sess.run(init) @@ -1122,6 +1241,9 @@ def main(_): (datetime.now(), i, train_accuracy * 100)) tf.logging.info('%s: Step %d: Cross entropy = %f' % (datetime.now(), i, cross_entropy_value)) + # TODO(suharshs): Make this use an eval graph, to avoid quantization + # moving averages being updated by the validation set, though in + # practice this makes a negligable difference. validation_bottlenecks, validation_ground_truth, _ = ( get_random_cached_bottlenecks( sess, image_lists, FLAGS.validation_batch_size, 'validation', @@ -1144,41 +1266,33 @@ def main(_): if (intermediate_frequency > 0 and (i % intermediate_frequency == 0) and i > 0): + # If we want to do an intermediate save, save a checkpoint of the train + # graph, to restore into the eval graph. + train_saver.save(sess, CHECKPOINT_NAME) intermediate_file_name = (FLAGS.intermediate_output_graphs_dir + 'intermediate_' + str(i) + '.pb') tf.logging.info('Save intermediate result to : ' + intermediate_file_name) - save_graph_to_file(sess, graph, intermediate_file_name) + save_graph_to_file(graph, intermediate_file_name, model_info, + class_count) + + # After training is complete, force one last save of the train checkpoint. + train_saver.save(sess, CHECKPOINT_NAME) # We've completed all our training, so run a final test evaluation on # some new images we haven't used before. - test_bottlenecks, test_ground_truth, test_filenames = ( - get_random_cached_bottlenecks( - sess, image_lists, FLAGS.test_batch_size, 'testing', - FLAGS.bottleneck_dir, FLAGS.image_dir, jpeg_data_tensor, - decoded_image_tensor, resized_image_tensor, bottleneck_tensor, - FLAGS.architecture)) - test_accuracy, predictions = sess.run( - [evaluation_step, prediction], - feed_dict={bottleneck_input: test_bottlenecks, - ground_truth_input: test_ground_truth}) - tf.logging.info('Final test accuracy = %.1f%% (N=%d)' % - (test_accuracy * 100, len(test_bottlenecks))) - - if FLAGS.print_misclassified_test_images: - tf.logging.info('=== MISCLASSIFIED TEST IMAGES ===') - for i, test_filename in enumerate(test_filenames): - if predictions[i] != test_ground_truth[i]: - tf.logging.info('%70s %s' % - (test_filename, - list(image_lists.keys())[predictions[i]])) + run_final_eval(sess, model_info, class_count, image_lists, jpeg_data_tensor, + decoded_image_tensor, resized_image_tensor, + bottleneck_tensor) # Write out the trained graph and labels with the weights stored as # constants. - save_graph_to_file(sess, graph, FLAGS.output_graph) + save_graph_to_file(graph, FLAGS.output_graph, model_info, class_count) with gfile.FastGFile(FLAGS.output_labels, 'w') as f: f.write('\n'.join(image_lists.keys()) + '\n') + export_model(model_info, class_count, FLAGS.saved_model_dir) + if __name__ == '__main__': parser = argparse.ArgumentParser() @@ -1358,9 +1472,15 @@ if __name__ == '__main__': form 'mobilenet__[_quantized]'. For example, 'mobilenet_1.0_224' will pick a model that is 17 MB in size and takes 224 pixel input images, while 'mobilenet_0.25_128_quantized' will choose a much - less accurate, but smaller and faster network that's 920 KB on disk and - takes 128x128 images. See https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html + smaller and less accurate model, taking 128x128 images, and instrumented + for eventual quantization via TensorFlow Lite. + See https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html for more information on Mobilenet.\ """) + parser.add_argument( + '--saved_model_dir', + type=str, + default='/tmp/saved_models/1/', + help='Where to save the exported graph.') FLAGS, unparsed = parser.parse_known_args() tf.app.run(main=main, argv=[sys.argv[0]] + unparsed) diff --git a/tensorflow/examples/image_retraining/retrain_test.py b/tensorflow/examples/image_retraining/retrain_test.py index 8b8dd45fd72e3d29bdb7f6291cc53b912adf3644..fb7324c58ac1be60baad840207f31a61ec6182be 100644 --- a/tensorflow/examples/image_retraining/retrain_test.py +++ b/tensorflow/examples/image_retraining/retrain_test.py @@ -67,22 +67,52 @@ class ImageRetrainingTest(test_util.TensorFlowTestCase): self.assertIsNotNone(sess.graph.get_tensor_by_name('DistortResult:0')) @tf.test.mock.patch.object(retrain, 'FLAGS', learning_rate=0.01) - def testAddFinalTrainingOps(self, flags_mock): + def testAddFinalRetrainOps(self, flags_mock): with tf.Graph().as_default(): with tf.Session() as sess: bottleneck = tf.placeholder(tf.float32, [1, 1024], name='bottleneck') - # Test creating final training op with quantization - retrain.add_final_training_ops(5, 'final', bottleneck, 1024, False) + # Test creating final training op with quantization. + retrain.add_final_retrain_ops(5, 'final', bottleneck, 1024, False, + False) self.assertIsNotNone(sess.graph.get_tensor_by_name('final:0')) @tf.test.mock.patch.object(retrain, 'FLAGS', learning_rate=0.01) - def testAddFinalTrainingOpsQuantized(self, flags_mock): - with tf.Graph().as_default(): + def testAddFinalRetrainOpsQuantized(self, flags_mock): + # Ensure that the training and eval graph for quantized models are correctly + # created. + with tf.Graph().as_default() as g: + with tf.Session() as sess: + bottleneck = tf.placeholder(tf.float32, [1, 1024], name='bottleneck') + # Test creating final training op with quantization, set is_training to + # true. + retrain.add_final_retrain_ops(5, 'final', bottleneck, 1024, True, True) + self.assertIsNotNone(sess.graph.get_tensor_by_name('final:0')) + found_fake_quant = 0 + for op in g.get_operations(): + if op.type == 'FakeQuantWithMinMaxVars': + found_fake_quant += 1 + # Ensure that the inputs of each FakeQuant operations has 2 Assign + # operations in the training graph (Assign[Min,Max]Last, + # Assign[Min,Max]Ema) + self.assertEqual(2, + len([i for i in op.inputs if 'Assign' in i.name])) + self.assertEqual(found_fake_quant, 2) + with tf.Graph().as_default() as g: with tf.Session() as sess: bottleneck = tf.placeholder(tf.float32, [1, 1024], name='bottleneck') - # Test creating final training op with quantization - retrain.add_final_training_ops(5, 'final', bottleneck, 1024, True) + # Test creating final training op with quantization, set is_training to + # false. + retrain.add_final_retrain_ops(5, 'final', bottleneck, 1024, True, False) self.assertIsNotNone(sess.graph.get_tensor_by_name('final:0')) + found_fake_quant = 0 + for op in g.get_operations(): + if op.type == 'FakeQuantWithMinMaxVars': + found_fake_quant += 1 + for i in op.inputs: + # Ensure that no operations are Assign operation since this is the + # evaluation graph. + self.assertTrue('Assign' not in i.name) + self.assertEqual(found_fake_quant, 2) def testAddEvaluationStep(self): with tf.Graph().as_default(): diff --git a/tensorflow/examples/ios/README.md b/tensorflow/examples/ios/README.md index 5bdaeb43ce143e36e78cfe301fd9b59e8b85b034..5d7bd36837b2a2c33ab4bc311a582c174666dcd5 100644 --- a/tensorflow/examples/ios/README.md +++ b/tensorflow/examples/ios/README.md @@ -119,11 +119,13 @@ rundown: `tensorflow/contrib/makefile/gen/lib` to the Library Search Paths setting. - You'll also need to add `libprotobuf.a` and `libprotobuf-lite.a` from - `tensorflow/contrib/makefile/gen/protobuf_ios/lib` to your _Build Stages_ and - _Library Search Paths_. + `tensorflow/contrib/makefile/gen/protobuf_ios/lib` + and `nsync.a` from `tensorflow/contrib/makefile/downloads/nsync/builds/lipo.ios.c++11` + to your _Build Stages_ and _Library Search Paths_. - The _Header Search_ paths needs to contain: - the root folder of tensorflow, + - `tensorflow/contrib/makefile/downloads/nsync/public` - `tensorflow/contrib/makefile/downloads/protobuf/src` - `tensorflow/contrib/makefile/downloads`, - `tensorflow/contrib/makefile/downloads/eigen`, and diff --git a/tensorflow/examples/label_image/label_image.py b/tensorflow/examples/label_image/label_image.py index 1c1bd57d715ae39539d9cec1347caa702a8261b9..fe5e0fc684abce08d3d7b7f3fa22bb5ba701c64a 100644 --- a/tensorflow/examples/label_image/label_image.py +++ b/tensorflow/examples/label_image/label_image.py @@ -18,7 +18,6 @@ from __future__ import division from __future__ import print_function import argparse -import sys import numpy as np import tensorflow as tf diff --git a/tensorflow/examples/learn/resnet.py b/tensorflow/examples/learn/resnet.py index 9542e552504580a6614f8bd2f43c38dfa795750f..c00de932a8707ad5717aaf1251cf5c88464a28b0 100755 --- a/tensorflow/examples/learn/resnet.py +++ b/tensorflow/examples/learn/resnet.py @@ -53,6 +53,8 @@ def res_net_model(features, labels, mode): ndim = int(sqrt(input_shape[1])) x = tf.reshape(x, [-1, ndim, ndim, 1]) + training = (mode == tf.estimator.ModeKeys.TRAIN) + # First convolution expands to 64 channels with tf.variable_scope('conv_layer1'): net = tf.layers.conv2d( @@ -60,7 +62,7 @@ def res_net_model(features, labels, mode): filters=64, kernel_size=7, activation=tf.nn.relu) - net = tf.layers.batch_normalization(net) + net = tf.layers.batch_normalization(net, training=training) # Max pool net = tf.layers.max_pooling2d( @@ -88,7 +90,7 @@ def res_net_model(features, labels, mode): kernel_size=1, padding='valid', activation=tf.nn.relu) - conv = tf.layers.batch_normalization(conv) + conv = tf.layers.batch_normalization(conv, training=training) with tf.variable_scope(name + '/conv_bottleneck'): conv = tf.layers.conv2d( @@ -97,7 +99,7 @@ def res_net_model(features, labels, mode): kernel_size=3, padding='same', activation=tf.nn.relu) - conv = tf.layers.batch_normalization(conv) + conv = tf.layers.batch_normalization(conv, training=training) # 1x1 convolution responsible for restoring dimension with tf.variable_scope(name + '/conv_out'): @@ -108,7 +110,7 @@ def res_net_model(features, labels, mode): kernel_size=1, padding='valid', activation=tf.nn.relu) - conv = tf.layers.batch_normalization(conv) + conv = tf.layers.batch_normalization(conv, training=training) # shortcut connections that turn the network into its counterpart # residual function (identity shortcut) @@ -154,7 +156,7 @@ def res_net_model(features, labels, mode): loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits) # Create training op. - if mode == tf.estimator.ModeKeys.TRAIN: + if training: optimizer = tf.train.AdagradOptimizer(learning_rate=0.01) train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step()) return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op) diff --git a/tensorflow/examples/speech_commands/label_wav_dir.py b/tensorflow/examples/speech_commands/label_wav_dir.py new file mode 100644 index 0000000000000000000000000000000000000000..a34db512dda86be138e07a4ffaa1963fe00a5cea --- /dev/null +++ b/tensorflow/examples/speech_commands/label_wav_dir.py @@ -0,0 +1,136 @@ +# 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. +# ============================================================================== +r"""Runs a trained audio graph against WAVE files and reports the results. + +The model, labels and .wav files specified in the arguments will be loaded, and +then the predictions from running the model against the audio data will be +printed to the console. This is a useful script for sanity checking trained +models, and as an example of how to use an audio model from Python. + +Here's an example of running it: + +python tensorflow/examples/speech_commands/label_wav_dir.py \ +--graph=/tmp/my_frozen_graph.pb \ +--labels=/tmp/speech_commands_train/conv_labels.txt \ +--wav_dir=/tmp/speech_dataset/left + +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import argparse +import glob +import sys + +import tensorflow as tf + +# pylint: disable=unused-import +from tensorflow.contrib.framework.python.ops import audio_ops as contrib_audio +# pylint: enable=unused-import + +FLAGS = None + + +def load_graph(filename): + """Unpersists graph from file as default graph.""" + with tf.gfile.FastGFile(filename, 'rb') as f: + graph_def = tf.GraphDef() + graph_def.ParseFromString(f.read()) + tf.import_graph_def(graph_def, name='') + + +def load_labels(filename): + """Read in labels, one label per line.""" + return [line.rstrip() for line in tf.gfile.GFile(filename)] + + +def run_graph(wav_dir, labels, input_layer_name, output_layer_name, + num_top_predictions): + """Runs the audio data through the graph and prints predictions.""" + with tf.Session() as sess: + # Feed the audio data as input to the graph. + # predictions will contain a two-dimensional array, where one + # dimension represents the input image count, and the other has + # predictions per class + for wav_path in glob.glob(wav_dir + '/*.wav'): + if not wav_path or not tf.gfile.Exists(wav_path): + tf.logging.fatal('Audio file does not exist %s', wav_path) + + with open(wav_path, 'rb') as wav_file: + wav_data = wav_file.read() + + softmax_tensor = sess.graph.get_tensor_by_name(output_layer_name) + predictions, = sess.run(softmax_tensor, {input_layer_name: wav_data}) + + # Sort to show labels in order of confidence + print('\n%s' % (wav_path.split('/')[-1])) + top_k = predictions.argsort()[-num_top_predictions:][::-1] + for node_id in top_k: + human_string = labels[node_id] + score = predictions[node_id] + print('%s (score = %.5f)' % (human_string, score)) + + return 0 + + +def label_wav(wav_dir, labels, graph, input_name, output_name, how_many_labels): + """Loads the model and labels, and runs the inference to print predictions.""" + if not labels or not tf.gfile.Exists(labels): + tf.logging.fatal('Labels file does not exist %s', labels) + + if not graph or not tf.gfile.Exists(graph): + tf.logging.fatal('Graph file does not exist %s', graph) + + labels_list = load_labels(labels) + + # load graph, which is stored in the default session + load_graph(graph) + + run_graph(wav_dir, labels_list, input_name, output_name, how_many_labels) + + +def main(_): + """Entry point for script, converts flags to arguments.""" + label_wav(FLAGS.wav_dir, FLAGS.labels, FLAGS.graph, FLAGS.input_name, + FLAGS.output_name, FLAGS.how_many_labels) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument( + '--wav_dir', type=str, default='', help='Audio file to be identified.') + parser.add_argument( + '--graph', type=str, default='', help='Model to use for identification.') + parser.add_argument( + '--labels', type=str, default='', help='Path to file containing labels.') + parser.add_argument( + '--input_name', + type=str, + default='wav_data:0', + help='Name of WAVE data input node in model.') + parser.add_argument( + '--output_name', + type=str, + default='labels_softmax:0', + help='Name of node outputting a prediction in the model.') + parser.add_argument( + '--how_many_labels', + type=int, + default=3, + help='Number of results to show.') + + FLAGS, unparsed = parser.parse_known_args() + tf.app.run(main=main, argv=[sys.argv[0]] + unparsed) diff --git a/tensorflow/examples/speech_commands/train.py b/tensorflow/examples/speech_commands/train.py index a4e80041f82191d7c58a3e52c929340eb604ec9d..f084931215261f183f1ecfc5517ea9a5126db039 100644 --- a/tensorflow/examples/speech_commands/train.py +++ b/tensorflow/examples/speech_commands/train.py @@ -357,12 +357,12 @@ if __name__ == '__main__': '--window_size_ms', type=float, default=30.0, - help='How long each spectrogram timeslice is',) + help='How long each spectrogram timeslice is.',) parser.add_argument( '--window_stride_ms', type=float, default=10.0, - help='How long each spectrogram timeslice is',) + help='How far to move in time between spectogram timeslices.',) parser.add_argument( '--dct_coefficient_count', type=int, diff --git a/tensorflow/examples/tutorials/mnist/input_data.py b/tensorflow/examples/tutorials/mnist/input_data.py index f1a7e1c4af57dba4f06326eb8b03c7eddae86b51..fa148ae3e6f44e140e3b4fb6a4204a601b6c0a24 100644 --- a/tensorflow/examples/tutorials/mnist/input_data.py +++ b/tensorflow/examples/tutorials/mnist/input_data.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +# pylint: disable=unused-import import gzip import os import tempfile @@ -27,3 +28,4 @@ from six.moves import urllib from six.moves import xrange # pylint: disable=redefined-builtin import tensorflow as tf from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets +# pylint: enable=unused-import diff --git a/tensorflow/examples/tutorials/word2vec/BUILD b/tensorflow/examples/tutorials/word2vec/BUILD index 42d6355b4f06258a3c22d0ef324bb31880f2d9a3..bfcf4592690a1692db67090c9b6d4e1e4832c45f 100644 --- a/tensorflow/examples/tutorials/word2vec/BUILD +++ b/tensorflow/examples/tutorials/word2vec/BUILD @@ -13,6 +13,9 @@ py_binary( "word2vec_basic.py", ], srcs_version = "PY2AND3", + tags = [ + "no-internal-py3", + ], deps = [ "//tensorflow:tensorflow_py", "//third_party/py/numpy", diff --git a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py index d055d157454d4cb351e8db59eec484f212893fe5..14ae7fbf35836ad7f5d56101ae0fc33a3f3fb9ba 100644 --- a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py +++ b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py @@ -131,7 +131,7 @@ def generate_batch(batch_size, num_skips, skip_window): batch = np.ndarray(shape=(batch_size), dtype=np.int32) labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32) span = 2 * skip_window + 1 # [ skip_window target skip_window ] - buffer = collections.deque(maxlen=span) + buffer = collections.deque(maxlen=span) # pylint: disable=redefined-builtin if data_index + span > len(data): data_index = 0 buffer.extend(data[data_index:data_index + span]) @@ -270,12 +270,6 @@ with tf.Session(graph=graph) as session: run_metadata=run_metadata) average_loss += loss_val - # Add returned summaries to writer in each step. - writer.add_summary(summary, step) - # Add metadata to visualize the graph for the last run. - if step == (num_steps - 1): - writer.add_run_metadata(run_metadata, 'step%d' % step) - # Add returned summaries to writer in each step. writer.add_summary(summary, step) # Add metadata to visualize the graph for the last run. diff --git a/tensorflow/examples/udacity/5_word2vec.ipynb b/tensorflow/examples/udacity/5_word2vec.ipynb index 18c456cad787b2ed5b39d5791de649874bbe7ae3..3b43d1fb55ee5d7f6a91754a221962755f04190c 100644 --- a/tensorflow/examples/udacity/5_word2vec.ipynb +++ b/tensorflow/examples/udacity/5_word2vec.ipynb @@ -455,7 +455,7 @@ " \n", " # Compute the similarity between minibatch examples and all embeddings.\n", " # We use the cosine distance:\n", - " norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keep_dims=True))\n", + " norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keepdims=True))\n", " normalized_embeddings = embeddings / norm\n", " valid_embeddings = tf.nn.embedding_lookup(\n", " normalized_embeddings, valid_dataset)\n", diff --git a/tensorflow/contrib/bayesflow/python/ops/csiszar_divergence.py b/tensorflow/experimental_api.py similarity index 50% rename from tensorflow/contrib/bayesflow/python/ops/csiszar_divergence.py rename to tensorflow/experimental_api.py index 9f7a95f138f7fd3e726f095dc16f41abb6182e17..63a8aa9cb1dc130a7999c3b248815633998c4cd0 100644 --- a/tensorflow/contrib/bayesflow/python/ops/csiszar_divergence.py +++ b/tensorflow/experimental_api.py @@ -1,4 +1,4 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -12,40 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Csiszar f-Divergence and helpers. -See ${python/contrib.bayesflow.csiszar_divergence}. -""" +# Bring in all of the public TensorFlow interface into this +# module. from __future__ import absolute_import from __future__ import division from __future__ import print_function -# go/tf-wildcard-import +# pylint: disable=g-bad-import-order +from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import # pylint: disable=wildcard-import -from tensorflow.contrib.bayesflow.python.ops.csiszar_divergence_impl import * +from tensorflow.tools.api.generator.api import * # pylint: disable=redefined-builtin # pylint: enable=wildcard-import -from tensorflow.python.util.all_util import remove_undocumented -_allowed_symbols = [ - 'amari_alpha', - 'arithmetic_geometric', - 'chi_square', - 'csiszar_vimco', - 'dual_csiszar_function', - 'jeffreys', - 'jensen_shannon', - 'kl_forward', - 'kl_reverse', - 'log1p_abs', - 'modified_gan', - 'monte_carlo_csiszar_f_divergence', - 'pearson', - 'squared_hellinger', - 'symmetrized_csiszar_function', - 'total_variation', - 't_power', - 'triangular', -] +from tensorflow.python.util.lazy_loader import LazyLoader +contrib = LazyLoader('contrib', globals(), 'tensorflow.contrib') +del LazyLoader -remove_undocumented(__name__, _allowed_symbols) +from tensorflow.python.platform import flags # pylint: disable=g-import-not-at-top +app.flags = flags # pylint: disable=undefined-variable + +del absolute_import +del division +del print_function diff --git a/tensorflow/go/genop/internal/api_def_map.go b/tensorflow/go/genop/internal/api_def_map.go index 07b689dbba23a3aa991983f3b373fa8445c673e1..8600452b476dee49292cbffe630026cf6077e22b 100644 --- a/tensorflow/go/genop/internal/api_def_map.go +++ b/tensorflow/go/genop/internal/api_def_map.go @@ -31,7 +31,7 @@ import ( "unsafe" "github.com/golang/protobuf/proto" - pb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/tensorflow/core/framework" + pb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/tensorflow/core/framework_go_proto" ) // Encapsulates a collection of API definitions. diff --git a/tensorflow/go/genop/internal/genop.go b/tensorflow/go/genop/internal/genop.go index 82f7510f2ed947e0a87e4d88cfce1ecaaa6362f8..fb8163121850cee36e1fcc652ca258b1fe2d42ff 100644 --- a/tensorflow/go/genop/internal/genop.go +++ b/tensorflow/go/genop/internal/genop.go @@ -47,7 +47,7 @@ import ( "unsafe" "github.com/golang/protobuf/proto" - pb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/tensorflow/core/framework" + pb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/tensorflow/core/framework_go_proto" ) // GenerateFunctionsForRegisteredOps writes a Go source code file to w @@ -359,13 +359,13 @@ type attrWrapper struct { api *pb.ApiDef_Attr } -func (a *attrWrapper) Name() string { return a.api.Name } -func (a *attrWrapper) RenameTo() string { return a.api.RenameTo } -func (a *attrWrapper) Description() string { return a.api.Description } -func (a *attrWrapper) Type() string { return a.op.Type } -func (a *attrWrapper) IsListAttr() bool { return isListAttr(a.op) } -func (a *attrWrapper) HasMinimum() bool { return a.op.HasMinimum } -func (a *attrWrapper) Minimum() int64 { return a.op.Minimum } +func (a *attrWrapper) Name() string { return a.api.Name } +func (a *attrWrapper) RenameTo() string { return a.api.RenameTo } +func (a *attrWrapper) Description() string { return a.api.Description } +func (a *attrWrapper) Type() string { return a.op.Type } +func (a *attrWrapper) IsListAttr() bool { return isListAttr(a.op) } +func (a *attrWrapper) HasMinimum() bool { return a.op.HasMinimum } +func (a *attrWrapper) Minimum() int64 { return a.op.Minimum } func (a *attrWrapper) DefaultValue() interface{} { return a.api.DefaultValue } type argWrapper struct { diff --git a/tensorflow/go/genop/internal/genop_test.go b/tensorflow/go/genop/internal/genop_test.go index b3a23dff102a690b1f7f08b675219929355f139f..d20d22e0c1502f92ade7ef5aa40985dce73b7552 100644 --- a/tensorflow/go/genop/internal/genop_test.go +++ b/tensorflow/go/genop/internal/genop_test.go @@ -22,7 +22,7 @@ import ( "testing" "github.com/golang/protobuf/proto" - pb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/tensorflow/core/framework" + pb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/tensorflow/core/framework_go_proto" ) // Creates an ApiDef based on opdef and applies overrides diff --git a/tensorflow/go/graph.go b/tensorflow/go/graph.go index fc087d9d995dfe031e61fd0fa15d649c2ee35cc9..08943a527cbdc072b12b066240c213be45ffd54c 100644 --- a/tensorflow/go/graph.go +++ b/tensorflow/go/graph.go @@ -173,7 +173,11 @@ type OpSpec struct { // operation. Attrs map[string]interface{} - // Other possible fields: Device, ColocateWith, ControlInputs. + // Operations that must be executed before executing the operation + // being added. + ControlDependencies []*Operation + + // Other possible fields: Device, ColocateWith. } // AddOperation adds an operation to g. @@ -204,6 +208,9 @@ func (g *Graph) AddOperation(args OpSpec) (*Operation, error) { } } } + for _, in := range args.ControlDependencies { + C.TF_AddControlInput(cdesc, in.c) + } status := newStatus() for name, value := range args.Attrs { if err := setAttr(cdesc, status, name, value); err != nil { diff --git a/tensorflow/go/op/scope.go b/tensorflow/go/op/scope.go index a9ec79463a00022bf85bf00032df9004648525ae..13de4294dc2ebdfff9bb68d277c09239d0bc8593 100644 --- a/tensorflow/go/op/scope.go +++ b/tensorflow/go/op/scope.go @@ -33,10 +33,11 @@ import ( // A Scope object and all its derivates (e.g., obtained from Scope.SubScope) // are not safe for concurrent use by multiple goroutines. type Scope struct { - graph *tf.Graph - namemap map[string]int - namespace string - err *scopeErr + graph *tf.Graph + namemap map[string]int + namespace string + controlDependencies []*tf.Operation + err *scopeErr } // scopeErr is used to share errors between all derivatives of a root scope. @@ -80,6 +81,7 @@ func (s *Scope) AddOperation(args tf.OpSpec) *tf.Operation { if s.namespace != "" { args.Name = s.namespace + "/" + args.Name } + args.ControlDependencies = append(args.ControlDependencies, s.controlDependencies...) op, err := s.graph.AddOperation(args) if err != nil { s.UpdateErr(args.Type, err) @@ -103,6 +105,28 @@ func (s *Scope) SubScope(namespace string) *Scope { } } +// WithControlDependencies returns a new Scope which will cause all operations +// added to the graph to execute only after all the provided operations have +// executed first (in addition to any other control dependencies in s). +func (s *Scope) WithControlDependencies(ops ...*tf.Operation) *Scope { + // Force a copy of the control dependencies into a new underlying array on + // every call. We cannot alias the same underlying array as `ops`, otherwise + // the user could modify that array after calling s.WithControlDependencies, + // which would be confusing. We cannot alias the same underlying array as the + // original `s.controlDependencies`, since Scopes form a logical tree, and + // other calls to s.WithControlDependencies could stomp on each other. + deps := make([]*tf.Operation, 0, len(s.controlDependencies)+len(ops)) + deps = append(deps, s.controlDependencies...) + deps = append(deps, ops...) + return &Scope{ + graph: s.graph, + namemap: s.namemap, + namespace: s.namespace, + controlDependencies: deps, + err: s.err, + } +} + // Err returns the error, if any, encountered during the construction // of the Graph managed by s. // diff --git a/tensorflow/go/op/scope_test.go b/tensorflow/go/op/scope_test.go index 6fb5d32e503c7c9a5a48747844da15be81b1de2d..b58a61de98b0f5b04959e1eca35c6b6c4d77e42b 100644 --- a/tensorflow/go/op/scope_test.go +++ b/tensorflow/go/op/scope_test.go @@ -69,6 +69,49 @@ func TestScopeSubScopeErrors(t *testing.T) { } } +func TestControlDependencies(t *testing.T) { + var ( + s = NewScope() + zero = Const(s.SubScope("zero"), int32(0)) + one = Const(s.SubScope("one"), int32(1)) + variable = VarHandleOp(s, tf.Int32, tf.ScalarShape()) + init = AssignVariableOp(s, variable, zero) + update = AssignAddVariableOp(s, variable, one) + readDeps = []*tf.Operation{update} + ) + // We intend for `read` to have a control dependency on `update`. + s = s.WithControlDependencies(readDeps...) + // Ensure that Scope.WithControlDependencies makes a copy of the underlying + // array, rather than just holding a slice reference to the same user-supplied + // underlying array. If the copy is correctly performed, overwriting + // readDeps[0] should have no effect on control dependencies for `read`. + readDeps[0] = init + read := ReadVariableOp(s, variable, tf.Int32) + + graph, err := s.Finalize() + if err != nil { + t.Fatal(err) + } + sess, err := tf.NewSession(graph, nil) + if err != nil { + t.Fatal(err) + } + if _, err = sess.Run(nil, nil, []*tf.Operation{init}); err != nil { + t.Fatal(err) + } + // Without the control dependency, the read operation may not see the + // update. + for i := int32(0); i < 10; i++ { + out, err := sess.Run(nil, []tf.Output{read}, nil) + if err != nil { + t.Fatal(err) + } + if got, want := out[0].Value().(int32), i+1; got != want { + t.Errorf("Got %d, want %d", got, want) + } + } +} + func TestScopeFinalize(t *testing.T) { var ( root = NewScope() diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go index 5b19c90238ef3bb1361a5e2476e94dd06e76d128..469d1e9adb73849e2b3cf963ac9b8fb2a0861c35 100644 --- a/tensorflow/go/op/wrappers.go +++ b/tensorflow/go/op/wrappers.go @@ -220,174 +220,97 @@ func CreateSummaryFileWriter(scope *Scope, writer tf.Output, logdir tf.Output, m return scope.AddOperation(opspec) } -// Partitions `data` into `num_partitions` tensors using indices from `partitions`. -// -// For each index tuple `js` of size `partitions.ndim`, the slice `data[js, ...]` -// becomes part of `outputs[partitions[js]]`. The slices with `partitions[js] = i` -// are placed in `outputs[i]` in lexicographic order of `js`, and the first -// dimension of `outputs[i]` is the number of entries in `partitions` equal to `i`. -// In detail, -// -// ```python -// outputs[i].shape = [sum(partitions == i)] + data.shape[partitions.ndim:] -// -// outputs[i] = pack([data[js, ...] for js if partitions[js] == i]) -// ``` -// -// `data.shape` must start with `partitions.shape`. -// -// For example: -// -// ```python -// # Scalar partitions. -// partitions = 1 -// num_partitions = 2 -// data = [10, 20] -// outputs[0] = [] # Empty with shape [0, 2] -// outputs[1] = [[10, 20]] -// -// # Vector partitions. -// partitions = [0, 0, 1, 1, 0] -// num_partitions = 2 -// data = [10, 20, 30, 40, 50] -// outputs[0] = [10, 20, 50] -// outputs[1] = [30, 40] -// ``` -// -// See `dynamic_stitch` for an example on how to merge partitions back. -// -//
-// -//
-// -// Arguments: -// -// partitions: Any shape. Indices in the range `[0, num_partitions)`. -// num_partitions: The number of partitions to output. -func DynamicPartition(scope *Scope, data tf.Output, partitions tf.Output, num_partitions int64) (outputs []tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"num_partitions": num_partitions} - opspec := tf.OpSpec{ - Type: "DynamicPartition", - Input: []tf.Input{ - data, partitions, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - if scope.Err() != nil { - return - } - var idx int - var err error - if outputs, idx, err = makeOutputList(op, idx, "outputs"); err != nil { - scope.UpdateErr("DynamicPartition", err) - return - } - return outputs -} - -// MutableHashTableOfTensorsV2Attr is an optional argument to MutableHashTableOfTensorsV2. -type MutableHashTableOfTensorsV2Attr func(optionalAttr) +// FakeQuantWithMinMaxVarsPerChannelGradientAttr is an optional argument to FakeQuantWithMinMaxVarsPerChannelGradient. +type FakeQuantWithMinMaxVarsPerChannelGradientAttr func(optionalAttr) -// MutableHashTableOfTensorsV2Container sets the optional container attribute to value. +// FakeQuantWithMinMaxVarsPerChannelGradientNumBits sets the optional num_bits attribute to value. // -// value: If non-empty, this table is placed in the given container. -// Otherwise, a default container is used. -// If not specified, defaults to "" -func MutableHashTableOfTensorsV2Container(value string) MutableHashTableOfTensorsV2Attr { +// value: The bitwidth of the quantization; between 2 and 8, inclusive. +// If not specified, defaults to 8 +func FakeQuantWithMinMaxVarsPerChannelGradientNumBits(value int64) FakeQuantWithMinMaxVarsPerChannelGradientAttr { return func(m optionalAttr) { - m["container"] = value + m["num_bits"] = value } } -// MutableHashTableOfTensorsV2SharedName sets the optional shared_name attribute to value. +// FakeQuantWithMinMaxVarsPerChannelGradientNarrowRange sets the optional narrow_range attribute to value. // -// value: If non-empty, this table is shared under the given name across -// multiple sessions. -// If not specified, defaults to "" -func MutableHashTableOfTensorsV2SharedName(value string) MutableHashTableOfTensorsV2Attr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// MutableHashTableOfTensorsV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value. +// value: Whether to quantize into 2^num_bits - 1 distinct values. // If not specified, defaults to false -func MutableHashTableOfTensorsV2UseNodeNameSharing(value bool) MutableHashTableOfTensorsV2Attr { - return func(m optionalAttr) { - m["use_node_name_sharing"] = value - } -} - -// MutableHashTableOfTensorsV2ValueShape sets the optional value_shape attribute to value. -// If not specified, defaults to <> -func MutableHashTableOfTensorsV2ValueShape(value tf.Shape) MutableHashTableOfTensorsV2Attr { +func FakeQuantWithMinMaxVarsPerChannelGradientNarrowRange(value bool) FakeQuantWithMinMaxVarsPerChannelGradientAttr { return func(m optionalAttr) { - m["value_shape"] = value + m["narrow_range"] = value } } -// Creates an empty hash table. -// -// This op creates a mutable hash table, specifying the type of its keys and -// values. Each value must be a vector. Data can be inserted into the table using -// the insert operations. It does not support the initialization operation. +// Compute gradients for a FakeQuantWithMinMaxVarsPerChannel operation. // // Arguments: -// key_dtype: Type of the table keys. -// value_dtype: Type of the table values. +// gradients: Backpropagated gradients above the FakeQuantWithMinMaxVars operation, +// shape one of: `[d]`, `[b, d]`, `[b, h, w, d]`. +// inputs: Values passed as inputs to the FakeQuantWithMinMaxVars operation, shape +// same as `gradients`. +// min, max: Quantization interval, floats of shape `[d]`. // -// Returns Handle to a table. -func MutableHashTableOfTensorsV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...MutableHashTableOfTensorsV2Attr) (table_handle tf.Output) { +// +// +// Returns Backpropagated gradients w.r.t. inputs, shape same as +// `inputs`: +// `gradients * (inputs >= min && inputs <= max)`.Backpropagated gradients w.r.t. min parameter, shape `[d]`: +// `sum_per_d(gradients * (inputs < min))`.Backpropagated gradients w.r.t. max parameter, shape `[d]`: +// `sum_per_d(gradients * (inputs > max))`. +func FakeQuantWithMinMaxVarsPerChannelGradient(scope *Scope, gradients tf.Output, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsPerChannelGradientAttr) (backprops_wrt_input tf.Output, backprop_wrt_min tf.Output, backprop_wrt_max tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype} + attrs := map[string]interface{}{} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "MutableHashTableOfTensorsV2", - + Type: "FakeQuantWithMinMaxVarsPerChannelGradient", + Input: []tf.Input{ + gradients, inputs, min, max, + }, Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2) } -// ResourceApplyProximalAdagradAttr is an optional argument to ResourceApplyProximalAdagrad. -type ResourceApplyProximalAdagradAttr func(optionalAttr) +// FakeQuantWithMinMaxVarsPerChannelAttr is an optional argument to FakeQuantWithMinMaxVarsPerChannel. +type FakeQuantWithMinMaxVarsPerChannelAttr func(optionalAttr) -// ResourceApplyProximalAdagradUseLocking sets the optional use_locking attribute to value. -// -// value: If True, updating of the var and accum tensors will be protected by -// a lock; otherwise the behavior is undefined, but may exhibit less contention. +// FakeQuantWithMinMaxVarsPerChannelNumBits sets the optional num_bits attribute to value. +// If not specified, defaults to 8 +func FakeQuantWithMinMaxVarsPerChannelNumBits(value int64) FakeQuantWithMinMaxVarsPerChannelAttr { + return func(m optionalAttr) { + m["num_bits"] = value + } +} + +// FakeQuantWithMinMaxVarsPerChannelNarrowRange sets the optional narrow_range attribute to value. // If not specified, defaults to false -func ResourceApplyProximalAdagradUseLocking(value bool) ResourceApplyProximalAdagradAttr { +func FakeQuantWithMinMaxVarsPerChannelNarrowRange(value bool) FakeQuantWithMinMaxVarsPerChannelAttr { return func(m optionalAttr) { - m["use_locking"] = value + m["narrow_range"] = value } } -// Update '*var' and '*accum' according to FOBOS with Adagrad learning rate. +// Fake-quantize the 'inputs' tensor of type float and one of the shapes: `[d]`, // -// accum += grad * grad -// prox_v = var - lr * grad * (1 / sqrt(accum)) -// var = sign(prox_v)/(1+lr*l2) * max{|prox_v|-lr*l1,0} +// `[b, d]` `[b, h, w, d]` via per-channel floats `min` and `max` of shape `[d]` +// to 'outputs' tensor of same shape as `inputs`. // -// Arguments: -// var_: Should be from a Variable(). -// accum: Should be from a Variable(). -// lr: Scaling factor. Must be a scalar. -// l1: L1 regularization. Must be a scalar. -// l2: L2 regularization. Must be a scalar. -// grad: The gradient. +// `[min; max]` define the clamping range for the `inputs` data. +// `inputs` values are quantized into the quantization range (`[0; 2^num_bits - 1]` +// when `narrow_range` is false and `[1; 2^num_bits - 1]` when it is true) and +// then de-quantized and output as floats in `[min; max]` interval. +// `num_bits` is the bitwidth of the quantization; between 2 and 8, inclusive. // -// Returns the created operation. -func ResourceApplyProximalAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, grad tf.Output, optional ...ResourceApplyProximalAdagradAttr) (o *tf.Operation) { +// This operation has a gradient and thus allows for training `min` and `max` +// values. +func FakeQuantWithMinMaxVarsPerChannel(scope *Scope, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsPerChannelAttr) (outputs tf.Output) { if scope.Err() != nil { return } @@ -396,209 +319,190 @@ func ResourceApplyProximalAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, a(attrs) } opspec := tf.OpSpec{ - Type: "ResourceApplyProximalAdagrad", + Type: "FakeQuantWithMinMaxVarsPerChannel", Input: []tf.Input{ - var_, accum, lr, l1, l2, grad, + inputs, min, max, }, Attrs: attrs, } - return scope.AddOperation(opspec) + op := scope.AddOperation(opspec) + return op.Output(0) } -// MutableHashTableV2Attr is an optional argument to MutableHashTableV2. -type MutableHashTableV2Attr func(optionalAttr) - -// MutableHashTableV2Container sets the optional container attribute to value. -// -// value: If non-empty, this table is placed in the given container. -// Otherwise, a default container is used. -// If not specified, defaults to "" -func MutableHashTableV2Container(value string) MutableHashTableV2Attr { - return func(m optionalAttr) { - m["container"] = value - } -} +// FakeQuantWithMinMaxVarsGradientAttr is an optional argument to FakeQuantWithMinMaxVarsGradient. +type FakeQuantWithMinMaxVarsGradientAttr func(optionalAttr) -// MutableHashTableV2SharedName sets the optional shared_name attribute to value. +// FakeQuantWithMinMaxVarsGradientNumBits sets the optional num_bits attribute to value. // -// value: If non-empty, this table is shared under the given name across -// multiple sessions. -// If not specified, defaults to "" -func MutableHashTableV2SharedName(value string) MutableHashTableV2Attr { +// value: The bitwidth of the quantization; between 2 and 8, inclusive. +// If not specified, defaults to 8 +func FakeQuantWithMinMaxVarsGradientNumBits(value int64) FakeQuantWithMinMaxVarsGradientAttr { return func(m optionalAttr) { - m["shared_name"] = value + m["num_bits"] = value } } -// MutableHashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value. +// FakeQuantWithMinMaxVarsGradientNarrowRange sets the optional narrow_range attribute to value. // -// value: If true and shared_name is empty, the table is shared -// using the node name. +// value: Whether to quantize into 2^num_bits - 1 distinct values. // If not specified, defaults to false -func MutableHashTableV2UseNodeNameSharing(value bool) MutableHashTableV2Attr { +func FakeQuantWithMinMaxVarsGradientNarrowRange(value bool) FakeQuantWithMinMaxVarsGradientAttr { return func(m optionalAttr) { - m["use_node_name_sharing"] = value + m["narrow_range"] = value } } -// Creates an empty hash table. -// -// This op creates a mutable hash table, specifying the type of its keys and -// values. Each value must be a scalar. Data can be inserted into the table using -// the insert operations. It does not support the initialization operation. +// Compute gradients for a FakeQuantWithMinMaxVars operation. // // Arguments: -// key_dtype: Type of the table keys. -// value_dtype: Type of the table values. +// gradients: Backpropagated gradients above the FakeQuantWithMinMaxVars operation. +// inputs: Values passed as inputs to the FakeQuantWithMinMaxVars operation. +// min, max: Quantization interval, scalar floats. // -// Returns Handle to a table. -func MutableHashTableV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...MutableHashTableV2Attr) (table_handle tf.Output) { +// +// +// Returns Backpropagated gradients w.r.t. inputs: +// `gradients * (inputs >= min && inputs <= max)`.Backpropagated gradients w.r.t. min parameter: +// `sum(gradients * (inputs < min))`.Backpropagated gradients w.r.t. max parameter: +// `sum(gradients * (inputs > max))`. +func FakeQuantWithMinMaxVarsGradient(scope *Scope, gradients tf.Output, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsGradientAttr) (backprops_wrt_input tf.Output, backprop_wrt_min tf.Output, backprop_wrt_max tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype} + attrs := map[string]interface{}{} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "MutableHashTableV2", - + Type: "FakeQuantWithMinMaxVarsGradient", + Input: []tf.Input{ + gradients, inputs, min, max, + }, Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2) } -// MapUnstageNoKeyAttr is an optional argument to MapUnstageNoKey. -type MapUnstageNoKeyAttr func(optionalAttr) +// FakeQuantWithMinMaxArgsGradientAttr is an optional argument to FakeQuantWithMinMaxArgsGradient. +type FakeQuantWithMinMaxArgsGradientAttr func(optionalAttr) -// MapUnstageNoKeyCapacity sets the optional capacity attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func MapUnstageNoKeyCapacity(value int64) MapUnstageNoKeyAttr { +// FakeQuantWithMinMaxArgsGradientMin sets the optional min attribute to value. +// If not specified, defaults to -6 +func FakeQuantWithMinMaxArgsGradientMin(value float32) FakeQuantWithMinMaxArgsGradientAttr { return func(m optionalAttr) { - m["capacity"] = value + m["min"] = value } } -// MapUnstageNoKeyMemoryLimit sets the optional memory_limit attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func MapUnstageNoKeyMemoryLimit(value int64) MapUnstageNoKeyAttr { +// FakeQuantWithMinMaxArgsGradientMax sets the optional max attribute to value. +// If not specified, defaults to 6 +func FakeQuantWithMinMaxArgsGradientMax(value float32) FakeQuantWithMinMaxArgsGradientAttr { return func(m optionalAttr) { - m["memory_limit"] = value + m["max"] = value } } -// MapUnstageNoKeyContainer sets the optional container attribute to value. -// If not specified, defaults to "" -func MapUnstageNoKeyContainer(value string) MapUnstageNoKeyAttr { +// FakeQuantWithMinMaxArgsGradientNumBits sets the optional num_bits attribute to value. +// If not specified, defaults to 8 +func FakeQuantWithMinMaxArgsGradientNumBits(value int64) FakeQuantWithMinMaxArgsGradientAttr { return func(m optionalAttr) { - m["container"] = value + m["num_bits"] = value } } -// MapUnstageNoKeySharedName sets the optional shared_name attribute to value. -// If not specified, defaults to "" -func MapUnstageNoKeySharedName(value string) MapUnstageNoKeyAttr { +// FakeQuantWithMinMaxArgsGradientNarrowRange sets the optional narrow_range attribute to value. +// If not specified, defaults to false +func FakeQuantWithMinMaxArgsGradientNarrowRange(value bool) FakeQuantWithMinMaxArgsGradientAttr { return func(m optionalAttr) { - m["shared_name"] = value + m["narrow_range"] = value } } -// Op removes and returns a random (key, value) +// Compute gradients for a FakeQuantWithMinMaxArgs operation. +// +// Arguments: +// gradients: Backpropagated gradients above the FakeQuantWithMinMaxArgs operation. +// inputs: Values passed as inputs to the FakeQuantWithMinMaxArgs operation. // -// from the underlying container. If the underlying container -// does not contain elements, the op will block until it does. -func MapUnstageNoKey(scope *Scope, indices tf.Output, dtypes []tf.DataType, optional ...MapUnstageNoKeyAttr) (key tf.Output, values []tf.Output) { +// Returns Backpropagated gradients below the FakeQuantWithMinMaxArgs operation: +// `gradients * (inputs >= min && inputs <= max)`. +func FakeQuantWithMinMaxArgsGradient(scope *Scope, gradients tf.Output, inputs tf.Output, optional ...FakeQuantWithMinMaxArgsGradientAttr) (backprops tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"dtypes": dtypes} + attrs := map[string]interface{}{} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "MapUnstageNoKey", + Type: "FakeQuantWithMinMaxArgsGradient", Input: []tf.Input{ - indices, + gradients, inputs, }, Attrs: attrs, } op := scope.AddOperation(opspec) - if scope.Err() != nil { - return - } - var idx int - var err error - key = op.Output(idx) - if values, idx, err = makeOutputList(op, idx, "values"); err != nil { - scope.UpdateErr("MapUnstageNoKey", err) - return - } - return key, values + return op.Output(0) } -// HashTableV2Attr is an optional argument to HashTableV2. -type HashTableV2Attr func(optionalAttr) +// FakeQuantWithMinMaxArgsAttr is an optional argument to FakeQuantWithMinMaxArgs. +type FakeQuantWithMinMaxArgsAttr func(optionalAttr) -// HashTableV2Container sets the optional container attribute to value. -// -// value: If non-empty, this table is placed in the given container. -// Otherwise, a default container is used. -// If not specified, defaults to "" -func HashTableV2Container(value string) HashTableV2Attr { +// FakeQuantWithMinMaxArgsMin sets the optional min attribute to value. +// If not specified, defaults to -6 +func FakeQuantWithMinMaxArgsMin(value float32) FakeQuantWithMinMaxArgsAttr { return func(m optionalAttr) { - m["container"] = value + m["min"] = value } } -// HashTableV2SharedName sets the optional shared_name attribute to value. -// -// value: If non-empty, this table is shared under the given name across -// multiple sessions. -// If not specified, defaults to "" -func HashTableV2SharedName(value string) HashTableV2Attr { +// FakeQuantWithMinMaxArgsMax sets the optional max attribute to value. +// If not specified, defaults to 6 +func FakeQuantWithMinMaxArgsMax(value float32) FakeQuantWithMinMaxArgsAttr { return func(m optionalAttr) { - m["shared_name"] = value + m["max"] = value } } -// HashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value. -// -// value: If true and shared_name is empty, the table is shared -// using the node name. +// FakeQuantWithMinMaxArgsNumBits sets the optional num_bits attribute to value. +// If not specified, defaults to 8 +func FakeQuantWithMinMaxArgsNumBits(value int64) FakeQuantWithMinMaxArgsAttr { + return func(m optionalAttr) { + m["num_bits"] = value + } +} + +// FakeQuantWithMinMaxArgsNarrowRange sets the optional narrow_range attribute to value. // If not specified, defaults to false -func HashTableV2UseNodeNameSharing(value bool) HashTableV2Attr { +func FakeQuantWithMinMaxArgsNarrowRange(value bool) FakeQuantWithMinMaxArgsAttr { return func(m optionalAttr) { - m["use_node_name_sharing"] = value + m["narrow_range"] = value } } -// Creates a non-initialized hash table. -// -// This op creates a hash table, specifying the type of its keys and values. -// Before using the table you will have to initialize it. After initialization the -// table will be immutable. +// Fake-quantize the 'inputs' tensor, type float to 'outputs' tensor of same type. // -// Arguments: -// key_dtype: Type of the table keys. -// value_dtype: Type of the table values. +// Attributes `[min; max]` define the clamping range for the `inputs` data. +// `inputs` values are quantized into the quantization range (`[0; 2^num_bits - 1]` +// when `narrow_range` is false and `[1; 2^num_bits - 1]` when it is true) and +// then de-quantized and output as floats in `[min; max]` interval. +// `num_bits` is the bitwidth of the quantization; between 2 and 8, inclusive. // -// Returns Handle to a table. -func HashTableV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...HashTableV2Attr) (table_handle tf.Output) { +// Quantization is called fake since the output is still in floating point. +func FakeQuantWithMinMaxArgs(scope *Scope, inputs tf.Output, optional ...FakeQuantWithMinMaxArgsAttr) (outputs tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype} + attrs := map[string]interface{}{} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "HashTableV2", - + Type: "FakeQuantWithMinMaxArgs", + Input: []tf.Input{ + inputs, + }, Attrs: attrs, } op := scope.AddOperation(opspec) @@ -629,6 +533,77 @@ func LookupTableImportV2(scope *Scope, table_handle tf.Output, keys tf.Output, v return scope.AddOperation(opspec) } +// MapPeekAttr is an optional argument to MapPeek. +type MapPeekAttr func(optionalAttr) + +// MapPeekCapacity sets the optional capacity attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func MapPeekCapacity(value int64) MapPeekAttr { + return func(m optionalAttr) { + m["capacity"] = value + } +} + +// MapPeekMemoryLimit sets the optional memory_limit attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func MapPeekMemoryLimit(value int64) MapPeekAttr { + return func(m optionalAttr) { + m["memory_limit"] = value + } +} + +// MapPeekContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func MapPeekContainer(value string) MapPeekAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// MapPeekSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func MapPeekSharedName(value string) MapPeekAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Op peeks at the values at the specified key. If the +// +// underlying container does not contain this key +// this op will block until it does. +func MapPeek(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.DataType, optional ...MapPeekAttr) (values []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"dtypes": dtypes} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MapPeek", + Input: []tf.Input{ + key, indices, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + if values, idx, err = makeOutputList(op, idx, "values"); err != nil { + scope.UpdateErr("MapPeek", err) + return + } + return values +} + // Returns (x - y)(x - y) element-wise. // // *NOTE*: `SquaredDifference` supports broadcasting. More about broadcasting @@ -1232,6 +1207,21 @@ func Sinh(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Computes rectified linear 6: `min(max(features, 0), 6)`. +func Relu6(scope *Scope, features tf.Output) (activations tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Relu6", + Input: []tf.Input{ + features, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Computes the sum along segments of a tensor. // // Read @{$math_ops#segmentation$the section on segmentation} for an explanation of @@ -1515,61 +1505,6 @@ func Igammac(scope *Scope, a tf.Output, x tf.Output) (z tf.Output) { return op.Output(0) } -// FakeQuantWithMinMaxVarsGradientAttr is an optional argument to FakeQuantWithMinMaxVarsGradient. -type FakeQuantWithMinMaxVarsGradientAttr func(optionalAttr) - -// FakeQuantWithMinMaxVarsGradientNumBits sets the optional num_bits attribute to value. -// -// value: The bitwidth of the quantization; between 2 and 8, inclusive. -// If not specified, defaults to 8 -func FakeQuantWithMinMaxVarsGradientNumBits(value int64) FakeQuantWithMinMaxVarsGradientAttr { - return func(m optionalAttr) { - m["num_bits"] = value - } -} - -// FakeQuantWithMinMaxVarsGradientNarrowRange sets the optional narrow_range attribute to value. -// -// value: Whether to quantize into 2^num_bits - 1 distinct values. -// If not specified, defaults to false -func FakeQuantWithMinMaxVarsGradientNarrowRange(value bool) FakeQuantWithMinMaxVarsGradientAttr { - return func(m optionalAttr) { - m["narrow_range"] = value - } -} - -// Compute gradients for a FakeQuantWithMinMaxVars operation. -// -// Arguments: -// gradients: Backpropagated gradients above the FakeQuantWithMinMaxVars operation. -// inputs: Values passed as inputs to the FakeQuantWithMinMaxVars operation. -// min, max: Quantization interval, scalar floats. -// -// -// -// Returns Backpropagated gradients w.r.t. inputs: -// `gradients * (inputs >= min && inputs <= max)`.Backpropagated gradients w.r.t. min parameter: -// `sum(gradients * (inputs < min))`.Backpropagated gradients w.r.t. max parameter: -// `sum(gradients * (inputs > max))`. -func FakeQuantWithMinMaxVarsGradient(scope *Scope, gradients tf.Output, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsGradientAttr) (backprops_wrt_input tf.Output, backprop_wrt_min tf.Output, backprop_wrt_max tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxVarsGradient", - Input: []tf.Input{ - gradients, inputs, min, max, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) -} - // LogUniformCandidateSamplerAttr is an optional argument to LogUniformCandidateSampler. type LogUniformCandidateSamplerAttr func(optionalAttr) @@ -2286,34 +2221,14 @@ func Conv2DBackpropFilter(scope *Scope, input tf.Output, filter_sizes tf.Output, // // Arguments: // reader_handle: Handle to a Reader. -func ReaderNumWorkUnitsCompletedV2(scope *Scope, reader_handle tf.Output) (units_completed tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "ReaderNumWorkUnitsCompletedV2", - Input: []tf.Input{ - reader_handle, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Returns x / y element-wise for real types. -// -// If `x` and `y` are reals, this will return the floating-point division. -// -// *NOTE*: `Div` supports broadcasting. More about broadcasting -// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) -func RealDiv(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { +func ReaderNumWorkUnitsCompletedV2(scope *Scope, reader_handle tf.Output) (units_completed tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "RealDiv", + Type: "ReaderNumWorkUnitsCompletedV2", Input: []tf.Input{ - x, y, + reader_handle, }, } op := scope.AddOperation(opspec) @@ -4022,21 +3937,6 @@ func TakeDataset(scope *Scope, input_dataset tf.Output, count tf.Output, output_ return op.Output(0) } -// Computes rectified linear 6: `min(max(features, 0), 6)`. -func Relu6(scope *Scope, features tf.Output) (activations tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Relu6", - Input: []tf.Input{ - features, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // Computes rectified linear gradients for a Relu operation. // // Arguments: @@ -4327,34 +4227,6 @@ func MaxPoolGradGradV2(scope *Scope, orig_input tf.Output, orig_output tf.Output return op.Output(0) } -// Fast Fourier transform. -// -// Computes the 1-dimensional discrete Fourier transform over the inner-most -// dimension of `input`. -// -// Arguments: -// input: A complex64 tensor. -// -// Returns A complex64 tensor of the same shape as `input`. The inner-most -// dimension of `input` is replaced with its 1D Fourier transform. -// -// @compatibility(numpy) -// Equivalent to np.fft.fft -// @end_compatibility -func FFT(scope *Scope, input tf.Output) (output tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "FFT", - Input: []tf.Input{ - input, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // MaxPoolAttr is an optional argument to MaxPool. type MaxPoolAttr func(optionalAttr) @@ -4468,47 +4340,6 @@ func MaxPoolGradWithArgmax(scope *Scope, input tf.Output, grad tf.Output, argmax return op.Output(0) } -// CriticalSectionOpAttr is an optional argument to CriticalSectionOp. -type CriticalSectionOpAttr func(optionalAttr) - -// CriticalSectionOpContainer sets the optional container attribute to value. -// -// value: the container this critical section is placed in. -// If not specified, defaults to "" -func CriticalSectionOpContainer(value string) CriticalSectionOpAttr { - return func(m optionalAttr) { - m["container"] = value - } -} - -// CriticalSectionOpSharedName sets the optional shared_name attribute to value. -// -// value: the name by which this critical section is referred to. -// If not specified, defaults to "" -func CriticalSectionOpSharedName(value string) CriticalSectionOpAttr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// Creates a handle to a CriticalSection resource. -func CriticalSectionOp(scope *Scope, optional ...CriticalSectionOpAttr) (resource tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "CriticalSectionOp", - - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // AvgPool3DAttr is an optional argument to AvgPool3D. type AvgPool3DAttr func(optionalAttr) @@ -4814,6 +4645,78 @@ func DepthwiseConv2dNative(scope *Scope, input tf.Output, filter tf.Output, stri return op.Output(0) } +// MaxPoolGradV2Attr is an optional argument to MaxPoolGradV2. +type MaxPoolGradV2Attr func(optionalAttr) + +// MaxPoolGradV2DataFormat sets the optional data_format attribute to value. +// +// value: Specify the data format of the input and output data. With the +// default format "NHWC", the data is stored in the order of: +// [batch, in_height, in_width, in_channels]. +// Alternatively, the format could be "NCHW", the data storage order of: +// [batch, in_channels, in_height, in_width]. +// If not specified, defaults to "NHWC" +func MaxPoolGradV2DataFormat(value string) MaxPoolGradV2Attr { + return func(m optionalAttr) { + m["data_format"] = value + } +} + +// Computes gradients of the maxpooling function. +// +// Arguments: +// orig_input: The original input tensor. +// orig_output: The original output tensor. +// grad: 4-D. Gradients w.r.t. the output of `max_pool`. +// ksize: The size of the window for each dimension of the input tensor. +// strides: The stride of the sliding window for each dimension of the +// input tensor. +// padding: The type of padding algorithm to use. +// +// Returns Gradients w.r.t. the input to `max_pool`. +func MaxPoolGradV2(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize tf.Output, strides tf.Output, padding string, optional ...MaxPoolGradV2Attr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"padding": padding} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MaxPoolGradV2", + Input: []tf.Input{ + orig_input, orig_output, grad, ksize, strides, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Restore a reader to a previously saved state. +// +// Not all Readers support being restored, so this can produce an +// Unimplemented error. +// +// Arguments: +// reader_handle: Handle to a Reader. +// state: Result of a ReaderSerializeState of a Reader with type +// matching reader_handle. +// +// Returns the created operation. +func ReaderRestoreStateV2(scope *Scope, reader_handle tf.Output, state tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ReaderRestoreStateV2", + Input: []tf.Input{ + reader_handle, state, + }, + } + return scope.AddOperation(opspec) +} + // TensorArrayGatherV3Attr is an optional argument to TensorArrayGatherV3. type TensorArrayGatherV3Attr func(optionalAttr) @@ -5478,117 +5381,12 @@ func QuantizedReluX(scope *Scope, features tf.Output, max_value tf.Output, min_f opspec := tf.OpSpec{ Type: "QuantizedReluX", Input: []tf.Input{ - features, max_value, min_features, max_features, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) -} - -// SummaryWriterAttr is an optional argument to SummaryWriter. -type SummaryWriterAttr func(optionalAttr) - -// SummaryWriterSharedName sets the optional shared_name attribute to value. -// If not specified, defaults to "" -func SummaryWriterSharedName(value string) SummaryWriterAttr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// SummaryWriterContainer sets the optional container attribute to value. -// If not specified, defaults to "" -func SummaryWriterContainer(value string) SummaryWriterAttr { - return func(m optionalAttr) { - m["container"] = value - } -} - -// Returns a handle to be used to access a summary writer. -// -// The summary writer is an in-graph resource which can be used by ops to write -// summaries to event files. -// -// Returns the summary writer resource. Scalar handle. -func SummaryWriter(scope *Scope, optional ...SummaryWriterAttr) (writer tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "SummaryWriter", - - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Computes gradients for SparseSegmentMean. -// -// Returns tensor "output" with same shape as grad, except for dimension 0 whose -// value is output_dim0. -// -// Arguments: -// grad: gradient propagated to the SparseSegmentMean op. -// indices: indices passed to the corresponding SparseSegmentMean op. -// segment_ids: segment_ids passed to the corresponding SparseSegmentMean op. -// output_dim0: dimension 0 of "data" passed to SparseSegmentMean op. -func SparseSegmentMeanGrad(scope *Scope, grad tf.Output, indices tf.Output, segment_ids tf.Output, output_dim0 tf.Output) (output tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "SparseSegmentMeanGrad", - Input: []tf.Input{ - grad, indices, segment_ids, output_dim0, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Applies softmax to a batched N-D `SparseTensor`. -// -// The inputs represent an N-D SparseTensor with logical shape `[..., B, C]` -// (where `N >= 2`), and with indices sorted in the canonical lexicographic order. -// -// This op is equivalent to applying the normal `tf.nn.softmax()` to each innermost -// logical submatrix with shape `[B, C]`, but with the catch that *the implicitly -// zero elements do not participate*. Specifically, the algorithm is equivalent -// to the following: -// -// (1) Applies `tf.nn.softmax()` to a densified view of each innermost submatrix -// with shape `[B, C]`, along the size-C dimension; -// (2) Masks out the original implicitly-zero locations; -// (3) Renormalizes the remaining elements. -// -// Hence, the `SparseTensor` result has exactly the same non-zero indices and -// shape. -// -// Arguments: -// sp_indices: 2-D. `NNZ x R` matrix with the indices of non-empty values in a -// SparseTensor, in canonical ordering. -// sp_values: 1-D. `NNZ` non-empty values corresponding to `sp_indices`. -// sp_shape: 1-D. Shape of the input SparseTensor. -// -// Returns 1-D. The `NNZ` values for the result `SparseTensor`. -func SparseSoftmax(scope *Scope, sp_indices tf.Output, sp_values tf.Output, sp_shape tf.Output) (output tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "SparseSoftmax", - Input: []tf.Input{ - sp_indices, sp_values, sp_shape, + features, max_value, min_features, max_features, }, + Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2) } // RandomPoissonAttr is an optional argument to RandomPoisson. @@ -5632,78 +5430,6 @@ func RandomPoisson(scope *Scope, shape tf.Output, rate tf.Output, optional ...Ra return op.Output(0) } -// MaxPoolGradV2Attr is an optional argument to MaxPoolGradV2. -type MaxPoolGradV2Attr func(optionalAttr) - -// MaxPoolGradV2DataFormat sets the optional data_format attribute to value. -// -// value: Specify the data format of the input and output data. With the -// default format "NHWC", the data is stored in the order of: -// [batch, in_height, in_width, in_channels]. -// Alternatively, the format could be "NCHW", the data storage order of: -// [batch, in_channels, in_height, in_width]. -// If not specified, defaults to "NHWC" -func MaxPoolGradV2DataFormat(value string) MaxPoolGradV2Attr { - return func(m optionalAttr) { - m["data_format"] = value - } -} - -// Computes gradients of the maxpooling function. -// -// Arguments: -// orig_input: The original input tensor. -// orig_output: The original output tensor. -// grad: 4-D. Gradients w.r.t. the output of `max_pool`. -// ksize: The size of the window for each dimension of the input tensor. -// strides: The stride of the sliding window for each dimension of the -// input tensor. -// padding: The type of padding algorithm to use. -// -// Returns Gradients w.r.t. the input to `max_pool`. -func MaxPoolGradV2(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize tf.Output, strides tf.Output, padding string, optional ...MaxPoolGradV2Attr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"padding": padding} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "MaxPoolGradV2", - Input: []tf.Input{ - orig_input, orig_output, grad, ksize, strides, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Restore a reader to a previously saved state. -// -// Not all Readers support being restored, so this can produce an -// Unimplemented error. -// -// Arguments: -// reader_handle: Handle to a Reader. -// state: Result of a ReaderSerializeState of a Reader with type -// matching reader_handle. -// -// Returns the created operation. -func ReaderRestoreStateV2(scope *Scope, reader_handle tf.Output, state tf.Output) (o *tf.Operation) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "ReaderRestoreStateV2", - Input: []tf.Input{ - reader_handle, state, - }, - } - return scope.AddOperation(opspec) -} - // ResourceSparseApplyFtrlV2Attr is an optional argument to ResourceSparseApplyFtrlV2. type ResourceSparseApplyFtrlV2Attr func(optionalAttr) @@ -6860,70 +6586,289 @@ func StringToHashBucketStrong(scope *Scope, input tf.Output, num_buckets int64, opspec := tf.OpSpec{ Type: "StringToHashBucketStrong", Input: []tf.Input{ - input, + input, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Generates values in an interval. +// +// A sequence of `num` evenly-spaced values are generated beginning at `start`. +// If `num > 1`, the values in the sequence increase by `stop - start / num - 1`, +// so that the last one is exactly `stop`. +// +// For example: +// +// ``` +// tf.linspace(10.0, 12.0, 3, name="linspace") => [ 10.0 11.0 12.0] +// ``` +// +// Arguments: +// start: First entry in the range. +// stop: Last entry in the range. +// num: Number of values to generate. +// +// Returns 1-D. The generated values. +func LinSpace(scope *Scope, start tf.Output, stop tf.Output, num tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "LinSpace", + Input: []tf.Input{ + start, stop, num, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// DestroyResourceOpAttr is an optional argument to DestroyResourceOp. +type DestroyResourceOpAttr func(optionalAttr) + +// DestroyResourceOpIgnoreLookupError sets the optional ignore_lookup_error attribute to value. +// +// value: whether to ignore the error when the resource +// doesn't exist. +// If not specified, defaults to true +func DestroyResourceOpIgnoreLookupError(value bool) DestroyResourceOpAttr { + return func(m optionalAttr) { + m["ignore_lookup_error"] = value + } +} + +// Deletes the resource specified by the handle. +// +// All subsequent operations using the resource will result in a NotFound +// error status. +// +// Arguments: +// resource: handle to the resource to delete. +// +// Returns the created operation. +func DestroyResourceOp(scope *Scope, resource tf.Output, optional ...DestroyResourceOpAttr) (o *tf.Operation) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "DestroyResourceOp", + Input: []tf.Input{ + resource, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) +} + +// SummaryWriterAttr is an optional argument to SummaryWriter. +type SummaryWriterAttr func(optionalAttr) + +// SummaryWriterSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func SummaryWriterSharedName(value string) SummaryWriterAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// SummaryWriterContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func SummaryWriterContainer(value string) SummaryWriterAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// Returns a handle to be used to access a summary writer. +// +// The summary writer is an in-graph resource which can be used by ops to write +// summaries to event files. +// +// Returns the summary writer resource. Scalar handle. +func SummaryWriter(scope *Scope, optional ...SummaryWriterAttr) (writer tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "SummaryWriter", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Computes gradients for SparseSegmentMean. +// +// Returns tensor "output" with same shape as grad, except for dimension 0 whose +// value is output_dim0. +// +// Arguments: +// grad: gradient propagated to the SparseSegmentMean op. +// indices: indices passed to the corresponding SparseSegmentMean op. +// segment_ids: segment_ids passed to the corresponding SparseSegmentMean op. +// output_dim0: dimension 0 of "data" passed to SparseSegmentMean op. +func SparseSegmentMeanGrad(scope *Scope, grad tf.Output, indices tf.Output, segment_ids tf.Output, output_dim0 tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "SparseSegmentMeanGrad", + Input: []tf.Input{ + grad, indices, segment_ids, output_dim0, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Applies softmax to a batched N-D `SparseTensor`. +// +// The inputs represent an N-D SparseTensor with logical shape `[..., B, C]` +// (where `N >= 2`), and with indices sorted in the canonical lexicographic order. +// +// This op is equivalent to applying the normal `tf.nn.softmax()` to each innermost +// logical submatrix with shape `[B, C]`, but with the catch that *the implicitly +// zero elements do not participate*. Specifically, the algorithm is equivalent +// to the following: +// +// (1) Applies `tf.nn.softmax()` to a densified view of each innermost submatrix +// with shape `[B, C]`, along the size-C dimension; +// (2) Masks out the original implicitly-zero locations; +// (3) Renormalizes the remaining elements. +// +// Hence, the `SparseTensor` result has exactly the same non-zero indices and +// shape. +// +// Arguments: +// sp_indices: 2-D. `NNZ x R` matrix with the indices of non-empty values in a +// SparseTensor, in canonical ordering. +// sp_values: 1-D. `NNZ` non-empty values corresponding to `sp_indices`. +// sp_shape: 1-D. Shape of the input SparseTensor. +// +// Returns 1-D. The `NNZ` values for the result `SparseTensor`. +func SparseSoftmax(scope *Scope, sp_indices tf.Output, sp_values tf.Output, sp_shape tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "SparseSoftmax", + Input: []tf.Input{ + sp_indices, sp_values, sp_shape, }, - Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// Generates values in an interval. +// Partitions `data` into `num_partitions` tensors using indices from `partitions`. // -// A sequence of `num` evenly-spaced values are generated beginning at `start`. -// If `num > 1`, the values in the sequence increase by `stop - start / num - 1`, -// so that the last one is exactly `stop`. +// For each index tuple `js` of size `partitions.ndim`, the slice `data[js, ...]` +// becomes part of `outputs[partitions[js]]`. The slices with `partitions[js] = i` +// are placed in `outputs[i]` in lexicographic order of `js`, and the first +// dimension of `outputs[i]` is the number of entries in `partitions` equal to `i`. +// In detail, // -// For example: +// ```python +// outputs[i].shape = [sum(partitions == i)] + data.shape[partitions.ndim:] // +// outputs[i] = pack([data[js, ...] for js if partitions[js] == i]) // ``` -// tf.linspace(10.0, 12.0, 3, name="linspace") => [ 10.0 11.0 12.0] +// +// `data.shape` must start with `partitions.shape`. +// +// For example: +// +// ```python +// # Scalar partitions. +// partitions = 1 +// num_partitions = 2 +// data = [10, 20] +// outputs[0] = [] # Empty with shape [0, 2] +// outputs[1] = [[10, 20]] +// +// # Vector partitions. +// partitions = [0, 0, 1, 1, 0] +// num_partitions = 2 +// data = [10, 20, 30, 40, 50] +// outputs[0] = [10, 20, 50] +// outputs[1] = [30, 40] // ``` // +// See `dynamic_stitch` for an example on how to merge partitions back. +// +//
+// +//
+// // Arguments: -// start: First entry in the range. -// stop: Last entry in the range. -// num: Number of values to generate. // -// Returns 1-D. The generated values. -func LinSpace(scope *Scope, start tf.Output, stop tf.Output, num tf.Output) (output tf.Output) { +// partitions: Any shape. Indices in the range `[0, num_partitions)`. +// num_partitions: The number of partitions to output. +func DynamicPartition(scope *Scope, data tf.Output, partitions tf.Output, num_partitions int64) (outputs []tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{"num_partitions": num_partitions} opspec := tf.OpSpec{ - Type: "LinSpace", + Type: "DynamicPartition", Input: []tf.Input{ - start, stop, num, + data, partitions, }, + Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + if scope.Err() != nil { + return + } + var idx int + var err error + if outputs, idx, err = makeOutputList(op, idx, "outputs"); err != nil { + scope.UpdateErr("DynamicPartition", err) + return + } + return outputs } -// DestroyResourceOpAttr is an optional argument to DestroyResourceOp. -type DestroyResourceOpAttr func(optionalAttr) +// ResourceApplyAdagradAttr is an optional argument to ResourceApplyAdagrad. +type ResourceApplyAdagradAttr func(optionalAttr) -// DestroyResourceOpIgnoreLookupError sets the optional ignore_lookup_error attribute to value. +// ResourceApplyAdagradUseLocking sets the optional use_locking attribute to value. // -// value: whether to ignore the error when the resource -// doesn't exist. -// If not specified, defaults to true -func DestroyResourceOpIgnoreLookupError(value bool) DestroyResourceOpAttr { +// value: If `True`, updating of the var and accum tensors will be protected +// by a lock; otherwise the behavior is undefined, but may exhibit less +// contention. +// If not specified, defaults to false +func ResourceApplyAdagradUseLocking(value bool) ResourceApplyAdagradAttr { return func(m optionalAttr) { - m["ignore_lookup_error"] = value + m["use_locking"] = value } } -// Deletes the resource specified by the handle. +// Update '*var' according to the adagrad scheme. // -// All subsequent operations using the resource will result in a NotFound -// error status. +// accum += grad * grad +// var -= lr * grad * (1 / sqrt(accum)) // // Arguments: -// resource: handle to the resource to delete. +// var_: Should be from a Variable(). +// accum: Should be from a Variable(). +// lr: Scaling factor. Must be a scalar. +// grad: The gradient. // // Returns the created operation. -func DestroyResourceOp(scope *Scope, resource tf.Output, optional ...DestroyResourceOpAttr) (o *tf.Operation) { +func ResourceApplyAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, grad tf.Output, optional ...ResourceApplyAdagradAttr) (o *tf.Operation) { if scope.Err() != nil { return } @@ -6932,9 +6877,9 @@ func DestroyResourceOp(scope *Scope, resource tf.Output, optional ...DestroyReso a(attrs) } opspec := tf.OpSpec{ - Type: "DestroyResourceOp", + Type: "ResourceApplyAdagrad", Input: []tf.Input{ - resource, + var_, accum, lr, grad, }, Attrs: attrs, } @@ -7825,63 +7770,6 @@ func ResourceScatterNdUpdate(scope *Scope, ref tf.Output, indices tf.Output, upd return scope.AddOperation(opspec) } -// StageSizeAttr is an optional argument to StageSize. -type StageSizeAttr func(optionalAttr) - -// StageSizeCapacity sets the optional capacity attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func StageSizeCapacity(value int64) StageSizeAttr { - return func(m optionalAttr) { - m["capacity"] = value - } -} - -// StageSizeMemoryLimit sets the optional memory_limit attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func StageSizeMemoryLimit(value int64) StageSizeAttr { - return func(m optionalAttr) { - m["memory_limit"] = value - } -} - -// StageSizeContainer sets the optional container attribute to value. -// If not specified, defaults to "" -func StageSizeContainer(value string) StageSizeAttr { - return func(m optionalAttr) { - m["container"] = value - } -} - -// StageSizeSharedName sets the optional shared_name attribute to value. -// If not specified, defaults to "" -func StageSizeSharedName(value string) StageSizeAttr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// Op returns the number of elements in the underlying container. -func StageSize(scope *Scope, dtypes []tf.DataType, optional ...StageSizeAttr) (size tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"dtypes": dtypes} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "StageSize", - - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // NonMaxSuppressionAttr is an optional argument to NonMaxSuppression. type NonMaxSuppressionAttr func(optionalAttr) @@ -8621,51 +8509,6 @@ func SparseReduceSum(scope *Scope, input_indices tf.Output, input_values tf.Outp return op.Output(0) } -// ResourceApplyAdagradAttr is an optional argument to ResourceApplyAdagrad. -type ResourceApplyAdagradAttr func(optionalAttr) - -// ResourceApplyAdagradUseLocking sets the optional use_locking attribute to value. -// -// value: If `True`, updating of the var and accum tensors will be protected -// by a lock; otherwise the behavior is undefined, but may exhibit less -// contention. -// If not specified, defaults to false -func ResourceApplyAdagradUseLocking(value bool) ResourceApplyAdagradAttr { - return func(m optionalAttr) { - m["use_locking"] = value - } -} - -// Update '*var' according to the adagrad scheme. -// -// accum += grad * grad -// var -= lr * grad * (1 / sqrt(accum)) -// -// Arguments: -// var_: Should be from a Variable(). -// accum: Should be from a Variable(). -// lr: Scaling factor. Must be a scalar. -// grad: The gradient. -// -// Returns the created operation. -func ResourceApplyAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, grad tf.Output, optional ...ResourceApplyAdagradAttr) (o *tf.Operation) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "ResourceApplyAdagrad", - Input: []tf.Input{ - var_, accum, lr, grad, - }, - Attrs: attrs, - } - return scope.AddOperation(opspec) -} - // Returns element-wise remainder of division. This emulates C semantics in that // // the result here is consistent with a truncating divide. E.g. `truncate(x / y) * @@ -8729,31 +8572,6 @@ func IRFFT2D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Out return op.Output(0) } -// Compute the pairwise cross product. -// -// `a` and `b` must be the same shape; they can either be simple 3-element vectors, -// or any shape where the innermost dimension is 3. In the latter case, each pair -// of corresponding 3-element vectors is cross-multiplied independently. -// -// Arguments: -// a: A tensor containing 3-element vectors. -// b: Another tensor, of same type and shape as `a`. -// -// Returns Pairwise cross product of the vectors in `a` and `b`. -func Cross(scope *Scope, a tf.Output, b tf.Output) (product tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Cross", - Input: []tf.Input{ - a, b, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // Transforms a vector of brain.Example protos (as strings) into typed tensors. // // Arguments: @@ -9128,11 +8946,39 @@ func SparseSoftmaxCrossEntropyWithLogits(scope *Scope, features tf.Output, label opspec := tf.OpSpec{ Type: "SparseSoftmaxCrossEntropyWithLogits", Input: []tf.Input{ - features, labels, + features, labels, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1) +} + +// Fast Fourier transform. +// +// Computes the 1-dimensional discrete Fourier transform over the inner-most +// dimension of `input`. +// +// Arguments: +// input: A complex64 tensor. +// +// Returns A complex64 tensor of the same shape as `input`. The inner-most +// dimension of `input` is replaced with its 1D Fourier transform. +// +// @compatibility(numpy) +// Equivalent to np.fft.fft +// @end_compatibility +func FFT(scope *Scope, input tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "FFT", + Input: []tf.Input{ + input, }, } op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1) + return op.Output(0) } // ResourceSparseApplyAdagradDAAttr is an optional argument to ResourceSparseApplyAdagradDA. @@ -9405,6 +9251,258 @@ func Conv3DBackpropInput(scope *Scope, input tf.Output, filter tf.Output, out_ba return op.Output(0) } +// ResourceApplyProximalAdagradAttr is an optional argument to ResourceApplyProximalAdagrad. +type ResourceApplyProximalAdagradAttr func(optionalAttr) + +// ResourceApplyProximalAdagradUseLocking sets the optional use_locking attribute to value. +// +// value: If True, updating of the var and accum tensors will be protected by +// a lock; otherwise the behavior is undefined, but may exhibit less contention. +// If not specified, defaults to false +func ResourceApplyProximalAdagradUseLocking(value bool) ResourceApplyProximalAdagradAttr { + return func(m optionalAttr) { + m["use_locking"] = value + } +} + +// Update '*var' and '*accum' according to FOBOS with Adagrad learning rate. +// +// accum += grad * grad +// prox_v = var - lr * grad * (1 / sqrt(accum)) +// var = sign(prox_v)/(1+lr*l2) * max{|prox_v|-lr*l1,0} +// +// Arguments: +// var_: Should be from a Variable(). +// accum: Should be from a Variable(). +// lr: Scaling factor. Must be a scalar. +// l1: L1 regularization. Must be a scalar. +// l2: L2 regularization. Must be a scalar. +// grad: The gradient. +// +// Returns the created operation. +func ResourceApplyProximalAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, grad tf.Output, optional ...ResourceApplyProximalAdagradAttr) (o *tf.Operation) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "ResourceApplyProximalAdagrad", + Input: []tf.Input{ + var_, accum, lr, l1, l2, grad, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) +} + +// MutableHashTableOfTensorsV2Attr is an optional argument to MutableHashTableOfTensorsV2. +type MutableHashTableOfTensorsV2Attr func(optionalAttr) + +// MutableHashTableOfTensorsV2Container sets the optional container attribute to value. +// +// value: If non-empty, this table is placed in the given container. +// Otherwise, a default container is used. +// If not specified, defaults to "" +func MutableHashTableOfTensorsV2Container(value string) MutableHashTableOfTensorsV2Attr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// MutableHashTableOfTensorsV2SharedName sets the optional shared_name attribute to value. +// +// value: If non-empty, this table is shared under the given name across +// multiple sessions. +// If not specified, defaults to "" +func MutableHashTableOfTensorsV2SharedName(value string) MutableHashTableOfTensorsV2Attr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// MutableHashTableOfTensorsV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value. +// If not specified, defaults to false +func MutableHashTableOfTensorsV2UseNodeNameSharing(value bool) MutableHashTableOfTensorsV2Attr { + return func(m optionalAttr) { + m["use_node_name_sharing"] = value + } +} + +// MutableHashTableOfTensorsV2ValueShape sets the optional value_shape attribute to value. +// If not specified, defaults to <> +func MutableHashTableOfTensorsV2ValueShape(value tf.Shape) MutableHashTableOfTensorsV2Attr { + return func(m optionalAttr) { + m["value_shape"] = value + } +} + +// Creates an empty hash table. +// +// This op creates a mutable hash table, specifying the type of its keys and +// values. Each value must be a vector. Data can be inserted into the table using +// the insert operations. It does not support the initialization operation. +// +// Arguments: +// key_dtype: Type of the table keys. +// value_dtype: Type of the table values. +// +// Returns Handle to a table. +func MutableHashTableOfTensorsV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...MutableHashTableOfTensorsV2Attr) (table_handle tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MutableHashTableOfTensorsV2", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// HashTableV2Attr is an optional argument to HashTableV2. +type HashTableV2Attr func(optionalAttr) + +// HashTableV2Container sets the optional container attribute to value. +// +// value: If non-empty, this table is placed in the given container. +// Otherwise, a default container is used. +// If not specified, defaults to "" +func HashTableV2Container(value string) HashTableV2Attr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// HashTableV2SharedName sets the optional shared_name attribute to value. +// +// value: If non-empty, this table is shared under the given name across +// multiple sessions. +// If not specified, defaults to "" +func HashTableV2SharedName(value string) HashTableV2Attr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// HashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value. +// +// value: If true and shared_name is empty, the table is shared +// using the node name. +// If not specified, defaults to false +func HashTableV2UseNodeNameSharing(value bool) HashTableV2Attr { + return func(m optionalAttr) { + m["use_node_name_sharing"] = value + } +} + +// Creates a non-initialized hash table. +// +// This op creates a hash table, specifying the type of its keys and values. +// Before using the table you will have to initialize it. After initialization the +// table will be immutable. +// +// Arguments: +// key_dtype: Type of the table keys. +// value_dtype: Type of the table values. +// +// Returns Handle to a table. +func HashTableV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...HashTableV2Attr) (table_handle tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "HashTableV2", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// MapUnstageNoKeyAttr is an optional argument to MapUnstageNoKey. +type MapUnstageNoKeyAttr func(optionalAttr) + +// MapUnstageNoKeyCapacity sets the optional capacity attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func MapUnstageNoKeyCapacity(value int64) MapUnstageNoKeyAttr { + return func(m optionalAttr) { + m["capacity"] = value + } +} + +// MapUnstageNoKeyMemoryLimit sets the optional memory_limit attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func MapUnstageNoKeyMemoryLimit(value int64) MapUnstageNoKeyAttr { + return func(m optionalAttr) { + m["memory_limit"] = value + } +} + +// MapUnstageNoKeyContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func MapUnstageNoKeyContainer(value string) MapUnstageNoKeyAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// MapUnstageNoKeySharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func MapUnstageNoKeySharedName(value string) MapUnstageNoKeyAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Op removes and returns a random (key, value) +// +// from the underlying container. If the underlying container +// does not contain elements, the op will block until it does. +func MapUnstageNoKey(scope *Scope, indices tf.Output, dtypes []tf.DataType, optional ...MapUnstageNoKeyAttr) (key tf.Output, values []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"dtypes": dtypes} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MapUnstageNoKey", + Input: []tf.Input{ + indices, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + key = op.Output(idx) + if values, idx, err = makeOutputList(op, idx, "values"); err != nil { + scope.UpdateErr("MapUnstageNoKey", err) + return + } + return key, values +} + // Inverse 2D fast Fourier transform. // // Computes the inverse 2-dimensional discrete Fourier transform over the @@ -11271,6 +11369,54 @@ func ResourceApplyMomentum(scope *Scope, var_ tf.Output, accum tf.Output, lr tf. return scope.AddOperation(opspec) } +// MaxPoolGradGradAttr is an optional argument to MaxPoolGradGrad. +type MaxPoolGradGradAttr func(optionalAttr) + +// MaxPoolGradGradDataFormat sets the optional data_format attribute to value. +// +// value: Specify the data format of the input and output data. With the +// default format "NHWC", the data is stored in the order of: +// [batch, in_height, in_width, in_channels]. +// Alternatively, the format could be "NCHW", the data storage order of: +// [batch, in_channels, in_height, in_width]. +// If not specified, defaults to "NHWC" +func MaxPoolGradGradDataFormat(value string) MaxPoolGradGradAttr { + return func(m optionalAttr) { + m["data_format"] = value + } +} + +// Computes second-order gradients of the maxpooling function. +// +// Arguments: +// orig_input: The original input tensor. +// orig_output: The original output tensor. +// grad: 4-D. Gradients of gradients w.r.t. the input of `max_pool`. +// ksize: The size of the window for each dimension of the input tensor. +// strides: The stride of the sliding window for each dimension of the +// input tensor. +// padding: The type of padding algorithm to use. +// +// Returns Gradients of gradients w.r.t. the input to `max_pool`. +func MaxPoolGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPoolGradGradAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MaxPoolGradGrad", + Input: []tf.Input{ + orig_input, orig_output, grad, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Returns the truth value of (x >= y) element-wise. // // *NOTE*: `GreaterEqual` supports broadcasting. More about broadcasting @@ -12053,6 +12199,63 @@ func MutableDenseHashTableV2(scope *Scope, empty_key tf.Output, value_dtype tf.D return op.Output(0) } +// StageSizeAttr is an optional argument to StageSize. +type StageSizeAttr func(optionalAttr) + +// StageSizeCapacity sets the optional capacity attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func StageSizeCapacity(value int64) StageSizeAttr { + return func(m optionalAttr) { + m["capacity"] = value + } +} + +// StageSizeMemoryLimit sets the optional memory_limit attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func StageSizeMemoryLimit(value int64) StageSizeAttr { + return func(m optionalAttr) { + m["memory_limit"] = value + } +} + +// StageSizeContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func StageSizeContainer(value string) StageSizeAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// StageSizeSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func StageSizeSharedName(value string) StageSizeAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Op returns the number of elements in the underlying container. +func StageSize(scope *Scope, dtypes []tf.DataType, optional ...StageSizeAttr) (size tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"dtypes": dtypes} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "StageSize", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Produces the max pool of the input tensor for quantized types. // // Arguments: @@ -12795,6 +12998,56 @@ func Neg(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// FakeQuantWithMinMaxVarsAttr is an optional argument to FakeQuantWithMinMaxVars. +type FakeQuantWithMinMaxVarsAttr func(optionalAttr) + +// FakeQuantWithMinMaxVarsNumBits sets the optional num_bits attribute to value. +// If not specified, defaults to 8 +func FakeQuantWithMinMaxVarsNumBits(value int64) FakeQuantWithMinMaxVarsAttr { + return func(m optionalAttr) { + m["num_bits"] = value + } +} + +// FakeQuantWithMinMaxVarsNarrowRange sets the optional narrow_range attribute to value. +// If not specified, defaults to false +func FakeQuantWithMinMaxVarsNarrowRange(value bool) FakeQuantWithMinMaxVarsAttr { + return func(m optionalAttr) { + m["narrow_range"] = value + } +} + +// Fake-quantize the 'inputs' tensor of type float via global float scalars `min` +// +// and `max` to 'outputs' tensor of same shape as `inputs`. +// +// `[min; max]` define the clamping range for the `inputs` data. +// `inputs` values are quantized into the quantization range (`[0; 2^num_bits - 1]` +// when `narrow_range` is false and `[1; 2^num_bits - 1]` when it is true) and +// then de-quantized and output as floats in `[min; max]` interval. +// `num_bits` is the bitwidth of the quantization; between 2 and 8, inclusive. +// +// This operation has a gradient and thus allows for training `min` and `max` +// values. +func FakeQuantWithMinMaxVars(scope *Scope, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsAttr) (outputs tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "FakeQuantWithMinMaxVars", + Input: []tf.Input{ + inputs, min, max, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Writes a `Summary` protocol buffer with a histogram. // // The generated @@ -14828,54 +15081,6 @@ func TensorArrayCloseV3(scope *Scope, handle tf.Output) (o *tf.Operation) { return scope.AddOperation(opspec) } -// MaxPoolGradGradAttr is an optional argument to MaxPoolGradGrad. -type MaxPoolGradGradAttr func(optionalAttr) - -// MaxPoolGradGradDataFormat sets the optional data_format attribute to value. -// -// value: Specify the data format of the input and output data. With the -// default format "NHWC", the data is stored in the order of: -// [batch, in_height, in_width, in_channels]. -// Alternatively, the format could be "NCHW", the data storage order of: -// [batch, in_channels, in_height, in_width]. -// If not specified, defaults to "NHWC" -func MaxPoolGradGradDataFormat(value string) MaxPoolGradGradAttr { - return func(m optionalAttr) { - m["data_format"] = value - } -} - -// Computes second-order gradients of the maxpooling function. -// -// Arguments: -// orig_input: The original input tensor. -// orig_output: The original output tensor. -// grad: 4-D. Gradients of gradients w.r.t. the input of `max_pool`. -// ksize: The size of the window for each dimension of the input tensor. -// strides: The stride of the sliding window for each dimension of the -// input tensor. -// padding: The type of padding algorithm to use. -// -// Returns Gradients of gradients w.r.t. the input to `max_pool`. -func MaxPoolGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPoolGradGradAttr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "MaxPoolGradGrad", - Input: []tf.Input{ - orig_input, orig_output, grad, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // RandomUniformIntAttr is an optional argument to RandomUniformInt. type RandomUniformIntAttr func(optionalAttr) @@ -15136,60 +15341,9 @@ func TruncatedNormal(scope *Scope, shape tf.Output, dtype tf.DataType, optional a(attrs) } opspec := tf.OpSpec{ - Type: "TruncatedNormal", - Input: []tf.Input{ - shape, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// FakeQuantWithMinMaxVarsPerChannelAttr is an optional argument to FakeQuantWithMinMaxVarsPerChannel. -type FakeQuantWithMinMaxVarsPerChannelAttr func(optionalAttr) - -// FakeQuantWithMinMaxVarsPerChannelNumBits sets the optional num_bits attribute to value. -// If not specified, defaults to 8 -func FakeQuantWithMinMaxVarsPerChannelNumBits(value int64) FakeQuantWithMinMaxVarsPerChannelAttr { - return func(m optionalAttr) { - m["num_bits"] = value - } -} - -// FakeQuantWithMinMaxVarsPerChannelNarrowRange sets the optional narrow_range attribute to value. -// If not specified, defaults to false -func FakeQuantWithMinMaxVarsPerChannelNarrowRange(value bool) FakeQuantWithMinMaxVarsPerChannelAttr { - return func(m optionalAttr) { - m["narrow_range"] = value - } -} - -// Fake-quantize the 'inputs' tensor of type float and one of the shapes: `[d]`, -// -// `[b, d]` `[b, h, w, d]` via per-channel floats `min` and `max` of shape `[d]` -// to 'outputs' tensor of same shape as `inputs`. -// -// `[min; max]` define the clamping range for the `inputs` data. -// `inputs` values are quantized into the quantization range (`[0; 2^num_bits - 1]` -// when `narrow_range` is false and `[1; 2^num_bits - 1]` when it is true) and -// then de-quantized and output as floats in `[min; max]` interval. -// `num_bits` is the bitwidth of the quantization; between 2 and 8, inclusive. -// -// This operation has a gradient and thus allows for training `min` and `max` -// values. -func FakeQuantWithMinMaxVarsPerChannel(scope *Scope, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsPerChannelAttr) (outputs tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxVarsPerChannel", + Type: "TruncatedNormal", Input: []tf.Input{ - inputs, min, max, + shape, }, Attrs: attrs, } @@ -16709,6 +16863,70 @@ func ResourceApplyPowerSign(scope *Scope, var_ tf.Output, m tf.Output, lr tf.Out return scope.AddOperation(opspec) } +// MutableHashTableV2Attr is an optional argument to MutableHashTableV2. +type MutableHashTableV2Attr func(optionalAttr) + +// MutableHashTableV2Container sets the optional container attribute to value. +// +// value: If non-empty, this table is placed in the given container. +// Otherwise, a default container is used. +// If not specified, defaults to "" +func MutableHashTableV2Container(value string) MutableHashTableV2Attr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// MutableHashTableV2SharedName sets the optional shared_name attribute to value. +// +// value: If non-empty, this table is shared under the given name across +// multiple sessions. +// If not specified, defaults to "" +func MutableHashTableV2SharedName(value string) MutableHashTableV2Attr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// MutableHashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value. +// +// value: If true and shared_name is empty, the table is shared +// using the node name. +// If not specified, defaults to false +func MutableHashTableV2UseNodeNameSharing(value bool) MutableHashTableV2Attr { + return func(m optionalAttr) { + m["use_node_name_sharing"] = value + } +} + +// Creates an empty hash table. +// +// This op creates a mutable hash table, specifying the type of its keys and +// values. Each value must be a scalar. Data can be inserted into the table using +// the insert operations. It does not support the initialization operation. +// +// Arguments: +// key_dtype: Type of the table keys. +// value_dtype: Type of the table values. +// +// Returns Handle to a table. +func MutableHashTableV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...MutableHashTableV2Attr) (table_handle tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MutableHashTableV2", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Deprecated. Disallowed in GraphDef version >= 2. // // DEPRECATED at GraphDef version 2: Use AdjustContrastv2 instead @@ -17594,23 +17812,6 @@ func SoftplusGrad(scope *Scope, gradients tf.Output, features tf.Output) (backpr return op.Output(0) } -// Creates a dataset that contains the unique elements of `input_dataset`. -func UniqueDataset(scope *Scope, input_dataset tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes} - opspec := tf.OpSpec{ - Type: "UniqueDataset", - Input: []tf.Input{ - input_dataset, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // SelfAdjointEigV2Attr is an optional argument to SelfAdjointEigV2. type SelfAdjointEigV2Attr func(optionalAttr) @@ -19855,130 +20056,38 @@ func SparseTensorSliceDataset(scope *Scope, indices tf.Output, values tf.Output, return op.Output(0) } -// Creates a dataset that concatenates `input_dataset` with `another_dataset`. -func ConcatenateDataset(scope *Scope, input_dataset tf.Output, another_dataset tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes} - opspec := tf.OpSpec{ - Type: "ConcatenateDataset", - Input: []tf.Input{ - input_dataset, another_dataset, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Creates a dataset that contains the elements of `input_dataset` ignoring errors. -func IgnoreErrorsDataset(scope *Scope, input_dataset tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes} - opspec := tf.OpSpec{ - Type: "IgnoreErrorsDataset", - Input: []tf.Input{ - input_dataset, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// CropAndResizeGradImageAttr is an optional argument to CropAndResizeGradImage. -type CropAndResizeGradImageAttr func(optionalAttr) - -// CropAndResizeGradImageMethod sets the optional method attribute to value. -// -// value: A string specifying the interpolation method. Only 'bilinear' is -// supported for now. -// If not specified, defaults to "bilinear" -func CropAndResizeGradImageMethod(value string) CropAndResizeGradImageAttr { - return func(m optionalAttr) { - m["method"] = value - } -} - -// Computes the gradient of the crop_and_resize op wrt the input image tensor. -// -// Arguments: -// grads: A 4-D tensor of shape `[num_boxes, crop_height, crop_width, depth]`. -// boxes: A 2-D tensor of shape `[num_boxes, 4]`. The `i`-th row of the tensor -// specifies the coordinates of a box in the `box_ind[i]` image and is specified -// in normalized coordinates `[y1, x1, y2, x2]`. A normalized coordinate value of -// `y` is mapped to the image coordinate at `y * (image_height - 1)`, so as the -// `[0, 1]` interval of normalized image height is mapped to -// `[0, image_height - 1] in image height coordinates. We do allow y1 > y2, in -// which case the sampled crop is an up-down flipped version of the original -// image. The width dimension is treated similarly. Normalized coordinates -// outside the `[0, 1]` range are allowed, in which case we use -// `extrapolation_value` to extrapolate the input image values. -// box_ind: A 1-D tensor of shape `[num_boxes]` with int32 values in `[0, batch)`. -// The value of `box_ind[i]` specifies the image that the `i`-th box refers to. -// image_size: A 1-D tensor with value `[batch, image_height, image_width, depth]` -// containing the original image size. Both `image_height` and `image_width` need -// to be positive. +// Returns x / y element-wise for real types. // +// If `x` and `y` are reals, this will return the floating-point division. // -// Returns A 4-D tensor of shape `[batch, image_height, image_width, depth]`. -func CropAndResizeGradImage(scope *Scope, grads tf.Output, boxes tf.Output, box_ind tf.Output, image_size tf.Output, T tf.DataType, optional ...CropAndResizeGradImageAttr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"T": T} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "CropAndResizeGradImage", - Input: []tf.Input{ - grads, boxes, box_ind, image_size, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Reads and outputs the entire contents of the input filename. -func ReadFile(scope *Scope, filename tf.Output) (contents tf.Output) { +// *NOTE*: `Div` supports broadcasting. More about broadcasting +// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) +func RealDiv(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "ReadFile", + Type: "RealDiv", Input: []tf.Input{ - filename, + x, y, }, } op := scope.AddOperation(opspec) return op.Output(0) } -// Concatenates tensors along one dimension. -// -// Arguments: -// values: List of `N` Tensors to concatenate. Their ranks and types must match, -// and their sizes must match in all dimensions except `concat_dim`. -// axis: 0-D. The dimension along which to concatenate. Must be in the -// range [-rank(values), rank(values)). -// -// Returns A `Tensor` with the concatenation of values stacked along the -// `concat_dim` dimension. This tensor's shape matches that of `values` except -// in `concat_dim` where it has the sum of the sizes. -func ConcatV2(scope *Scope, values []tf.Output, axis tf.Output) (output tf.Output) { +// Creates a dataset that concatenates `input_dataset` with `another_dataset`. +func ConcatenateDataset(scope *Scope, input_dataset tf.Output, another_dataset tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes} opspec := tf.OpSpec{ - Type: "ConcatV2", + Type: "ConcatenateDataset", Input: []tf.Input{ - tf.OutputList(values), axis, + input_dataset, another_dataset, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) @@ -20698,70 +20807,8 @@ func TFRecordDataset(scope *Scope, filenames tf.Output, compression_type tf.Outp opspec := tf.OpSpec{ Type: "TFRecordDataset", Input: []tf.Input{ - filenames, compression_type, buffer_size, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// FakeQuantWithMinMaxArgsGradientAttr is an optional argument to FakeQuantWithMinMaxArgsGradient. -type FakeQuantWithMinMaxArgsGradientAttr func(optionalAttr) - -// FakeQuantWithMinMaxArgsGradientMin sets the optional min attribute to value. -// If not specified, defaults to -6 -func FakeQuantWithMinMaxArgsGradientMin(value float32) FakeQuantWithMinMaxArgsGradientAttr { - return func(m optionalAttr) { - m["min"] = value - } -} - -// FakeQuantWithMinMaxArgsGradientMax sets the optional max attribute to value. -// If not specified, defaults to 6 -func FakeQuantWithMinMaxArgsGradientMax(value float32) FakeQuantWithMinMaxArgsGradientAttr { - return func(m optionalAttr) { - m["max"] = value - } -} - -// FakeQuantWithMinMaxArgsGradientNumBits sets the optional num_bits attribute to value. -// If not specified, defaults to 8 -func FakeQuantWithMinMaxArgsGradientNumBits(value int64) FakeQuantWithMinMaxArgsGradientAttr { - return func(m optionalAttr) { - m["num_bits"] = value - } -} - -// FakeQuantWithMinMaxArgsGradientNarrowRange sets the optional narrow_range attribute to value. -// If not specified, defaults to false -func FakeQuantWithMinMaxArgsGradientNarrowRange(value bool) FakeQuantWithMinMaxArgsGradientAttr { - return func(m optionalAttr) { - m["narrow_range"] = value - } -} - -// Compute gradients for a FakeQuantWithMinMaxArgs operation. -// -// Arguments: -// gradients: Backpropagated gradients above the FakeQuantWithMinMaxArgs operation. -// inputs: Values passed as inputs to the FakeQuantWithMinMaxArgs operation. -// -// Returns Backpropagated gradients below the FakeQuantWithMinMaxArgs operation: -// `gradients * (inputs >= min && inputs <= max)`. -func FakeQuantWithMinMaxArgsGradient(scope *Scope, gradients tf.Output, inputs tf.Output, optional ...FakeQuantWithMinMaxArgsGradientAttr) (backprops tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxArgsGradient", - Input: []tf.Input{ - gradients, inputs, + filenames, compression_type, buffer_size, }, - Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) @@ -21290,6 +21337,31 @@ func StatsAggregatorSummary(scope *Scope, iterator tf.Output) (summary tf.Output return op.Output(0) } +// Compute the pairwise cross product. +// +// `a` and `b` must be the same shape; they can either be simple 3-element vectors, +// or any shape where the innermost dimension is 3. In the latter case, each pair +// of corresponding 3-element vectors is cross-multiplied independently. +// +// Arguments: +// a: A tensor containing 3-element vectors. +// b: Another tensor, of same type and shape as `a`. +// +// Returns Pairwise cross product of the vectors in `a` and `b`. +func Cross(scope *Scope, a tf.Output, b tf.Output) (product tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Cross", + Input: []tf.Input{ + a, b, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Performs a padding as a preprocess during a convolution. // // Similar to FusedResizeAndPadConv2d, this op allows for an optimized @@ -22254,6 +22326,171 @@ func TensorArrayCloseV2(scope *Scope, handle tf.Output) (o *tf.Operation) { return scope.AddOperation(opspec) } +// CropAndResizeGradImageAttr is an optional argument to CropAndResizeGradImage. +type CropAndResizeGradImageAttr func(optionalAttr) + +// CropAndResizeGradImageMethod sets the optional method attribute to value. +// +// value: A string specifying the interpolation method. Only 'bilinear' is +// supported for now. +// If not specified, defaults to "bilinear" +func CropAndResizeGradImageMethod(value string) CropAndResizeGradImageAttr { + return func(m optionalAttr) { + m["method"] = value + } +} + +// Computes the gradient of the crop_and_resize op wrt the input image tensor. +// +// Arguments: +// grads: A 4-D tensor of shape `[num_boxes, crop_height, crop_width, depth]`. +// boxes: A 2-D tensor of shape `[num_boxes, 4]`. The `i`-th row of the tensor +// specifies the coordinates of a box in the `box_ind[i]` image and is specified +// in normalized coordinates `[y1, x1, y2, x2]`. A normalized coordinate value of +// `y` is mapped to the image coordinate at `y * (image_height - 1)`, so as the +// `[0, 1]` interval of normalized image height is mapped to +// `[0, image_height - 1] in image height coordinates. We do allow y1 > y2, in +// which case the sampled crop is an up-down flipped version of the original +// image. The width dimension is treated similarly. Normalized coordinates +// outside the `[0, 1]` range are allowed, in which case we use +// `extrapolation_value` to extrapolate the input image values. +// box_ind: A 1-D tensor of shape `[num_boxes]` with int32 values in `[0, batch)`. +// The value of `box_ind[i]` specifies the image that the `i`-th box refers to. +// image_size: A 1-D tensor with value `[batch, image_height, image_width, depth]` +// containing the original image size. Both `image_height` and `image_width` need +// to be positive. +// +// +// Returns A 4-D tensor of shape `[batch, image_height, image_width, depth]`. +func CropAndResizeGradImage(scope *Scope, grads tf.Output, boxes tf.Output, box_ind tf.Output, image_size tf.Output, T tf.DataType, optional ...CropAndResizeGradImageAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"T": T} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "CropAndResizeGradImage", + Input: []tf.Input{ + grads, boxes, box_ind, image_size, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Reads and outputs the entire contents of the input filename. +func ReadFile(scope *Scope, filename tf.Output) (contents tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ReadFile", + Input: []tf.Input{ + filename, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Concatenates tensors along one dimension. +// +// Arguments: +// values: List of `N` Tensors to concatenate. Their ranks and types must match, +// and their sizes must match in all dimensions except `concat_dim`. +// axis: 0-D. The dimension along which to concatenate. Must be in the +// range [-rank(values), rank(values)). +// +// Returns A `Tensor` with the concatenation of values stacked along the +// `concat_dim` dimension. This tensor's shape matches that of `values` except +// in `concat_dim` where it has the sum of the sizes. +func ConcatV2(scope *Scope, values []tf.Output, axis tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ConcatV2", + Input: []tf.Input{ + tf.OutputList(values), axis, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Forwards the value of an available tensor from `inputs` to `output`. +// +// `Merge` waits for at least one of the tensors in `inputs` to become available. +// It is usually combined with `Switch` to implement branching. +// +// `Merge` forwards the first tensor to become available to `output`, and sets +// `value_index` to its index in `inputs`. +// +// Arguments: +// inputs: The input tensors, exactly one of which will become available. +// +// Returns Will be set to the available input tensor.The index of the chosen input tensor in `inputs`. +func Merge(scope *Scope, inputs []tf.Output) (output tf.Output, value_index tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Merge", + Input: []tf.Input{ + tf.OutputList(inputs), + }, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1) +} + +// QueueCloseV2Attr is an optional argument to QueueCloseV2. +type QueueCloseV2Attr func(optionalAttr) + +// QueueCloseV2CancelPendingEnqueues sets the optional cancel_pending_enqueues attribute to value. +// +// value: If true, all pending enqueue requests that are +// blocked on the given queue will be canceled. +// If not specified, defaults to false +func QueueCloseV2CancelPendingEnqueues(value bool) QueueCloseV2Attr { + return func(m optionalAttr) { + m["cancel_pending_enqueues"] = value + } +} + +// Closes the given queue. +// +// This operation signals that no more elements will be enqueued in the +// given queue. Subsequent Enqueue(Many) operations will fail. +// Subsequent Dequeue(Many) operations will continue to succeed if +// sufficient elements remain in the queue. Subsequent Dequeue(Many) +// operations that would block will fail immediately. +// +// Arguments: +// handle: The handle to a queue. +// +// Returns the created operation. +func QueueCloseV2(scope *Scope, handle tf.Output, optional ...QueueCloseV2Attr) (o *tf.Operation) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "QueueCloseV2", + Input: []tf.Input{ + handle, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) +} + // Computes inverse hyperbolic tangent of x element-wise. func Atanh(scope *Scope, x tf.Output) (y tf.Output) { if scope.Err() != nil { @@ -23708,81 +23945,18 @@ func Conv2DDilations(value []int64) Conv2DAttr { // // Returns A 4-D tensor. The dimension order is determined by the value of // `data_format`, see below for details. -func Conv2D(scope *Scope, input tf.Output, filter tf.Output, strides []int64, padding string, optional ...Conv2DAttr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"strides": strides, "padding": padding} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "Conv2D", - Input: []tf.Input{ - input, filter, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// FakeQuantWithMinMaxArgsAttr is an optional argument to FakeQuantWithMinMaxArgs. -type FakeQuantWithMinMaxArgsAttr func(optionalAttr) - -// FakeQuantWithMinMaxArgsMin sets the optional min attribute to value. -// If not specified, defaults to -6 -func FakeQuantWithMinMaxArgsMin(value float32) FakeQuantWithMinMaxArgsAttr { - return func(m optionalAttr) { - m["min"] = value - } -} - -// FakeQuantWithMinMaxArgsMax sets the optional max attribute to value. -// If not specified, defaults to 6 -func FakeQuantWithMinMaxArgsMax(value float32) FakeQuantWithMinMaxArgsAttr { - return func(m optionalAttr) { - m["max"] = value - } -} - -// FakeQuantWithMinMaxArgsNumBits sets the optional num_bits attribute to value. -// If not specified, defaults to 8 -func FakeQuantWithMinMaxArgsNumBits(value int64) FakeQuantWithMinMaxArgsAttr { - return func(m optionalAttr) { - m["num_bits"] = value - } -} - -// FakeQuantWithMinMaxArgsNarrowRange sets the optional narrow_range attribute to value. -// If not specified, defaults to false -func FakeQuantWithMinMaxArgsNarrowRange(value bool) FakeQuantWithMinMaxArgsAttr { - return func(m optionalAttr) { - m["narrow_range"] = value - } -} - -// Fake-quantize the 'inputs' tensor, type float to 'outputs' tensor of same type. -// -// Attributes `[min; max]` define the clamping range for the `inputs` data. -// `inputs` values are quantized into the quantization range (`[0; 2^num_bits - 1]` -// when `narrow_range` is false and `[1; 2^num_bits - 1]` when it is true) and -// then de-quantized and output as floats in `[min; max]` interval. -// `num_bits` is the bitwidth of the quantization; between 2 and 8, inclusive. -// -// Quantization is called fake since the output is still in floating point. -func FakeQuantWithMinMaxArgs(scope *Scope, inputs tf.Output, optional ...FakeQuantWithMinMaxArgsAttr) (outputs tf.Output) { +func Conv2D(scope *Scope, input tf.Output, filter tf.Output, strides []int64, padding string, optional ...Conv2DAttr) (output tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{} + attrs := map[string]interface{}{"strides": strides, "padding": padding} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxArgs", + Type: "Conv2D", Input: []tf.Input{ - inputs, + input, filter, }, Attrs: attrs, } @@ -24203,147 +24377,6 @@ func MapStage(scope *Scope, key tf.Output, indices tf.Output, values []tf.Output return scope.AddOperation(opspec) } -// MapPeekAttr is an optional argument to MapPeek. -type MapPeekAttr func(optionalAttr) - -// MapPeekCapacity sets the optional capacity attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func MapPeekCapacity(value int64) MapPeekAttr { - return func(m optionalAttr) { - m["capacity"] = value - } -} - -// MapPeekMemoryLimit sets the optional memory_limit attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func MapPeekMemoryLimit(value int64) MapPeekAttr { - return func(m optionalAttr) { - m["memory_limit"] = value - } -} - -// MapPeekContainer sets the optional container attribute to value. -// If not specified, defaults to "" -func MapPeekContainer(value string) MapPeekAttr { - return func(m optionalAttr) { - m["container"] = value - } -} - -// MapPeekSharedName sets the optional shared_name attribute to value. -// If not specified, defaults to "" -func MapPeekSharedName(value string) MapPeekAttr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// Op peeks at the values at the specified key. If the -// -// underlying container does not contain this key -// this op will block until it does. -func MapPeek(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.DataType, optional ...MapPeekAttr) (values []tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"dtypes": dtypes} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "MapPeek", - Input: []tf.Input{ - key, indices, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - if scope.Err() != nil { - return - } - var idx int - var err error - if values, idx, err = makeOutputList(op, idx, "values"); err != nil { - scope.UpdateErr("MapPeek", err) - return - } - return values -} - -// QueueCloseV2Attr is an optional argument to QueueCloseV2. -type QueueCloseV2Attr func(optionalAttr) - -// QueueCloseV2CancelPendingEnqueues sets the optional cancel_pending_enqueues attribute to value. -// -// value: If true, all pending enqueue requests that are -// blocked on the given queue will be canceled. -// If not specified, defaults to false -func QueueCloseV2CancelPendingEnqueues(value bool) QueueCloseV2Attr { - return func(m optionalAttr) { - m["cancel_pending_enqueues"] = value - } -} - -// Closes the given queue. -// -// This operation signals that no more elements will be enqueued in the -// given queue. Subsequent Enqueue(Many) operations will fail. -// Subsequent Dequeue(Many) operations will continue to succeed if -// sufficient elements remain in the queue. Subsequent Dequeue(Many) -// operations that would block will fail immediately. -// -// Arguments: -// handle: The handle to a queue. -// -// Returns the created operation. -func QueueCloseV2(scope *Scope, handle tf.Output, optional ...QueueCloseV2Attr) (o *tf.Operation) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "QueueCloseV2", - Input: []tf.Input{ - handle, - }, - Attrs: attrs, - } - return scope.AddOperation(opspec) -} - -// Forwards the value of an available tensor from `inputs` to `output`. -// -// `Merge` waits for at least one of the tensors in `inputs` to become available. -// It is usually combined with `Switch` to implement branching. -// -// `Merge` forwards the first tensor to become available to `output`, and sets -// `value_index` to its index in `inputs`. -// -// Arguments: -// inputs: The input tensors, exactly one of which will become available. -// -// Returns Will be set to the available input tensor.The index of the chosen input tensor in `inputs`. -func Merge(scope *Scope, inputs []tf.Output) (output tf.Output, value_index tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Merge", - Input: []tf.Input{ - tf.OutputList(inputs), - }, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1) -} - // MapUnstageAttr is an optional argument to MapUnstage. type MapUnstageAttr func(optionalAttr) @@ -28247,111 +28280,3 @@ func QuantizedInstanceNorm(scope *Scope, x tf.Output, x_min tf.Output, x_max tf. op := scope.AddOperation(opspec) return op.Output(0), op.Output(1), op.Output(2) } - -// FakeQuantWithMinMaxVarsAttr is an optional argument to FakeQuantWithMinMaxVars. -type FakeQuantWithMinMaxVarsAttr func(optionalAttr) - -// FakeQuantWithMinMaxVarsNumBits sets the optional num_bits attribute to value. -// If not specified, defaults to 8 -func FakeQuantWithMinMaxVarsNumBits(value int64) FakeQuantWithMinMaxVarsAttr { - return func(m optionalAttr) { - m["num_bits"] = value - } -} - -// FakeQuantWithMinMaxVarsNarrowRange sets the optional narrow_range attribute to value. -// If not specified, defaults to false -func FakeQuantWithMinMaxVarsNarrowRange(value bool) FakeQuantWithMinMaxVarsAttr { - return func(m optionalAttr) { - m["narrow_range"] = value - } -} - -// Fake-quantize the 'inputs' tensor of type float via global float scalars `min` -// -// and `max` to 'outputs' tensor of same shape as `inputs`. -// -// `[min; max]` define the clamping range for the `inputs` data. -// `inputs` values are quantized into the quantization range (`[0; 2^num_bits - 1]` -// when `narrow_range` is false and `[1; 2^num_bits - 1]` when it is true) and -// then de-quantized and output as floats in `[min; max]` interval. -// `num_bits` is the bitwidth of the quantization; between 2 and 8, inclusive. -// -// This operation has a gradient and thus allows for training `min` and `max` -// values. -func FakeQuantWithMinMaxVars(scope *Scope, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsAttr) (outputs tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxVars", - Input: []tf.Input{ - inputs, min, max, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// FakeQuantWithMinMaxVarsPerChannelGradientAttr is an optional argument to FakeQuantWithMinMaxVarsPerChannelGradient. -type FakeQuantWithMinMaxVarsPerChannelGradientAttr func(optionalAttr) - -// FakeQuantWithMinMaxVarsPerChannelGradientNumBits sets the optional num_bits attribute to value. -// -// value: The bitwidth of the quantization; between 2 and 8, inclusive. -// If not specified, defaults to 8 -func FakeQuantWithMinMaxVarsPerChannelGradientNumBits(value int64) FakeQuantWithMinMaxVarsPerChannelGradientAttr { - return func(m optionalAttr) { - m["num_bits"] = value - } -} - -// FakeQuantWithMinMaxVarsPerChannelGradientNarrowRange sets the optional narrow_range attribute to value. -// -// value: Whether to quantize into 2^num_bits - 1 distinct values. -// If not specified, defaults to false -func FakeQuantWithMinMaxVarsPerChannelGradientNarrowRange(value bool) FakeQuantWithMinMaxVarsPerChannelGradientAttr { - return func(m optionalAttr) { - m["narrow_range"] = value - } -} - -// Compute gradients for a FakeQuantWithMinMaxVarsPerChannel operation. -// -// Arguments: -// gradients: Backpropagated gradients above the FakeQuantWithMinMaxVars operation, -// shape one of: `[d]`, `[b, d]`, `[b, h, w, d]`. -// inputs: Values passed as inputs to the FakeQuantWithMinMaxVars operation, shape -// same as `gradients`. -// min, max: Quantization interval, floats of shape `[d]`. -// -// -// -// Returns Backpropagated gradients w.r.t. inputs, shape same as -// `inputs`: -// `gradients * (inputs >= min && inputs <= max)`.Backpropagated gradients w.r.t. min parameter, shape `[d]`: -// `sum_per_d(gradients * (inputs < min))`.Backpropagated gradients w.r.t. max parameter, shape `[d]`: -// `sum_per_d(gradients * (inputs > max))`. -func FakeQuantWithMinMaxVarsPerChannelGradient(scope *Scope, gradients tf.Output, inputs tf.Output, min tf.Output, max tf.Output, optional ...FakeQuantWithMinMaxVarsPerChannelGradientAttr) (backprops_wrt_input tf.Output, backprop_wrt_min tf.Output, backprop_wrt_max tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxVarsPerChannelGradient", - Input: []tf.Input{ - gradients, inputs, min, max, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) -} diff --git a/tensorflow/java/maven/libtensorflow/pom.xml b/tensorflow/java/maven/libtensorflow/pom.xml index a9ce5372aeb32b6957359fdcaa9da01c732c9f9f..7f3a83b195308b3261049c2845d018280dd61601 100644 --- a/tensorflow/java/maven/libtensorflow/pom.xml +++ b/tensorflow/java/maven/libtensorflow/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.5.0 + 1.7.0-rc0 ../ libtensorflow diff --git a/tensorflow/java/maven/libtensorflow_jni/pom.xml b/tensorflow/java/maven/libtensorflow_jni/pom.xml index fe34ca83ff30373fa4f3c4f345323bad40a8754e..cc436ff8409a5a17d69dbe0fa0ce9c6333349987 100644 --- a/tensorflow/java/maven/libtensorflow_jni/pom.xml +++ b/tensorflow/java/maven/libtensorflow_jni/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.5.0 + 1.7.0-rc0 ../ libtensorflow_jni diff --git a/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml b/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml index 390152808eb0c4abebe093ed5db39faf37fcafe3..47f678382aad910559395cd0b304321c30c92ede 100644 --- a/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml +++ b/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.5.0 + 1.7.0-rc0 ../ libtensorflow_jni_gpu diff --git a/tensorflow/java/maven/pom.xml b/tensorflow/java/maven/pom.xml index 524ec45f48bb91d09dfb5fca3cc19256d45587fb..42d32810a263df7ba1044e1ff3e7c4e8124d478e 100644 --- a/tensorflow/java/maven/pom.xml +++ b/tensorflow/java/maven/pom.xml @@ -6,7 +6,7 @@ 4.0.0 org.tensorflow parentpom - 1.5.0 + 1.7.0-rc0 pom https://www.tensorflow.org diff --git a/tensorflow/java/maven/proto/pom.xml b/tensorflow/java/maven/proto/pom.xml index 9cf3217f51f73184a02a58ead1a2735c5a44fd26..463893ce62dc7e72f4cd1c909ae56d097d060035 100644 --- a/tensorflow/java/maven/proto/pom.xml +++ b/tensorflow/java/maven/proto/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.5.0 + 1.7.0-rc0 ../ proto diff --git a/tensorflow/java/maven/tensorflow-android/pom-android.xml.template b/tensorflow/java/maven/tensorflow-android/pom-android.xml.template index 5cbd0c898dc52ec5dfb72f0a2ac893d492a7d4be..37d2372d7b09f6f144e7abb145cb75bf98356615 100644 --- a/tensorflow/java/maven/tensorflow-android/pom-android.xml.template +++ b/tensorflow/java/maven/tensorflow-android/pom-android.xml.template @@ -20,10 +20,8 @@ UTF-8 - ${build_number} ${build_commit_id} ${build_type} - ${build_url} diff --git a/tensorflow/java/maven/tensorflow-android/update.py b/tensorflow/java/maven/tensorflow-android/update.py index 7c250718347f5fdd65aaf8003aad75a87a19c96a..2206d800ca1fe82c5596ff39e56518bc5aea6211 100644 --- a/tensorflow/java/maven/tensorflow-android/update.py +++ b/tensorflow/java/maven/tensorflow-android/update.py @@ -45,6 +45,9 @@ def get_json(url): def get_commit_id(build_info): """Fetch the git commit id from the build info json object.""" + release_commit_id = build_info.get('build_commit_id') + if release_commit_id: + return release_commit_id actions = build_info.get('actions') build_data = next( a for a in actions @@ -95,20 +98,12 @@ def main(): release_prefix = 'https://storage.googleapis.com/tensorflow/libtensorflow' info_url = '%s/android_buildinfo-%s.json' % (release_prefix, args.version) aar_url = '%s/tensorflow-%s.aar' % (release_prefix, args.version) - build_type = 'release-matrix-android' + build_type = 'release-android' # Retrieve build information build_info = get_json(info_url) # Check all required build info is present - if build_info.get('result') != 'SUCCESS': - raise ValueError('Invalid json: %s' % build_info) - build_url = build_info.get('url') - if not build_url: - raise ValueError('Missing url: %s' % build_info) - build_number = build_info.get('number') - if not build_number: - raise ValueError('Missing build number: %s' % build_info) build_commit_id = get_commit_id(build_info) if not build_commit_id: raise ValueError('Missing commit id: %s' % build_info) @@ -119,9 +114,7 @@ def main(): f.write( template.substitute({ 'build_commit_id': build_commit_id, - 'build_number': build_number, 'build_type': build_type, - 'build_url': build_url, 'version': args.version })) diff --git a/tensorflow/java/maven/tensorflow/pom.xml b/tensorflow/java/maven/tensorflow/pom.xml index d619f986a9a03ac67f5de6bbe80e686a05ce5d42..60e7f3c199f1451a054f6a779a63a98ff428bd87 100644 --- a/tensorflow/java/maven/tensorflow/pom.xml +++ b/tensorflow/java/maven/tensorflow/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.5.0 + 1.7.0-rc0 ../ tensorflow diff --git a/tensorflow/java/src/main/java/org/tensorflow/NativeLibrary.java b/tensorflow/java/src/main/java/org/tensorflow/NativeLibrary.java index 499757e8cf4d6166e425d801ce20335bd8ad83e8..cf773e1686dea97f62f432be43f2c10b69fa8e24 100644 --- a/tensorflow/java/src/main/java/org/tensorflow/NativeLibrary.java +++ b/tensorflow/java/src/main/java/org/tensorflow/NativeLibrary.java @@ -88,7 +88,7 @@ final class NativeLibrary { // Deletions are in the reverse order of requests, so we need to request that the directory be // deleted first, so that it is empty when the request is fulfilled. tempPath.deleteOnExit(); - final String tempDirectory = tempPath.toString(); + final String tempDirectory = tempPath.getCanonicalPath(); if (frameworkResource != null) { extractResource(frameworkResource, frameworkLibName, tempDirectory); } else { diff --git a/tensorflow/java/src/main/java/org/tensorflow/package-info.java b/tensorflow/java/src/main/java/org/tensorflow/package-info.java index dd4859e1b14045e4123e7f15fbaff98e14d0b377..521c5c610c1f775cf9174664f5b786786ce1181d 100644 --- a/tensorflow/java/src/main/java/org/tensorflow/package-info.java +++ b/tensorflow/java/src/main/java/org/tensorflow/package-info.java @@ -35,5 +35,9 @@ limitations under the License. *
  • Graph execution: Using a Session to execute the graphs and find the best label for an * image. * + * + *

    Additional examples can be found in the tensorflow/models + * GitHub repository. */ package org.tensorflow; diff --git a/tensorflow/java/src/main/native/tensor_jni.cc b/tensorflow/java/src/main/native/tensor_jni.cc index 745abec244d1528e918464473e5d3fb19ad5082c..7e3cf4a88aac5acd4721a07c8316d8d124dce001 100644 --- a/tensorflow/java/src/main/native/tensor_jni.cc +++ b/tensorflow/java/src/main/native/tensor_jni.cc @@ -400,7 +400,13 @@ size_t nonScalarTF_STRINGTensorSize(JNIEnv* env, jarray value, int num_dims) { for (jsize i = 0; i < len; ++i) { jarray elem = static_cast( env->GetObjectArrayElement(static_cast(value), i)); + if (elem == nullptr) { + throwException(env, kNullPointerException, + "null entries in provided array"); + return ret; + } ret += nonScalarTF_STRINGTensorSize(env, elem, num_dims - 1); + if (env->ExceptionCheck()) return ret; } return ret; } @@ -421,8 +427,8 @@ void fillNonScalarTF_STRINGTensorData(JNIEnv* env, jarray value, int num_dims, for (jsize i = 0; i < len; ++i) { jarray elem = static_cast( env->GetObjectArrayElement(static_cast(value), i)); - if (TF_GetCode(status) != TF_OK) return; fillNonScalarTF_STRINGTensorData(env, elem, num_dims - 1, writer, status); + if (TF_GetCode(status) != TF_OK) return; } } } // namespace @@ -444,6 +450,7 @@ JNIEXPORT jlong JNICALL Java_org_tensorflow_Tensor_allocateNonScalarBytes( } const size_t encoded_size = nonScalarTF_STRINGTensorSize(env, value, num_dims); + if (env->ExceptionCheck()) return 0; TF_Tensor* t = TF_AllocateTensor(TF_STRING, dims, num_dims, 8 * num_elements + encoded_size); if (t == nullptr) { diff --git a/tensorflow/java/src/test/java/org/tensorflow/TensorTest.java b/tensorflow/java/src/test/java/org/tensorflow/TensorTest.java index 6538359d11a95eae698cc5aac8430e74ab1ed74c..1bd00a763ddff2f067183f57cfa80fdcbed84fd2 100644 --- a/tensorflow/java/src/test/java/org/tensorflow/TensorTest.java +++ b/tensorflow/java/src/test/java/org/tensorflow/TensorTest.java @@ -432,7 +432,7 @@ public class TensorTest { try (Tensor t = Tensor.create(vector, Integer.class)) { fail("Tensor.create() should fail because it was given an array of boxed values"); } catch (IllegalArgumentException e) { - // The expected exception + // The expected exception } } @@ -536,4 +536,15 @@ public class TensorTest { assertArrayEquals(matrix, cpy.copyTo(new float[2][3])); } } + + @Test + public void gracefullyFailCreationFromNullArrayForStringTensor() { + // Motivated by: https://github.com/tensorflow/tensorflow/issues/17130 + byte[][] array = new byte[1][]; + try { + Tensors.create(array); + } catch (NullPointerException e) { + // expected. + } + } } diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index ca79383d43931760414350634967b5ab00ffbf2f..3cbeb34c5471ccb6ff8f36f2de2de57d67f974f8 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -1,5 +1,8 @@ # Description: # Python support for TensorFlow. +# +# Public targets: +# ":platform" - Low-level and platform-specific Python code. package( default_visibility = [ @@ -25,6 +28,7 @@ load("//tensorflow:tensorflow.bzl", "py_tests") load("//tensorflow:tensorflow.bzl", "tf_py_build_info_genrule") load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") load("//tensorflow:tensorflow.bzl", "tf_cc_shared_object") +load("//tensorflow:tensorflow.bzl", "tf_custom_op_library_additional_deps_impl") load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "cuda_py_tests") load("//tensorflow/core:platform/default/build_config.bzl", "pyx_library") @@ -83,39 +87,39 @@ py_library( ":ops", ":platform", ":pywrap_tensorflow", + ":saver_test_utils", ":script_ops", ":session_ops", ":sets", ":sparse_ops", ":spectral_ops", + ":spectral_ops_test_util", ":standard_ops", ":state_ops", ":string_ops", + ":subscribe", ":summary", ":tensor_array_ops", - ":training", - ":saver_test_utils", - ":subscribe", ":test_ops", # TODO: Break testing code out into separate rule. - ":tf_item", ":tf_cluster", + ":tf_item", ":tf_optimizer", + ":training", ":util", ":weights_broadcast_ops", - "//third_party/py/numpy", + "//tensorflow/contrib:contrib_py", "//tensorflow/core:protos_all_py", "//tensorflow/python/data", "//tensorflow/python/estimator:estimator_py", "//tensorflow/python/feature_column:feature_column_py", "//tensorflow/python/keras", - "//tensorflow/python/ops/losses", "//tensorflow/python/ops/distributions", "//tensorflow/python/ops/linalg", + "//tensorflow/python/ops/losses", "//tensorflow/python/profiler", "//tensorflow/python/saved_model", - ] + if_not_windows([ - "//tensorflow/contrib:contrib_py", - ]), + "//third_party/py/numpy", + ], ) tf_py_build_info_genrule() @@ -132,6 +136,7 @@ py_library( ], ) + ["platform/build_info.py"], srcs_version = "PY2AND3", + visibility = ["//visibility:public"], deps = [ ":lib", ":pywrap_tensorflow", @@ -577,6 +582,7 @@ py_library( ":pywrap_tensorflow", ":random_seed", ":sparse_tensor", + ":tensor_spec", ":tensor_util", ":util", "//tensorflow/python/eager:context", @@ -760,6 +766,31 @@ py_library( ], ) +py_library( + name = "smart_cond", + srcs = ["framework/smart_cond.py"], + srcs_version = "PY2AND3", + deps = [ + ":control_flow_ops", + ":tensor_util", + ], +) + +py_test( + name = "smart_cond_test", + size = "small", + srcs = ["framework/smart_cond_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":client_testlib", + ":constant_op", + ":framework_ops", + ":math_ops", + ":session", + ":smart_cond", + ], +) + py_library( name = "sparse_tensor", srcs = ["framework/sparse_tensor.py"], @@ -781,6 +812,18 @@ py_library( ], ) +py_library( + name = "tensor_spec", + srcs = ["framework/tensor_spec.py"], + srcs_version = "PY2AND3", + deps = [ + ":common_shapes", + ":dtypes", + ":tensor_shape", + "//third_party/py/numpy", + ], +) + py_library( name = "tensor_util", srcs = ["framework/tensor_util.py"], @@ -904,7 +947,6 @@ py_test( srcs = ["framework/contrib_test.py"], main = "framework/contrib_test.py", srcs_version = "PY2AND3", - tags = ["no_windows"], deps = [ "//tensorflow:tensorflow_py", "//tensorflow/python:client_testlib", @@ -1149,6 +1191,21 @@ py_test( ], ) +py_test( + name = "framework_tensor_spec_test", + size = "small", + srcs = ["framework/tensor_spec_test.py"], + main = "framework/tensor_spec_test.py", + srcs_version = "PY2AND3", + deps = [ + ":framework_for_generated_wrappers", + ":framework_test_lib", + ":platform_test", + ":tensor_spec", + "//third_party/py/numpy", + ], +) + py_test( name = "framework_sparse_tensor_test", size = "small", @@ -1254,7 +1311,6 @@ py_test( srcs = ["framework/dtypes_test.py"], main = "framework/dtypes_test.py", srcs_version = "PY2AND3", - tags = ["no_windows"], deps = [ ":framework_for_generated_wrappers", ":framework_test_lib", @@ -1596,7 +1652,6 @@ py_test( size = "small", srcs = ["ops/clip_ops_test.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], deps = [ ":client_testlib", ":clip_ops", @@ -1718,6 +1773,7 @@ py_library( py_library( name = "gradients", srcs = [ + "ops/custom_gradient.py", "ops/gradients.py", "ops/gradients_impl.py", ], @@ -1731,6 +1787,7 @@ py_library( ":control_flow_util", ":framework", ":framework_for_generated_wrappers", + ":framework_ops", ":functional_ops", ":image_grad", ":linalg_grad", @@ -1743,6 +1800,8 @@ py_library( ":platform", ":spectral_grad", ":util", + "//tensorflow/python/eager:context", + "//tensorflow/python/eager:tape", "//third_party/py/numpy", "@six_archive//:six", ], @@ -1785,13 +1844,16 @@ py_library( ":control_flow_ops", ":framework", ":framework_for_generated_wrappers", + ":gradients", ":image_ops_gen", ":math_ops", + ":nn", ":nn_ops_gen", ":random_ops", ":string_ops", ":util", ":variables", + "//third_party/py/numpy", ], ) @@ -2486,6 +2548,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":array_ops", + ":checkpointable", ":control_flow_ops", ":dtypes", ":framework_ops", @@ -2655,7 +2718,6 @@ cuda_py_test( ], data = ["//tensorflow/core:image_testdata"], shard_count = 5, - tags = ["no_windows"], ) cuda_py_test( @@ -2819,6 +2881,30 @@ py_library( ], ) +py_library( + name = "checkpointable", + srcs = ["training/checkpointable.py"], + srcs_version = "PY2AND3", + deps = [ + ":array_ops", + ":dtypes", + ":io_ops_gen", + ":ops", + ":util", + "//tensorflow/python/eager:context", + ], +) + +py_test( + name = "checkpointable_test", + srcs = ["training/checkpointable_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":checkpointable", + ":client_testlib", + ], +) + py_test( name = "evaluation_test", size = "small", @@ -2994,7 +3080,6 @@ tf_proto_library( "framework/cpp_shape_inference.proto", ], ), - go_api_version = 2, ) tf_proto_library_py( @@ -3169,6 +3254,10 @@ tf_py_wrap_cc( "util/transform_graph.i", "util/util.i", ], + win_def_file = select({ + "//tensorflow:windows": ":pywrap_tensorflow_filtered_def_file", + "//conditions:default": None, + }), deps = [ ":bfloat16_lib", ":cost_analyzer_lib", @@ -3212,6 +3301,65 @@ tf_py_wrap_cc( tf_additional_gdr_deps()), ) +# ** Targets for Windows build (start) ** +# We need the following targets to expose symbols from _pywrap_tensorflow.dll + +# Build a cc_binary from tf_custom_op_library_additional_deps_impl, +# it contains all object code from its dependencies. +cc_binary( + name = "tf_custom_op_library_additional_deps.so", + linkshared = 1, + linkstatic = 1, + deps = tf_custom_op_library_additional_deps_impl(), +) + +# Get a DEF file generated by parsing all object files +# of tf_custom_op_library_additional_deps.so +filegroup( + name = "pywrap_tensorflow_def_file", + srcs = [":tf_custom_op_library_additional_deps.so"], + output_group = "def_file", +) + +# Filter the DEF file to reduce the number of symbols to 64K or less. +# Note that we also write the name of the pyd file into DEF file so that +# the dynamic libraries of custom ops can find it at runtime. +genrule( + name = "pywrap_tensorflow_filtered_def_file", + srcs = [":pywrap_tensorflow_def_file"], + outs = ["pywrap_tensorflow_filtered_def_file.def"], + cmd = select({ + "//tensorflow:windows": """ + $(location @local_config_def_file_filter//:def_file_filter) \\ + --input $(location :pywrap_tensorflow_def_file) \\ + --output $@ \\ + --target _pywrap_tensorflow_internal.pyd + """, + "//conditions:default": "touch $@", # Just a placeholder for Unix platforms + }), + tools = ["@local_config_def_file_filter//:def_file_filter"], +) + +# Get the import library of _pywrap_tensorflow_internal.dll +filegroup( + name = "pywrap_tensorflow_import_lib_file", + srcs = [":_pywrap_tensorflow_internal.so"], + output_group = "interface_library", +) + +# Create a cc_import rule for the import library of _pywrap_tensorflow_internal.dll +# so that custom ops' dynamic libraries can link against it. +cc_import( + name = "pywrap_tensorflow_import_lib", + interface_library = select({ + "//tensorflow:windows": ":pywrap_tensorflow_import_lib_file", + "//conditions:default": "not_exsiting_on_unix.lib", # Just a placeholder for Unix platforms + }), + system_provided = 1, +) + +# ** Targets for Windows build (end) ** + py_library( name = "lib", srcs = [ @@ -3573,6 +3721,7 @@ py_test( ":framework_for_generated_wrappers", ":math_ops", ":state_ops_gen", + ":variable_scope", ":variables", "//tensorflow/core:protos_all_py", ], @@ -3583,7 +3732,6 @@ py_test( size = "small", srcs = ["lib/core/bfloat16_test.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], deps = [ ":client_testlib", ":lib", @@ -3863,7 +4011,13 @@ py_test( size = "small", srcs = ["training/checkpoint_utils_test.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], + tags = [ + "manual", + "no_cuda_on_cpu_tap", + "no_oss", + "no_windows", + "notap", + ], deps = [ ":client", ":client_testlib", @@ -3872,6 +4026,7 @@ py_test( ":partitioned_variables", ":platform", ":pywrap_tensorflow", + ":resource_variable_ops", ":state_ops", ":training", ":variable_scope", @@ -3884,7 +4039,6 @@ py_test( size = "small", srcs = ["training/checkpoint_ops_test.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], deps = [ ":checkpoint_ops_gen", ":client", @@ -3901,15 +4055,31 @@ py_test( ], ) +py_test( + name = "warm_starting_util_test", + size = "small", + srcs = ["training/warm_starting_util_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":array_ops", + ":client_testlib", + ":dtypes", + ":framework_ops", + ":init_ops", + ":training", + ":variable_scope", + ":variables", + "//tensorflow/python/feature_column", + "//third_party/py/numpy", + ], +) + py_test( name = "monitored_session_test", size = "medium", srcs = ["training/monitored_session_test.py"], srcs_version = "PY2AND3", - tags = [ - "no_windows", - "notsan", # b/67945581 - ], + tags = ["notsan"], # b/67945581 deps = [ ":array_ops", ":client_testlib", @@ -4035,6 +4205,7 @@ py_library( ":control_flow_ops", ":framework_for_generated_wrappers", ":platform", + ":smart_cond", ":tensor_util", ":util", ":variable_scope", @@ -4051,8 +4222,6 @@ py_library( "layers/convolutional.py", "layers/core.py", "layers/layers.py", - "layers/maxout.py", - "layers/network.py", "layers/normalization.py", "layers/pooling.py", ], @@ -4105,25 +4274,6 @@ py_test( ], ) -py_test( - name = "layers_network_test", - size = "small", - srcs = ["layers/network_test.py"], - main = "layers/network_test.py", - srcs_version = "PY2AND3", - deps = [ - ":array_ops", - ":client_testlib", - ":framework_for_generated_wrappers", - ":framework_test_lib", - ":layers", - ":layers_base", - ":sparse_ops", - "//tensorflow/python/eager:context", - "//third_party/py/numpy", - ], -) - py_test( name = "layers_core_test", size = "small", @@ -4162,22 +4312,6 @@ py_test( ], ) -py_test( - name = "layers_maxout_test", - size = "small", - srcs = ["layers/maxout_test.py"], - main = "layers/maxout_test.py", - srcs_version = "PY2AND3", - deps = [ - ":client_testlib", - ":framework_for_generated_wrappers", - ":layers", - ":math_ops", - ":nn_ops", - ":random_ops", - ], -) - py_test( name = "layers_utils_test", size = "small", @@ -4266,12 +4400,6 @@ filegroup( visibility = ["//tensorflow:__subpackages__"], ) -filegroup( - name = "hidden_ops", - srcs = ["ops/hidden_ops.txt"], - visibility = ["//tensorflow:__subpackages__"], -) - cuda_py_test( name = "accumulate_n_benchmark", size = "large", @@ -4579,6 +4707,34 @@ py_test( ], ) +py_library( + name = "graph_placer", + srcs = [ + "grappler/controller.py", + "grappler/graph_placer.py", + "grappler/hierarchical_controller.py", + ], + deps = [ + ":python", + "//third_party/py/numpy", + ], +) + +py_test( + name = "graph_placer_test", + size = "large", + srcs = ["grappler/graph_placer_test.py"], + tags = [ + "grappler", + "no_pip", # graph_placer is not available in pip. + ], + deps = [ + ":client_testlib", + ":graph_placer", + "//tensorflow/python:math_ops", + ], +) + py_test( name = "memory_optimizer_test", size = "medium", @@ -4670,6 +4826,7 @@ py_test( srcs_version = "PY2AND3", tags = [ "grappler", + "no_cuda_on_cpu_tap", "no_pip", ], deps = [ diff --git a/tensorflow/python/__init__.py b/tensorflow/python/__init__.py index ea7604d30f127e51d573a5251420faa22cbb3468..5a9cd7531db6482313005feebaadf953c6da3bb4 100644 --- a/tensorflow/python/__init__.py +++ b/tensorflow/python/__init__.py @@ -60,7 +60,7 @@ from tensorflow.core.protobuf.tensorflow_server_pb2 import * from tensorflow.core.util.event_pb2 import * # Framework -from tensorflow.python.framework.framework_lib import * +from tensorflow.python.framework.framework_lib import * # pylint: disable=redefined-builtin from tensorflow.python.framework.versions import * from tensorflow.python.framework import errors from tensorflow.python.framework import graph_util @@ -116,6 +116,7 @@ from tensorflow.python.platform import test from tensorflow.python.util.all_util import remove_undocumented from tensorflow.python.util.all_util import make_all +from tensorflow.python.util.tf_export import tf_export # Import modules whose docstrings contribute, for use by remove_undocumented # below. @@ -138,6 +139,10 @@ from tensorflow.python.ops import state_ops from tensorflow.python.ops import string_ops from tensorflow.python.ops import tensor_array_ops +# Eager execution +from tensorflow.python.eager.context import executing_eagerly +from tensorflow.python.framework.ops import enable_eager_execution + # Symbols whitelisted for export without documentation. # TODO(cwhipkey): review these and move to contrib, expose through # documentation, or remove. @@ -167,18 +172,39 @@ _allowed_symbols = [ 'TensorInfo', # Used for tf.saved_model functionality. ] +# Export protos +# pylint: disable=undefined-variable +tf_export('AttrValue')(AttrValue) +tf_export('ConfigProto')(ConfigProto) +tf_export('Event', 'summary.Event')(Event) +tf_export('GPUOptions')(GPUOptions) +tf_export('GraphDef')(GraphDef) +tf_export('GraphOptions')(GraphOptions) +tf_export('HistogramProto')(HistogramProto) +tf_export('LogMessage')(LogMessage) +tf_export('MetaGraphDef')(MetaGraphDef) +tf_export('NameAttrList')(NameAttrList) +tf_export('NodeDef')(NodeDef) +tf_export('OptimizerOptions')(OptimizerOptions) +tf_export('RunMetadata')(RunMetadata) +tf_export('RunOptions')(RunOptions) +tf_export('SessionLog', 'summary.SessionLog')(SessionLog) +tf_export('Summary', 'summary.Summary')(Summary) +tf_export('summary.SummaryDescription')(SummaryDescription) +tf_export('SummaryMetadata')(SummaryMetadata) +tf_export('summary.TaggedRunMetadata')(TaggedRunMetadata) +tf_export('TensorInfo')(TensorInfo) +# pylint: enable=undefined-variable + + # The following symbols are kept for compatibility. It is our plan # to remove them in the future. _allowed_symbols.extend([ 'arg_max', 'arg_min', - 'mul', # use tf.multiply instead. - 'neg', # use tf.negative instead. - 'sub', # use tf.subtract instead. 'create_partitioned_variables', 'deserialize_many_sparse', 'lin_space', - 'list_diff', # Use tf.listdiff instead. 'listdiff', # Use tf.listdiff instead. 'parse_single_sequence_example', 'serialize_many_sparse', @@ -268,6 +294,12 @@ _allowed_symbols.extend([ 'MONOLITHIC_BUILD', ]) +# Eager execution +_allowed_symbols.extend([ + 'enable_eager_execution', + 'executing_eagerly', +]) + # Remove all extra symbols that don't have a docstring or are not explicitly # referenced in the whitelist. remove_undocumented(__name__, _allowed_symbols, [ diff --git a/tensorflow/python/build_defs.bzl b/tensorflow/python/build_defs.bzl index 7f29adc06fcc5922114b7cd2bde8a8df5b1e0665..b9056f86e6d0465a8521f054a459c06eb5aeb37c 100644 --- a/tensorflow/python/build_defs.bzl +++ b/tensorflow/python/build_defs.bzl @@ -22,7 +22,6 @@ def tf_gen_op_wrapper_private_py(name, out=None, deps=[], bare_op_name = name[:-4] # Strip off the _gen tf_gen_op_wrapper_py(name=bare_op_name, out=out, - hidden_file="ops/hidden_ops.txt", visibility=visibility, deps=deps, require_shape_functions=require_shape_functions, diff --git a/tensorflow/python/client/events_writer.i b/tensorflow/python/client/events_writer.i index de030fcb4282912475ed8853bae9d41cde2c085d..c72b76b8fa4a05588841466a836bc189bb64d154 100644 --- a/tensorflow/python/client/events_writer.i +++ b/tensorflow/python/client/events_writer.i @@ -23,6 +23,9 @@ limitations under the License. %nodefaultctor EventsWriter; +%ignore tensorflow::Status::operator=; +%include "tensorflow/core/lib/core/status.h" + %ignoreall %unignore tensorflow; %unignore tensorflow::EventsWriter; diff --git a/tensorflow/python/client/session.py b/tensorflow/python/client/session.py index e6f94396b85eb4d0ab0774a53484089f735be940..29f06c8f227a75b72d08030ec1afd08b4483b42d 100644 --- a/tensorflow/python/client/session.py +++ b/tensorflow/python/client/session.py @@ -21,6 +21,7 @@ from __future__ import print_function import functools import re import threading +import warnings import numpy as np @@ -35,6 +36,7 @@ from tensorflow.python.ops import session_ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import compat from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export class SessionInterface(object): @@ -1084,7 +1086,10 @@ class BaseSession(SessionInterface): if isinstance(subfeed_val, ops.Tensor): raise TypeError('The value of a feed cannot be a tf.Tensor object. ' 'Acceptable feed values include Python scalars, ' - 'strings, lists, numpy ndarrays, or TensorHandles.') + 'strings, lists, numpy ndarrays, or TensorHandles.' + 'For reference, the tensor object was ' + + str(feed_val) + ' which was passed to the ' + 'feed with key ' + str(feed) + '.') subfeed_dtype = subfeed_t.dtype.as_numpy_dtype if isinstance(subfeed_val, int) and _convert_to_numpy_obj( @@ -1216,19 +1221,12 @@ class BaseSession(SessionInterface): compat.as_bytes(options.SerializeToString())) if options else None run_metadata_ptr = tf_session.TF_NewBuffer() if run_metadata else None try: - with errors.raise_exception_on_not_ok_status() as status: - if self._created_with_new_api: - results = tf_session.TF_SessionRun_wrapper( - self._session, options_ptr, {}, fetch_list, target_list, - run_metadata_ptr, status) - else: - results = tf_session.TF_Run(self._session, options_ptr, {}, - fetch_list, target_list, status, - run_metadata_ptr) - if fetch_handler: - results = fetch_handler.build_results(self, results) - else: - results = results[0] if results else None + results = self._call_tf_sessionrun( + options_ptr, {}, fetch_list, target_list, run_metadata_ptr) + if fetch_handler: + results = fetch_handler.build_results(self, results) + else: + results = results[0] if results else None if run_metadata: proto_data = tf_session.TF_GetBuffer(run_metadata_ptr) run_metadata.ParseFromString(compat.as_bytes(proto_data)) @@ -1249,13 +1247,7 @@ class BaseSession(SessionInterface): assert len(target_list) == 1 def _single_operation_run(): - with errors.raise_exception_on_not_ok_status() as status: - if self._created_with_new_api: - tf_session.TF_SessionRun_wrapper(self._session, None, {}, [], - target_list, None, status) - else: - tf_session.TF_Run(self._session, None, {}, [], target_list, status, - None) + self._call_tf_sessionrun(None, {}, [], target_list, None) return _single_operation_run elif isinstance(fetches, ops.Tensor): @@ -1265,13 +1257,7 @@ class BaseSession(SessionInterface): assert not target_list def _single_tensor_run(): - with errors.raise_exception_on_not_ok_status() as status: - if self._created_with_new_api: - results = tf_session.TF_SessionRun_wrapper( - self._session, None, {}, fetch_list, [], None, status) - else: - results = tf_session.TF_Run(self._session, None, {}, fetch_list, [], - status, None) + results = self._call_tf_sessionrun(None, {}, fetch_list, [], None) return results[0] return _single_tensor_run @@ -1279,13 +1265,8 @@ class BaseSession(SessionInterface): # In all other cases, we must use `fetch_handler` to build the # results for us. def _fetch_handler_run(): - with errors.raise_exception_on_not_ok_status() as status: - if self._created_with_new_api: - results = tf_session.TF_SessionRun_wrapper( - self._session, None, {}, fetch_list, target_list, None, status) - else: - results = tf_session.TF_Run(self._session, None, {}, fetch_list, - target_list, status, None) + results = self._call_tf_sessionrun( + None, {}, fetch_list, target_list, None) return fetch_handler.build_results(self, results) return _fetch_handler_run @@ -1325,35 +1306,22 @@ class BaseSession(SessionInterface): fetches = _name_list(fetch_list) targets = _name_list(target_list) - def _run_fn(session, feed_dict, fetch_list, target_list, options, - run_metadata): + def _run_fn(feed_dict, fetch_list, target_list, options, run_metadata): # Ensure any changes to the graph are reflected in the runtime. self._extend_graph() - with errors.raise_exception_on_not_ok_status() as status: - if self._created_with_new_api: - return tf_session.TF_SessionRun_wrapper(session, options, feed_dict, - fetch_list, target_list, - run_metadata, status) - else: - return tf_session.TF_Run(session, options, feed_dict, fetch_list, - target_list, status, run_metadata) + return self._call_tf_sessionrun( + options, feed_dict, fetch_list, target_list, run_metadata) - def _prun_fn(session, handle, feed_dict, fetch_list): + def _prun_fn(handle, feed_dict, fetch_list): if target_list: raise RuntimeError('partial_run() requires empty target_list.') - with errors.raise_exception_on_not_ok_status() as status: - if self._created_with_new_api: - return tf_session.TF_SessionPRun_wrapper(session, handle, feed_dict, - fetch_list, status) - else: - return tf_session.TF_PRun(session, handle, feed_dict, fetch_list, - status) + return self._call_tf_sessionprun(handle, feed_dict, fetch_list) if handle is None: - return self._do_call(_run_fn, self._session, feeds, fetches, targets, - options, run_metadata) + return self._do_call(_run_fn, feeds, fetches, targets, options, + run_metadata) else: - return self._do_call(_prun_fn, self._session, handle, feeds, fetches) + return self._do_call(_prun_fn, handle, feeds, fetches) def _do_call(self, fn, *args): try: @@ -1373,23 +1341,23 @@ class BaseSession(SessionInterface): raise type(e)(node_def, op, message) def _extend_graph(self): - # Nothing to do if we're using the new session interface - # TODO(skyewm): remove this function altogether eventually if self._created_with_new_api: - return - - # Ensure any changes to the graph are reflected in the runtime. - with self._extend_lock: - if self._graph.version > self._current_version: - # pylint: disable=protected-access - graph_def, self._current_version = self._graph._as_graph_def( - from_version=self._current_version, add_shapes=self._add_shapes) - # pylint: enable=protected-access - + with self._graph._lock: # pylint: disable=protected-access with errors.raise_exception_on_not_ok_status() as status: - tf_session.TF_ExtendGraph(self._session, - graph_def.SerializeToString(), status) - self._opened = True + tf_session.ExtendSession(self._session, status) + else: + # Ensure any changes to the graph are reflected in the runtime. + with self._extend_lock: + if self._graph.version > self._current_version: + # pylint: disable=protected-access + graph_def, self._current_version = self._graph._as_graph_def( + from_version=self._current_version, add_shapes=self._add_shapes) + # pylint: enable=protected-access + + with errors.raise_exception_on_not_ok_status() as status: + tf_session.TF_ExtendGraph(self._session, + graph_def.SerializeToString(), status) + self._opened = True # The threshold to run garbage collection to delete dead tensors. _DEAD_HANDLES_THRESHOLD = 10 @@ -1440,7 +1408,29 @@ class BaseSession(SessionInterface): feed_dict[feed_tensor] = np_val return handles + def _call_tf_sessionrun(self, options, feed_dict, fetch_list, target_list, + run_metadata): + with errors.raise_exception_on_not_ok_status() as status: + if self._created_with_new_api: + return tf_session.TF_SessionRun_wrapper( + self._session, options, feed_dict, fetch_list, target_list, + run_metadata, status) + else: + return tf_session.TF_Run( + self._session, options, feed_dict, fetch_list, target_list, + status, run_metadata) + def _call_tf_sessionprun(self, handle, feed_dict, fetch_list): + with errors.raise_exception_on_not_ok_status() as status: + if self._created_with_new_api: + return tf_session.TF_SessionPRun_wrapper( + self._session, handle, feed_dict, fetch_list, status) + else: + return tf_session.TF_PRun( + self._session, handle, feed_dict, fetch_list, status) + + +@tf_export('Session') class Session(BaseSession): """A class for running TensorFlow operations. @@ -1537,8 +1527,22 @@ class Session(BaseSession): def __exit__(self, exec_type, exec_value, exec_tb): if exec_type is errors.OpError: logging.error('Session closing due to OpError: %s', (exec_value,)) - self._default_session_context_manager.__exit__(exec_type, exec_value, - exec_tb) + try: + self._default_session_context_manager.__exit__(exec_type, exec_value, + exec_tb) + except RuntimeError as error: + if error == exec_value: + # NOTE(skyewm): for some reason, in Python3, + # _default_session_context_manager.__exit__ will re-raise the "not + # re-entrant" exception raised in __enter__ above (note that if we're + # here, we're in the outer session context manager, since __exit__ is + # not called when __enter__ raises an exception). We still want to + # continue cleaning up this context manager before the exception is + # further propagated, so we ignore it here (note that it'll continue + # being propagated after this method completes). + pass + else: + raise self._default_graph_context_manager.__exit__(exec_type, exec_value, exec_tb) self._default_session_context_manager = None @@ -1581,6 +1585,7 @@ class Session(BaseSession): tf_session.TF_Reset(target, containers, config) +@tf_export('InteractiveSession') class InteractiveSession(BaseSession): """A TensorFlow `Session` for use in interactive contexts, such as a shell. @@ -1620,6 +1625,9 @@ class InteractiveSession(BaseSession): ``` """ + _count_lock = threading.Lock() + _active_session_count = 0 # GUARDED_BY(_count_lock) + def __init__(self, target='', graph=None, config=None): """Creates a new interactive TensorFlow session. @@ -1648,6 +1656,15 @@ class InteractiveSession(BaseSession): config.graph_options.place_pruned_graph = True super(InteractiveSession, self).__init__(target, graph, config) + with InteractiveSession._count_lock: + if InteractiveSession._active_session_count > 0: + warnings.warn('An interactive session is already active. This can ' + 'cause out-of-memory errors in some cases. You must ' + 'explicitly call `InteractiveSession.close()` to release ' + 'resources held by the other session(s).') + InteractiveSession._active_session_count += 1 + self._closed = False + self._default_session = self.as_default() self._default_session.enforce_nesting = False self._default_session.__enter__() @@ -1660,6 +1677,12 @@ class InteractiveSession(BaseSession): def close(self): """Closes an `InteractiveSession`.""" super(InteractiveSession, self).close() + with InteractiveSession._count_lock: + if not self._closed: + InteractiveSession._active_session_count -= 1 + self._closed = True + else: + return if self._explicit_graph is not None: self._default_graph.__exit__(None, None, None) self._default_session.__exit__(None, None, None) diff --git a/tensorflow/python/client/session_benchmark.py b/tensorflow/python/client/session_benchmark.py index 06e9a099267938e53e377e65ed975f7f4b8b966b..da74855193dbfe3019f23c542d86c5e493e9ac7a 100644 --- a/tensorflow/python/client/session_benchmark.py +++ b/tensorflow/python/client/session_benchmark.py @@ -22,7 +22,7 @@ import time import numpy as np -from six.moves import xrange +from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.python.client import session from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops diff --git a/tensorflow/python/client/session_test.py b/tensorflow/python/client/session_test.py index 768a5db88aa647609dba1c479a5aca68cd26652a..6c7339f3d83328903dee4667e64b4db699c7e35a 100644 --- a/tensorflow/python/client/session_test.py +++ b/tensorflow/python/client/session_test.py @@ -22,6 +22,7 @@ import os import sys import threading import time +import warnings import numpy as np import six @@ -31,13 +32,13 @@ from tensorflow.core.framework import attr_value_pb2 from tensorflow.core.framework import types_pb2 from tensorflow.core.lib.core import error_codes_pb2 from tensorflow.core.protobuf import config_pb2 -from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.client import session from tensorflow.python.framework import common_shapes 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 importer from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_util @@ -46,8 +47,9 @@ from tensorflow.python.framework import versions 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 gen_control_flow_ops +from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import math_ops -from tensorflow.python.ops import random_ops # Import resource_variable_ops for the variables-to-tensor implicit conversion. from tensorflow.python.ops import resource_variable_ops # pylint: disable=unused-import from tensorflow.python.ops import state_ops @@ -61,9 +63,14 @@ from tensorflow.python.util import compat ops.RegisterShape('ConstructionFails')(common_shapes.unknown_shape) -@test_util.with_c_api +# TODO(skyewm): reenable when this works with _USE_C_SHAPES=False +# @test_util.with_c_api class SessionTest(test_util.TensorFlowTestCase): + def setUp(self): + super(SessionTest, self).setUp() + warnings.simplefilter('always') + def testUseExistingGraph(self): with ops.Graph().as_default() as g, ops.device('/cpu:0'): a = constant_op.constant(6.0, shape=[1, 1]) @@ -188,12 +195,10 @@ class SessionTest(test_util.TensorFlowTestCase): a = constant_op.constant(0.0, shape=[2, 3]) # NOTE(mrry): The original_op is nonsense, but used here to test that the # errors are reported correctly. - # pylint: disable=protected-access with sess.graph._original_op(a.op): b = array_ops.identity(a, name='id') with sess.graph._original_op(b.op): c = array_ops.placeholder(dtypes.float32) - # pylint: enable=protected-access def exc_predicate(e): return (e.op == c.op and e.op._original_op == b.op and @@ -1053,6 +1058,43 @@ class SessionTest(test_util.TensorFlowTestCase): for t in threads: t.join() + def testParallelRunAndBuild(self): + with session.Session() as sess: + c = constant_op.constant(5.0) + stop = threading.Event() + + def run_loop(): + while not stop.is_set(): + self.assertEqual(sess.run(c), 5.0) + + threads = [self.checkedThread(target=run_loop) for _ in range(100)] + for t in threads: + t.start() + + # Do some graph construction. Try to exercise non-trivial paths. + graph = ops.get_default_graph() + gdef = None + for _ in range(10): + x = array_ops.placeholder(dtype=dtypes.float32) + with ops.colocate_with(x): + y = array_ops.placeholder(dtype=dtypes.float32) + with ops.device('/cpu:0'): + z = control_flow_ops.while_loop( + lambda x, y: x < 10, lambda x, y: (x + 1, x * y), [x, y]) + with graph._attr_scope({'_a': attr_value_pb2.AttrValue(b=False)}): + gradients_impl.gradients(z, [x, y]) + if gdef is None: + gdef = graph.as_graph_def() + else: + # NOTE(skyewm): import_graph_def breaks the running threads without + # the C API enabled. This is not a regression so I didn't fix it. + if ops._USE_C_API: + importer.import_graph_def(gdef, name='import') + + stop.set() + for t in threads: + t.join() + def testRunFeedDict(self): with session.Session() as s: x = array_ops.zeros([2]) @@ -1154,6 +1196,32 @@ class SessionTest(test_util.TensorFlowTestCase): self.assertAllEqual([[24.0]], e.eval()) sess.close() + def testMultipleInteractiveSessionsWarning(self): + # Reinitialize the global state to ensure that the expected warnings will + # be emitted. + session.InteractiveSession._active_session_count = 0 # pylint: disable=protected-access + + sess = session.InteractiveSession() + sess.close() + # Opening and closing interactive sessions serially should not warn. + with warnings.catch_warnings(record=True) as w: + sess = session.InteractiveSession() + sess.close() + self.assertEqual(0, len(w)) + + with warnings.catch_warnings(record=True) as w: + sess = session.InteractiveSession() + self.assertEqual(0, len(w)) + with warnings.catch_warnings(record=True) as w: + sess2 = session.InteractiveSession() + self.assertEqual(1, len(w)) + self.assertTrue('An interactive session is already active. This can cause ' + 'out-of-memory errors in some cases. You must explicitly ' + 'call `InteractiveSession.close()` to release resources ' + 'held by the other session(s).' in str(w[0].message)) + sess2.close() + sess.close() + def testInteractivePlacePrunedGraph(self): sess = session.InteractiveSession() @@ -1745,8 +1813,10 @@ class SessionTest(test_util.TensorFlowTestCase): def runTestBuildGraphError(self, sess): # Ensure that errors from building the graph get propagated. data = array_ops.placeholder(dtypes.float32, shape=[]) - enter_1 = control_flow_ops.enter(data, 'foo_1', False) - enter_2 = control_flow_ops.enter(data, 'foo_2', False) + # pylint: disable=protected-access + enter_1 = gen_control_flow_ops.enter(data, 'foo_1', False) + enter_2 = gen_control_flow_ops.enter(data, 'foo_2', False) + # pylint: enable=protected-access res = math_ops.add(enter_1, enter_2) with self.assertRaisesOpError('has inputs from different frames'): sess.run(res, feed_dict={data: 1.0}) diff --git a/tensorflow/python/client/tf_session.i b/tensorflow/python/client/tf_session.i index 1fd488e7b6388f7953a279dca8f93ab57a85f63d..e88fc0c01a8bb7534f47e2a0389965c102bbad7b 100644 --- a/tensorflow/python/client/tf_session.i +++ b/tensorflow/python/client/tf_session.i @@ -719,6 +719,11 @@ def TF_Reset(target, containers=None, config=None): $1 = &types_local; } +%unignore SetRequireShapeInferenceFns; +%unignore TF_TryEvaluateConstant_wrapper; +%noexception TF_TryEvaluateConstant_wrapper; +%unignore ExtendSession; + %include "tensorflow/python/client/tf_session_helper.h" %unignoreall diff --git a/tensorflow/python/client/tf_session_helper.cc b/tensorflow/python/client/tf_session_helper.cc index 361dbc22b097a9bc82f656d7416b88c4a3a1ec2d..a8ab91749a86749a1eef25e2674634334682d0f3 100644 --- a/tensorflow/python/client/tf_session_helper.cc +++ b/tensorflow/python/client/tf_session_helper.cc @@ -493,4 +493,19 @@ std::vector TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper( return input_strs; } +PyObject* TF_TryEvaluateConstant_wrapper(TF_Graph* graph, TF_Output output, + TF_Status* status) { + TF_Tensor* result_tensor; + bool evaluated = + TF_TryEvaluateConstant(graph, output, &result_tensor, status); + if (!evaluated || TF_GetCode(status) != TF_OK) Py_RETURN_NONE; + + Safe_TF_TensorPtr safe_result_tensor(result_tensor); + PyObject* out; + Status s = TF_TensorToPyArray(std::move(safe_result_tensor), &out); + Set_TF_Status_from_Status(status, s); + if (!s.ok()) Py_RETURN_NONE; + return out; +} + } // namespace tensorflow diff --git a/tensorflow/python/client/tf_session_helper.h b/tensorflow/python/client/tf_session_helper.h index 29d5b28f40a7c07c199eec8c8cd85de626f6b068..83318dc178f6da3828a8dc41e81b7fc3e2e19e22 100644 --- a/tensorflow/python/client/tf_session_helper.h +++ b/tensorflow/python/client/tf_session_helper.h @@ -213,6 +213,11 @@ std::vector TF_GraphGetTensorShape_wrapper(TF_Graph* graph, std::vector TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper( TF_ImportGraphDefResults* results); +// If evaluation was possible, returns the numpy ndarray of the evaluated +// result. Otherwise returns None. +PyObject* TF_TryEvaluateConstant_wrapper(TF_Graph* graph, TF_Output output, + TF_Status* status); + } // namespace tensorflow #endif // TENSORFLOW_PYTHON_CLIENT_TF_SESSION_HELPER_H_ diff --git a/tensorflow/python/client/timeline_test.py b/tensorflow/python/client/timeline_test.py index 9641b8b7f2735e2e0477aec59edd539e999fa969..5e6b5acdb02e4c8c167485520a8d84ac43db7511 100644 --- a/tensorflow/python/client/timeline_test.py +++ b/tensorflow/python/client/timeline_test.py @@ -155,9 +155,12 @@ class TimelineTest(test.TestCase): ctf = step_analysis.chrome_trace.format_to_string() self._validateTrace(ctf) maximums = step_analysis.allocator_maximums - self.assertTrue('cpu' in maximums) + cpuname = 'cpu' + if 'mklcpu' in maximums: + cpuname = 'mkl' + cpuname + self.assertTrue(cpuname in maximums) cpu_max = maximums[ - 'cuda_host_bfc'] if 'cuda_host_bfc' in maximums else maximums['cpu'] + 'cuda_host_bfc'] if 'cuda_host_bfc' in maximums else maximums[cpuname] # At least num1 + num2, both float32s (4 bytes each) self.assertGreater(cpu_max.num_bytes, 8) self.assertGreater(cpu_max.timestamp, 0) diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD index 43cbde69d9db20d85c55e071d8393074a78a4a1b..8b8adefa65a5c54d40bc28d8f50953513cfd3605 100644 --- a/tensorflow/python/data/kernel_tests/BUILD +++ b/tensorflow/python/data/kernel_tests/BUILD @@ -357,6 +357,9 @@ tf_py_test( "//tensorflow/python:session", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/data/ops:iterator_ops", + "//tensorflow/python:constant_op", + "//tensorflow/python:string_ops", + "//tensorflow/python:lookup_ops", ], grpc_enabled = True, tags = [ diff --git a/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py b/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py index b71652c980f233ce116ea89544fcb38ad1d816d1..02720a2e985914d3a6774dc6f64d1316890c46bf 100644 --- a/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py @@ -28,6 +28,7 @@ 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.ops import variables from tensorflow.python.platform import test @@ -202,44 +203,45 @@ class FilesystemCacheDatasetTest(test.TestCase): 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)) + with ops.device("cpu:0"): + 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) + 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() + # 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: + with self.test_session() as sess: - sess.run(repeat_count.initializer) - sess.run(cached_iterator.initializer) - sess.run(uncached_iterator.initializer) + 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) + 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)) + sess.run(repeat_count.assign(0)) - # The uncached iterator should now be empty. - with self.assertRaises(errors.OutOfRangeError): - sess.run(uncached_next) + # 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 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) + # 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]), diff --git a/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py b/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py index 14627810b57f68fd96e3e3cc7b51b4fbf7365299..ea5b41e5d819743ad03f3148d654329aea51dab7 100644 --- a/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py +++ b/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py @@ -263,7 +263,7 @@ class DatasetConstructorTest(test.TestCase): for i in range(3): results = sess.run(get_next) for component, result_component in zip( - (zip(*components[:3])[i] + expected[i]), results): + (list(zip(*components[:3]))[i] + expected[i]), results): if sparse_tensor.is_sparse(component): self.assertSparseValuesEqual(component, result_component) else: diff --git a/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py b/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py index f129d07b57b96b7869c84467aeb2276c93531ef8..6aabad2f574551cbdc152fe378eb9dc0f5f71995 100644 --- a/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py +++ b/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py @@ -21,9 +21,12 @@ import threading import numpy as np +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 +from tensorflow.python.ops import script_ops from tensorflow.python.platform import test @@ -302,6 +305,89 @@ class DatasetConstructorTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testFromGeneratorStopShort(self): + + def generator(): + yield 0 + yield 1 + yield 2 + + iterator = ( + dataset_ops.Dataset.from_generator( + generator, output_types=dtypes.int64).make_initializable_iterator()) + init_op = iterator.initializer + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(init_op) + self.assertAllEqual(0, sess.run(get_next)) + self.assertAllEqual(1, sess.run(get_next)) + + def testFromGeneratorDestructorCalled(self): + # Use an `Event` to signal that the generator has been deleted. + event = threading.Event() + + class GeneratorWrapper(object): + + def __iter__(self): + return self + + def next(self): + return self.__next__() + + def __next__(self): + return 42 + + def __del__(self): + event.set() + + iterator = dataset_ops.Dataset.from_generator( + GeneratorWrapper, + output_types=dtypes.int64).take(2).make_initializable_iterator() + init_op = iterator.initializer + get_next = iterator.get_next() + + with session.Session() as sess: + sess.run(init_op) + self.assertAllEqual(42, sess.run(get_next)) + self.assertAllEqual(42, sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + # Test that `GeneratorWrapper` object is destroyed when the + # iterator terminates (and the generator iterator is deleted). + self.assertTrue(event.is_set()) + + def testGeneratorDatasetFinalizeFunctionCalled(self): + # NOTE(mrry): This test tests the internal `_GeneratorDataset`, + # which affords more control over what the finalize function can do than + # the `Dataset.from_generator()` wrapper. + + # Use an `Event` to signal that the generator has been deleted. + event = threading.Event() + + def finalize_fn(_): + def finalize_py_func(): + event.set() + return 0 + return script_ops.py_func(finalize_py_func, [], [dtypes.int64], + stateful=True) + + dummy = constant_op.constant(37) + iterator = (dataset_ops._GeneratorDataset(dummy, lambda x: x, + lambda x: x, finalize_fn) + .take(2) + .make_initializable_iterator()) + init_op = iterator.initializer + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(init_op) + self.assertAllEqual(37, sess.run(get_next)) + self.assertAllEqual(37, sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + self.assertTrue(event.is_set()) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py b/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py index b9258b720edd4ecd620c61eed18f6f975cb7f439..4f2216f0a340acb582c2d09523b0c78af99bdd90 100644 --- a/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py @@ -17,11 +17,15 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import time + import numpy as np +from tensorflow.python.client import session 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.framework import sparse_tensor from tensorflow.python.ops import array_ops from tensorflow.python.ops import functional_ops @@ -156,6 +160,65 @@ class FilterDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testReturnComponent(self): + iterator = ( + dataset_ops.Dataset.zip( + (dataset_ops.Dataset.range(10), + dataset_ops.Dataset.from_tensors(True).repeat(None))) + .filter(lambda x, y: y).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): + self.assertEqual((i, True), sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + def testParallelFilters(self): + dataset = dataset_ops.Dataset.range(10).filter( + lambda x: math_ops.equal(x % 2, 0)) + iterators = [dataset.make_one_shot_iterator() for _ in range(10)] + next_elements = [iterator.get_next() for iterator in iterators] + with self.test_session() as sess: + self.assertEqual([0 for _ in range(10)], sess.run(next_elements)) + + +class FilterDatasetBenchmark(test.Benchmark): + + def _benchmark(self, predicate, name): + with ops.Graph().as_default(): + dataset = ( + dataset_ops.Dataset.from_tensors(True).repeat(None).filter(predicate)) + iterator = dataset.make_one_shot_iterator() + next_element = iterator.get_next() + + with session.Session() as sess: + for _ in range(5): + sess.run(next_element.op) + deltas = [] + for _ in range(100): + start = time.time() + for _ in range(100): + sess.run(next_element.op) + end = time.time() + deltas.append(end - start) + + median_wall_time = np.median(deltas) / 100 + print("Filter dataset using %s. Median wall time: %f" % + (name, median_wall_time)) + self.report_benchmark( + iters=100, + wall_time=median_wall_time, + name="benchmark_filter_dataset_%s" % name) + + def benchmarkSimpleFunction(self): + self._benchmark(array_ops.identity, "simple_function") + + def benchmarkReturnComponentOptimization(self): + self._benchmark(lambda x: x, "return_component") + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py b/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py index 28cb50c00208f95e64bb11ae80656383b1f41e1e..7dbf7268d74a2a18af551de64ced03daab264799 100644 --- a/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py @@ -201,6 +201,20 @@ class InterleaveDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testEmptyInput(self): + iterator = ( + dataset_ops.Dataset.from_tensor_slices([]) + .repeat(None) + .interleave(dataset_ops.Dataset.from_tensors, cycle_length=2) + .make_initializable_iterator()) + init_op = iterator.initializer + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(init_op) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/kernel_tests/iterator_ops_cluster_test.py b/tensorflow/python/data/kernel_tests/iterator_ops_cluster_test.py index 45dfa13720b09c7bba979b72a339c13dcd2d827b..25c91b42dc65f849a680e65fc7fc2548c1cea8ea 100644 --- a/tensorflow/python/data/kernel_tests/iterator_ops_cluster_test.py +++ b/tensorflow/python/data/kernel_tests/iterator_ops_cluster_test.py @@ -17,10 +17,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + 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.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 @@ -28,6 +31,9 @@ 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 lookup_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import string_ops from tensorflow.python.platform import test @@ -103,6 +109,67 @@ class IteratorClusterTest(test.TestCase): "/job:worker/replica:0/task:1/cpu:0", workers[0].target) + def testCaptureHashTableInSharedIterator(self): + worker, _ = test_util.create_local_cluster(1, 1) + + # 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, + shared_name="shared_table") + + 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(shared_name="shared_iterator")) + init_op = iterator.initializer + get_next = iterator.get_next() + + with session.Session(worker[0].target) as sess: + sess.run(table.init) + sess.run(init_op) + self.assertAllEqual([0, 0, -1, 1, 2], sess.run(get_next)) + + with session.Session(worker[0].target) as sess: + self.assertAllEqual([2, 0], sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + 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(None) -> PrefetchDataset(100). + worker, _ = test_util.create_local_cluster(1, 1) + + components = (np.arange(1000), + np.array([[1, 2, 3]]) * np.arange(1000)[:, np.newaxis], + np.array(37.0) * np.arange(1000)) + + def _map_fn(x, y, z): + return math_ops.square(x), math_ops.square(y), math_ops.square(z) + + dataset = ( + dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) + .repeat(None).prefetch(10000)) + + iterator = dataset.make_initializable_iterator() + init_op = iterator.initializer + get_next = iterator.get_next() + + with session.Session(worker[0].target) as sess: + sess.run(init_op) + for _ in range(3): + sess.run(get_next) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/kernel_tests/iterator_ops_test.py b/tensorflow/python/data/kernel_tests/iterator_ops_test.py index 23c6d7385f8d4a12019fa514f349f2598d9629de..4a14a915bdb33f1ac6e8fc1839b32bc81fa8de05 100644 --- a/tensorflow/python/data/kernel_tests/iterator_ops_test.py +++ b/tensorflow/python/data/kernel_tests/iterator_ops_test.py @@ -22,6 +22,7 @@ import warnings import numpy as np +from tensorflow.core.protobuf import cluster_pb2 from tensorflow.core.protobuf import config_pb2 from tensorflow.python.client import session from tensorflow.python.data.ops import dataset_ops @@ -44,6 +45,7 @@ from tensorflow.python.ops import script_ops from tensorflow.python.ops import variables from tensorflow.python.platform import test from tensorflow.python.training import server_lib +from tensorflow.python.util import compat class IteratorTest(test.TestCase): @@ -63,8 +65,9 @@ class IteratorTest(test.TestCase): def testCapturingStateInOneShotRaisesException(self): var = variables.Variable(37.0, name="myvar") - dataset = (dataset_ops.Dataset.from_tensor_slices([0.0, 1.0, 2.0]) - .map(lambda x: x + var)) + dataset = ( + dataset_ops.Dataset.from_tensor_slices([0.0, 1.0, 2.0]) + .map(lambda x: x + var)) with self.assertRaisesRegexp( ValueError, r"`Dataset.make_one_shot_iterator\(\)` does not support " "datasets that capture stateful objects.+myvar"): @@ -78,8 +81,9 @@ class IteratorTest(test.TestCase): 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()) + 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], @@ -103,8 +107,9 @@ class IteratorTest(test.TestCase): 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()) + 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], @@ -125,10 +130,13 @@ class IteratorTest(test.TestCase): 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()) + + 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() @@ -159,8 +167,8 @@ class IteratorTest(test.TestCase): # 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) + 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): @@ -169,6 +177,7 @@ class IteratorTest(test.TestCase): # 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)) @@ -177,7 +186,8 @@ class IteratorTest(test.TestCase): num_threads = 8 threads = [ - self.checkedThread(consumer_thread) for _ in range(num_threads)] + self.checkedThread(consumer_thread) for _ in range(num_threads) + ] for t in threads: t.start() for t in threads: @@ -205,24 +215,24 @@ class IteratorTest(test.TestCase): 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)] + 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) - ) + 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() @@ -231,9 +241,10 @@ class IteratorTest(test.TestCase): # 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")) + 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() @@ -269,8 +280,9 @@ class IteratorTest(test.TestCase): 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()) + iterator = ( + dataset_ops.Dataset.from_tensors(components) + .make_initializable_iterator()) get_next = iterator.get_next() with self.test_session() as sess: @@ -320,8 +332,8 @@ class IteratorTest(test.TestCase): 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]) + iterator = iterator_ops.Iterator.from_structure( + (dtypes.int64, dtypes.float64), [None]) # Test validation of dataset argument. iterator = iterator_ops.Iterator.from_structure((dtypes.int64, @@ -337,18 +349,18 @@ class IteratorTest(test.TestCase): # 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)))) + 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)))) + 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]) @@ -370,33 +382,40 @@ class IteratorTest(test.TestCase): 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})) + 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}) + 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}) + sess.run( + next_element, feed_dict={handle_placeholder: iterator_4_handle}) def testIteratorStringHandleReuseTensorObject(self): dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) @@ -427,8 +446,8 @@ class IteratorTest(test.TestCase): self.assertIsNot(handle_with_name, handle_with_same_name) def testIteratorStringHandleError(self): - dataset_int_scalar = (dataset_ops.Dataset.from_tensor_slices([1, 2, - 3]).repeat()) + 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=[]) @@ -522,6 +541,58 @@ class IteratorTest(test.TestCase): target_placeholder: "/job:localhost/replica:0/task:0/cpu:1" }) + def testRemoteIteratorUsingRemoteCallOpMultiWorkers(self): + s1 = server_lib.Server.create_local_server() + s2 = server_lib.Server.create_local_server() + s3 = server_lib.Server.create_local_server() + + cluster_def = cluster_pb2.ClusterDef() + workers = cluster_def.job.add() + workers.name = "worker" + workers.tasks[0] = s1.target[len("grpc://"):] + workers.tasks[1] = s2.target[len("grpc://"):] + client = cluster_def.job.add() + client.name = "client" + client.tasks[0] = s3.target[len("grpc://"):] + config = config_pb2.ConfigProto(cluster_def=cluster_def) + + worker_devices = [ + "/job:worker/replica:0/task:%d/cpu:0" % i for i in range(2) + ] + itr_handles = [] + for device in worker_devices: + with ops.device(device): + src = dataset_ops.Dataset.from_tensor_slices([device]) + itr = src.make_one_shot_iterator() + itr_handles.append(itr.string_handle()) + + targets = dataset_ops.Dataset.from_tensor_slices(worker_devices) + handles = dataset_ops.Dataset.from_tensor_slices(itr_handles) + + @function.Defun(dtypes.string) + def loading_func(h): + remote_itr = iterator_ops.Iterator.from_string_handle( + h, itr.output_types, itr.output_shapes) + return remote_itr.get_next() + + def map_fn(target, handle): + return functional_ops.remote_call( + args=[handle], Tout=[dtypes.string], f=loading_func, target=target) + + with ops.device("/job:client"): + client_dataset = dataset_ops.Dataset.zip((targets, handles)).map(map_fn) + itr = client_dataset.make_initializable_iterator() + n = itr.get_next() + + with session.Session(s3.target, config=config) as sess: + sess.run(itr.initializer) + expected_values = worker_devices + for expected in expected_values: + self.assertEqual((compat.as_bytes(expected),), sess.run(n)) + + with self.assertRaises(errors.OutOfRangeError): + sess.run(n) + def testRemoteIteratorUsingRemoteCallOpDirectSessionGPUCPU(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") @@ -641,8 +712,7 @@ class IteratorTest(test.TestCase): with warnings.catch_warnings(record=True) as w: for _ in range(100): iterator.get_next() - self.assertEqual(100 - iterator_ops.GET_NEXT_CALL_WARNING_THRESHOLD, - len(w)) + self.assertEqual(100 - iterator_ops.GET_NEXT_CALL_WARNING_THRESHOLD, len(w)) for warning in w: self.assertTrue( iterator_ops.GET_NEXT_CALL_WARNING_MESSAGE in str(warning.message)) diff --git a/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py b/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py index 4e7691ee8144a19a62476281d86fb5df46dd3e4b..6442eb9ff554e61829796fb904342072d1846a32 100644 --- a/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py @@ -46,8 +46,9 @@ class ListFilesDatasetOpTest(test.TestCase): dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*')) with self.test_session() as sess: itr = dataset.make_one_shot_iterator() + next_element = itr.get_next() with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) + sess.run(next_element) def testSimpleDirectory(self): filenames = ['a', 'b', 'c'] @@ -56,13 +57,14 @@ class ListFilesDatasetOpTest(test.TestCase): dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*')) with self.test_session() as sess: itr = dataset.make_one_shot_iterator() + next_element = itr.get_next() 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()))) + produced_filenames.append(compat.as_bytes(sess.run(next_element))) self.assertItemsEqual(full_filenames, produced_filenames) with self.assertRaises(errors.OutOfRangeError): sess.run(itr.get_next()) @@ -73,12 +75,13 @@ class ListFilesDatasetOpTest(test.TestCase): with self.test_session() as sess: itr = dataset.make_initializable_iterator() + next_element = itr.get_next() sess.run( itr.initializer, feed_dict={filename_placeholder: path.join(self.tmp_dir, '*')}) with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) + sess.run(next_element) def testSimpleDirectoryInitializer(self): filenames = ['a', 'b', 'c'] @@ -89,6 +92,7 @@ class ListFilesDatasetOpTest(test.TestCase): with self.test_session() as sess: itr = dataset.make_initializable_iterator() + next_element = itr.get_next() sess.run( itr.initializer, feed_dict={filename_placeholder: path.join(self.tmp_dir, '*')}) @@ -98,7 +102,7 @@ class ListFilesDatasetOpTest(test.TestCase): 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()))) + produced_filenames.append(compat.as_bytes(sess.run(next_element))) self.assertItemsEqual(full_filenames, produced_filenames) @@ -114,6 +118,7 @@ class ListFilesDatasetOpTest(test.TestCase): with self.test_session() as sess: itr = dataset.make_initializable_iterator() + next_element = itr.get_next() sess.run( itr.initializer, feed_dict={filename_placeholder: path.join(self.tmp_dir, '*.py')}) @@ -123,7 +128,7 @@ class ListFilesDatasetOpTest(test.TestCase): 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()))) + produced_filenames.append(compat.as_bytes(sess.run(next_element))) self.assertItemsEqual(full_filenames, produced_filenames) with self.assertRaises(errors.OutOfRangeError): @@ -138,6 +143,7 @@ class ListFilesDatasetOpTest(test.TestCase): with self.test_session() as sess: itr = dataset.make_initializable_iterator() + next_element = itr.get_next() sess.run( itr.initializer, feed_dict={filename_placeholder: path.join(self.tmp_dir, '*.py*')}) @@ -147,13 +153,44 @@ class ListFilesDatasetOpTest(test.TestCase): 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()))) + produced_filenames.append(compat.as_bytes(sess.run(next_element))) self.assertItemsEqual(full_filenames, produced_filenames) with self.assertRaises(errors.OutOfRangeError): sess.run(itr.get_next()) + def testNoShuffle(self): + filenames = ['a', 'b', 'c'] + self._touchTempFiles(filenames) + + # Repeat the list twice and ensure that the order is the same each time. + # NOTE(mrry): This depends on an implementation detail of `list_files()`, + # which is that the list of files is captured when the iterator is + # initialized. Otherwise, or if e.g. the iterator were initialized more than + # once, it's possible that the non-determinism of `tf.matching_files()` + # would cause this test to fail. However, it serves as a useful confirmation + # that the `shuffle=False` argument is working as intended. + # TODO(b/73959787): Provide some ordering guarantees so that this test is + # more meaningful. + dataset = dataset_ops.Dataset.list_files( + path.join(self.tmp_dir, '*'), shuffle=False).repeat(2) + with self.test_session() as sess: + itr = dataset.make_one_shot_iterator() + next_element = itr.get_next() + + full_filenames = [] + produced_filenames = [] + for filename in filenames * 2: + full_filenames.append( + compat.as_bytes(path.join(self.tmp_dir, filename))) + produced_filenames.append(compat.as_bytes(sess.run(next_element))) + with self.assertRaises(errors.OutOfRangeError): + sess.run(itr.get_next()) + self.assertItemsEqual(full_filenames, produced_filenames) + self.assertEqual(produced_filenames[:len(filenames)], + produced_filenames[len(filenames):]) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/data/kernel_tests/map_dataset_op_test.py b/tensorflow/python/data/kernel_tests/map_dataset_op_test.py index 04d1abdb254feea1df6f1b8cfc5a512802107224..0791c614fa88700fdf2d0d673e168fc9784731a5 100644 --- a/tensorflow/python/data/kernel_tests/map_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/map_dataset_op_test.py @@ -602,6 +602,28 @@ class MapDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testParallelMapOutOfRangeError(self): + def raising_py_func(i): + if i == 100: + raise StopIteration() + else: + return i + + iterator = ( + dataset_ops.Dataset.range(105) + .map(lambda x: script_ops.py_func(raising_py_func, [x], dtypes.int64), + num_parallel_calls=2) + .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(100): + self.assertEqual(i, sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + class MapDatasetBenchmark(test.Benchmark): diff --git a/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py b/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py index d7140088c310767d40bd2cf3413c899375acab15..1ddedfda4e1c9d6b6949f796be1870f167435763 100644 --- a/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py +++ b/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py @@ -21,6 +21,7 @@ import gzip import os import zlib +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops from tensorflow.python.data.ops import readers from tensorflow.python.framework import constant_op @@ -736,12 +737,43 @@ class TFRecordDatasetTest(test.TestCase): one_mebibyte = 2**20 d = readers.TFRecordDataset(self.test_filenames, buffer_size=one_mebibyte) iterator = d.make_one_shot_iterator() + next_element = iterator.get_next() 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())) + self.assertAllEqual(self._record(j, i), sess.run(next_element)) with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) + sess.run(next_element) + + def testReadFromDatasetOfFiles(self): + files = dataset_ops.Dataset.from_tensor_slices(self.test_filenames) + d = readers.TFRecordDataset(files) + iterator = d.make_one_shot_iterator() + next_element = iterator.get_next() + 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(next_element)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + + def testReadTenEpochsFromDatasetOfFilesInParallel(self): + files = dataset_ops.Dataset.from_tensor_slices( + self.test_filenames).repeat(10) + d = readers.TFRecordDataset(files, num_parallel_reads=4) + iterator = d.make_one_shot_iterator() + next_element = iterator.get_next() + expected = [] + actual = [] + with self.test_session() as sess: + for _ in range(10): + for j in range(self._num_files): + for i in range(self._num_records): + expected.append(self._record(j, i)) + actual.append(sess.run(next_element)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + self.assertEqual(sorted(expected), sorted(actual)) if __name__ == "__main__": diff --git a/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py b/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py index ae08032e191487c38d73876374b24e8f6eefbc80..1d27b036eb804aa301b916b7ed0b7884f75e1a0f 100644 --- a/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py @@ -201,9 +201,7 @@ class SequenceDatasetTest(test.TestCase): with self.test_session() as sess: sess.run(init_op) - with self.assertRaisesRegexp( - errors.OutOfRangeError, - "Attempted to repeat an empty dataset infinitely."): + with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) diff --git a/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py b/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py index c089fb08c1082c1cf74d492796550980d6755591..5fcc48831f3ca744e015c92760f12ea4dbef2ff7 100644 --- a/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py @@ -132,6 +132,33 @@ class ShuffleDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testSeedZero(self): + """Test for same behavior when the seed is a Python or Tensor zero.""" + iterator = ( + dataset_ops.Dataset.range(10).shuffle(10, seed=0) + .make_one_shot_iterator()) + get_next = iterator.get_next() + + elems = [] + with self.test_session() as sess: + for _ in range(10): + elems.append(sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + seed_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) + iterator = ( + dataset_ops.Dataset.range(10).shuffle(10, seed=seed_placeholder) + .make_initializable_iterator()) + get_next = iterator.get_next() + + with self.test_session() as sess: + sess.run(iterator.initializer, feed_dict={seed_placeholder: 0}) + for elem in elems: + self.assertEqual(elem, sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + def testDefaultArguments(self): components = [0, 1, 2, 3, 4] iterator = (dataset_ops.Dataset.from_tensor_slices(components).shuffle(5) diff --git a/tensorflow/python/data/ops/BUILD b/tensorflow/python/data/ops/BUILD index f12b358a7dc35c18338171e489fa88ba1a82d11b..3119ab003794cb9bc0c748dfeb47597e0877f5fd 100644 --- a/tensorflow/python/data/ops/BUILD +++ b/tensorflow/python/data/ops/BUILD @@ -23,6 +23,7 @@ py_library( "//tensorflow/python:tensor_util", "//tensorflow/python:util", "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:random_seed", "//tensorflow/python/data/util:sparse", "//third_party/py/numpy", ], @@ -34,6 +35,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":dataset_ops", + "//tensorflow/python:array_ops", "//tensorflow/python:dataset_ops_gen", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", @@ -50,9 +52,11 @@ py_library( "//tensorflow/python:dataset_ops_gen", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", + "//tensorflow/python:resource_variable_ops", "//tensorflow/python:tensor_shape", "//tensorflow/python/data/util:nest", "//tensorflow/python/data/util:sparse", + "//tensorflow/python/eager:context", ], ) diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index 9c62d5700cc8b7197bde2877e48d474f93f7b940..390ce852b1337bf3ae387660e4b83b0b678f737f 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -26,23 +26,26 @@ import six from tensorflow.python.data.ops import iterator_ops from tensorflow.python.data.util import nest +from tensorflow.python.data.util import random_seed 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 from tensorflow.python.framework import function from tensorflow.python.framework import ops -from tensorflow.python.framework import random_seed from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_dataset_ops from tensorflow.python.ops import gen_io_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import script_ops from tensorflow.python.util import deprecation +from tensorflow.python.util.tf_export import tf_export +@tf_export("data.Dataset") class Dataset(object): """Represents a potentially large set of elements. @@ -88,7 +91,7 @@ class Dataset(object): Raises: RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "dataset.make_initializable_iterator is not supported when eager " "execution is enabled.") @@ -108,11 +111,11 @@ class Dataset(object): self.output_types, self.output_shapes, self.output_classes) - def make_one_shot_iterator(self): + def __iter__(self): """Creates an `Iterator` for enumerating the elements of this dataset. - Note: The returned iterator will be initialized automatically. - A "one-shot" iterator does not currently support re-initialization. + The returned iterator implements the Python iterator protocol and therefore + can only be used in eager mode. Returns: An `Iterator` over the elements of this dataset. @@ -120,10 +123,23 @@ class Dataset(object): Raises: RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): - raise RuntimeError( - "dataset.make_one_shot_iterator is not supported when eager " - "execution is enabled.") + if context.executing_eagerly(): + return iterator_ops.EagerIterator(self) + else: + raise RuntimeError("dataset.__iter__() is only supported when eager " + "execution is enabled.") + + def make_one_shot_iterator(self): + """Creates an `Iterator` for enumerating the elements of this dataset. + + Note: The returned iterator will be initialized automatically. + A "one-shot" iterator does not currently support re-initialization. + + Returns: + An `Iterator` over the elements of this dataset. + """ + if context.executing_eagerly(): + return iterator_ops.EagerIterator(self) # NOTE(mrry): We capture by value here to ensure that `_make_dataset()` is # a 0-argument function. @function.Defun(capture_by_value=True) @@ -329,10 +345,10 @@ class Dataset(object): generator_state = Dataset._GeneratorState(generator) - def get_iterator_id_map_fn(unused_dummy): + def get_iterator_id_fn(unused_dummy): """Creates a unique `iterator_id` for each pass over the dataset. - The "iterator_id" disambiguates between multiple concurrently + The returned `iterator_id` disambiguates between multiple concurrently existing iterators. Args: @@ -345,7 +361,7 @@ class Dataset(object): return script_ops.py_func( generator_state.get_next_id, [], dtypes.int64, stateful=True) - def generator_map_fn(iterator_id_t): + def generator_next_fn(iterator_id_t): """Generates the next element from iterator with ID `iterator_id_t`. We map this function across an infinite repetition of the @@ -361,11 +377,9 @@ class Dataset(object): def generator_py_func(iterator_id): """A `py_func` that will be called to invoke the iterator.""" - try: - values = next(generator_state.get_iterator(iterator_id)) - except StopIteration: - generator_state.iterator_completed(iterator_id) - raise StopIteration("Iteration finished.") + # `next()` raises `StopIteration` when there are no more + # elements remaining to be generated. + values = next(generator_state.get_iterator(iterator_id)) # Use the same _convert function from the py_func() implementation to # convert the returned values to arrays early, so that we can inspect @@ -406,17 +420,31 @@ class Dataset(object): return nest.pack_sequence_as(output_types, flat_values) + def finalize_fn(iterator_id_t): + """Releases host-side state for the iterator with ID `iterator_id_t`.""" + + def finalize_py_func(iterator_id): + generator_state.iterator_completed(iterator_id) + # We return a dummy value so that the `finalize_fn` has a valid + # signature. + # NOTE(mrry): Explicitly create an array of `np.int64` because implicit + # casting in `py_func()` will create an array of `np.int32` on Windows, + # leading to a runtime error. + return np.array(0, dtype=np.int64) + + return script_ops.py_func( + finalize_py_func, [iterator_id_t], dtypes.int64, stateful=True) + # This function associates each traversal of `generator` with a unique # iterator ID. - def flat_map_fn(iterator_id_t): - # First, generate an infinite dataset containing the iterator ID repeated - # forever. - repeated_id = Dataset.from_tensors(iterator_id_t).repeat(None) - - # The `generator_map_fn` gets the next element from the iterator with the - # relevant ID, and raises StopIteration when that iterator contains no + def flat_map_fn(dummy_arg): + # The `get_iterator_id_fn` gets a unique ID for the current instance of + # of the generator. + # The `generator_next_fn` gets the next element from the iterator with the + # given ID, and raises StopIteration when that iterator contains no # more elements. - return repeated_id.map(generator_map_fn) + return _GeneratorDataset(dummy_arg, get_iterator_id_fn, generator_next_fn, + finalize_fn) # A single-element dataset that, each time it is evaluated, contains a # freshly-generated and unique (for the returned dataset) int64 @@ -424,7 +452,7 @@ class Dataset(object): # is encapsulated in `generator_state`, and captured in # `get_iterator_id_map_fn`. dummy = 0 - id_dataset = Dataset.from_tensors(dummy).map(get_iterator_id_map_fn) + id_dataset = Dataset.from_tensors(dummy) # A dataset that contains all of the elements generated by a # single iterator created from `generator`, identified by the @@ -543,7 +571,7 @@ class Dataset(object): return PrefetchDataset(self, buffer_size) @staticmethod - def list_files(file_pattern): + def list_files(file_pattern, shuffle=None): """A dataset of all files matching a pattern. Example: @@ -556,14 +584,31 @@ class Dataset(object): - /path/to/dir/b.py - /path/to/dir/c.py + NOTE: The order of the file names returned can be non-deterministic even + when `shuffle` is `False`. + Args: file_pattern: A string or scalar string `tf.Tensor`, representing the filename pattern that will be matched. + shuffle: (Optional.) If `True`, the file names will be shuffled randomly. + Defaults to `True`. Returns: Dataset: A `Dataset` of strings corresponding to file names. """ - return Dataset.from_tensor_slices(gen_io_ops.matching_files(file_pattern)) + # TODO(b/73959787): Add a `seed` argument and make the `shuffle=False` + # behavior deterministic (e.g. by sorting the filenames). + if shuffle is None: + shuffle = True + matching_files = gen_io_ops.matching_files(file_pattern) + dataset = Dataset.from_tensor_slices(matching_files) + if shuffle: + # NOTE(mrry): The shuffle buffer size must be greater than zero, but the + # list of files might be empty. + buffer_size = math_ops.maximum( + array_ops.shape(matching_files, out_type=dtypes.int64)[0], 1) + dataset = dataset.shuffle(buffer_size) + return dataset def repeat(self, count=None): """Repeats this dataset `count` times. @@ -742,11 +787,31 @@ class Dataset(object): def padded_batch(self, batch_size, padded_shapes, padding_values=None): """Combines consecutive elements of this dataset into padded batches. - Like `Dataset.dense_to_sparse_batch()`, this method combines - multiple consecutive elements of this dataset, which might have - different shapes, into a single element. The tensors in the - resulting element have an additional outer dimension, and are - padded to the respective shape in `padded_shapes`. + This transformation combines multiple consecutive elements of the input + dataset into a single element. Like @{tf.data.Dataset.batch}, the tensors + in the resulting element have an additional outer dimension, which will be + `batch_size` for all but the last element, and `N % batch_size` for the + last element (where `N` is the number of elements in this dataset). Unlike + @{tf.data.Dataset.batch}, the elements may have different shapes for some + of their components, and this transformation will pad each component to + the respective shape in `padding_shapes`. The `padding_shapes` argument + determines the resulting shape for each dimension of each component in an + output element: + + * If the dimension is a constant (e.g. `tf.Dimension(37)`), the component + will be padded out to that length in that dimension. + * If the dimension is unknown (e.g. `tf.Dimension(None)`), the component + will be padded out to the maximum length of all elements in that + dimension. + + NOTE: If the number of elements (`N`) in this dataset is not an exact + multiple of `batch_size`, the final batch contain smaller tensors with + shape `N % batch_size` in the batch dimension. If your program depends on + the batches having the same shape, consider using the + @{tf.contrib.data.padded_batch_and_drop_remainder} transformation instead. + + See also @{tf.contrib.data.dense_to_sparse_batch}, which combines elements + that may have different shapes into a @{tf.SparseTensor}. Args: batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of @@ -899,10 +964,11 @@ class Dataset(object): Args: transformation_func: A function that takes one `Dataset` argument and - returns a `Dataset`. + returns a `Dataset`. Returns: - Dataset: The `Dataset` returned by applying `transformation_func` to this dataset. + Dataset: The `Dataset` returned by applying `transformation_func` to this + dataset. """ dataset = transformation_func(self) if not isinstance(dataset, Dataset): @@ -1028,6 +1094,196 @@ class SparseTensorSliceDataset(Dataset): return (dtypes.int64, self._sparse_tensor.dtype, dtypes.int64) +class _GeneratorDataset(Dataset): + """A `Dataset` that generates elements by invoking a function.""" + + def __init__(self, init_args, init_func, next_func, finalize_func): + """Constructs a `_GeneratorDataset`. + + Args: + init_args: A nested structure representing the arguments to `init_func`. + init_func: A TensorFlow function that will be called on `init_args` each + time a C++ iterator over this dataset is constructed. Returns a nested + structure representing the "state" of the dataset. + next_func: A TensorFlow function that will be called on the result of + `init_func` to produce each element, and that raises `OutOfRangeError` + to terminate iteration. + finalize_func: A TensorFlow function that will be called on the result of + `init_func` immediately before a C++ iterator over this dataset is + destroyed. The return value is ignored. + """ + super(_GeneratorDataset, self).__init__() + # These members will be initialized by `tf_init_func`. + self._state_classes = None + self._state_shapes = None + self._state_types = None + + self._init_args = init_args + + init_args_classes = sparse.get_classes(init_args) + init_args_shapes = nest.pack_sequence_as( + init_args, [t.get_shape() for t in nest.flatten(init_args)]) + init_args_types = nest.pack_sequence_as( + init_args, [t.dtype for t in nest.flatten(init_args)]) + + @function.Defun(*nest.flatten( + sparse.as_dense_types(init_args_types, init_args_classes))) + def tf_init_func(*args): + """A wrapper for Defun that facilitates shape inference.""" + dense_shapes = sparse.as_dense_shapes(init_args_shapes, init_args_classes) + for arg, shape in zip(args, nest.flatten(dense_shapes)): + arg.set_shape(shape) + + nested_args = nest.pack_sequence_as(init_args_classes, args) + nested_args = sparse.deserialize_sparse_tensors( + nested_args, init_args_types, init_args_shapes, init_args_classes) + if _should_unpack_args(nested_args): + ret = init_func(*nested_args) + else: + ret = init_func(nested_args) + + # If `init_func` returns a list of tensors, `nest.flatten()` and + # `ops.convert_to_tensor()` would conspire to attempt to stack + # those tensors into a single tensor, because the customized + # version of `nest.flatten()` does not recurse into lists. Since + # it is more likely that the list arose from returning the + # result of an operation (such as `tf.py_func()`) that returns a + # list of not-necessarily-stackable tensors, we treat the + # returned value is a `tuple` instead. A user wishing to pack + # the return value into a single tensor can use an explicit + # `tf.stack()` before returning. + if isinstance(ret, list): + ret = tuple(ret) + + # Convert any `SparseTensorValue`s to `SparseTensor`s. + ret = nest.pack_sequence_as(ret, [ + sparse_tensor_lib.SparseTensor.from_value(t) + if sparse_tensor_lib.is_sparse(t) else t for t in nest.flatten(ret) + ]) + + self._state_classes = sparse.get_classes(ret) + self._state_shapes = nest.pack_sequence_as( + ret, [t.get_shape() for t in nest.flatten(ret)]) + self._state_types = nest.pack_sequence_as( + ret, [t.dtype for t in nest.flatten(ret)]) + + # Serialize any sparse tensors and convert result to tensors. + ret = nest.pack_sequence_as(ret, [ + ops.convert_to_tensor(t) + for t in nest.flatten(sparse.serialize_sparse_tensors(ret)) + ]) + return nest.flatten(ret) + + self._init_func = tf_init_func + self._init_func.add_to_graph(ops.get_default_graph()) + + # These members will be initialized by `tf_next_func`. + self._output_classes = None + self._output_shapes = None + self._output_types = None + + @function.Defun(*nest.flatten( + sparse.as_dense_types(self._state_types, self._state_classes))) + def tf_next_func(*args): + """A wrapper for Defun that facilitates shape inference.""" + # Pass in shape information from the input_dataset. + dense_shapes = sparse.as_dense_shapes(self._state_shapes, + self._state_classes) + for arg, shape in zip(args, nest.flatten(dense_shapes)): + arg.set_shape(shape) + + nested_args = nest.pack_sequence_as(self._state_classes, args) + nested_args = sparse.deserialize_sparse_tensors( + nested_args, self._state_types, self._state_shapes, + self._state_classes) + if _should_unpack_args(nested_args): + ret = next_func(*nested_args) + else: + ret = next_func(nested_args) + + # If `next_func` returns a list of tensors, `nest.flatten()` and + # `ops.convert_to_tensor()` would conspire to attempt to stack + # those tensors into a single tensor, because the customized + # version of `nest.flatten()` does not recurse into lists. Since + # it is more likely that the list arose from returning the + # result of an operation (such as `tf.py_func()`) that returns a + # list of not-necessarily-stackable tensors, we treat the + # returned value is a `tuple` instead. A user wishing to pack + # the return value into a single tensor can use an explicit + # `tf.stack()` before returning. + if isinstance(ret, list): + ret = tuple(ret) + + # Convert any `SparseTensorValue`s to `SparseTensor`s. + ret = nest.pack_sequence_as(ret, [ + sparse_tensor_lib.SparseTensor.from_value(t) + if sparse_tensor_lib.is_sparse(t) else t for t in nest.flatten(ret) + ]) + + self._output_classes = sparse.get_classes(ret) + self._output_shapes = nest.pack_sequence_as( + ret, [t.get_shape() for t in nest.flatten(ret)]) + self._output_types = nest.pack_sequence_as( + ret, [t.dtype for t in nest.flatten(ret)]) + + # Serialize any sparse tensors and convert result to tensors. + ret = nest.pack_sequence_as(ret, [ + ops.convert_to_tensor(t) + for t in nest.flatten(sparse.serialize_sparse_tensors(ret)) + ]) + return nest.flatten(ret) + + self._next_func = tf_next_func + self._next_func.add_to_graph(ops.get_default_graph()) + + @function.Defun(*nest.flatten( + sparse.as_dense_types(self._state_types, self._state_classes))) + def tf_finalize_func(*args): + """A wrapper for Defun that facilitates shape inference.""" + # Pass in shape information from the state. + dense_shapes = sparse.as_dense_shapes(self._state_shapes, + self._state_classes) + for arg, shape in zip(args, nest.flatten(dense_shapes)): + arg.set_shape(shape) + + nested_args = nest.pack_sequence_as(self._state_classes, args) + nested_args = sparse.deserialize_sparse_tensors( + nested_args, self._state_types, self._state_shapes, + self._state_classes) + if _should_unpack_args(nested_args): + return finalize_func(*nested_args) + else: + return finalize_func(nested_args) + + self._finalize_func = tf_finalize_func + self._finalize_func.add_to_graph(ops.get_default_graph()) + + def _as_variant_tensor(self): + return gen_dataset_ops.generator_dataset( + nest.flatten(self._init_args) + self._init_func.captured_inputs, + self._next_func.captured_inputs, + self._finalize_func.captured_inputs, + init_func=self._init_func, + next_func=self._next_func, + finalize_func=self._finalize_func, + 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_classes(self): + return self._output_classes + + @property + def output_shapes(self): + return self._output_shapes + + @property + def output_types(self): + return self._output_types + + class ZipDataset(Dataset): """A `Dataset` that zips its inputs together.""" @@ -1277,16 +1533,7 @@ class ShuffleDataset(Dataset): self._input_dataset = input_dataset self._buffer_size = ops.convert_to_tensor( buffer_size, dtype=dtypes.int64, name="buffer_size") - seed, seed2 = random_seed.get_seed(seed) - if seed is None: - self._seed = constant_op.constant(0, dtype=dtypes.int64, name="seed") - else: - self._seed = ops.convert_to_tensor(seed, dtype=dtypes.int64, name="seed") - if seed2 is None: - self._seed2 = constant_op.constant(0, dtype=dtypes.int64, name="seed2") - else: - self._seed2 = ops.convert_to_tensor( - seed2, dtype=dtypes.int64, name="seed2") + self._seed, self._seed2 = random_seed.get_seed(seed) if reshuffle_each_iteration is None: self._reshuffle_each_iteration = True else: @@ -1454,6 +1701,19 @@ def _padding_value_to_tensor(value, output_type): return value +def _default_padding(input_dataset): + + def make_zero(t): + if t.base_dtype == dtypes.string: + return "" + elif t.base_dtype == dtypes.variant: + raise TypeError("Unable to create padding for field of type 'variant'") + else: + return np.zeros_like(t.as_numpy_dtype()) + + return nest.map_structure(make_zero, input_dataset.output_types) + + class PaddedBatchDataset(Dataset): """A `Dataset` that batches and pads contiguous elements from its input.""" @@ -1469,23 +1729,13 @@ class PaddedBatchDataset(Dataset): batch_size, dtype=dtypes.int64, name="batch_size") padding_values = ( padding_values - if padding_values is not None else self._default_padding(input_dataset)) + if padding_values is not None else _default_padding(input_dataset)) self._padded_shapes = nest.map_structure_up_to( input_dataset.output_shapes, _partial_shape_to_tensor, padded_shapes) self._padding_values = nest.map_structure_up_to( input_dataset.output_shapes, _padding_value_to_tensor, padding_values, input_dataset.output_types) - def _default_padding(self, input_dataset): - - def make_zero(t): - if t.base_dtype == dtypes.string: - return "" - else: - return np.zeros_like(t.as_numpy_dtype()) - - return nest.map_structure(make_zero, input_dataset.output_types) - def _as_variant_tensor(self): return gen_dataset_ops.padded_batch_dataset( self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access diff --git a/tensorflow/python/data/ops/iterator_ops.py b/tensorflow/python/data/ops/iterator_ops.py index 53a3244ce1948803be5d8ee0f7db12fa40c8a32f..d79b9d6011b6ebd00a47d572165cdbba8a31bd32 100644 --- a/tensorflow/python/data/ops/iterator_ops.py +++ b/tensorflow/python/data/ops/iterator_ops.py @@ -17,14 +17,19 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import threading import warnings 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 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 gen_dataset_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.util.tf_export import tf_export # NOTE(mrry): It is legitimate to call `Iterator.get_next()` multiple @@ -43,10 +48,12 @@ GET_NEXT_CALL_WARNING_MESSAGE = ( "This often indicates that `Iterator.get_next()` is being called inside " "a training loop, which will cause gradual slowdown and eventual resource " "exhaustion. If this is the case, restructure your code to call " - "`next_element = iterator.get_next() once outside the loop, and use " - "`next_element` inside the loop.") + "`next_element = iterator.get_next()` once outside the loop, and use " + "`next_element` as the input to some computation that is invoked inside " + "the loop.") +@tf_export("data.Iterator") class Iterator(object): """Represents the state of iterating through a `Dataset`.""" @@ -301,7 +308,42 @@ class Iterator(object): dataset._as_variant_tensor(), self._iterator_resource, name=name) # pylint: disable=protected-access def get_next(self, name=None): - """Returns a nested structure of `tf.Tensor`s containing the next element. + """Returns a nested structure of `tf.Tensor`s representing the next element. + + In graph mode, you should typically call this method *once* and use its + result as the input to another computation. A typical loop will then call + @{tf.Session.run} on the result of that computation. The loop will terminate + when the `Iterator.get_next()` operation raises + @{tf.errors.OutOfRangeError}. The following skeleton shows how to use + this method when building a training loop: + + ```python + dataset = ... # A `tf.data.Dataset` object. + iterator = dataset.make_initializable_iterator() + next_element = iterator.get_next() + + # Build a TensorFlow graph that does something with each element. + loss = model_function(next_element) + optimizer = ... # A `tf.train.Optimizer` object. + train_op = optimizer.minimize(loss) + + with tf.Session() as sess: + try: + while True: + sess.run(train_op) + except tf.errors.OutOfRangeError: + pass + ``` + + NOTE: It is legitimate to call `Iterator.get_next()` multiple times, e.g. + when you are distributing different elements to multiple devices in a single + step. However, a common pitfall arises when users call `Iterator.get_next()` + in each iteration of their training loop. `Iterator.get_next()` adds ops to + the graph, and executing each op allocates resources (including threads); as + a consequence, invoking it in every iteration of a training loop causes + slowdown and eventual resource exhaustion. To guard against this outcome, we + log a warning when the number of uses crosses a fixed threshold of + suspiciousness. Args: name: (Optional.) A name for the created operation. @@ -374,3 +416,147 @@ class Iterator(object): of an element of this dataset. """ return self._output_types + + +_uid_counter = 0 +_uid_lock = threading.Lock() + + +def _generate_shared_name(prefix): + with _uid_lock: + global _uid_counter + uid = _uid_counter + _uid_counter += 1 + return "{}{}".format(prefix, uid) + + +class EagerIterator(object): + """An iterator producing tf.Tensor objects from a tf.data.Dataset.""" + + def __init__(self, dataset): + """Creates a new iterator over the given dataset. + + For example: + ```python + dataset = tf.data.Dataset.range(4) + for x in Iterator(dataset): + print(x) + ``` + + Tensors produced will be placed on the device on which this iterator object + was created. + + Args: + dataset: A `tf.data.Dataset` object. + + Raises: + RuntimeError: When invoked without eager execution enabled. + """ + + if not context.executing_eagerly(): + raise RuntimeError( + "{} objects can only be used when eager execution is enabled, use " + "tf.data.Dataset.make_initializable_iterator or " + "tf.data.Dataset.make_one_shot_iterator for graph construction". + 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( + 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( + 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) + # Delete the resource when this object is deleted + self._resource_deleter = resource_variable_ops.EagerResourceDeleter( + handle=self._resource, handle_device="/device:CPU:0") + self._device = context.context().device_name + + def __iter__(self): + return self + + def __next__(self): # For Python 3 compatibility + return self.next() + + def _next_internal(self): + """Returns a nested structure of `tf.Tensor`s containing the next element. + """ + with ops.device(self._device): + # 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/python/data/ops/readers.py b/tensorflow/python/data/ops/readers.py index 830dc5cec4a54469d001f0ba57d1adc7bc5efd11..6c493d8163b051b2e724335923d7b4c721523083 100644 --- a/tensorflow/python/data/ops/readers.py +++ b/tensorflow/python/data/ops/readers.py @@ -17,19 +17,25 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.data.ops.dataset_ops import Dataset +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import convert +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 function 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.util.tf_export import tf_export # TODO(b/64974358): Increase default buffer size to 256 MB. _DEFAULT_READER_BUFFER_SIZE_BYTES = 256 * 1024 # 256 KB -class TextLineDataset(Dataset): +@tf_export("data.TextLineDataset") +class TextLineDataset(dataset_ops.Dataset): """A `Dataset` comprising lines from one or more text files.""" def __init__(self, filenames, compression_type=None, buffer_size=None): @@ -71,7 +77,7 @@ class TextLineDataset(Dataset): return dtypes.string -class TFRecordDataset(Dataset): +class _TFRecordDataset(dataset_ops.Dataset): """A `Dataset` comprising records from one or more TFRecord files.""" def __init__(self, filenames, compression_type=None, buffer_size=None): @@ -84,7 +90,7 @@ class TFRecordDataset(Dataset): buffer_size: (Optional.) A `tf.int64` scalar representing the number of bytes in the read buffer. 0 means no buffering. """ - super(TFRecordDataset, self).__init__() + super(_TFRecordDataset, self).__init__() # Force the type to string even if filenames is an empty list. self._filenames = ops.convert_to_tensor( filenames, dtypes.string, name="filenames") @@ -115,7 +121,159 @@ class TFRecordDataset(Dataset): return dtypes.string -class FixedLengthRecordDataset(Dataset): +class ParallelInterleaveDataset(dataset_ops.Dataset): + """A `Dataset` that maps a function over its input and flattens the result.""" + + def __init__(self, input_dataset, map_func, cycle_length, block_length, + sloppy, buffer_output_elements, prefetch_input_elements): + """See `tf.contrib.data.parallel_interleave()` for details.""" + super(ParallelInterleaveDataset, self).__init__() + self._input_dataset = input_dataset + + @function.Defun(*nest.flatten( + sparse.as_dense_types(input_dataset.output_types, + input_dataset.output_classes))) + def tf_map_func(*args): + """A wrapper for Defun that facilitates shape inference.""" + # Pass in shape information from the input_dataset. + dense_shapes = sparse.as_dense_shapes(input_dataset.output_shapes, + input_dataset.output_classes) + for arg, shape in zip(args, nest.flatten(dense_shapes)): + arg.set_shape(shape) + + nested_args = nest.pack_sequence_as(input_dataset.output_types, args) + nested_args = sparse.deserialize_sparse_tensors( + nested_args, input_dataset.output_types, input_dataset.output_shapes, + input_dataset.output_classes) + if dataset_ops._should_unpack_args(nested_args): # pylint: disable=protected-access + dataset = map_func(*nested_args) + else: + dataset = map_func(nested_args) + + if not isinstance(dataset, dataset_ops.Dataset): + raise TypeError("`map_func` must return a `Dataset` object.") + + self._output_classes = dataset.output_classes + self._output_types = dataset.output_types + self._output_shapes = dataset.output_shapes + + return dataset._as_variant_tensor() # pylint: disable=protected-access + + self._map_func = tf_map_func + self._map_func.add_to_graph(ops.get_default_graph()) + + self._cycle_length = ops.convert_to_tensor( + cycle_length, dtype=dtypes.int64, name="cycle_length") + self._block_length = ops.convert_to_tensor( + block_length, dtype=dtypes.int64, name="block_length") + self._sloppy = ops.convert_to_tensor( + sloppy, dtype=dtypes.bool, name="sloppy") + self._buffer_output_elements = convert.optional_param_to_tensor( + "buffer_output_elements", + buffer_output_elements, + argument_default=2 * block_length) + self._prefetch_input_elements = convert.optional_param_to_tensor( + "prefetch_input_elements", + prefetch_input_elements, + argument_default=2 * cycle_length) + + def _as_variant_tensor(self): + return gen_dataset_ops.parallel_interleave_dataset( + self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access + self._map_func.captured_inputs, + self._cycle_length, + self._block_length, + self._sloppy, + self._buffer_output_elements, + self._prefetch_input_elements, + f=self._map_func, + 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_classes(self): + return self._output_classes + + @property + def output_shapes(self): + return self._output_shapes + + @property + def output_types(self): + return self._output_types + + +@tf_export("data.TFRecordDataset") +class TFRecordDataset(dataset_ops.Dataset): + """A `Dataset` comprising records from one or more TFRecord files.""" + + def __init__(self, filenames, compression_type=None, buffer_size=None, + num_parallel_reads=None): + """Creates a `TFRecordDataset` to read for one or more TFRecord files. + + NOTE: The `num_parallel_reads` argument can be used to improve performance + when reading from a remote filesystem. + + Args: + filenames: A `tf.string` tensor or `tf.data.Dataset` 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. + num_parallel_reads: (Optional.) A `tf.int64` scalar representing the + number of files to read in parallel. Defaults to reading files + sequentially. + + Raises: + TypeError: If any argument does not have the expected type. + ValueError: If any argument does not have the expected shape. + """ + super(TFRecordDataset, self).__init__() + if isinstance(filenames, dataset_ops.Dataset): + if filenames.output_types != dtypes.string: + raise TypeError( + "`filenames` must be a `tf.data.Dataset` of `tf.string` elements.") + if not filenames.output_shapes.is_compatible_with(tensor_shape.scalar()): + raise ValueError( + "`filenames` must be a `tf.data.Dataset` of scalar `tf.string` " + "elements.") + else: + filenames = ops.convert_to_tensor(filenames, dtype=dtypes.string) + filenames = array_ops.reshape(filenames, [-1], name="flat_filenames") + filenames = dataset_ops.Dataset.from_tensor_slices(filenames) + + def read_one_file(filename): + return _TFRecordDataset(filename, compression_type, buffer_size) + + if num_parallel_reads is None: + self._impl = filenames.flat_map(read_one_file) + else: + self._impl = ParallelInterleaveDataset( + filenames, read_one_file, cycle_length=num_parallel_reads, + block_length=1, sloppy=False, buffer_output_elements=None, + prefetch_input_elements=None) + + def _as_variant_tensor(self): + return self._impl._as_variant_tensor() # pylint: disable=protected-access + + @property + def output_classes(self): + return self._impl.output_classes + + @property + def output_shapes(self): + return self._impl.output_shapes + + @property + def output_types(self): + return self._impl.output_types + + +@tf_export("data.FixedLengthRecordDataset") +class FixedLengthRecordDataset(dataset_ops.Dataset): """A `Dataset` of fixed-length records from one or more binary files.""" def __init__(self, diff --git a/tensorflow/python/data/util/BUILD b/tensorflow/python/data/util/BUILD index e32c7b54a48dd887c2748897c3ce3661aab9f497..b1bdbdab37b63667b475c732df7a47d9e57f2b19 100644 --- a/tensorflow/python/data/util/BUILD +++ b/tensorflow/python/data/util/BUILD @@ -86,6 +86,30 @@ py_test( ], ) +py_library( + name = "random_seed", + srcs = ["random_seed.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:constant_op", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework", + ], +) + +py_test( + name = "random_seed_test", + size = "small", + srcs = ["random_seed_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":random_seed", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:util", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/python/data/util/nest.py b/tensorflow/python/data/util/nest.py index df5498be5f0643dd533cb522423e5728d389d7fb..e90ce3fb40af68fb68d6ee8bac6892848d8c5a79 100644 --- a/tensorflow/python/data/util/nest.py +++ b/tensorflow/python/data/util/nest.py @@ -383,8 +383,8 @@ def assert_shallow_structure(shallow_tree, input_tree, check_types=True): "structure has keys %s, while shallow structure has keys %s." % (list(_six.iterkeys(input_tree)), list(_six.iterkeys(shallow_tree)))) - input_tree = list(_six.iteritems(input_tree)) - shallow_tree = list(_six.iteritems(shallow_tree)) + input_tree = list(sorted(_six.iteritems(input_tree))) + shallow_tree = list(sorted(_six.iteritems(shallow_tree))) for shallow_branch, input_branch in zip(shallow_tree, input_tree): assert_shallow_structure(shallow_branch, input_branch, @@ -479,8 +479,8 @@ def map_structure_up_to(shallow_tree, func, *inputs): The `inputs`, can be thought of as having the same structure as `shallow_tree`, but with leaf nodes that are themselves tree structures. - This function, therefore, will return something with the same base structure as - `shallow_tree`. + This function, therefore, will return something with the same base structure + as `shallow_tree`. Examples: diff --git a/tensorflow/python/data/util/nest_test.py b/tensorflow/python/data/util/nest_test.py index 90dd7dfe7775b2f10611e5579784fbda63fc9669..ff380815a4a32192de621888199e66355f9b4635 100644 --- a/tensorflow/python/data/util/nest_test.py +++ b/tensorflow/python/data/util/nest_test.py @@ -277,6 +277,10 @@ class NestTest(test.TestCase): with self.assertRaisesRegexp(ValueError, expected_message): nest.assert_shallow_structure(inp_ab2, inp_ab1) + inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))]) + inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)]) + nest.assert_shallow_structure(inp_ab, inp_ba) + def testFlattenUpTo(self): input_tree = (((2, 2), (3, 3)), ((4, 9), (5, 5))) shallow_tree = ((True, True), (False, True)) diff --git a/tensorflow/python/data/util/random_seed.py b/tensorflow/python/data/util/random_seed.py new file mode 100644 index 0000000000000000000000000000000000000000..e2c9d8672f94587fd3164f25f97b44a97526be07 --- /dev/null +++ b/tensorflow/python/data/util/random_seed.py @@ -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. +# ============================================================================== +"""Utilities for generating Tensor-valued random seeds.""" + +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.framework import random_seed +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops + + +def get_seed(seed): + """Returns the local seeds an operation should use given an op-specific seed. + + See @{tf.get_seed} for more details. This wrapper adds support for the case + where `seed` may be a tensor. + + Args: + seed: An integer or a @{tf.int64} scalar tensor. + + Returns: + A tuple of two @{tf.int64} scalar tensors that should be used for the local + seed of the calling dataset. + """ + seed, seed2 = random_seed.get_seed(seed) + if seed is None: + seed = constant_op.constant(0, dtype=dtypes.int64, name="seed") + else: + seed = ops.convert_to_tensor(seed, dtype=dtypes.int64, name="seed") + if seed2 is None: + seed2 = constant_op.constant(0, dtype=dtypes.int64, name="seed2") + else: + with ops.name_scope("seed2") as scope: + seed2 = ops.convert_to_tensor(seed2, dtype=dtypes.int64) + seed2 = array_ops.where( + math_ops.logical_and( + math_ops.equal(seed, 0), math_ops.equal(seed2, 0)), + constant_op.constant(2**31 - 1, dtype=dtypes.int64), + seed2, + name=scope) + return seed, seed2 diff --git a/tensorflow/python/data/util/random_seed_test.py b/tensorflow/python/data/util/random_seed_test.py new file mode 100644 index 0000000000000000000000000000000000000000..33227e82afe6fe1c748693d107d4e9844abb8e09 --- /dev/null +++ b/tensorflow/python/data/util/random_seed_test.py @@ -0,0 +1,83 @@ +# 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 utilities working with arbitrarily nested structures.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.data.util import random_seed as data_random_seed +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 random_seed +from tensorflow.python.framework import test_util +from tensorflow.python.platform import test + + +class RandomSeedTest(test.TestCase): + + @test_util.run_in_graph_and_eager_modes() + def testRandomSeed(self): + zero_t = constant_op.constant(0, dtype=dtypes.int64, name='zero') + one_t = constant_op.constant(1, dtype=dtypes.int64, name='one') + intmax_t = constant_op.constant( + 2**31 - 1, dtype=dtypes.int64, name='intmax') + test_cases = [ + # Each test case is a tuple with input to get_seed: + # (input_graph_seed, input_op_seed) + # and output from get_seed: + # (output_graph_seed, output_op_seed) + ((None, None), (0, 0)), + ((None, 1), (random_seed.DEFAULT_GRAPH_SEED, 1)), + ((1, 1), (1, 1)), + ((0, 0), (0, 2**31 - 1)), # Avoid nondeterministic (0, 0) output + ((2**31 - 1, 0), (0, 2**31 - 1)), # Don't wrap to (0, 0) either + ((0, 2**31 - 1), (0, 2**31 - 1)), # Wrapping for the other argument + # Once more, with tensor-valued arguments + ((None, one_t), (random_seed.DEFAULT_GRAPH_SEED, 1)), + ((1, one_t), (1, 1)), + ((0, zero_t), (0, 2**31 - 1)), # Avoid nondeterministic (0, 0) output + ((2**31 - 1, zero_t), (0, 2**31 - 1)), # Don't wrap to (0, 0) either + ((0, intmax_t), (0, 2**31 - 1)), # Wrapping for the other argument + ] + for tc in test_cases: + tinput, toutput = tc[0], tc[1] + random_seed.set_random_seed(tinput[0]) + g_seed, op_seed = data_random_seed.get_seed(tinput[1]) + g_seed = self.evaluate(g_seed) + op_seed = self.evaluate(op_seed) + msg = 'test_case = {0}, got {1}, want {2}'.format( + tinput, (g_seed, op_seed), toutput) + self.assertEqual((g_seed, op_seed), toutput, msg=msg) + random_seed.set_random_seed(None) + + if not context.executing_eagerly(): + random_seed.set_random_seed(1) + tinput = (1, None) + toutput = (1, ops.get_default_graph()._last_id) # pylint: disable=protected-access + random_seed.set_random_seed(tinput[0]) + g_seed, op_seed = data_random_seed.get_seed(tinput[1]) + g_seed = self.evaluate(g_seed) + op_seed = self.evaluate(op_seed) + msg = 'test_case = {0}, got {1}, want {2}'.format(1, (g_seed, op_seed), + toutput) + self.assertEqual((g_seed, op_seed), toutput, msg=msg) + random_seed.set_random_seed(None) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/debug/BUILD b/tensorflow/python/debug/BUILD index f0e90f67772d114142ccc218ed9f42b723a1b556..c60f6923900fb1c43e6028c844260ee2aa981ff2 100644 --- a/tensorflow/python/debug/BUILD +++ b/tensorflow/python/debug/BUILD @@ -913,6 +913,7 @@ cuda_py_test( "//tensorflow/python:util", "//tensorflow/python:variables", ], + tags = ["no_windows"], # TODO: needs investigation on Windows ) py_test( @@ -967,7 +968,6 @@ cuda_py_test( ":grpc_wrapper", ":hooks", ":session_debug_testlib", - "//third_party/py/numpy", "//tensorflow/python:client", "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", @@ -983,6 +983,29 @@ cuda_py_test( ], ) +cuda_py_test( + name = "grpc_large_data_test", + size = "medium", + srcs = ["lib/grpc_large_data_test.py"], + additional_deps = [ + ":dumping_wrapper", + ":grpc_debug_test_server", + ":grpc_wrapper", + ":session_debug_testlib", + "//third_party/py/numpy", + "//tensorflow/python:client", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + "//tensorflow/python:variables", + ], + tags = [ + "no_oss", # Test flaky due to port collisions. + "no_windows", + "oss_serial", + ], +) + # TODO(cais): Run the test in OSS, perhaps through a sh_test. cuda_py_test( name = "dist_session_debug_grpc_test", diff --git a/tensorflow/python/debug/README.md b/tensorflow/python/debug/README.md index a2273b050bb1ecd5a35938c3de57fb8562f1d26d..269bbb19bdb898d1d81d0b9c618a284a437e68b9 100644 --- a/tensorflow/python/debug/README.md +++ b/tensorflow/python/debug/README.md @@ -37,12 +37,18 @@ models: * Association of nodes and tensors in graphs with Python source lines * Profiling of models at the level of graph nodes and Python source lines. (Omitted internal-only feature) +* A [gRPC](https://grpc.io/)-based remote debugging protocol, which allows us to + build a browser-based graphical user interface (GUI) for TFDBG: the + [TensorBoard Debugger Plugin](https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/debugger/README.md). ## How to use TFDBG? * For a walkthrough of TFDBG command-line interface, see https://www.tensorflow.org/programmers_guide/debugger. +* For information on the web GUI of TFDBG (TensorBoard Debugger Plugin), see + [this README](https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/debugger/README.md). * For programmatic use of the API of TFDBG, see https://www.tensorflow.org/api_docs/python/tfdbg. + ## Related Publications * Cai, S., Breck E., Nielsen E., Salib M., Sculley D. (2016) TensorFlow Debugger: diff --git a/tensorflow/python/debug/cli/cli_shared.py b/tensorflow/python/debug/cli/cli_shared.py index a0fe6066acd1462a94e93d6091db237d01cfede3..dea019fef58015fbd7982a81319dcabe4e5f4930 100644 --- a/tensorflow/python/debug/cli/cli_shared.py +++ b/tensorflow/python/debug/cli/cli_shared.py @@ -175,7 +175,7 @@ def format_tensor(tensor, include_numeric_summary: Whether a text summary of the numeric values (if applicable) will be included. write_path: A path to save the tensor value (after any slicing) to - (optinal). `numpy.save()` is used to save the value. + (optional). `numpy.save()` is used to save the value. Returns: An instance of `debugger_cli_common.RichTextLines` representing the diff --git a/tensorflow/python/debug/cli/tensor_format.py b/tensorflow/python/debug/cli/tensor_format.py index d4aea76d652e7606939f3d8a89ff0378da0774d2..9ba84e3f2261de277361d503e9189583494a5084 100644 --- a/tensorflow/python/debug/cli/tensor_format.py +++ b/tensorflow/python/debug/cli/tensor_format.py @@ -134,7 +134,7 @@ def format_tensor(tensor, if include_metadata: lines.append(" dtype: %s" % str(tensor.dtype)) - lines.append(" shape: %s" % str(tensor.shape)) + lines.append(" shape: %s" % str(tensor.shape).replace("L", "")) if lines: lines.append("") @@ -535,7 +535,7 @@ def numeric_summary(tensor): if not isinstance(tensor, np.ndarray) or not np.size(tensor): return debugger_cli_common.RichTextLines([ "No numeric summary available due to empty tensor."]) - elif (np.issubdtype(tensor.dtype, np.float) or + elif (np.issubdtype(tensor.dtype, np.floating) or np.issubdtype(tensor.dtype, np.complex) or np.issubdtype(tensor.dtype, np.integer)): counts = [ diff --git a/tensorflow/python/debug/lib/debug_data.py b/tensorflow/python/debug/lib/debug_data.py index c4b13a1045dac4966b0e841155a2932216881d34..8d355aa27f6fa10a1889420a9087800be12a81ce 100644 --- a/tensorflow/python/debug/lib/debug_data.py +++ b/tensorflow/python/debug/lib/debug_data.py @@ -222,7 +222,7 @@ def has_inf_or_nan(datum, tensor): # Also return False for data types that cannot be represented as numpy # arrays. return False - elif (np.issubdtype(tensor.dtype, np.float) or + elif (np.issubdtype(tensor.dtype, np.floating) or np.issubdtype(tensor.dtype, np.complex) or np.issubdtype(tensor.dtype, np.integer)): return np.any(np.isnan(tensor)) or np.any(np.isinf(tensor)) diff --git a/tensorflow/python/debug/lib/debug_gradients.py b/tensorflow/python/debug/lib/debug_gradients.py index 16f51a4b32f711b97077643cec669bb8970e0b21..589a13db7f798aef3bb82dfbd442deabfbcf2a41 100644 --- a/tensorflow/python/debug/lib/debug_gradients.py +++ b/tensorflow/python/debug/lib/debug_gradients.py @@ -156,11 +156,12 @@ class GradientsDebugger(object): # TODO(cais): Implement value_stack. grad_debug_op_name = _tensor_to_grad_debug_op_name(input_tensor, self._uuid) # pylint: disable=protected-access - identity_op = (gen_array_ops._debug_gradient_ref_identity - if input_tensor.dtype._is_ref_dtype - else gen_array_ops._debug_gradient_identity) - debug_grad_identity = identity_op(input_tensor, name=grad_debug_op_name) + identity_op = ( + gen_array_ops.debug_gradient_ref_identity + if input_tensor.dtype._is_ref_dtype else + gen_array_ops.debug_gradient_identity) # pylint: enable=protected-access + debug_grad_identity = identity_op(input_tensor, name=grad_debug_op_name) assert debug_grad_identity.dtype == input_tensor.dtype if debug_grad_identity.op.name != grad_debug_op_name: raise ValueError( diff --git a/tensorflow/python/debug/lib/debug_gradients_test.py b/tensorflow/python/debug/lib/debug_gradients_test.py index c1e9869d978e4f5ddfd3cd5f1abd7f5c97b7ca88..01867fc69d0782b34edb1e8eb873b19f5dfc8529 100644 --- a/tensorflow/python/debug/lib/debug_gradients_test.py +++ b/tensorflow/python/debug/lib/debug_gradients_test.py @@ -40,6 +40,7 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): def setUp(self): rewriter_config = rewriter_config_pb2.RewriterConfig( + disable_model_pruning=True, dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF) graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_config) config = config_pb2.ConfigProto(graph_options=graph_options) @@ -117,8 +118,8 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): def testCallingIdentifyGradientTwiceWithTheSameGradientsDebuggerErrors(self): grad_debugger = debug_gradients.GradientsDebugger() grad_debugger.identify_gradient(self.w) - with self.assertRaisesRegexp( - ValueError, "The graph already contains an op named .*"): + with self.assertRaisesRegexp(ValueError, + "The graph already contains an op named .*"): grad_debugger.identify_gradient(self.w) def testIdentifyGradientWorksOnMultipleLosses(self): @@ -144,10 +145,10 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): self.assertIsNot(dz1_dy, dz2_dy) self.sess.run(variables.global_variables_initializer()) - self.assertAllClose(5.0 ** 2, self.sess.run(z1)) - self.assertAllClose(5.0 ** 0.5, self.sess.run(z2)) + self.assertAllClose(5.0**2, self.sess.run(z1)) + self.assertAllClose(5.0**0.5, self.sess.run(z2)) self.assertAllClose(2.0 * 5.0, self.sess.run(dz1_dy)) - self.assertAllClose(0.5 * (5.0 ** -0.5), self.sess.run(dz2_dy)) + self.assertAllClose(0.5 * (5.0**-0.5), self.sess.run(dz2_dy)) def testIdentifyGradientRaisesLookupErrorForUnknownXTensor(self): grad_debugger_1 = debug_gradients.GradientsDebugger() @@ -259,8 +260,8 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): self.sess.run(variables.global_variables_initializer()) self.assertAllClose(3.0, self.sess.run(u_grad)) self.assertAllClose(2.0, self.sess.run(v_grad)) - self.assertAllClose( - 3.0, self.sess.run(grad_debugger.gradient_tensor("u:0"))) + self.assertAllClose(3.0, self.sess.run( + grad_debugger.gradient_tensor("u:0"))) def testWatchGradientsWorksOnMultipleTensors(self): y = math_ops.add(self.w, -1.0, name="y") @@ -277,10 +278,10 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): self.assertIsInstance(grad_debugger.gradient_tensor("w:0"), ops.Tensor) self.sess.run(variables.global_variables_initializer()) - self.assertAllClose( - 1.0, self.sess.run(grad_debugger.gradient_tensor("w:0"))) - self.assertAllClose( - 3.0, self.sess.run(grad_debugger.gradient_tensor("u:0"))) + self.assertAllClose(1.0, self.sess.run( + grad_debugger.gradient_tensor("w:0"))) + self.assertAllClose(3.0, self.sess.run( + grad_debugger.gradient_tensor("u:0"))) def testWatchGradientsByXTensorsWorks(self): y = math_ops.add(self.w, -1.0, name="foo/y") @@ -290,8 +291,8 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): # But we can still get the gradient tensors by using # watch_gradients_by_x_tensors(). grad_debugger = debug_gradients.GradientsDebugger() - with grad_debugger.watch_gradients_by_tensors( - self.sess.graph, [self.w, self.u, y]): + with grad_debugger.watch_gradients_by_tensors(self.sess.graph, + [self.w, self.u, y]): gradient_descent.GradientDescentOptimizer(0.1).minimize(z) self.assertEqual(3, len(grad_debugger.gradient_tensors())) @@ -324,18 +325,18 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): self.assertIsNot(dz1_dy, dz2_dy) self.sess.run(variables.global_variables_initializer()) - self.assertAllClose(5.0 ** 2, self.sess.run(z1)) - self.assertAllClose(5.0 ** 0.5, self.sess.run(z2)) + self.assertAllClose(5.0**2, self.sess.run(z1)) + self.assertAllClose(5.0**0.5, self.sess.run(z2)) self.assertAllClose(2.0 * 5.0, self.sess.run(dz1_dy)) - self.assertAllClose(0.5 * (5.0 ** -0.5), self.sess.run(dz2_dy)) + self.assertAllClose(0.5 * (5.0**-0.5), self.sess.run(dz2_dy)) def testGradientsValuesFromDumpWorks(self): y = math_ops.add(self.w, -1.0, name="y") z = math_ops.square(y, name="z") grad_debugger = debug_gradients.GradientsDebugger() - with grad_debugger.watch_gradients_by_tensors( - self.sess.graph, [self.w, self.u, y]): + with grad_debugger.watch_gradients_by_tensors(self.sess.graph, + [self.w, self.u, y]): train_op = gradient_descent.GradientDescentOptimizer(0.1).minimize(z) self.sess.run(variables.global_variables_initializer()) @@ -343,10 +344,7 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase): run_options = config_pb2.RunOptions(output_partition_graphs=True) dump_dir = tempfile.mkdtemp() debug_url = "file://" + dump_dir - debug_utils.watch_graph( - run_options, - self.sess.graph, - debug_urls=debug_url) + debug_utils.watch_graph(run_options, self.sess.graph, debug_urls=debug_url) run_metadata = config_pb2.RunMetadata() self.assertAllClose(2.0, self.sess.run(self.u)) self.sess.run(train_op, options=run_options, run_metadata=run_metadata) diff --git a/tensorflow/python/debug/lib/grpc_large_data_test.py b/tensorflow/python/debug/lib/grpc_large_data_test.py new file mode 100644 index 0000000000000000000000000000000000000000..5bc477a9baeb7116530fc9122b926458c1a6c08e --- /dev/null +++ b/tensorflow/python/debug/lib/grpc_large_data_test.py @@ -0,0 +1,210 @@ +# 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 sending large-size data through tfdbg grpc channels. + +"Large-size data" includes large GraphDef protos and large Tensor protos. +""" +from __future__ import absolute_import +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.debug.lib import grpc_debug_test_server +from tensorflow.python.debug.lib import session_debug_testlib +from tensorflow.python.debug.wrappers import framework +from tensorflow.python.debug.wrappers import grpc_wrapper +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.ops import array_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import googletest +from tensorflow.python.platform import test +from tensorflow.python.platform import tf_logging + + +class LargeGraphAndLargeTensorsDebugTest(test_util.TensorFlowTestCase): + + @classmethod + def setUpClass(cls): + (cls.debug_server_port, cls.debug_server_url, _, cls.debug_server_thread, + cls.debug_server + ) = grpc_debug_test_server.start_server_on_separate_thread( + dump_to_filesystem=False) + tf_logging.info("debug server url: %s", cls.debug_server_url) + + @classmethod + def tearDownClass(cls): + cls.debug_server.stop_server().wait() + cls.debug_server_thread.join() + + def tearDown(self): + ops.reset_default_graph() + self.debug_server.clear_data() + + def testSendingLargeGraphDefsWorks(self): + with self.test_session( + use_gpu=True, + config=session_debug_testlib.no_rewrite_session_config()) as sess: + u = variables.Variable(42.0, name="original_u") + for _ in xrange(50 * 1000): + u = array_ops.identity(u) + sess.run(variables.global_variables_initializer()) + + def watch_fn(fetches, feeds): + del fetches, feeds + return framework.WatchOptions( + debug_ops=["DebugIdentity"], + node_name_regex_whitelist=r"original_u") + sess = grpc_wrapper.GrpcDebugWrapperSession( + sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) + self.assertAllClose(42.0, sess.run(u)) + + self.assertAllClose( + [42.0], + self.debug_server.debug_tensor_values["original_u:0:DebugIdentity"]) + self.assertEqual(2 if test.is_gpu_available() else 1, + len(self.debug_server.partition_graph_defs)) + max_graph_def_size = max([ + len(graph_def.SerializeToString()) + for graph_def in self.debug_server.partition_graph_defs]) + self.assertGreater(max_graph_def_size, 4 * 1024 * 1024) + + def testSendingLargeFloatTensorWorks(self): + with self.test_session( + use_gpu=True, + config=session_debug_testlib.no_rewrite_session_config()) as sess: + u_init_val_array = list(xrange(1200 * 1024)) + # Size: 4 * 1200 * 1024 = 4800k > 4M + + u_init = constant_op.constant( + u_init_val_array, dtype=dtypes.float32, name="u_init") + u = variables.Variable(u_init, name="u") + + def watch_fn(fetches, feeds): + del fetches, feeds # Unused by this watch_fn. + return framework.WatchOptions( + debug_ops=["DebugIdentity"], + node_name_regex_whitelist=r"u_init") + sess = grpc_wrapper.GrpcDebugWrapperSession( + sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) + sess.run(u.initializer) + + self.assertAllEqual( + u_init_val_array, + self.debug_server.debug_tensor_values["u_init:0:DebugIdentity"][0]) + + def testSendingStringTensorWithAlmostTooLargeStringsWorks(self): + with self.test_session( + use_gpu=True, + config=session_debug_testlib.no_rewrite_session_config()) as sess: + u_init_val = [ + b"", b"spam", b"A" * 2500 * 1024, b"B" * 2500 * 1024, b"egg", b""] + u_init = constant_op.constant( + u_init_val, dtype=dtypes.string, name="u_init") + u = variables.Variable(u_init, name="u") + + def watch_fn(fetches, feeds): + del fetches, feeds + return framework.WatchOptions( + debug_ops=["DebugIdentity"], + node_name_regex_whitelist=r"u_init") + sess = grpc_wrapper.GrpcDebugWrapperSession( + sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) + sess.run(u.initializer) + + self.assertAllEqual( + u_init_val, + self.debug_server.debug_tensor_values["u_init:0:DebugIdentity"][0]) + + def testSendingLargeStringTensorWorks(self): + with self.test_session( + use_gpu=True, + config=session_debug_testlib.no_rewrite_session_config()) as sess: + strs_total_size_threshold = 5000 * 1024 + cum_size = 0 + u_init_val_array = [] + while cum_size < strs_total_size_threshold: + strlen = np.random.randint(200) + u_init_val_array.append(b"A" * strlen) + cum_size += strlen + + u_init = constant_op.constant( + u_init_val_array, dtype=dtypes.string, name="u_init") + u = variables.Variable(u_init, name="u") + + def watch_fn(fetches, feeds): + del fetches, feeds + return framework.WatchOptions( + debug_ops=["DebugIdentity"], + node_name_regex_whitelist=r"u_init") + sess = grpc_wrapper.GrpcDebugWrapperSession( + sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) + sess.run(u.initializer) + + self.assertAllEqual( + u_init_val_array, + self.debug_server.debug_tensor_values["u_init:0:DebugIdentity"][0]) + + def testSendingEmptyFloatTensorWorks(self): + with self.test_session( + use_gpu=True, + config=session_debug_testlib.no_rewrite_session_config()) as sess: + u_init = constant_op.constant( + [], dtype=dtypes.float32, shape=[0], name="u_init") + u = variables.Variable(u_init, name="u") + + def watch_fn(fetches, feeds): + del fetches, feeds + return framework.WatchOptions( + debug_ops=["DebugIdentity"], + node_name_regex_whitelist=r"u_init") + sess = grpc_wrapper.GrpcDebugWrapperSession( + sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) + sess.run(u.initializer) + + u_init_value = self.debug_server.debug_tensor_values[ + "u_init:0:DebugIdentity"][0] + self.assertEqual(np.float32, u_init_value.dtype) + self.assertEqual(0, len(u_init_value)) + + def testSendingEmptyStringTensorWorks(self): + with self.test_session( + use_gpu=True, + config=session_debug_testlib.no_rewrite_session_config()) as sess: + u_init = constant_op.constant( + [], dtype=dtypes.string, shape=[0], name="u_init") + u = variables.Variable(u_init, name="u") + + def watch_fn(fetches, feeds): + del fetches, feeds + return framework.WatchOptions( + debug_ops=["DebugIdentity"], + node_name_regex_whitelist=r"u_init") + sess = grpc_wrapper.GrpcDebugWrapperSession( + sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) + sess.run(u.initializer) + + u_init_value = self.debug_server.debug_tensor_values[ + "u_init:0:DebugIdentity"][0] + self.assertEqual(np.object, u_init_value.dtype) + self.assertEqual(0, len(u_init_value)) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/python/debug/lib/session_debug_file_test.py b/tensorflow/python/debug/lib/session_debug_file_test.py index 1a6bedbbcbf94eb95e49d43e2d03c85b53bebb7b..ba0f15b4e2ff23295eae764088144a3d1b533f01 100644 --- a/tensorflow/python/debug/lib/session_debug_file_test.py +++ b/tensorflow/python/debug/lib/session_debug_file_test.py @@ -22,7 +22,6 @@ import shutil import tempfile from tensorflow.core.protobuf import config_pb2 -from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.client import session from tensorflow.python.debug.lib import debug_data from tensorflow.python.debug.lib import debug_utils @@ -36,13 +35,6 @@ from tensorflow.python.platform import googletest class SessionDebugFileTest(session_debug_testlib.SessionDebugTestBase): - def _no_rewrite_session_config(self): - rewriter_config = rewriter_config_pb2.RewriterConfig( - disable_model_pruning=True, - arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF) - graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_config) - return config_pb2.ConfigProto(graph_options=graph_options) - def _debug_urls(self, run_number=None): return ["file://%s" % self._debug_dump_dir(run_number=run_number)] @@ -55,7 +47,8 @@ class SessionDebugFileTest(session_debug_testlib.SessionDebugTestBase): def testAllowsDifferentWatchesOnDifferentRuns(self): """Test watching different tensors on different runs of the same graph.""" - with session.Session(config=self._no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: u_init_val = [[5.0, 3.0], [-1.0, 0.0]] v_init_val = [[2.0], [-1.0]] diff --git a/tensorflow/python/debug/lib/session_debug_grpc_test.py b/tensorflow/python/debug/lib/session_debug_grpc_test.py index b623ee31c5dc59894373ec7952e53acd0f6e1126..ff49b6954776264ccb2eceeceab7da5a881081f0 100644 --- a/tensorflow/python/debug/lib/session_debug_grpc_test.py +++ b/tensorflow/python/debug/lib/session_debug_grpc_test.py @@ -24,11 +24,9 @@ from __future__ import print_function import os import shutil -import numpy as np from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.core.protobuf import config_pb2 -from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.client import session from tensorflow.python.debug.lib import debug_data from tensorflow.python.debug.lib import debug_utils @@ -38,28 +36,15 @@ from tensorflow.python.debug.wrappers import framework from tensorflow.python.debug.wrappers import grpc_wrapper from tensorflow.python.debug.wrappers import hooks 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.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variables from tensorflow.python.platform import googletest -from tensorflow.python.platform import test -from tensorflow.python.platform import tf_logging from tensorflow.python.training import monitored_session -def no_rewrite_session_config(): - rewriter_config = rewriter_config_pb2.RewriterConfig( - disable_model_pruning=True, - arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF, - dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF) - graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_config) - return config_pb2.ConfigProto(graph_options=graph_options) - - class GrpcDebugServerTest(test_util.TensorFlowTestCase): def testRepeatedRunServerRaisesException(self): @@ -142,19 +127,22 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): return os.path.join(self._dump_root, "run_%d" % run_number) def testConstructGrpcDebugWrapperSessionWithInvalidTypeRaisesException(self): - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) with self.assertRaisesRegexp( TypeError, "Expected type str or list in grpc_debug_server_addresses"): grpc_wrapper.GrpcDebugWrapperSession(sess, 1337) def testConstructGrpcDebugWrapperSessionWithInvalidTypeRaisesException2(self): - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) with self.assertRaisesRegexp( TypeError, "Expected type str in list grpc_debug_server_addresses"): grpc_wrapper.GrpcDebugWrapperSession(sess, ["localhost:1337", 1338]) def testUseInvalidWatchFnTypeWithGrpcDebugWrapperSessionRaisesException(self): - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) with self.assertRaises(TypeError): grpc_wrapper.GrpcDebugWrapperSession( sess, "localhost:%d" % self._server_port, watch_fn="foo") @@ -164,7 +152,8 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): v = variables.Variable(20.0, name="v") w = math_ops.multiply(u, v, name="w") - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) sess.run(u.initializer) sess.run(v.initializer) @@ -190,7 +179,8 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): v = variables.Variable(20.0, name="v") w = math_ops.multiply(u, v, name="w") - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) sess.run(u.initializer) sess.run(v.initializer) @@ -223,7 +213,8 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): v = variables.Variable(20.0, name="v") w = math_ops.multiply(u, v, name="w") - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) sess.run(u.initializer) sess.run(v.initializer) @@ -254,7 +245,8 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): v = variables.Variable(20.0, name="v") w = math_ops.multiply(u, v, name="w") - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) sess.run(u.initializer) sess.run(v.initializer) @@ -298,7 +290,8 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): v = variables.Variable(20.0, name="v") w = math_ops.multiply(u, v, name="w") - sess = session.Session(config=no_rewrite_session_config()) + sess = session.Session( + config=session_debug_testlib.no_rewrite_session_config()) sess.run(variables.global_variables_initializer()) grpc_debug_hook = hooks.TensorBoardDebugHook( @@ -324,168 +317,6 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase): hooks.GrpcDebugHook(["foo:42424"]) -class LargeGraphAndLargeTensorsDebugTest(test_util.TensorFlowTestCase): - - @classmethod - def setUpClass(cls): - (cls.debug_server_port, cls.debug_server_url, _, cls.debug_server_thread, - cls.debug_server - ) = grpc_debug_test_server.start_server_on_separate_thread( - dump_to_filesystem=False) - tf_logging.info("debug server url: %s", cls.debug_server_url) - - @classmethod - def tearDownClass(cls): - cls.debug_server.stop_server().wait() - cls.debug_server_thread.join() - - def tearDown(self): - ops.reset_default_graph() - self.debug_server.clear_data() - - def testSendingLargeGraphDefsWorks(self): - with self.test_session( - use_gpu=True, config=no_rewrite_session_config()) as sess: - u = variables.Variable(42.0, name="original_u") - for _ in xrange(50 * 1000): - u = array_ops.identity(u) - sess.run(variables.global_variables_initializer()) - - def watch_fn(fetches, feeds): - del fetches, feeds - return framework.WatchOptions( - debug_ops=["DebugIdentity"], - node_name_regex_whitelist=r"original_u") - sess = grpc_wrapper.GrpcDebugWrapperSession( - sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) - self.assertAllClose(42.0, sess.run(u)) - - self.assertAllClose( - [42.0], - self.debug_server.debug_tensor_values["original_u:0:DebugIdentity"]) - self.assertEqual(2 if test.is_gpu_available() else 1, - len(self.debug_server.partition_graph_defs)) - max_graph_def_size = max([ - len(graph_def.SerializeToString()) - for graph_def in self.debug_server.partition_graph_defs]) - self.assertGreater(max_graph_def_size, 4 * 1024 * 1024) - - def testSendingLargeFloatTensorWorks(self): - with self.test_session( - use_gpu=True, config=no_rewrite_session_config()) as sess: - u_init_val_array = list(xrange(1200 * 1024)) - # Size: 4 * 1200 * 1024 = 4800k > 4M - - u_init = constant_op.constant( - u_init_val_array, dtype=dtypes.float32, name="u_init") - u = variables.Variable(u_init, name="u") - - def watch_fn(fetches, feeds): - del fetches, feeds # Unused by this watch_fn. - return framework.WatchOptions( - debug_ops=["DebugIdentity"], - node_name_regex_whitelist=r"u_init") - sess = grpc_wrapper.GrpcDebugWrapperSession( - sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) - sess.run(u.initializer) - - self.assertAllEqual( - u_init_val_array, - self.debug_server.debug_tensor_values["u_init:0:DebugIdentity"][0]) - - def testSendingStringTensorWithAlmostTooLargeStringsWorks(self): - with self.test_session( - use_gpu=True, config=no_rewrite_session_config()) as sess: - u_init_val = [ - b"", b"spam", b"A" * 2500 * 1024, b"B" * 2500 * 1024, b"egg", b""] - u_init = constant_op.constant( - u_init_val, dtype=dtypes.string, name="u_init") - u = variables.Variable(u_init, name="u") - - def watch_fn(fetches, feeds): - del fetches, feeds - return framework.WatchOptions( - debug_ops=["DebugIdentity"], - node_name_regex_whitelist=r"u_init") - sess = grpc_wrapper.GrpcDebugWrapperSession( - sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) - sess.run(u.initializer) - - self.assertAllEqual( - u_init_val, - self.debug_server.debug_tensor_values["u_init:0:DebugIdentity"][0]) - - def testSendingLargeStringTensorWorks(self): - with self.test_session( - use_gpu=True, config=no_rewrite_session_config()) as sess: - strs_total_size_threshold = 5000 * 1024 - cum_size = 0 - u_init_val_array = [] - while cum_size < strs_total_size_threshold: - strlen = np.random.randint(200) - u_init_val_array.append(b"A" * strlen) - cum_size += strlen - - u_init = constant_op.constant( - u_init_val_array, dtype=dtypes.string, name="u_init") - u = variables.Variable(u_init, name="u") - - def watch_fn(fetches, feeds): - del fetches, feeds - return framework.WatchOptions( - debug_ops=["DebugIdentity"], - node_name_regex_whitelist=r"u_init") - sess = grpc_wrapper.GrpcDebugWrapperSession( - sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) - sess.run(u.initializer) - - self.assertAllEqual( - u_init_val_array, - self.debug_server.debug_tensor_values["u_init:0:DebugIdentity"][0]) - - def testSendingEmptyFloatTensorWorks(self): - with self.test_session( - use_gpu=True, config=no_rewrite_session_config()) as sess: - u_init = constant_op.constant( - [], dtype=dtypes.float32, shape=[0], name="u_init") - u = variables.Variable(u_init, name="u") - - def watch_fn(fetches, feeds): - del fetches, feeds - return framework.WatchOptions( - debug_ops=["DebugIdentity"], - node_name_regex_whitelist=r"u_init") - sess = grpc_wrapper.GrpcDebugWrapperSession( - sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) - sess.run(u.initializer) - - u_init_value = self.debug_server.debug_tensor_values[ - "u_init:0:DebugIdentity"][0] - self.assertEqual(np.float32, u_init_value.dtype) - self.assertEqual(0, len(u_init_value)) - - def testSendingEmptyStringTensorWorks(self): - with self.test_session( - use_gpu=True, config=no_rewrite_session_config()) as sess: - u_init = constant_op.constant( - [], dtype=dtypes.string, shape=[0], name="u_init") - u = variables.Variable(u_init, name="u") - - def watch_fn(fetches, feeds): - del fetches, feeds - return framework.WatchOptions( - debug_ops=["DebugIdentity"], - node_name_regex_whitelist=r"u_init") - sess = grpc_wrapper.GrpcDebugWrapperSession( - sess, "localhost:%d" % self.debug_server_port, watch_fn=watch_fn) - sess.run(u.initializer) - - u_init_value = self.debug_server.debug_tensor_values[ - "u_init:0:DebugIdentity"][0] - self.assertEqual(np.object, u_init_value.dtype) - self.assertEqual(0, len(u_init_value)) - - class SessionDebugConcurrentTest( session_debug_testlib.DebugConcurrentRunCallsTest): @@ -548,7 +379,8 @@ class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase): self._server_2.clear_data() def testToggleEnableTwoDebugWatchesNoCrosstalkBetweenDebugNodes(self): - with session.Session(config=no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: v_1 = variables.Variable(50.0, name="v_1") v_2 = variables.Variable(-50.0, name="v_1") delta_1 = constant_op.constant(5.0, name="delta_1") @@ -617,7 +449,8 @@ class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase): ("toggled_2", 0, "DebugIdentity")]) self._servers_and_threads.append((server, server_thread)) - with session.Session(config=no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: v_1 = variables.Variable(50.0, name="v_1") v_2 = variables.Variable(-50.0, name="v_1") # These two nodes have names that match those in the @@ -656,7 +489,8 @@ class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase): self.assertEqual(0, len(server.debug_tensor_values)) def testToggleEnableTwoDebugWatchesNoCrosstalkBetweenServers(self): - with session.Session(config=no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: v = variables.Variable(50.0, name="v") delta = constant_op.constant(5.0, name="delta") inc_v = state_ops.assign_add(v, delta, name="inc_v") @@ -698,7 +532,8 @@ class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase): self.assertEqual(0, len(self._server_2.debug_tensor_values)) def testToggleBreakpointsWorks(self): - with session.Session(config=no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: v_1 = variables.Variable(50.0, name="v_1") v_2 = variables.Variable(-50.0, name="v_2") delta_1 = constant_op.constant(5.0, name="delta_1") @@ -755,7 +590,8 @@ class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase): self.assertSetEqual(set(), self._server_1.breakpoints) def testTensorBoardDebuggerWrapperToggleBreakpointsWorks(self): - with session.Session(config=no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: v_1 = variables.Variable(50.0, name="v_1") v_2 = variables.Variable(-50.0, name="v_2") delta_1 = constant_op.constant(5.0, name="delta_1") @@ -827,7 +663,8 @@ class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase): self._server_1.query_source_file_line(__file__, 1) def testTensorBoardDebuggerWrapperDisablingTracebackSourceSendingWorks(self): - with session.Session(config=no_rewrite_session_config()) as sess: + with session.Session( + config=session_debug_testlib.no_rewrite_session_config()) as sess: v_1 = variables.Variable(50.0, name="v_1") v_2 = variables.Variable(-50.0, name="v_2") delta_1 = constant_op.constant(5.0, name="delta_1") diff --git a/tensorflow/python/debug/lib/session_debug_testlib.py b/tensorflow/python/debug/lib/session_debug_testlib.py index 20a40018bf9c67c5b743963489c8fc5616efa2db..f4fac1401918ccacd38aae5ad2ef8d686c9204b9 100644 --- a/tensorflow/python/debug/lib/session_debug_testlib.py +++ b/tensorflow/python/debug/lib/session_debug_testlib.py @@ -988,7 +988,7 @@ class SessionDebugTestBase(test_util.TensorFlowTestCase): def testWatchingVariableUpdateOpsSeesUpdatedValues(self): """Watch output slots on Variable-updating ops, with no emitted edges.""" - with session.Session() as sess: + with session.Session(config=no_rewrite_session_config()) as sess: u_init = constant_op.constant(10.0) u = variables.Variable(u_init, name="gdo/u") v_init = constant_op.constant(20.0) diff --git a/tensorflow/python/debug/wrappers/grpc_wrapper.py b/tensorflow/python/debug/wrappers/grpc_wrapper.py index 74d7c2b9e242f947a33c0bdb6508847808d69c0b..fb9494f57636e46e54ef230cf4803dbb6ccad0c7 100644 --- a/tensorflow/python/debug/wrappers/grpc_wrapper.py +++ b/tensorflow/python/debug/wrappers/grpc_wrapper.py @@ -17,6 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import signal +import sys import traceback # Google-internal import(s). @@ -137,6 +139,29 @@ class GrpcDebugWrapperSession(framework.NonInteractiveDebugWrapperSession): if not address.startswith(common.GRPC_URL_PREFIX) else address) +def _signal_handler(unused_signal, unused_frame): + try: + input_func = raw_input + except NameError: + # Python 3 does not have raw_input. + input_func = input + + while True: + response = input_func("\nSIGINT received. Quit program? (Y/n): ").strip() + if response in ("", "Y", "y"): + sys.exit(0) + elif response in ("N", "n"): + break + + +def register_signal_handler(): + try: + signal.signal(signal.SIGINT, _signal_handler) + except ValueError: + # This can happen if we are not in the MainThread. + pass + + class TensorBoardDebugWrapperSession(GrpcDebugWrapperSession): """A tfdbg Session wrapper that can be used with TensorBoard Debugger Plugin. @@ -185,6 +210,8 @@ class TensorBoardDebugWrapperSession(GrpcDebugWrapperSession): # sent to the debug servers. self._sent_graph_version = -1 + register_signal_handler() + def run(self, fetches, feed_dict=None, diff --git a/tensorflow/python/debug/wrappers/hooks.py b/tensorflow/python/debug/wrappers/hooks.py index 0204254ccab109f4844f077df78902872d1156d5..6705cd31e291d2eab7aa8179e9b2b829f8970c18 100644 --- a/tensorflow/python/debug/wrappers/hooks.py +++ b/tensorflow/python/debug/wrappers/hooks.py @@ -345,6 +345,7 @@ class TensorBoardDebugHook(GrpcDebugHook): self._grpc_debug_server_addresses = grpc_debug_server_addresses self._send_traceback_and_source_code = send_traceback_and_source_code self._sent_graph_version = -1 + grpc_wrapper.register_signal_handler() def before_run(self, run_context): if self._send_traceback_and_source_code: diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD index 9e3382d4f301529cd2b476bc76efe7dfd2be9298..5bedf9c6fdfcf630a72ef8a34bca38417b738d4f 100644 --- a/tensorflow/python/eager/BUILD +++ b/tensorflow/python/eager/BUILD @@ -42,7 +42,6 @@ py_library( ":backprop", ":context", ":core", - ":custom_gradient", ":execute", ":function", ":graph_callable", @@ -103,7 +102,6 @@ cuda_py_test( additional_deps = [ ":backprop", ":context", - ":custom_gradient", ":test", "//tensorflow/python:embedding_ops", "//tensorflow/python:array_ops", @@ -206,44 +204,6 @@ cc_library( ], ) -cc_library( - name = "python_eager_op_gen_main", - srcs = [ - "python_eager_op_gen_main.cc", - ], - visibility = ["//visibility:public"], - deps = [ - ":python_eager_op_gen", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:op_gen_lib", - "//tensorflow/core:protos_all_cc", - ], -) - -tf_cc_binary( - name = "python_eager_op_gen_demo", - deps = [ - ":python_eager_op_gen_main", - "//tensorflow/core:ops", - ], -) - -py_library( - name = "custom_gradient", - srcs = ["custom_gradient.py"], - srcs_version = "PY2AND3", - visibility = ["//tensorflow:internal"], - deps = [ - ":context", - ":tape", - "//tensorflow/python:array_ops", - "//tensorflow/python:framework_ops", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:util", - ], -) - py_library( name = "graph_only_ops", srcs = ["graph_only_ops.py"], @@ -387,7 +347,6 @@ py_test( deps = [ ":backprop", ":context", - ":custom_gradient", ":test", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py index d79d1fc0a6400a894293f3254d5cac5a10661e13..88de1a951ffd85b4ca97a084c8616d804dc10b45 100644 --- a/tensorflow/python/eager/backprop.py +++ b/tensorflow/python/eager/backprop.py @@ -116,110 +116,6 @@ _gradient_functions_lock = threading.Lock() _tracing = False -# TODO(apassos) replace this with a mechanism which can happen at the op -# gradient function registration site, to be less error-prone -# TODO(apassos) add ops other than those in nn_grad and math_grad -_ops_which_dont_need_outputs = set([ - "Identity", - "MatMul", - "Conv2DBackpropInput", - "Conv2DBackpropFilter", - "Conv3D", - "Conv3DBackpropInputV2", - "AvgPool3D", - "AvgPool3DGrad", - "MaxPool3D", - "MaxPool3DGrad", - "MaxPool3DGradGrad", - "BiasAdd", - "BiasAddV1", - "BiasAddGrad", - "Relu6", - "Softplus", - "SoftplusGrad", - "Softsign", - "ReluGrad", - "Conv2D", - "DepthwiseConv2dNative", - "Dilation2D", - "AvgPool", - "AvgPoolGrad", - "BatchNormWithGlobalNormalization", - "L2Loss", - "Sum", - "Prod", - "SegmentSum", - "SegmentMean", - "SparseSegmentSum", - "SparseSegmentMean", - "SparseSegmentSqrtN", - "SegmentMin", - "SegmentMax", - "UnsortedSegmentSum", - "UnsortedSegmentMax", - "Abs", - "Neg", - "ReciprocalGrad", - "Square", - "Expm1", - "Log", - "Log1p", - "TanhGrad", - "SigmoidGrad", - "Sign", - "Sin", - "Cos", - "Tan", - "Add", - "Sub", - "Mul", - "Div", - "RealDiv", - "Maximum", - "Minimum", - "SquaredDifference", - "Select", - "SparseMatMul", - "BatchMatMul", - "Complex", - "Real", - "Imag", - "Angle", - "Conj", - "Cast", - "Cross", - "Cumsum", - "Cumprod", - "ReadVariableOp", - "VarHandleOp", - "Shape", -]) - -_ops_which_dont_need_inputs = set([ - "Identity", - "Softmax", - "LogSoftmax", - "BiasAdd", - "Relu", - "Elu", - "Selu", - "SparseSoftmaxCrossEntropyWithLogits", - "Neg", - "Inv", - "Reciprocal", - "Sqrt", - "Exp", - "Tanh", - "Sigmoid", - "Real", - "Imag", - "Conj", - "ReadVariableOp", - "VarHandleOp", - "Shape", -]) - - # TODO(agarwal): use an automatic mechanism for handling None arguments to # gradient functions. # Some gradient functions can accept None arguments for gradients. The following @@ -238,57 +134,25 @@ _grad_fn_accepts_none_for_indices = { } -def _record_gradient(op_name, inputs, attrs, results, name): - """Records gradients for a TensorFlow operation. - - Args: - op_name: Name of the TensorFlow operation (see REGISTER_OP in C++ code) to - execute. - inputs: A flat list of Tensor object inputs to the operation. - attrs: A tuple with alternating string attr names and attr values for this - operation. - results: The results of the operation (as a flat list). - name: Customized name for the operation. - - Returns: - A list of maybe-wrapped results. Either Tensors or TensorNodes. - - Raises: - An exception on error. - """ - if not tape.could_possibly_record(): - return - - if op_name in _ops_which_dont_need_outputs: - op_outputs = None - else: - # TODO(apassos) this line creates a weak circular reference where the - # backprop function keeps an output alive which in turn keeps the tape entry - # alive which keeps the backprop function alive. Figure out how to break - # this up without breaking second derivatives of ops like Exp whose - # gradients depend only on the outputs. - op_outputs = results - - if op_name in _ops_which_dont_need_inputs: - op_inputs = None - else: - op_inputs = inputs - - num_inputs = len(inputs) +def _get_backward_fn(op_name, attrs, num_inputs, op_inputs, op_outputs): def grad_fn(*orig_outputs): - """Generated gradient function.""" result = _magic_gradient_function(op_name, attrs, num_inputs, op_inputs, op_outputs, orig_outputs) if _tracing: - print("Gradient for", (name if name else op_name), "inputs", op_inputs, - "output_grads", orig_outputs, "gradients", result) + print("Gradient for", op_name, "inputs", op_inputs, "output_grads", + orig_outputs, "gradients", result) return nest.flatten(result) - tape.record_operation(op_name, results, inputs, grad_fn) - if _tracing: - print("Computed op", (name if name else op_name), "inputs", inputs, - "outputs", results) + return grad_fn + + +pywrap_tensorflow.TFE_Py_RegisterBackwardFunctionGetter(_get_backward_fn) + + +def _record_gradient(op_name, inputs, attrs, results, name): + return pywrap_tensorflow.TFE_Py_RecordGradient(op_name, inputs, attrs, + results, name) execute.record_gradient = _record_gradient @@ -355,6 +219,7 @@ def implicit_val_and_grad(f): tape.pop_tape(this_tape) # Sorting variables by id, which is monotonically increasing in construction # order. This ensures unique order across executions. + # TODO(josh11b): Move the sort to the C++ implementation in pywrap_tfe_src.cc. variables = list(sorted(this_tape.watched_variables(), key=lambda v: v.handle._id)) # pylint: disable=protected-access sources = [x.handle for x in variables] @@ -616,7 +481,7 @@ def val_and_grad_function(f, params=None): return decorated -def make_vjp(f, params=None): +def make_vjp(f, params=None, persistent=True): """Returns a function that computes f and is vjp w.r.t. params. The term "vjp" here is an abbreviation for vector-jacobian product. @@ -625,6 +490,8 @@ def make_vjp(f, params=None): f: the function to be differentiated. params: the parameters (numbers or names) to differentiate with respect to. A value of None will differentiate with respect to all parameters. + persistent: Boolean controlling whether the VJP function can be re-used. + Must be True or False. Returns: A function, which when called, returns a tuple (value, vjp), where: @@ -652,7 +519,7 @@ def make_vjp(f, params=None): """Computes the value and gradient of the decorated function.""" parameter_positions = _get_arg_spec(f, params, args) assert not kwds, "The gradient function can't take keyword arguments." - this_tape = tape.push_new_tape() + this_tape = tape.push_new_tape(persistent=persistent) try: sources = [] args = [ @@ -734,8 +601,7 @@ def _num_elements(grad): raise ValueError("`grad` not a Tensor or IndexedSlices.") -_last_zero_shape_dtype = [None, None] -_last_zero = [None] +_zeros_cache = context._TensorCache() # pylint: disable=protected-access def _fast_fill(value, shape, dtype): @@ -744,14 +610,17 @@ def _fast_fill(value, shape, dtype): def _zeros(shape, dtype): """Wraps array_ops.zeros to cache last zero for a given shape and dtype.""" + device = context.context().device_name if dtype == dtypes.variant: # TODO(apassos): need to save enough information about variant tensors to do # a zeros return None - if [shape, dtype] != _last_zero_shape_dtype: - _last_zero_shape_dtype[:] = [shape, dtype] - _last_zero[0] = _fast_fill(0, shape, dtype) - return _last_zero[0] + cache_key = shape, dtype, device + cached = _zeros_cache.get(cache_key) + if cached is None: + cached = _fast_fill(0, shape, dtype) + _zeros_cache.put(cache_key, cached) + return cached def _ones(shape, dtype): @@ -858,13 +727,22 @@ class GradientTape(object): t = t.handle tape.watch(t) - def gradient(self, target, sources): + def watched_variables(self): + # Sorting variables by id, which is monotonically increasing in construction + # order. This ensures unique order across executions. + # TODO(josh11b): Move the sort to the C++ implementation in pywrap_tfe_src.cc. + return list(sorted(self._tape.watched_variables(), + key=lambda v: v.handle._id)) # pylint: disable=protected-access + + def gradient(self, target, sources, output_gradients=None): """Computes the gradient using information traced by the tape. Args: target: the tensor to be differentiated. sources: a list of Tensors or Variables, the target will be differentiated with respect to the sources. + output_gradients: a list of gradients, one for each element of + target. Defaults to None. Returns: a list of Tensors (or IndexedSlices, or None), one for each element in @@ -882,7 +760,8 @@ class GradientTape(object): else x for x in sources] grad = imperative_grad.imperative_grad( - _default_vspace, self._tape, [target], sources) + _default_vspace, self._tape, [target], sources, + output_gradients=output_gradients) if not self._persistent: self._tape = None return grad diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py index a12113893ab3eac671e8138472bc95e9d8b89499..5934293dfc59ef11b0f253d548ccfa2eff53fe38 100644 --- a/tensorflow/python/eager/backprop_test.py +++ b/tensorflow/python/eager/backprop_test.py @@ -23,7 +23,6 @@ import numpy as np from tensorflow.python import pywrap_tensorflow from tensorflow.python.eager import backprop from tensorflow.python.eager import context -from tensorflow.python.eager import custom_gradient from tensorflow.python.eager import tape from tensorflow.python.eager import test from tensorflow.python.framework import constant_op @@ -32,6 +31,7 @@ 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 custom_gradient from tensorflow.python.ops import embedding_ops from tensorflow.python.ops import gradients from tensorflow.python.ops import math_ops @@ -115,6 +115,19 @@ class BackpropTest(test.TestCase): with self.assertRaises(RuntimeError): backprop.gradients_function(f)(constant_op.constant(1.0)) + def testGradientsFunctionInCustomGradient(self): + + @custom_gradient.custom_gradient + def f(x): + (y,) = backprop.gradients_function(lambda x: x * x)(x) + + def grad(dy): + return [2 * dy] + + return y, grad + + self.assertAllEqual(f(1.0), 2.0) + def testImplicitGradOverEmbeddingLookup(self): batch_size = 8 embedding_size = 512 @@ -182,6 +195,17 @@ class BackpropTest(test.TestCase): g, = backprop.gradients_function(loss, [0])(logits, labels) self.assertAllEqual(g.numpy(), [[-0.5, 0.5]]) + def testGradientWithinTapeBlock(self): + v1 = resource_variable_ops.ResourceVariable(1.) + with backprop.GradientTape() as t: + loss = 2 * v1 + with self.assertRaises(RuntimeError): + t.gradient(loss, [v1]) + with backprop.GradientTape(persistent=True) as t: + loss = 2 * v1 + grad = t.gradient(loss, [v1]) + self.assertAllEqual(grad[0], 2.0) + @test_util.assert_no_new_tensors def testSecondGrad(self): @@ -205,11 +229,22 @@ class BackpropTest(test.TestCase): def f(x): return x * x - wrapped_fn = backprop.make_vjp(f) + wrapped_fn = backprop.make_vjp(f, persistent=False) result, vjp = wrapped_fn(constant_op.constant(3.0)) self.assertAllEqual(result, 9.0) self.assertAllEqual(vjp(2.0)[0], 12.0) + def testPersistentMakeVJP(self): + + def f(x): + return x * x + + wrapped_fn = backprop.make_vjp(f, persistent=True) + _, vjp = wrapped_fn(constant_op.constant(3.0)) + vjp_result1 = vjp(2.0)[0] + vjp_result2 = vjp(2.0)[0] + self.assertAllEqual(vjp_result1, vjp_result2, 12.0) + @test_util.assert_no_new_tensors def testGradGrad(self): diff --git a/tensorflow/python/eager/benchmarks_test.py b/tensorflow/python/eager/benchmarks_test.py index 75526ba9c139e78dfe9e3de271f1316924539371..9ca5041c38ed07b39fd73b9f110ab06e8e903251 100644 --- a/tensorflow/python/eager/benchmarks_test.py +++ b/tensorflow/python/eager/benchmarks_test.py @@ -28,11 +28,13 @@ from __future__ import print_function import time import numpy as np +import six from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.python import pywrap_tensorflow from tensorflow.python.eager import backprop # pylint: disable=unused-import from tensorflow.python.eager import context +from tensorflow.python.eager import core from tensorflow.python.eager import function from tensorflow.python.eager import test from tensorflow.python.framework import dtypes @@ -41,24 +43,31 @@ from tensorflow.python.ops import gen_array_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops +from tensorflow.python.ops import resource_variable_ops CPU = "/device:CPU:0" GPU = "/device:GPU:0" -def record_gradient_callback(inputs, attrs, results): - return backprop._record_gradient("MatMul", inputs, attrs, results, None) - - -def c_tfe_py_fastpath_execute(a, b, transpose_a=False, transpose_b=False): +def c_tfe_py_fastpath_execute(a, + b, + transpose_a=False, + transpose_b=False, + name=None): ctx = context.context() - assert not ctx.in_graph_mode( + assert ctx.executing_eagerly( ), "The prototype doesn't contain C code for graph construction" - ctx_handle = ctx._handle # pylint: disable=protected-access - - return pywrap_tensorflow.TFE_Py_FastPathExecute( - ctx_handle, None, "MatMul", record_gradient_callback, a, b, - "transpose_a", transpose_a, "transpose_b", transpose_b)[0] + try: + return pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", name, + ctx._post_execution_callbacks, a, b, "transpose_a", transpose_a, + "transpose_b", transpose_b) + except core._NotOkStatusException as e: + if name is not None: + message = e.message + " name: " + name + else: + message = e.message + six.raise_from(core._status_to_exception(e.code, message), None) class MicroBenchmarks(test.Benchmark): @@ -73,16 +82,24 @@ class MicroBenchmarks(test.Benchmark): self._num_iters_2_by_2 = 30000 self._num_iters_100_by_784 = 1000 - def _run(self, func, num_iters): + def _run(self, func, num_iters, execution_mode=None): # call func to maybe warm up the GPU - func() - start = time.time() - for _ in xrange(num_iters): + ctx = context.context() + with ctx.execution_mode(execution_mode): func() - end = time.time() - mean_us = (end - start) * 1e6 / num_iters - self.report_benchmark(iters=num_iters, wall_time=mean_us, - extras={"examples_per_sec": num_iters/(end-start)}) + if execution_mode == context.ASYNC: + ctx.async_wait() + start = time.time() + for _ in xrange(num_iters): + func() + if execution_mode == context.ASYNC: + ctx.async_wait() + end = time.time() + mean_us = (end - start) * 1e6 / num_iters + self.report_benchmark( + iters=num_iters, + wall_time=mean_us, + extras={"examples_per_sec": num_iters / (end - start)}) def benchmark_create_np_array(self): func = lambda: np.array([3.0]) @@ -227,13 +244,15 @@ class MicroBenchmarks(test.Benchmark): func = lambda: np.dot(a, b) self._run(func, num_iters) - def _benchmark_tf_matmul(self, m, transpose_b, num_iters): + def _benchmark_tf_matmul(self, m, transpose_b, num_iters, + execution_mode=None): func = lambda: math_ops.matmul(m, m, transpose_b=transpose_b) - self._run(func, num_iters) + self._run(func, num_iters, execution_mode=execution_mode) def _benchmark_gen_math_ops_matmul(self, m, transpose_b, num_iters): def func(): - gen_math_ops._mat_mul(m, m, transpose_b=transpose_b) + gen_math_ops.mat_mul(m, m, transpose_b=transpose_b) + self._run(func, num_iters) def _benchmark_tfe_py_fastpath_execute_matmul(self, m, transpose_b, @@ -257,10 +276,32 @@ class MicroBenchmarks(test.Benchmark): self._run(func, num_iters) - def _benchmark_defun_matmul(self, m, transpose_b, num_iters): + def _benchmark_defun_matmul(self, + m, + transpose_b, + num_iters, + execution_mode=None): f = function.defun(math_ops.matmul) func = lambda: f(m, m, transpose_b) - self._run(func, num_iters) + self._run(func, num_iters, execution_mode=execution_mode) + + def _benchmark_read_variable(self, m, num_iters): + self._run(m.value, num_iters) + + def _benchmark_matmul_read_variable(self, m, num_iters): + self._benchmark_gen_math_ops_matmul( + m, transpose_b=False, num_iters=num_iters) + + def _benchmark_matmul_read_variable_with_tape(self, m, num_iters): + with backprop.GradientTape() as tape: + tape.watch(m) + self._benchmark_gen_math_ops_matmul( + m, transpose_b=False, num_iters=num_iters) + + def _benchmark_read_variable_with_tape(self, m, num_iters): + with backprop.GradientTape() as tape: + tape.watch(m) + self._run(m.value, num_iters) # Benchmarks for A^2, A of dimension 2 by 2. def benchmark_np_matmul_2_by_2(self): @@ -273,6 +314,15 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_tf_matmul( m, transpose_b=False, num_iters=self._num_iters_2_by_2) + def benchmark_tf_matmul_2_by_2_CPU_async(self): + with context.device(CPU): + m = self._m_2_by_2.cpu() + self._benchmark_tf_matmul( + m, + transpose_b=False, + num_iters=self._num_iters_2_by_2, + execution_mode=context.ASYNC) + def benchmark_gen_math_ops_matmul_2_by_2_CPU(self): with context.device(CPU): m = self._m_2_by_2.cpu() @@ -297,6 +347,15 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_defun_matmul( m, transpose_b=False, num_iters=self._num_iters_2_by_2) + def benchmark_defun_matmul_2_by_2_CPU_async(self): + with context.device(CPU): + m = self._m_2_by_2.cpu() + self._benchmark_defun_matmul( + m, + transpose_b=False, + num_iters=self._num_iters_2_by_2, + execution_mode=context.ASYNC) + def benchmark_tf_matmul_2_by_2_GPU(self): if not context.num_gpus(): return @@ -305,6 +364,17 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_tf_matmul( m, transpose_b=False, num_iters=self._num_iters_2_by_2) + def benchmark_tf_matmul_2_by_2_GPU_async(self): + if not context.num_gpus(): + return + with context.device(GPU): + m = self._m_2_by_2.gpu() + self._benchmark_tf_matmul( + m, + transpose_b=False, + num_iters=self._num_iters_2_by_2, + execution_mode=context.ASYNC) + def benchmark_gen_math_ops_matmul_2_by_2_GPU(self): if not context.num_gpus(): return @@ -329,6 +399,17 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_defun_matmul( m, transpose_b=False, num_iters=self._num_iters_2_by_2) + def benchmark_defun_matmul_2_by_2_GPU_async(self): + if not context.num_gpus(): + return + with context.device(GPU): + m = self._m_2_by_2.gpu() + self._benchmark_defun_matmul( + m, + transpose_b=False, + num_iters=self._num_iters_2_by_2, + execution_mode=context.ASYNC) + # Benchmarks for AA.T, A of dimension 100 by 784. def benchmark_np_matmul_100_by_784(self): self._benchmark_np_matmul( @@ -342,6 +423,15 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_tf_matmul( m, transpose_b=True, num_iters=self._num_iters_100_by_784) + def benchmark_tf_matmul_100_by_784_CPU_async(self): + with context.device(CPU): + m = self._m_100_by_784.cpu() + self._benchmark_tf_matmul( + m, + transpose_b=True, + num_iters=self._num_iters_100_by_784, + execution_mode=context.ASYNC) + def benchmark_gen_math_ops_matmul_100_by_784_CPU(self): with context.device(CPU): m = self._m_100_by_784.cpu() @@ -374,6 +464,17 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_tf_matmul( m, transpose_b=True, num_iters=self._num_iters_100_by_784) + def benchmark_tf_matmul_100_by_784_GPU_async(self): + if not context.num_gpus(): + return + with context.device(GPU): + m = self._m_100_by_784.gpu() + self._benchmark_tf_matmul( + m, + transpose_b=True, + num_iters=self._num_iters_100_by_784, + execution_mode=context.ASYNC) + def benchmark_gen_math_ops_matmul_100_by_784_GPU(self): if not context.num_gpus(): return @@ -398,6 +499,43 @@ class MicroBenchmarks(test.Benchmark): self._benchmark_defun_matmul( m, transpose_b=True, num_iters=self._num_iters_100_by_784) + def benchmark_matmul_read_variable_op_2_by_2_CPU(self): + with context.device(CPU): + m = resource_variable_ops.ResourceVariable(self._m_2_by_2) + self._benchmark_matmul_read_variable(m, num_iters=self._num_iters_2_by_2) + + def benchmark_matmul_read_variable_op_with_tape_2_by_2_CPU(self): + with context.device(CPU): + m = resource_variable_ops.ResourceVariable(self._m_2_by_2) + self._benchmark_matmul_read_variable_with_tape( + m, num_iters=self._num_iters_2_by_2) + + def benchmark_read_variable_op_2_by_2_CPU(self): + with context.device(CPU): + m = resource_variable_ops.ResourceVariable(self._m_2_by_2) + self._benchmark_read_variable(m, num_iters=self._num_iters_2_by_2) + + def benchmark_read_variable_op_2_by_2_GPU(self): + if not context.num_gpus(): + return + with context.device(GPU): + m = resource_variable_ops.ResourceVariable(self._m_2_by_2.gpu()) + self._benchmark_read_variable(m, num_iters=self._num_iters_2_by_2) + + def benchmark_read_variable_op_with_tape_2_by_2_CPU(self): + with context.device(CPU): + m = resource_variable_ops.ResourceVariable(self._m_2_by_2) + self._benchmark_read_variable_with_tape( + m, num_iters=self._num_iters_2_by_2) + + def benchmark_read_variable_op_with_tape_2_by_2_GPU(self): + if not context.num_gpus(): + return + with context.device(GPU): + m = resource_variable_ops.ResourceVariable(self._m_2_by_2.gpu()) + self._benchmark_read_variable_with_tape( + m, num_iters=self._num_iters_2_by_2) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py index b6c7d823237231a138f6a25bb9d03954b69d58d9..6c9a14730c0db4bdf23fc10b23d63b758349bdc1 100644 --- a/tensorflow/python/eager/context.py +++ b/tensorflow/python/eager/context.py @@ -30,7 +30,9 @@ from tensorflow.python.framework import c_api_util from tensorflow.python.framework import device as pydev from tensorflow.python.framework import errors from tensorflow.python.util import compat +from tensorflow.python.util import is_in_graph_mode from tensorflow.python.util import tf_contextlib +from tensorflow.python.util.tf_export import tf_export GRAPH_MODE = 0 EAGER_MODE = 1 @@ -51,6 +53,28 @@ DEVICE_PLACEMENT_WARN = pywrap_tensorflow.TFE_DEVICE_PLACEMENT_WARN DEVICE_PLACEMENT_SILENT = pywrap_tensorflow.TFE_DEVICE_PLACEMENT_SILENT DEVICE_PLACEMENT_SILENT_FOR_INT32 = ( pywrap_tensorflow.TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32) +SYNC = 0 +ASYNC = 1 + + +class _TensorCache(object): + """Simple cache which evicts items based on length in a FIFO manner.""" + + def __init__(self, max_items=256): + self._data = collections.OrderedDict() + self._max_items = max_items if max_items else 256 + + def put(self, key, value): + self._data[key] = value + + if len(self._data) > self._max_items: + self._data.popitem(last=False) + + def get(self, key): + return self._data.get(key, None) + + def flush(self): + self._data = {} # TODO(agarwal): better name ? @@ -66,24 +90,36 @@ class _EagerContext(threading.local): self.recording_summaries = False self.summary_writer_resource = None self.scalar_cache = {} + self.ones_rank_cache = _TensorCache() + self.execution_mode = None -ContextStackEntry = collections.namedtuple( - "ContextStackEntry", ["is_building_function", "enter_context_fn"]) +ContextSwitch = collections.namedtuple( + "ContextSwitch", ["is_building_function", "enter_context_fn"]) -class ContextStack(threading.local): +# `_ContextSwitchStack` is a `threading.local` to match the semantics of +# ``DefaultGraphStack`, which is also a `threading.local`. +class _ContextSwitchStack(threading.local): """A thread-local stack of context switches.""" - def __init__(self): - super(ContextStack, self).__init__() + def __init__(self, eager): + super(_ContextSwitchStack, self).__init__() self.stack = [] + if eager: + # Initialize the stack with a pointer to enter the eager context; this + # ensures that the fact that eager execution was enabled is propagated + # across threads, since (1) `enable_eager_execution` modifies a + # process-level flag (`_default_mode`) and (2) `__init__` is called each + # time a threading.local object is used in a separate thread. + self.push(is_building_function=False, enter_context_fn=eager_mode) def push(self, is_building_function, enter_context_fn): """Push metadata about a context switch onto the stack. A context switch can take one of two forms: installing a graph as the - default graph, or entering the eager context. + default graph, or entering the eager context. For each context switch, + we record whether or not the entered context is building a function. Args: is_building_function: (bool.) Whether the context is building a function. @@ -92,7 +128,7 @@ class ContextStack(threading.local): """ self.stack.append( - ContextStackEntry(is_building_function, enter_context_fn)) + ContextSwitch(is_building_function, enter_context_fn)) def pop(self): """Pop the stack.""" @@ -100,34 +136,49 @@ class ContextStack(threading.local): self.stack.pop() -context_stack = ContextStack() - - # TODO(agarwal): rename to EagerContext / EagerRuntime ? # TODO(agarwal): consider keeping the corresponding Graph here. class Context(object): """Environment in which eager operations execute.""" - def __init__(self, config=None, device_policy=None): + # TODO(agarwal): create and link in some documentation for `execution_mode`. + # pylint: disable=redefined-outer-name + def __init__(self, config=None, device_policy=None, execution_mode=None): """Creates a new Context. Args: config: (Optional.) A `ConfigProto` protocol buffer with configuration - options for the Context. Note that a lot of these options may be - currently unimplemented or irrelevant when eager execution is enabled. + options for the Context. Note that a lot of these options may be + currently unimplemented or irrelevant when eager execution is enabled. device_policy: (Optional.) What policy to use when trying to run an - operation on a device with inputs which are not on that device. - Valid values: - tfe.DEVICE_PLACEMENT_EXPLICIT: raises an error if the placement is not - correct. - tfe.DEVICE_PLACEMENT_WARN: copies the tensors which are not on the + operation on a device with inputs which are not on that device. + When set to None, an appropriate value will be picked automatically. + The value picked may change between TensorFlow releases. + + Defaults to tf.contrib.eager.DEVICE_PLACEMENT_SILENT_FOR_INT32. + Valid values: + - tfe.DEVICE_PLACEMENT_EXPLICIT: raises an error if the placement is + not correct. + - tfe.DEVICE_PLACEMENT_WARN: copies the tensors which are not on the right device but raises a warning. - tfe.DEVICE_PLACEMENT_SILENT: silently copies the tensors. This might + - tfe.DEVICE_PLACEMENT_SILENT: silently copies the tensors. This might hide performance problems. - tfe.DEVICE_PLACEMENT_SILENT_FOR_INT32: silently copies int32 tensors, + - tfe.DEVICE_PLACEMENT_SILENT_FOR_INT32: silently copies int32 tensors, raising errors on the other ones. + execution_mode: (Optional.) Policy controlling how operations dispatched + are actually executed. When set to None, an appropriate value will be + picked automatically. The value picked may change between TensorFlow + releases. + Valid values: + - tf.contrib.eager.SYNC: executes each operation synchronously. + - tf.contrib.eager.ASYNC: executes each operation asynchronously. These + operations may return "non-ready" handles. + + Raises: + ValueError: If execution_mode is not valid. """ self._eager_context = _EagerContext() + self._context_switches = _ContextSwitchStack(self.executing_eagerly()) self._context_handle = None self._context_devices = None self._post_execution_callbacks = [] @@ -135,6 +186,14 @@ class Context(object): self._seed = None self._initialize_lock = threading.Lock() self._device_policy = device_policy + if execution_mode not in (None, SYNC, ASYNC): + raise ValueError( + "execution_mode should be None/SYNC/ASYNC. Got %s" % execution_mode) + if execution_mode is None: + execution_mode = SYNC + self._execution_mode = execution_mode + + # pylint: enable=redefined-outer-name def _set_global_seed(self, seed): """Set a global eager mode seed for random ops.""" @@ -172,6 +231,8 @@ class Context(object): if self._device_policy is not None: pywrap_tensorflow.TFE_ContextOptionsSetDevicePlacementPolicy( opts, self._device_policy) + if self._execution_mode == ASYNC: + pywrap_tensorflow.TFE_ContextOptionsSetAsync(True) self._context_handle = pywrap_tensorflow.TFE_NewContext(opts, status) finally: pywrap_tensorflow.TFE_DeleteContextOptions(opts) @@ -230,26 +291,29 @@ class Context(object): old_mode = ctx.mode ctx.mode = mode if mode == EAGER_MODE: - context_stack.push(False, eager_mode) + # Entering graph mode does not provide us with sufficient information to + # record a context switch; graph-based context switches are only logged + # when a graph is registered as the default graph. + self.context_switches.push(False, eager_mode) try: yield finally: ctx.mode = old_mode if mode == EAGER_MODE: - context_stack.pop() - - def in_graph_mode(self): - """Returns True if current thread is in GRAPH mode.""" - return self._eager_context.mode == GRAPH_MODE + self.context_switches.pop() - def in_eager_mode(self): - """Returns True if current thread is in EAGER mode.""" + def executing_eagerly(self): + """Returns True if current thread has eager executing enabled.""" return self._eager_context.mode == EAGER_MODE def scalar_cache(self): """Per-device cache for scalars.""" return self._eager_context.scalar_cache + def ones_rank_cache(self): + """Per-device cache for scalars.""" + return self._eager_context.ones_rank_cache + @property def scope_name(self): """Returns scope name for the current thread.""" @@ -333,6 +397,43 @@ class Context(object): """List of the names of devices available to execute operations.""" return self._devices + def get_execution_mode(self): + mode = self._eager_context.execution_mode + if mode is None: + mode = self._execution_mode + return mode + + def set_execution_mode(self, mode): + """Sets execution mode for current thread.""" + if mode not in (None, SYNC, ASYNC): + raise ValueError( + "Execution mode should be None/SYNC/ASYNC. Got %s" % mode) + if mode is None: + mode = SYNC + self._eager_context.execution_mode = mode + with errors.raise_exception_on_not_ok_status() as status: + pywrap_tensorflow.TFE_ContextSetAsyncForThread(self._handle, + mode == ASYNC, status) + + @tf_contextlib.contextmanager + def execution_mode(self, mode): + """Context manager for setting execution mode for current thread.""" + old_mode = self.get_execution_mode() + try: + self.set_execution_mode(mode) + yield + finally: + self.set_execution_mode(old_mode) + + def async_wait(self): + """Waits for ops dispatched in ASYNC mode to finish.""" + with errors.raise_exception_on_not_ok_status() as status: + pywrap_tensorflow.TFE_ContextAsyncWait(self._handle, status) + + def async_clear_error(self): + """Clears errors raised during ASYNC execution.""" + pywrap_tensorflow.TFE_ContextAsyncClearError(self._handle) + def num_gpus(self): """The number of GPUs available to execute operations.""" self._initialize_handle_and_devices() @@ -455,6 +556,11 @@ class Context(object): run_metadata.ParseFromString(compat.as_bytes(proto_data)) return run_metadata + @property + def context_switches(self): + """Returns a stack of context switches.""" + return self._context_switches + _context = None _context_lock = threading.Lock() @@ -496,23 +602,29 @@ def internal_operation_seed(): return context()._internal_operation_seed() # pylint: disable=protected-access -def in_graph_mode(): - """Returns True if current thread is in GRAPH mode for default context.""" - return context().in_graph_mode() +@tf_export("executing_eagerly") +def executing_eagerly(): + """Returns True if the current thread has eager execution enabled. + + Eager execution is typically enabled via @{tf.enable_eager_execution}, + but may also be enabled within the context of a Python function via + tf.contrib.eager.py_func. + """ + return context().executing_eagerly() def in_eager_mode(): - """Returns True if current thread is in EAGER mode for default context.""" - return context().in_eager_mode() + """Use executing_eagerly() instead. This function will be removed.""" + return executing_eagerly() def graph_mode(): - """Context-manager to enable GRAPH mode for current thread.""" + """Context-manager to disable eager execution for the current thread.""" return context()._mode(GRAPH_MODE) # pylint: disable=protected-access def eager_mode(): - """Context-manager to enable EAGER mode for current thread.""" + """Context-manager to enable eager execution for the current thread.""" return context()._mode(EAGER_MODE) # pylint: disable=protected-access @@ -566,6 +678,26 @@ def list_devices(): return context().devices() +def set_execution_mode(mode): + """Sets execution mode for the current thread.""" + context().set_execution_mode(mode) + + +def execution_mode(mode): + """Context manager for setting execution mode for current thread.""" + return context().execution_mode(mode) + + +def async_wait(): + """Waits for ops dispatched in ASYNC mode to finish.""" + return context().async_wait() + + +def async_clear_error(): + """Clears errors raised during ASYNC execution mode.""" + return context().async_clear_error() + + def num_gpus(): """Get the number of available GPU devices. @@ -599,3 +731,14 @@ def export_run_metadata(): A RunMetadata protocol buffer. """ return context().export_run_metadata() + + +# Not every user creates a Context via context.context() +# (for example, enable_eager_execution in python/framework/ops.py), +# but they do all import this file. Note that IS_IN_GRAPH_MODE and +# in_graph_mode are both parameterless functions. +def _tmp_in_graph_mode(): + return not executing_eagerly() + + +is_in_graph_mode.IS_IN_GRAPH_MODE = _tmp_in_graph_mode diff --git a/tensorflow/python/eager/core.py b/tensorflow/python/eager/core.py index 483b7172107838a0069831f2347b0c644c05c000..8fb69300209d74a164c38654d737432cdfb7884a 100644 --- a/tensorflow/python/eager/core.py +++ b/tensorflow/python/eager/core.py @@ -47,3 +47,17 @@ class _NotOkStatusException(Exception): pywrap_tensorflow.TFE_Py_RegisterExceptionClass(_NotOkStatusException) + + +class _FallbackException(Exception): + """Exception class to handle fallback from the fastpath. + + The fastpath that we refer to here is the one implemented to reduce per-op + overheads (TFE_Py_FastPathExecute_C). If the conditions for executing the op + on the fastpath are not met, we fallback to a safer (and more complete) + slowpath, and this Exception is raised to signal that transition. + """ + pass + + +pywrap_tensorflow.TFE_Py_RegisterFallbackExceptionClass(_FallbackException) diff --git a/tensorflow/python/eager/core_test.py b/tensorflow/python/eager/core_test.py index ee3c10633e1cb849e319f2f5490e5beb5dd15c80..6ebf5b24819d48ba4a17d6059510eee7affe40ea 100644 --- a/tensorflow/python/eager/core_test.py +++ b/tensorflow/python/eager/core_test.py @@ -33,7 +33,10 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors 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 gen_resource_variable_ops from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import resource_variable_ops def execute(op_name, num_outputs, inputs, attrs=None): @@ -54,13 +57,22 @@ class TFETest(test_util.TensorFlowTestCase): def testContext(self): ctx = context.Context() - self.assertFalse(ctx.in_graph_mode()) - self.assertTrue(ctx.in_eager_mode()) + self.assertTrue(ctx.executing_eagerly()) self.assertEqual('', ctx.scope_name) ctx.scope_name = 'foo' self.assertEqual('foo', ctx.scope_name) + self.assertEqual(context.SYNC, ctx.get_execution_mode()) + ctx.set_execution_mode(context.ASYNC) + self.assertEqual(context.ASYNC, ctx.get_execution_mode()) + ctx.set_execution_mode(context.SYNC) + self.assertEqual(context.SYNC, ctx.get_execution_mode()) + with ctx.execution_mode(context.ASYNC): + self.assertEqual(context.ASYNC, ctx.get_execution_mode()) + ctx.set_execution_mode(context.SYNC) + self.assertEqual(context.SYNC, ctx.get_execution_mode()) + self.assertIsNone(ctx.summary_writer_resource) ctx.summary_writer_resource = 'mock' self.assertEqual('mock', ctx.summary_writer_resource) @@ -99,19 +111,30 @@ class TFETest(test_util.TensorFlowTestCase): self.assertEqual(len(cpu_stats.node_stats), 1) self.assertEqual(cpu_stats.node_stats[0].node_name, 'Add') - def testContextStackContainsEagerMode(self): - # Eager execution has been enabled, and no other context - # switch has occurred, so `context_stack` should contain - # exactly one entry. - self.assertEqual(len(context.context_stack.stack), 1) - stack_entry = context.context_stack.stack[0] + def testShouldCopy(self): + if not context.context().num_gpus(): + self.skipTest('No devices other than CPUs found') + with ops.device('gpu:0'): + x = constant_op.constant(1.0) + y = array_ops.identity(x) + # The value we're testing y.device against will depend on what the behavior + # of not explicitly specifying a device in the context is. This behavior is + # subject to change (for example, in the future we may want to use GPUs, if + # available, when no device is explicitly provided) + self.assertEqual(y.device, '/job:localhost/replica:0/task:0/device:CPU:0') + + def testContextSwitchStackContainsEagerMode(self): + # Eager execution has been enabled, and no other context switch has + # occurred, so `context_switches` should contain exactly one entry. + self.assertEqual(len(context.context().context_switches.stack), 1) + switch = context.context().context_switches.stack[0] # The entry should log that eager mode was entered. - self.assertIs(stack_entry.enter_context_fn, context.eager_mode) + self.assertIs(switch.enter_context_fn, context.eager_mode) # It is not possible to build a graph function when eager execution # is enabled; the stack entry should reflect this fact. - self.assertFalse(stack_entry.is_building_function) + self.assertFalse(switch.is_building_function) def testInt32GPU(self): if not context.context().num_gpus(): @@ -135,9 +158,9 @@ class TFETest(test_util.TensorFlowTestCase): def get_context_values(ctx): return [ - ctx.in_graph_mode(), - ctx.in_eager_mode(), ctx.scope_name, ctx.summary_writer_resource, - ctx.device_name, ctx.num_gpus() + ctx.executing_eagerly(), ctx.scope_name, ctx.summary_writer_resource, + ctx.device_name, + ctx.num_gpus() ] def get_values(ctx, values): @@ -168,6 +191,18 @@ class TFETest(test_util.TensorFlowTestCase): attrs=('T', x.dtype.as_datatype_enum))[0].cpu().numpy() self.assertEqual(3, result) + def testResourceTensorPlacement(self): + if not context.context().num_gpus(): + self.skipTest('No GPUs found') + + with context.device('gpu:0'): + v = resource_variable_ops.ResourceVariable(1.0) + with context.device('cpu:0'): + # Check that even though we specified the cpu device we'll run the read op + # in the device where the handle is. + self.assertAllEqual( + gen_resource_variable_ops.read_variable_op(v.handle, v.dtype), 1.0) + def testCopyBetweenDevices(self): if not context.context().num_gpus(): self.skipTest('No GPUs found') @@ -182,6 +217,23 @@ class TFETest(test_util.TensorFlowTestCase): with self.assertRaises(RuntimeError): x.gpu(context.context().num_gpus() + 1) + def testCopyBetweenDevicesAsync(self): + if not context.context().num_gpus(): + self.skipTest('No GPUs found') + with context.execution_mode(context.ASYNC): + x = constant_op.constant([[1., 2.], [3., 4.]]) + x = x.cpu() + x = x.gpu() + x = x.gpu() + x = x.cpu() + context.async_wait() + + # Invalid device + with self.assertRaises(RuntimeError): + x.gpu(context.context().num_gpus() + 1) + context.async_wait() + context.async_clear_error() + def testCopyScope(self): if not context.context().num_gpus(): self.skipTest('No GPUs found') @@ -222,16 +274,49 @@ class TFETest(test_util.TensorFlowTestCase): attrs=('T', three.dtype.as_datatype_enum))[0] self.assertAllEqual(15, product) + def testExecuteBasicAsync(self): + with context.execution_mode(context.ASYNC): + three = constant_op.constant(3) + five = constant_op.constant(5) + product = execute( + b'Mul', + num_outputs=1, + inputs=[three, five], + attrs=('T', three.dtype.as_datatype_enum))[0] + self.assertAllEqual(15, product) + # Error: Invalid arguments + context.set_execution_mode(context.ASYNC) + with self.assertRaises(errors.InvalidArgumentError): + execute( + b'MatMul', + num_outputs=1, + inputs=[three, five], + attrs=('transpose_a', False, 'transpose_b', False, 'T', + three.dtype.as_datatype_enum)) + context.async_wait() + context.async_clear_error() + context.set_execution_mode(context.SYNC) + def testExecuteTooManyNumOutputs(self): # num_outputs provided is 50, but only one output is produced. - # That should be okay. product = execute( b'Mul', num_outputs=50, - inputs=[constant_op.constant(3), constant_op.constant(5)], + inputs=[constant_op.constant(3), + constant_op.constant(5)], attrs=('T', dtypes.int32.as_datatype_enum))[0] self.assertAllEqual(15, product) + def testExecuteTooFewNumOutputs(self): + # num_outputs provided is 0, but one output is produced. + with self.assertRaises(errors.InvalidArgumentError): + _ = execute( + b'Mul', + num_outputs=0, + inputs=[constant_op.constant(3), + constant_op.constant(5)], + attrs=('T', dtypes.int32.as_datatype_enum))[0] + def testMatMulGPU(self): if not context.context().num_gpus(): self.skipTest('No GPUs found') @@ -519,5 +604,61 @@ class TFETest(test_util.TensorFlowTestCase): self.assertIsInstance(t, ops.EagerTensor) +class SendRecvTest(test_util.TensorFlowTestCase): + + cpu_device = '/job:localhost/replica:0/task:0/device:CPU:0' + + def _send(self, tensor, tensor_name, to_device): + return execute( + b'_Send', num_outputs=0, inputs=[tensor], + attrs=('T', tensor.dtype.as_datatype_enum, + 'tensor_name', tensor_name, + 'send_device', tensor.device, + 'send_device_incarnation', 0, + 'recv_device', to_device, + 'client_terminated', True)) + + def _recv(self, dtype, tensor_name, from_device): + device_name = context.context().device_name + if not device_name: + device_name = self.cpu_device + return execute( + b'_Recv', num_outputs=1, inputs=[], + attrs=('tensor_type', dtype.as_datatype_enum, + 'tensor_name', tensor_name, + 'send_device', from_device, + 'send_device_incarnation', 0, + 'recv_device', device_name, + 'client_terminated', False))[0] + + def testBasic(self): + t0 = constant_op.constant(1.0) + t1 = constant_op.constant(2.0) + self._send(t0, 't0', self.cpu_device) + self._send(t1, 't1', self.cpu_device) + self.assertAllEqual( + self._recv(dtypes.float32, 't0', self.cpu_device), + 1.0) + self.assertAllEqual( + self._recv(dtypes.float32, 't1', self.cpu_device), + 2.0) + + def testLocalCrossDevice(self): + if not context.context().num_gpus(): + self.skipTest('No GPUs found') + gpu_device_name = '/job:localhost/replica:0/task:0/device:GPU:0' + with ops.device('GPU:0'): + t0 = constant_op.constant(1.0) + self._send(t0, 't0', self.cpu_device) + self.assertAllEqual( + self._recv(dtypes.float32, 't0', gpu_device_name), + 1.0) + self._send(constant_op.constant(2.0), 't1', gpu_device_name) + with ops.device('GPU:0'): + self.assertAllEqual( + self._recv(dtypes.float32, 't1', self.cpu_device), + 2.0) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/eager/custom_gradient.py b/tensorflow/python/eager/custom_gradient.py deleted file mode 100644 index 05460ff9968312528d87f5fc2ad0495b4da2ad1a..0000000000000000000000000000000000000000 --- a/tensorflow/python/eager/custom_gradient.py +++ /dev/null @@ -1,91 +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. -# ============================================================================== -"""Decorator to overrides the gradient for a function.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.python.eager import context -from tensorflow.python.eager import tape -from tensorflow.python.framework import ops as tf_ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gen_array_ops -from tensorflow.python.util import nest -from tensorflow.python.util import tf_decorator - - -def custom_gradient(f): - """Decorator to define a function with a custom gradient. - - The input function is expected to return the tuple - (results, gradient_function). - - The output function will return results while possibly recording the - gradient_function and inputs in the tape. - - Args: - f: function to be decorated. - - Returns: - decorated function. - """ - - def decorated(*args, **kwargs): - """Decorated function with custom gradient.""" - if context.in_graph_mode(): - if kwargs: - raise ValueError( - "custom_gradient in graph mode doesn't support keyword arguments.") - name = "CustomGradient-%s" % tf_ops.uid() - args = [tf_ops.convert_to_tensor(x) for x in args] - result, grad_fn = f(*args) - flat_result = nest.flatten(result) - all_tensors = flat_result + args - - @tf_ops.RegisterGradient(name) - def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable - gradients = nest.flatten(grad_fn(*result_grads[:len(flat_result)])) - # Need to return one value per input to the IdentityN, so pad the - # gradients of the inputs of the custom_gradient function with the - # gradients of the outputs as well. - return ([None] * len(flat_result)) + gradients - - with tf_ops.get_default_graph().gradient_override_map( - {"IdentityN": name}): - all_tensors = array_ops.identity_n(all_tensors) - return nest.pack_sequence_as( - structure=result, flat_sequence=all_tensors[:len(flat_result)]) - - input_tensors = [tf_ops.convert_to_tensor(x) for x in args] - - with tape.stop_recording(): - result, grad_fn = f(*args, **kwargs) - flat_result = nest.flatten(result) - # TODO(apassos) consider removing the identity below. - flat_result = [gen_array_ops.identity(x) for x in flat_result] - - def actual_grad_fn(*outputs): - return nest.flatten(grad_fn(*outputs)) - - tape.record_operation( - f.__name__, - flat_result, - input_tensors, - actual_grad_fn) - flat_result = list(flat_result) - return nest.pack_sequence_as(result, flat_result) - - return tf_decorator.make_decorator(f, decorated) diff --git a/tensorflow/python/eager/execute.py b/tensorflow/python/eager/execute.py index 306cf07aabe1c214d02da5f077a57043cc1f4089..2ff5b8d8f489731c14d8abb81652a17026ed4935 100644 --- a/tensorflow/python/eager/execute.py +++ b/tensorflow/python/eager/execute.py @@ -72,7 +72,7 @@ def execute_with_callbacks(op_name, num_outputs, inputs, attrs, ctx, name=None): """Monkey-patch to execute to enable execution callbacks.""" tensors = quick_execute(op_name, num_outputs, inputs, attrs, ctx, name) for callback in ctx.post_execution_callbacks: - callback(op_name, name, attrs, inputs, tensors) + callback(op_name, inputs, attrs, tensors, name) return tensors diff --git a/tensorflow/python/eager/execution_callbacks.py b/tensorflow/python/eager/execution_callbacks.py index 2f1654dda499583fe4766cbe2e330399defc96fd..535361498a9dd33003d0479051e97d7ff2553067 100644 --- a/tensorflow/python/eager/execution_callbacks.py +++ b/tensorflow/python/eager/execution_callbacks.py @@ -104,10 +104,10 @@ class InfOrNanError(Exception): def inf_nan_callback(op_type, - op_name, - attrs, inputs, + attrs, outputs, + op_name, check_inf=True, check_nan=True, action=_DEFAULT_CALLBACK_ACTION): @@ -121,14 +121,14 @@ def inf_nan_callback(op_type, Args: op_type: Name of the TFE operation type (e.g., `MatMul`). - op_name: Name of the TFE operation. This name is set by client and can be - `None` if it unset. - attrs: Attributes of the TFE operation, as a tuple of alternating attribute - names and attribute values. inputs: The `list` of input tensors to the operation, currently unused by this callback. + attrs: Attributes of the TFE operation, as a tuple of alternating attribute + names and attribute values. outputs: The `list` of output tensors from the operation, checked by this callback for `inf` and `nan` values. + op_name: Name of the TFE operation. This name is set by client and can be + `None` if it unset. check_inf: (`bool`) Whether this callback should check for `inf` values in the output tensor values. check_nan: (`bool`) Whether this callback should check for `nan` values in @@ -153,7 +153,7 @@ def inf_nan_callback(op_type, continue numpy_dtype = output.dtype.as_numpy_dtype - if (np.issubdtype(numpy_dtype, np.float) or + if (np.issubdtype(numpy_dtype, np.floating) or np.issubdtype(numpy_dtype, np.complex) or np.issubdtype(numpy_dtype, np.integer)): try: @@ -187,26 +187,38 @@ def inf_nan_callback(op_type, def inf_callback(op_type, - op_name, - attrs, inputs, + attrs, outputs, + op_name, action=_DEFAULT_CALLBACK_ACTION): """A specialization of `inf_nan_callback` that checks for `inf`s only.""" inf_nan_callback( - op_type, op_name, attrs, inputs, outputs, check_inf=True, check_nan=False, + op_type, + inputs, + attrs, + outputs, + op_name, + check_inf=True, + check_nan=False, action=action) def nan_callback(op_type, - op_name, - attrs, inputs, + attrs, outputs, + op_name, action=_DEFAULT_CALLBACK_ACTION): """A specialization of `inf_nan_callback` that checks for `nan`s only.""" inf_nan_callback( - op_type, op_name, attrs, inputs, outputs, check_inf=False, check_nan=True, + op_type, + inputs, + attrs, + outputs, + op_name, + check_inf=False, + check_nan=True, action=action) diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py index 81b1f6f12a1899ddccb711a81122905bfd363748..343012e552592a6f8bb1255118add3e938aa443c 100644 --- a/tensorflow/python/eager/function.py +++ b/tensorflow/python/eager/function.py @@ -36,6 +36,8 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes as dtypes_module 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 from tensorflow.python.ops import gradients_impl from tensorflow.python.util import compat from tensorflow.python.util import nest @@ -110,7 +112,7 @@ def _convert_to_graph_tensor(value, dtype=None, name=None, as_ref=False): """ del as_ref # Unused. - if context.in_eager_mode(): + if context.executing_eagerly(): return value default_graph = ops.get_default_graph() @@ -161,31 +163,15 @@ class CapturingGraph(ops.Graph): op_def=None, compute_shapes=True, compute_device=True): - # TODO(apassos) probably control flow has to be handled delicately here as - # in if a resource is accessed inside a control flow context we need the - # control dependency to point to something outside the context which is - # guaranteed to happen after the access. - # # TODO(apassos) this should do some form of alias analysis as ops which # forward the resources such as Identity and Switch can cause serialization # to fail. - resource_inputs = set() - control_inputs = set() for i, inp in enumerate(inputs): if inp.graph is not self: inputs[i] = capture_value(self.captures, inp, inp.dtype, inp.op.name) - inp = inputs[i] - if inp.dtype == dtypes_module.resource: - if inp.name in self._last_op_using_resource_tensor: - control_inputs.add(self._last_op_using_resource_tensor[inp.name]) - resource_inputs.add(inp.name) - with self.control_dependencies(list(control_inputs)): - op = super(CapturingGraph, self).create_op( - op_type, inputs, dtypes, input_types, name, attrs, op_def, - compute_shapes, compute_device) - for name in resource_inputs: - self._last_op_using_resource_tensor[name] = op - return op + return super(CapturingGraph, self).create_op( + op_type, inputs, dtypes, input_types, name, attrs, op_def, + compute_shapes, compute_device) # TODO(apassos): it'd be really nice if we could scope this registration. @@ -195,33 +181,66 @@ ops.register_tensor_conversion_function( ops.EagerTensor, _convert_to_graph_tensor, priority=-1) -class _CapturingContext(object): - """Tracks references to Tensors outside this context while it is active.""" +# pylint: disable=invalid-name +class HelperContext(object): + """ControlFlowContext with a customizable AddOp method.""" - def __init__(self): - # known_ops are ops which are created while this context is active - self.known_ops = set() + def __init__(self, add_op_internal): + self._add_op_internal = add_op_internal + self._values = set() # control flow code sometimes updates this. + + def _AddOpInternal(self, op): + self._add_op_internal(op) + + @property + def outer_context(self): + return self._outer_context + + def GetWhileContext(self): + if self._outer_context: + return self._outer_context.GetWhileContext() + + def IsWhileContext(self): + return False - # captured_tensors are all tensors referenced to by ops in this context but - # not produced in it - self.captured_tensors = set() + def IsCondContext(self): + return False + + def IsXLAContext(self): + return False def AddOp(self, op): # pylint: disable=invalid-name - if op.type in ["Variable", "VariableV2", "VarHandleOp"]: - raise ValueError("tfe.defun cannot capture variables created without " - "using tf.get_variable. Op: %s" % op) - self.known_ops.add(op) - for i in op.inputs: - if i.op not in self.known_ops: - self.captured_tensors.add(i) + self._AddOpInternal(op) + if self._outer_context: + self._outer_context.AddOp(op) + + def AddName(self, _): + pass + + def AddInnerOp(self, op): + self._AddOpInternal(op) + if self._outer_context: + self._outer_context.AddInnerOp(op) + + def AddValue(self, val): + if self._outer_context: + return self._outer_context.AddValue(val) + else: + return val def __enter__(self): + # pylint: disable=protected-access self._g = ops.get_default_graph() - self._old = self._g._get_control_flow_context() # pylint: disable=protected-access - self._g._set_control_flow_context(self) # pylint: disable=protected-access + self._outer_context = self._g._get_control_flow_context() + self._g._set_control_flow_context(self) + self._nested_contexts = ( + self._outer_context._nested_contexts + if self._outer_context is not None else None) + # pylint: enable=protected-access - def __exit__(self, _, __, ___): # pylint: disable=invalid-name - self._g._set_control_flow_context(self._old) # pylint: disable=protected-access + def __exit__(self, *_): + self._g._set_control_flow_context(self._outer_context) # pylint: disable=protected-access +# pylint: enable=invalid-name def _forward_name(n): @@ -276,7 +295,7 @@ class _EagerDefinedFunction(object): proto_data = pywrap_tensorflow.TF_GetBuffer(buffer_) function_def = function_pb2.FunctionDef() function_def.ParseFromString(compat.as_bytes(proto_data)) - if context.in_eager_mode(): + if context.executing_eagerly(): _register(fn) self.definition = function_def self.name = function_def.signature.name @@ -292,6 +311,22 @@ def _map_sequence_obj_to_idx(sequence): return {id(x): i for i, x in enumerate(sequence)} +def _flatten(sequence): + """A wrapper around `nest.flatten` that also unpacks `IndexedSlices`.""" + # TODO(akshayka): Support `SparseTensor` in a similar fashion. + flat_sequence = nest.flatten(sequence) + outputs = [] + for item in flat_sequence: + if isinstance(item, ops.IndexedSlices): + if item.dense_shape is not None: + outputs.extend([item.values, item.indices, item.dense_shape]) + else: + outputs.extend([item.values, item.indices]) + else: + outputs.append(item) + return outputs + + class GraphModeFunction(object): """Callable object representing a graph-mode function. @@ -333,14 +368,14 @@ class GraphModeFunction(object): self._input_placeholders = input_placeholders self._extra_inputs = list(extra_inputs) self._graph = graph - self._has_backprop = False + self._backward_function = None self._func_name = name self._function_def = defined_function self._num_outputs = len(defined_function.signature.output_arg) self._ops = operations self._func_outputs = func_outputs self._returns = [func_outputs] if isinstance( - func_outputs, (ops.Tensor, type(None))) else list(func_outputs) + func_outputs, (ops.Tensor, type(None))) else _flatten(func_outputs) self._output_shapes = output_shapes self._variables = variables if variables is not None else [] @@ -348,11 +383,23 @@ class GraphModeFunction(object): def variables(self): return self._variables - def _compute_backprop(self): - """Computes the backprop function object for this function.""" - self._has_backprop = True + def _construct_backprop_function(self): + """Constructs the backprop function object for this function.""" with self._graph.as_default(), context.graph_mode(): - c = _CapturingContext() + c_known_ops = set() + c_captured_tensors = set() + + def add_op_internal(op): + if op.type in ["Variable", "VariableV2", "VarHandleOp"]: + raise ValueError("tfe.defun cannot capture variables created without " + "using tf.get_variable. Op: %s" % op) + c_known_ops.add(op) + for i in op.inputs: + if i.op not in c_known_ops: + c_captured_tensors.add(i) + + c = HelperContext(add_op_internal) + with c: filtered_outputs = [x for x in self._returns if x is not None] self._out_grad_placeholders = [ @@ -361,13 +408,16 @@ class GraphModeFunction(object): filtered_outputs, self._input_placeholders, grad_ys=self._out_grad_placeholders) - shapes = tuple(x.shape for x in in_gradients if x is not None) - captures = list(sorted(c.captured_tensors, key=lambda x: x.name)) + + backward_outputs = tuple( + grad for grad in _flatten(in_gradients) if grad is not None) + output_shapes = tuple(grad.shape for grad in backward_outputs) + + captures = list(sorted(c_captured_tensors, key=lambda x: x.name)) forward_name = _forward_name(self._func_name) self._forward_fdef = _EagerDefinedFunction( forward_name, self._graph, self._ops, self._input_placeholders, filtered_outputs + captures) - backward_outputs = tuple(x for x in in_gradients if x is not None) all_inputs = self._out_grad_placeholders + captures # Excluding input ops from the body as we do not intend to execute these # operations when the function is executed. @@ -376,19 +426,26 @@ class GraphModeFunction(object): # means rerunning the function-defining code will always define the same # function, which is useful if we serialize this etc. function_def_ops = tuple(x - for x in sorted(c.known_ops, key=lambda x: x.name) + for x in sorted(c_known_ops, key=lambda x: x.name) if x not in all_ignored_ops) bname = _backward_name(self._func_name) self._backward_function = GraphModeFunction( bname, all_inputs, [], self._graph, function_def_ops, - backward_outputs, in_gradients, shapes) + backward_outputs, in_gradients, output_shapes) def _backprop_call(self, args): """Calls the wrapped function and records the result on a tape.""" all_args = args + self._extra_inputs signature = self._forward_fdef.signature ctx = context.context() - if ctx.in_graph_mode(): + if ctx.executing_eagerly(): + outputs = execute.execute( + str(signature.name), + num_outputs=len(signature.output_arg), + inputs=all_args, + attrs=None, + ctx=ctx) + else: g = ops.get_default_graph() g._add_function(self._forward_fdef) # pylint: disable=protected-access op = g.create_op( @@ -403,13 +460,6 @@ class GraphModeFunction(object): outputs, (ops.Tensor, type(None))) else list(outputs) for i, s in enumerate(self._output_shapes): outputs[i].set_shape(s) - else: - outputs = execute.execute( - str(signature.name), - num_outputs=len(signature.output_arg), - inputs=all_args, - attrs=None, - ctx=ctx) real_outputs = outputs[:len(self._returns)] side_outputs = outputs[len(self._returns):] @@ -426,9 +476,24 @@ class GraphModeFunction(object): @property def output_shapes(self): + """The function's output shapes.""" # TODO(ebrevdo): Should we only keep the output shapes associated # with len(self._returns) outputs? - return nest.pack_sequence_as(self._func_outputs, self._output_shapes) + outputs_list = nest.flatten(self._func_outputs) + j = 0 + for i, o in enumerate(outputs_list): + if o is not None: + if isinstance(o, ops.IndexedSlices): + # Extract the shape of the `IndexedSlices` object's `values` field. + outputs_list[i] = self._output_shapes[j] # the `values` shape + if o.dense_shape is not None: + j += 3 # skip over shapes for `values`, `indices`, `dense_shape` + else: + j += 2 # skip over shapes for `values`, `indices` + else: + outputs_list[i] = self._output_shapes[j] + j += 1 + return nest.pack_sequence_as(self._func_outputs, outputs_list) @property def output_dtypes(self): @@ -457,16 +522,22 @@ class GraphModeFunction(object): if v._trainable: # pylint: disable=protected-access tape.watch_variable(v) - tensor_inputs = [x for x in nest.flatten(args) - if isinstance(x, ops.Tensor)] + tensor_inputs = [x for x in nest.flatten(args) if isinstance(x, ops.Tensor)] if tape.should_record(tensor_inputs) or tape.should_record( self._extra_inputs): - if not self._has_backprop: - self._compute_backprop() + if self._backward_function is None: + self._construct_backprop_function() return self._backprop_call(tensor_inputs) ctx = context.context() - if ctx.in_graph_mode(): + if ctx.executing_eagerly(): + result = execute.execute( + str(self._func_name), + num_outputs=self._num_outputs, + inputs=tensor_inputs + self._extra_inputs, + attrs=None, + ctx=ctx) + else: g = ops.get_default_graph() self.add_to_graph(g) signature = self._function_def.definition.signature @@ -483,13 +554,6 @@ class GraphModeFunction(object): return op for i, s in enumerate(self._output_shapes): result[i].set_shape(s) - else: - result = execute.execute( - str(self._func_name), - num_outputs=self._num_outputs, - inputs=tensor_inputs + self._extra_inputs, - attrs=None, - ctx=ctx) return self._build_call_outputs(result) @@ -503,13 +567,30 @@ class GraphModeFunction(object): """ if self._func_outputs is None: return None + # Use `nest.flatten` instead of `_flatten` in order to preserve any + # IndexedSlices in `self._func_outputs`. outputs_list = nest.flatten(self._func_outputs) j = 0 for i, o in enumerate(outputs_list): if o is not None: - outputs_list[i] = result[j] - j += 1 - return nest.pack_sequence_as(self._func_outputs, outputs_list) + if isinstance(o, ops.IndexedSlices): + # Repack Tensors for IndexedSlices. + if o.dense_shape is not None: + outputs_list[i] = ops.IndexedSlices( + values=result[j], + indices=result[j + 1], + dense_shape=result[j + 2]) + j += 3 + else: + outputs_list[i] = ops.IndexedSlices( + values=result[j], + indices=result[j + 1]) + j += 2 + else: + outputs_list[i] = result[j] + j += 1 + ret = nest.pack_sequence_as(self._func_outputs, outputs_list) + return ret def _get_defun_inputs(args): @@ -526,15 +607,13 @@ def _get_defun_inputs(args): def _defun_internal(name, func, args, kwds): """Defines and returns graph-mode version of func.""" - container_prefix = ops.get_default_graph()._container_prefix # pylint: disable=protected-access + graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access with context.graph_mode(): captures = {} tmp_graph = CapturingGraph(captures) - # Inherit the container prefix, since this is used for error checking when - # isolating eager execution (the container prefix at creation must match the - # container prefix when used, and variables accessed in the defun will be - # used in the outside context). - tmp_graph._container_prefix = container_prefix # pylint: disable=protected-access + # Inherit the graph key, since this is used for matching variables in + # optimizers. + tmp_graph._graph_key = graph_key # pylint: disable=protected-access # Copy the graph collections to ensure summaries and other things work. This # lets the function access (but not mutate) collections of the containing # graph, such as the global step and the summary writer collections. @@ -542,20 +621,28 @@ def _defun_internal(name, func, args, kwds): for collection in curr_graph.collections: tmp_graph.get_collection_ref(collection)[:] = curr_graph.get_collection( collection) - with tmp_graph.as_default(): + with tmp_graph.as_default(), AutomaticControlDependencies() as a: func_inputs = _get_defun_inputs(args) + def convert(x): + if x is None: + return None + x = ops.convert_to_tensor_or_indexed_slices(x) + x = a.mark_as_return(x) + return x + with capture_tensors(captures): this_tape = tape.push_new_tape() try: func_outputs = func(*func_inputs, **kwds) + func_outputs = nest.map_structure(convert, func_outputs) finally: tape.pop_tape(this_tape) variables = this_tape.watched_variables() # Returning a closed-over tensor as an output does not trigger a # call to convert_to_tensor, so we manually capture all such tensors. - outputs_list = nest.flatten(func_outputs) + outputs_list = _flatten(func_outputs) func_def_outputs = [ _convert_to_graph_tensor(x) for x in outputs_list if x is not None ] @@ -579,7 +666,7 @@ def _defun_internal(name, func, args, kwds): if x not in all_ignored_ops) # Register any other functions defined in the graph # TODO(ashankar): Oh lord, forgive me for this lint travesty. - if context.in_eager_mode(): + if context.executing_eagerly(): for f in tmp_graph._functions.values(): # pylint: disable=protected-access # TODO(ashankar): What about the gradient registry? _register(f._c_func) # pylint: disable=protected-access @@ -600,6 +687,18 @@ def _cache_key(x): """Cache key for tfe functions.""" if isinstance(x, ops.Tensor): return _TensorDtype(x.dtype, x._shape_tuple()) # pylint: disable=protected-access + if isinstance(x, ops.IndexedSlices): + if x.dense_shape is not None: + return tuple([ + _TensorDtype(x.values.dtype, x.values._shape_tuple()), # pylint: disable=protected-access + _TensorDtype(x.indices.dtype, x.indices._shape_tuple()), # pylint: disable=protected-access + _TensorDtype(x.dense_shape.dtype, x.dense_shape._shape_tuple()) # pylint: disable=protected-access + ]) + else: + return tuple([ + _TensorDtype(x.values.dtype, x.values._shape_tuple()), # pylint: disable=protected-access + _TensorDtype(x.indices.dtype, x.indices._shape_tuple()) # pylint: disable=protected-access + ]) if isinstance(x, np.ndarray): return ("array", x.shape, tuple(x.reshape(-1))) if isinstance(x, (list, tuple)): @@ -697,7 +796,11 @@ def defun(func): or more Tensor objects). """ # TODO(apassos): deal with captured global state. Deal with control flow. - return tf_decorator.make_decorator(func, named_defun(func, func.__name__)) + try: + name = func.__name__ + except AttributeError: + name = "function" + return tf_decorator.make_decorator(func, named_defun(func, name)) def make_defun_op(func, *args, **kwds): @@ -750,3 +853,238 @@ def make_defun_op(func, *args, **kwds): if any(isinstance(x, ops.EagerTensor) for x in kwds.values()): raise ValueError("Tensor keyword arguments are not supported.") return _defun_internal(name, func, args, kwds) + + +class AutomaticControlDependencies(object): + """Context manager to automatically add control dependencies. + + Code under this context manager will act as if a sensible set of control + dependencies were present. More specifically: + 1. All stateful ops in the scope will execute + 2. Stateful ops which modify the same resource will execute in program order + + Note: creating variables in an automatic control dependencies context is not + supported (the value of the variables will never change as they will keep + getting reinitialized). + + NOT THREAD SAFE + """ + + def __init__(self): + self._returned_tensors = set() + + def mark_as_return(self, tensor): + """Acts like identity but marks the `Tensor` as a return value. + + This will possibly return a copy of the `Tensor`. Usage: + + ``` + with AutomaticControlDependencies() as a: + ... + t = a.mark_as_return(t) + _ = ...(t...) # i.e. it's safe to use t here + ``` + + Args: + tensor: the `Tensor` to be marked + + Returns: + a copy of the `Tensor`. + """ + if isinstance(tensor, ops.IndexedSlices): + values = array_ops.identity(tensor.values) + indices = array_ops.identity(tensor.indices) + self._returned_tensors.add(indices) + self._returned_tensors.add(values) + return ops.IndexedSlices(values, indices, dense_shape=tensor.dense_shape) + # We want to make the return values depend on the stateful operations, but + # we don't want to introduce a cycle, so we make the return value the result + # of a new identity operation that the stateful operations definitely don't + # depend on. + tensor = array_ops.identity(tensor) + self._returned_tensors.add(tensor) + return tensor + + def __enter__(self): + if context.executing_eagerly(): + return self + # This code assumes no other thread is adding ops to the graph while + # we're adding ops to the graph. + # TODO(apassos): Fix this by locking the graph or using a temporary + # graph (but that would mess up devices and collections at least, + # probably other things as well). + self._graph = ops.get_default_graph() + self._n_operations = len(self._graph.get_operations()) + return self + + def _process_switch(self, switch_op, ops_which_must_run, + last_op_using_resource_tensor, merge_for_resource): + """Processes a switch node for a resource input. + + When tensorflow creates a cond, it creates a control flow context for each + branch of the cond. Each external tensor accessed by that branch is routed + through a switch op, which gets created in the graph _after_ the op which + uses that tensor get created. + + If the resource comes from another switch op we process that one first. + + _process_switch creates a corresponding merge node for the switch node. This + merge node is added to the outer control flow context of the switch + node. We also ensure that: + + 1. The switch node executes after the previous op which used the resource + tensor + + 2. Any op which uses a resource output of the switch node executes before + the merge for the switch node. + + 3. The next op which uses the input resource to the switch node (which + might be another switch node for the other branch of the conditional) + will execute after the merge node is done. + + 4. The merge node is marked as must_run so it will run even if no + subsequent operation uses the resource. + + Args: + switch_op: the switch op to be processed + ops_which_must_run: the set of ops which must run + last_op_using_resource_tensor: map from resource tensor to last op using + it + merge_for_resource: map from resource tensor to merge which must follow + all usages of it. + """ + inp = switch_op.inputs[0] + if inp.dtype == dtypes_module.resource and inp.op.type == "Switch": + self._process_switch(inp.op, ops_which_must_run, + last_op_using_resource_tensor, merge_for_resource) + if switch_op.outputs[0] in merge_for_resource: + return + new_merge = control_flow_ops.merge(switch_op.outputs, + name="artificial_merge") + new_merge[0].op._control_flow_context = ( # pylint: disable=protected-access + switch_op._control_flow_context.outer_context) # pylint: disable=protected-access + # Ensures the merge always runs + ops_which_must_run.add(new_merge[0].op) + if inp in last_op_using_resource_tensor: + # Ensures the switch exectutes after the previous op using the resource. + switch_op._add_control_input(last_op_using_resource_tensor[inp]) # pylint: disable=protected-access + # Ensure the next op outside the cond happens after the merge. + last_op_using_resource_tensor[inp] = new_merge[0].op + if inp in merge_for_resource: + merge_for_resource[inp]._add_control_input(new_merge[0].op) # pylint: disable=protected-access + for o in switch_op.outputs: + # Ensures the merge will execute after all ops inside the cond + merge_for_resource[o] = new_merge[0].op + + def __exit__(self, unused_type, unused_value, unused_traceback): + if context.executing_eagerly(): + return + + if self._graph is not ops.get_default_graph(): + raise RuntimeError( + "Graph changed while trying to add control dependencies.") + + # map from resource tensor to the last op which used it + last_op_using_resource_tensor = {} + # set of conditional and loop exits + ops_which_must_run = set() + # merge which must depend on ops which use this resource + merge_for_resource = {} + + new_operations = self._graph.get_operations()[self._n_operations:] + + # Ensures that uses of resource tensors get serialized properly and all + # execute. This is done by keeping a map from resource tensor to the last op + # in graph-construction order which used it (last_op_using_resource_tensor). + # + # Conditionals are written in TensorFlow such that every external tensor + # accessed in the conditional goes through a switch op and every return + # tensor (it's guaranteed that there will be at least one) goes through a + # merge op. + # + # To handle conditionals, switches are handled in a special way (see + # comments for _process_switch). Merge nodes created by TF's conditional + # logic (as opposed to by _process_switch) are forced to run and also get a + # control dependency added to them to ensure all stateful ops inside their + # control flow context run. + # + # We also ensure that if an op is using a resource output by a switch node + # (that is, a resource tensor for which there's a value in + # merge_for_resource) this op will run before the merge for that resource. + # + # We try to add control inputs to nodes respecting their control flow + # contexts to avoid dead nodes propagating everywhere and leading to + # "retval[0] doesn't have value" errors. If a node gets a control dependency + # on a dead node (i.e. a note from an untaken control flow branch) that node + # will be marked as dead unless it's a merge node. + # + # TODO(apassos): serialize non-resource-taking stateful ops as well, and + # test that it works. Support while loops. Support init_scope escaping from + # this. + for op in new_operations: + control_inputs = set() + # Ensure stateful ops run + if (op.type not in self._graph._registered_ops # pylint: disable=protected-access + or self._graph._registered_ops[op.type].is_stateful): # pylint: disable=protected-access + ops_which_must_run.add(op) + # Ignore switches (they're handled separately) + if op.type == "Switch" and op.inputs[0].dtype == dtypes_module.resource: + continue + # Make merges trigger all other computation which must run + if op.type == "Merge": + for o in ops_which_must_run: + op._add_control_input(o) # pylint: disable=protected-access + for inp in o.inputs: + if inp in last_op_using_resource_tensor: + last_op_using_resource_tensor[inp] = op + ops_which_must_run = set([op]) + continue + for inp in op.inputs: + if inp.dtype == dtypes_module.resource: + # Deal with switches, finally. + if inp.op.type == "Switch": + self._process_switch(inp.op, ops_which_must_run, + last_op_using_resource_tensor, + merge_for_resource) + # Ensure uses of resources are serialized + if inp in last_op_using_resource_tensor: + if (last_op_using_resource_tensor[inp]._control_flow_context # pylint: disable=protected-access + is op._control_flow_context): # pylint: disable=protected-access + control_inputs.add(last_op_using_resource_tensor[inp]) + # Ensure merges happen after the closing of a cond block + if inp in merge_for_resource: + merge_for_resource[inp]._add_control_input(op) # pylint: disable=protected-access + last_op_using_resource_tensor[inp] = op + control_inputs = [c for c in control_inputs + if c._control_flow_context is op._control_flow_context] # pylint: disable=protected-access + op._add_control_inputs(control_inputs) # pylint: disable=protected-access + + # Ensure all ops which must run do run + for r in self._returned_tensors: + if ops_which_must_run: + r.op._add_control_inputs( # pylint: disable=protected-access + [o for o in ops_which_must_run + if o._control_flow_context is r.op._control_flow_context]) # pylint: disable=protected-access + + +def automatic_control_dependencies(f): + """Wraps f to automatically insert control dependencies. + + The inserted dependencies ensure that: + 1. All stateful ops in f run when the result of f runs + 2. Updates to the same resources happen in order. + + Args: + f: the function to be wrapped. + + Returns: + The wrapped function. + """ + + def wrapper(*args, **kwds): + with AutomaticControlDependencies() as a: + result = f(*args, **kwds) + result_flat = [a.mark_as_return(t) for t in nest.flatten(result)] + return nest.pack_sequence_as(result, result_flat) + + return tf_decorator.make_decorator(f, wrapper) diff --git a/tensorflow/python/eager/function_test.py b/tensorflow/python/eager/function_test.py index 2cb2cfb76c7c1713c66a060fc227331bb1acaf71..b9cde16867d498d73715535f028a5eb2bea97ea6 100644 --- a/tensorflow/python/eager/function_test.py +++ b/tensorflow/python/eager/function_test.py @@ -32,6 +32,7 @@ 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 clip_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variable_scope @@ -374,6 +375,78 @@ class FunctionTest(test.TestCase): self.assertAllEqual(f(constant_op.constant(1.0)), 2.0) + def testGradientOfGatherWithDefun(self): + + v = resource_variable_ops.ResourceVariable([0.0, 1.0, 2.0]) + + def sum_gather(): + return math_ops.reduce_sum(array_ops.gather(v, [1, 2])) + + grad_fn = backprop.implicit_grad(sum_gather) + gradient = grad_fn() + defun_grad_fn = backprop.implicit_grad(function.defun(sum_gather)) + defun_gradient = defun_grad_fn() + self.assertEqual(len(gradient), len(defun_gradient)) + + gradient = gradient[0][0] + defun_gradient = defun_gradient[0][0] + self.assertAllEqual(gradient.values, defun_gradient.values) + self.assertAllEqual(gradient.indices, defun_gradient.indices) + self.assertAllEqual(gradient.dense_shape, defun_gradient.dense_shape) + + def testReturningIndexedSlicesWithDefun(self): + + def validate(indexed_slice): + def f(): + return indexed_slice + + output = function.defun(f)() + self.assertTrue(isinstance(output, ops.IndexedSlices)) + self.assertAllEqual(indexed_slice.values, output.values) + self.assertAllEqual(indexed_slice.indices, output.indices) + self.assertAllEqual(indexed_slice.dense_shape, output.dense_shape) + + self.assertEqual( + function.make_defun_op(f).output_shapes, indexed_slice.values.shape) + + arg = ops.IndexedSlices( + values=constant_op.constant([1, 2]), + indices=constant_op.constant([0, 1]), + dense_shape=constant_op.constant([2])) + validate(arg) + + arg = ops.IndexedSlices( + values=constant_op.constant([1, 2]), + indices=constant_op.constant([0, 1]), + dense_shape=None) + validate(arg) + + def testIndexedSliceAsArgumentWithDefun(self): + + @function.defun + def f(indexed_slice): + return indexed_slice + + def validate(arg): + output = f(arg) + self.assertTrue(isinstance(output, ops.IndexedSlices)) + self.assertAllEqual(arg.values, output.values) + self.assertAllEqual(arg.indices, output.indices) + self.assertAllEqual(arg.dense_shape, output.dense_shape) + + indexed_slice = ops.IndexedSlices( + values=constant_op.constant([1]), + indices=constant_op.constant([0]), + dense_shape=constant_op.constant([1])) + validate(indexed_slice) + + # Test that `f` works even when `dense_shape` is None. + indexed_slice = ops.IndexedSlices( + values=constant_op.constant([1]), + indices=constant_op.constant([0]), + dense_shape=None) + validate(indexed_slice) + def testFunctionOnDevice(self): if not context.context().num_gpus(): self.skipTest('No GPUs found') @@ -523,5 +596,172 @@ class FunctionTest(test.TestCase): create_variable() +class AutomaticControlDependenciesTest(test.TestCase): + + def testBasic(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + with function.AutomaticControlDependencies() as c: + v.assign(v + 1) + v.assign(2 * v) + val = v.read_value() + val = c.mark_as_return(val) + self.assertAllEqual(val.eval(), 4.0) + + def testCondMustRun(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + p = array_ops.placeholder(dtype=dtypes.bool) + with function.AutomaticControlDependencies() as c: + + def true_fn(): + v.assign(v + 1) + return 0.0 + + def false_fn(): + v.assign(v + 4) + return 1.0 + + control_flow_ops.cond(p, true_fn, false_fn) + val = v.read_value() + val = c.mark_as_return(val) + self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0) + self.assertAllEqual(val.eval(feed_dict={p: True}), 6.0) + + def testCondMustRunSeparateRead(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + p = array_ops.placeholder(dtype=dtypes.bool) + with function.AutomaticControlDependencies() as c: + + def true_fn(): + v.assign(v + 1) + return 0.0 + + def false_fn(): + v.assign(v + 4) + return 1.0 + + control_flow_ops.cond(p, true_fn, false_fn) + one = constant_op.constant(1.0) + one = c.mark_as_return(one) + one.eval(feed_dict={p: False}) + self.assertAllEqual(v.read_value().eval(), 5.0) + one.eval(feed_dict={p: True}) + self.assertAllEqual(v.read_value().eval(), 6.0) + + def testCondNested(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + p = array_ops.placeholder(dtype=dtypes.bool) + q = array_ops.placeholder(dtype=dtypes.bool) + with function.AutomaticControlDependencies() as c: + + def true_fn(): + v.assign(v + 1, name='true') + return 1.0 + + def false_fn(): + + def inner_true_fn(): + v.assign(v * 2, name='false_true') + return 2.0 + + def inner_false_fn(): + v.assign(v * 3, name='false_false') + return 3.0 + + control_flow_ops.cond(q, inner_true_fn, inner_false_fn) + return 1.0 + + control_flow_ops.cond(p, true_fn, false_fn) + with ops.name_scope('final'): + val = v.read_value() + val = c.mark_as_return(val) + self.assertAllEqual(val.eval(feed_dict={p: False, q: False}), 3.0) + self.assertAllEqual(val.eval(feed_dict={p: False, q: True}), 6.0) + self.assertAllEqual(val.eval(feed_dict={p: True, q: True}), 7.0) + self.assertAllEqual(val.eval(feed_dict={p: True, q: False}), 8.0) + + def testCondOneBranch(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + p = array_ops.placeholder(dtype=dtypes.bool) + with function.AutomaticControlDependencies() as c: + + def true_fn(): + return 0.0 + + def false_fn(): + v.assign(v + 4) + return 1.0 + + control_flow_ops.cond(p, true_fn, false_fn) + val = v.read_value() + val = c.mark_as_return(val) + self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0) + self.assertAllEqual(val.eval(feed_dict={p: True}), 5.0) + + def testCondOneBranchUpdateBefore(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + p = array_ops.placeholder(dtype=dtypes.bool) + with function.AutomaticControlDependencies() as c: + v.assign(v * 2) + + def true_fn(): + return 0.0 + + def false_fn(): + v.assign(v + 4) + return 1.0 + + control_flow_ops.cond(p, true_fn, false_fn) + val = v.read_value() + val = c.mark_as_return(val) + self.assertAllEqual(val.eval(feed_dict={p: False}), 6.0) + self.assertAllEqual(val.eval(feed_dict={p: True}), 12.0) + + def testCondOneBranchUpdateAfter(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + p = array_ops.placeholder(dtype=dtypes.bool) + with function.AutomaticControlDependencies() as c: + + def true_fn(): + return 0.0 + + def false_fn(): + v.assign(v + 4) + return 1.0 + + control_flow_ops.cond(p, true_fn, false_fn) + v.assign(v * 2) + val = v.read_value() + val = c.mark_as_return(val) + self.assertAllEqual(val.eval(feed_dict={p: False}), 10.0) + self.assertAllEqual(val.eval(feed_dict={p: True}), 20.0) + + def testDecorator(self): + with context.graph_mode(), self.test_session(): + v = resource_variable_ops.ResourceVariable(1.0) + variables.global_variables_initializer().run() + + @function.automatic_control_dependencies + def f(): + v.assign(v + 1) + v.assign(2 * v) + return v.read_value() + + self.assertAllEqual(f().eval(), 4.0) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/eager/gen_op.bzl b/tensorflow/python/eager/gen_op.bzl deleted file mode 100644 index 8bc1d6c10a60b89a026cb34dbf6fd98d29e909c2..0000000000000000000000000000000000000000 --- a/tensorflow/python/eager/gen_op.bzl +++ /dev/null @@ -1,65 +0,0 @@ -"""For eager-mode Python.""" - -load("//tensorflow:tensorflow.bzl", - "clean_dep", - "tf_binary_additional_srcs", - "tf_copts", - "tf_cc_binary") - -def tfe_gen_op_wrapper_py(name, - out=None, - visibility=None, - deps=[], - generated_target_name=None, - # ApiDefs will be loaded in the order specified in this list. - api_def_srcs=[]): - """Generate an eager-mode Python op wrapper for an op library.""" - # Construct a cc_binary containing the specified ops. - tool_name = "gen_" + name + "_py_wrappers_cc" - if not deps: - deps = [str(Label("//tensorflow/core:" + name + "_op_lib"))] - tf_cc_binary( - name=tool_name, - linkopts=["-lm"], - copts=tf_copts(), - linkstatic=1, - deps=([ - clean_dep("//tensorflow/python/eager:python_eager_op_gen_main") - ] + deps), - visibility=[clean_dep("//visibility:public")],) - - # Invoke the previous cc_binary to generate a python file. - if not out: - out = "gen_" + name + ".py" - - if not api_def_srcs: - api_def_args_str = "," - else: - api_def_args = [] - for api_def_src in api_def_srcs: - # Add directory of the first ApiDef source to args. - # We are assuming all ApiDefs in a single api_def_src are in the - # same directory. - api_def_args.append( - "$$(dirname $$(echo $(locations " + api_def_src + - ") | cut -d\" \" -f1))") - api_def_args_str = ",".join(api_def_args) - - native.genrule( - name=name + "_pygenrule", - outs=[out], - srcs=api_def_srcs, - tools=[tool_name] + tf_binary_additional_srcs(), - cmd=("$(location " + tool_name + ") " + api_def_args_str + " > $@")) - - # Make a py_library out of the generated python file. - if not generated_target_name: - generated_target_name = name - native.py_library( - name=generated_target_name, - srcs=[out], - srcs_version="PY2AND3", - visibility=visibility, - deps=[ - clean_dep("//tensorflow/python/eager:framework_for_generated_wrappers"), - ],) diff --git a/tensorflow/python/eager/graph_callable.py b/tensorflow/python/eager/graph_callable.py index 5c13ea89081a7d060c0ed1201f0169b739a204c2..ee5d87f0835a8e70e0ce14537a51ea5418db41b9 100644 --- a/tensorflow/python/eager/graph_callable.py +++ b/tensorflow/python/eager/graph_callable.py @@ -252,21 +252,17 @@ def _graph_callable_internal(func, shape_and_dtypes): Callable graph object. """ container = tf_ops.get_default_graph()._container # pylint: disable=protected-access - container_prefix = tf_ops.get_default_graph()._container_prefix # pylint: disable=protected-access + graph_key = tf_ops.get_default_graph()._graph_key # pylint: disable=protected-access with context.graph_mode(): # This graph will store both the initialization and the call version of the # wrapped function. It will later be used by the backprop code to build the # backprop graph, if necessary. captures = {} tmp_graph = function.CapturingGraph(captures) - # Inherit the container from the original graph to create resources at user - # expected containers. Also inherits the container prefix, since this is - # used for error checking when isolating Eager execution (the container - # prefix at creation must match the container prefix when used, and - # variables returned from the graph callable will be used in the outside - # context). + # Inherit the graph key from the original graph to ensure optimizers don't + # misbehave. tmp_graph._container = container # pylint: disable=protected-access - tmp_graph._container_prefix = container_prefix # pylint: disable=protected-access + tmp_graph._graph_key = graph_key # pylint: disable=protected-access with tmp_graph.as_default(): # Placeholders for the non-variable inputs. func_inputs = _get_graph_callable_inputs(shape_and_dtypes) @@ -283,9 +279,12 @@ def _graph_callable_internal(func, shape_and_dtypes): # scope's view of which variables exist. variable_captures = _VariableCapturingScope() with variable_captures.initializing_scope(), function.capture_tensors( - captures): + captures), function.AutomaticControlDependencies() as a: func_outputs = func(*func_inputs) - outputs_list = nest.flatten(func_outputs) + outputs_list = nest.flatten(func_outputs) + for i, x in enumerate(outputs_list): + if x is not None: + outputs_list[i] = a.mark_as_return(x) if len(outputs_list) == 1 and outputs_list[0] is None: outputs_list = [] output_shapes = [x.shape for x in outputs_list] @@ -298,9 +297,12 @@ def _graph_callable_internal(func, shape_and_dtypes): # knows about all variables. tmp_graph.clear_resource_control_flow_state() with variable_captures.capturing_scope(), function.capture_tensors( - captures): + captures), function.AutomaticControlDependencies() as a: captured_outputs = func(*func_inputs) captured_outlist = nest.flatten(captured_outputs) + for i, x in enumerate(captured_outlist): + if x is not None: + captured_outlist[i] = a.mark_as_return(x) capturing_operations = tmp_graph.get_operations()[ len(initializing_operations):] @@ -404,7 +406,7 @@ def graph_callable(shape_and_dtypes): A callable graph object. """ # TODO(alive,apassos): support initialized_value and friends from tf.Variable. - assert context.in_eager_mode(), ( + assert context.executing_eagerly(), ( "graph_callable can only be used when Eager execution is enabled.") def decorator(func): return tf_decorator.make_decorator(func, diff --git a/tensorflow/python/eager/ops_test.py b/tensorflow/python/eager/ops_test.py index f2e70341d975fb06bce7f2ce6cba7d8c3bc9826c..fc76ede4c502ae8b554c925a921e419bf003c40c 100644 --- a/tensorflow/python/eager/ops_test.py +++ b/tensorflow/python/eager/ops_test.py @@ -17,8 +17,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import threading import numpy as np +from tensorflow.core.protobuf import config_pb2 from tensorflow.python.eager import context from tensorflow.python.eager import execute from tensorflow.python.eager import test @@ -130,8 +132,12 @@ class OpsTest(test_util.TensorFlowTestCase): dtype=dtypes.int64) values = constant_op.constant([2, 3, 5, 7, 11]) shape = constant_op.constant([2, 7], dtype=dtypes.int64) - result = sparse_ops.gen_sparse_ops._sparse_split( # pylint: disable=protected-access - split_dim, indices, values, shape, num_split=2) + result = sparse_ops.gen_sparse_ops.sparse_split( + split_dim, + indices, + values, + shape, + num_split=2) output_indices, output_values, output_shape = result self.assertEqual(2, len(output_indices)) self.assertEqual(2, len(output_values)) @@ -277,6 +283,25 @@ class OpsTest(test_util.TensorFlowTestCase): context._context = context.Context() # pylint: enable=protected-access + def testSoftPlacement(self): + if not context.context().num_gpus(): + self.skipTest('No GPUs found') + # Temporarily replace the context + # pylint: disable=protected-access + del context._context + try: + context._context = context.Context( + device_policy=context.DEVICE_PLACEMENT_SILENT, + config=config_pb2.ConfigProto(allow_soft_placement=True)) + cpu_tensor = constant_op.constant(1.0) + result = cpu_tensor + cpu_tensor + self.assertEqual(result.device, + '/job:localhost/replica:0/task:0/device:GPU:0') + finally: + del context._context + context._context = context.Context() + # pylint: enable=protected-access + def testRandomUniform(self): scalar_shape = constant_op.constant([], dtype=dtypes.int32) @@ -352,6 +377,22 @@ class OpsTest(test_util.TensorFlowTestCase): def testNoOpIsNone(self): self.assertTrue(control_flow_ops.no_op() is None) + def testEagerContextPreservedAcrossThreads(self): + def init_fn(): + self.assertTrue(context.executing_eagerly()) + with ops.init_scope(): + self.assertTrue(context.executing_eagerly()) + context_switches = context.context().context_switches + self.assertEqual(len(context_switches.stack), 1) + self.assertFalse(context_switches.stack[0].is_building_function) + self.assertEqual(context_switches.stack[0].enter_context_fn, + context.eager_mode) + + self.assertTrue(context.executing_eagerly()) + t1 = threading.Thread(target=init_fn) + t1.start() + t1.join() + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/eager/python_eager_op_gen.cc b/tensorflow/python/eager/python_eager_op_gen.cc index 90a8779ff845b2fd63d1ba1019e8601fef257e42..c2ce8efd7f70c6ba93b6d444f88ddbb9aa51ccdb 100644 --- a/tensorflow/python/eager/python_eager_op_gen.cc +++ b/tensorflow/python/eager/python_eager_op_gen.cc @@ -42,6 +42,8 @@ namespace { const int kRightMargin = 78; +constexpr char kEagerFallbackSuffix[] = "_eager_fallback"; + string AttrVarName(const string& attr_name, std::unordered_map* attr_expressions) { const string var = strings::StrCat("_attr_", attr_name); @@ -49,11 +51,12 @@ string AttrVarName(const string& attr_name, return var; } -void AddInferredAttr(const string& attr_name, const string& value_expression, - string* result, +void AddInferredAttr(const string& indentation, const string& attr_name, + const string& value_expression, string* result, std::unordered_map* attr_expressions) { - strings::StrAppend(result, " ", AttrVarName(attr_name, attr_expressions), - " = ", value_expression, "\n"); + strings::StrAppend(result, indentation, + AttrVarName(attr_name, attr_expressions), " = ", + value_expression, "\n"); } string VectorToTuple(const std::vector& l) { @@ -121,11 +124,33 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp { string Code() override; protected: - void ExpectListArg(const string& arg_name); - void AddEagerInferredAttrs(); - void AddEagerInputCasts(); - void AddEagerAttrs(); - void AddEagerExecute(const string& num_outputs_expr); + void HandleGraphMode(const string& function_setup); + + string GetEagerNotAllowedError(); + void ExpectListArg(const string& indentation, const string& arg_name, + string* output); + bool GetEagerFunctionSetup(const string& indentation, string* function_setup); + void GetOutputSizesAndNumOutputsExpr(std::vector* output_sizes, + string* num_outputs_expr); + + void AddEagerFunctionTeardown(const string& indentation, + const std::vector& output_sizes, + bool execute_record_gradient); + + bool AddEagerFastPathAndGraphCode(const string& parameters, + const std::vector& output_sizes, + const string& eager_not_allowed_error); + bool AddEagerFallbackCode(const string& parameters, + const std::vector& output_sizes, + const string& num_outputs_expr, + const string& eager_not_allowed_error); + void AddEagerFastPathExecute(); + + void AddEagerInferredAttrs(const string& indentation); + void AddEagerInputCasts(const string& indentation); + void AddEagerAttrs(const string& indentation); + void AddEagerExecute(const string& indentation, + const string& num_outputs_expr); void AddAttrForArg(const string& attr, int arg_index) { gtl::InsertIfNotPresent(&inferred_attrs_, attr, @@ -148,6 +173,13 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp { typedef std::unordered_map> AttrToArgMap; AttrToArgMap attr_to_args_; std::unordered_map attr_expressions_; + // This has all the input args followed by those attrs that don't have + // defaults. + std::vector params_no_default_; + // The parameters with defaults (these have to be listed after those without). + // No input args are included, just attrs. + std::vector> + params_with_default_; }; string GetEagerPythonOp(const OpDef& op_def, const ApiDef& api_def, @@ -207,18 +239,12 @@ string GenEagerPythonOp::Code() { if (api_def_.visibility() == ApiDef::SKIP) { return ""; } - // This has all the input args followed by those attrs that don't have - // defaults. - std::vector params_no_default; - // The parameters with defaults (these have to be listed after those without). - // No input args are included, just attrs. - std::vector> - params_with_default; for (int i = 0; i < api_def_.arg_order_size(); ++i) { const auto& arg = *FindInputArg(api_def_.arg_order(i), op_def_); const auto& api_def_arg = *FindInputArg(api_def_.arg_order(i), api_def_); - params_no_default.emplace_back(api_def_arg.name(), api_def_arg.rename_to()); + params_no_default_.emplace_back(api_def_arg.name(), + api_def_arg.rename_to()); if (!arg.type_attr().empty()) { AddAttrForArg(arg.type_attr(), i); } else if (!arg.type_list_attr().empty()) { @@ -235,7 +261,7 @@ string GenEagerPythonOp::Code() { if (inferred_attrs_.find(attr.name()) == inferred_attrs_.end()) { if (api_def_attr.has_default_value()) { if (attr.type() == "tensor") { - params_with_default.emplace_back( + params_with_default_.emplace_back( python_op_gen_internal::ParamNames(api_def_attr.name(), api_def_attr.rename_to()), strings::StrCat( @@ -247,22 +273,22 @@ string GenEagerPythonOp::Code() { for (const auto& pb : api_def_attr.default_value().list().tensor()) { pbtxt.emplace_back(TensorPBString(pb)); } - params_with_default.emplace_back( + params_with_default_.emplace_back( python_op_gen_internal::ParamNames(api_def_attr.name(), api_def_attr.rename_to()), strings::StrCat("[_execute.make_tensor(_pb, \"", api_def_attr.rename_to(), "\") for _pb in ", VectorToTuple(pbtxt), "]")); } else { - params_with_default.emplace_back( + params_with_default_.emplace_back( python_op_gen_internal::ParamNames(api_def_attr.name(), api_def_attr.rename_to()), python_op_gen_internal::AttrValueToPython( attr.type(), api_def_attr.default_value(), "_dtypes.")); } } else { - params_no_default.emplace_back(api_def_attr.name(), - api_def_attr.rename_to()); + params_no_default_.emplace_back(api_def_attr.name(), + api_def_attr.rename_to()); } } } @@ -270,29 +296,29 @@ string GenEagerPythonOp::Code() { // Save the list of attr parameters (attrs that won't be inferred), // those with defaults go at the end. // Get the attrs in the order we want by taking the attrs without defaults - // from the end of params_no_default, and adding params_no_default. - attrs_.reserve(params_no_default.size() - op_def_.input_arg_size() + - params_with_default.size()); - for (int i = op_def_.input_arg_size(); i < params_no_default.size(); ++i) { - attrs_.push_back(params_no_default[i].GetName()); + // from the end of params_no_default_, and adding params_no_default_. + attrs_.reserve(params_no_default_.size() - op_def_.input_arg_size() + + params_with_default_.size()); + for (int i = op_def_.input_arg_size(); i < params_no_default_.size(); ++i) { + attrs_.push_back(params_no_default_[i].GetName()); } - for (const auto& p : params_with_default) { + for (const auto& p : params_with_default_) { attrs_.push_back(p.first.GetName()); } - param_names_.reserve(params_no_default.size() + params_with_default.size()); - param_names_.insert(param_names_.begin(), params_no_default.begin(), - params_no_default.end()); - for (const auto& param_and_default : params_with_default) { + param_names_.reserve(params_no_default_.size() + params_with_default_.size()); + param_names_.insert(param_names_.begin(), params_no_default_.begin(), + params_no_default_.end()); + for (const auto& param_and_default : params_with_default_) { param_names_.push_back(param_and_default.first); } string parameters; - for (const auto& param : params_no_default) { + for (const auto& param : params_no_default_) { if (!parameters.empty()) strings::StrAppend(¶meters, ", "); strings::StrAppend(¶meters, param.GetRenameTo()); } - for (const auto& param_and_default : params_with_default) { + for (const auto& param_and_default : params_with_default_) { if (!parameters.empty()) strings::StrAppend(¶meters, ", "); strings::StrAppend(¶meters, param_and_default.first.GetRenameTo(), "=", param_and_default.second); @@ -300,19 +326,125 @@ string GenEagerPythonOp::Code() { if (!parameters.empty()) strings::StrAppend(¶meters, ", "); strings::StrAppend(¶meters, "name=None"); - AddExport(); - AddDefLine(parameters); - AddDocStringDescription(); - AddDocStringArgs(); - AddDocStringInputs(); - AddDocStringAttrs(); - AddDocStringNameArg(); - AddOutputGlobals(); - AddDocStringOutputs(); - strings::StrAppend(&result_, " \"\"\"\n"); + // Add attr_expressions_ for attrs that are params. + for (int i = 0; i < attrs_.size(); ++i) { + const string& attr_name = attrs_[i]; + const string& attr_api_name = + param_names_[i + op_def_.input_arg_size()].GetRenameTo(); + attr_expressions_[attr_name] = attr_api_name; + } + // Add attr_expressions_ for attrs that are inferred. + for (int i = 0; i < op_def_.attr_size(); ++i) { + const auto& attr(op_def_.attr(i)); + if (attr.type() == "int") { + auto arg_list = attr_to_args_.find(attr.name()); + if (arg_list != attr_to_args_.end()) { + AttrVarName(attr.name(), &attr_expressions_); + } + } + } + + string num_outputs_expr; + std::vector output_sizes(num_outs_); + GetOutputSizesAndNumOutputsExpr(&output_sizes, &num_outputs_expr); + + string eager_not_allowed_error = GetEagerNotAllowedError(); - // Function body. + if (!AddEagerFastPathAndGraphCode(parameters, output_sizes, + eager_not_allowed_error)) { + return result_; + } + if (!AddEagerFallbackCode(parameters, output_sizes, num_outputs_expr, + eager_not_allowed_error)) { + return result_; + } + + return prelude_ + result_; +} + +void GenEagerPythonOp::HandleGraphMode(const string& function_setup) { + // Handle graph-mode case + strings::StrAppend(&result_, + " _ctx = _context.context()\n" + " if not _ctx.executing_eagerly():\n", + function_setup, + " _, _, _op = _op_def_lib._apply_op_helper(\n"); + AddBodyNoReturn(" "); + if (num_outs_ > 0) { + strings::StrAppend(&result_, " _result = _op.outputs[:]\n"); + // Special case handling for stateful op with single list output + // that might be empty. + if (num_outs_ == 1 && op_def_.is_stateful() && + (!op_def_.output_arg(0).number_attr().empty() || + !op_def_.output_arg(0).type_list_attr().empty())) { + // TODO(josh11b): Can skip this if the number_attr/type_list_attr has + // a constraint indicating that this can never be empty. + strings::StrAppend(&result_, + " if not _result:\n" + " return _op\n"); + } + strings::StrAppend(&result_, " _inputs_flat = _op.inputs\n"); + + // Compute graph-mode attrs. + if (op_def_.attr_size() > 0) { + string attr_values; + for (int i = 0; i < op_def_.attr_size(); ++i) { + if (i > 0) strings::StrAppend(&attr_values, ", "); + const auto& attr_name(op_def_.attr(i).name()); + strings::StrAppend(&attr_values, "\"", attr_name, "\", _op.get_attr(\"", + attr_name, "\")"); + } + strings::StrAppend(&attr_values, ")"); + strings::StrAppend(&result_, + WordWrap(" _attrs = (", attr_values, kRightMargin), + "\n"); + } else { + strings::StrAppend(&result_, " _attrs = None\n"); + } + } else { + strings::StrAppend(&result_, " return _op\n"); + } +} + +string GenEagerPythonOp::GetEagerNotAllowedError() { + bool eager_allowed = true; + string ref_arg; + for (int i = 0; i < op_def_.input_arg_size(); ++i) { + const auto& arg = op_def_.input_arg(i); + if (arg.is_ref()) { + eager_allowed = false; + DCHECK_EQ(op_def_.input_arg(i).name(), api_def_.in_arg(i).name()); + ref_arg = api_def_.in_arg(i).rename_to(); + } + } + for (int i = 0; i < op_def_.output_arg_size(); ++i) { + const auto& arg = op_def_.output_arg(i); + if (arg.is_ref()) { + eager_allowed = false; + DCHECK_EQ(op_def_.output_arg(i).name(), api_def_.out_arg(i).name()); + ref_arg = api_def_.out_arg(i).rename_to(); + } + } + + if (eager_allowed) return ""; + + return strings::StrCat("raise RuntimeError(\"", op_name_, + " op does not support eager execution. ", "Arg '", + ref_arg, "' is a ref.\")\n"); +} + +void GenEagerPythonOp::ExpectListArg(const string& indentation, + const string& arg_name, string* output) { + strings::StrAppend(output, indentation, "if not isinstance(", arg_name, + ", (list, tuple)):\n", indentation, " raise TypeError(\n", + indentation, " \"Expected list for '", arg_name, + "' argument to \"\n", indentation, " \"'", op_name_, + "' Op, not %r.\" % ", arg_name, ")\n"); +} + +bool GenEagerPythonOp::GetEagerFunctionSetup(const string& indentation, + string* function_setup) { // Validate list inputs, infer length attrs. for (int i = 0; i < op_def_.attr_size(); ++i) { const auto& attr(op_def_.attr(i)); @@ -324,32 +456,27 @@ string GenEagerPythonOp::Code() { for (auto iter = arg_list->second.begin(); iter != arg_list->second.end(); ++iter) { const string& arg_api_name = param_names_[*iter].GetRenameTo(); - ExpectListArg(arg_api_name); + ExpectListArg(indentation, arg_api_name, function_setup); if (iter == arg_list->second.begin()) { - AddInferredAttr(attr.name(), + AddInferredAttr(indentation, attr.name(), strings::StrCat("len(", arg_api_name, ")"), - &result_, &attr_expressions_); + function_setup, &attr_expressions_); } else { const auto& attr_var = attr_expressions_[attr.name()]; - strings::StrAppend(&result_, " if len(", arg_api_name, - ") != ", attr_var, - ":\n" - " raise ValueError(\n" - " \"List argument '", - arg_api_name, "' to '", op_name_, - "' Op with length %d \"\n" - " \"must match length %d of argument '", - inferred_attrs_[attr.name()], - "'.\" %\n" - " (len(", - arg_api_name, "), ", attr_var, "))\n"); + strings::StrAppend( + function_setup, indentation, "if len(", arg_api_name, + ") != ", attr_var, ":\n", indentation, " raise ValueError(\n", + indentation, " \"List argument '", arg_api_name, "' to '", + op_name_, "' Op with length %d \"\n", indentation, + " \"must match length %d of argument '", + inferred_attrs_[attr.name()], "'.\" %\n", indentation, + " (len(", arg_api_name, "), ", attr_var, "))\n"); } } } } } - // Values for non-inferred attrs. for (int i = 0; i < attrs_.size(); ++i) { const string& attr_name = attrs_[i]; const auto& param = param_names_[i + op_def_.input_arg_size()]; @@ -357,241 +484,300 @@ string GenEagerPythonOp::Code() { const string& attr_api_name = param.GetRenameTo(); StringPiece attr_type = attr.type(); attr_expressions_[attr_name] = attr_api_name; - const int default_index = i - (attrs_.size() - params_with_default.size()); + const int default_index = i - (attrs_.size() - params_with_default_.size()); if (default_index >= 0) { - const string& default_value = params_with_default[default_index].second; - strings::StrAppend(&result_, " if ", attr_api_name, " is None:\n"); - strings::StrAppend(&result_, " ", attr_api_name, " = ", default_value, - "\n"); + const string& default_value = params_with_default_[default_index].second; + strings::StrAppend(function_setup, indentation, "if ", attr_api_name, + " is None:\n"); + strings::StrAppend(function_setup, indentation, " ", attr_api_name, + " = ", default_value, "\n"); } if (attr_type.starts_with("list(")) { - ExpectListArg(attr_api_name); + ExpectListArg(indentation, attr_api_name, function_setup); } if (attr_type == "string") { - strings::StrAppend(&result_, " ", attr_api_name, " = _execute.make_str(", - attr_api_name, ", \"", attr_api_name, "\")\n"); + strings::StrAppend(function_setup, indentation, attr_api_name, + " = _execute.make_str(", attr_api_name, ", \"", + attr_api_name, "\")\n"); } else if (attr_type == "list(string)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_str(_s, \"", attr_api_name, "\") for _s in ", attr_api_name, "]\n"); } else if (attr_type == "int") { - strings::StrAppend(&result_, " ", attr_api_name, " = _execute.make_int(", - attr_api_name, ", \"", attr_api_name, "\")\n"); + strings::StrAppend(function_setup, indentation, attr_api_name, + " = _execute.make_int(", attr_api_name, ", \"", + attr_api_name, "\")\n"); } else if (attr_type == "list(int)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_int(_i, \"", attr_api_name, "\") for _i in ", attr_api_name, "]\n"); } else if (attr_type == "float") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = _execute.make_float(", attr_api_name, ", \"", attr_api_name, "\")\n"); } else if (attr_type == "list(float)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_float(_f, \"", attr_api_name, "\") for _f in ", attr_api_name, "]\n"); } else if (attr_type == "bool") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = _execute.make_bool(", attr_api_name, ", \"", attr_api_name, "\")\n"); } else if (attr_type == "list(bool)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_bool(_b, \"", attr_api_name, "\") for _b in ", attr_api_name, "]\n"); } else if (attr_type == "type") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = _execute.make_type(", attr_api_name, ", \"", attr_api_name, "\")\n"); } else if (attr_type == "list(type)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_type(_t, \"", attr_api_name, "\") for _t in ", attr_api_name, "]\n"); } else if (attr_type == "shape") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = _execute.make_shape(", attr_api_name, ", \"", attr_api_name, "\")\n"); } else if (attr_type == "list(shape)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_shape(_s, \"", attr_api_name, "\") for _s in ", attr_api_name, "]\n"); } else if (attr_type == "tensor") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = _execute.make_tensor(", attr_api_name, ", \"", attr_api_name, "\")\n"); } else if (attr_type == "list(tensor)") { - strings::StrAppend(&result_, " ", attr_api_name, + strings::StrAppend(function_setup, indentation, attr_api_name, " = [_execute.make_tensor(_t, \"", attr_api_name, "\") for _t in ", attr_api_name, "]\n"); } else if (attr_type != "func") { - return strings::StrCat("# No definition for ", function_name_, - " since we don't support attrs with type\n" - "# '", - attr_type, "' right now.\n\n"); - } - } - - // Figure out the list of inputs. - const string inputs = FlattenInputs(nullptr, nullptr); - - // Handle graph-mode case - strings::StrAppend(&result_, - " _ctx = _context.context()\n" - - " if _ctx.in_graph_mode():\n" - " _, _, _op = _op_def_lib._apply_op_helper(\n"); - AddBodyNoReturn(" "); - if (num_outs_ > 0) { - strings::StrAppend(&result_, " _result = _op.outputs[:]\n"); - // Special case handling for stateful op with single list output - // that might be empty. - if (num_outs_ == 1 && op_def_.is_stateful() && - (!op_def_.output_arg(0).number_attr().empty() || - !op_def_.output_arg(0).type_list_attr().empty())) { - // TODO(josh11b): Can skip this if the number_attr/type_list_attr has - // a constraint indicating that this can never be empty. - strings::StrAppend(&result_, - " if not _result:\n" - " return _op\n"); - } - strings::StrAppend(&result_, " _inputs_flat = _op.inputs\n"); - - // Compute graph-mode attrs. - if (op_def_.attr_size() > 0) { - string attr_values; - for (int i = 0; i < op_def_.attr_size(); ++i) { - if (i > 0) strings::StrAppend(&attr_values, ", "); - const auto& attr_name(op_def_.attr(i).name()); - strings::StrAppend(&attr_values, "\"", attr_name, "\", _op.get_attr(\"", - attr_name, "\")"); - } - strings::StrAppend(&attr_values, ")"); - strings::StrAppend(&result_, - WordWrap(" _attrs = (", attr_values, kRightMargin), - "\n"); - } else { - strings::StrAppend(&result_, " _attrs = None\n"); + *function_setup = + strings::StrCat("# No definition for ", function_name_, + " since we don't support attrs with type\n" + "# '", + attr_type, "' right now.\n\n"); + return false; } - } else { - strings::StrAppend(&result_, " return _op\n"); } + return true; +} - // Handle eager-mode case - strings::StrAppend(&result_, " else:\n"); - +// If output i is list output, output_sizes[i] will be set to a +// string with the python expression that will evaluate to its +// length. output_sizes[i] is empty for non-list outputs. +void GenEagerPythonOp::GetOutputSizesAndNumOutputsExpr( + std::vector* output_sizes, string* num_outputs_expr) { // Expression representing the number of outputs. int num_fixed_outputs = 0; - string num_outputs_expr; - // If output i is list output, output_sizes[i] will be set to a - // string with the python expression that will evaluate to its - // length. output_sizes[i] is empty for non-list outputs. - std::vector output_sizes(num_outs_); for (int i = 0; i < num_outs_; ++i) { const auto& arg(op_def_.output_arg(i)); if (!arg.number_attr().empty()) { - if (!num_outputs_expr.empty()) { - strings::StrAppend(&num_outputs_expr, " + "); + if (!num_outputs_expr->empty()) { + strings::StrAppend(num_outputs_expr, " + "); } - output_sizes[i] = attr_expressions_[arg.number_attr()]; - strings::StrAppend(&num_outputs_expr, output_sizes[i]); + (*output_sizes)[i] = attr_expressions_[arg.number_attr()]; + strings::StrAppend(num_outputs_expr, (*output_sizes)[i]); } else if (!arg.type_list_attr().empty()) { - if (!num_outputs_expr.empty()) { - strings::StrAppend(&num_outputs_expr, " + "); + if (!num_outputs_expr->empty()) { + strings::StrAppend(num_outputs_expr, " + "); } // Have to be careful to use an expression that works in both // graph and eager paths here. const auto iter = inferred_attrs_.find(arg.type_list_attr()); if (iter == inferred_attrs_.end()) { - output_sizes[i] = strings::StrCat( + (*output_sizes)[i] = strings::StrCat( "len(", attr_expressions_[arg.type_list_attr()], ")"); } else { - output_sizes[i] = strings::StrCat("len(", iter->second, ")"); + (*output_sizes)[i] = strings::StrCat("len(", iter->second, ")"); } - strings::StrAppend(&num_outputs_expr, output_sizes[i]); + strings::StrAppend(num_outputs_expr, (*output_sizes)[i]); } else { ++num_fixed_outputs; } } if (num_fixed_outputs > 0) { - if (!num_outputs_expr.empty()) { - strings::StrAppend(&num_outputs_expr, " + "); + if (!num_outputs_expr->empty()) { + strings::StrAppend(num_outputs_expr, " + "); } - strings::StrAppend(&num_outputs_expr, num_fixed_outputs); - } else if (num_outputs_expr.empty()) { - num_outputs_expr = "0"; - } - - bool eager_allowed = true; - string ref_arg; - for (int i = 0; i < op_def_.input_arg_size(); ++i) { - const auto& arg = op_def_.input_arg(i); - if (arg.is_ref()) { - eager_allowed = false; - DCHECK_EQ(op_def_.input_arg(i).name(), api_def_.in_arg(i).name()); - ref_arg = api_def_.in_arg(i).rename_to(); - } - } - for (int i = 0; i < op_def_.output_arg_size(); ++i) { - const auto& arg = op_def_.output_arg(i); - if (arg.is_ref()) { - eager_allowed = false; - DCHECK_EQ(op_def_.output_arg(i).name(), api_def_.out_arg(i).name()); - ref_arg = api_def_.out_arg(i).rename_to(); - } - } - - if (eager_allowed) { - AddEagerInferredAttrs(); - AddEagerInputCasts(); - strings::StrAppend(&result_, " _inputs_flat = ", inputs, "\n"); - AddEagerAttrs(); - AddEagerExecute(num_outputs_expr); - } else { - strings::StrAppend(&result_, - " raise RuntimeError(\n" - " \"", - op_name_, " op does not support eager execution. ", - "Arg '", ref_arg, "'' is a ref.\")\n"); + strings::StrAppend(num_outputs_expr, num_fixed_outputs); + } else if (num_outputs_expr->empty()) { + *num_outputs_expr = "0"; } +} +void GenEagerPythonOp::AddEagerFunctionTeardown( + const string& indentation, const std::vector& output_sizes, + bool execute_record_gradient) { if (num_outs_ > 0) { - strings::StrAppend(&result_, " _execute.record_gradient(\n", " \"", - op_def_.name(), - "\", _inputs_flat, _attrs, _result, name)\n"); + if (execute_record_gradient) { + strings::StrAppend(&result_, indentation, "_execute.record_gradient(\n", + " \"", op_def_.name(), + "\", _inputs_flat, _attrs, _result, name)\n"); + } if (num_outs_ == 1 && !output_sizes[0].empty()) { // Single list result. } else if (num_outs_ == 1) { // Execute returns a single-element list which we need to destructure. - strings::StrAppend(&result_, " _result, = _result\n"); + strings::StrAppend(&result_, indentation, "_result, = _result\n"); } else { // Have multiple outputs, so we will need to reformat the return // value of execute() to be a list with one entry per op output // (that entry will be a list of tensors if that output is of list // type). // For list outputs, convert the right subrange of _result into a list. - Unflatten(" ", output_sizes, "_result", &result_); + Unflatten(indentation, output_sizes, "_result", &result_); // Convert to a named tuple. - strings::StrAppend(&result_, " _result = _", op_def_.name(), + strings::StrAppend(&result_, indentation, "_result = _", op_def_.name(), "Output._make(_result)\n"); } } else { - strings::StrAppend(&result_, " _result = None\n"); + strings::StrAppend(&result_, indentation, "_result = None\n"); } - strings::StrAppend(&result_, " return _result\n\n"); - return prelude_ + result_; + strings::StrAppend(&result_, indentation, "return _result\n\n"); } -void GenEagerPythonOp::ExpectListArg(const string& arg_name) { - strings::StrAppend(&result_, " if not isinstance(", arg_name, - ", (list, tuple)):\n" - " raise TypeError(\n" - " \"Expected list for '", - arg_name, - "' argument to \"\n" - " \"'", - op_name_, "' Op, not %r.\" % ", arg_name, ")\n"); +bool GenEagerPythonOp::AddEagerFastPathAndGraphCode( + const string& parameters, const std::vector& output_sizes, + const string& eager_not_allowed_error) { + AddExport(); + AddDefLine(function_name_, parameters); + AddDocStringDescription(); + AddDocStringArgs(); + AddDocStringInputs(); + AddDocStringAttrs(); + AddDocStringNameArg(); + AddOutputGlobals(); // Added to prelude_ + AddDocStringOutputs(); + strings::StrAppend(&result_, " \"\"\"\n"); + + // Handle graph-mode case + string function_setup; + if (!GetEagerFunctionSetup(" ", &function_setup)) { + result_ = function_setup; + return false; + } + HandleGraphMode(function_setup); + AddEagerFunctionTeardown(" ", output_sizes, + true /* execute_record_gradient */); + + // Handle eager-mode case + strings::StrAppend(&result_, " else:\n"); + + if (eager_not_allowed_error.empty()) { + AddEagerFastPathExecute(); + } else { + strings::StrAppend(&result_, " ", eager_not_allowed_error); + } + + strings::StrAppend(&result_, "\n\n"); + return true; } -void GenEagerPythonOp::AddEagerInferredAttrs() { +bool GenEagerPythonOp::AddEagerFallbackCode( + const string& parameters, const std::vector& output_sizes, + const string& num_outputs_expr, const string& eager_not_allowed_error) { + if (!eager_not_allowed_error.empty()) { + strings::StrAppend(&result_, " ", eager_not_allowed_error); + return true; + } + + AddDefLine(strings::StrCat(function_name_, kEagerFallbackSuffix), parameters); + strings::StrAppend( + &result_, " r\"\"\"This is the slowpath function for Eager mode.\n"); + strings::StrAppend(&result_, " This is for function ", function_name_, + "\n \"\"\"\n"); + + strings::StrAppend(&result_, " _ctx = _context.context()\n"); + + string function_setup; + if (!GetEagerFunctionSetup(" ", &function_setup)) { + result_ = function_setup; + return false; + } + strings::StrAppend(&result_, function_setup); + + AddEagerInferredAttrs(" "); + AddEagerInputCasts(" "); + strings::StrAppend( + &result_, " _inputs_flat = ", FlattenInputs(nullptr, nullptr), "\n"); + AddEagerAttrs(" "); + AddEagerExecute(" ", num_outputs_expr); + + AddEagerFunctionTeardown(" ", output_sizes, + true /* execute_record_gradient */); + + return true; +} + +void GenEagerPythonOp::AddEagerFastPathExecute() { + string fastpath_execute_params = + strings::StrCat("_ctx._handle, _ctx.device_name, \"", op_def_.name(), + "\", ", "name, _ctx._post_execution_callbacks"); + string fallback_params; + + for (int i = 0; i < api_def_.in_arg_size(); i++) { + const string param_name = param_names_[i].GetRenameTo(); + strings::StrAppend(&fastpath_execute_params, ", ", param_name); + if (!fallback_params.empty()) strings::StrAppend(&fallback_params, ", "); + strings::StrAppend(&fallback_params, param_name); + } + + for (const auto& attr : api_def_.attr()) { + if (inferred_attrs_.find(attr.name()) == inferred_attrs_.end()) { + strings::StrAppend(&fastpath_execute_params, ", \"", attr.name(), "\", ", + attr.rename_to()); + + if (!fallback_params.empty()) strings::StrAppend(&fallback_params, ", "); + strings::StrAppend(&fallback_params, attr.rename_to(), "=", + attr.rename_to()); + } + } + + if (!fallback_params.empty()) strings::StrAppend(&fallback_params, ", "); + strings::StrAppend(&fallback_params, "name=name"); + + strings::StrAppend(&result_, " try:\n"); + strings::StrAppend( + &result_, " ", + "_result = _pywrap_tensorflow.TFE_Py_FastPathExecute(\n", + WordWrap(strings::StrCat(" "), + strings::StrCat(fastpath_execute_params, ")"), kRightMargin), + "\n"); + + if (op_def_.output_arg_size() > 1) { + const string output_tuple_name = + strings::StrCat("_", op_def_.name(), "Output"); + strings::StrAppend(&result_, " ", "_result = ", output_tuple_name, + "._make(_result)\n"); + } + strings::StrAppend(&result_, " ", "return _result\n"); + + // Handle fallback. + strings::StrAppend(&result_, " ", "except _core._FallbackException:\n"); + strings::StrAppend( + &result_, " ", "return ", function_name_, kEagerFallbackSuffix, + "(\n", + WordWrap(strings::StrCat(" "), + strings::StrCat(fallback_params, ")"), kRightMargin), + "\n"); + + // Any errors thrown from execute need to be unwrapped from + // _NotOkStatusException. + strings::StrAppend(&result_, " ", + "except _core._NotOkStatusException as e:\n"); + strings::StrAppend(&result_, " ", "if name is not None:\n"); + strings::StrAppend(&result_, " ", + "message = e.message + \" name: \" + name\n"); + strings::StrAppend(&result_, " ", "else:\n"); + strings::StrAppend(&result_, " ", "message = e.message\n"); + strings::StrAppend( + &result_, " ", + "_six.raise_from(_core._status_to_exception(e.code, message), None)\n"); +} + +void GenEagerPythonOp::AddEagerInferredAttrs(const string& indentation) { // Figure out values for inferred attrs, and cast to eager tensors. for (int i = 0; i < op_def_.attr_size(); ++i) { const auto& attr(op_def_.attr(i)); @@ -618,24 +804,24 @@ void GenEagerPythonOp::AddEagerInferredAttrs() { const string inputs_var = param_names_[arg_list->second.front()].GetRenameTo(); if (output_sizes.front().empty()) { - strings::StrAppend(&result_, " ", var_name, ", (", inputs_var, - ",) = ", conversion, "\n"); + strings::StrAppend(&result_, indentation, var_name, ", (", + inputs_var, ",) = ", conversion, "\n"); } else { - strings::StrAppend(&result_, " ", var_name, ", ", inputs_var, - " = ", conversion, "\n"); + strings::StrAppend(&result_, indentation, var_name, ", ", + inputs_var, " = ", conversion, "\n"); } } else { const string inputs_var = strings::StrCat("_inputs_", attr.name()); - strings::StrAppend(&result_, " ", var_name, ", ", inputs_var, + strings::StrAppend(&result_, indentation, var_name, ", ", inputs_var, " = ", conversion, "\n"); // Convert from a flat list of eager tensors back to the // parameter variables. - Unflatten(" ", output_sizes, inputs_var, &result_); + Unflatten(indentation, output_sizes, inputs_var, &result_); std::vector p; for (int j : arg_list->second) { p.emplace_back(param_names_[j].GetRenameTo()); } - strings::StrAppend(&result_, " ", VectorToTuple(p), " = ", + strings::StrAppend(&result_, indentation, VectorToTuple(p), " = ", inputs_var, "\n"); } } else if (attr.type() == "list(type)") { @@ -662,14 +848,14 @@ void GenEagerPythonOp::AddEagerInferredAttrs() { inputs_var = param_names_[arg_list->second.front()].GetRenameTo(); conversion = "_execute.convert_to_mixed_eager_tensors"; } - strings::StrAppend(&result_, " ", var_name, ", ", inputs_var, " = ", - conversion, "(", inputs_var, ", _ctx)\n"); + strings::StrAppend(&result_, indentation, var_name, ", ", inputs_var, + " = ", conversion, "(", inputs_var, ", _ctx)\n"); } } } } -void GenEagerPythonOp::AddEagerInputCasts() { +void GenEagerPythonOp::AddEagerInputCasts(const string& indentation) { // Cast remaining args to eager tensors for (int i = 0; i < op_def_.input_arg_size(); ++i) { const auto& arg(op_def_.input_arg(i)); @@ -678,12 +864,12 @@ void GenEagerPythonOp::AddEagerInputCasts() { const string fn = arg.number_attr().empty() ? "" : "n_"; const string dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes."); - strings::StrAppend(&result_, " ", param, " = _ops.convert_", fn, + strings::StrAppend(&result_, indentation, param, " = _ops.convert_", fn, "to_tensor(", param, ", ", dtype, ")\n"); } } -void GenEagerPythonOp::AddEagerAttrs() { +void GenEagerPythonOp::AddEagerAttrs(const string& indentation) { // Compute eager attrs if (op_def_.attr_size() > 0) { string attr_values; @@ -695,14 +881,19 @@ void GenEagerPythonOp::AddEagerAttrs() { } strings::StrAppend(&attr_values, ")"); strings::StrAppend( - &result_, WordWrap(" _attrs = (", attr_values, kRightMargin), "\n"); + &result_, + WordWrap(indentation, strings::StrCat("_attrs = (", attr_values), + kRightMargin), + "\n"); } else { - strings::StrAppend(&result_, " _attrs = None\n"); + strings::StrAppend(&result_, indentation, "_attrs = None\n"); } } -void GenEagerPythonOp::AddEagerExecute(const string& num_outputs_expr) { - const string return_prefix = " _result = _execute.execute("; +void GenEagerPythonOp::AddEagerExecute(const string& indentation, + const string& num_outputs_expr) { + const string return_prefix = + strings::StrCat(indentation, "_result = _execute.execute("); const string return_args = strings::StrCat( "b\"", op_def_.name(), "\", ", num_outputs_expr, ", inputs=_inputs_flat, attrs=_attrs, ctx=_ctx, name=name)"); @@ -723,8 +914,8 @@ string GetEagerPythonOps(const OpList& ops, const ApiDefMap& api_defs, This file is MACHINE GENERATED! Do not edit. )"); - // Mention the original source file so someone tracing back through generated - // Python code will know where to look next. + // Mention the original source file so someone tracing back through + // generated Python code will know where to look next. if (!source_file_name.empty()) { strings::StrAppend(&result, "Original C++ source file: "); strings::StrAppend(&result, source_file_name); @@ -734,11 +925,14 @@ This file is MACHINE GENERATED! Do not edit. strings::StrAppend(&result, R"(""" import collections as _collections +import six as _six -from tensorflow.python.eager import execute as _execute +from tensorflow.python import pywrap_tensorflow as _pywrap_tensorflow from tensorflow.python.eager import context as _context from tensorflow.python.eager import core as _core +from tensorflow.python.eager import execute as _execute from tensorflow.python.framework import dtypes as _dtypes +from tensorflow.python.framework import errors as _errors from tensorflow.python.framework import tensor_shape as _tensor_shape from tensorflow.core.framework import op_def_pb2 as _op_def_pb2 @@ -756,28 +950,46 @@ from tensorflow.python.util.tf_export import tf_export auto out = cleaned_ops.mutable_op(); out->Reserve(ops.op_size()); for (const auto& op_def : ops.op()) { - bool is_hidden = false; - for (const string& hidden : hidden_ops) { - if (op_def.name() == hidden) { - is_hidden = true; - break; + const auto* api_def = api_defs.GetApiDef(op_def.name()); + + if (api_def->visibility() == ApiDef::SKIP) { + continue; + } + // An op is hidden if either its ApiDef visibility is HIDDEN + // or it is in the hidden_ops list. + bool is_hidden = api_def->visibility() == ApiDef::HIDDEN; + bool hidden_by_api_def = is_hidden; + if (!is_hidden) { + for (const string& hidden : hidden_ops) { + if (op_def.name() == hidden) { + is_hidden = true; + break; + } } } string function_name; python_op_gen_internal::GenerateLowerCaseOpName(op_def.name(), &function_name); - if (is_hidden) function_name = strings::StrCat("_", function_name); - - // When users create custom python wrappers, they may link in the - // default op registry by accident, and because they can't - // enumerate all 'hidden' symbols, this guard is to prevent - // instantiating a python reserved word in their wrapper. - if (python_op_gen_internal::IsPythonReserved(function_name)) { + bool is_reserved = python_op_gen_internal::IsPythonReserved(function_name); + + // Prefix an op with underscore if the op is listed in hidden_ops or + // name is reserved or it is of the exceptions in IsOpWithUnderscorePrefix. + // Do not add underscores to ops set to HIDDEN in ApiDef otherwise. + // TODO(annarev): don't prefix with underscores even if op is in hidden_ops. + if (is_hidden) { + if (!hidden_by_api_def || is_reserved || + python_op_gen_internal::IsOpWithUnderscorePrefix(function_name)) { + function_name = strings::StrCat("_", function_name); + } + } else if (is_reserved) { + // When users create custom python wrappers, they may link in the + // default op registry by accident, and because they can't + // enumerate all 'hidden' symbols, this guard is to prevent + // instantiating a python reserved word in their wrapper. continue; } - const auto* api_def = api_defs.GetApiDef(op_def.name()); strings::StrAppend(&result, GetEagerPythonOp(op_def, *api_def, function_name)); diff --git a/tensorflow/python/eager/python_eager_op_gen_main.cc b/tensorflow/python/eager/python_eager_op_gen_main.cc deleted file mode 100644 index 05351bd8b115ae07482b82166974e86758bc7712..0000000000000000000000000000000000000000 --- a/tensorflow/python/eager/python_eager_op_gen_main.cc +++ /dev/null @@ -1,72 +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/python/eager/python_eager_op_gen.h" - -#include -#include -#include - -#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/lib/io/path.h" -#include "tensorflow/core/lib/strings/str_util.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/init_main.h" - -namespace tensorflow { -namespace { - -void PrintAllPythonOps(const std::vector& hidden_ops, - const std::vector& api_def_dirs) { - OpList ops; - OpRegistry::Global()->Export(false, &ops); - - ApiDefMap api_def_map(ops); - if (!api_def_dirs.empty()) { - Env* env = Env::Default(); - - for (const auto& api_def_dir : api_def_dirs) { - std::vector api_files; - TF_CHECK_OK(env->GetMatchingPaths(io::JoinPath(api_def_dir, "*.pbtxt"), - &api_files)); - TF_CHECK_OK(api_def_map.LoadFileList(env, api_files)); - } - api_def_map.UpdateDocs(); - } - - PrintEagerPythonOps(ops, api_def_map, hidden_ops, true /* require_shapes */); -} - -} // namespace -} // namespace tensorflow - -int main(int argc, char* argv[]) { - tensorflow::port::InitMain(argv[0], &argc, &argv); - - // Usage: - // python_eager_op_gen_main api_def_dir1,api_def_dir2,... - if (argc == 1) { - tensorflow::PrintAllPythonOps({}, {}); - } else if (argc == 2) { - const std::vector api_def_dirs = - tensorflow::str_util::Split(argv[1], ",", - tensorflow::str_util::SkipEmpty()); - tensorflow::PrintAllPythonOps({}, api_def_dirs); - } else { - return -1; - } - return 0; -} diff --git a/tensorflow/python/eager/pywrap_tensor.cc b/tensorflow/python/eager/pywrap_tensor.cc index 6fa076507d11ab9c88891cbeb0a4fb3959e4e99d..105c09e81fefc09bca630534cbe044027dac860b 100644 --- a/tensorflow/python/eager/pywrap_tensor.cc +++ b/tensorflow/python/eager/pywrap_tensor.cc @@ -185,6 +185,16 @@ typedef struct EagerTensor { // This stores `_keras_mask` object and is set by Tensorflow layers. PyObject* keras_mask; + + // This stores `_tensor_shape`, a cached `TensorShape` object, and is set the + // first time that `_EagerTensorBase`'s `shape` property is called. + PyObject* tensor_shape; + + // We store a status object here as an optimization to avoid allocating a new + // Status objects on different functions that operate on EagerTensor and need + // to use a TF_Status object. However note that accesses to `status` are not + // thread-safe. + TF_Status* status; } EagerTensor; // tp_init for EagerTensor. @@ -195,6 +205,9 @@ int EagerTensor_init(EagerTensor* self, PyObject* args, PyObject* kwds) { self->handle_data = Py_None; Py_INCREF(Py_None); self->keras_mask = Py_None; + Py_INCREF(Py_None); + self->tensor_shape = Py_None; + self->status = TF_NewStatus(); PyObject* value; PyObject* context = nullptr; PyObject* device = nullptr; @@ -269,17 +282,17 @@ int EagerTensor_init(EagerTensor* self, PyObject* args, PyObject* kwds) { } TF_DataType handle_dtype = TFE_TensorHandleDataType(handle.get()); if (desired_dtype >= 0 && desired_dtype != handle_dtype) { - auto out_status = tensorflow::make_safe(TF_NewStatus()); handle = tensorflow::make_safe( EagerCast(GetContext(context), handle.get(), handle_dtype, - static_cast(desired_dtype), out_status.get())); - if (TF_GetCode(out_status.get()) != TF_OK) { - PyErr_SetString( - PyExc_ValueError, - tensorflow::strings::StrCat("Error while casting from DataType ", - handle_dtype, " to ", desired_dtype, ". ", - TF_Message(out_status.get())) - .c_str()); + static_cast(desired_dtype), self->status)); + if (TF_GetCode(self->status) != TF_OK) { + PyErr_SetString(PyExc_ValueError, + tensorflow::strings::StrCat( + "Error while casting from DataType ", handle_dtype, + " to ", desired_dtype, ". ", TF_Message(self->status)) + .c_str()); + // Cleanup self->status before returning. + TF_SetStatus(self->status, TF_OK, ""); return -1; } handle_dtype = TFE_TensorHandleDataType(handle.get()); @@ -323,10 +336,14 @@ int EagerTensor_init(EagerTensor* self, PyObject* args, PyObject* kwds) { // tp_dealloc for EagerTensor. void EagerTensor_dealloc(EagerTensor* self) { + TF_DeleteStatus(self->status); Py_DECREF(self->handle_data); Py_DECREF(self->keras_mask); - TFE_DeleteTensorHandle(self->handle); - self->handle = nullptr; + Py_DECREF(self->tensor_shape); + if (self->handle != nullptr) { + TFE_DeleteTensorHandle(self->handle); + self->handle = nullptr; + } // We have the global interpreter lock, so use this chance to perform delayed // refcount decrements. tensorflow::ClearDecrefCache(); @@ -348,12 +365,21 @@ static PyObject* EagerTensor_datatype_enum(EagerTensor* self) { // Getter for `_shape_tuple`. static PyObject* EagerTensor_shape_tuple(EagerTensor* self) { auto handle = self->handle; - int n = TFE_TensorHandleNumDims(handle); + int n = TFE_TensorHandleNumDims(handle, self->status); + if (MaybeRaiseExceptionFromTFStatus(self->status, PyExc_ValueError)) { + // Cleanup self->status before returning. + TF_SetStatus(self->status, TF_OK, ""); + return nullptr; + } PyObject* shape = PyTuple_New(n); if (PyErr_Occurred()) return nullptr; for (int i = 0; i < n; ++i) { - PyObject* dim = PyLong_FromLongLong(TFE_TensorHandleDim(handle, i)); - if (dim == nullptr || PyTuple_SetItem(shape, i, dim) != 0) { + PyObject* dim = + PyLong_FromLongLong(TFE_TensorHandleDim(handle, i, self->status)); + if (MaybeRaiseExceptionFromTFStatus(self->status, PyExc_ValueError) || + dim == nullptr || PyTuple_SetItem(shape, i, dim) != 0) { + // Cleanup self->status before returning. + TF_SetStatus(self->status, TF_OK, ""); Py_DECREF(shape); if (dim != nullptr) Py_DECREF(dim); PyErr_SetString(PyExc_RuntimeError, "Error while creating shape"); @@ -365,10 +391,16 @@ static PyObject* EagerTensor_shape_tuple(EagerTensor* self) { // Getter for `_rank`. static PyObject* EagerTensor_rank(EagerTensor* self) { + int num_dims = TFE_TensorHandleNumDims(self->handle, self->status); + if (MaybeRaiseExceptionFromTFStatus(self->status, PyExc_ValueError)) { + // Cleanup self->status before returning. + TF_SetStatus(self->status, TF_OK, ""); + return nullptr; + } #if PY_MAJOR_VERSION < 3 - return PyInt_FromLong(TFE_TensorHandleNumDims(self->handle)); + return PyInt_FromLong(num_dims); #else - return PyLong_FromLong(TFE_TensorHandleNumDims(self->handle)); + return PyLong_FromLong(num_dims); #endif } @@ -397,6 +429,19 @@ static int EagerTensor_setkeras_mask(EagerTensor* self, PyObject* value, self->keras_mask = value; return 0; } + +static PyObject* EagerTensor_tensor_shape(EagerTensor* self, void* unused) { + Py_INCREF(self->tensor_shape); + return self->tensor_shape; +} + +static int EagerTensor_settensor_shape(EagerTensor* self, PyObject* value, + void* unused) { + Py_DECREF(self->tensor_shape); + Py_INCREF(value); + self->tensor_shape = value; + return 0; +} // Function `_copy_to_device`. static PyObject* EagerTensor_copy_to_device(EagerTensor* self, PyObject* args, PyObject* kwds) { @@ -437,10 +482,16 @@ static PyObject* EagerTensor_numpy(EagerTensor* self) { // Getter `device`. static PyObject* EagerTensor_device(EagerTensor* self) { + const char* device = TFE_TensorHandleDeviceName(self->handle, self->status); + if (MaybeRaiseExceptionFromTFStatus(self->status, PyExc_ValueError)) { + // Cleanup self->status before returning. + TF_SetStatus(self->status, TF_OK, ""); + return nullptr; + } #if PY_MAJOR_VERSION >= 3 - return PyUnicode_FromString(TFE_TensorHandleDeviceName(self->handle)); + return PyUnicode_FromString(device); #else - return PyBytes_FromString(TFE_TensorHandleDeviceName(self->handle)); + return PyBytes_FromString(device); #endif } @@ -455,6 +506,9 @@ static PyGetSetDef EagerTensor_getseters[] = { {const_cast("_keras_mask"), (getter)EagerTensor_keras_mask, (setter)EagerTensor_setkeras_mask, const_cast("_keras_mask"), nullptr}, + {const_cast("_tensor_shape"), (getter)EagerTensor_tensor_shape, + (setter)EagerTensor_settensor_shape, const_cast("_tensor_shape"), + nullptr}, {nullptr} /* Sentinel */ }; @@ -491,16 +545,11 @@ PyTypeObject* EagerTensorType = nullptr; #if PY_MAJOR_VERSION >= 3 static PyType_Slot EagerTensor_Type_slots[] = { - Py_tp_dealloc, - reinterpret_cast(EagerTensor_dealloc), - Py_tp_methods, - reinterpret_cast(EagerTensor_methods), - Py_tp_getset, - reinterpret_cast(EagerTensor_getseters), - Py_tp_init, - reinterpret_cast(EagerTensor_init), - 0, - nullptr, + {Py_tp_dealloc, reinterpret_cast(EagerTensor_dealloc)}, + {Py_tp_methods, reinterpret_cast(EagerTensor_methods)}, + {Py_tp_getset, reinterpret_cast(EagerTensor_getseters)}, + {Py_tp_init, reinterpret_cast(EagerTensor_init)}, + {0, nullptr}, }; PyType_Spec EagerTensor_Type_spec = {"EagerTensor", sizeof(EagerTensor), 0, @@ -575,7 +624,10 @@ PyObject* EagerTensorFromHandle(TFE_TensorHandle* handle) { t->handle_data = Py_None; Py_INCREF(Py_None); t->keras_mask = Py_None; + Py_INCREF(Py_None); + t->tensor_shape = Py_None; t->handle = handle; + t->status = TF_NewStatus(); } return reinterpret_cast(t); } @@ -673,6 +725,7 @@ PyObject* TFE_Py_TensorShapeSlice(PyObject* tensor_list, int slice_dim) { auto tensor = tensorflow::make_safe(TF_AllocateTensor( TF_INT32, &num_tensors_int, /*num_dims=*/1, /*len=*/4 * num_tensors_int)); int32_t* data = reinterpret_cast(TF_TensorData(tensor.get())); + auto status = tensorflow::make_safe(TF_NewStatus()); for (Py_ssize_t i = 0; i < num_tensors; ++i) { PyObject* tensor_obj = PyList_GET_ITEM(tensor_list, i); if (!EagerTensor_CheckExact(tensor_obj)) { @@ -687,21 +740,27 @@ PyObject* TFE_Py_TensorShapeSlice(PyObject* tensor_list, int slice_dim) { EagerTensor* t = reinterpret_cast(tensor_obj); TFE_TensorHandle* handle = t->handle; - if (slice_dim >= TFE_TensorHandleNumDims(handle)) { - PyErr_SetString(PyExc_IndexError, - tensorflow::strings::StrCat( - "Slice dimension (", slice_dim, - ") must be smaller than rank of all " - "tensors, but tensor at index ", - i, " has rank ", TFE_TensorHandleNumDims(handle)) - .c_str()); + int num_dims = TFE_TensorHandleNumDims(handle, status.get()); + if (MaybeRaiseExceptionFromTFStatus(status.get(), PyExc_ValueError)) { + return nullptr; + } + if (slice_dim >= num_dims) { + PyErr_SetString( + PyExc_IndexError, + tensorflow::strings::StrCat("Slice dimension (", slice_dim, + ") must be smaller than rank of all " + "tensors, but tensor at index ", + i, " has rank ", num_dims) + .c_str()); + return nullptr; + } + int64_t dim = TFE_TensorHandleDim(handle, slice_dim, status.get()); + if (MaybeRaiseExceptionFromTFStatus(status.get(), PyExc_ValueError)) { return nullptr; } - int64_t dim = TFE_TensorHandleDim(handle, slice_dim); data[i] = dim; } - auto status = tensorflow::make_safe(TF_NewStatus()); TFE_TensorHandle* handle = TFE_NewTensorHandle(tensor.get(), status.get()); if (TF_GetCode(status.get()) != TF_OK) { PyErr_SetString( diff --git a/tensorflow/python/eager/pywrap_tfe.h b/tensorflow/python/eager/pywrap_tfe.h index 4aea134fa9df845fe2a84f32d56a17a8766bde9b..32d731d0f68910b8e41a57cb32ae60c3ea6742f7 100644 --- a/tensorflow/python/eager/pywrap_tfe.h +++ b/tensorflow/python/eager/pywrap_tfe.h @@ -47,8 +47,34 @@ void TFE_Py_Execute(TFE_Context* ctx, const char* device_name, // Registers e as the Exception class for handling not ok Status. Returns // Py_None if registration succeeds, else throws a TypeError and returns NULL. +// +// This function is not thread-safe. PyObject* TFE_Py_RegisterExceptionClass(PyObject* e); +// Registers e as the type of the ResourceVariable class. +// Returns Py_None if registration succeeds, else throws a TypeError and returns +// NULL. +// +// This function is not thread-safe. +PyObject* TFE_Py_RegisterResourceVariableType(PyObject* e); + +// Registers e as the Exception to be raised when the conditions of +// TFE_Py_FastPathExecute_C have not been met. When this exception is set, it +// is a signal to the calling code that it should fall back to the safer (and +// more complete) code path. +// +// This function is not thread-safe. +PyObject* TFE_Py_RegisterFallbackExceptionClass(PyObject* e); + +// Registers e as the backward_function_getter. +// The registered function creates a backward function (a function that can +// return the gradient of the inputs an op given the gradient of it's outputs). +// The registered function will be passed the following arguments: +// op_name, attrs, num_inputs, op_inputs, op_outputs +// +// This function is not thread-safe. +PyObject* TFE_Py_RegisterBackwardFunctionGetter(PyObject* e); + // Returns 0 if 'status' is TF_OK. Otherwise, raises an exception (using // `exception` if not nullptr, else using the class registered via // TFE_Py_RegisterExceptionClass), and returns -1. @@ -141,11 +167,10 @@ PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* vspace, // Item 2: device_name: Name of the device on which to execute the operation, // or NULL for automatic selection. // Item 3: op_name: Name of the TensorFlow op to execute. -// Item 4: record_gradient_callback: Callback that records the gradient of the -// result. -// The callback takes (inputs, attrs, result) - all sequences and -// records the gradient. -// Item 5 onwards: inputs - This is a list of inputs followed by a list of +// Item 4: name: An optional name for the operation. +// Item 5: List representing all callbacks to execute after successful +// op execute. +// Item 6 onwards: inputs - This is a list of inputs followed by a list of // attrs. It is not necessary for type attrs to be present. // // This is named _C since there doesn't seem to be any way to make it visible @@ -153,6 +178,11 @@ PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* vspace, // directive. PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args); +// Record the gradient for a given op. +PyObject* TFE_Py_RecordGradient(PyObject* op_name, PyObject* inputs, + PyObject* attrs, PyObject* results, + PyObject* name); + // Returns the set of variables watched by the given tape. PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape); diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc index 836998cfdc39ced78732a29b6a11329077239ca5..701f68b8f7a4b4342e7912c7a115f630b4cab627 100644 --- a/tensorflow/python/eager/pywrap_tfe_src.cc +++ b/tensorflow/python/eager/pywrap_tfe_src.cc @@ -23,17 +23,38 @@ limitations under the License. #include "tensorflow/c/eager/tape.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/gtl/compactptrset.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" #include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/protobuf.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/python/eager/pywrap_tensor.h" +#include "tensorflow/python/lib/core/safe_ptr.h" using tensorflow::string; using tensorflow::strings::Printf; namespace { +struct FastPathOpExecInfo { + TFE_Context* ctx; + const char* device_name; + // The op def of the main op being executed. + const tensorflow::OpDef* op_def; + + bool run_callbacks; + bool run_post_exec_callbacks; + bool run_gradient_callback; + + // The op name of the main op being executed. + PyObject* name; + // The op type name of the main op being executed. + PyObject* op_name; + PyObject* callbacks; +}; + #define PARSE_VALUE(fn_name, type, check_fn, parse_fn) \ bool fn_name(const string& key, PyObject* py_value, TF_Status* status, \ type* value) { \ @@ -56,10 +77,50 @@ PARSE_VALUE(ParseInt64Value, int64_t, PyLong_Check, PyLong_AsLong) #else PARSE_VALUE(ParseIntValue, int, PyInt_Check, PyInt_AsLong) PARSE_VALUE(ParseInt64Value, int64_t, PyInt_Check, PyInt_AsLong) +PARSE_VALUE(ParseInt64LongValue, int64_t, PyLong_Check, PyLong_AsLong) #endif PARSE_VALUE(ParseFloatValue, float, PyFloat_Check, PyFloat_AsDouble) #undef PARSE_VALUE +Py_ssize_t TensorShapeNumDims(PyObject* value) { + const auto size = PySequence_Size(value); + if (size == -1) { + // TensorShape.__len__ raises an error in the scenario where the shape is an + // unknown, which needs to be cleared. + // TODO(nareshmodi): ensure that this is actually a TensorShape. + PyErr_Clear(); + } + return size; +} + +bool IsInteger(PyObject* py_value) { +#if PY_MAJOR_VERSION >= 3 + return PyLong_Check(py_value); +#else + return PyInt_Check(py_value); +#endif +} + +bool ParseDimensionValue(const string& key, PyObject* py_value, + TF_Status* status, int64_t* value) { + if (IsInteger(py_value)) { + return ParseInt64Value(key, py_value, status, value); + } + + tensorflow::Safe_PyObjectPtr dimension_value( + PyObject_GetAttrString(py_value, "_value")); + if (dimension_value == nullptr) { + TF_SetStatus( + status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat("Expecting a Dimension for attr ", key, + ", got ", py_value->ob_type->tp_name) + .c_str()); + return false; + } + + return ParseInt64Value(key, dimension_value.get(), status, value); +} + bool ParseStringValue(const string& key, PyObject* py_value, TF_Status* status, const char** value) { if (PyBytes_Check(py_value)) { @@ -86,8 +147,32 @@ bool ParseBoolValue(const string& key, PyObject* py_value, TF_Status* status, return true; } -bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, - TF_AttrType type, TF_Status* status) { +// The passed in py_value is expected to be an object of the python type +// dtypes.DType or an int. +bool ParseTypeValue(const string& key, PyObject* py_value, TF_Status* status, + int* value) { + if (IsInteger(py_value)) { + return ParseIntValue(key, py_value, status, value); + } + + tensorflow::Safe_PyObjectPtr py_type_enum( + PyObject_GetAttrString(py_value, "_type_enum")); + if (py_type_enum == nullptr) { + TF_SetStatus( + status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat("Expecting a DType.dtype for attr ", key, + ", got ", py_value->ob_type->tp_name) + .c_str()); + return false; + } + + return ParseIntValue(key, py_type_enum.get(), status, value); +} + +bool SetOpAttrList( + TFE_Op* op, const char* key, PyObject* py_list, TF_AttrType type, + tensorflow::gtl::FlatMap* attr_list_sizes, + TF_Status* status) { if (!PySequence_Check(py_list)) { TF_SetStatus( status, TF_INVALID_ARGUMENT, @@ -97,12 +182,13 @@ bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, return false; } const int num_values = PySequence_Size(py_list); + if (attr_list_sizes != nullptr) (*attr_list_sizes)[key] = num_values; -#define PARSE_LIST(c_type, parse_fn) \ - std::unique_ptr values(new c_type[num_values]); \ - for (int i = 0; i < num_values; ++i) { \ - auto py_value = PySequence_ITEM(py_list, i); \ - if (!parse_fn(key, py_value, status, &values[i])) return false; \ +#define PARSE_LIST(c_type, parse_fn) \ + std::unique_ptr values(new c_type[num_values]); \ + for (int i = 0; i < num_values; ++i) { \ + tensorflow::Safe_PyObjectPtr py_value(PySequence_ITEM(py_list, i)); \ + if (!parse_fn(key, py_value.get(), status, &values[i])) return false; \ } if (type == TF_ATTR_STRING) { @@ -118,7 +204,7 @@ bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, PARSE_LIST(unsigned char, ParseBoolValue); TFE_OpSetAttrBoolList(op, key, values.get(), num_values); } else if (type == TF_ATTR_TYPE) { - PARSE_LIST(int, ParseIntValue); + PARSE_LIST(int, ParseTypeValue); TFE_OpSetAttrTypeList(op, key, reinterpret_cast(values.get()), num_values); @@ -127,9 +213,9 @@ bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, // dims across all the input lists. int total_dims = 0; for (int i = 0; i < num_values; ++i) { - auto py_value = PySequence_ITEM(py_list, i); - if (py_value != Py_None) { - if (!PySequence_Check(py_value)) { + tensorflow::Safe_PyObjectPtr py_value(PySequence_ITEM(py_list, i)); + if (py_value.get() != Py_None) { + if (!PySequence_Check(py_value.get())) { TF_SetStatus( status, TF_INVALID_ARGUMENT, tensorflow::strings::StrCat( @@ -138,8 +224,10 @@ bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, .c_str()); return false; } - const auto size = PySequence_Size(py_value); - total_dims += size; + const auto size = TensorShapeNumDims(py_value.get()); + if (size >= 0) { + total_dims += size; + } } } // Allocate a buffer that can fit all of the dims together. @@ -150,19 +238,26 @@ bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, std::unique_ptr num_dims(new int[num_values]); int64_t* offset = buffer.get(); for (int i = 0; i < num_values; ++i) { - auto py_value = PySequence_ITEM(py_list, i); - if (py_value == Py_None) { + tensorflow::Safe_PyObjectPtr py_value(PySequence_ITEM(py_list, i)); + if (py_value.get() == Py_None) { dims[i] = nullptr; num_dims[i] = -1; } else { - const auto size = PySequence_Size(py_value); + const auto size = TensorShapeNumDims(py_value.get()); + if (size == -1) { + dims[i] = nullptr; + num_dims[i] = -1; + continue; + } dims[i] = offset; num_dims[i] = size; for (int j = 0; j < size; ++j) { - auto inner_py_value = PySequence_ITEM(py_value, j); - if (inner_py_value == Py_None) { + tensorflow::Safe_PyObjectPtr inner_py_value( + PySequence_ITEM(py_value.get(), j)); + if (inner_py_value.get() == Py_None) { *offset = -1; - } else if (!ParseInt64Value(key, inner_py_value, status, offset)) { + } else if (!ParseDimensionValue(key, inner_py_value.get(), status, + offset)) { return false; } ++offset; @@ -183,8 +278,114 @@ bool SetOpAttrList(TFE_Op* op, const char* key, PyObject* py_list, return true; } -bool SetOpAttrScalar(TFE_Context* ctx, TFE_Op* op, const char* key, - PyObject* py_value, TF_AttrType type, TF_Status* status) { +TFE_Op* GetFunc(TFE_Context* ctx, const tensorflow::NameAttrList& func, + TF_Status* status) { + TFE_Op* func_op = TFE_NewOp(ctx, func.name().data(), status); + for (const auto& attr : func.attr()) { + if (TF_GetCode(status) != TF_OK) return nullptr; + SetOpAttrValueScalar(ctx, func_op, attr.second, attr.first.data(), status); + if (TF_GetCode(status) != TF_OK) return nullptr; + } + return func_op; +} + +void SetOpAttrListDefault( + TFE_Context* ctx, TFE_Op* op, const tensorflow::OpDef::AttrDef& attr, + const char* key, TF_AttrType type, + tensorflow::gtl::FlatMap* attr_list_sizes, + TF_Status* status) { + if (type == TF_ATTR_STRING) { + int num_values = attr.default_value().list().s_size(); + std::unique_ptr values(new const char*[num_values]); + (*attr_list_sizes)[key] = num_values; + for (int i = 0; i < num_values; i++) { + values[i] = attr.default_value().list().s(i).data(); + } + TFE_OpSetAttrStringList(op, key, values.get(), num_values); + } else if (type == TF_ATTR_INT) { + int num_values = attr.default_value().list().i_size(); + std::unique_ptr values(new int64_t[num_values]); + (*attr_list_sizes)[key] = num_values; + for (int i = 0; i < num_values; i++) { + values[i] = attr.default_value().list().i(i); + } + TFE_OpSetAttrIntList(op, key, values.get(), num_values); + } else if (type == TF_ATTR_FLOAT) { + int num_values = attr.default_value().list().f_size(); + std::unique_ptr values(new float[num_values]); + (*attr_list_sizes)[key] = num_values; + for (int i = 0; i < num_values; i++) { + values[i] = attr.default_value().list().f(i); + } + TFE_OpSetAttrFloatList(op, key, values.get(), num_values); + } else if (type == TF_ATTR_BOOL) { + int num_values = attr.default_value().list().b_size(); + std::unique_ptr values(new unsigned char[num_values]); + (*attr_list_sizes)[key] = num_values; + for (int i = 0; i < num_values; i++) { + values[i] = attr.default_value().list().b(i); + } + TFE_OpSetAttrBoolList(op, key, values.get(), num_values); + } else if (type == TF_ATTR_TYPE) { + int num_values = attr.default_value().list().type_size(); + std::unique_ptr values(new int[num_values]); + (*attr_list_sizes)[key] = num_values; + for (int i = 0; i < num_values; i++) { + values[i] = attr.default_value().list().type(i); + } + TFE_OpSetAttrTypeList(op, key, + reinterpret_cast(values.get()), + attr.default_value().list().type_size()); + } else if (type == TF_ATTR_SHAPE) { + int num_values = attr.default_value().list().shape_size(); + (*attr_list_sizes)[key] = num_values; + int total_dims = 0; + for (int i = 0; i < num_values; ++i) { + if (!attr.default_value().list().shape(i).unknown_rank()) { + total_dims += attr.default_value().list().shape(i).dim_size(); + } + } + // Allocate a buffer that can fit all of the dims together. + std::unique_ptr buffer(new int64_t[total_dims]); + // Copy the input dims into the buffer and set dims to point to + // the start of each list's dims. + std::unique_ptr dims(new const int64_t*[num_values]); + std::unique_ptr num_dims(new int[num_values]); + int64_t* offset = buffer.get(); + for (int i = 0; i < num_values; ++i) { + const auto& shape = attr.default_value().list().shape(i); + if (shape.unknown_rank()) { + dims[i] = nullptr; + num_dims[i] = -1; + } else { + for (int j = 0; j < shape.dim_size(); j++) { + *offset = shape.dim(j).size(); + ++offset; + } + } + } + TFE_OpSetAttrShapeList(op, key, dims.get(), num_dims.get(), num_values, + status); + } else if (type == TF_ATTR_FUNC) { + int num_values = attr.default_value().list().func_size(); + (*attr_list_sizes)[key] = num_values; + std::unique_ptr funcs(new const TFE_Op*[num_values]); + for (int i = 0; i < num_values; i++) { + funcs[i] = GetFunc(ctx, attr.default_value().list().func(i), status); + } + TFE_OpSetAttrFunctionList(op, key, funcs.get(), num_values); + } else { + TF_SetStatus(status, TF_UNIMPLEMENTED, + "Lists of tensors are not yet implemented for default valued " + "attributes for an operation."); + } +} + +bool SetOpAttrScalar( + TFE_Context* ctx, TFE_Op* op, const char* key, PyObject* py_value, + TF_AttrType type, + tensorflow::gtl::FlatMap* attr_list_sizes, + TF_Status* status) { if (type == TF_ATTR_STRING) { const char* value; if (!ParseStringValue(key, py_value, status, &value)) return false; @@ -193,6 +394,10 @@ bool SetOpAttrScalar(TFE_Context* ctx, TFE_Op* op, const char* key, int64_t value; if (!ParseInt64Value(key, py_value, status, &value)) return false; TFE_OpSetAttrInt(op, key, value); + // attr_list_sizes is set for all int attributes (since at this point we are + // not aware if that attribute might be used to calculate the size of an + // output list or not). + if (attr_list_sizes != nullptr) (*attr_list_sizes)[key] = value; } else if (type == TF_ATTR_FLOAT) { float value; if (!ParseFloatValue(key, py_value, status, &value)) return false; @@ -203,7 +408,7 @@ bool SetOpAttrScalar(TFE_Context* ctx, TFE_Op* op, const char* key, TFE_OpSetAttrBool(op, key, value); } else if (type == TF_ATTR_TYPE) { int value; - if (!ParseIntValue(key, py_value, status, &value)) return false; + if (!ParseTypeValue(key, py_value, status, &value)) return false; TFE_OpSetAttrType(op, key, static_cast(value)); } else if (type == TF_ATTR_SHAPE) { if (py_value == Py_None) { @@ -217,13 +422,19 @@ bool SetOpAttrScalar(TFE_Context* ctx, TFE_Op* op, const char* key, .c_str()); return false; } - const auto num_dims = PySequence_Size(py_value); + const auto num_dims = TensorShapeNumDims(py_value); + if (num_dims == -1) { + TFE_OpSetAttrShape(op, key, nullptr, -1, status); + return true; + } std::unique_ptr dims(new int64_t[num_dims]); for (int i = 0; i < num_dims; ++i) { - auto inner_py_value = PySequence_ITEM(py_value, i); - if (inner_py_value == Py_None) { + tensorflow::Safe_PyObjectPtr inner_py_value( + PySequence_ITEM(py_value, i)); + if (inner_py_value.get() == Py_None) { dims[i] = -1; - } else if (!ParseInt64Value(key, inner_py_value, status, &dims[i])) { + } else if (!ParseDimensionValue(key, inner_py_value.get(), status, + &dims[i])) { return false; } } @@ -269,6 +480,17 @@ bool SetOpAttrScalar(TFE_Context* ctx, TFE_Op* op, const char* key, return true; } +void SetOpAttrScalarDefault( + TFE_Context* ctx, TFE_Op* op, const tensorflow::AttrValue& default_value, + const char* attr_name, + tensorflow::gtl::FlatMap* attr_list_sizes, + TF_Status* status) { + SetOpAttrValueScalar(ctx, op, default_value, attr_name, status); + if (default_value.value_case() == tensorflow::AttrValue::kI) { + (*attr_list_sizes)[attr_name] = default_value.i(); + } +} + // start_index is the index at which the Tuple/List attrs will start getting // processed. void SetOpAttrs(TFE_Context* ctx, TFE_Op* op, PyObject* attrs, int start_index, @@ -294,9 +516,40 @@ void SetOpAttrs(TFE_Context* ctx, TFE_Op* op, PyObject* attrs, int start_index, const TF_AttrType type = TFE_OpGetAttrType(op, key, &is_list, out_status); if (TF_GetCode(out_status) != TF_OK) return; if (is_list != 0) { - if (!SetOpAttrList(op, key, py_value, type, out_status)) return; + if (!SetOpAttrList(op, key, py_value, type, nullptr, out_status)) return; } else { - if (!SetOpAttrScalar(ctx, op, key, py_value, type, out_status)) return; + if (!SetOpAttrScalar(ctx, op, key, py_value, type, nullptr, out_status)) + return; + } + } +} + +// This function will set the op attrs required. If an attr has the value of +// None, then it will read the AttrDef to get the default value and set that +// instead. Any failure in this function will simply fall back to the slow +// path. +void SetOpAttrWithDefaults( + TFE_Context* ctx, TFE_Op* op, const tensorflow::OpDef::AttrDef& attr, + const char* attr_name, PyObject* attr_value, + tensorflow::gtl::FlatMap* attr_list_sizes, + TF_Status* status) { + unsigned char is_list = 0; + const TF_AttrType type = TFE_OpGetAttrType(op, attr_name, &is_list, status); + if (TF_GetCode(status) != TF_OK) return; + if (attr_value == Py_None) { + if (is_list != 0) { + SetOpAttrListDefault(ctx, op, attr, attr_name, type, attr_list_sizes, + status); + } else { + SetOpAttrScalarDefault(ctx, op, attr.default_value(), attr_name, + attr_list_sizes, status); + } + } else { + if (is_list != 0) { + SetOpAttrList(op, attr_name, attr_value, type, attr_list_sizes, status); + } else { + SetOpAttrScalar(ctx, op, attr_name, attr_value, type, attr_list_sizes, + status); } } } @@ -305,6 +558,14 @@ void SetOpAttrs(TFE_Context* ctx, TFE_Op* op, PyObject* attrs, int start_index, tensorflow::mutex exception_class_mutex(tensorflow::LINKER_INITIALIZED); PyObject* exception_class GUARDED_BY(exception_class_mutex) = nullptr; +// Python subclass of Exception that is created to signal fallback. +PyObject* fallback_exception_class = nullptr; + +// Python function that returns a backward_function. +PyObject* backward_function_getter = nullptr; + +PyTypeObject* resource_variable_type = nullptr; + tensorflow::mutex _uid_mutex(tensorflow::LINKER_INITIALIZED); tensorflow::int64 _uid GUARDED_BY(_uid_mutex) = 0; @@ -353,13 +614,78 @@ PyObject* TFE_Py_RegisterExceptionClass(PyObject* e) { "TFE_Py_RegisterExceptionClass: " "Registered class should be subclass of Exception."); return nullptr; + } + + Py_INCREF(e); + exception_class = e; + Py_RETURN_NONE; +} + +PyObject* TFE_Py_RegisterResourceVariableType(PyObject* e) { + if (!PyType_Check(e)) { + PyErr_SetString( + PyExc_TypeError, + "TFE_Py_RegisterResourceVariableType: Need to register a type."); + return nullptr; + } + + if (resource_variable_type != nullptr) { + Py_DECREF(resource_variable_type); + } + + Py_INCREF(e); + resource_variable_type = reinterpret_cast(e); + Py_RETURN_NONE; +} + +PyObject* TFE_Py_RegisterFallbackExceptionClass(PyObject* e) { + if (fallback_exception_class != nullptr) { + Py_DECREF(fallback_exception_class); + } + if (PyObject_IsSubclass(e, PyExc_Exception) <= 0) { + fallback_exception_class = nullptr; + PyErr_SetString(PyExc_TypeError, + "TFE_Py_RegisterFallbackExceptionClass: " + "Registered class should be subclass of Exception."); + return nullptr; + } else { + Py_INCREF(e); + fallback_exception_class = e; + Py_RETURN_NONE; + } +} + +PyObject* TFE_Py_RegisterBackwardFunctionGetter(PyObject* e) { + if (backward_function_getter != nullptr) { + Py_DECREF(backward_function_getter); + } + if (!PyCallable_Check(e)) { + backward_function_getter = nullptr; + PyErr_SetString(PyExc_TypeError, + "TFE_Py_RegisterBackwardFunctionGetter: " + "Registered object should be function."); + return nullptr; } else { Py_INCREF(e); - exception_class = e; + backward_function_getter = e; Py_RETURN_NONE; } } +void RaiseFallbackException(const char* message) { + if (fallback_exception_class != nullptr) { + PyErr_SetObject(fallback_exception_class, Py_BuildValue("s", message)); + return; + } + + PyErr_SetString( + PyExc_RuntimeError, + tensorflow::strings::StrCat( + "Fallback exception type not set, attempting to fallback due to ", + message) + .data()); +} + int MaybeRaiseExceptionFromTFStatus(TF_Status* status, PyObject* exception) { if (TF_GetCode(status) == TF_OK) return 0; const char* msg = TF_Message(status); @@ -528,6 +854,34 @@ tensorflow::gtl::CompactPointerSet* GetTapeSet() { return tape_set; } +// A safe copy of the current tapeset. Does not get affected by other python +// threads changing the set of active tapes. +class SafeTapeSet { + public: + SafeTapeSet() : tape_set_(*GetTapeSet()) { + for (auto* tape : tape_set_) { + Py_INCREF(tape); + } + } + + ~SafeTapeSet() { + for (auto* tape : tape_set_) { + Py_DECREF(tape); + } + } + + tensorflow::gtl::CompactPointerSet::const_iterator begin() { + return tape_set_.begin(); + } + + tensorflow::gtl::CompactPointerSet::const_iterator end() { + return tape_set_.end(); + } + + private: + tensorflow::gtl::CompactPointerSet tape_set_; +}; + // xcode 7 doesn't define thread_local, so for compatibility we implement our // own. TODO(apassos) remove once we can deprecate xcode 7. #ifndef __APPLE__ @@ -658,7 +1012,14 @@ static tensorflow::eager::TapeTensor TapeTensorFromTensor(PyObject* tensor) { if (EagerTensor_CheckExact(tensor)) { TFE_TensorHandle* t = EagerTensor_Handle(tensor); tensorflow::int64 id = EagerTensor_id(tensor); - return tensorflow::eager::TapeTensor{id, t->t.dtype(), t->t.shape()}; + const tensorflow::Tensor* tensor = nullptr; + const tensorflow::Status status = t->Tensor(&tensor); + if (MaybeRaiseExceptionFromStatus(status, nullptr)) { + return tensorflow::eager::TapeTensor{id, t->dtype, + tensorflow::TensorShape({})}; + } else { + return tensorflow::eager::TapeTensor{id, t->dtype, tensor->shape()}; + } } tensorflow::int64 id = FastTensorId(tensor); if (PyErr_Occurred()) { @@ -718,10 +1079,7 @@ void TFE_Py_TapeSetWatchVariable(PyObject* variable) { if (*ThreadTapeIsStopped()) { return; } - // Note: making a copy because watching a variable can trigger a change to the - // set of tapes by allowing python's garbage collector to run. - auto tape_set = *GetTapeSet(); - for (TFE_Py_Tape* tape : tape_set) { + for (TFE_Py_Tape* tape : SafeTapeSet()) { tape->tape->WatchVariable(variable); } } @@ -736,16 +1094,10 @@ PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape) { return result; } -void TFE_Py_TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors, - PyObject* input_tensors, - PyObject* backward_function) { - if (GetTapeSet()->empty() || *ThreadTapeIsStopped()) { - return; - } - std::vector input_ids = MakeTensorIDList(input_tensors); - if (PyErr_Occurred()) { - return; - } +namespace { +void TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors, + const std::vector& input_ids, + PyObject* backward_function) { std::vector output_info; PyObject* seq = PySequence_Fast(output_tensors, "expected a sequence of integer tensor ids"); @@ -777,20 +1129,29 @@ void TFE_Py_TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors, return; } - auto set = *GetTapeSet(); - for (TFE_Py_Tape* tape : set) { + for (TFE_Py_Tape* tape : SafeTapeSet()) { Py_INCREF(backward_function); tape->tape->RecordOperation( op_type_str, output_info, input_ids, backward_function, [backward_function]() { Py_DECREF(backward_function); }); } } +} // namespace + +void TFE_Py_TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors, + PyObject* input_tensors, + PyObject* backward_function) { + if (GetTapeSet()->empty() || *ThreadTapeIsStopped()) { + return; + } + std::vector input_ids = MakeTensorIDList(input_tensors); + if (PyErr_Occurred()) return; + + TapeSetRecordOperation(op_type, output_tensors, input_ids, backward_function); +} void TFE_Py_TapeSetDeleteTrace(tensorflow::int64 tensor_id) { - // Note: making a copy because deleting the trace can trigger a change to the - // set of tapes by allowing python's garbage collector to run. - auto tape_set = *GetTapeSet(); - for (TFE_Py_Tape* tape : tape_set) { + for (TFE_Py_Tape* tape : SafeTapeSet()) { tape->tape->DeleteTrace(tensor_id); } } @@ -962,6 +1323,16 @@ std::vector MakeTensorList(PyObject* tensors) { PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* vspace, PyObject* target, PyObject* sources, PyObject* output_gradients, TF_Status* status) { + TFE_Py_Tape* tape_obj = reinterpret_cast(tape); + if (!tape_obj->tape->IsPersistent()) { + auto* tape_set = GetTapeSet(); + if (tape_set->find(tape_obj) != tape_set->end()) { + PyErr_SetString(PyExc_RuntimeError, + "Trying to call tape.gradient on a non-persistent tape " + "while it is still active."); + return nullptr; + } + } PyVSpace c_vspace(vspace); if (!c_vspace.Initialize().ok()) { return nullptr; @@ -987,7 +1358,6 @@ PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* vspace, Py_INCREF(tensor); } } - TFE_Py_Tape* tape_obj = reinterpret_cast(tape); std::vector result; status->status = tape_obj->tape->ComputeGradient( c_vspace, target_vec, sources_vec, outgrad_vec, &result); @@ -1010,17 +1380,320 @@ PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* vspace, } return py_result; } - Py_INCREF(Py_None); - return Py_None; + return PyList_New(0); } namespace { -static const int kFastPathExecuteInputStartIndex = 4; +static const int kFastPathExecuteInputStartIndex = 5; + +PyObject* GetPythonObjectFromString(const char* s) { +#if PY_MAJOR_VERSION >= 3 + return PyUnicode_FromString(s); +#else + return PyBytes_FromString(s); +#endif +} + +bool CheckResourceVariable(PyObject* item) { + return PyObject_TypeCheck(item, resource_variable_type); +} + +bool CheckInputsOk(PyObject* seq, int start_index, + const tensorflow::OpDef& op_def) { + for (int i = 0; i < op_def.input_arg_size(); i++) { + PyObject* item = PyTuple_GET_ITEM(seq, i + start_index); + if (!op_def.input_arg(i).number_attr().empty() || + !op_def.input_arg(i).type_list_attr().empty()) { + // This item should be a list input. + if (!PyList_Check(item)) return false; + for (Py_ssize_t j = 0; j < PyList_Size(item); j++) { + PyObject* inner_item = PyList_GET_ITEM(item, j); + if (!EagerTensor_CheckExact(inner_item) && + !CheckResourceVariable(inner_item)) { + return false; + } + } + } else if (!EagerTensor_CheckExact(item) && !CheckResourceVariable(item)) { + return false; + } + } + + return true; +} + +bool OpDoesntRequireOutput(const string& op_name) { + static tensorflow::gtl::FlatSet* ops_that_dont_require_outputs = + new tensorflow::gtl::FlatSet({ + "Identity", + "MatMul", + "Conv2DBackpropInput", + "Conv2DBackpropFilter", + "Conv3D", + "Conv3DBackpropInputV2", + "AvgPool3D", + "AvgPool3DGrad", + "MaxPool3D", + "MaxPool3DGrad", + "MaxPool3DGradGrad", + "BiasAdd", + "BiasAddV1", + "BiasAddGrad", + "Relu6", + "Softplus", + "SoftplusGrad", + "Softsign", + "ReluGrad", + "Conv2D", + "DepthwiseConv2dNative", + "Dilation2D", + "AvgPool", + "AvgPoolGrad", + "BatchNormWithGlobalNormalization", + "L2Loss", + "Sum", + "Prod", + "SegmentSum", + "SegmentMean", + "SparseSegmentSum", + "SparseSegmentMean", + "SparseSegmentSqrtN", + "SegmentMin", + "SegmentMax", + "UnsortedSegmentSum", + "UnsortedSegmentMax", + "Abs", + "Neg", + "ReciprocalGrad", + "Square", + "Expm1", + "Log", + "Log1p", + "TanhGrad", + "SigmoidGrad", + "Sign", + "Sin", + "Cos", + "Tan", + "Add", + "Sub", + "Mul", + "Div", + "RealDiv", + "Maximum", + "Minimum", + "SquaredDifference", + "Select", + "SparseMatMul", + "BatchMatMul", + "Complex", + "Real", + "Imag", + "Angle", + "Conj", + "Cast", + "Cross", + "Cumsum", + "Cumprod", + "ReadVariableOp", + "VarHandleOp", + "Shape", + }); + + return ops_that_dont_require_outputs->find(op_name) != + ops_that_dont_require_outputs->end(); +} + +bool OpDoesntRequireInput(const string& op_name) { + static tensorflow::gtl::FlatSet* ops_that_dont_require_inputs = + new tensorflow::gtl::FlatSet({ + "Identity", + "Softmax", + "LogSoftmax", + "BiasAdd", + "Relu", + "Elu", + "Selu", + "SparseSoftmaxCrossEntropyWithLogits", + "Neg", + "Inv", + "Reciprocal", + "Sqrt", + "Exp", + "Tanh", + "Sigmoid", + "Real", + "Imag", + "Conj", + "ReadVariableOp", + "VarHandleOp", + "Shape", + }); + + return ops_that_dont_require_inputs->find(op_name) != + ops_that_dont_require_inputs->end(); +} + +PyObject* RecordGradient(PyObject* op_name, PyObject* inputs, PyObject* attrs, + PyObject* results, PyObject* name) { + std::vector input_ids = MakeTensorIDList(inputs); + if (PyErr_Occurred()) return nullptr; + + bool should_record = false; + for (TFE_Py_Tape* tape : SafeTapeSet()) { + if (tape->tape->ShouldRecord(input_ids)) { + should_record = true; + break; + } + } + if (!should_record) Py_RETURN_NONE; + + string c_op_name = TFE_GetPythonString(op_name); + PyObject* op_outputs; + if (OpDoesntRequireOutput(c_op_name)) { + op_outputs = Py_None; + } else { + op_outputs = results; + } + + PyObject* op_inputs; + if (OpDoesntRequireInput(c_op_name)) { + op_inputs = Py_None; + } else { + op_inputs = inputs; + } + + PyObject* num_inputs = PyLong_FromLong(PySequence_Size(inputs)); + PyObject* callback_args = + Py_BuildValue("OOOOO", op_name, attrs, num_inputs, op_inputs, op_outputs); + + PyObject* backward_function = + PyObject_CallObject(backward_function_getter, callback_args); + Py_DECREF(callback_args); + if (backward_function == nullptr) return nullptr; + + TapeSetRecordOperation(op_name, results, input_ids, backward_function); + + Py_DECREF(backward_function); + + Py_RETURN_NONE; +} + +void MaybeWatchVariable(PyObject* input) { + DCHECK(CheckResourceVariable(input)); + DCHECK(PyObject_HasAttrString(input, "_trainable")); + + tensorflow::Safe_PyObjectPtr trainable( + PyObject_GetAttrString(input, "_trainable")); + if (trainable.get() == Py_False) return; + TFE_Py_TapeSetWatchVariable(input); +} + +bool ReadVariableOp(const FastPathOpExecInfo& parent_op_exec_info, + PyObject* input, tensorflow::Safe_PyObjectPtr* output, + TF_Status* status) { + MaybeWatchVariable(input); + + TFE_Op* op = TFE_NewOp(parent_op_exec_info.ctx, "ReadVariableOp", status); + auto cleaner = tensorflow::gtl::MakeCleanup([op] { TFE_DeleteOp(op); }); + if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) return false; + + // Set dtype + DCHECK(PyObject_HasAttrString(input, "_dtype")); + tensorflow::Safe_PyObjectPtr dtype(PyObject_GetAttrString(input, "_dtype")); + int value; + if (!ParseTypeValue("_dtype", dtype.get(), status, &value)) { + return false; + } + TFE_OpSetAttrType(op, "dtype", static_cast(value)); + + TFE_OpSetDevice(op, parent_op_exec_info.device_name, status); + if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) return false; + + // Get handle + tensorflow::Safe_PyObjectPtr handle(PyObject_GetAttrString(input, "_handle")); + if (!EagerTensor_CheckExact(handle.get())) return false; + TFE_OpAddInput(op, EagerTensor_Handle(handle.get()), status); + if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) return false; + + int num_retvals = 1; + TFE_TensorHandle* output_handle; + TFE_Execute(op, &output_handle, &num_retvals, status); + if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) return false; + + // Always create the py object (and correctly DECREF it) from the returned + // value, else the data will leak. + output->reset(EagerTensorFromHandle(output_handle)); + + // TODO(nareshmodi): Should we run post exec callbacks here? + if (parent_op_exec_info.run_gradient_callback) { + tensorflow::Safe_PyObjectPtr inputs(PyTuple_New(1)); + PyTuple_SET_ITEM(inputs.get(), 0, handle.release()); + + tensorflow::Safe_PyObjectPtr outputs(PyTuple_New(1)); + Py_INCREF(output->get()); // stay alive after since tuple steals. + PyTuple_SET_ITEM(outputs.get(), 0, output->get()); + + if (!RecordGradient(GetPythonObjectFromString("ReadVariableOp"), + inputs.get(), Py_None, outputs.get(), Py_None)) { + return false; + } + } + + return true; +} + +// Supports only 2 cases at the moment: +// i) input is an EagerTensor +// ii) input is a ResourceVariable - in this case, the is_variable param is set +// to true. +bool ConvertToTensor(const FastPathOpExecInfo& op_exec_info, PyObject* input, + tensorflow::Safe_PyObjectPtr* output_handle, + TF_Status* status) { + if (CheckResourceVariable(input)) { + return ReadVariableOp(op_exec_info, input, output_handle, status); + } + + Py_INCREF(input); + output_handle->reset(input); + + return true; +} + +// Adds input and type attr to the op, and to the list of flattened +// inputs/attrs. +bool AddInputToOp(const FastPathOpExecInfo& op_exec_info, PyObject* input, + const tensorflow::OpDef::ArgDef* input_arg, + std::vector* flattened_attrs, + std::vector* flattened_inputs, + TFE_Op* op, TF_Status* status) { + // py_eager_tensor's ownership is transferred to flattened_inputs if it is + // required, else the object is destroyed and DECREF'd when the object goes + // out of scope in this function. + tensorflow::Safe_PyObjectPtr py_eager_tensor = nullptr; + + if (!ConvertToTensor(op_exec_info, input, &py_eager_tensor, status)) { + return false; + } + + TFE_TensorHandle* input_handle = EagerTensor_Handle(py_eager_tensor.get()); + + if (input_arg != nullptr && !input_arg->type_attr().empty()) { + auto dtype = TFE_TensorHandleDataType(input_handle); + TFE_OpSetAttrType(op, input_arg->type_attr().data(), dtype); + if (flattened_attrs != nullptr) { + flattened_attrs->emplace_back( + GetPythonObjectFromString(input_arg->type_attr().data())); + flattened_attrs->emplace_back(PyLong_FromLong(dtype)); + } + } + + if (flattened_inputs != nullptr) { + flattened_inputs->emplace_back(std::move(py_eager_tensor)); + } -bool CheckEagerTensors(PyObject* seq, int start_index, int num_to_check) { - for (int i = start_index; i < start_index + num_to_check; i++) { - PyObject* item = PyTuple_GET_ITEM(seq, i); - if (!EagerTensor_CheckExact(item)) return false; + TFE_OpAddInput(op, input_handle, status); + if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) { + return false; } return true; @@ -1052,59 +1725,89 @@ const char* GetDeviceName(PyObject* py_device_name) { return nullptr; } -bool MaybeRunRecordGradientCallback(const tensorflow::OpDef* op_def, - PyObject* args, PyObject* result, - PyObject* record_gradient_callback) { - if (*ThreadTapeIsStopped() || GetTapeSet()->empty() || - record_gradient_callback == Py_None) { - return true; - } - if (!PyCallable_Check(record_gradient_callback)) { - PyErr_SetString( - PyExc_TypeError, - Printf( - "expected a function for record_gradient_callback, got %s instead", - record_gradient_callback->ob_type->tp_name) - .c_str()); +bool RaiseIfNotPyList(PyObject* list, const string& attr_name) { + if (!PyList_Check(list)) { + PyErr_SetString(PyExc_TypeError, + Printf("expected a list for attr %s, got %s instead", + attr_name.data(), list->ob_type->tp_name) + .data()); + return false; } + return true; +} - PyObject* inputs = PyTuple_New(op_def->input_arg_size()); - for (int i = 0; i < op_def->input_arg_size(); i++) { - auto* input = PyTuple_GET_ITEM(args, kFastPathExecuteInputStartIndex + i); +bool RunCallbacks( + const FastPathOpExecInfo& op_exec_info, PyObject* args, + const std::vector& flattened_inputs, + const std::vector& flattened_attrs, + PyObject* flattened_result) { + if (!op_exec_info.run_callbacks) return true; + + tensorflow::Safe_PyObjectPtr inputs(PyTuple_New(flattened_inputs.size())); + for (int i = 0; i < flattened_inputs.size(); i++) { + PyObject* input = flattened_inputs[i].get(); Py_INCREF(input); - PyTuple_SET_ITEM(inputs, i, input); + PyTuple_SET_ITEM(inputs.get(), i, input); } - int args_size = PyTuple_GET_SIZE(args); - int num_attrs = - args_size - op_def->input_arg_size() - kFastPathExecuteInputStartIndex; - PyObject* attrs = PyTuple_New(num_attrs); - for (int i = 0; i < num_attrs; i++) { - auto* attr = PyTuple_GET_ITEM( - args, kFastPathExecuteInputStartIndex + op_def->input_arg_size() + i); + int num_non_inferred_attrs = PyTuple_GET_SIZE(args) - + op_exec_info.op_def->input_arg_size() - + kFastPathExecuteInputStartIndex; + int num_attrs = flattened_attrs.size() + num_non_inferred_attrs; + tensorflow::Safe_PyObjectPtr attrs(PyTuple_New(num_attrs)); + + for (int i = 0; i < num_non_inferred_attrs; i++) { + auto* attr = + PyTuple_GET_ITEM(args, kFastPathExecuteInputStartIndex + + op_exec_info.op_def->input_arg_size() + i); Py_INCREF(attr); - PyTuple_SET_ITEM(attrs, i, attr); + PyTuple_SET_ITEM(attrs.get(), i, attr); + } + for (int i = num_non_inferred_attrs; i < num_attrs; i++) { + PyObject* attr_or_name = + flattened_attrs.at(i - num_non_inferred_attrs).get(); + Py_INCREF(attr_or_name); + PyTuple_SET_ITEM(attrs.get(), i, attr_or_name); } - PyObject* callback_args = Py_BuildValue("OOO", inputs, attrs, result); - PyObject_CallObject(record_gradient_callback, callback_args); + if (op_exec_info.run_gradient_callback) { + if (!RecordGradient(op_exec_info.op_name, inputs.get(), attrs.get(), + flattened_result, op_exec_info.name)) { + return false; + } + } + + if (op_exec_info.run_post_exec_callbacks) { + tensorflow::Safe_PyObjectPtr callback_args( + Py_BuildValue("OOOOO", op_exec_info.op_name, inputs.get(), attrs.get(), + flattened_result, op_exec_info.name)); + for (Py_ssize_t i = 0; i < PyList_Size(op_exec_info.callbacks); i++) { + PyObject* callback_fn = PyList_GET_ITEM(op_exec_info.callbacks, i); + if (!PyCallable_Check(callback_fn)) { + PyErr_SetString( + PyExc_TypeError, + Printf("expected a function for " + "post execution callback in index %ld, got %s instead", + i, callback_fn->ob_type->tp_name) + .c_str()); + return false; + } + PyObject* callback_result = + PyObject_CallObject(callback_fn, callback_args.get()); + if (!callback_result) { + return false; + } + Py_DECREF(callback_result); + } + } - Py_DECREF(inputs); - Py_DECREF(callback_args); - Py_DECREF(attrs); return true; } + } // namespace PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { - TFE_Context* ctx = reinterpret_cast( - PyCapsule_GetPointer(PyTuple_GET_ITEM(args, 0), nullptr)); - const tensorflow::OpDef* op_def = GetOpDef(PyTuple_GET_ITEM(args, 2)); - if (op_def == nullptr) return nullptr; - const char* device_name = GetDeviceName(PyTuple_GET_ITEM(args, 1)); - PyObject* record_gradient_callback = PyTuple_GET_ITEM(args, 3); - Py_ssize_t args_size = PyTuple_GET_SIZE(args); if (args_size < kFastPathExecuteInputStartIndex) { PyErr_SetString( @@ -1115,6 +1818,31 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { return nullptr; } + FastPathOpExecInfo op_exec_info; + + op_exec_info.ctx = reinterpret_cast( + PyCapsule_GetPointer(PyTuple_GET_ITEM(args, 0), nullptr)); + op_exec_info.device_name = GetDeviceName(PyTuple_GET_ITEM(args, 1)); + op_exec_info.op_name = PyTuple_GET_ITEM(args, 2); + op_exec_info.op_def = GetOpDef(op_exec_info.op_name); + if (op_exec_info.op_def == nullptr) return nullptr; + op_exec_info.name = PyTuple_GET_ITEM(args, 3); + op_exec_info.callbacks = PyTuple_GET_ITEM(args, 4); + + const tensorflow::OpDef* op_def = op_exec_info.op_def; + + // TODO(nareshmodi): Add a benchmark for the fast-path with gradient callbacks + // (similar to benchmark_tf_gradient_function_*). Also consider using an + // InlinedVector for flattened_attrs and flattened_inputs if the benchmarks + // point out problems with heap allocs. + op_exec_info.run_gradient_callback = + !*ThreadTapeIsStopped() && !GetTapeSet()->empty(); + op_exec_info.run_post_exec_callbacks = + op_exec_info.callbacks != Py_None && + PyList_Size(op_exec_info.callbacks) > 0; + op_exec_info.run_callbacks = op_exec_info.run_gradient_callback || + op_exec_info.run_post_exec_callbacks; + if (args_size < kFastPathExecuteInputStartIndex + op_def->input_arg_size()) { PyErr_SetString( PyExc_ValueError, @@ -1126,18 +1854,15 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { return nullptr; } - if (!CheckEagerTensors(args, kFastPathExecuteInputStartIndex, - op_def->input_arg_size())) { - // TODO(nareshmodi): Maybe some other way of signalling that this should - // fall back? - PyErr_SetString(PyExc_NotImplementedError, - "This function does not handle the case of the path where " - "all inputs are not already EagerTensors."); + if (!CheckInputsOk(args, kFastPathExecuteInputStartIndex, *op_def)) { + RaiseFallbackException( + "This function does not handle the case of the path where " + "all inputs are not already EagerTensors."); return nullptr; } TF_Status* status = TF_NewStatus(); - TFE_Op* op = TFE_NewOp(ctx, op_def->name().c_str(), status); + TFE_Op* op = TFE_NewOp(op_exec_info.ctx, op_def->name().c_str(), status); auto cleaner = tensorflow::gtl::MakeCleanup([status, op] { TF_DeleteStatus(status); TFE_DeleteOp(op); @@ -1146,62 +1871,248 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { return nullptr; } - TFE_OpSetDevice(op, device_name, status); - if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) { - return nullptr; + // Mapping of attr name to size - used to calculate the number of values + // to be expected by the TFE_Execute run. + tensorflow::gtl::FlatMap attr_list_sizes; + + // Set non-inferred attrs, including setting defaults if the attr is passed in + // as None. + for (int i = kFastPathExecuteInputStartIndex + op_def->input_arg_size(); + i < args_size; i += 2) { + PyObject* py_attr_name = PyTuple_GET_ITEM(args, i); + const tensorflow::StringPiece attr_name(TFE_GetPythonString(py_attr_name)); + PyObject* py_attr_value = PyTuple_GET_ITEM(args, i + 1); + + // Not creating an index since most of the time there are not more than a + // few attrs. + // TODO(nareshmodi): Maybe include the index as part of the + // OpRegistrationData. + for (const auto& attr : op_def->attr()) { + if (attr_name == attr.name()) { + SetOpAttrWithDefaults(op_exec_info.ctx, op, attr, attr_name.data(), + py_attr_value, &attr_list_sizes, status); + + if (TF_GetCode(status) != TF_OK) { + RaiseFallbackException(TF_Message(status)); + return nullptr; + } + + break; + } + } } - // Add non-type attrs. - SetOpAttrs(ctx, op, args, - kFastPathExecuteInputStartIndex + op_def->input_arg_size(), - status); + TFE_OpSetDevice(op, op_exec_info.device_name, status); if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) { return nullptr; } - // Add type attrs and inputs. + // Flat attrs and inputs as required by the record_gradient call. The attrs + // here only contain inferred attrs (non-inferred attrs are added directly + // from the input args). + // All items in flattened_attrs and flattened_inputs contain + // Safe_PyObjectPtr - any time something steals a reference to this, it must + // INCREF. + // TODO(nareshmodi): figure out why PyList_New/PyList_Append don't work + // directly. + std::unique_ptr> flattened_attrs = + nullptr; + std::unique_ptr> flattened_inputs = + nullptr; + + // TODO(nareshmodi): Encapsulate callbacks information into a struct. + if (op_exec_info.run_callbacks) { + flattened_attrs.reset(new std::vector); + flattened_inputs.reset(new std::vector); + } + + // Add inferred attrs and inputs. + // The following code might set duplicate type attrs. This will result in + // the CacheKey for the generated AttrBuilder possibly differing from + // those where the type attrs are correctly set. Inconsistent CacheKeys + // for ops means that there might be unnecessarily duplicated kernels. + // TODO(nareshmodi): Fix this. for (int i = 0; i < op_def->input_arg_size(); i++) { const auto& input_arg = op_def->input_arg(i); PyObject* input = PyTuple_GET_ITEM(args, kFastPathExecuteInputStartIndex + i); - TFE_TensorHandle* input_handle = EagerTensor_Handle(input); - - // The following code might set duplicate type attrs. This will result in - // the CacheKey for the generated AttrBuilder possibly differing from those - // where the type attrs are correctly set. Inconsistent CacheKeys for ops - // means that there might be unnecessarily duplicated kernels. - // TODO(nareshmodi): Fix this. - if (!input_arg.type_attr().empty()) { - TFE_OpSetAttrType(op, input_arg.type_attr().data(), - TFE_TensorHandleDataType(input_handle)); + if (!input_arg.number_attr().empty()) { + // The item is a homogeneous list. + if (!RaiseIfNotPyList(input, input_arg.number_attr())) return nullptr; + Py_ssize_t len = PyList_Size(input); + + TFE_OpSetAttrInt(op, input_arg.number_attr().data(), len); + if (op_exec_info.run_callbacks) { + flattened_attrs->emplace_back( + GetPythonObjectFromString(input_arg.number_attr().data())); + flattened_attrs->emplace_back(PyLong_FromLong(len)); + } + attr_list_sizes[input_arg.number_attr()] = len; + + if (len > 0) { + // First item adds the type attr. + if (!AddInputToOp(op_exec_info, PyList_GET_ITEM(input, 0), &input_arg, + flattened_attrs.get(), flattened_inputs.get(), op, + status)) { + return nullptr; + } + + for (Py_ssize_t j = 1; j < len; j++) { + // Since the list is homogeneous, we don't need to re-add the attr. + if (!AddInputToOp(op_exec_info, PyList_GET_ITEM(input, j), + nullptr /* input_arg */, + nullptr /* flattened_attrs */, + flattened_inputs.get(), op, status)) { + return nullptr; + } + } + } + } else if (!input_arg.type_list_attr().empty()) { + // The item is a heterogeneous list. + if (!RaiseIfNotPyList(input, input_arg.type_list_attr())) return nullptr; + const string& attr_name = input_arg.type_list_attr(); + Py_ssize_t len = PyList_Size(input); + tensorflow::gtl::InlinedVector attr_value(len); + PyObject* py_attr_value = nullptr; + if (op_exec_info.run_callbacks) { + py_attr_value = PyTuple_New(len); + } + for (Py_ssize_t j = 0; j < len; j++) { + PyObject* py_input = PyList_GET_ITEM(input, j); + tensorflow::Safe_PyObjectPtr py_eager_tensor; + if (!ConvertToTensor(op_exec_info, py_input, &py_eager_tensor, + status)) { + return nullptr; + } + + TFE_TensorHandle* input_handle = + EagerTensor_Handle(py_eager_tensor.get()); + + attr_value[j] = TFE_TensorHandleDataType(input_handle); + + TFE_OpAddInput(op, input_handle, status); + if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) { + return nullptr; + } + + if (op_exec_info.run_callbacks) { + flattened_inputs->emplace_back(std::move(py_eager_tensor)); + + PyTuple_SET_ITEM(py_attr_value, j, PyLong_FromLong(attr_value[j])); + } + } + if (op_exec_info.run_callbacks) { + flattened_attrs->emplace_back( + GetPythonObjectFromString(attr_name.data())); + flattened_attrs->emplace_back(py_attr_value); + } + TFE_OpSetAttrTypeList(op, attr_name.data(), attr_value.data(), + attr_value.size()); + attr_list_sizes[attr_name] = len; + } else { + // The item is a single item. + if (!AddInputToOp(op_exec_info, input, &input_arg, flattened_attrs.get(), + flattened_inputs.get(), op, status)) { + return nullptr; + } } + } - TFE_OpAddInput(op, input_handle, status); - if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) { - return nullptr; + int num_retvals = 0; + for (int i = 0; i < op_def->output_arg_size(); i++) { + const auto& output_arg = op_def->output_arg(i); + if (!output_arg.number_attr().empty()) { + num_retvals += attr_list_sizes[output_arg.number_attr()]; + } else if (!output_arg.type_list_attr().empty()) { + num_retvals += attr_list_sizes[output_arg.type_list_attr()]; + } else { + num_retvals++; } } - int num_retvals = op_def->output_arg_size(); tensorflow::gtl::InlinedVector retvals(num_retvals); Py_BEGIN_ALLOW_THREADS; TFE_Execute(op, retvals.data(), &num_retvals, status); Py_END_ALLOW_THREADS; - if (MaybeRaiseExceptionFromTFStatus(status, nullptr)) { + + if (TF_GetCode(status) != TF_OK) { + // Augment the status with the op_name for easier debugging similar to + // TFE_Py_Execute. + TF_SetStatus(status, TF_GetCode(status), + tensorflow::strings::StrCat( + TF_Message(status), + " [Op:", TFE_GetPythonString(op_exec_info.op_name), "]") + .c_str()); + + MaybeRaiseExceptionFromTFStatus(status, nullptr); return nullptr; } - PyObject* result = PyTuple_New(num_retvals); + tensorflow::Safe_PyObjectPtr flat_result(PyList_New(num_retvals)); for (int i = 0; i < num_retvals; ++i) { - PyTuple_SET_ITEM(result, i, EagerTensorFromHandle(retvals[i])); + PyList_SET_ITEM(flat_result.get(), i, EagerTensorFromHandle(retvals[i])); } - if (!MaybeRunRecordGradientCallback(op_def, args, result, - record_gradient_callback)) { + if (!RunCallbacks(op_exec_info, args, *flattened_inputs, *flattened_attrs, + flat_result.get())) { return nullptr; } + // Unflatten results. + if (op_def->output_arg_size() == 0) { + Py_RETURN_NONE; + } + + if (op_def->output_arg_size() == 1) { + if (!op_def->output_arg(0).number_attr().empty() || + !op_def->output_arg(0).type_list_attr().empty()) { + return flat_result.release(); + } else { + auto* result = PyList_GET_ITEM(flat_result.get(), 0); + Py_INCREF(result); + return result; + } + } + + // Correctly output the results that are made into a namedtuple. + PyObject* result = PyList_New(op_def->output_arg_size()); + int flat_result_index = 0; + for (int i = 0; i < op_def->output_arg_size(); i++) { + if (!op_def->output_arg(i).number_attr().empty()) { + int list_length = attr_list_sizes[op_def->output_arg(i).number_attr()]; + PyObject* inner_list = PyList_New(list_length); + for (int j = 0; j < list_length; j++) { + PyObject* obj = PyList_GET_ITEM(flat_result.get(), flat_result_index++); + Py_INCREF(obj); + PyList_SET_ITEM(inner_list, j, obj); + } + PyList_SET_ITEM(result, i, inner_list); + } else if (!op_def->output_arg(i).type_list_attr().empty()) { + int list_length = attr_list_sizes[op_def->output_arg(i).type_list_attr()]; + PyObject* inner_list = PyList_New(list_length); + for (int j = 0; j < list_length; j++) { + PyObject* obj = PyList_GET_ITEM(flat_result.get(), flat_result_index++); + Py_INCREF(obj); + PyList_SET_ITEM(inner_list, j, obj); + } + PyList_SET_ITEM(result, i, inner_list); + } else { + PyObject* obj = PyList_GET_ITEM(flat_result.get(), flat_result_index++); + Py_INCREF(obj); + PyList_SET_ITEM(result, i, obj); + } + } return result; } + +PyObject* TFE_Py_RecordGradient(PyObject* op_name, PyObject* inputs, + PyObject* attrs, PyObject* results, + PyObject* name) { + if (*ThreadTapeIsStopped() || GetTapeSet()->empty()) { + Py_RETURN_NONE; + } + + return RecordGradient(op_name, inputs, attrs, results, name); +} diff --git a/tensorflow/python/eager/pywrap_tfe_test.py b/tensorflow/python/eager/pywrap_tfe_test.py index d4f4ed592fb99e475af4652a33e5364d9abeea1a..faaae40b3f1ef02984a7a75c23ae4acae65ac335 100644 --- a/tensorflow/python/eager/pywrap_tfe_test.py +++ b/tensorflow/python/eager/pywrap_tfe_test.py @@ -24,23 +24,10 @@ from tensorflow.python.eager import context from tensorflow.python.eager import test from tensorflow.python.framework import constant_op from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops - - -def record_gradient_callback(inputs, attrs, results): - return backprop._record_gradient("MatMul", inputs, attrs, results, None) - - -def c_tfe_py_fastpath_execute(a, b, transpose_a=False, transpose_b=False): - ctx = context.context() - assert not ctx.in_graph_mode( - ), "The prototype doesn't contain C code for graph construction" - ctx_handle = ctx._handle # pylint: disable=protected-access - - return pywrap_tensorflow.TFE_Py_FastPathExecute( - ctx_handle, ctx.device_name, "MatMul", record_gradient_callback, a, b, - "transpose_a", transpose_a, "transpose_b", transpose_b)[0] +from tensorflow.python.ops import resource_variable_ops class Tests(test.TestCase): @@ -54,55 +41,154 @@ class Tests(test.TestCase): a_100_by_784 = random_ops.random_uniform((100, 784)) b_100_by_784 = random_ops.random_uniform((100, 784)) + ctx = context.context() + self.assertAllClose( math_ops.matmul(a_2_by_2, b_2_by_2), - c_tfe_py_fastpath_execute(a_2_by_2, b_2_by_2)) + pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", None, None, a_2_by_2, + b_2_by_2, "transpose_a", False, "transpose_b", False)) self.assertAllClose( math_ops.matmul(a_100_by_784, b_100_by_784, transpose_b=True), - c_tfe_py_fastpath_execute(a_100_by_784, b_100_by_784, transpose_b=True)) + pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", None, None, a_100_by_784, + b_100_by_784, "transpose_a", False, "transpose_b", True)) + + @test_util.assert_no_new_tensors + @test_util.assert_no_garbage_created + def testFastpathExecute_ResourceVariableMatMulCorrectResponse(self): + ctx = context.context() + a_2_by_2 = constant_op.constant(1.0, shape=[2, 2]) + m = resource_variable_ops.ResourceVariable(a_2_by_2) + x = pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", None, None, m, m, "transpose_a", + False, "transpose_b", False) + y = pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", None, None, a_2_by_2, a_2_by_2, + "transpose_a", False, "transpose_b", False) + + self.assertAllEqual(x, y) @test_util.assert_no_new_tensors @test_util.assert_no_garbage_created def testFastpathExecute_TapeWrite(self): + ctx = context.context() with backprop.GradientTape(persistent=True) as tape: a_2_by_2 = constant_op.constant(1.0, shape=[2, 2]) tape.watch(a_2_by_2) - z = c_tfe_py_fastpath_execute(a_2_by_2, a_2_by_2) + z = pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", None, None, a_2_by_2, + a_2_by_2, "transpose_a", False, "transpose_b", False) dz_dy = tape.gradient(z, [a_2_by_2])[0] self.assertAllEqual(dz_dy.numpy(), constant_op.constant(4.0, shape=[2, 2]).numpy()) @test_util.assert_no_new_tensors @test_util.assert_no_garbage_created - def testFastpathExecute_MatMulSlowPath(self): - a_2_by_2 = random_ops.random_uniform((2, 2)).cpu().numpy() + def testFastpathExecute_ResourceVariableTapeWrite(self): + ctx = context.context() + with backprop.GradientTape(persistent=True) as tape: + a_2_by_2 = constant_op.constant(1.0, shape=[2, 2]) + m = resource_variable_ops.ResourceVariable(a_2_by_2) + tape.watch(m) + z = pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "MatMul", None, None, m, m, + "transpose_a", False, "transpose_b", False) + dz_dy = tape.gradient(z, [m])[0] + self.assertAllEqual(dz_dy.numpy(), + constant_op.constant(4.0, shape=[2, 2]).numpy()) - with self.assertRaises(NotImplementedError): - c_tfe_py_fastpath_execute(a_2_by_2, a_2_by_2) + # Tests homogeneous list op + @test_util.assert_no_new_tensors + @test_util.assert_no_garbage_created + def testFastpathExecute_AddNCorrectResponse(self): + ctx = context.context() + a_2_by_2 = random_ops.random_uniform((2, 2)) + b_2_by_2 = random_ops.random_uniform((2, 2)) + + self.assertAllClose( + math_ops.add_n([a_2_by_2, b_2_by_2]), + pywrap_tensorflow.TFE_Py_FastPathExecute(ctx._handle, ctx.device_name, + "AddN", None, None, + [a_2_by_2, b_2_by_2])) + + # Tests homogeneous list op + @test_util.assert_no_new_tensors + @test_util.assert_no_garbage_created + def testFastpathExecute_AddNTapeWrite(self): + ctx = context.context() + a_2_by_2 = random_ops.random_uniform((2, 2)) + b_2_by_2 = random_ops.random_uniform((2, 2)) + + with backprop.GradientTape(persistent=True) as tape: + tape.watch(a_2_by_2) + tape.watch(b_2_by_2) + z1 = pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "AddN", None, None, + [a_2_by_2, b_2_by_2]) + z2 = math_ops.add_n([a_2_by_2, b_2_by_2]) + dz1_dy = tape.gradient(z1, [a_2_by_2])[0] + dz2_dy = tape.gradient(z2, [a_2_by_2])[0] + self.assertAllEqual(dz1_dy.numpy(), dz2_dy.numpy()) + + # Tests heterogeneous list op + @test_util.assert_no_new_tensors + @test_util.assert_no_garbage_created + def testFastpathExecute_IdentityNCorrectResponse(self): + ctx = context.context() + a_2_by_2 = random_ops.random_uniform((2, 2)) + b_2_by_2 = random_ops.random_uniform((2, 2)) + + self.assertAllClose( + array_ops.identity_n([a_2_by_2, b_2_by_2]), + pywrap_tensorflow.TFE_Py_FastPathExecute(ctx._handle, ctx.device_name, + "IdentityN", None, None, + [a_2_by_2, b_2_by_2])) + + # Tests heterogeneous list op + @test_util.assert_no_new_tensors + @test_util.assert_no_garbage_created + def testFastpathExecute_IdentityNTapeWrite(self): + ctx = context.context() + a_2_by_2 = random_ops.random_uniform((2, 2)) + b_2_by_2 = random_ops.random_uniform((2, 2)) + + with backprop.GradientTape(persistent=True) as tape: + tape.watch(a_2_by_2) + tape.watch(b_2_by_2) + z1 = pywrap_tensorflow.TFE_Py_FastPathExecute( + ctx._handle, ctx.device_name, "IdentityN", None, None, + [a_2_by_2, b_2_by_2]) + z2 = array_ops.identity_n([a_2_by_2, b_2_by_2]) + dz1_dy = tape.gradient(z1[0], [a_2_by_2])[0] + dz2_dy = tape.gradient(z2[0], [a_2_by_2])[0] + self.assertAllEqual(dz1_dy.numpy(), dz2_dy.numpy()) @test_util.assert_no_new_tensors @test_util.assert_no_garbage_created def testFastpathExecute_InvalidInputs(self): a_2_by_2 = random_ops.random_uniform((2, 2)) ctx = context.context() - assert not ctx.in_graph_mode( + assert ctx.executing_eagerly( ), "The prototype doesn't contain C code for graph construction" ctx_handle = ctx._handle # pylint: disable=protected-access + # Not enough base params with self.assertRaisesRegexp(ValueError, - "at least 4 items in the input tuple"): + "at least 5 items in the input tuple"): pywrap_tensorflow.TFE_Py_FastPathExecute(ctx_handle, ctx.device_name, "Identity") + # Not enough inputs with self.assertRaisesRegexp(ValueError, - "Expected to be at least 5, was 4"): - pywrap_tensorflow.TFE_Py_FastPathExecute( - ctx_handle, ctx_handle, "Identity", record_gradient_callback) + "Expected to be at least 6, was 5"): + pywrap_tensorflow.TFE_Py_FastPathExecute(ctx_handle, ctx_handle, + "Identity", None, []) + # Bad type with self.assertRaisesRegexp(TypeError, "expected a string for op_name"): - pywrap_tensorflow.TFE_Py_FastPathExecute( - ctx_handle, ctx.device_name, ctx_handle, record_gradient_callback, - a_2_by_2) + pywrap_tensorflow.TFE_Py_FastPathExecute(ctx_handle, ctx.device_name, + ctx_handle, None, [], a_2_by_2) if __name__ == "__main__": diff --git a/tensorflow/python/eager/tape_test.py b/tensorflow/python/eager/tape_test.py index b490bac66db03b0a61a8852f45f1f558cccaf121..4326d5efa3d362e883815eb2d3dafb27df25afd4 100644 --- a/tensorflow/python/eager/tape_test.py +++ b/tensorflow/python/eager/tape_test.py @@ -21,11 +21,11 @@ from __future__ import print_function from tensorflow.python.eager import backprop from tensorflow.python.eager import context -from tensorflow.python.eager import custom_gradient from tensorflow.python.eager import test 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 custom_gradient from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import math_ops # Importing nn_grad for the registration functions. @@ -165,21 +165,6 @@ class TapeTest(test.TestCase): g, = backprop.gradients_function(fn, [0])(t) self.assertAllEqual(g, 1.0) - def testCustomGradientGraphMode(self): - with context.graph_mode(), self.test_session(): - - @custom_gradient.custom_gradient - def f(x): - - def grad(dresult): - return dresult * 10.0 - - return x, grad - - inp = constant_op.constant(1.0) - grad = gradients_impl.gradients(f(inp), inp) - self.assertAllEqual(grad[0].eval(), 10.0) - if __name__ == '__main__': test.main() diff --git a/tensorflow/python/estimator/BUILD b/tensorflow/python/estimator/BUILD index c519fd557a9319d6ef5522b26198e5b4202917fc..04fcbb0e874bcc5583cd5f607de3a3764f4ba6b0 100644 --- a/tensorflow/python/estimator/BUILD +++ b/tensorflow/python/estimator/BUILD @@ -7,6 +7,7 @@ package( licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "py_test") +load("//tensorflow:tensorflow.bzl", "cuda_py_test") filegroup( name = "all_files", @@ -35,9 +36,9 @@ py_library( ":linear", ":model_fn", ":parsing_utils", + ":replicate_model_fn", ":run_config", ":training", - ":warm_starting_util", "//tensorflow/python:util", ], ) @@ -278,12 +279,12 @@ py_library( srcs = ["canned/dnn_testing_utils.py"], srcs_version = "PY2AND3", deps = [ + ":estimator", ":head", ":metric_keys", ":model_fn", ":numpy_io", ":prediction_keys", - ":warm_starting_util", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", @@ -427,7 +428,6 @@ py_library( ":model_fn", ":run_config", ":util", - ":warm_starting_util", "//tensorflow/core:protos_all_py", "//tensorflow/python:client", "//tensorflow/python:control_flow_ops", @@ -870,37 +870,64 @@ py_test( ) py_library( - name = "warm_starting_util", - srcs = ["warm_starting_util.py"], + name = "replicate_model_fn", + srcs = [ + "replicate_model_fn.py", + ], srcs_version = "PY2AND3", deps = [ + ":export_output", + ":model_fn", + ":util", + "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:device", + "//tensorflow/python:device_lib", "//tensorflow/python:framework_ops", + "//tensorflow/python:math_ops", "//tensorflow/python:platform", + "//tensorflow/python:sparse_ops", + "//tensorflow/python:sparse_tensor", "//tensorflow/python:state_ops", "//tensorflow/python:training", "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "//tensorflow/python/feature_column", + "//tensorflow/python/ops/losses", + "@six_archive//:six", ], ) -py_test( - name = "warm_starting_util_test", - size = "small", - srcs = ["warm_starting_util_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":warm_starting_util", +cuda_py_test( + name = "replicate_model_fn_test", + size = "medium", + srcs = ["replicate_model_fn_test.py"], + additional_deps = [ + "//tensorflow/python/estimator", + ":dnn", + ":export_export", + ":export_output", + ":model_fn", + ":numpy_io", + ":optimizers", + ":prediction_keys", + "//tensorflow/python/feature_column", + "//tensorflow/python/ops/losses", + "//tensorflow/python/saved_model:signature_constants", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:init_ops", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:metrics", + "//tensorflow/python:platform", + "//tensorflow/python:summary", "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", - "//tensorflow/python/feature_column", - "//third_party/py/numpy", + ":replicate_model_fn", + ], + tags = [ + "multi_gpu", ], ) diff --git a/tensorflow/python/estimator/canned/baseline.py b/tensorflow/python/estimator/canned/baseline.py index 138152ac1c6b2d7e399218208dd7bdf2d8136f5e..3e92a77543e3d2162497e9f995f3adc2a01cb4dd 100644 --- a/tensorflow/python/estimator/canned/baseline.py +++ b/tensorflow/python/estimator/canned/baseline.py @@ -59,6 +59,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops.losses import losses from tensorflow.python.training import training_util +from tensorflow.python.util.tf_export import tf_export # The default learning rate of 0.3 is a historical artifact of the initial # implementation, but seems a reasonable choice. @@ -173,6 +174,7 @@ def _baseline_model_fn(features, labels, mode, head, optimizer, train_op_fn=train_op_fn) +@tf_export('estimator.BaselineClassifier') class BaselineClassifier(estimator.Estimator): """A classifier that can establish a simple baseline. @@ -275,6 +277,7 @@ class BaselineClassifier(estimator.Estimator): config=config) +@tf_export('estimator.BaselineRegressor') class BaselineRegressor(estimator.Estimator): """A regressor that can establish a simple baseline. diff --git a/tensorflow/python/estimator/canned/dnn.py b/tensorflow/python/estimator/canned/dnn.py index 0f274a23c03426fc431c15ac0a14617a4a65bb79..7043da8de036e5be27d223271c37e065d9ffbcdd 100644 --- a/tensorflow/python/estimator/canned/dnn.py +++ b/tensorflow/python/estimator/canned/dnn.py @@ -33,6 +33,7 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.ops.losses import losses from tensorflow.python.summary import summary from tensorflow.python.training import training_util +from tensorflow.python.util.tf_export import tf_export # The default learning rate of 0.05 is a historical artifact of the initial # implementation, but seems a reasonable choice. @@ -149,9 +150,7 @@ def _dnn_model_fn(features, config: `RunConfig` object to configure the runtime settings. Returns: - predictions: A dict of `Tensor` objects. - loss: A scalar containing the loss of the step. - train_op: The op for training. + An `EstimatorSpec` instance. Raises: ValueError: If features has the wrong type. @@ -198,6 +197,7 @@ def _dnn_model_fn(features, logits=logits) +@tf_export('estimator.DNNClassifier') class DNNClassifier(estimator.Estimator): """A classifier for TensorFlow DNN models. @@ -358,6 +358,7 @@ class DNNClassifier(estimator.Estimator): warm_start_from=warm_start_from) +@tf_export('estimator.DNNRegressor') class DNNRegressor(estimator.Estimator): """A regressor for TensorFlow DNN models. diff --git a/tensorflow/python/estimator/canned/dnn_linear_combined.py b/tensorflow/python/estimator/canned/dnn_linear_combined.py index 1a0f4c5c3931a6b41026470f30e7bdd381e5b37a..6d0fb96057ee93964ee3571bae3b878faad88882 100644 --- a/tensorflow/python/estimator/canned/dnn_linear_combined.py +++ b/tensorflow/python/estimator/canned/dnn_linear_combined.py @@ -37,6 +37,7 @@ from tensorflow.python.ops.losses import losses from tensorflow.python.summary import summary from tensorflow.python.training import sync_replicas_optimizer from tensorflow.python.training import training_util +from tensorflow.python.util.tf_export import tf_export # The default learning rates are a historical artifact of the initial # implementation. @@ -116,7 +117,7 @@ def _dnn_linear_combined_model_fn(features, config: `RunConfig` object to configure the runtime settings. Returns: - `ModelFnOps` + An `EstimatorSpec` instance. Raises: ValueError: If both `linear_feature_columns` and `dnn_features_columns` @@ -225,6 +226,7 @@ def _dnn_linear_combined_model_fn(features, logits=logits) +@tf_export('estimator.DNNLinearCombinedClassifier') class DNNLinearCombinedClassifier(estimator.Estimator): """An estimator for TensorFlow Linear and DNN joined classification models. @@ -405,6 +407,7 @@ class DNNLinearCombinedClassifier(estimator.Estimator): warm_start_from=warm_start_from) +@tf_export('estimator.DNNLinearCombinedRegressor') class DNNLinearCombinedRegressor(estimator.Estimator): """An estimator for TensorFlow Linear and DNN joined models for regression. diff --git a/tensorflow/python/estimator/canned/dnn_linear_combined_test.py b/tensorflow/python/estimator/canned/dnn_linear_combined_test.py index 84675bf2a4a1655026bbba37c5d7a63d2f788c46..d275695eb319117cf94aefd7038ab5ee685e05a9 100644 --- a/tensorflow/python/estimator/canned/dnn_linear_combined_test.py +++ b/tensorflow/python/estimator/canned/dnn_linear_combined_test.py @@ -26,7 +26,7 @@ import six from tensorflow.core.example import example_pb2 from tensorflow.core.example import feature_pb2 -from tensorflow.python.estimator import warm_starting_util +from tensorflow.python.estimator import estimator from tensorflow.python.estimator.canned import dnn_linear_combined from tensorflow.python.estimator.canned import dnn_testing_utils from tensorflow.python.estimator.canned import linear_testing_utils @@ -866,7 +866,7 @@ class DNNLinearCombinedWarmStartingTest(test.TestCase): learning_rate=0.0), # The provided regular expression will only warm-start the deep # portion of the model. - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=dnn_lc_classifier.model_dir, vars_to_warm_start='.*(dnn).*'))) diff --git a/tensorflow/python/estimator/canned/dnn_testing_utils.py b/tensorflow/python/estimator/canned/dnn_testing_utils.py index 706575985ff9e0fef94f110825ec11af33031ea3..9a7d088778e440b8aa4d9bc4f05a72aa86e7c106 100644 --- a/tensorflow/python/estimator/canned/dnn_testing_utils.py +++ b/tensorflow/python/estimator/canned/dnn_testing_utils.py @@ -27,8 +27,8 @@ import six from tensorflow.core.framework import summary_pb2 from tensorflow.python.client import session as tf_session +from tensorflow.python.estimator import estimator from tensorflow.python.estimator import model_fn -from tensorflow.python.estimator import warm_starting_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 @@ -828,7 +828,7 @@ class BaseDNNWarmStartingTest(object): optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0), # The provided regular expression will only warm-start the city # embedding, not the kernels and biases of the hidden weights. - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=dnn_classifier.model_dir, vars_to_warm_start='.*(city).*')) @@ -892,7 +892,7 @@ class BaseDNNWarmStartingTest(object): dimension=2) # We can create our VocabInfo object from the new and old occupation # FeatureColumn's. - occupation_vocab_info = warm_starting_util.VocabInfo( + occupation_vocab_info = estimator.VocabInfo( new_vocab=new_occupation.categorical_column.vocabulary_file, new_vocab_size=new_occupation.categorical_column.vocabulary_size, num_oov_buckets=new_occupation.categorical_column.num_oov_buckets, @@ -907,7 +907,7 @@ class BaseDNNWarmStartingTest(object): feature_columns=[occupation], n_classes=4, optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0), - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=dnn_classifier.model_dir, var_name_to_vocab_info={ OCCUPATION_EMBEDDING_NAME: occupation_vocab_info @@ -978,7 +978,7 @@ class BaseDNNWarmStartingTest(object): optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0), # The 'city' variable correspond to the 'locality' variable in the # previous model. - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=dnn_classifier.model_dir, var_name_to_prev_var_name={ CITY_EMBEDDING_NAME: diff --git a/tensorflow/python/estimator/canned/head.py b/tensorflow/python/estimator/canned/head.py index cb9e3fc6ca116ac0f48a37cea92fa4119754f324..8d742a2c6147e86619d4c0aad59b69459384bd4d 100644 --- a/tensorflow/python/estimator/canned/head.py +++ b/tensorflow/python/estimator/canned/head.py @@ -1156,6 +1156,7 @@ def _regression_head_with_mean_squared_error_loss( label_dimension=1, loss_reduction=losses.Reduction.SUM, loss_fn=None, + inverse_link_fn=None, name=None): """Creates a `_Head` for regression using the `mean_squared_error` loss. @@ -1174,10 +1175,16 @@ def _regression_head_with_mean_squared_error_loss( `[D0, D1, ... DN]`, `[D0, D1, ... DN, 1]` or `[D0, D1, ... DN, label_dimension]`. - Also supports custom `loss_fn`. `loss_fn` takes `(labels, logits)` or + 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]`. + Also supports custom `inverse_link_fn`, also known as 'mean function'. + `inverse_link_fn` takes `logits` as argument and returns predicted values. + This function is the inverse of the link function defined in + https://en.wikipedia.org/wiki/Generalized_linear_model#Link_function + Namely, for poisson regression, set `inverse_link_fn=tf.exp`. + Args: weight_column: A string or a `_NumericColumn` created by `tf.feature_column.numeric_column` defining feature column representing @@ -1188,7 +1195,9 @@ def _regression_head_with_mean_squared_error_loss( `[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. + loss_fn: Optional loss function. Defaults to `mean_squared_error`. + inverse_link_fn: Optional inverse link function, also known as 'mean + function'. Defaults to identity. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. @@ -1208,6 +1217,7 @@ def _regression_head_with_mean_squared_error_loss( label_dimension=label_dimension, loss_reduction=loss_reduction, loss_fn=loss_fn, + inverse_link_fn=inverse_link_fn, name=name) @@ -1220,6 +1230,7 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head): weight_column=None, loss_reduction=losses.Reduction.SUM, loss_fn=None, + inverse_link_fn=None, name=None): """`Head` for regression.""" if label_dimension < 1: @@ -1228,6 +1239,7 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head): self._weight_column = weight_column self._loss_reduction = loss_reduction self._loss_fn = loss_fn + self._inverse_link_fn = inverse_link_fn self._name = name @property @@ -1294,9 +1306,19 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head): # Predict. with ops.name_scope(self._name, 'head'): logits = _check_logits_final_dim(logits, self._logits_dimension) - predictions = {prediction_keys.PredictionKeys.PREDICTIONS: logits} + if self._inverse_link_fn: + predicted_value = self._inverse_link_fn(logits) + predictions = { + prediction_keys.PredictionKeys.PREDICTIONS: predicted_value, + prediction_keys.PredictionKeys.LOGITS: logits, + } + else: + predicted_value = logits + predictions = { + prediction_keys.PredictionKeys.PREDICTIONS: predicted_value} if mode == model_fn.ModeKeys.PREDICT: - regression_output = export_output.RegressionOutput(value=logits) + regression_output = export_output.RegressionOutput( + value=predicted_value) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, diff --git a/tensorflow/python/estimator/canned/head_test.py b/tensorflow/python/estimator/canned/head_test.py index 3a03770af498981a054c3df9155e83a60c7f0350..b40758f8fef13bad939b3fc2f296e848ade46f63 100644 --- a/tensorflow/python/estimator/canned/head_test.py +++ b/tensorflow/python/estimator/canned/head_test.py @@ -300,7 +300,12 @@ class MultiClassHeadWithSoftmaxCrossEntropyLoss(test.TestCase): features = {'x': values_2x3} # Static shape. - with self.assertRaisesRegexp(ValueError, 'Dimensions must be equal'): + with self.assertRaisesRegexp( + ValueError, + r'Shape mismatch: The shape of labels \(received \(3,\)\) should equal ' + r'the shape of logits except for the last dimension ' + r'\(received \(2, 3\)\)\.' + ): head.create_loss( features=features, mode=model_fn.ModeKeys.EVAL, @@ -2703,10 +2708,9 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase): self.assertIsNone(spec.loss) self.assertEqual({}, spec.eval_metric_ops) self.assertIsNone(spec.train_op) + default_serving_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY self.assertItemsEqual( - (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, - 'predict', - 'regression'), + (default_serving_key, 'predict', 'regression'), spec.export_outputs.keys()) _assert_no_hooks(self, spec) @@ -2714,6 +2718,54 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase): with self.test_session(): _initialize_variables(self, spec.scaffold) self.assertAllClose(logits, spec.predictions[prediction_key].eval()) + self.assertAllClose( + logits, spec.export_outputs[default_serving_key].value.eval()) + self.assertAllClose( + logits, spec.export_outputs['regression'].value.eval()) + self.assertAllClose( + logits, spec.export_outputs['predict'].outputs['predictions'].eval()) + + def test_predict_with_inverse_link_fn(self): + def _inverse_link_fn(logits): + return logits - 10. + head = head_lib._regression_head_with_mean_squared_error_loss( + inverse_link_fn=_inverse_link_fn) + + # Create estimator spec. + logits = np.array(((45,), (41,),), dtype=np.int32) + expected_predictions = np.array(((35,), (31,),), dtype=np.int32) + spec = head.create_estimator_spec( + features={'x': np.array(((42.,),), dtype=np.int32)}, + mode=model_fn.ModeKeys.PREDICT, + logits=logits) + + # Assert spec contains expected tensors. + keys = prediction_keys.PredictionKeys + self.assertItemsEqual( + (keys.PREDICTIONS, keys.LOGITS), spec.predictions.keys()) + self.assertEqual(dtypes.float32, spec.predictions[keys.PREDICTIONS].dtype) + self.assertEqual(dtypes.float32, spec.predictions[keys.LOGITS].dtype) + default_serving_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY + self.assertItemsEqual( + (default_serving_key, 'predict', 'regression'), + spec.export_outputs.keys()) + + # Assert predictions. + with self.test_session(): + _initialize_variables(self, spec.scaffold) + self.assertAllClose( + expected_predictions, spec.predictions[keys.PREDICTIONS].eval()) + self.assertAllClose(logits, spec.predictions[keys.LOGITS].eval()) + self.assertAllClose( + expected_predictions, + spec.export_outputs[default_serving_key].value.eval()) + self.assertAllClose( + expected_predictions, spec.export_outputs['regression'].value.eval()) + self.assertAllClose( + expected_predictions, + spec.export_outputs['predict'].outputs['predictions'].eval()) + self.assertAllClose( + logits, spec.export_outputs['predict'].outputs['logits'].eval()) def test_eval_create_loss(self): head = head_lib._regression_head_with_mean_squared_error_loss() diff --git a/tensorflow/python/estimator/canned/linear.py b/tensorflow/python/estimator/canned/linear.py index a5b1172e729240a2ea02fa1d4330420786c2686c..a2f24ef27044680fe93b176b5207593165d0d109 100644 --- a/tensorflow/python/estimator/canned/linear.py +++ b/tensorflow/python/estimator/canned/linear.py @@ -34,6 +34,7 @@ from tensorflow.python.ops.losses import losses from tensorflow.python.summary import summary from tensorflow.python.training import ftrl from tensorflow.python.training import training_util +from tensorflow.python.util.tf_export import tf_export # The default learning rate of 0.2 is a historical artifact of the initial @@ -170,6 +171,7 @@ def _linear_model_fn(features, labels, mode, head, feature_columns, optimizer, logits=logits) +@tf_export('estimator.LinearClassifier') class LinearClassifier(estimator.Estimator): """Linear classifier model. @@ -322,6 +324,7 @@ class LinearClassifier(estimator.Estimator): warm_start_from=warm_start_from) +@tf_export('estimator.LinearRegressor') class LinearRegressor(estimator.Estimator): """An estimator for TensorFlow Linear regression problems. diff --git a/tensorflow/python/estimator/canned/linear_testing_utils.py b/tensorflow/python/estimator/canned/linear_testing_utils.py index 3e9183cf1b633757074377472e9b4cac953e04a1..8e506a763133c3f8c15f49f253f84de5470088c5 100644 --- a/tensorflow/python/estimator/canned/linear_testing_utils.py +++ b/tensorflow/python/estimator/canned/linear_testing_utils.py @@ -31,7 +31,6 @@ from tensorflow.core.example import feature_pb2 from tensorflow.python.client import session as tf_session from tensorflow.python.estimator import estimator from tensorflow.python.estimator import run_config -from tensorflow.python.estimator import warm_starting_util from tensorflow.python.estimator.canned import linear from tensorflow.python.estimator.canned import metric_keys from tensorflow.python.estimator.export import export @@ -1968,7 +1967,7 @@ class BaseLinearWarmStartingTest(object): optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0), # The provided regular expression will only warm-start the age variable # and not the bias. - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=linear_classifier.model_dir, vars_to_warm_start='.*(age).*')) @@ -2016,7 +2015,7 @@ class BaseLinearWarmStartingTest(object): vocabulary_size=len(new_vocab_list)) # We can create our VocabInfo object from the new and old occupation # FeatureColumn's. - occupation_vocab_info = warm_starting_util.VocabInfo( + occupation_vocab_info = estimator.VocabInfo( new_vocab=new_occupation.vocabulary_file, new_vocab_size=new_occupation.vocabulary_size, num_oov_buckets=new_occupation.num_oov_buckets, @@ -2030,7 +2029,7 @@ class BaseLinearWarmStartingTest(object): feature_columns=[occupation], n_classes=4, optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0), - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=linear_classifier.model_dir, var_name_to_vocab_info={ OCCUPATION_WEIGHT_NAME: occupation_vocab_info @@ -2082,7 +2081,7 @@ class BaseLinearWarmStartingTest(object): optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0), # The 'age' variable correspond to the 'age_in_years' variable in the # previous model. - warm_start_from=warm_starting_util.WarmStartSettings( + warm_start_from=estimator.WarmStartSettings( ckpt_to_initialize_from=linear_classifier.model_dir, var_name_to_prev_var_name={ AGE_WEIGHT_NAME: AGE_WEIGHT_NAME.replace('age', 'age_in_years') diff --git a/tensorflow/python/estimator/canned/parsing_utils.py b/tensorflow/python/estimator/canned/parsing_utils.py index f153272947ca427b25b00e6df4741d7ada5790df..74e5e5a1bed80229c68daa3ff33ee7af4004bf47 100644 --- a/tensorflow/python/estimator/canned/parsing_utils.py +++ b/tensorflow/python/estimator/canned/parsing_utils.py @@ -23,8 +23,10 @@ import six from tensorflow.python.feature_column import feature_column as fc from tensorflow.python.framework import dtypes from tensorflow.python.ops import parsing_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export('estimator.classifier_parse_example_spec') def classifier_parse_example_spec(feature_columns, label_key, label_dtype=dtypes.int64, @@ -164,6 +166,7 @@ def classifier_parse_example_spec(feature_columns, return parsing_spec +@tf_export('estimator.regressor_parse_example_spec') def regressor_parse_example_spec(feature_columns, label_key, label_dtype=dtypes.float32, diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py index 78d74b63d3ebea29e2dc3ab8f655efcc7ab8e130..5245a050a1ed87af157dbd2008b52122798f9491 100644 --- a/tensorflow/python/estimator/estimator.py +++ b/tensorflow/python/estimator/estimator.py @@ -19,6 +19,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import copy import os import tempfile @@ -35,7 +36,6 @@ from tensorflow.python.eager import context from tensorflow.python.estimator import model_fn as model_fn_lib from tensorflow.python.estimator import run_config from tensorflow.python.estimator import util -from tensorflow.python.estimator import warm_starting_util from tensorflow.python.estimator.export.export import build_all_signature_defs from tensorflow.python.estimator.export.export import get_temp_export_dir from tensorflow.python.estimator.export.export import get_timestamped_export_dir @@ -49,20 +49,24 @@ from tensorflow.python.saved_model import builder as saved_model_builder from tensorflow.python.saved_model import tag_constants from tensorflow.python.summary import summary from tensorflow.python.summary.writer import writer_cache +from tensorflow.python.training import device_setter from tensorflow.python.training import evaluation from tensorflow.python.training import monitored_session from tensorflow.python.training import saver from tensorflow.python.training import training from tensorflow.python.training import training_util +from tensorflow.python.training import warm_starting_util from tensorflow.python.util import compat from tensorflow.python.util import compat_internal from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export _VALID_MODEL_FN_ARGS = set( ['features', 'labels', 'mode', 'params', 'self', 'config']) +@tf_export('estimator.Estimator') class Estimator(object): """Estimator class to train and evaluate TensorFlow models. @@ -135,8 +139,8 @@ class Estimator(object): to configure Estimators from hyper parameter tuning. * `config`: Optional configuration object. Will receive what is passed to Estimator in `config` parameter, or the default `config`. - Allows updating things in your model_fn based on configuration - such as `num_ps_replicas`, or `model_dir`. + Allows updating things in your `model_fn` based on + configuration such as `num_ps_replicas`, or `model_dir`. * Returns: `EstimatorSpec` @@ -163,7 +167,7 @@ class Estimator(object): ValueError: if this is called via a subclass and if that class overrides a member of `Estimator`. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( 'Estimators are not supported when eager execution is enabled.') @@ -214,8 +218,8 @@ class Estimator(object): self._params = copy.deepcopy(params or {}) # pylint: disable=protected-access - self._warm_start_settings = ( - warm_starting_util._get_default_warm_start_settings(warm_start_from)) + self._warm_start_settings = _get_default_warm_start_settings( + warm_start_from) # pylint: enable=protected-access @property @@ -297,11 +301,11 @@ class Estimator(object): * A 'tf.data.Dataset' object: Outputs of `Dataset` object must be a tuple (features, labels) with same constraints as below. - * A tuple (features, labels): Where features is a `Tensor` or a - dictionary of string feature name to `Tensor` and labels is a + * A tuple (features, labels): Where `features` is a `Tensor` or a + dictionary of string feature name to `Tensor` and `labels` is a `Tensor` or a dictionary of string label name to `Tensor`. Both - features and labels are consumed by `model_fn`. They should satisfy - the expectation of `model_fn` from inputs. + `features` and `labels` are consumed by `model_fn`. They should + satisfy the expectation of `model_fn` from inputs. hooks: List of `SessionRunHook` subclass instances. Used for callbacks inside the training loop. @@ -377,11 +381,11 @@ class Estimator(object): * A 'tf.data.Dataset' object: Outputs of `Dataset` object must be a tuple (features, labels) with same constraints as below. - * A tuple (features, labels): Where features is a `Tensor` or a - dictionary of string feature name to `Tensor` and labels is a + * A tuple (features, labels): Where `features` is a `Tensor` or a + dictionary of string feature name to `Tensor` and `labels` is a `Tensor` or a dictionary of string label name to `Tensor`. Both - features and labels are consumed by `model_fn`. They should satisfy - the expectation of `model_fn` from inputs. + `features` and `labels` are consumed by `model_fn`. They should + satisfy the expectation of `model_fn` from inputs. steps: Number of steps for which to evaluate model. If `None`, evaluates until `input_fn` raises an end-of-input exception. @@ -425,7 +429,8 @@ class Estimator(object): input_fn, predict_keys=None, hooks=None, - checkpoint_path=None): + checkpoint_path=None, + yield_single_examples=True): """Yields predictions for given features. Args: @@ -451,13 +456,18 @@ class Estimator(object): inside the prediction call. checkpoint_path: Path of a specific checkpoint to predict. If `None`, the latest checkpoint in `model_dir` is used. + yield_single_examples: If False, yield the whole batch as returned by the + `model_fn` instead of decomposing the batch into individual elements. + This is useful if `model_fn` returns some tensors whose first dimension + is not equal to the batch size. Yields: Evaluated values of `predictions` tensors. Raises: - ValueError: Could not find a trained model in model_dir. - ValueError: if batch length of predictions are not same. + ValueError: Could not find a trained model in `model_dir`. + ValueError: If batch length of predictions is not the same and + `yield_single_examples` is True. ValueError: If there is a conflict between `predict_keys` and `predictions`. For example if `predict_keys` is not `None` but `EstimatorSpec.predictions` is not a `dict`. @@ -490,7 +500,9 @@ class Estimator(object): hooks=all_hooks) as mon_sess: while not mon_sess.should_stop(): preds_evaluated = mon_sess.run(predictions) - if not isinstance(predictions, dict): + if not yield_single_examples: + yield preds_evaluated + elif not isinstance(predictions, dict): for pred in preds_evaluated: yield pred else: @@ -502,9 +514,11 @@ class Estimator(object): def _assert_members_are_not_overridden(self): """Asserts members of `Estimator` are not overridden.""" - allowed_overrides = set(['_call_input_fn', '_create_global_step', - '_convert_train_steps_to_hooks', - '_convert_eval_steps_to_hooks']) + allowed_overrides = set([ + '_call_input_fn', '_create_global_step', + '_convert_train_steps_to_hooks', '_convert_eval_steps_to_hooks', + '_tf_api_names', '_validate_features_in_predict_input' + ]) estimator_members = set([m for m in Estimator.__dict__.keys() if not m.startswith('__')]) subclass_members = set(self.__class__.__dict__.keys()) @@ -558,7 +572,7 @@ class Estimator(object): export_dir_base: A string containing a directory in which to create timestamped subdirectories containing exported SavedModels. serving_input_receiver_fn: A function that takes no argument and - returns a `ServingInputReceiver`. + returns a `ServingInputReceiver` or `TensorServingInputReceiver`. assets_extra: A dict specifying how to populate the assets.extra directory within the exported SavedModel, or `None` if no extra assets are needed. as_text: whether to write the SavedModel proto in text format. @@ -613,7 +627,6 @@ class Estimator(object): sharded=True) saver_for_restore.restore(session, checkpoint_path) - # TODO(b/36111876): replace legacy_init_op with main_op mechanism # pylint: disable=protected-access local_init_op = ( estimator_spec.scaffold.local_init_op or @@ -657,11 +670,14 @@ class Estimator(object): # Unconditionally drop the label (the second element of result). result = result[0] + self._validate_features_in_predict_input(result) + return result, input_hooks + + def _validate_features_in_predict_input(self, result): if not _has_dataset_or_queue_runner(result): logging.warning('Input graph does not use tf.data.Dataset or contain a ' 'QueueRunner. That means predict yields forever. ' 'This is probably a mistake.') - return result, input_hooks def _get_features_and_labels_from_input_fn(self, input_fn, mode): """Extracts the `features` and labels from return values of `input_fn`.""" @@ -709,7 +725,7 @@ class Estimator(object): """Creates the global step tensor in graph. The global step tensor must be an integer type with name 'global_step' and - be added to the collection ${tf.GraphKeys.GLOBAL_STEP}. + be added to the collection @{tf.GraphKeys.GLOBAL_STEP}. Args: graph: The graph in which to create the global step tensor. @@ -815,7 +831,7 @@ class Estimator(object): logging.info('Warm-starting with WarmStartSettings: %s' % (self._warm_start_settings,)) # pylint: disable=protected-access - warm_starting_util._warm_start(self._warm_start_settings) + warm_starting_util.warm_start(*self._warm_start_settings) # pylint: enable=protected-access # Check if the user created a loss summary, and add one if they didn't. # We assume here that the summary is called 'loss'. If it is not, we will @@ -996,13 +1012,6 @@ def _get_replica_device_setter(config): Returns: A replica device setter, or None. """ - ps_ops = [ - 'Variable', 'VariableV2', 'AutoReloadVariable', 'MutableHashTable', - 'MutableHashTableV2', 'MutableHashTableOfTensors', - 'MutableHashTableOfTensorsV2', 'MutableDenseHashTable', - 'MutableDenseHashTableV2' - ] - if config.task_type: worker_device = '/job:%s/task:%d' % (config.task_type, config.task_id) else: @@ -1013,7 +1022,7 @@ def _get_replica_device_setter(config): ps_tasks=config.num_ps_replicas, worker_device=worker_device, merge_devices=True, - ps_ops=ps_ops, + ps_ops=list(device_setter.STANDARD_PS_OPS), cluster=config.cluster_spec) else: return None @@ -1103,11 +1112,11 @@ def _write_dict_to_summary(output_dir, isinstance(dictionary[key], np.int32) or isinstance(dictionary[key], int)): summary_proto.value.add(tag=key, simple_value=int(dictionary[key])) - elif isinstance(dictionary[key], six.string_types): + elif isinstance(dictionary[key], six.binary_type): try: summ = summary_pb2.Summary.FromString(dictionary[key]) for i, _ in enumerate(summ.value): - summ.value[i].tag = key + summ.value[i].tag = '%s/%d' % (key, i) summary_proto.value.extend(summ.value) except message.DecodeError: logging.warn('Skipping summary for %s, cannot parse string to Summary.', @@ -1144,3 +1153,187 @@ class _DatasetInitializerHook(training.SessionRunHook): def after_create_session(self, session, coord): del coord session.run(self._initializer) + +VocabInfo = warm_starting_util.VocabInfo # pylint: disable=invalid-name + + +@tf_export('estimator.WarmStartSettings') +class WarmStartSettings( + collections.namedtuple('WarmStartSettings', [ + 'ckpt_to_initialize_from', + 'vars_to_warm_start', + 'var_name_to_vocab_info', + 'var_name_to_prev_var_name', + ])): + """Settings for warm-starting in Estimators. + + Example Use with canned `DNNEstimator`: + + ``` + emb_vocab_file = tf.feature_column.embedding_column( + tf.feature_column.categorical_column_with_vocabulary_file( + "sc_vocab_file", "new_vocab.txt", vocab_size=100), + dimension=8) + emb_vocab_list = tf.feature_column.embedding_column( + tf.feature_column.categorical_column_with_vocabulary_list( + "sc_vocab_list", vocabulary_list=["a", "b"]), + dimension=8) + estimator = tf.estimator.DNNClassifier( + hidden_units=[128, 64], feature_columns=[emb_vocab_file, emb_vocab_list], + warm_start_from=ws) + ``` + + where `ws` could be defined as: + + Warm-start all weights in the model (input layer and hidden weights). + Either the directory or a specific checkpoint can be provided (in the case + of the former, the latest checkpoint will be used): + + ``` + ws = WarmStartSettings(ckpt_to_initialize_from="/tmp") + ws = WarmStartSettings(ckpt_to_initialize_from="/tmp/model-1000") + ``` + + Warm-start only the embeddings (input layer): + + ``` + ws = WarmStartSettings(ckpt_to_initialize_from="/tmp", + vars_to_warm_start=".*input_layer.*") + ``` + + Warm-start all weights but the embedding parameters corresponding to + `sc_vocab_file` have a different vocab from the one used in the current + model: + + ``` + vocab_info = tf.estimator.VocabInfo( + new_vocab=sc_vocab_file.vocabulary_file, + new_vocab_size=sc_vocab_file.vocabulary_size, + num_oov_buckets=sc_vocab_file.num_oov_buckets, + old_vocab="old_vocab.txt" + ) + ws = WarmStartSettings( + ckpt_to_initialize_from="/tmp", + var_name_to_vocab_info={ + "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info + }) + ``` + + Warm-start only `sc_vocab_file` embeddings (and no other variables), which + have a different vocab from the one used in the current model: + + ``` + vocab_info = tf.estimator.VocabInfo( + new_vocab=sc_vocab_file.vocabulary_file, + new_vocab_size=sc_vocab_file.vocabulary_size, + num_oov_buckets=sc_vocab_file.num_oov_buckets, + old_vocab="old_vocab.txt" + ) + ws = WarmStartSettings( + ckpt_to_initialize_from="/tmp", + vars_to_warm_start=None, + var_name_to_vocab_info={ + "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info + }) + ``` + + Warm-start all weights but the parameters corresponding to `sc_vocab_file` + have a different vocab from the one used in current checkpoint, and only + 100 of those entries were used: + + ``` + vocab_info = tf.estimator.VocabInfo( + new_vocab=sc_vocab_file.vocabulary_file, + new_vocab_size=sc_vocab_file.vocabulary_size, + num_oov_buckets=sc_vocab_file.num_oov_buckets, + old_vocab="old_vocab.txt", + old_vocab_size=100 + ) + ws = WarmStartSettings( + ckpt_to_initialize_from="/tmp", + var_name_to_vocab_info={ + "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info + }) + ``` + + Warm-start all weights but the parameters corresponding to `sc_vocab_file` + have a different vocab from the one used in current checkpoint and the + parameters corresponding to `sc_vocab_list` have a different name from the + current checkpoint: + + ``` + vocab_info = tf.estimator.VocabInfo( + new_vocab=sc_vocab_file.vocabulary_file, + new_vocab_size=sc_vocab_file.vocabulary_size, + num_oov_buckets=sc_vocab_file.num_oov_buckets, + old_vocab="old_vocab.txt", + old_vocab_size=100 + ) + ws = WarmStartSettings( + ckpt_to_initialize_from="/tmp", + var_name_to_vocab_info={ + "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info + }, + var_name_to_prev_var_name={ + "input_layer/sc_vocab_list_embedding/embedding_weights": + "old_tensor_name" + }) + ``` + + Attributes: + ckpt_to_initialize_from: [Required] A string specifying the directory with + checkpoint file(s) or path to checkpoint from which to warm-start the + model parameters. + vars_to_warm_start: [Optional] A regular expression that captures which + variables to warm-start (see tf.get_collection). Defaults to `'.*'`, + which warm-starts all variables. If `None` is explicitly given, only + variables specified in `var_name_to_vocab_info` will be warm-started. + var_name_to_vocab_info: [Optional] Dict of variable names (strings) to + VocabInfo. The variable names should be "full" variables, not the names + of the partitions. If not explicitly provided, the variable is assumed to + have no vocabulary. + var_name_to_prev_var_name: [Optional] Dict of variable names (strings) to + name of the previously-trained variable in `ckpt_to_initialize_from`. If + not explicitly provided, the name of the variable is assumed to be same + between previous checkpoint and current model. + """ + + def __new__(cls, + ckpt_to_initialize_from, + vars_to_warm_start='.*', + var_name_to_vocab_info=None, + var_name_to_prev_var_name=None): + if not ckpt_to_initialize_from: + raise ValueError( + '`ckpt_to_initialize_from` MUST be set in WarmStartSettings') + return super(WarmStartSettings, cls).__new__( + cls, + ckpt_to_initialize_from, + vars_to_warm_start, + var_name_to_vocab_info or {}, + var_name_to_prev_var_name or {}, + ) + + +def _get_default_warm_start_settings(warm_start_from): + """Returns default WarmStartSettings. + + Args: + warm_start_from: Either a string representing the filepath of a checkpoint + to initialize from, or an instance of WarmStartSettings. + + Returns: + Either None or an instance of WarmStartSettings. + + Raises: + ValueError: If warm_start_from is not None but is neither a string nor an + instance of WarmStartSettings. + """ + if warm_start_from is None: + return None + if isinstance(warm_start_from, six.string_types): + return WarmStartSettings(ckpt_to_initialize_from=warm_start_from) + elif isinstance(warm_start_from, WarmStartSettings): + return warm_start_from + else: + raise ValueError('warm_start_from must be a string or a WarmStartSettings') diff --git a/tensorflow/python/estimator/estimator_lib.py b/tensorflow/python/estimator/estimator_lib.py index 01699e7399c4089281e9ece76e534e1f82692257..be8930b3cbcd89dbb31dffde0a7a5ecfb64fcd8b 100644 --- a/tensorflow/python/estimator/estimator_lib.py +++ b/tensorflow/python/estimator/estimator_lib.py @@ -30,6 +30,8 @@ from tensorflow.python.estimator.canned.linear import LinearRegressor from tensorflow.python.estimator.canned.parsing_utils import classifier_parse_example_spec from tensorflow.python.estimator.canned.parsing_utils import regressor_parse_example_spec from tensorflow.python.estimator.estimator import Estimator +from tensorflow.python.estimator.estimator import VocabInfo +from tensorflow.python.estimator.estimator import WarmStartSettings from tensorflow.python.estimator.export import export_lib as export from tensorflow.python.estimator.exporter import Exporter from tensorflow.python.estimator.exporter import FinalExporter @@ -41,8 +43,6 @@ from tensorflow.python.estimator.run_config import RunConfig from tensorflow.python.estimator.training import EvalSpec from tensorflow.python.estimator.training import train_and_evaluate from tensorflow.python.estimator.training import TrainSpec -from tensorflow.python.estimator.warm_starting_util import VocabInfo -from tensorflow.python.estimator.warm_starting_util import WarmStartSettings from tensorflow.python.util.all_util import remove_undocumented diff --git a/tensorflow/python/estimator/estimator_test.py b/tensorflow/python/estimator/estimator_test.py index 39a5b998ebdcccfbeddf0fc96dab44dc91a289fa..f4255091bf6c44916789a182e60e583171ad5e6b 100644 --- a/tensorflow/python/estimator/estimator_test.py +++ b/tensorflow/python/estimator/estimator_test.py @@ -48,6 +48,7 @@ from tensorflow.python.ops import check_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import lookup_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics as metrics_lib from tensorflow.python.ops import parsing_ops from tensorflow.python.ops import state_ops @@ -80,18 +81,18 @@ def dummy_model_fn(features, labels, params): _, _, _ = features, labels, params -def check_eventfile_for_keyword(keyword, est): +def check_eventfile_for_keyword(keyword, dir_): """Checks event files for the keyword.""" writer_cache.FileWriterCache.clear() # Get last Event written. - event_paths = glob.glob(os.path.join(est.model_dir, 'events*')) + event_paths = glob.glob(os.path.join(dir_, 'events*')) last_event = None for last_event in summary_iterator.summary_iterator(event_paths[-1]): if last_event.summary is not None: - if last_event.summary.value: - if keyword in last_event.summary.value[0].tag: + for value in last_event.summary.value: + if keyword in value.tag: return True return False @@ -610,7 +611,7 @@ class EstimatorTrainTest(test.TestCase): # Make sure nothing is stuck in limbo. writer_cache.FileWriterCache.clear() - if check_eventfile_for_keyword('loss', est): + if check_eventfile_for_keyword('loss', est.model_dir): return self.fail('{} should be part of reported summaries.'.format('loss')) @@ -1268,10 +1269,10 @@ class EstimatorEvaluateTest(test.TestCase): _, _ = features, labels global_step = training.get_global_step() - image = array_ops.zeros([1, 3, 3, 1]) + image = array_ops.zeros([5, 3, 3, 1]) eval_metric_ops = { - 'image': (summary.image('image', image, max_outputs=1), - constant_op.constant(1)) + 'foo': (summary.image('image', image, max_outputs=3), + constant_op.constant(1)) } return model_fn_lib.EstimatorSpec( mode, @@ -1290,10 +1291,11 @@ class EstimatorEvaluateTest(test.TestCase): # Make sure nothing is stuck in limbo. writer_cache.FileWriterCache.clear() - # Get last Event written. - if check_eventfile_for_keyword('image', est): - return - self.fail('{} should be part of reported summaries.'.format('image')) + # Get last evaluation Event written. + for key in ['foo/0', 'foo/1', 'foo/2']: + self.assertTrue( + check_eventfile_for_keyword(key, os.path.join(est.model_dir, 'eval')), + '{} should be part of reported summaries.'.format(key)) class EstimatorPredictTest(test.TestCase): @@ -1472,6 +1474,27 @@ class EstimatorPredictTest(test.TestCase): 'Batch length of predictions should be same'): next(est.predict(dummy_input_fn)) + def test_iterate_batches(self): + + def _model_fn(features, labels, mode): + _, _ = features, labels + return model_fn_lib.EstimatorSpec( + mode, + loss=constant_op.constant(0.), + train_op=state_ops.assign_add(training.get_global_step(), 1), + predictions={ + # First dim is different but the prediction should still work + 'y1': array_ops.zeros(shape=[3]), + 'y2': array_ops.zeros(shape=[5, 3]) + }) + + est = estimator.Estimator(model_fn=_model_fn) + est.train(dummy_input_fn, steps=1) + + predictions = next(est.predict(dummy_input_fn, yield_single_examples=False)) + self.assertAllEqual(predictions['y1'].shape, [3]) + self.assertAllEqual(predictions['y2'].shape, [5, 3]) + def test_predict_keys_defined_for_tensor(self): def _model_fn(features, labels, mode): @@ -1914,6 +1937,60 @@ class EstimatorExportTest(test.TestCase): # cleanup gfile.DeleteRecursively(tmpdir) + def test_export_savedmodel_tensor_features(self): + """Test that models accepting a single raw Tensor can be exported. + + See https://github.com/tensorflow/tensorflow/issues/11674 + + If the model_fn and receiver_fn accept raw tensors rather than dictionaries + as input, export_savedmodel should be okay with that, too. + + """ + + tmpdir = tempfile.mkdtemp() + + def _input_fn_tensor_features(): + t = array_ops.constant([1, 2, 3], dtype=dtypes.float32, shape=[1, 3]) + return (t, None) + + def _model_fn_tensor_features(features, labels, mode): + _ = labels + prediction = math_ops.matmul(features, features, transpose_b=True) + + return model_fn_lib.EstimatorSpec( + mode, + predictions=prediction, + loss=constant_op.constant(1.), + train_op=state_ops.assign_add(training.get_global_step(), 1), + export_outputs={ + 'test': export_output.PredictOutput({'prediction': prediction}) + }) + + def _serving_input_receiver_fn(): + feat = array_ops.placeholder(dtype=dtypes.float32) + return export.TensorServingInputReceiver( + features=feat, receiver_tensors=feat) + + est = estimator.Estimator(model_fn=_model_fn_tensor_features) + est.train(input_fn=_input_fn_tensor_features, steps=1) + + # Perform the export. + export_dir_base = os.path.join( + compat.as_bytes(tmpdir), compat.as_bytes('export')) + export_dir = est.export_savedmodel( + export_dir_base, _serving_input_receiver_fn) + + # Restore, to validate that the export was well-formed. + with ops.Graph().as_default() as graph: + with session.Session(graph=graph) as sess: + loader.load(sess, [tag_constants.SERVING], export_dir) + graph_ops = [x.name.lower() for x in graph.get_operations()] + self.assertTrue('const' in graph_ops) + self.assertTrue('matmul' in graph_ops) + + # Clean up. + gfile.DeleteRecursively(tmpdir) + def test_scaffold_is_used_for_saver(self): tmpdir = tempfile.mkdtemp() diff --git a/tensorflow/python/estimator/export/export.py b/tensorflow/python/estimator/export/export.py index 51075731ddc52a55799958c3bfa6140f77404541..9206a4964b3b7a6e3cc1e0f9e965a197be78c4ba 100644 --- a/tensorflow/python/estimator/export/export.py +++ b/tensorflow/python/estimator/export/export.py @@ -21,27 +21,28 @@ from __future__ import print_function import collections import os -import time import six +from tensorflow.python.estimator import util from tensorflow.python.framework import dtypes 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 parsing_ops -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.util import compat +from tensorflow.python.util.tf_export import tf_export _SINGLE_FEATURE_DEFAULT_NAME = 'feature' _SINGLE_RECEIVER_DEFAULT_NAME = 'input' +@tf_export('estimator.export.ServingInputReceiver') class ServingInputReceiver(collections.namedtuple( 'ServingInputReceiver', ['features', 'receiver_tensors', 'receiver_tensors_alternatives'])): @@ -118,6 +119,63 @@ class ServingInputReceiver(collections.namedtuple( receiver_tensors_alternatives=receiver_tensors_alternatives) +@tf_export('estimator.export.TensorServingInputReceiver') +class TensorServingInputReceiver(collections.namedtuple( + 'TensorServingInputReceiver', + ['features', 'receiver_tensors', 'receiver_tensors_alternatives'])): + """A return type for a serving_input_receiver_fn. + + This is for use with models that expect a single `Tensor` or `SparseTensor` + as an input feature, as opposed to a dict of features. + + The normal `ServingInputReceiver` always returns a feature dict, even if it + contains only one entry, and so can be used only with models that accept such + a dict. For models that accept only a single raw feature, the + `serving_input_receiver_fn` provided to `Estimator.export_savedmodel()` should + return this `TensorServingInputReceiver` instead. See: + https://github.com/tensorflow/tensorflow/issues/11674 + + Note that the receiver_tensors and receiver_tensor_alternatives arguments + will be automatically converted to the dict representation in either case, + because the SavedModel format requires each input `Tensor` to have a name + (provided by the dict key). + + The expected return values are: + features: A single `Tensor` or `SparseTensor`, representing the feature + to be passed to the model. + receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying + input nodes where this receiver expects to be fed by default. Typically, + this is a single placeholder expecting serialized `tf.Example` protos. + receiver_tensors_alternatives: a dict of string to additional + groups of receiver tensors, each of which may be a `Tensor` or a dict of + string to `Tensor`. These named receiver tensor alternatives generate + additional serving signatures, which may be used to feed inputs at + different points within the input receiver subgraph. A typical usage is + to allow feeding raw feature `Tensor`s *downstream* of the + tf.parse_example() op. Defaults to None. + """ + + def __new__(cls, features, receiver_tensors, + receiver_tensors_alternatives=None): + if features is None: + raise ValueError('features must be defined.') + if not (isinstance(features, ops.Tensor) + or isinstance(features, sparse_tensor.SparseTensor)): + raise ValueError('feature must be a Tensor or SparseTensor.') + + receiver = ServingInputReceiver( + features=features, + receiver_tensors=receiver_tensors, + receiver_tensors_alternatives=receiver_tensors_alternatives) + + return super(TensorServingInputReceiver, cls).__new__( + cls, + features=receiver.features[_SINGLE_FEATURE_DEFAULT_NAME], + receiver_tensors=receiver.receiver_tensors, + receiver_tensors_alternatives=receiver.receiver_tensors_alternatives) + + +@tf_export('estimator.export.build_parsing_serving_input_receiver_fn') def build_parsing_serving_input_receiver_fn(feature_spec, default_batch_size=None): """Build a serving_input_receiver_fn expecting fed tf.Examples. @@ -146,6 +204,7 @@ def build_parsing_serving_input_receiver_fn(feature_spec, return serving_input_receiver_fn +@tf_export('estimator.export.build_raw_serving_input_receiver_fn') def build_raw_serving_input_receiver_fn(features, default_batch_size=None): """Build a serving_input_receiver_fn expecting feature Tensors. @@ -269,13 +328,6 @@ def _log_signature_report(signature_def_map, excluded_signatures): logging.warn('Export includes no default signature!') -# When we create a timestamped directory, there is a small chance that the -# directory already exists because another worker is also writing exports. -# In this case we just wait one second to get a new timestamp and try again. -# If this fails several times in a row, then something is seriously wrong. -MAX_DIRECTORY_CREATION_ATTEMPTS = 10 - - def get_timestamped_export_dir(export_dir_base): """Builds a path to a new subdirectory within the base directory. @@ -294,25 +346,7 @@ def get_timestamped_export_dir(export_dir_base): RuntimeError: if repeated attempts fail to obtain a unique timestamped directory name. """ - attempts = 0 - while attempts < MAX_DIRECTORY_CREATION_ATTEMPTS: - export_timestamp = int(time.time()) - - export_dir = os.path.join( - compat.as_bytes(export_dir_base), - compat.as_bytes(str(export_timestamp))) - if not gfile.Exists(export_dir): - # Collisions are still possible (though extremely unlikely): this - # directory is not actually created yet, but it will be almost - # instantly on return from this function. - return export_dir - time.sleep(1) - attempts += 1 - logging.warn( - 'Export directory {} already exists; retrying (attempt {}/{})'.format( - export_dir, attempts, MAX_DIRECTORY_CREATION_ATTEMPTS)) - raise RuntimeError('Failed to obtain a unique export directory name after ' - '{} attempts.'.format(MAX_DIRECTORY_CREATION_ATTEMPTS)) + return util.get_timestamped_dir(export_dir_base) def get_temp_export_dir(timestamped_export_dir): diff --git a/tensorflow/python/estimator/export/export_lib.py b/tensorflow/python/estimator/export/export_lib.py index 99cd81d678bc04e7ed52de721a1fdf799c116795..226fc97fd3a3aefe61c4b88088873ce7489168c7 100644 --- a/tensorflow/python/estimator/export/export_lib.py +++ b/tensorflow/python/estimator/export/export_lib.py @@ -22,6 +22,7 @@ from __future__ import print_function from tensorflow.python.estimator.export.export import build_parsing_serving_input_receiver_fn from tensorflow.python.estimator.export.export import build_raw_serving_input_receiver_fn from tensorflow.python.estimator.export.export import ServingInputReceiver +from tensorflow.python.estimator.export.export import TensorServingInputReceiver from tensorflow.python.estimator.export.export_output import ClassificationOutput from tensorflow.python.estimator.export.export_output import ExportOutput from tensorflow.python.estimator.export.export_output import PredictOutput @@ -34,6 +35,7 @@ _allowed_symbols = [ 'build_parsing_serving_input_receiver_fn', 'build_raw_serving_input_receiver_fn', 'ServingInputReceiver', + 'TensorServingInputReceiver', 'ClassificationOutput', 'ExportOutput', 'PredictOutput', diff --git a/tensorflow/python/estimator/export/export_output.py b/tensorflow/python/estimator/export/export_output.py index 863af6d41d985043542b03375372fe564c283b82..87b964be37197dac99b8ce4398cbdaf3b4989c7f 100644 --- a/tensorflow/python/estimator/export/export_output.py +++ b/tensorflow/python/estimator/export/export_output.py @@ -26,8 +26,10 @@ import six from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.saved_model import signature_def_utils +from tensorflow.python.util.tf_export import tf_export +@tf_export('estimator.export.ExportOutput') class ExportOutput(object): """Represents an output of a model that can be served. @@ -50,6 +52,7 @@ class ExportOutput(object): pass +@tf_export('estimator.export.ClassificationOutput') class ClassificationOutput(ExportOutput): """Represents the output of a classification head. @@ -118,6 +121,7 @@ class ClassificationOutput(ExportOutput): examples, self.classes, self.scores) +@tf_export('estimator.export.RegressionOutput') class RegressionOutput(ExportOutput): """Represents the output of a regression head.""" @@ -153,6 +157,7 @@ class RegressionOutput(ExportOutput): _SINGLE_OUTPUT_DEFAULT_NAME = 'output' +@tf_export('estimator.export.PredictOutput') class PredictOutput(ExportOutput): """Represents the output of a generic prediction head. diff --git a/tensorflow/python/estimator/export/export_test.py b/tensorflow/python/estimator/export/export_test.py index 8442bf04accbd0bc15f5958069bf3060debd42bc..eb9688bc973666554b6057f5f546b9a2d18461d6 100644 --- a/tensorflow/python/estimator/export/export_test.py +++ b/tensorflow/python/estimator/export/export_test.py @@ -385,5 +385,67 @@ class ExportTest(test_util.TensorFlowTestCase): self.assertTrue(int(time_2) < int(time_3)) +class TensorServingReceiverTest(test_util.TensorFlowTestCase): + + def test_tensor_serving_input_receiver_constructor(self): + features = constant_op.constant([0]) + receiver_tensors = { + "example0": array_ops.placeholder(dtypes.string, name="example0"), + u"example1": array_ops.placeholder(dtypes.string, name="example1"), + } + r = export.TensorServingInputReceiver(features, receiver_tensors) + self.assertTrue(isinstance(r.features, ops.Tensor)) + self.assertTrue(isinstance(r.receiver_tensors, dict)) + + def test_tensor_serving_input_receiver_sparse(self): + features = sparse_tensor.SparseTensor( + indices=[[0, 0]], values=[1], dense_shape=[1, 1]) + receiver_tensors = { + "example0": array_ops.placeholder(dtypes.string, name="example0"), + u"example1": array_ops.placeholder(dtypes.string, name="example1"), + } + r = export.TensorServingInputReceiver(features, receiver_tensors) + self.assertTrue(isinstance(r.features, sparse_tensor.SparseTensor)) + self.assertTrue(isinstance(r.receiver_tensors, dict)) + + def test_serving_input_receiver_features_invalid(self): + receiver_tensors = { + "example0": array_ops.placeholder(dtypes.string, name="example0"), + u"example1": array_ops.placeholder(dtypes.string, name="example1"), + } + + with self.assertRaisesRegexp(ValueError, "features must be defined"): + export.TensorServingInputReceiver( + features=None, + receiver_tensors=receiver_tensors) + + with self.assertRaisesRegexp(ValueError, "feature must be a Tensor"): + export.TensorServingInputReceiver( + features={"1": constant_op.constant([1])}, + receiver_tensors=receiver_tensors) + + def test_serving_input_receiver_receiver_tensors_invalid(self): + features = constant_op.constant([0]) + + with self.assertRaisesRegexp( + ValueError, "receiver_tensors must be defined"): + export.TensorServingInputReceiver( + features=features, + receiver_tensors=None) + + with self.assertRaisesRegexp( + ValueError, "receiver_tensors keys must be strings"): + export.TensorServingInputReceiver( + features=features, + receiver_tensors={ + 1: array_ops.placeholder(dtypes.string, name="example0")}) + + with self.assertRaisesRegexp( + ValueError, "receiver_tensor example1 must be a Tensor"): + export.TensorServingInputReceiver( + features=features, + receiver_tensors={"example1": [1]}) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/estimator/exporter.py b/tensorflow/python/estimator/exporter.py index ba522f396d0eda1bb3d13b21acfddcc3d593e21b..a3f04626d1e5ed7ca7fb09a5dcc2457a0cf5ab82 100644 --- a/tensorflow/python/estimator/exporter.py +++ b/tensorflow/python/estimator/exporter.py @@ -25,8 +25,10 @@ from tensorflow.python.estimator import gc from tensorflow.python.framework import errors_impl from tensorflow.python.platform import gfile from tensorflow.python.platform import tf_logging +from tensorflow.python.util.tf_export import tf_export +@tf_export('estimator.Exporter') class Exporter(object): """A class representing a type of model export.""" @@ -123,6 +125,7 @@ class _SavedModelExporter(Exporter): return export_result +@tf_export('estimator.FinalExporter') class FinalExporter(Exporter): """This class exports the serving graph and checkpoints in the end. @@ -174,6 +177,7 @@ class FinalExporter(Exporter): is_the_final_export) +@tf_export('estimator.LatestExporter') class LatestExporter(Exporter): """This class regularly exports the serving graph and checkpoints. diff --git a/tensorflow/python/estimator/inputs/numpy_io.py b/tensorflow/python/estimator/inputs/numpy_io.py index c4c2e30e8771c5cb1e492fed751c71583dcf477b..a6f471291008e3c27dea1aeea5865e334f76e5c8 100644 --- a/tensorflow/python/estimator/inputs/numpy_io.py +++ b/tensorflow/python/estimator/inputs/numpy_io.py @@ -24,6 +24,7 @@ import numpy as np from six import string_types from tensorflow.python.estimator.inputs.queues import feeding_functions +from tensorflow.python.util.tf_export import tf_export # Key name to pack the target into dict of `features`. See # `_get_unique_target_key` for details. @@ -86,6 +87,7 @@ def _validate_and_convert_features(x): return ordered_dict_data +@tf_export('estimator.inputs.numpy_input_fn') def numpy_input_fn(x, y=None, batch_size=128, diff --git a/tensorflow/python/estimator/inputs/pandas_io.py b/tensorflow/python/estimator/inputs/pandas_io.py index 90d6145377d8f931b94793f8a912f77f1620f16e..bd06843021f47f81fc0c22d0fcee43530dc10098 100644 --- a/tensorflow/python/estimator/inputs/pandas_io.py +++ b/tensorflow/python/estimator/inputs/pandas_io.py @@ -21,6 +21,7 @@ from __future__ import print_function import numpy as np from tensorflow.python.estimator.inputs.queues import feeding_functions +from tensorflow.python.util.tf_export import tf_export try: # pylint: disable=g-import-not-at-top @@ -34,6 +35,7 @@ except ImportError: HAS_PANDAS = False +@tf_export('estimator.inputs.pandas_input_fn') def pandas_input_fn(x, y=None, batch_size=128, diff --git a/tensorflow/python/estimator/model_fn.py b/tensorflow/python/estimator/model_fn.py index b08f83fc569b1bb1ea6e5c93c57be7b5bb96f0a5..8111ab564c017175b3f7bc1020d850db74587958 100644 --- a/tensorflow/python/estimator/model_fn.py +++ b/tensorflow/python/estimator/model_fn.py @@ -31,8 +31,10 @@ from tensorflow.python.saved_model import signature_constants from tensorflow.python.training import monitored_session from tensorflow.python.training import session_run_hook from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export +@tf_export('estimator.ModeKeys') class ModeKeys(object): """Standard names for model modes. @@ -52,6 +54,7 @@ LOSS_METRIC_KEY = 'loss' AVERAGE_LOSS_METRIC_KEY = 'average_loss' +@tf_export('estimator.EstimatorSpec') class EstimatorSpec( collections.namedtuple('EstimatorSpec', [ 'mode', 'predictions', 'loss', 'train_op', 'eval_metric_ops', diff --git a/tensorflow/python/estimator/replicate_model_fn.py b/tensorflow/python/estimator/replicate_model_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..144d89abf3444062927d9261301fe50f4a63b280 --- /dev/null +++ b/tensorflow/python/estimator/replicate_model_fn.py @@ -0,0 +1,824 @@ +# 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 to replicate model_fn's over local GPUs. + +This file contains util that allow to replicate `Estimator.model_fn` over +GPUs. Replicated version of a `model_fn` is returned that can subsequently +be used with `Estimator`. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from collections import defaultdict +from contextlib import contextmanager +import copy + +import six + +from tensorflow.core.framework import node_def_pb2 +from tensorflow.python.client import device_lib +from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.estimator import util +from tensorflow.python.estimator.export import export_output as export_output_lib +from tensorflow.python.framework import device as framework_device +from tensorflow.python.framework import ops as ops_lib +from tensorflow.python.framework import sparse_tensor +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 sparse_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops.losses import losses +from tensorflow.python.platform import tf_logging +from tensorflow.python.training import device_setter as device_setter_lib +from tensorflow.python.training import optimizer as optimizer_lib + + +def _replicate_model_fn(model_fn, + devices=None): + """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 + (a.k.a. a tower). The input features and labels get sharded into the chunks + that correspond to the number of GPUs. Each tower computes a loss based + on its input. For each such loss, gradients are computed. After that, the + available losses are aggregated to form aggregated loss. Available + gradients are summed. Then, they update weights using the specified + optimizer. + + If `devices` are `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. + + Two modes of local replication over available GPUs are supported: + 1) If exactly 1 GPU is detected, then variables and operations are placed + 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. + + Here is an example of how one might use their `model_fn` to run over GPUs: + ```python + ... + def model_fn(...): # See `model_fn` in `Estimator`. + loss = ... + optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) + 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, + train_op=optimizer.minimize(loss)) + + # No change for `ModeKeys.EVAL` or `ModeKeys.PREDICT`. + return EstimatorSpec(...) + ... + classifier = tf.estimator.Estimator( + 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 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 + split across the dimension 0. Features and labels need to be batch major. + + On reduction algorithms: + Certain algorithms were chosen for aggregating results of computations on + multiple towers: + - Losses from all towers are reduced according to `loss_reduction` argument + to TowerOptimizer.. + - Gradients from all towers are reduced according to the `loss_reduction` + for each trainable variable. + - `eval_metrics_ops` are reduced per metric using `reduce_mean`. + - `EstimatorSpec.predictions` and `EstimatorSpec.export_outputs` are + reduced using concatenation. + - For all other fields of `EstimatorSpec` the values of the first tower + are taken. + + On distribution of variables: + Variables are not duplicated between towers. Instead, they are placed on a + single device as defined above and shared across towers. + + 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 + the train_op argument of `EstimatorSpec`. + devices: Optional list of devices to replicate the model across. This + argument can be used to replice only on the subset of available GPUs. + 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. + + 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 + instead of the supplied `model_fn`. + """ + return _replicate_model_fn_with_mode( + model_fn, + devices, + # TODO(isaprykin): Query the system configuration to choose modes other + # than `SHARED_LOCAL_PARAMETER_SERVER`, even though it is often + # appropriate. + mode=_VariableDistributionMode.SHARED_LOCAL_PARAMETER_SERVER) + + +class _VariableDistributionMode(object): + """Modes for variable distribution used for forcing a particular one. + + Forcing a mode is meant for performance experimentation purposes rather than + for general use cases. + """ + + SHARED_LOCAL_PARAMETER_SERVER = 1 + """Variables are placed on a single device and shared across all devices. + + Two ways to achieve this distribution over available GPUs are supported: + 1) If exactly 1 GPU is detected, then variables and operations are placed + onto 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. + """ + + SHARED_ROUND_ROBIN = 2 + """Variables are placed on all devices in a round-robin fashion. + + Every subsequent variable is placed on the next device. There is only one + copy of each variable that is shared across all devices. + """ + + +def _replicate_model_fn_with_mode( + model_fn, + devices=None, + mode=_VariableDistributionMode.SHARED_LOCAL_PARAMETER_SERVER): + """A version of `replicate_model_fn` that allows to specify a `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].upper() + consolidation_device = devices[0] if is_a_single_gpu_case else '/CPU:0' + + ps_devices = [consolidation_device] + if mode == _VariableDistributionMode.SHARED_ROUND_ROBIN: + ps_devices = devices + + tf_logging.info('Replicating the `model_fn` across {}. Variables are going ' + '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, + 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( + features, labels, len(devices), device=consolidation_device) + tower_specs = _get_loss_towers( + model_fn=model_fn, + mode=mode, + features=feature_shards, + labels=label_shards, + params=params, + config=config, + devices=devices, + local_ps_devices=ps_devices) + + if mode == model_fn_lib.ModeKeys.TRAIN: + train_op = _minimize_towers(tower_specs) + return _train_spec( + tower_specs, train_op, aggregation_device=consolidation_device) + elif mode == model_fn_lib.ModeKeys.EVAL: + return _eval_spec(tower_specs, aggregation_device=consolidation_device) + elif mode == model_fn_lib.ModeKeys.PREDICT: + return _predict_spec(tower_specs, aggregation_device=consolidation_device) + + if len(devices) == 1: + return single_device_model_fn + else: + return replicated_model_fn + + +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' + + def __init__(self, optimizer_or_optimizer_fn, + loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE): + """Wrap an existing optimizer for gathering gradients across towers. + + 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. + + 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 + instance is going to be used for optimizer-specific logic. This can + also be a no-argument function that returns such an optimizer instance. + loss_reduction: controls whether losses are summed or averaged. + """ + self._optimizer_or_optimizer_fn = optimizer_or_optimizer_fn + self._loss_reduction = loss_reduction + + @staticmethod + def has_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) + + def get_slot_names(self, *args, **kwargs): + return self._get_optimizer().get_slot_names(*args, **kwargs) + + def get_name(self, *args, **kwargs): + return self._get_optimizer().get_name(*args, **kwargs) + + def variables(self, *args, **kwargs): + return self._get_optimizer().variables(*args, **kwargs) + + def compute_gradients(self, loss, *args, **kwargs): + """Compute gradients, but first, if needed, scale the loss.""" + _TowerOptimizer._graph_state().set_loss_reduction(self._loss_reduction) + loss = _scale_loss(loss, + self._loss_reduction, + self._graph_state().number_of_towers) + return self._get_optimizer().compute_gradients(loss, *args, **kwargs) + + def apply_gradients(self, grads_and_vars, global_step=None, **kwargs): + """Collect gradients updates to apply them with the last tower.""" + 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: + 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. + var_scope, name_scope = self._graph_state().scopes_of_the_first_tower + with variable_scope.variable_scope(var_scope): + with ops_lib.name_scope(name_scope): + return self._apply_gathered_gradients(global_step, **kwargs) + + def _apply_gathered_gradients(self, global_step, **kwargs): + graph_state = self._graph_state() + optimizer = self._get_optimizer() + + grad_lists = {} + for grad, var in graph_state.get_latest_gradients_from_all_towers(): + if grad is not None: + grad_lists.setdefault(var, []).append(grad) + + aggregated_grads = [] + with ops_lib.name_scope('gradient_aggregating'): + for var, grads in six.iteritems(grad_lists): + grad = _compute_sum_on_device(grads, var.device) + aggregated_grads.append((grad, var)) + return optimizer.apply_gradients( + aggregated_grads, global_step=global_step, **kwargs) + + def _get_optimizer(self): + 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. 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_tower_optimizer_been_used = True + return self._optimizer_or_optimizer_fn + + def _construct_no_op_train_op(self): + return control_flow_ops.no_op(name='train_op_placeholder') + + @staticmethod + def _graph_state(): + graph_states = ops_lib.get_default_graph().get_collection_ref( + _TowerOptimizer.COLLECTION_FOR_GRAPH_STATES) + if not graph_states: + 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 the Graph collection will prevent _PerGraphState from being + # serialized. + ops_lib.get_default_graph().clear_collection( + _TowerOptimizer.COLLECTION_FOR_GRAPH_STATES) + + class _PerGraphState(object): + """Gradient reduction related state of a Tensorflow graph.""" + + def __init__(self): + 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 _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_number_of_towers(self, number_of_towers): + self._number_of_towers = number_of_towers + + def set_loss_reduction(self, loss_reduction): + self._loss_reduction = loss_reduction + + @contextmanager + def tower(self, tower_id, var_scope, name_scope): + if tower_id == 0: + self._variable_scope = var_scope + self._name_scope = name_scope + self._current_tower_index = tower_id + yield + + @property + def scopes_of_the_first_tower(self): + return self._variable_scope, self._name_scope + + @property + def is_the_last_tower(self): + return self._current_tower_index == (self._number_of_towers - 1) + + @property + def number_of_towers(self): + return self._number_of_towers + + @property + def loss_reduction(self): + return self._loss_reduction + + @property + def has_tower_optimizer_been_used(self): + return self._has_tower_optimizer_been_used + + @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): + local_device_protos = device_lib.list_local_devices() + return [ + device.name + for device in local_device_protos + if device.device_type == 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)] + for name, tensor in six.iteritems(dictionary): + if isinstance(tensor, sparse_tensor.SparseTensor): + for i, shard in enumerate( + sparse_ops.sparse_split( + 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 + + with ops_lib.name_scope('split_inputs'): + with ops_lib.device(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: + label_shards = None + 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 + + +_DEFAULT_NAME_SCOPE_PATTERN = 'tower_{}' + + +def _get_loss_towers(model_fn, + mode, + features, + labels, + params, + config, + devices, + local_ps_devices, + name_scope_pattern=_DEFAULT_NAME_SCOPE_PATTERN): + """Replicate the loss computation across devices.""" + tower_specs = [] + + model_fn_args = util.fn_args(model_fn) + optional_params = {} + if 'params' in model_fn_args: + optional_params['params'] = copy.deepcopy(params) + if 'config' in model_fn_args: + optional_params['config'] = copy.deepcopy(config) + + # pylint: disable=protected-access + round_robin_strategy = device_setter_lib._RoundRobinStrategy( + num_tasks=len(local_ps_devices)) + _TowerOptimizer._graph_state().set_number_of_towers(len(devices)) + + for i, device in enumerate(devices): + is_the_first_tower = (i == 0) + + device_setter = _local_device_setter( + worker_device=device, + ps_devices=local_ps_devices, + ps_strategy=round_robin_strategy) + + # We would like to preserve the names of the variables and ops that the user + # might be relying on. Names without a prefix are going to resolve to + # variables and ops of the first tower. + name_scope = name_scope_pattern + if is_the_first_tower: + name_scope = '' + + 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 _TowerOptimizer._graph_state().tower( + tower_id=i, var_scope=var_scope, name_scope=name_scope): + with ops_lib.device(device_setter): + labels_shard = None + if labels: + labels_shard = labels[i] + + tower_spec = model_fn( + mode=mode, + features=features[i], + labels=labels_shard, + **optional_params) + + 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 _TowerOptimizer. + tower_spec = _scale_tower_loss( + tower_spec, + _TowerOptimizer._graph_state().loss_reduction, + number_of_towers=len(devices)) + tower_specs.append(tower_spec) + + 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 + + +def _local_device_setter(worker_device, ps_devices, ps_strategy): + """A device setter that puts distributes Var/Ops to PS/workers.""" + ps_ops = ['Variable', 'VariableV2', 'VarHandleOp'] + + def local_device_chooser(op): + current_device = framework_device.DeviceSpec.from_string(op.device or '') + + node_def = op if isinstance(op, node_def_pb2.NodeDef) else op.node_def + if node_def.op in ps_ops: + ps_device_spec = framework_device.DeviceSpec.from_string( + '{}'.format(ps_devices[ps_strategy(op)])) + + ps_device_spec.merge_from(current_device) + return ps_device_spec.to_string() + else: + worker_device_spec = framework_device.DeviceSpec.from_string( + worker_device or '') + worker_device_spec.merge_from(current_device) + return worker_device_spec.to_string() + + return local_device_chooser + + +def _scale_tower_loss(tower_spec, loss_reduction, number_of_towers): + """Produce an EstimatorSpec with approproriately scaled loss.""" + if tower_spec.loss is None: + return tower_spec + + estimator_spec = _asdict(tower_spec) + estimator_spec['loss'] = _scale_loss( + tower_spec.loss, + loss_reduction, + number_of_towers, + reduced_loss_name='averaged_loss') + return model_fn_lib.EstimatorSpec(**estimator_spec) + + +def _scale_loss(loss, loss_reduction, number_of_towers, reduced_loss_name=None): + """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.NONE: + raise ValueError('Tower losses need to be reduced in some way, yet {} ' + 'reduction is specified.'.format(loss_reduction)) + + if loss_reduction != losses.Reduction.SUM: + return math_ops.div(loss, 1.0 * number_of_towers, name=reduced_loss_name) + else: + return loss + + +def _minimize_towers(tower_specs): + """`train_op` of the last tower applies aggregated gradients.""" + return tower_specs[-1].train_op + + +def _compute_sum_on_device(values, device, name=None): + with ops_lib.device(device): + if isinstance(values[0], ops_lib.IndexedSlices): + if name: + raise ValueError('The name {} is not expected to be given to ' + 'IndexedSlices {}'.format(name, values)) + + values_concat = array_ops.concat([v.values for v in values], axis=0) + indices_concat = array_ops.concat([v.indices for v in values], axis=0) + return ops_lib.IndexedSlices(values_concat, indices_concat, + values[0].dense_shape) + else: + return math_ops.add_n(values, name=name) + + +def _train_spec(tower_specs, + train_op, + aggregation_device, + aggregated_loss_name='loss'): + """Populate replicated EstimatorSpec for `GraphKeys.TRAIN`.""" + # 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( + [spec.loss for spec in tower_specs], aggregation_device, + aggregated_loss_name) + return model_fn_lib.EstimatorSpec(**estimator_spec) + + +def _eval_spec(tower_specs, aggregation_device, aggregated_loss_name='loss'): + """Populate replicated EstimatorSpec for `GraphKeys.EVAL`.""" + estimator_spec = _asdict(tower_specs[0]) + estimator_spec['mode'] = model_fn_lib.ModeKeys.EVAL + estimator_spec['loss'] = _compute_sum_on_device( + [spec.loss for spec in tower_specs], aggregation_device, + aggregated_loss_name) + + update_ops = [] + for tower_spec in tower_specs: + for name, (_, update_op) in six.iteritems(tower_spec.eval_metric_ops): + update_ops.append(update_op) + + with ops_lib.control_dependencies(update_ops): + reduced_update_op = _reduce_metric_variables(len(tower_specs)) + + eval_metric_ops = {} + for name, (metric_tensor, _) in six.iteritems(tower_specs[0].eval_metric_ops): + eval_metric_ops[name] = (metric_tensor, reduced_update_op) + estimator_spec['eval_metric_ops'] = eval_metric_ops + return model_fn_lib.EstimatorSpec(**estimator_spec) + + +def _reduce_metric_variables(number_of_towers): + """Aggregate local variables used in metrics into the first tower.""" + if number_of_towers == 1: + return control_flow_ops.no_op(name='no_eval_metric_reduction') + + metric_variables = ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES) + variables_per_tower = len(metric_variables) // number_of_towers + + if len(metric_variables) % number_of_towers != 0: + raise ValueError( + 'Different `EstimatorSpec.eval_metric_ops` across `model_fn()` calls.' + ' Expected {} local variables, but got {} instead.'.format( + variables_per_tower * number_of_towers, len(metric_variables))) + + # `metric_variables` has the size of `variables_per_tower` x + # number_of_towers. Each tower is produced by calling the same model_fn. + # First `variables_per_tower` correspond to the first tower. Each such + # variable has an replica at the `(variables_per_tower * i)` position, where + # `i` is `[1.. number_of_towers]`. We are going to add values from replicas + # to each variable of the first tower. We then zero out replica values, so + # that `_reduce_metric_variables` operation is idempotent. If a metric + # is then computed based on local variables from the first tower, then the + # resulting metric is an estimate for all `number_of_towers` towers. + ops = [] + for i in range(0, variables_per_tower): + next_replica_id = i + variables_per_tower + replicas = [ + metric_variables[replica_id] + for replica_id in range(next_replica_id, len(metric_variables), + variables_per_tower) + ] # `replicas` doesn't contain the first-tower variable. + + reduce_op = state_ops.assign_add(metric_variables[i], + math_ops.add_n(replicas)) + + with ops_lib.control_dependencies([reduce_op]): + for replica in replicas: + zeros_for_replica = array_ops.zeros( + array_ops.shape(replica), dtype=replica.dtype) + zero_out_replica_op = state_ops.assign(replica, zeros_for_replica) + ops.append(zero_out_replica_op) + + return control_flow_ops.group(*ops) + + +def _predict_spec(tower_specs, aggregation_device): + """Populate replicated EstimatorSpec for `GraphKeys.PREDICT`.""" + estimator_spec = _asdict(tower_specs[0]) + estimator_spec['mode'] = model_fn_lib.ModeKeys.PREDICT + + with ops_lib.device(aggregation_device): + estimator_spec['predictions'] = _concat_tensor_dicts( + *[tower_spec.predictions for tower_spec in tower_specs]) + + export_outputs_dict = _dict_concat( + *[tower_spec.export_outputs for tower_spec in tower_specs]) + + export_outputs = {} + for name, export_output_list in six.iteritems(export_outputs_dict): + if isinstance(export_output_list[0], export_output_lib.PredictOutput): + export_outputs[name] = export_output_lib.PredictOutput( + outputs=_concat_tensor_dicts(*[ + export_output.outputs for export_output in export_output_list + ])) + elif isinstance(export_output_list[0], + export_output_lib.RegressionOutput): + export_outputs[name] = export_output_lib.RegressionOutput( + value=array_ops.concat( + [export_output.value for export_output in export_output_list], + axis=0)) + elif isinstance(export_output_list[0], + export_output_lib.ClassificationOutput): + scores = None + if export_output_list[0].scores is not None: + scores = array_ops.concat( + [export_output.scores for export_output in export_output_list], + axis=0) + + classes = None + if export_output_list[0].classes is not None: + classes = array_ops.stack( + [export_output.classes for export_output in export_output_list], + axis=0) + + export_outputs[name] = export_output_lib.ClassificationOutput( + scores=scores, classes=classes) + + estimator_spec['export_outputs'] = export_outputs + return model_fn_lib.EstimatorSpec(**estimator_spec) + + +def _concat_tensor_dicts(*tensor_dicts): + return { + name: array_ops.concat(tensors, axis=0, name=name) + for name, tensors in six.iteritems(_dict_concat(*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: + if d is None: + continue + + for k, v in six.iteritems(d): + list_dict.setdefault(k, []).append(v) + return list_dict + + +def _asdict(namedtuple): + """Returns a namedtuple as a dictionary. + + This is required because `_asdict()` in Python 3.x.x is broken in classes + that inherit from `collections.namedtuple`. See + https://bugs.python.org/issue24931 for more details. + + Args: + namedtuple: An object that inherits from `collections.namedtuple`. + + Returns: + A dictionary version of the tuple. + """ + return {k: getattr(namedtuple, k) for k in namedtuple._fields} diff --git a/tensorflow/python/estimator/replicate_model_fn_test.py b/tensorflow/python/estimator/replicate_model_fn_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ad1f9c02b92d7b1ce929494f4b6fbf636762a7fd --- /dev/null +++ b/tensorflow/python/estimator/replicate_model_fn_test.py @@ -0,0 +1,1739 @@ +# 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 utilities that replicate `Estimator.model_fn` over GPUs.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import re +import shutil +import tempfile +import numpy as np +import six + +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 replicate_model_fn +from tensorflow.python.estimator.canned import dnn +from tensorflow.python.estimator.canned import optimizers +from tensorflow.python.estimator.canned import prediction_keys +from tensorflow.python.estimator.export import export +from tensorflow.python.estimator.export import export_output +from tensorflow.python.estimator.inputs import numpy_io +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 +from tensorflow.python.ops import losses +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import metrics as metrics_lib +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.ops.losses import losses +from tensorflow.python.platform import gfile +from tensorflow.python.platform import test +from tensorflow.python.saved_model import signature_constants +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 +# replicate_model_fn._VariableDistributionMode when it's supported. +class DNNClassifierIntegrationTest(test_util.TensorFlowTestCase): + + def setUp(self): + self._model_dir = tempfile.mkdtemp() + + def test_complete_flow_with_public_version(self): + return self._complete_flow_with_mode(mode=None) + + def test_complete_flow_with_mode_local_ps_server(self): + return self._complete_flow_with_mode( + replicate_model_fn._VariableDistributionMode. + SHARED_LOCAL_PARAMETER_SERVER) + + def test_complete_flow_with_mode_round_robin(self): + return self._complete_flow_with_mode( + replicate_model_fn._VariableDistributionMode.SHARED_ROUND_ROBIN) + + def _complete_flow_with_mode(self, mode): + n_classes = 3 + input_dimension = 2 + batch_size = 12 + + data = np.linspace( + 0., n_classes - 1., batch_size * input_dimension, dtype=np.float32) + x_data = data.reshape(batch_size, input_dimension) + categorical_data = np.random.random_integers( + 0, len(x_data), size=len(x_data)) + y_data = np.reshape(self._as_label(data[:batch_size]), (batch_size, 1)) + train_input_fn = numpy_io.numpy_input_fn( + x={'x': x_data, + 'categories': categorical_data}, + y=y_data, + batch_size=batch_size, + num_epochs=None, + shuffle=True) + eval_input_fn = numpy_io.numpy_input_fn( + x={'x': x_data, + 'categories': categorical_data}, + y=y_data, + batch_size=batch_size, + shuffle=False) + predict_input_fn = numpy_io.numpy_input_fn( + x={'x': x_data, + 'categories': categorical_data}, + batch_size=batch_size, + shuffle=False) + + feature_columns = [ + feature_column.numeric_column('x', shape=(input_dimension,)), + feature_column.embedding_column( + feature_column.categorical_column_with_vocabulary_list( + 'categories', + vocabulary_list=np.linspace( + 0., len(x_data), len(x_data), dtype=np.int64)), 1) + ] + + def optimizer_fn(): + return optimizers.get_optimizer_instance('Adagrad', learning_rate=0.05) + + estimator = dnn.DNNClassifier( + hidden_units=(2, 2), + # Adagrad is configured with `get_optimizer_instance`, so the function + # form of `TowerOptimizer.__init__` is used. + optimizer=replicate_model_fn._TowerOptimizer( + optimizer_fn, loss_reduction=losses.Reduction.SUM), + feature_columns=feature_columns, + n_classes=n_classes, + model_dir=self._model_dir) + + if not mode: # Use the public `replicate_model_fn`. + model_fn = replicate_model_fn._replicate_model_fn( + 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 = estimator_lib.Estimator( + model_fn=model_fn, + model_dir=estimator.model_dir, + config=estimator.config, + params=estimator.params) + + num_steps = 10 + estimator.train(train_input_fn, steps=num_steps) + + scores = estimator.evaluate(eval_input_fn) + self.assertEqual(num_steps, scores[ops_lib.GraphKeys.GLOBAL_STEP]) + self.assertIn('loss', six.iterkeys(scores)) + + predicted_proba = np.array([ + x[prediction_keys.PredictionKeys.PROBABILITIES] + for x in estimator.predict(predict_input_fn) + ]) + self.assertAllEqual((batch_size, n_classes), predicted_proba.shape) + + feature_spec = feature_column.make_parse_example_spec(feature_columns) + serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn( + feature_spec) + export_dir = estimator.export_savedmodel(tempfile.mkdtemp(), + serving_input_receiver_fn) + self.assertTrue(gfile.Exists(export_dir)) + + # 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._TowerOptimizer.COLLECTION_FOR_GRAPH_STATES)) + + def _as_label(self, data_in_float): + return np.rint(data_in_float).astype(np.int64) + + def tearDown(self): + if self._model_dir: + writer_cache.FileWriterCache.clear() + shutil.rmtree(self._model_dir) + + +class ReplicateModelTest(test_util.TensorFlowTestCase): + + def create_model_fn_with_loss_reduction(self, loss_reduction): + + def model_fn(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 = replicate_model_fn._TowerOptimizer( + gradient_descent.GradientDescentOptimizer(params['learning_rate']), + loss_reduction=loss_reduction) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops=metrics, + predictions={'probabilities': predictions}, + train_op=optimizer.minimize(loss)) + + return model_fn + + @property + def params(self): + params = {} + params['learning_rate'] = 1.0 + return params + + def test_train(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.create_model_fn_with_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()) + + # 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)) + + # derivative of loss = (1*c - 1) + (2*c - 2) 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)) + + def test_train_with_mean_reduction(self): + features = np.array([[1.0], [2.0]]) + 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.create_model_fn_with_loss_reduction(losses.Reduction.MEAN), + 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()) + + # loss = feature * c - label + total_loss = ((1.0 * 10 - 1.0) + (2.0 * 10 - 2.0)) / 2.0 + self.assertEqual(total_loss, session.run(estimator_spec.loss)) + + # derivative of loss = (1*c - 1)/2 + (2*c - 2)/2 is 1.5. + # It's the same computation as without mean reduction, but the + # loss from every tower is scaled by 1/. + # new value of c = 10 - learning rate * 1.5 = 8.5 + session.run(estimator_spec.train_op) + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual(8.5, session.run(c)) + + def test_train_two_steps_collected_gradients_are_reset_between_steps(self): + with ops_lib.Graph().as_default(): + features = array_ops.placeholder(dtypes.float64) + labels = array_ops.placeholder(dtypes.float64) + + feature_inputs = np.array([[1.0], [2.0]]), np.array([[1.5], [2.5]]) + label_inputs = np.array([[1.0], [2.0]]), np.array([[1.5], [2.5]]) + + # loss = feature * c - label + expected_losses = ((1.0 * 10 - 1.0) + (2.0 * 10 - 2.0), + (1.5 * 7.0 - 1.5) + (2.5 * 7.0 - 2.5)) + # Derivative of the loss is 1.0 + 2.0 for the first step and 1.5 + 2.5 + # for the second. + expected_c = 10.0 - 3.0, 7.0 - 4.0 + + 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.create_model_fn_with_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()) + + for feature_input, label_input, loss, weight in zip( + feature_inputs, label_inputs, expected_losses, expected_c): + feeds = {features: feature_input, labels: label_input} + + self.assertEqual(loss, session.run(estimator_spec.loss, feeds)) + + session.run(estimator_spec.train_op, feeds) + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual(weight, session.run(c, feeds)) + + def test_eval(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session() as session: + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.create_model_fn_with_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()) + session.run(variables.global_variables_initializer()) + + accuracy, a = estimator_spec.eval_metric_ops['accuracy'] + auc, b = estimator_spec.eval_metric_ops['auc'] + + session.run([a, b]) + accuracy = session.run(accuracy) + auc = session.run(auc) + + # loss[i] = features[i] * 10 - labels[i]. + # Accuracy is 0.0 (no match) in the first tower. + # Accuracy is 1.0 (match) in the second tower, since the feature + # times weight "c" happened to be equal to the label. + total_loss = ((0.01 * 10 - 0.01) + (0.002 * 10 - 0.02)) + + self.assertNear((0.0 + 1.0) / 2.0, accuracy, 0.01) + self.assertEqual(0, auc) + self.assertNear(total_loss, session.run(estimator_spec.loss), 0.01) + + def test_eval_with_mean_reduction(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session() as session: + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.create_model_fn_with_loss_reduction(losses.Reduction.MEAN), + 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()) + session.run(variables.global_variables_initializer()) + + accuracy, a = estimator_spec.eval_metric_ops['accuracy'] + auc, b = estimator_spec.eval_metric_ops['auc'] + + session.run([a, b]) + accuracy = session.run(accuracy) + auc = session.run(auc) + + # loss[i] = features[i] * 10 - labels[i]. + # Accuracy is 0.0 (no match) in the first tower. + # Accuracy is 1.0 (match) in the second tower, since the feature + # times weight "c" happened to be equal to the label. + total_loss = ((0.01 * 10 - 0.01) + (0.002 * 10 - 0.02)) / 2.0 + + self.assertNear((0.0 + 1.0) / 2.0, accuracy, 0.01) + self.assertEqual(0, auc) + self.assertNear(total_loss, session.run(estimator_spec.loss), 0.01) + + def test_predict(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session() as session: + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + devices=['/gpu:0', '/gpu:1']) + estimator_spec = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.PREDICT, self.params) + session.run(variables.global_variables_initializer()) + + self.assertAllClose({ + 'probabilities': np.array([[0.1], [0.02]]) + }, session.run(estimator_spec.predictions)) + + 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.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + 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)) + + def test_eval_single_tower(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session() as session: + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + devices=['/gpu:0']) + estimator_spec = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.EVAL, self.params) + session.run(variables.local_variables_initializer()) + session.run(variables.global_variables_initializer()) + + accuracy, a = estimator_spec.eval_metric_ops['accuracy'] + auc, b = estimator_spec.eval_metric_ops['auc'] + + session.run([a, b]) + accuracy = session.run(accuracy) + auc = session.run(auc) + + # Accuracy is 0.0 (no match) in the first tower. + # Accuracy is 1.0 (match) in the second tower, since the feature + # times weight "c" happened to be equal to the label. + total_loss = ((0.01 * 10 - 0.01) + (0.002 * 10 - 0.02)) + + self.assertNear((0.0 + 1.0) / 2.0, accuracy, 0.01) + self.assertEqual(0, auc) + self.assertNear(total_loss, session.run(estimator_spec.loss), 0.01) + + def test_predict_single_tower(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session() as session: + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + devices=['/gpu:0']) + estimator_spec = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.PREDICT, self.params) + session.run(variables.global_variables_initializer()) + + self.assertAllClose({ + '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.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + devices=['/gpu:0', '/gpu:1']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + + def test_unsupported_loss_reduction(self): + features = np.array([[1.0], [2.0], [3.0]]) + labels = np.array([[1.0], [2.0], [3.0]]) + + with self.assertRaisesRegexp(ValueError, + '.+none.+reduction.+is.+specified.+'): + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.create_model_fn_with_loss_reduction(losses.Reduction.NONE), + devices=['/gpu:0', '/gpu:1', '/gpu:2']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + + 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.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + 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.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + 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, loss_reduction=losses.Reduction.SUM) + + 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, + 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): + + 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) + + side_effects = variable_scope.get_variable( + 'side_effects', + initializer=constant_op.constant(0, dtype=dtypes.float64), + dtype=dtypes.float64, + use_resource=True, + trainable=False) + + 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) + } + + first_optimizer = replicate_model_fn._TowerOptimizer( + gradient_descent.GradientDescentOptimizer(1.0), + loss_reduction=losses.Reduction.SUM) + second_optimizer = replicate_model_fn._TowerOptimizer( + adam.AdamOptimizer(1.0), loss_reduction=losses.Reduction.SUM) + + 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.apply_gradients(first_grads_and_vars), + second_optimizer.minimize(loss)]) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops=metrics, + predictions={'probabilities': predictions}, + train_op=train_op) + + def test_train(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', '/gpu:1']) + estimator_spec = replicated_model_fn(features, labels, + model_fn_lib.ModeKeys.TRAIN, {}) + 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. + # Adam subtracts another ~1. + session.run(estimator_spec.train_op) + with variable_scope.variable_scope('', reuse=True): + 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): + + 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( + 'c', + initializer=constant_op.constant(10, dtype=dtypes.float64), + dtype=dtypes.float64) + d = variable_scope.get_variable( + 'd', + initializer=constant_op.constant(2, 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) + + another_predictions = math_ops.multiply(features, d) + another_loss = losses.absolute_difference( + labels=labels, + predictions=another_predictions, + reduction=losses.Reduction.SUM) + another_loss = math_ops.reduce_sum(another_loss) + + total_loss = math_ops.add(loss, another_loss) + + metrics = { + 'accuracy': metrics_lib.accuracy(labels, predictions), + 'auc': metrics_lib.auc(labels, predictions) + } + + train_ops = [] + + optimizer = replicate_model_fn._TowerOptimizer( + gradient_descent.GradientDescentOptimizer(1.0), + loss_reduction=losses.Reduction.SUM) + 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), + loss_reduction=losses.Reduction.SUM) + 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, + eval_metric_ops=metrics, + predictions={'probabilities': predictions}, + train_op=train_op) + + def test_train(self): + features = np.array([[1.0], [2.0]]) + labels = np.array([[1.0], [2.0]]) + + with ops_lib.Graph().as_default(), self.test_session() as session: + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.model_fn, + devices=['/gpu:0', '/gpu:1']) + estimator_spec = replicated_model_fn(features, labels, + model_fn_lib.ModeKeys.TRAIN, {}) + session.run(variables.global_variables_initializer()) + + # For each tower, loss = (feature * c - label) + (feature * d - label). + total_loss = (1.0 * 10 - 1.0 + 1.0 * 2.0 - 1.0) + ( + 2.0 * 10 - 2.0 + 2.0 * 2.0 - 2.0) + self.assertEqual(total_loss, session.run(estimator_spec.loss)) + + session.run(estimator_spec.train_op) + + # loss' of c or loss' of d is 3. + # new value of c = 10 - learning rate * 3 = 7.0. + # new value of d = 2 - learning rate * 3 = -1.0. + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertNear(7.0, session.run(c), 0.000001) + 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): + + 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(1.0) + train_op = optimizer.minimize(loss) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops=metrics, + predictions={'probabilities': predictions}, + train_op=train_op) + + def test_train(self): + features = np.array([[1.0], [2.0]]) + labels = np.array([[1.0], [2.0]]) + + with self.test_session(): + with self.assertRaisesRegexp(ValueError, + '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, + {}) + + +class GetLossTowersTest(test_util.TensorFlowTestCase): + + def create_model_fn_with_loss_reduction(self, loss_reduction): + + def model_fn(mode, features, labels, params): + del params + c = variable_scope.get_variable( + 'c', + initializer=constant_op.constant(0.25, dtype=dtypes.float64), + dtype=dtypes.float64) + + predictions = math_ops.add(np.array([0.1, 0.2, 0.3, features[0]]), c) + labels = np.array([0.1, 0.2, 0.3, labels[0]]) + + loss = losses.absolute_difference( + labels=labels, + predictions=predictions, + reduction=losses.Reduction.SUM) + + optimizer = replicate_model_fn._TowerOptimizer( + gradient_descent.GradientDescentOptimizer(1.0), + loss_reduction) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=math_ops.reduce_sum(loss), + train_op=optimizer.minimize(loss)) + + return model_fn + + def test_gradients_are_computed(self): + with self.test_session() as session: + tower_specs = replicate_model_fn._get_loss_towers( + self.create_model_fn_with_loss_reduction(losses.Reduction.SUM), + mode=None, + features=[[0.6], [1.6]], + labels=[[0.6], [0.6]], + params=None, + config=None, + devices=['/gpu:0', '/gpu:1'], + local_ps_devices=['/gpu:0'], + name_scope_pattern='test_tower_{}') + session.run(variables.global_variables_initializer()) + + self.assertEqual(len(tower_specs), 2) + + self.assertEqual('/device:GPU:0', tower_specs[0].loss.device) + self.assertEqual('Sum:0', tower_specs[0].loss.name) + self.assertEqual(1.0, session.run(tower_specs[0].loss)) + + self.assertEqual('/device:GPU:1', tower_specs[1].loss.device) + self.assertEqual('test_tower_1/Sum:0', tower_specs[1].loss.name) + # The input batch for the second tower had a loss that is 1.0 + # bigger: 0.6 vs 1.6. + self.assertEqual(2.0, session.run(tower_specs[1].loss)) + + self.assertEqual(1, len(variables.global_variables())) + self.assertEqual(1, len(variables.trainable_variables())) + + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual(0.25, session.run(c)) + + def test_gradients_are_computed_with_mean_reduction(self): + with self.test_session() as session: + tower_specs = replicate_model_fn._get_loss_towers( + self.create_model_fn_with_loss_reduction(losses.Reduction.MEAN), + mode=model_fn_lib.ModeKeys.EVAL, + features=[[0.6], [1.6]], + labels=[[0.6], [0.6]], + params=None, + config=None, + devices=['/gpu:0', '/gpu:1'], + local_ps_devices=['/gpu:0'], + name_scope_pattern='test_tower_{}') + session.run(variables.global_variables_initializer()) + + self.assertEqual(len(tower_specs), 2) + + self.assertEqual('/device:GPU:0', tower_specs[0].loss.device) + self.assertEqual('averaged_loss:0', tower_specs[0].loss.name) + self.assertEqual(0.5, session.run(tower_specs[0].loss)) + + self.assertEqual('/device:GPU:1', tower_specs[1].loss.device) + self.assertEqual('test_tower_1/averaged_loss:0', tower_specs[1].loss.name) + # The input batch for the second tower had a loss that is 1.0 + # bigger: 0.6 vs 1.6. + self.assertEqual(1.0, session.run(tower_specs[1].loss)) + + self.assertEqual(1, len(variables.global_variables())) + self.assertEqual(1, len(variables.trainable_variables())) + + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual(0.25, session.run(c)) + + def test_variables_are_round_robined_correctly(self): + """Test that creates multiple variables and tests round-robin placement.""" + + def model_fn(mode, features, labels, params): + del params + for variable_name in ['a', 'b', 'c', 'd']: + c = variable_scope.get_variable( + variable_name, + initializer=constant_op.constant(0.25, dtype=dtypes.float64), + dtype=dtypes.float64) + + predictions = math_ops.add(np.array([0.1, 0.2, 0.3, features[0]]), c) + labels = np.array([0.1, 0.2, 0.3, labels[0]]) + loss = losses.absolute_difference( + labels=labels, + predictions=predictions, + reduction=losses.Reduction.SUM) + return model_fn_lib.EstimatorSpec( + mode=mode, loss=math_ops.reduce_sum(loss)) + + with self.test_session() as session: + tower_specs = replicate_model_fn._get_loss_towers( + model_fn, + mode=None, + features=[[0.6], [1.6], [2.6]], + labels=[[0.6], [0.6], [2.6]], + params=None, + config=None, + devices=['/gpu:0', '/gpu:1', '/gpu:3'], + local_ps_devices=['/gpu:0', '/gpu:1', '/gpu:3'], + name_scope_pattern='test_tower_{}') + session.run(variables.global_variables_initializer()) + + self.assertEqual(len(tower_specs), 3) + self.assertEqual('/device:GPU:0', tower_specs[0].loss.device) + self.assertEqual('/device:GPU:1', tower_specs[1].loss.device) + self.assertEqual('/device:GPU:3', tower_specs[2].loss.device) + + with variable_scope.variable_scope('', reuse=True): + a = variable_scope.get_variable('a', dtype=dtypes.float64) + self.assertEqual('/device:GPU:0', a.device) + b = variable_scope.get_variable('b', dtype=dtypes.float64) + self.assertEqual('/device:GPU:1', b.device) + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual('/device:GPU:3', c.device) + d = variable_scope.get_variable('d', dtype=dtypes.float64) + self.assertEqual('/device:GPU:0', d.device) + + +class SplitBatchTest(test_util.TensorFlowTestCase): + + def evaluate_shards(self, first_list, second_list): + evaluate_items = lambda x: x.eval() + 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(): + 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( + features, labels, 2, device='/gpu:0') + + feature_shards, label_shards = self.evaluate_shards( + feature_shards, label_shards) + + self.assertAllEqual([[0.0, 1.0], [2.0, 3.0]], feature_shards) + self.assertAllEqual([[10.0, 11.0], [12.0, 13.0]], label_shards) + + def test_to_each_their_own(self): + 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( + features, labels, 4, device='/gpu:0') + + feature_shards, label_shards = self.evaluate_shards( + feature_shards, label_shards) + + self.assertAllEqual([[0.0], [1.0], [2.0], [3.0]], feature_shards) + self.assertAllEqual([[10.0], [11.0], [12.0], [13.0]], label_shards) + + def test_one_batch(self): + 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( + features, labels, 1, device='/gpu:0') + + feature_shards, label_shards = self.evaluate_shards( + feature_shards, label_shards) + + self.assertAllEqual([[0.0, 1.0, 2.0, 3.0]], feature_shards) + self.assertAllEqual([[10.0, 11.0, 12.0, 13.0]], label_shards) + + def test_half_split_in_dictionary(self): + 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] + + feature_shards, label_shards = replicate_model_fn._split_batch( + features, labels, 2, device='/gpu:0') + + self.assertAllEqual([0.0, 1.0], feature_shards[0]['first'].eval()) + self.assertAllEqual([4.0, 5.0], feature_shards[0]['second'].eval()) + self.assertAllEqual([2.0, 3.0], feature_shards[1]['first'].eval()) + self.assertAllEqual([6.0, 7.0], feature_shards[1]['second'].eval()) + 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]} + labels = [10.0, 11.0, 12.0, 13.0] + + feature_shards, label_shards = replicate_model_fn._split_batch( + features, labels, 1, device='/gpu:0') + + self.assertAllEqual([0.0, 1.0, 2.0, 3.0], + feature_shards[0]['first'].eval()) + self.assertAllEqual([4.0, 5.0, 6.0, 7.0], + feature_shards[0]['second'].eval()) + self.assertAllEqual([10.0, 11.0, 12.0, 13.0], label_shards[0].eval()) + + def test_feature_and_label_dictionaries(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]} + labels = {'first': [10.0, 11.0], 'second': [12.0, 13.0]} + + feature_shards, label_shards = replicate_model_fn._split_batch( + features, labels, 2, device='/gpu:0') + + self.assertAllEqual([0.0, 1.0], feature_shards[0]['first'].eval()) + self.assertAllEqual([4.0, 5.0], feature_shards[0]['second'].eval()) + self.assertAllEqual([2.0, 3.0], feature_shards[1]['first'].eval()) + self.assertAllEqual([6.0, 7.0], feature_shards[1]['second'].eval()) + self.assertAllEqual([10.0], label_shards[0]['first'].eval()) + self.assertAllEqual([12.0], label_shards[0]['second'].eval()) + self.assertAllEqual([11], label_shards[1]['first'].eval()) + self.assertAllEqual([13.0], label_shards[1]['second'].eval()) + + +class TrainSpecTest(test_util.TensorFlowTestCase): + + expected_predictions = {} + + def create_estimator_spec(self, loss): + return model_fn_lib.EstimatorSpec( + mode=model_fn_lib.ModeKeys.TRAIN, + loss=loss, + train_op=loss, # Not used; currently required. + predictions=self.expected_predictions) + + def create_constant_loss(self, loss_value): + return constant_op.constant(loss_value, dtype=dtypes.float64) + + def test_example(self): + with self.test_session() as session: + tower_losses = list(map(self.create_constant_loss, [2, 4, 6])) + tower_specs = list(map(self.create_estimator_spec, tower_losses)) + + expected_train_op = tower_losses[1] + + estimator_spec = replicate_model_fn._train_spec( + tower_specs, expected_train_op, aggregation_device='/gpu:0') + + self.assertEqual(expected_train_op, estimator_spec.train_op) + self.assertEqual(2 + 4 + 6, session.run(estimator_spec.loss)) + self.assertEqual(self.expected_predictions, estimator_spec.predictions) + + +class EvalSpecTest(test_util.TensorFlowTestCase): + + def create_estimator_spec(self, loss, metrics): + return model_fn_lib.EstimatorSpec( + mode=model_fn_lib.ModeKeys.EVAL, loss=loss, eval_metric_ops=metrics) + + def create_constant_loss(self, loss_value): + return constant_op.constant(loss_value, dtype=dtypes.float64) + + def create_eval_metrics(self, noise): + predictions = np.array([0.1, 0.2, 0.3, 0.6 + noise]) + labels = np.array([0.1, 0.2, 0.3, 0.6]) + + metrics = { + 'accuracy': metrics_lib.accuracy(labels, predictions), + 'auc': metrics_lib.auc(labels, predictions) + } + return metrics + + def test_example(self): + with self.test_session() as session: + tower_losses = map(self.create_constant_loss, [2, 4, 6]) + tower_metrics = map(self.create_eval_metrics, [0, 0.2, 0.3]) + tower_specs = [ + self.create_estimator_spec(l, m) + for l, m in zip(tower_losses, tower_metrics) + ] + session.run(variables.local_variables_initializer()) + + estimator_spec = replicate_model_fn._eval_spec( + tower_specs, aggregation_device='/device:GPU:0') + + accuracy, a = estimator_spec.eval_metric_ops['accuracy'] + auc, b = estimator_spec.eval_metric_ops['auc'] + + self.assertEqual('/device:CPU:0', accuracy.device) + self.assertEqual('/device:CPU:0', auc.device) + + session.run([a, b]) + accuracy, auc = session.run([accuracy, auc]) + + self.assertNear((12 - 2) / 12, accuracy, 0.01) + self.assertEqual(0, auc) + self.assertEqual(2 + 4 + 6, session.run(estimator_spec.loss)) + + def test_handles_single_tower(self): + with self.test_session() as session: + tower_losses = map(self.create_constant_loss, [5]) + tower_metrics = map(self.create_eval_metrics, [0.2]) + tower_specs = [ + self.create_estimator_spec(l, m) + for l, m in zip(tower_losses, tower_metrics) + ] + session.run(variables.local_variables_initializer()) + + estimator_spec = replicate_model_fn._eval_spec( + tower_specs, aggregation_device='/device:GPU:0') + + accuracy, a = estimator_spec.eval_metric_ops['accuracy'] + auc, b = estimator_spec.eval_metric_ops['auc'] + + self.assertEqual('/device:CPU:0', accuracy.device) + self.assertEqual('/device:CPU:0', auc.device) + + session.run([a, b]) + accuracy = session.run(accuracy) + auc = session.run(auc) + + self.assertNear((4 - 1) / 4, accuracy, 0.01) + self.assertEqual(0, auc) + self.assertEqual(5, session.run(estimator_spec.loss)) + + +class PredictSpecTest(test_util.TensorFlowTestCase): + + def model_fn(self, mode, features, labels, params): + c = variable_scope.get_variable( + 'c', + initializer=constant_op.constant(0.25, dtype=dtypes.float64), + dtype=dtypes.float64) + + predictions = math_ops.add(np.array([features[0], features[0]]), c) + + return model_fn_lib.EstimatorSpec( + mode=model_fn_lib.ModeKeys.PREDICT, + predictions={ + 'probabilities': predictions + }) + + def test_example(self): + with self.test_session() as session: + tower_specs = replicate_model_fn._get_loss_towers( + self.model_fn, + mode=None, + features=[[0.1], [0.2]], + labels=[[], []], + params=None, + config=None, + devices=['/gpu:0', '/gpu:1'], + local_ps_devices=['/gpu:0'], + ) + session.run(variables.global_variables_initializer()) + + estimator_spec = replicate_model_fn._predict_spec( + tower_specs, aggregation_device='/gpu:0') + + self.assertEqual('/device:GPU:0', + estimator_spec.predictions['probabilities'].device) + self.assertAllClose({ + 'probabilities': np.array([0.35, 0.35, 0.45, 0.45]) + }, session.run(estimator_spec.predictions)) + + +class ReduceMetricVariablesTest(test_util.TensorFlowTestCase): + + def create_metric_variable(self, initial_value, name): + return variable_scope.variable( + initial_value, + trainable=False, + collections=[ops_lib.GraphKeys.METRIC_VARIABLES], + validate_shape=True, + name=name) + + def create_tower_metrics(self, tower_id): + with variable_scope.variable_scope('', reuse=(tower_id != 0)): + self.create_metric_variable(1.3 * (tower_id + 1), 'total') + self.create_metric_variable(2.3 * (tower_id + 1), 'count') + self.create_metric_variable( + np.array([3.3, 3.5, 3.7]) * (tower_id + 1), 'total') + + def test_example(self): + with self.test_session() as session: + for tower_id in range(3): + self.create_tower_metrics(tower_id) + + session.run( + variables.variables_initializer( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES))) + + session.run( + replicate_model_fn._reduce_metric_variables(number_of_towers=3)) + + # 1st tower = 1.3, 2.3, [3.3, 3.5, 3.7] + # 2nd tower = 2.6, 4.6, [6.6, 7.0, 7.4] + # 3rd tower = 3.9, 6.9, [9.9, 10.5, 11.1] + # Reduced = 7.8, 13.8, [19.8, 21.0, 22.2] + # Towers are accumulated in the first tower. + local_metrics = session.run( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES)) + + self.assertNear(7.8, local_metrics[0], 0.01) + self.assertNear(13.8, local_metrics[1], 0.01) + self.assertAllClose([19.8, 21., 22.1], local_metrics[2], 0.01) + self.assertNear(0.0, local_metrics[3], 0.01) + self.assertNear(0.0, local_metrics[4], 0.01) + self.assertAllClose([0.0, 0.0, 0.0], local_metrics[5], 0.01) + self.assertNear(0.0, local_metrics[6], 0.01) + self.assertNear(0.0, local_metrics[7], 0.01) + self.assertAllClose([0.0, 0.0, 0.0], local_metrics[8], 0.01) + + def test_reduce_is_idempotent(self): + with self.test_session() as session: + for tower_id in range(3): + self.create_tower_metrics(tower_id) + + session.run( + variables.variables_initializer( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES))) + + for _ in range(20): + session.run( + replicate_model_fn._reduce_metric_variables(number_of_towers=3)) + + local_metrics = session.run( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES)) + + self.assertNear(7.8, local_metrics[0], 0.01) + self.assertNear(13.8, local_metrics[1], 0.01) + self.assertAllClose([19.8, 21., 22.1], local_metrics[2], 0.01) + self.assertNear(0.0, local_metrics[3], 0.01) + self.assertNear(0.0, local_metrics[4], 0.01) + self.assertAllClose([0.0, 0.0, 0.0], local_metrics[5], 0.01) + self.assertNear(0.0, local_metrics[6], 0.01) + self.assertNear(0.0, local_metrics[7], 0.01) + self.assertAllClose([0.0, 0.0, 0.0], local_metrics[8], 0.01) + + def test_handles_single_tower(self): + with self.test_session() as session: + self.create_tower_metrics(0) + session.run( + variables.variables_initializer( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES))) + + session.run( + replicate_model_fn._reduce_metric_variables(number_of_towers=1)) + + local_metrics = session.run( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES)) + + self.assertNear(1.3, local_metrics[0], 0.01) + self.assertNear(2.3, local_metrics[1], 0.01) + self.assertAllClose([3.3, 3.5, 3.7], local_metrics[2], 0.01) + + def test_doesnt_accept_uneven_number_of_variables(self): + with self.test_session() as session: + for tower_id in range(3): + self.create_tower_metrics(tower_id) + self.create_metric_variable(-1.0, 'oddball') + + session.run( + variables.variables_initializer( + ops_lib.get_collection(ops_lib.GraphKeys.METRIC_VARIABLES))) + + with self.assertRaisesRegexp( + ValueError, '.+Expected.+local.+variables.+but.+got.+instead.+'): + session.run( + replicate_model_fn._reduce_metric_variables(number_of_towers=3)) + + +class MergeExportOutputsTest(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 = {'probabilities': math_ops.multiply(features, c)} + loss = losses.absolute_difference( + labels=labels, + predictions=predictions['probabilities'], + reduction=losses.Reduction.SUM) + + metrics = { + 'accuracy': metrics_lib.accuracy(labels, predictions['probabilities']), + 'auc': metrics_lib.auc(labels, predictions['probabilities']) + } + tensor_string_repr = str(features) + classes = constant_op.constant( + re.search('(split_inputs/split:[0-9])', tensor_string_repr).group(1), + dtype=dtypes.string) + + export_outputs = { + signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: + export_output.PredictOutput(predictions), + 'classification_output': + export_output.ClassificationOutput(predictions['probabilities'], + classes), + 'classification_scores': + export_output.ClassificationOutput( + scores=predictions['probabilities']), + 'classification_classes': + export_output.ClassificationOutput(classes=classes), + 'regression_output': + export_output.RegressionOutput(predictions['probabilities']), + } + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=math_ops.reduce_sum(loss), + eval_metric_ops=metrics, + predictions=predictions, + export_outputs=export_outputs) + + def replicate_estimator_spec(self, session): + features = np.array([0.01, 0.002]) + labels = np.array([0.01, 0.02]) + + replicated_model_fn = replicate_model_fn._replicate_model_fn( + self.model_fn, devices=['/gpu:0', '/gpu:1']) + estimator_spec = replicated_model_fn(features, labels, + model_fn_lib.ModeKeys.PREDICT, {}) + session.run(variables.global_variables_initializer()) + return estimator_spec + + def test_merge_predict_output(self): + with self.test_session() as session: + estimator_spec = self.replicate_estimator_spec(session) + self.assertAllClose( + { + 'probabilities': np.array([0.1, 0.02]) + }, + session.run(estimator_spec.export_outputs[ + signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].outputs)) + + def test_merge_classification_output_scores_classes(self): + with self.test_session() as session: + estimator_spec = self.replicate_estimator_spec(session) + self.assertAllClose( + [0.1, 0.02], + session.run( + estimator_spec.export_outputs['classification_output'].scores)) + self.assertAllEqual( + [b'split_inputs/split:0', b'split_inputs/split:1'], + session.run( + estimator_spec.export_outputs['classification_output'].classes)) + + def test_merge_classification_output_scores(self): + with self.test_session() as session: + estimator_spec = self.replicate_estimator_spec(session) + self.assertAllClose( + [0.1, 0.02], + session.run( + estimator_spec.export_outputs['classification_scores'].scores)) + self.assertEqual( + None, estimator_spec.export_outputs['classification_scores'].classes) + + def test_merge_classification_output_classes(self): + with self.test_session() as session: + estimator_spec = self.replicate_estimator_spec(session) + self.assertAllEqual( + [b'split_inputs/split:0', b'split_inputs/split:1'], + session.run( + estimator_spec.export_outputs['classification_classes'].classes)) + self.assertEqual( + None, estimator_spec.export_outputs['classification_classes'].scores) + + def test_merge_regression_output(self): + with self.test_session() as session: + estimator_spec = self.replicate_estimator_spec(session) + self.assertAllClose( + [0.1, 0.02], + session.run(estimator_spec.export_outputs['regression_output'].value)) + + +class GetLocalDevicesTest(test_util.TensorFlowTestCase): + + def test_there_is_at_least_a_cpu(self): + self.assertTrue(replicate_model_fn._get_local_devices('CPU')) + + def test_there_is_no_xpu(self): + self.assertFalse( + replicate_model_fn._get_local_devices('XPU')) # XPU doesn't exist. + + def test_whether_there_is_a_gpu(self): + if test.is_gpu_available(): + self.assertTrue(len(replicate_model_fn._get_local_devices('GPU'))) + + +class LocalDeviceSetterTest(test_util.TensorFlowTestCase): + + def test_vars_are_on_ps_but_ops_are_on_workers(self): + ps_devices = ['/device:GPU:3'] + round_robin = device_setter._RoundRobinStrategy(num_tasks=len(ps_devices)) + + local_device_setter = replicate_model_fn._local_device_setter( + ps_devices=ps_devices, + ps_strategy=round_robin, + worker_device='/device:GPU:2') + + with ops_lib.device(local_device_setter): + a = variables.Variable(0.01) + self.assertEqual('/device:GPU:3', a.device) + + b = variables.Variable(0.02) + self.assertEqual('/device:GPU:3', b.device) + + c = variables.Variable(0.03) + self.assertEqual('/device:GPU:3', c.device) + + a_op = array_ops.concat(a, axis=0) + self.assertEqual('/device:GPU:2', a_op.device) + + b_op = array_ops.concat(b, axis=0) + self.assertEqual('/device:GPU:2', b_op.device) + + def test_round_robin_placement(self): + ps_devices = [ + '/device:GPU:0', '/device:GPU:1', '/device:GPU:3', '/device:GPU:4' + ] + round_robin = device_setter._RoundRobinStrategy(num_tasks=len(ps_devices)) + + local_device_setter = replicate_model_fn._local_device_setter( + ps_devices=ps_devices, + ps_strategy=round_robin, + worker_device='/device:GPU:2') + + with ops_lib.device(local_device_setter): + a = variables.Variable(0.01) + self.assertEqual('/device:GPU:0', a.device) + + b = variables.Variable(0.02) + self.assertEqual('/device:GPU:1', b.device) + + c = variables.Variable(0.03) + self.assertEqual('/device:GPU:3', c.device) + + a_op = array_ops.concat(a, axis=0) + self.assertEqual('/device:GPU:2', a_op.device) + + b_op = array_ops.concat(b, axis=0) + self.assertEqual('/device:GPU:2', b_op.device) + + c = variables.Variable(0.03) + self.assertEqual('/device:GPU:4', c.device) + + d = variables.Variable(0.03) + self.assertEqual('/device:GPU:0', d.device) + + c_op = array_ops.concat(c, axis=0) + self.assertEqual('/device:GPU:2', c_op.device) + + +class ComputeSumWithDevicePlacementTest(test_util.TensorFlowTestCase): + + def test_vectors(self): + with self.test_session() as session: + total = replicate_model_fn._compute_sum_on_device( + [1.0, 2.0, 3.0, 4.0], device='/device:GPU:0', name='test_sum') + + self.assertEqual('/device:GPU:0', total.device) + self.assertEqual('test_sum', total.op.name) + self.assertEqual(10.0, session.run(total)) + + def test_tensors(self): + with self.test_session() as session: + total = replicate_model_fn._compute_sum_on_device( + [[1.0, 2.0], [3.0, 4.0]], device='/device:GPU:0', name='test_sum') + + self.assertEqual('/device:GPU:0', total.device) + self.assertEqual('test_sum', total.op.name) + self.assertAllEqual([4.0, 6.0], session.run(total)) + + def test_indexedslices(self): + with self.test_session() as session: + a = ops_lib.IndexedSlices( + constant_op.constant([1.0, 2.0]), [0, 1], + dense_shape=constant_op.constant([2])) + b = ops_lib.IndexedSlices(constant_op.constant([3.0, 4.0]), [0, 1]) + + total = replicate_model_fn._compute_sum_on_device( + [a, b], device='/device:GPU:0') + + self.assertEqual('/device:GPU:0', total.device) + self.assertAllEqual([4.0, 6.0], + session.run(ops_lib.convert_to_tensor(total))) + + def test_indexedslices_higher_dimensions(self): + with self.test_session() as session: + a = ops_lib.IndexedSlices( + constant_op.constant([[1.0, 5.0], [2.0, 6.0]]), [0, 1], + dense_shape=constant_op.constant([2, 4])) + b = ops_lib.IndexedSlices( + constant_op.constant([[3.0, 7.0], [4.0, 8.0]]), [0, 1]) + + total = replicate_model_fn._compute_sum_on_device( + [a, b], device='/device:GPU:0') + + self.assertEqual('/device:GPU:0', total.device) + self.assertAllEqual([[4.0, 12.0], [6.0, 14.0]], + session.run(ops_lib.convert_to_tensor(total))) + + def test_indexedslices_some_dont_overlap(self): + with self.test_session() as session: + a = ops_lib.IndexedSlices( + constant_op.constant([1.0, 2.0]), [0, 3], + dense_shape=constant_op.constant([4])) + b = ops_lib.IndexedSlices(constant_op.constant([3.0, 4.0]), [0, 1]) + + total = replicate_model_fn._compute_sum_on_device( + [a, b], device='/device:GPU:0') + + self.assertEqual('/device:GPU:0', total.device) + self.assertAllEqual([4.0, 4.0, 0.0, 2.0], + session.run(ops_lib.convert_to_tensor(total))) + + def test_no_name_for_indexslices(self): + a = ops_lib.IndexedSlices( + constant_op.constant([1.0, 2.0]), [0, 1], + dense_shape=constant_op.constant([2])) + b = ops_lib.IndexedSlices(constant_op.constant([3.0, 4.0]), [0, 1]) + + with self.assertRaisesRegexp(ValueError, '.+name.+not.+expected.+'): + _ = replicate_model_fn._compute_sum_on_device( + [a, b], device='/device:GPU:0', name='cant_name_indexslices') + + +class ConcatTensorDictsTest(test_util.TensorFlowTestCase): + + def test_example(self): + tensor_dicts = [ + { + 'a': np.array([1.0, 2.0]), + 'b': np.array([11.0]), + 'c': np.array([21.0]), + }, + { + 'a': np.array([3.0]), + 'b': np.array([12.0, 13.0]), + }, + { + 'b': np.array([14.0]), + }, + ] + + with self.test_session() as session: + self.assertAllClose({ + 'a': np.array([1.0, 2.0, 3.0]), + 'b': np.array([11.0, 12.0, 13.0, 14.0]), + 'c': np.array([21.0]), + }, session.run(replicate_model_fn._concat_tensor_dicts(*tensor_dicts))) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/estimator/run_config.py b/tensorflow/python/estimator/run_config.py index e446b3e03a262e0d4abe69df73cb8604b0dab9f9..62f035bce558f57c5fd39d60b44cf8eb0130ce38 100644 --- a/tensorflow/python/estimator/run_config.py +++ b/tensorflow/python/estimator/run_config.py @@ -27,8 +27,8 @@ import six from tensorflow.core.protobuf import config_pb2 from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import server_lib -from tensorflow.python.util import compat from tensorflow.python.util import compat_internal +from tensorflow.python.util.tf_export import tf_export _USE_DEFAULT = object() @@ -287,6 +287,7 @@ class TaskType(object): EVALUATOR = 'evaluator' +@tf_export('estimator.RunConfig') class RunConfig(object): """This class specifies the configurations for an `Estimator` run.""" @@ -344,7 +345,7 @@ class RunConfig(object): os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'worker', 'index': 1}}) - config = ClusterConfig() + config = RunConfig() assert config.master == 'host4:2222' assert config.task_id == 1 assert config.num_ps_replicas == 2 @@ -362,7 +363,7 @@ class RunConfig(object): os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'chief', 'index': 0}}) - config = ClusterConfig() + config = RunConfig() assert config.master == 'host0:2222' assert config.task_id == 0 assert config.num_ps_replicas == 2 @@ -380,7 +381,7 @@ class RunConfig(object): os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'evaluator', 'index': 0}}) - config = ClusterConfig() + config = RunConfig() assert config.master == '' assert config.evaluator_master == '' assert config.task_id == 0 diff --git a/tensorflow/python/estimator/training.py b/tensorflow/python/estimator/training.py index 2e84c5014f6e17b34b38d2dfe5711b5b654553bb..e38b765da52a7b6957a4fb8a02087c5d1fd5a781 100644 --- a/tensorflow/python/estimator/training.py +++ b/tensorflow/python/estimator/training.py @@ -35,6 +35,7 @@ from tensorflow.python.training import basic_session_run_hooks from tensorflow.python.training import server_lib from tensorflow.python.training import session_run_hook from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export _MAX_DELAY_SECS = 60 _DELAY_SECS_PER_WORKER = 5 @@ -114,6 +115,7 @@ def _is_google_env(): return tf_config.get(_ENVIRONMENT_KEY) == _ENVIRONMENT_GOOGLE_VALUE +@tf_export('estimator.TrainSpec') class TrainSpec( collections.namedtuple('TrainSpec', ['input_fn', 'max_steps', 'hooks'])): """Configuration for the "train" part for the `train_and_evaluate` call. @@ -126,9 +128,16 @@ class TrainSpec( """Creates a validated `TrainSpec` instance. Args: - input_fn: Training input function returning a tuple of: - features - `Tensor` or dictionary of string feature name to `Tensor`. - labels - `Tensor` or dictionary of `Tensor` with labels. + input_fn: A function that provides input data for training as minibatches. + See @{$get_started/premade_estimators#create_input_functions} for more + information. The function should construct and return one of + the following: + * A 'tf.data.Dataset' object: Outputs of `Dataset` object must be a + tuple (features, labels) with same constraints as below. + * A tuple (features, labels): Where features is a `Tensor` or a + dictionary of string feature name to `Tensor` and labels is a + `Tensor` or a dictionary of string label name to `Tensor`. + max_steps: Int. Positive number of total steps for which to train model. If `None`, train forever. The training `input_fn` is not expected to generate `OutOfRangeError` or `StopIteration` exceptions. See the @@ -158,6 +167,7 @@ class TrainSpec( cls, input_fn=input_fn, max_steps=max_steps, hooks=hooks) +@tf_export('estimator.EvalSpec') class EvalSpec( collections.namedtuple('EvalSpec', [ 'input_fn', 'steps', 'name', 'hooks', 'exporters', 'start_delay_secs', @@ -182,9 +192,16 @@ class EvalSpec( """Creates a validated `EvalSpec` instance. Args: - input_fn: Evaluation input function returning a tuple of: - features - `Tensor` or dictionary of string feature name to `Tensor`. - labels - `Tensor` or dictionary of `Tensor` with labels. + input_fn: A function that constructs the input data for evaluation. + See @{$get_started/premade_estimators#create_input_functions} for more + information. The function should construct and return one of + the following: + * A 'tf.data.Dataset' object: Outputs of `Dataset` object must be a + tuple (features, labels) with same constraints as below. + * A tuple (features, labels): Where features is a `Tensor` or a + dictionary of string feature name to `Tensor` and labels is a + `Tensor` or a dictionary of string label name to `Tensor`. + steps: Int. Positive number of steps for which to evaluate model. If `None`, evaluates until `input_fn` raises an end-of-input exception. See `Estimator.evaluate` for details. @@ -246,6 +263,7 @@ class EvalSpec( throttle_secs=throttle_secs) +@tf_export('estimator.train_and_evaluate') def train_and_evaluate(estimator, train_spec, eval_spec): """Train and evaluate the `estimator`. @@ -451,15 +469,21 @@ class _TrainingExecutor(object): train_hooks=None, continuous_eval_listener=None): if not isinstance(estimator, estimator_lib.Estimator): - raise TypeError('`estimator` must have type `tf.estimator.Estimator`.') + raise TypeError( + '`estimator` must have type `tf.estimator.Estimator`. ' + 'Got: {}'.format(type(estimator))) self._estimator = estimator if not isinstance(train_spec, TrainSpec): - raise TypeError('`train_spec` must have type `tf.estimator.TrainSpec`.') + raise TypeError( + '`train_spec` must have type `tf.estimator.TrainSpec`. ' + 'Got: {}'.format(type(train_spec))) self._train_spec = train_spec if not isinstance(eval_spec, EvalSpec): - raise TypeError('`eval_spec` must have type `tf.estimator.EvalSpec`.') + raise TypeError( + '`eval_spec` must have type `tf.estimator.EvalSpec`. ' + 'Got: {}'.format(type(eval_spec))) self._eval_spec = eval_spec self._train_hooks = _validate_hooks(train_hooks) diff --git a/tensorflow/python/estimator/util.py b/tensorflow/python/estimator/util.py index b7ba76d8714e6b13551bb3e18083f45e53d2afc3..bb4bdd3fdfb2e19dbc1c581d7771f2e1ac4442ba 100644 --- a/tensorflow/python/estimator/util.py +++ b/tensorflow/python/estimator/util.py @@ -20,11 +20,18 @@ from __future__ import division from __future__ import print_function import functools +import os +import time +from tensorflow.python.platform import gfile +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util import compat +from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect def _is_bounded_method(fn): + _, fn = tf_decorator.unwrap(fn) return tf_inspect.ismethod(fn) and (fn.__self__ is not None) @@ -54,3 +61,48 @@ def fn_args(fn): if _is_bounded_method(fn): args.remove('self') return tuple(args) + + +# When we create a timestamped directory, there is a small chance that the +# directory already exists because another process is also creating these +# directories. In this case we just wait one second to get a new timestamp and +# try again. If this fails several times in a row, then something is seriously +# wrong. +MAX_DIRECTORY_CREATION_ATTEMPTS = 10 + + +def get_timestamped_dir(dir_base): + """Builds a path to a new subdirectory within the base directory. + + The subdirectory will be named using the current time. + This guarantees monotonically increasing directory numbers even across + multiple runs of the pipeline. + The timestamp used is the number of seconds since epoch UTC. + + Args: + dir_base: A string containing a directory to create the subdirectory under. + + Returns: + The full path of the new subdirectory (which is not actually created yet). + + Raises: + RuntimeError: if repeated attempts fail to obtain a unique timestamped + directory name. + """ + attempts = 0 + while attempts < MAX_DIRECTORY_CREATION_ATTEMPTS: + timestamp = int(time.time()) + + result_dir = os.path.join( + compat.as_bytes(dir_base), compat.as_bytes(str(timestamp))) + if not gfile.Exists(result_dir): + # Collisions are still possible (though extremely unlikely): this + # directory is not actually created yet, but it will be almost + # instantly on return from this function. + return result_dir + time.sleep(1) + attempts += 1 + logging.warn('Directory {} already exists; retrying (attempt {}/{})'.format( + result_dir, attempts, MAX_DIRECTORY_CREATION_ATTEMPTS)) + raise RuntimeError('Failed to obtain a unique export directory name after ' + '{} attempts.'.format(MAX_DIRECTORY_CREATION_ATTEMPTS)) diff --git a/tensorflow/python/feature_column/BUILD b/tensorflow/python/feature_column/BUILD index a758f8a4fc4898713772c4e919acda48b0f6ad0b..238a90b67d9d0039c25a6f3800aad25a2db9e36f 100644 --- a/tensorflow/python/feature_column/BUILD +++ b/tensorflow/python/feature_column/BUILD @@ -74,7 +74,10 @@ py_test( srcs = ["feature_column_test.py"], data = [":vocabulary_testdata"], srcs_version = "PY2AND3", - tags = ["no_pip"], + tags = [ + "no_cuda_on_cpu_tap", + "no_pip", + ], deps = [ ":feature_column", ":feature_column_py", diff --git a/tensorflow/python/feature_column/feature_column.py b/tensorflow/python/feature_column/feature_column.py index 7feb209cc49c4be70387c44168dbdeea6d108d66..7d99fcb3e79318c2fecabaa9bdd0347aa67cf309 100644 --- a/tensorflow/python/feature_column/feature_column.py +++ b/tensorflow/python/feature_column/feature_column.py @@ -16,7 +16,7 @@ FeatureColumns provide a high level abstraction for ingesting and representing features. FeatureColumns are also the primary way of encoding features for -canned ${tf.estimator.Estimator}s. +canned @{tf.estimator.Estimator}s. When using FeatureColumns with `Estimators`, the type of feature column you should choose depends on (1) the feature type and (2) the model type. @@ -157,6 +157,8 @@ from tensorflow.python.platform import gfile from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import checkpoint_utils from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export +from tensorflow.python.util.tf_export import tf_export def _internal_input_layer(features, @@ -209,6 +211,7 @@ def _internal_input_layer(features, return array_ops.concat(output_tensors, 1) +@tf_export('feature_column.input_layer') def input_layer(features, feature_columns, weight_collections=None, @@ -329,6 +332,7 @@ class InputLayer(object): return self._input_layer_template.weights +@tf_export('feature_column.linear_model') def linear_model(features, feature_columns, units=1, @@ -498,6 +502,7 @@ def _transform_features(features, feature_columns): return outputs +@tf_export('feature_column.make_parse_example_spec') def make_parse_example_spec(feature_columns): """Creates parsing spec dictionary from input feature_columns. @@ -507,6 +512,7 @@ def make_parse_example_spec(feature_columns): ```python # Define features and transformations + feature_a = categorical_column_with_vocabulary_file(...) feature_b = numeric_column(...) feature_c_bucketized = bucketized_column(numeric_column("feature_c"), ...) feature_a_x_feature_c = crossed_column( @@ -557,6 +563,7 @@ def make_parse_example_spec(feature_columns): return result +@tf_export('feature_column.embedding_column') def embedding_column( categorical_column, dimension, combiner='mean', initializer=None, ckpt_to_load_from=None, tensor_name_in_ckpt=None, max_norm=None, @@ -657,6 +664,7 @@ def embedding_column( trainable=trainable) +@tf_export('feature_column.shared_embedding_columns') def shared_embedding_columns( categorical_columns, dimension, combiner='mean', initializer=None, shared_embedding_collection_name=None, ckpt_to_load_from=None, @@ -807,6 +815,7 @@ def shared_embedding_columns( return result +@tf_export('feature_column.numeric_column') def numeric_column(key, shape=(1,), default_value=None, @@ -881,6 +890,7 @@ def numeric_column(key, normalizer_fn=normalizer_fn) +@tf_export('feature_column.bucketized_column') def bucketized_column(source_column, boundaries): """Represents discretized dense input. @@ -970,6 +980,7 @@ def _assert_string_or_int(dtype, prefix): '{} dtype must be string or integer. dtype: {}.'.format(prefix, dtype)) +@tf_export('feature_column.categorical_column_with_hash_bucket') def categorical_column_with_hash_bucket(key, hash_bucket_size, dtype=dtypes.string): @@ -1026,6 +1037,7 @@ def categorical_column_with_hash_bucket(key, return _HashedCategoricalColumn(key, hash_bucket_size, dtype) +@tf_export('feature_column.categorical_column_with_vocabulary_file') def categorical_column_with_vocabulary_file(key, vocabulary_file, vocabulary_size=None, @@ -1145,6 +1157,7 @@ def categorical_column_with_vocabulary_file(key, dtype=dtype) +@tf_export('feature_column.categorical_column_with_vocabulary_list') def categorical_column_with_vocabulary_list( key, vocabulary_list, dtype=None, default_value=-1, num_oov_buckets=0): """A `_CategoricalColumn` with in-memory vocabulary. @@ -1255,6 +1268,7 @@ def categorical_column_with_vocabulary_list( default_value=default_value, num_oov_buckets=num_oov_buckets) +@tf_export('feature_column.categorical_column_with_identity') def categorical_column_with_identity(key, num_buckets, default_value=None): """A `_CategoricalColumn` that returns identity values. @@ -1322,6 +1336,7 @@ def categorical_column_with_identity(key, num_buckets, default_value=None): key=key, num_buckets=num_buckets, default_value=default_value) +@tf_export('feature_column.indicator_column') def indicator_column(categorical_column): """Represents multi-hot representation of given categorical column. @@ -1350,6 +1365,7 @@ def indicator_column(categorical_column): return _IndicatorColumn(categorical_column) +@tf_export('feature_column.weighted_categorical_column') def weighted_categorical_column( categorical_column, weight_feature_key, dtype=dtypes.float32): """Applies weight values to a `_CategoricalColumn`. @@ -1424,6 +1440,7 @@ def weighted_categorical_column( dtype=dtype) +@tf_export('feature_column.crossed_column') def crossed_column(keys, hash_bucket_size, hash_key=None): """Returns a column for performing crosses of categorical features. @@ -1609,7 +1626,7 @@ class _FeatureColumn(object): It is used for get_parsing_spec for `tf.parse_example`. Returned spec is a dict from keys ('string') to `VarLenFeature`, `FixedLenFeature`, and other - supported objects. Please check documentation of ${tf.parse_example} for all + supported objects. Please check documentation of @{tf.parse_example} for all supported spec objects. Let's say a Feature column depends on raw feature ('raw') and another @@ -1660,7 +1677,7 @@ class _DenseColumn(_FeatureColumn): weight_collections: List of graph collections to which Variables (if any will be created) are added. trainable: If `True` also add variables to the graph collection - `GraphKeys.TRAINABLE_VARIABLES` (see ${tf.Variable}). + `GraphKeys.TRAINABLE_VARIABLES` (see @{tf.Variable}). Returns: `Tensor` of shape [batch_size] + `_variable_shape`. @@ -1718,7 +1735,7 @@ class _CategoricalColumn(_FeatureColumn): WARNING: Do not subclass this layer unless you know what you are doing: the API is subject to future changes. - A categorical feature typically handled with a ${tf.SparseTensor} of IDs. + A categorical feature typically handled with a @{tf.SparseTensor} of IDs. """ __metaclass__ = abc.ABCMeta @@ -1753,7 +1770,7 @@ class _CategoricalColumn(_FeatureColumn): weight_collections: List of graph collections to which variables (if any will be created) are added. trainable: If `True` also add variables to the graph collection - `GraphKeys.TRAINABLE_VARIABLES` (see ${tf.get_variable}). + `GraphKeys.TRAINABLE_VARIABLES` (see @{tf.get_variable}). """ pass @@ -1787,6 +1804,21 @@ def _create_categorical_column_weighted_sum( name='weighted_sum') +class _SequenceDenseColumn(_FeatureColumn): + """Represents dense sequence data.""" + + __metaclass__ = abc.ABCMeta + + TensorSequenceLengthPair = collections.namedtuple( # pylint: disable=invalid-name + 'TensorSequenceLengthPair', ['dense_tensor', 'sequence_length']) + + @abc.abstractmethod + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + """Returns a `TensorSequenceLengthPair`.""" + pass + + class _LazyBuilder(object): """Handles caching of transformations while building the model. @@ -1857,12 +1889,12 @@ class _LazyBuilder(object): self._feature_tensors[key] = feature_tensor return feature_tensor - if not isinstance(key, (str, _FeatureColumn)): - raise TypeError('"key" must be either a "str" or "_FeatureColumn". ' - 'Provided: {}'.format(key)) + if isinstance(key, str): + raise ValueError('Feature {} is not in features dictionary.'.format(key)) if not isinstance(key, _FeatureColumn): - raise ValueError('Feature {} is not in features dictionary.'.format(key)) + raise TypeError('"key" must be either a "str" or "_FeatureColumn". ' + 'Provided: {}'.format(key)) column = key logging.debug('Transforming feature_column %s.', column) @@ -2135,7 +2167,7 @@ class _BucketizedColumn(_DenseColumn, _CategoricalColumn, class _EmbeddingColumn( - _DenseColumn, + _DenseColumn, _SequenceDenseColumn, collections.namedtuple('_EmbeddingColumn', ( 'categorical_column', 'dimension', 'combiner', 'initializer', 'ckpt_to_load_from', 'tensor_name_in_ckpt', 'max_norm', 'trainable' @@ -2161,7 +2193,9 @@ class _EmbeddingColumn( self._shape = tensor_shape.vector(self.dimension) return self._shape - def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): + def _get_dense_tensor_internal( + self, inputs, weight_collections=None, trainable=None): + """Private method that follows the signature of _get_dense_tensor.""" # Get sparse IDs and weights. sparse_tensors = self.categorical_column._get_sparse_tensors( # pylint: disable=protected-access inputs, weight_collections=weight_collections, trainable=trainable) @@ -2193,6 +2227,43 @@ class _EmbeddingColumn( name='%s_weights' % self.name, max_norm=self.max_norm) + def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): + if isinstance(self.categorical_column, _SequenceCategoricalColumn): + raise ValueError( + 'In embedding_column: {}. ' + 'categorical_column must not be of type _SequenceCategoricalColumn. ' + 'Suggested fix A: If you wish to use input_layer, use a ' + 'non-sequence categorical_column_with_*. ' + 'Suggested fix B: If you wish to create sequence input, use ' + 'sequence_input_layer instead of input_layer. ' + 'Given (type {}): {}'.format( + self.name, type(self.categorical_column), + self.categorical_column)) + return self._get_dense_tensor_internal( + inputs=inputs, weight_collections=weight_collections, + trainable=trainable) + + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + if not isinstance(self.categorical_column, _SequenceCategoricalColumn): + raise ValueError( + 'In embedding_column: {}. ' + 'categorical_column must be of type _SequenceCategoricalColumn ' + 'to use sequence_input_layer. ' + 'Suggested fix: Use one of sequence_categorical_column_with_*. ' + 'Given (type {}): {}'.format( + self.name, type(self.categorical_column), + self.categorical_column)) + dense_tensor = self._get_dense_tensor_internal( # pylint: disable=protected-access + inputs=inputs, + weight_collections=weight_collections, + trainable=trainable) + sparse_tensors = self.categorical_column._get_sparse_tensors(inputs) # pylint: disable=protected-access + sequence_length = _sequence_length_from_sparse_tensor( + sparse_tensors.id_tensor) + return _SequenceDenseColumn.TensorSequenceLengthPair( + dense_tensor=dense_tensor, sequence_length=sequence_length) + class _SharedEmbeddingColumn( _DenseColumn, @@ -2873,7 +2944,7 @@ def _prune_invalid_ids(sparse_ids, sparse_weights): return sparse_ids, sparse_weights -class _IndicatorColumn(_DenseColumn, +class _IndicatorColumn(_DenseColumn, _SequenceDenseColumn, collections.namedtuple('_IndicatorColumn', ['categorical_column'])): """Represents a one-hot column for use in deep networks. @@ -2949,15 +3020,53 @@ class _IndicatorColumn(_DenseColumn, Returns: Dense `Tensor` created within `_transform_feature`. + + Raises: + ValueError: If `categorical_column` is a `_SequenceCategoricalColumn`. """ # Do nothing with weight_collections and trainable since no variables are # created in this function. del weight_collections del trainable + if isinstance(self.categorical_column, _SequenceCategoricalColumn): + raise ValueError( + 'In indicator_column: {}. ' + 'categorical_column must not be of type _SequenceCategoricalColumn. ' + 'Suggested fix A: If you wish to use input_layer, use a ' + 'non-sequence categorical_column_with_*. ' + 'Suggested fix B: If you wish to create sequence input, use ' + 'sequence_input_layer instead of input_layer. ' + 'Given (type {}): {}'.format( + self.name, type(self.categorical_column), + self.categorical_column)) # Feature has been already transformed. Return the intermediate # representation created by _transform_feature. return inputs.get(self) + def _get_sequence_dense_tensor( + self, inputs, weight_collections=None, trainable=None): + # Do nothing with weight_collections and trainable since no variables are + # created in this function. + del weight_collections + del trainable + if not isinstance(self.categorical_column, _SequenceCategoricalColumn): + raise ValueError( + 'In indicator_column: {}. ' + 'categorical_column must be of type _SequenceCategoricalColumn ' + 'to use sequence_input_layer. ' + 'Suggested fix: Use one of sequence_categorical_column_with_*. ' + 'Given (type {}): {}'.format( + self.name, type(self.categorical_column), + self.categorical_column)) + # Feature has been already transformed. Return the intermediate + # representation created by _transform_feature. + dense_tensor = inputs.get(self) + sparse_tensors = self.categorical_column._get_sparse_tensors(inputs) # pylint: disable=protected-access + sequence_length = _sequence_length_from_sparse_tensor( + sparse_tensors.id_tensor) + return _SequenceDenseColumn.TensorSequenceLengthPair( + dense_tensor=dense_tensor, sequence_length=sequence_length) + def _verify_static_batch_size_equality(tensors, columns): # bath_size is a tf.Dimension object. @@ -2973,3 +3082,68 @@ def _verify_static_batch_size_equality(tensors, columns): 'Batch size of columns ({}, {}): ({}, {})'.format( columns[bath_size_column_index].name, columns[i].name, expected_batch_size, tensors[i].shape[0])) + + +def _sequence_length_from_sparse_tensor(sp_tensor, num_elements=1): + """Returns a [batch_size] Tensor with per-example sequence length.""" + with ops.name_scope(None, 'sequence_length') as name_scope: + row_ids = sp_tensor.indices[:, 0] + column_ids = sp_tensor.indices[:, 1] + column_ids += array_ops.ones_like(column_ids) + seq_length = math_ops.to_int64( + math_ops.segment_max(column_ids, segment_ids=row_ids) / num_elements) + # If the last n rows do not have ids, seq_length will have shape + # [batch_size - n]. Pad the remaining values with zeros. + n_pad = array_ops.shape(sp_tensor)[:1] - array_ops.shape(seq_length)[:1] + padding = array_ops.zeros(n_pad, dtype=seq_length.dtype) + return array_ops.concat([seq_length, padding], axis=0, name=name_scope) + + +class _SequenceCategoricalColumn( + _CategoricalColumn, + collections.namedtuple( + '_SequenceCategoricalColumn', ['categorical_column'])): + """Represents sequences of categorical data.""" + + @property + def name(self): + return self.categorical_column.name + + @property + def _parse_example_spec(self): + return self.categorical_column._parse_example_spec # pylint: disable=protected-access + + def _transform_feature(self, inputs): + return self.categorical_column._transform_feature(inputs) # pylint: disable=protected-access + + @property + def _num_buckets(self): + return self.categorical_column._num_buckets # pylint: disable=protected-access + + def _get_sparse_tensors(self, inputs, weight_collections=None, + trainable=None): + sparse_tensors = self.categorical_column._get_sparse_tensors(inputs) # pylint: disable=protected-access + id_tensor = sparse_tensors.id_tensor + weight_tensor = sparse_tensors.weight_tensor + # Expands final dimension, so that embeddings are not combined during + # embedding lookup. + check_id_rank = check_ops.assert_equal( + array_ops.rank(id_tensor), 2, + data=[ + 'Column {} expected ID tensor of rank 2. '.format(self.name), + 'id_tensor shape: ', array_ops.shape(id_tensor)]) + with ops.control_dependencies([check_id_rank]): + id_tensor = sparse_ops.sparse_reshape( + id_tensor, + shape=array_ops.concat([id_tensor.dense_shape, [1]], axis=0)) + if weight_tensor is not None: + check_weight_rank = check_ops.assert_equal( + array_ops.rank(weight_tensor), 2, + data=[ + 'Column {} expected weight tensor of rank 2.'.format(self.name), + 'weight_tensor shape:', array_ops.shape(weight_tensor)]) + with ops.control_dependencies([check_weight_rank]): + weight_tensor = sparse_ops.sparse_reshape( + weight_tensor, + shape=array_ops.concat([weight_tensor.dense_shape, [1]], axis=0)) + return _CategoricalColumn.IdWeightPair(id_tensor, weight_tensor) diff --git a/tensorflow/python/framework/common_shapes.py b/tensorflow/python/framework/common_shapes.py index 3b1092f923112dbd9a081942d40162ae384bf167..3c5aebbce8af117aa1e216f1ef07ded181c997ea 100644 --- a/tensorflow/python/framework/common_shapes.py +++ b/tensorflow/python/framework/common_shapes.py @@ -34,7 +34,7 @@ def scalar_shape(unused_op): def unchanged_shape(op): - """Shape function for ops that output an tensor like their first input.""" + """Shape function for ops that output a tensor like their first input.""" return [op.inputs[0].get_shape()] diff --git a/tensorflow/python/framework/constant_op.py b/tensorflow/python/framework/constant_op.py index d3d8c9c154fbfcc9613acce4e1bdab7df2e7d56d..782b505d6c1d0b576b7734f088c4d2c9625f4be2 100644 --- a/tensorflow/python/framework/constant_op.py +++ b/tensorflow/python/framework/constant_op.py @@ -181,7 +181,7 @@ def constant(value, dtype=None, shape=None, name="Const", verify_shape=False): TypeError: if shape is incorrectly specified or unsupported. """ ctx = context.context() - if not ctx.in_graph_mode(): + if ctx.executing_eagerly(): t = convert_to_eager_tensor(value, ctx, dtype) if shape is None: return t diff --git a/tensorflow/python/framework/dtypes.py b/tensorflow/python/framework/dtypes.py index c825114483697f63163d5a33365049b7c2b539a3..0edae92fd4a86e7d10a180ce64364d3ea552bf60 100644 --- a/tensorflow/python/framework/dtypes.py +++ b/tensorflow/python/framework/dtypes.py @@ -238,9 +238,9 @@ class DType(object): min, max : tuple Lower and upper intensity limits. """ - min, max = dtype_range[self.as_numpy_dtype] + min, max = dtype_range[self.as_numpy_dtype] # pylint: disable=redefined-builtin if clip_negative: - min = 0 + min = 0 # pylint: disable=redefined-builtin return min, max def is_compatible_with(self, other): @@ -343,7 +343,9 @@ tf_export("uint8").export_constant(__name__, "uint8") uint16 = DType(types_pb2.DT_UINT16) tf_export("uint16").export_constant(__name__, "uint16") uint32 = DType(types_pb2.DT_UINT32) +tf_export("uint32").export_constant(__name__, "uint32") uint64 = DType(types_pb2.DT_UINT64) +tf_export("uint64").export_constant(__name__, "uint32") int16 = DType(types_pb2.DT_INT16) tf_export("int16").export_constant(__name__, "int16") int8 = DType(types_pb2.DT_INT8) @@ -356,7 +358,7 @@ complex128 = DType(types_pb2.DT_COMPLEX128) tf_export("complex128").export_constant(__name__, "complex128") int64 = DType(types_pb2.DT_INT64) tf_export("int64").export_constant(__name__, "int64") -bool = DType(types_pb2.DT_BOOL) +bool = DType(types_pb2.DT_BOOL) # pylint: disable=redefined-builtin tf_export("bool").export_constant(__name__, "bool") qint8 = DType(types_pb2.DT_QINT8) tf_export("qint8").export_constant(__name__, "qint8") diff --git a/tensorflow/python/framework/framework_lib.py b/tensorflow/python/framework/framework_lib.py index d16fe979e6ef9a41063c3a2b3e8a3e18de2aa9d7..392a4f65c6e62c3cb70f8e02a9b24f015a09f649 100644 --- a/tensorflow/python/framework/framework_lib.py +++ b/tensorflow/python/framework/framework_lib.py @@ -48,6 +48,7 @@ ## Graph collections @@add_to_collection +@@add_to_collections @@get_collection @@get_collection_ref @@GraphKeys @@ -92,6 +93,7 @@ from tensorflow.python.framework.ops import get_default_graph from tensorflow.python.framework.ops import reset_default_graph from tensorflow.python.framework.ops import GraphKeys from tensorflow.python.framework.ops import add_to_collection +from tensorflow.python.framework.ops import add_to_collections from tensorflow.python.framework.ops import get_collection from tensorflow.python.framework.ops import get_collection_ref from tensorflow.python.framework.ops import convert_to_tensor @@ -118,7 +120,7 @@ from tensorflow.python.framework.ops import register_tensor_conversion_function # go/tf-wildcard-import # pylint: disable=wildcard-import -from tensorflow.python.framework.dtypes import * +from tensorflow.python.framework.dtypes import * # pylint: disable=redefined-builtin # Load a TensorFlow plugin from tensorflow.python.framework.load_library import * diff --git a/tensorflow/python/framework/function.py b/tensorflow/python/framework/function.py index cba225e749d88a45c43266e45172a7335a8e0b71..14d72d8a3de7e22bee4f9961c2f66044c217f641 100644 --- a/tensorflow/python/framework/function.py +++ b/tensorflow/python/framework/function.py @@ -353,8 +353,10 @@ class _DefinedFunction(object): outputs = (outputs,) if any([_ is None for _ in outputs]): raise ValueError("Function can not return None.") - # Ensures each output is a Tensor. - outputs = [ops.convert_to_tensor(_) for _ in outputs] + # Ensures each output is a Tensor in the function graph. + outputs = [ops.convert_to_tensor(t) for t in outputs] + outputs = [temp_graph.capture(t) if t.graph is not temp_graph else t + for t in outputs] self._extra_inputs = temp_graph.extra_inputs inputs.extend(temp_graph.extra_args) # pylint: disable=protected-access @@ -487,10 +489,10 @@ class _DefinedFunction(object): # Adds this function into 'g'. # pylint: disable=protected-access - if context.in_graph_mode(): - g._add_function(self) - else: + if context.executing_eagerly(): context.context().add_function_def(self.definition) + else: + g._add_function(self) # pylint: enable=protected-access # Ensures related sub-routines are defined in 'g', too. @@ -683,28 +685,34 @@ class _FuncGraph(ops.Graph): def create_op(self, op_type, inputs, data_types, **kwargs): for i, x in enumerate(inputs): if isinstance(x, ops.EagerTensor) or x.graph is not self: - # Referring to a tensor from other graph. - if x in self._captured: - # Captured already. - inputs[i] = self._captured[x] - elif self._capture_by_value: - inputs[i] = self._add_tensor_and_parents(x) - else: - # Substitute with a placeholder. - self.extra_inputs.append(x) - # Hoist the new input placeholder out of any control flow context - # we're currently in. - with ops.control_dependencies(None): - ph = array_ops.placeholder(x.dtype, shape=x.get_shape()) - # pylint: disable=protected-access - ph._handle_data = x._handle_data - # pylint: enable=protected-access - inputs[i] = ph - self._captured[x] = ph - self.extra_args.append(ph) + inputs[i] = self.capture(x) return super(_FuncGraph, self).create_op(op_type, inputs, data_types, **kwargs) + def capture(self, tensor): + """Adds the given tensor to this graph and returns the captured tensor.""" + if tensor in self._captured: + # Captured already. + return self._captured[tensor] + elif self._capture_by_value: + return self._add_tensor_and_parents(tensor) + else: + return self._capture_tensor_as_extra_input(tensor) + + def _capture_tensor_as_extra_input(self, tensor): + # Substitute with a placeholder. + self.extra_inputs.append(tensor) + # Hoist the new input placeholder out of any control flow context + # we're currently in. + with ops.control_dependencies(None): + ph = array_ops.placeholder(tensor.dtype, shape=tensor.get_shape()) + # pylint: disable=protected-access + ph._handle_data = tensor._handle_data + # pylint: enable=protected-access + self._captured[tensor] = ph + self.extra_args.append(ph) + return ph + def _add_tensor_and_parents(self, tensor): op = self._add_op_and_parents(tensor.op) return op.outputs[tensor.value_index] diff --git a/tensorflow/python/framework/function_test.py b/tensorflow/python/framework/function_test.py index a4ca3f9a89bd4cce2240d90895c43dda1acb849b..65ca801cbe922b36e3bc72bc2fbcd88f66aa5290 100644 --- a/tensorflow/python/framework/function_test.py +++ b/tensorflow/python/framework/function_test.py @@ -19,8 +19,8 @@ from __future__ import division from __future__ import print_function import re -import time import sys +import time import numpy as np @@ -86,6 +86,21 @@ class FunctionTest(test.TestCase): with session.Session() as sess: self.assertAllEqual([18.0], sess.run(call)) + def testIdentityImplicitDeref(self): + + @function.Defun(dtypes.float32, func_name="MyIdentity") + def MyIdentityFunc(a): + return a + + with ops.Graph().as_default(): + var = variables.Variable([18.0]) + call = MyIdentityFunc(var._ref()) # pylint: disable=protected-access + self.assertEqual("MyIdentity", call.op.name) + for cfg in _OptimizerOptions(): + with session.Session(config=cfg) as sess: + sess.run(var.initializer) + self.assertAllEqual([18.0], sess.run(call)) + def testIdentityOutputName(self): @function.Defun( @@ -178,7 +193,7 @@ class FunctionTest(test.TestCase): @function.Defun(dtypes.float32, dtypes.float32) def XSquarePlusOneGrad(x, dy): - dx = functional_ops._symbolic_gradient( + dx = functional_ops.symbolic_gradient( input=[x, dy], Tout=[dtypes.float32], f="XSquarePlusOneFn", name="dx") return dx @@ -280,7 +295,7 @@ class FunctionTest(test.TestCase): # gradient function is (x, y, dz) -> (dx, dy). dx's shape # should be the same as x's; and dy's shape should be the same # as y's. - dx, dy = functional_ops._symbolic_gradient( + dx, dy = functional_ops.symbolic_gradient( input=[x, y, dz], Tout=[dtypes.float32] * 2, f="Foo") self.assertEqual(x.get_shape(), dx.get_shape()) self.assertEqual(y.get_shape(), dy.get_shape()) @@ -710,9 +725,16 @@ class FunctionTest(test.TestCase): y = Foo(constant_op.constant([[10.]])) + @function.Defun() + def Bar(): + return w + + z = Bar() + with self.test_session(graph=g): variables.global_variables_initializer().run() self.assertAllEqual(y.eval(), [[12.0]]) + self.assertAllEqual(z.eval(), [[1.0]]) def testCaptureControls(self): g = ops.Graph() @@ -771,7 +793,7 @@ class FunctionTest(test.TestCase): # We added more randomness to function names in C API. # TODO(iga): Remove this if statement when we switch to C API. if ops._USE_C_API: # pylint: disable=protected-access - if sys.byteorder == 'big': + if sys.byteorder == "big": self.assertEqual("Foo_kEdkAG8SJvg", Foo.instantiate([dtypes.float32] * 3).name) else: @@ -1443,7 +1465,7 @@ class FunctionInlineControlTest(test.TestCase): def Cell(v): # If v is a vector [n, 1], x is a big square matrix. x = math_ops.tanh(v + array_ops.transpose(v, [1, 0])) - return math_ops.reduce_sum(x, 1, keep_dims=True) + return math_ops.reduce_sum(x, 1, keepdims=True) @function.Defun(dtype) def Forward(x): diff --git a/tensorflow/python/framework/graph_util_impl.py b/tensorflow/python/framework/graph_util_impl.py index 5a543317e665a940841714fd72d834a430f8406a..910364364c8be84b1a629dbdaae5e69443d07e75 100644 --- a/tensorflow/python/framework/graph_util_impl.py +++ b/tensorflow/python/framework/graph_util_impl.py @@ -235,7 +235,7 @@ def convert_variables_to_constants(sess, variable_names = [] variable_dict_names = [] for node in inference_graph.node: - if node.op in ["Variable", "VariableV2"]: + if node.op in ["Variable", "VariableV2", "VarHandleOp"]: variable_name = node.name if ((variable_names_whitelist is not None and variable_name not in variable_names_whitelist) or @@ -243,7 +243,10 @@ def convert_variables_to_constants(sess, variable_name in variable_names_blacklist)): continue variable_dict_names.append(variable_name) - variable_names.append(variable_name + ":0") + if node.op == "VarHandleOp": + variable_names.append(variable_name + "/Read/ReadVariableOp:0") + else: + variable_names.append(variable_name + ":0") if variable_names: returned_variables = sess.run(variable_names) else: @@ -266,6 +269,17 @@ def convert_variables_to_constants(sess, tensor=tensor_util.make_tensor_proto( data, dtype=dtype.type, shape=data.shape))) how_many_converted += 1 + elif input_node.op == "ReadVariableOp" and ( + input_node.input[0] in found_variables): + # The preceding branch converts all VarHandleOps of ResourceVariables to + # constants, so we need to convert the associated ReadVariableOps to + # Identity ops. + output_node.op = "Identity" + output_node.name = input_node.name + output_node.input.extend([input_node.input[0]]) + output_node.attr["T"].CopyFrom(input_node.attr["dtype"]) + if "_class" in input_node.attr: + output_node.attr["_class"].CopyFrom(input_node.attr["_class"]) else: output_node.CopyFrom(input_node) output_graph_def.node.extend([output_node]) diff --git a/tensorflow/python/framework/graph_util_test.py b/tensorflow/python/framework/graph_util_test.py index 0421837d49de753d642aed59d1524619a243dcb8..b618152b0256fd043dc7259960d867278ba55b0a 100644 --- a/tensorflow/python/framework/graph_util_test.py +++ b/tensorflow/python/framework/graph_util_test.py @@ -32,6 +32,7 @@ from tensorflow.python.framework import tensor_util from tensorflow.python.ops import gen_state_ops from tensorflow.python.ops import math_ops # pylint: disable=unused-import from tensorflow.python.ops import math_ops as math_ops_lib +from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test @@ -47,46 +48,46 @@ class DeviceFunctionsTest(test.TestCase): def testTwoDeviceFunctions(self): with ops.Graph().as_default() as g: - var_0 = gen_state_ops._variable( + var_0 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_0", container="", shared_name="") with g.device(test_device_func_pin_variable_to_cpu): - var_1 = gen_state_ops._variable( + var_1 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_1", container="", shared_name="") - var_2 = gen_state_ops._variable( + var_2 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_2", container="", shared_name="") - var_3 = gen_state_ops._variable( + var_3 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_3", container="", shared_name="") with g.device(test_device_func_pin_variable_to_cpu): - var_4 = gen_state_ops._variable( + var_4 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_4", container="", shared_name="") with g.device("/device:GPU:0"): - var_5 = gen_state_ops._variable( + var_5 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_5", container="", shared_name="") - var_6 = gen_state_ops._variable( + var_6 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="var_6", @@ -226,52 +227,62 @@ class DeviceFunctionsTest(test.TestCase): constant_graph_def.library) def testConvertVariablesToConsts(self): - with ops.Graph().as_default(): - variable_node = variables.Variable(1.0, name="variable_node") - _ = variables.Variable(1.0, name="unused_variable_node") - output_node = math_ops_lib.multiply( - variable_node, 2.0, name="output_node") - with session.Session() as sess: - init = variables.initialize_variables([variable_node]) - sess.run(init) - output = sess.run(output_node) - self.assertNear(2.0, output, 0.00001) - variable_graph_def = sess.graph.as_graph_def() - # First get the constant_graph_def when variable_names_whitelist is set, - # note that if variable_names_whitelist is not set an error will be - # thrown because unused_variable_node is not initialized. - constant_graph_def = graph_util.convert_variables_to_constants( - sess, - variable_graph_def, ["output_node"], - variable_names_whitelist=set(["variable_node"])) + self._test_variable_to_const_conversion(use_resource=False) - # Then initialize the unused variable, and get another - # constant_graph_def when variable_names_whitelist is not set. - sess.run(variables.global_variables_initializer()) - constant_graph_def_without_variable_whitelist = ( - graph_util.convert_variables_to_constants(sess, variable_graph_def, - ["output_node"])) - - # The unused variable should be cleared so the two graphs should be - # equivalent. - self.assertEqual( - str(constant_graph_def), - str(constant_graph_def_without_variable_whitelist)) - - # Test variable name black list. This should result in the variable not - # being a const. - sess.run(variables.global_variables_initializer()) - constant_graph_def_with_blacklist = ( - graph_util.convert_variables_to_constants( - sess, - variable_graph_def, ["output_node"], - variable_names_blacklist=set(["variable_node"]))) - variable_node = None - for node in constant_graph_def_with_blacklist.node: - if node.name == "variable_node": - variable_node = node - self.assertIsNotNone(variable_node) - self.assertEqual(variable_node.op, "VariableV2") + def testConvertResourceVariablesToConsts(self): + self._test_variable_to_const_conversion(use_resource=True) + + def _test_variable_to_const_conversion(self, use_resource): + with ops.Graph().as_default(): + with variable_scope.variable_scope("", use_resource=use_resource): + variable_node = variable_scope.get_variable( + "variable_node", initializer=1.0) + another_variable = variable_scope.get_variable( + "unused_variable_node", initializer=1.0) + output_node = math_ops_lib.multiply( + variable_node, 2.0, name="output_node") + with session.Session() as sess: + sess.run(variable_node.initializer) + output = sess.run(output_node) + self.assertNear(2.0, output, 0.00001) + variable_graph_def = sess.graph.as_graph_def() + # First get the constant_graph_def when variable_names_whitelist is + # set, note that if variable_names_whitelist is not set an error will + # be thrown because unused_variable_node is not initialized. + constant_graph_def = graph_util.convert_variables_to_constants( + sess, + variable_graph_def, ["output_node"], + variable_names_whitelist=set(["variable_node"])) + + # Then initialize the unused variable, and get another + # constant_graph_def when variable_names_whitelist is not set. + sess.run(another_variable.initializer) + constant_graph_def_without_variable_whitelist = ( + graph_util.convert_variables_to_constants( + sess, variable_graph_def, ["output_node"])) + + # The unused variable should be cleared so the two graphs should be + # equivalent. + self.assertEqual( + str(constant_graph_def), + str(constant_graph_def_without_variable_whitelist)) + + # Test variable name black list. This should result in the variable + # not being a const. + constant_graph_def_with_blacklist = ( + graph_util.convert_variables_to_constants( + sess, + variable_graph_def, ["output_node"], + variable_names_blacklist=set(["variable_node"]))) + variable_node = None + for node in constant_graph_def_with_blacklist.node: + if node.name == "variable_node": + variable_node = node + self.assertIsNotNone(variable_node) + if use_resource: + self.assertEqual(variable_node.op, "VarHandleOp") + else: + self.assertEqual(variable_node.op, "VariableV2") # Now we make sure the variable is now a constant, and that the graph still # produces the expected result. @@ -279,8 +290,9 @@ class DeviceFunctionsTest(test.TestCase): _ = importer.import_graph_def(constant_graph_def, name="") self.assertEqual(4, len(constant_graph_def.node)) for node in constant_graph_def.node: - self.assertNotEqual("Variable", node.op) - self.assertNotEqual("VariableV2", node.op) + self.assertNotIn( + node.op, + ["Variable", "VariableV2", "VarHandleOp", "ReadVariableOp"]) with session.Session() as sess: output_node = sess.graph.get_tensor_by_name("output_node:0") output = sess.run(output_node) diff --git a/tensorflow/python/framework/importer.py b/tensorflow/python/framework/importer.py index c26644362c86798012b5812f3554662a6d1d680b..783e9259ad1fae4eb05fc7beeb5be83ae3f7fdce 100644 --- a/tensorflow/python/framework/importer.py +++ b/tensorflow/python/framework/importer.py @@ -150,7 +150,7 @@ def _MaybeDevice(device): yield -def _ProcessGraphDefParam(graph_def): +def _ProcessGraphDefParam(graph_def, op_dict): """Type-checks and possibly canonicalizes `graph_def`.""" if not isinstance(graph_def, graph_pb2.GraphDef): # `graph_def` could be a dynamically-created message, so try a duck-typed @@ -161,6 +161,22 @@ def _ProcessGraphDefParam(graph_def): graph_def.MergeFrom(old_graph_def) except TypeError: raise TypeError('graph_def must be a GraphDef proto.') + else: + # If we're using the graph_def provided by the caller, modify graph_def + # in-place to add attr defaults to the NodeDefs (this is visible to the + # caller). + # NOTE(skyewm): this is undocumented behavior that at least meta_graph.py + # depends on. It might make sense to move this to meta_graph.py and have + # import_graph_def not modify the graph_def argument (we'd have to make sure + # this doesn't break anything else.) + for node in graph_def.node: + if node.op not in op_dict: + # Assume unrecognized ops are functions for now. TF_ImportGraphDef will + # report an error if the op is actually missing. + continue + op_def = op_dict[node.op] + _SetDefaultAttrValues(node, op_def) + return graph_def @@ -254,7 +270,6 @@ def _PopulateTFImportGraphDefOptions(options, prefix, input_map, """Populates the TF_ImportGraphDefOptions `options`.""" c_api.TF_ImportGraphDefOptionsSetPrefix(options, prefix) c_api.TF_ImportGraphDefOptionsSetUniquifyNames(options, True) - c_api.TF_ImportGraphDefOptionsSetUniquifyPrefix(options, True) for input_src, input_dst in input_map.items(): input_src = compat.as_str(input_src) @@ -369,6 +384,17 @@ def _GatherReturnElements(requested_return_elements, graph, results): return combined_return_elements +def _SetDefaultAttrValues(node_def, op_def): + """Set any default attr values in `node_def` that aren't present.""" + assert node_def.op == op_def.name + for attr_def in op_def.attr: + key = attr_def.name + if attr_def.HasField('default_value'): + value = node_def.attr[key] + if value is None or value.WhichOneof('value') is None: + node_def.attr[key].CopyFrom(attr_def.default_value) + + @tf_export('import_graph_def') @deprecated_args(None, 'Please file an issue at ' 'https://github.com/tensorflow/tensorflow/issues if you depend' @@ -420,12 +446,12 @@ def import_graph_def(graph_def, do not appear in `graph_def`, or `graph_def` is not well-formed (e.g. it refers to an unknown tensor). """ - graph_def = _ProcessGraphDefParam(graph_def) + op_dict = op_def_registry.get_registered_ops() + + graph_def = _ProcessGraphDefParam(graph_def, op_dict) input_map = _ProcessInputMapParam(input_map) return_elements = _ProcessReturnElementsParam(return_elements) - op_dict = op_def_registry.get_registered_ops() - if producer_op_list is not None: # TODO(skyewm): make a copy of graph_def so we're not mutating the argument? _RemoveDefaultAttrs(op_dict, producer_op_list, graph_def) @@ -449,16 +475,21 @@ def import_graph_def(graph_def, _PopulateTFImportGraphDefOptions(options, prefix, input_map, return_elements) - with c_api_util.tf_buffer(graph_def.SerializeToString()) as serialized: - try: - with errors.raise_exception_on_not_ok_status() as status: - results = c_api.TF_GraphImportGraphDefWithResults( - graph._c_graph, serialized, options, status) # pylint: disable=protected-access - except errors.InvalidArgumentError as e: - # Convert to ValueError for backwards compatibility. - raise ValueError(str(e)) + # _ProcessNewOps mutates the new operations. _lock ensures a Session.run + # call cannot occur between creating the TF_Operations in the + # TF_GraphImportGraphDefWithResults call and mutating the them in + # _ProcessNewOps. + with graph._lock: # pylint: disable=protected-access + with c_api_util.tf_buffer(graph_def.SerializeToString()) as serialized: + try: + with errors.raise_exception_on_not_ok_status() as status: + results = c_api.TF_GraphImportGraphDefWithResults( + graph._c_graph, serialized, options, status) # pylint: disable=protected-access + except errors.InvalidArgumentError as e: + # Convert to ValueError for backwards compatibility. + raise ValueError(str(e)) - _ProcessNewOps(graph) + _ProcessNewOps(graph) # Create _DefinedFunctions for any imported functions. # @@ -535,16 +566,9 @@ def import_graph_def(graph_def, # Check to see if this op's name matches a previously seen op if node.name in name_to_op: raise ValueError('Duplicate name \'%s\' in GraphDef.' % node.name) - # Set any default attr values that aren't present. if node.op not in op_dict: raise ValueError('No op named %s in defined operations.' % node.op) op_def = op_dict[node.op] - for attr_def in op_def.attr: - key = attr_def.name - if attr_def.HasField('default_value'): - value = node.attr[key] - if value is None or value.WhichOneof('value') is None: - node.attr[key].CopyFrom(attr_def.default_value) output_types = _OutputTypes(node, op_dict) name_to_op[node.name] = g.create_op( diff --git a/tensorflow/python/framework/importer_test.py b/tensorflow/python/framework/importer_test.py index acaec37f810cb00daa9bae17ffbcb675648b9fe1..c39191e6d90f5f00806eba021ea41dd6ebe86f98 100644 --- a/tensorflow/python/framework/importer_test.py +++ b/tensorflow/python/framework/importer_test.py @@ -31,7 +31,6 @@ from tensorflow.python.framework import function from tensorflow.python.framework import importer from tensorflow.python.framework import ops from tensorflow.python.framework import test_ops # pylint: disable=unused-import -from tensorflow.python.framework import test_util from tensorflow.python.framework import versions from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -44,7 +43,8 @@ import tensorflow.python.ops.nn_grad # pylint: disable=unused-import from tensorflow.python.platform import test -@test_util.with_c_api +# TODO(skyewm): reenable when this works with _USE_C_SHAPES=False +# @test_util.with_c_api class ImportGraphDefTest(test.TestCase): def _MakeGraphDef(self, @@ -154,6 +154,25 @@ class ImportGraphDefTest(test.TestCase): self.assertEqual(b3.name, "A_3/B") self.assertEqual(list(b3.inputs), [a3.outputs[0]]) + # Import with an already-used name but with a '/' to indicate an + # "absolute" name scope (see the Graph.name_scope docstring). + a_a, a_b = importer.import_graph_def( + graph_def, + return_elements=["A", "B"], + name="A/") + self.assertEqual(a_a.name, "A/A") + self.assertEqual(a_b.name, "A/B") + self.assertEqual(list(a_b.inputs), [a_a.outputs[0]]) + + # Repeat the same import. + a_a1, a_b1 = importer.import_graph_def( + graph_def, + return_elements=["A", "B"], + name="A/") + self.assertEqual(a_a1.name, "A/A_1") + self.assertEqual(a_b1.name, "A/B_1") + self.assertEqual(list(a_b1.inputs), [a_a1.outputs[0]]) + # Import with existing de-duped node names a1_1, b1_1 = importer.import_graph_def( self._MakeGraphDef(""" diff --git a/tensorflow/python/framework/load_library.py b/tensorflow/python/framework/load_library.py index c997ead829855f33efdb3efe947c3f59b5dbe76c..1f2aa264c110930b318f30e3a24010a96ebce47e 100644 --- a/tensorflow/python/framework/load_library.py +++ b/tensorflow/python/framework/load_library.py @@ -21,10 +21,10 @@ from __future__ import print_function import hashlib import imp import sys -import threading +import threading # pylint: disable=unused-import from tensorflow.core.framework import op_def_pb2 -from tensorflow.core.lib.core import error_codes_pb2 +from tensorflow.core.lib.core import error_codes_pb2 # pylint: disable=unused-import from tensorflow.python import pywrap_tensorflow as py_tf from tensorflow.python.framework import errors_impl from tensorflow.python.util import compat diff --git a/tensorflow/python/framework/meta_graph.py b/tensorflow/python/framework/meta_graph.py index fc1a82361ba59cddc02a65a96da98283d871fd2c..4bb9941bb778ae9b4022ef9d376cb031223ddb1c 100644 --- a/tensorflow/python/framework/meta_graph.py +++ b/tensorflow/python/framework/meta_graph.py @@ -87,6 +87,10 @@ def _node_def(from_node_def, export_scope, unbound_inputs, clear_devices=False): compat.as_str(s).split("@")[1].startswith(export_scope)] node_def.attr[k].CopyFrom(attr_value_pb2.AttrValue( list=attr_value_pb2.AttrValue.ListValue(s=new_s))) + elif node_def.op in ("Enter", "RefEnter") and k == "frame_name": + if not export_scope or compat.as_str(v.s).startswith(export_scope): + new_s = compat.as_bytes(ops.strip_name_scope(v.s, export_scope)) + node_def.attr[k].CopyFrom(attr_value_pb2.AttrValue(s=new_s)) else: node_def.attr[k].CopyFrom(v) @@ -691,7 +695,7 @@ def import_scoped_meta_graph(meta_graph_or_file, Raises: ValueError: If the graph_def contains unbound inputs. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError("Exporting/importing meta graphs is not supported when " "eager execution is enabled.") if isinstance(meta_graph_or_file, meta_graph_pb2.MetaGraphDef): @@ -737,6 +741,7 @@ def import_scoped_meta_graph(meta_graph_or_file, producer_op_list=producer_op_list) # Restores all the other collections. + variable_objects = {} for key, col_def in sorted(meta_graph_def.collection_def.items()): # Don't add unbound_inputs to the new graph. if key == unbound_inputs_col_name: @@ -752,11 +757,23 @@ def import_scoped_meta_graph(meta_graph_or_file, from_proto = ops.get_from_proto_function(key) if from_proto and kind == "bytes_list": proto_type = ops.get_collection_proto_type(key) - for value in col_def.bytes_list.value: - proto = proto_type() - proto.ParseFromString(value) - graph.add_to_collection( - key, from_proto(proto, import_scope=scope_to_prepend_to_names)) + if key in ops.GraphKeys._VARIABLE_COLLECTIONS: # pylint: disable=protected-access + for value in col_def.bytes_list.value: + variable = variable_objects.get(value, None) + if variable is None: + proto = proto_type() + proto.ParseFromString(value) + variable = from_proto( + proto, import_scope=scope_to_prepend_to_names) + variable_objects[value] = variable + graph.add_to_collection(key, variable) + else: + for value in col_def.bytes_list.value: + proto = proto_type() + proto.ParseFromString(value) + graph.add_to_collection( + key, from_proto( + proto, import_scope=scope_to_prepend_to_names)) else: field = getattr(col_def, kind) if key in _COMPAT_COLLECTION_LIST: @@ -839,7 +856,7 @@ def export_scoped_meta_graph(filename=None, Raises: ValueError: When the `GraphDef` is larger than 2GB. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError("Exporting/importing meta graphs is not supported when " "Eager Execution is enabled.") graph = graph or ops.get_default_graph() @@ -959,5 +976,3 @@ def copy_scoped_meta_graph(from_scope, to_scope, graph=to_graph, import_scope=to_scope) return var_list - - diff --git a/tensorflow/python/framework/meta_graph_test.py b/tensorflow/python/framework/meta_graph_test.py index b5ed1352843eac31b3e34eb96385acd13a5bc7a9..06cec504e407301c7cfbceaa04a07674fe87c712 100644 --- a/tensorflow/python/framework/meta_graph_test.py +++ b/tensorflow/python/framework/meta_graph_test.py @@ -25,6 +25,7 @@ import shutil from tensorflow.core.framework import graph_pb2 from tensorflow.core.protobuf import meta_graph_pb2 +from tensorflow.python.client import session from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import function @@ -34,6 +35,7 @@ from tensorflow.python.framework import test_util 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 gradients_impl from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics from tensorflow.python.ops import nn_ops @@ -259,8 +261,32 @@ class SimpleMetaGraphTest(test.TestCase): self.assertEqual(node_def.attr["attr_1"].i, 1) self.assertTrue(meta_graph_def.meta_info_def.stripped_default_attrs) - -@test_util.with_c_api + def testVariableObjectsAreSharedAmongCollections(self): + with ops.Graph().as_default() as graph1: + v = variables.Variable(3.0) + # A single instance of Variable is shared among the collections: + global_vars = graph1.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + trainable_vars = graph1.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) + self.assertEqual(len(global_vars), 1) + self.assertEqual(len(trainable_vars), 1) + self.assertIs(global_vars[0], trainable_vars[0]) + self.assertIs(v, global_vars[0]) + + orig_meta_graph, _ = meta_graph.export_scoped_meta_graph(graph=graph1) + del graph1 # To avoid accidental references in code involving graph2. + + with ops.Graph().as_default() as graph2: + meta_graph.import_scoped_meta_graph(orig_meta_graph) + global_vars = graph2.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + trainable_vars = graph2.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) + self.assertEqual(len(global_vars), 1) + self.assertEqual(len(trainable_vars), 1) + # A single instance of Variable is shared among the collections: + self.assertIs(global_vars[0], trainable_vars[0]) + + +# TODO(skyewm): reenable when this works with _USE_C_SHAPES=False +# @test_util.with_c_api class ScopedMetaGraphTest(test.TestCase): def _testScopedExport(self, test_dir, exported_filenames): @@ -447,6 +473,56 @@ class ScopedMetaGraphTest(test.TestCase): del b.collection_def["unbound_inputs"] test_util.assert_meta_graph_protos_equal(self, a, b) + def testWhileLoopGradients(self): + # Create a simple while loop. + with ops.Graph().as_default(): + with ops.name_scope("export"): + var = variables.Variable(0) + var_name = var.name + _, output = control_flow_ops.while_loop(lambda i, x: i < 5, + lambda i, x: (i + 1, x + i), + [0, var]) + output_name = output.name + + # Generate a MetaGraphDef containing the while loop with an export scope. + meta_graph_def, _ = meta_graph.export_scoped_meta_graph( + export_scope="export") + + # Build and run the gradients of the while loop. We use this below to + # verify that the gradients are correct with the imported MetaGraphDef. + init_op = variables.global_variables_initializer() + grad = gradients_impl.gradients([output], [var]) + with session.Session() as sess: + sess.run(init_op) + expected_grad_value = sess.run(grad) + + # Restore the MetaGraphDef into a new Graph with an import scope. + with ops.Graph().as_default(): + meta_graph.import_scoped_meta_graph(meta_graph_def, import_scope="import") + + # Re-export and make sure we get the same MetaGraphDef. + new_meta_graph_def, _ = meta_graph.export_scoped_meta_graph( + export_scope="import") + test_util.assert_meta_graph_protos_equal( + self, meta_graph_def, new_meta_graph_def) + + # Make sure we can still build gradients and get the same result. + + def new_name(tensor_name): + base_tensor_name = tensor_name.replace("export/", "") + return "import/" + base_tensor_name + + var = ops.get_default_graph().get_tensor_by_name(new_name(var_name)) + output = ops.get_default_graph().get_tensor_by_name(new_name(output_name)) + grad = gradients_impl.gradients([output], [var]) + + init_op = variables.global_variables_initializer() + + with session.Session() as sess: + sess.run(init_op) + actual_grad_value = sess.run(grad) + self.assertEqual(expected_grad_value, actual_grad_value) + def testScopedImportUnderNameScope(self): graph = ops.Graph() with graph.as_default(): @@ -830,22 +906,12 @@ class ExportImportAcrossScopesTest(test.TestCase): with variable_scope.variable_scope("importA/keepA"): graph_fn(use_resource=use_resource) - if use_resource: - # Bringing in a collection that contains ResourceVariables adds ops - # to the graph, so mimic the same behavior. - for collection_key in sorted([ - ops.GraphKeys.GLOBAL_VARIABLES, - ops.GraphKeys.TRAINABLE_VARIABLES, - ]): - for var in expected_graph.get_collection(collection_key): - var._read_variable_op() - result = meta_graph.export_scoped_meta_graph(graph=imported_graph)[0] expected = meta_graph.export_scoped_meta_graph(graph=expected_graph)[0] if use_resource: # Clear all shared_name attributes before comparing, since they are - # supposed to be orthogonal to scopes. + # orthogonal to scopes and are not updated on export/import. for meta_graph_def in [result, expected]: for node in meta_graph_def.graph_def.node: shared_name_attr = "shared_name" diff --git a/tensorflow/python/framework/op_def_library_test.py b/tensorflow/python/framework/op_def_library_test.py index 817007ce6c18e11d19038e09d77a8f27bd7eca91..84ca062ade3b32c37212ba2d5b7eb9c64fb1dfa5 100644 --- a/tensorflow/python/framework/op_def_library_test.py +++ b/tensorflow/python/framework/op_def_library_test.py @@ -42,7 +42,7 @@ class OpDefLibraryTest(test_util.TensorFlowTestCase): def setUp(self): self._lib = test_ops._op_def_lib - def _add_op(self, ascii): + def _add_op(self, ascii): # pylint: disable=redefined-builtin op_def = op_def_pb2.OpDef() text_format.Merge(ascii, op_def) self._lib.add_op(op_def) @@ -1336,7 +1336,7 @@ class OpDefLibraryGraphTest(test_util.TensorFlowTestCase): def setUp(self): self._lib = test_ops._op_def_lib - def _add_op(self, ascii): + def _add_op(self, ascii): # pylint: disable=redefined-builtin op_def = op_def_pb2.OpDef() text_format.Merge(ascii, op_def) self._lib.add_op(op_def) diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py index d5786cac68dc31210f45f0af9ff6c347d93c026f..01a0e03be20648fd65b0d4f545c3c53114d99a71 100644 --- a/tensorflow/python/framework/ops.py +++ b/tensorflow/python/framework/ops.py @@ -63,6 +63,7 @@ from tensorflow.python.util.tf_export import tf_export # in code or via the environment variable. This will be removed once all # functionality is supported and there's no performance penalty with it enabled. _USE_C_API = os.getenv("TF_C_API_GRAPH_CONSTRUCTION", "0") is not "0" +_USE_C_SHAPES = os.getenv("TF_C_API_GRAPH_CONSTRUCTION_SHAPES", "0") is not "0" def tensor_id(tensor): @@ -368,8 +369,8 @@ class Tensor(_TensorLike): A `TensorShape` representing the shape of this tensor. """ - if _USE_C_API: - graph = self._op._graph._c_graph # pylint: disable=protected-access + graph = self._op._graph._c_graph # pylint: disable=protected-access + if graph and _USE_C_SHAPES: with errors.raise_exception_on_not_ok_status() as status: num_dims = c_api.TF_GraphGetTensorNumDims(graph, self._as_tf_output(), status) @@ -395,10 +396,10 @@ class Tensor(_TensorLike): "Tensor._shape cannot be assigned, use Tensor.set_shape instead.") def __iter__(self): - if context.in_graph_mode(): + if not context.executing_eagerly(): raise TypeError( - "`Tensor` objects are not iterable when eager execution is not " - "enabled. To iterate over this tensor use `tf.map_fn`.") + "Tensor objects are not iterable when eager execution is not " + "enabled. To iterate over this tensor use tf.map_fn.") shape = self._shape_tuple() if shape is None: raise TypeError("Cannot iterate over a tensor with unknown shape.") @@ -466,9 +467,13 @@ class Tensor(_TensorLike): ValueError: If `shape` is not compatible with the current shape of this tensor. """ - if not _USE_C_API: + if not _USE_C_SHAPES: # pylint: disable=protected-access self._shape_val = self._shape_val.merge_with(shape) - return + + if not self._op._graph._c_graph: return + + # Update C shape even if _USE_C_SHAPES = False, since we still want + # set_shape to be reflected in the C API graph for when we run it. if not isinstance(shape, tensor_shape.TensorShape): shape = tensor_shape.TensorShape(shape) dim_list = [] @@ -772,7 +777,7 @@ class _EagerTensorBase(Tensor): six.raise_from(core._status_to_exception(e.code, e.message), None) # Record the copy on tape and define backprop copy as well. - if not context.in_graph_mode(): + if context.executing_eagerly(): self_device = self.device def grad_fun(dresult): return [dresult._copy(device_name=self_device)] @@ -782,7 +787,11 @@ class _EagerTensorBase(Tensor): @property def shape(self): - return tensor_shape.TensorShape(self._shape_tuple()) + if self._tensor_shape is None: # pylint: disable=access-member-before-definition + # `_tensor_shape` is declared and defined in the definition of + # `EagerTensor`, in C. + self._tensor_shape = tensor_shape.TensorShape(self._shape_tuple()) + return self._tensor_shape def get_shape(self): """Alias of Tensor.shape.""" @@ -989,7 +998,7 @@ def internal_convert_to_tensor(value, """ if ctx is None: ctx = context.context() - if ctx.in_eager_mode(): + if ctx.executing_eagerly(): # Fast path for EagerTensors that don't need any conversion. if isinstance(value, EagerTensor): # Note that we don't check that value's dtype matches the dtype @@ -1618,7 +1627,7 @@ class Operation(object): for i, x in zip(inputs, input_types)): raise TypeError("In op '%s', input types (%s) are not compatible " "with expected types (%s)" % - (self.node_def.name, [i.dtype for i in inputs], + (node_def.name, [i.dtype for i in inputs], input_types)) # Build the list of control inputs. @@ -1657,7 +1666,7 @@ class Operation(object): self._c_op = c_op elif self._graph._c_graph: # pylint: disable=protected-access if op_def is None: - op_def = self._graph._registered_ops[node_def.op] + op_def = self._graph._get_op_def(node_def.op) # TODO(skyewm): op_def_library.apply_op() flattens the incoming inputs. # Refactor so we don't have to do this here. grouped_inputs = self._reconstruct_sequence_inputs( @@ -2164,16 +2173,7 @@ class Operation(object): """ # pylint: enable=line-too-long if self._c_op: - with c_api_util.tf_buffer() as buf: - with errors.raise_exception_on_not_ok_status() as status: - # pylint: disable=protected-access - c_api.TF_GraphGetOpDef(self._graph._c_graph, - compat.as_bytes(self.type), buf, status) - # pylint: enable=protected-access - data = c_api.TF_GetBuffer(buf) - op_def = op_def_pb2.OpDef() - op_def.ParseFromString(compat.as_bytes(data)) - return op_def + return self._graph._get_op_def(self.type) else: return self._op_def_val @@ -2495,7 +2495,7 @@ def _set_shapes_for_outputs(op): def set_shapes_for_outputs(op): """Set the shapes for op's outputs.""" - if op._c_op: # pylint: disable=protected-access + if op._c_op and _USE_C_SHAPES: # pylint: disable=protected-access return _set_shapes_for_outputs_c_api(op) else: return _set_shapes_for_outputs(op) @@ -2699,15 +2699,20 @@ class Graph(object): def __init__(self): """Creates a new, empty Graph.""" - # Protects the core state that may be accessed by multiple readers. - # Only state that can be returned via public accessors (`as_graph_def()`, - # `get_operations()`, `as_graph_element()`, `get_collection()`, and - # `get_collection_ref()`) is by the lock. Thread-safety is provided on a - # best-effort basis to support buggy programs, and is not guaranteed by the - # public `tf.Graph` API. + # Protects core state that can be returned via public accessors, as well as + # synchronizes Session.run calls with methods that create and mutate ops + # (e.g. Graph.create_op()). This synchronization is necessary because it's + # illegal to modify an operation after it's been run. Thread-safety is + # provided on a best-effort basis to support buggy programs, and is not + # guaranteed by the public `tf.Graph` API. + # + # The lock must be reentrant because create_op can be called recursively due + # to control flow. Without a reentrant lock, many methods would also need a + # "locked" version or parameter (including generated code). + # # NOTE(mrry): This does not protect the various stacks. A warning will # be reported if these are used from multiple threads - self._lock = threading.Lock() + self._lock = threading.RLock() self._nodes_by_id = dict() # GUARDED_BY(self._lock) self._next_id_counter = 0 # GUARDED_BY(self._lock) self._nodes_by_name = dict() # GUARDED_BY(self._lock) @@ -2716,15 +2721,21 @@ class Graph(object): self._name_stack = "" # Maps a name used in the graph to the next id to use for that name. self._names_in_use = {} + self._stack_state_is_thread_local = False + self._thread_local = threading.local() # Functions that will be applied to choose a device if none is specified. - self._device_function_stack = [] + # After switch_to_thread_local(), self._thread_local._device_function_stack + # is used instead. + self._graph_device_function_stack = [] # Default original_op applied to new ops. self._default_original_op = None # Current control flow context. It could be either CondContext or # WhileContext defined in ops/control_flow_ops.py self._control_flow_context = None # A new node will depend of the union of all of the nodes in the stack. - self._control_dependencies_stack = [] + # After switch_to_thread_local(), + # self._thread_local._control_dependencies_stack is used instead. + self._graph_control_dependencies_stack = [] # Arbitrary collections of objects. self._collections = {} # The graph-level random seed @@ -2746,8 +2757,9 @@ class Graph(object): producer=versions.GRAPH_DEF_VERSION, min_consumer=versions.GRAPH_DEF_VERSION_MIN_CONSUMER) self._building_function = False - # Stack of colocate_with ops - self._colocation_stack = [] + # Stack of colocate_with ops. After switch_to_thread_local(), + # self._thread_local._colocation_stack is used instead. + self._graph_colocation_stack = [] # Set of tensors that are dangerous to feed! self._unfeedable_tensors = set() # Set of operations that are dangerous to fetch! @@ -2760,29 +2772,29 @@ class Graph(object): self._handle_movers = {} # A map from tensor handle to its delete op. self._handle_deleters = {} - # Resource container. - if context.in_graph_mode(): - self._container_prefix = "" - else: - # In Eager mode, isolate resources (particularly ResourceVariables) in - # Graphs by default. This prevents unintended variable sharing. Graph mode - # gets this kind of isolation from Sessions. - self._container_prefix = "eager-execution-%d/" % (uid(),) - self._container = self._container_prefix + # Allow optimizers and other objects to pseudo-uniquely key graphs (this key + # will be shared when defining function graphs, for example, so optimizers + # being called inside function definitions behave as if they were seeing the + # actual outside graph). + self._graph_key = "grap-key-%d/" % (uid(),) + self._container = "" self._registered_ops = op_def_registry.get_registered_ops() # TODO(skyewm): fold as much of the above as possible into the C # implementation - if _USE_C_API or self._use_c_api_hack(): + if self._use_c_api_hack(): self._scoped_c_graph = c_api_util.ScopedTFGraph() + # The C API requires all ops to have shape functions. Disable this + # requirement (many custom ops do not have shape functions, and we don't + # want to break these existing cases). + c_api.SetRequireShapeInferenceFns(self._c_graph, False) else: self._scoped_c_graph = None - self._variable_creator_stack = [] # TODO(apassos) remove once the C API is used by default. def _use_c_api_hack(self): """Temporary hack; can be overridden to force C API usage.""" - return False + return _USE_C_API def _convert_stack(self, stack, include_func_start_lineno=False): """Converts a stack extracted using _extract_stack() to a traceback stack. @@ -2818,17 +2830,26 @@ class Graph(object): # frozen, and this functionality is still not ready for public visibility. @tf_contextlib.contextmanager def _variable_creator_scope(self, creator): + # This step makes a copy of the existing stack, and it also initializes + # self._thread_local._variable_creator_stack if it doesn't exist yet. old = list(self._variable_creator_stack) - self._variable_creator_stack.append(creator) + self._thread_local._variable_creator_stack.append(creator) try: yield finally: - self._variable_creator_stack = old + self._thread_local._variable_creator_stack = old # Note: this method is private because the API of tf.Graph() is public and # frozen, and this functionality is still not ready for public visibility. - def _get_variable_creator_stack(self): - return list(self._variable_creator_stack) + @property + def _variable_creator_stack(self): + if not hasattr(self._thread_local, "_variable_creator_stack"): + self._thread_local._variable_creator_stack = [] + return list(self._thread_local._variable_creator_stack) + + @_variable_creator_stack.setter + def _variable_creator_stack(self, variable_creator_stack): + self._thread_local._variable_creator_stack = variable_creator_stack def _extract_stack(self): """A lightweight, extensible re-implementation of traceback.extract_stack. @@ -3042,7 +3063,7 @@ class Graph(object): """ # pylint: enable=line-too-long - if _USE_C_API: + if self._c_graph: with self._lock: with c_api_util.tf_buffer() as buf: with errors.raise_exception_on_not_ok_status() as status: @@ -3260,17 +3281,20 @@ class Graph(object): input_ops = set([t.op for t in inputs]) control_inputs = self._control_dependencies_for_inputs(input_ops) - ret = Operation( - node_def, - self, - inputs=inputs, - output_types=dtypes, - control_inputs=control_inputs, - input_types=input_types, - original_op=self._default_original_op, - op_def=op_def) - self._create_op_helper(ret, compute_shapes=compute_shapes, - compute_device=compute_device) + # _create_op_helper mutates the new Operation. _lock ensures a Session.run + # call cannot occur between creating and mutating the op. + with self._lock: + ret = Operation( + node_def, + self, + inputs=inputs, + output_types=dtypes, + control_inputs=control_inputs, + input_types=input_types, + original_op=self._default_original_op, + op_def=op_def) + self._create_op_helper(ret, compute_shapes=compute_shapes, + compute_device=compute_device) return ret def _create_op_from_tf_operation(self, c_op, compute_device=True): @@ -3309,7 +3333,7 @@ class Graph(object): # always computes shape information (even for function calls, which the # original Python shape inference code doesn't handle). Deprecate the # compute_shapes argument. - if op._c_op or compute_shapes: # pylint: disable=protected-access + if (op._c_op and _USE_C_SHAPES) or compute_shapes: # pylint: disable=protected-access set_shapes_for_outputs(op) # TODO(b/XXXX): move to Operation.__init__ once _USE_C_API flag is removed. self._add_op(op) @@ -3362,9 +3386,9 @@ class Graph(object): if (op.device and pydev.canonical_name(op.device) != pydev.canonical_name(colocation_op.device)): logging.warning("Tried to colocate %s with an op %s that had " - "a different device: %s vs %s. " - "Ignoring colocation property.", op.name, - colocation_op.name, op.device, + "a different device: %s vs %s. Postponing " + "error-checking until all devices are assigned.", + op.name, colocation_op.name, op.device, colocation_op.device) else: op._set_device(colocation_op.device) # pylint: disable=protected-access @@ -3380,8 +3404,8 @@ class Graph(object): # (2) "is_stateful" is set in OpDef # (3) "container" attribute is in OpDef # (4) "container" attribute is None - if (self._container and op.type in self._registered_ops and - self._registered_ops[op.type].is_stateful): + # TODO(skyewm): remove op.op_def check when _USE_C_API is removed. + if self._container and op.op_def and op.op_def.is_stateful: try: container_attr = op.get_attr("container") except ValueError: @@ -3661,6 +3685,22 @@ class Graph(object): def _last_id(self): return self._next_id_counter + def _get_op_def(self, type): # pylint: disable=redefined-builtin + """Returns the `OpDef` proto for `type`. `type` is a string.""" + if self._c_graph: + with c_api_util.tf_buffer() as buf: + with errors.raise_exception_on_not_ok_status() as status: + # pylint: disable=protected-access + c_api.TF_GraphGetOpDef(self._c_graph, + compat.as_bytes(type), buf, status) + # pylint: enable=protected-access + data = c_api.TF_GetBuffer(buf) + op_def = op_def_pb2.OpDef() + op_def.ParseFromString(compat.as_bytes(data)) + return op_def + else: + return self._registered_ops[type] + def as_default(self): """Returns a context manager that makes this `Graph` the default graph. @@ -4229,7 +4269,7 @@ class Graph(object): """ original_container = self._container try: - self._container = self._container_prefix + container_name + self._container = container_name yield self._container finally: self._container = original_container @@ -4665,6 +4705,79 @@ class Graph(object): else: return tensor_or_op not in self._unfetchable_ops + def switch_to_thread_local(self): + """Make device, colocation and dependencies stacks thread-local. + + Device, colocation and dependencies stacks are not thread-local be default. + If multiple threads access them, then the state is shared. This means that + one thread may affect the behavior of another thread. + + After this method is called, the stacks become thread-local. If multiple + threads access them, then the state is not shared. Each thread uses its own + value; a thread doesn't affect other threads by mutating such a stack. + + The initial value for every thread's stack is set to the current value + of the stack when `switch_to_thread_local()` was first called. + """ + if not self._stack_state_is_thread_local: + self._stack_state_is_thread_local = True + + @property + def _device_function_stack(self): + if self._stack_state_is_thread_local: + # This may be called from a thread where device_function_stack doesn't yet + # exist. + if not hasattr(self._thread_local, "_device_function_stack"): + self._thread_local._device_function_stack = ( + self._graph_device_function_stack[:]) + return self._thread_local._device_function_stack + else: + return self._graph_device_function_stack + + @_device_function_stack.setter + def _device_function_stack(self, device_function_stack): + if self._stack_state_is_thread_local: + self._thread_local._device_function_stack = device_function_stack + else: + self._graph_device_function_stack = device_function_stack + + @property + def _colocation_stack(self): + if self._stack_state_is_thread_local: + # This may be called from a thread where colocation_stack doesn't yet + # exist. + if not hasattr(self._thread_local, "_colocation_stack"): + self._thread_local._colocation_stack = self._graph_colocation_stack[:] + return self._thread_local._colocation_stack + else: + return self._graph_colocation_stack + + @_colocation_stack.setter + def _colocation_stack(self, colocation_stack): + if self._stack_state_is_thread_local: + self._thread_local._colocation_stack = colocation_stack + else: + self._graph_colocation_stack = colocation_stack + + @property + def _control_dependencies_stack(self): + if self._stack_state_is_thread_local: + # This may be called from a thread where control_dependencies_stack + # doesn't yet exist. + if not hasattr(self._thread_local, "_control_dependencies_stack"): + self._thread_local._control_dependencies_stack = ( + self._graph_control_dependencies_stack[:]) + return self._thread_local._control_dependencies_stack + else: + return self._graph_control_dependencies_stack + + @_control_dependencies_stack.setter + def _control_dependencies_stack(self, control_dependencies): + if self._stack_state_is_thread_local: + self._thread_local._control_dependencies_stack = control_dependencies + else: + self._graph_control_dependencies_stack = control_dependencies + # TODO(agarwal): currently device directives in an outer eager scope will not # apply to inner graph mode code. Fix that. @@ -4689,15 +4802,15 @@ def device(device_name_or_function): Raises: RuntimeError: If eager execution is enabled and a function is passed in. """ - if context.in_graph_mode(): - return get_default_graph().device(device_name_or_function) - else: + if context.executing_eagerly(): # TODO(agarwal): support device functions in EAGER mode. if callable(device_name_or_function): raise RuntimeError( "tf.device does not support functions when eager execution " "is enabled.") return context.device(device_name_or_function) + else: + return get_default_graph().device(device_name_or_function) @tf_export("container") @@ -4716,13 +4829,20 @@ def container(container_name): @tf_export("colocate_with") def colocate_with(op, ignore_existing=False): - if context.in_graph_mode(): - return get_default_graph().colocate_with(op, ignore_existing) - else: + if context.executing_eagerly(): if op is not None: return device(op.device) else: return _NullContextmanager() + else: + default_graph = get_default_graph() + if isinstance(op, EagerTensor): + if default_graph.building_function: + op = internal_convert_to_tensor(op) + else: + raise ValueError("Encountered an Eager-defined Tensor during graph " + "construction, but a function was not being built.") + return default_graph.colocate_with(op, ignore_existing) @tf_export("control_dependencies") @@ -4732,20 +4852,29 @@ def control_dependencies(control_inputs): See @{tf.Graph.control_dependencies} for more details. + When eager execution is enabled, any callable object in the `control_inputs` + list will be called. + Args: control_inputs: A list of `Operation` or `Tensor` objects which must be executed or computed before running the operations defined in the context. Can also be `None` to clear the control - dependencies. + dependencies. If eager execution is enabled, any callable object in the + `control_inputs` list will be called. Returns: A context manager that specifies control dependencies for all operations constructed within the context. """ - if context.in_graph_mode(): - return get_default_graph().control_dependencies(control_inputs) - else: + if context.executing_eagerly(): + if control_inputs: + # Excute any pending callables. + for control in control_inputs: + if callable(control): + control() return _NullContextmanager() + else: + return get_default_graph().control_dependencies(control_inputs) class _DefaultStack(threading.local): @@ -4966,11 +5095,12 @@ class _DefaultGraphStack(_DefaultStack): # pylint: disable=protected-access @tf_contextlib.contextmanager def get_controller(self, default): try: - context.context_stack.push(default.building_function, default.as_default) + context.context().context_switches.push(default.building_function, + default.as_default) with super(_DefaultGraphStack, self).get_controller(default) as g: yield g finally: - context.context_stack.pop() + context.context().context_switches.pop() _default_graph_stack = _DefaultGraphStack() @@ -4996,87 +5126,130 @@ def init_scope(): graph function. Here, a context is defined as either a graph or an eager context. Every context switch, i.e., every installation of a graph as the default graph and every switch into eager mode, is logged in a - thread-local stack called the `context_stack`; the log entry for a + thread-local stack called `context_switches`; the log entry for a context switch is popped from the stack when the context is exited. - Entering an `init_scope` is equivalent to crawling up the - `context_stack`, finding the first context that is not building a graph - function, and entering it. A caveat is that if graph mode is enabled - but the default graph stack is empty, then entering an `init_scope` - will simply install a fresh graph as the default one. + Entering an `init_scope` is equivalent to crawling up + `context_switches`, finding the first context that is not building a + graph function, and entering it. A caveat is that if graph mode is + enabled but the default graph stack is empty, then entering an + `init_scope` will simply install a fresh graph as the default one. (3) The gradient tape is paused while the scope is active. """ # pylint: enable=g-doc-return-or-yield,line-too-long - in_graph_mode = context.in_graph_mode() - # Retrieve the active name scope: entering an `init_scope` preserves - # the name scope of the current context. - if in_graph_mode: + if context.executing_eagerly(): + # Fastpath. + with tape.stop_recording(): + yield + else: + # Retrieve the active name scope: entering an `init_scope` preserves + # the name scope of the current context. default_graph = get_default_graph() scope = default_graph.get_name_scope() - else: - scope = context.context().scope_name - if scope and scope[-1] != '/': - # Names that end with trailing slashes are treated by `name_scope` as - # absolute. - scope = scope + '/' - - outer_context = None - if in_graph_mode and not _default_graph_stack.stack: - outer_context = default_graph.as_default - else: - for stack_entry in reversed(context.context_stack.stack): - if not stack_entry.is_building_function: - outer_context = stack_entry.enter_context_fn - break + if scope and scope[-1] != '/': + # Names that end with trailing slashes are treated by `name_scope` as + # absolute. + scope = scope + '/' + + outer_context = None + if not _default_graph_stack.stack: + # If the default graph stack is empty, then we cannot be building a + # function. Install the global graph (which, in this case, is also the + # default graph) as the outer context. + if default_graph.building_function: + raise RuntimeError("The global graph is building a function.") + outer_context = default_graph.as_default + else: + # Find a context that is not building a function. + for stack_entry in reversed(context.context().context_switches.stack): + if not stack_entry.is_building_function: + outer_context = stack_entry.enter_context_fn + break + + if outer_context is None: + # As a last resort, obtain the global default graph; this graph doesn't + # necessarily live on the graph stack (and hence it doesn't necessarily + # live on the context stack), but it is stored in the graph stack's + # encapsulating object. + outer_context = _default_graph_stack._GetGlobalDefaultGraph().as_default # pylint: disable=protected-access - if outer_context is None: - raise AssertionError("All graphs are building functions, and no " + if outer_context is None: + # Sanity check; this shouldn't be triggered. + raise RuntimeError("All graphs are building functions, and no " "eager context was previously active.") - try: with outer_context(), name_scope(scope), control_dependencies( None), tape.stop_recording(): yield - finally: - pass -def enable_eager_execution(config=None, device_policy=None): - """Enables, for the rest of the lifetime of this program, eager execution. +@tf_export("enable_eager_execution") +def enable_eager_execution(config=None, device_policy=None, + execution_mode=None): + """Enables eager execution for the lifetime of this program. - If not called immediately on startup risks creating breakage and bugs. + Eager execution provides an imperative interface to TensorFlow. With eager + execution enabled, TensorFlow functions execute operations immediately (as + opposed to adding to a graph to be executed later in a @{tf.Session}) and + return concrete values (as opposed to symbolic references to a node in a + computational graph). - Example: + For example: ```python - tfe.enable_eager_execution() + tf.enable_eager_execution() # After eager execution is enabled, operations are executed as they are - # defined and `Tensor`s hold concrete values, which can be accessed as - # `numpy.ndarray`s through the `numpy()` method. + # defined and Tensor objects hold concrete values, which can be accessed as + # numpy.ndarray`s through the numpy() method. assert tf.multiply(6, 7).numpy() == 42 ``` + Eager execution cannot be enabled after TensorFlow APIs have been used to + create or execute graphs. It is typically recommended to invoke this function + at program startup and not in a library (as most libraries should be usable + both with and without eager execution). + Args: - config: (Optional.) A `ConfigProto` protocol buffer with configuration - options for the Context. Note that a lot of these options may be - currently unimplemented or irrelevant when eager execution is enabled. - device_policy: (Optional.) What policy to use when trying to run an - operation on a device with inputs which are not on that device. + config: (Optional.) A @{tf.ConfigProto} to use to configure the environment + in which operations are executed. Note that @{tf.ConfigProto} is also + used to configure graph execution (via @{tf.Session}) and many options + within `tf.ConfigProto` are not implemented (or are irrelevant) when + eager execution is enabled. + device_policy: (Optional.) Policy controlling how operations requiring + inputs on a specific device (e.g., a GPU 0) handle inputs on a different + device (e.g. GPU 1 or CPU). When set to None, an appropriate value will be + picked automatically. The value picked may change between TensorFlow + releases. Valid values: - tfe.DEVICE_PLACEMENT_EXPLICIT: raises an error if the placement is not - correct. - tfe.DEVICE_PLACEMENT_WARN: copies the tensors which are not on the - right device but raises a warning. - tfe.DEVICE_PLACEMENT_SILENT: silently copies the tensors. This might - hide performance problems. - tfe.DEVICE_PLACEMENT_SILENT_FOR_INT32: silently copies int32 tensors, - raising errors on the other ones. + + - tf.contrib.eager.DEVICE_PLACEMENT_EXPLICIT: raises an error if the + placement is not correct. + + - tf.contrib.eager.DEVICE_PLACEMENT_WARN: copies the tensors which are not + on the right device but logs a warning. + + - tf.contrib.eager.DEVICE_PLACEMENT_SILENT: silently copies the tensors. + Note that this may hide performance problems as there is no notification + provided when operations are blocked on the tensor being copied between + devices. + + - tf.contrib.eager.DEVICE_PLACEMENT_SILENT_FOR_INT32: silently copies + int32 tensors, raising errors on the other ones. + execution_mode: (Optional.) Policy controlling how operations dispatched are + actually executed. When set to None, an appropriate value will be picked + automatically. The value picked may change between TensorFlow releases. + Valid values: + + - tf.contrib.eager.SYNC: executes each operation synchronously. + + - tf.contrib.eager.ASYNC: executes each operation asynchronously. These + operations may return "non-ready" handles. Raises: - ValueError: If trying to create a context after using graph operations - or if trying to create a context with nontrivial options which differ - from those of the existing context. + ValueError: If eager execution is enabled after creating/executing a + TensorFlow graph, or if options provided conflict with a previous call + to this function. """ if config is not None and not isinstance(config, config_pb2.ConfigProto): raise TypeError( @@ -5086,8 +5259,12 @@ def enable_eager_execution(config=None, device_policy=None): context.DEVICE_PLACEMENT_SILENT, context.DEVICE_PLACEMENT_SILENT_FOR_INT32): raise ValueError( - "device_policy must be one of None, tfe.DEVICE_PLACEMENT_*" + "device_policy must be one of None, tf.contrib.eager.DEVICE_PLACEMENT_*" ) + if execution_mode not in (None, context.SYNC, context.ASYNC): + raise ValueError( + "execution_mode must be one of None, tf.contrib.eager.SYNC, " + "tf.contrib.eager.ASYNC") # pylint: disable=protected-access if context._default_mode == context.GRAPH_MODE: graph_mode_has_been_used = ( @@ -5095,30 +5272,29 @@ def enable_eager_execution(config=None, device_policy=None): _default_graph_stack._global_default_graph is not None) if graph_mode_has_been_used: raise ValueError( - "tfe.enable_eager_execution has to be called at program startup.") + "tf.enable_eager_execution must be called at program startup.") context._default_mode = context.EAGER_MODE if context._context is None: - context._context = context.Context(config=config, - device_policy=device_policy) - if context.context_stack.stack: - raise AssertionError("Invariant violated: The context stack must " - "be empty when eager execution is enabled.") - # Log that eager execution has been enabled by pushing an entry onto the - # context stack; this entry won't ever be popped, as it's impossible to - # disable eager execution - context.context_stack.push(False, context.eager_mode) - elif ((config is not None and config is not context._context._config) - or (device_policy is not None - and device_policy is not context._context._device_policy)): + context._context = context.Context( + config=config, + device_policy=device_policy, + execution_mode=execution_mode) + elif ((config is not None and config is not context._context._config) or + (device_policy is not None and + device_policy is not context._context._device_policy) or + (execution_mode is not None and + execution_mode is not context._context._execution_mode)): raise ValueError("Trying to change the options of an active eager" " execution. Context config: %s, specified config:" - " %s. Context device policy: %s; specified device" - " policy: %s." % (config, context._context._config, - device_policy, - context._context._device_policy)) + " %s. Context device policy: %s, specified device" + " policy: %s. Context execution mode: %s, " + " specified execution mode %s." % + (context._context._config, config, + context._context._device_policy, device_policy, + context._context._execution_mode, execution_mode)) else: raise ValueError( - "tfe.enable_eager_execution has to be called at program startup.") + "tf.enable_eager_execution must be called at program startup.") def eager_run(main=None, argv=None): @@ -5202,6 +5378,8 @@ def get_name_scope(): Returns: A string representing the current name scope. """ + if context.in_eager_mode(): + return context.context().scope_name.rstrip("/") return get_default_graph().get_name_scope() @@ -5456,7 +5634,7 @@ def add_to_collection(name, value): """ get_default_graph().add_to_collection(name, value) - +@tf_export("add_to_collections") def add_to_collections(names, value): """Wrapper for `Graph.add_to_collections()` using the default graph. @@ -5533,6 +5711,9 @@ def get_all_collection_keys(): return get_default_graph().get_all_collection_keys() +name_scope_cache = {} + + # Named like a function for backwards compatibility with the # @tf_contextlib.contextmanager version, which was switched to a class to avoid # some object creation overhead. @@ -5575,7 +5756,7 @@ class name_scope(object): # pylint: disable=invalid-name self._default_name = default_name self._values = values self._ctx = context.context() - self._in_eager_mode = self._ctx.in_eager_mode() + self._in_eager_mode = self._ctx.executing_eagerly() def __enter__(self): """Start the scope block. @@ -5592,7 +5773,11 @@ class name_scope(object): # pylint: disable=invalid-name if not self._name: scope_name = "" else: - if self._name[-1] == "/": + cache_key = self._name, self._old_name, self._default_name + if cache_key in name_scope_cache: + self._ctx.scope_name = name_scope_cache[cache_key] + return self._ctx.scope_name + elif self._name[-1] == "/": # A trailing slash breaks out of nested name scopes, indicating a # fully specified scope name, for compatibility with Graph.name_scope. scope_name = self._name @@ -5601,6 +5786,7 @@ class name_scope(object): # pylint: disable=invalid-name scope_name = ( self._old_name + name_with_trailing_slash if self._old_name else name_with_trailing_slash) + name_scope_cache[cache_key] = scope_name self._ctx.scope_name = scope_name return scope_name else: @@ -5749,7 +5935,7 @@ def get_from_proto_function(collection_name): def _assert_collection_is_ok(collection_name): - if context.in_eager_mode(): + if context.executing_eagerly(): if collection_name in GraphKeys._VARIABLE_COLLECTIONS: # pylint: disable=protected-access raise ValueError("When Eager Execution is enabled, variable " "collections are not supported.") diff --git a/tensorflow/python/framework/ops_test.py b/tensorflow/python/framework/ops_test.py index c5e177d521013417d28fb01a975208ba5e4cf90b..d96e0708f8e5152b1a14e7bd5d6ca19784860bd4 100644 --- a/tensorflow/python/framework/ops_test.py +++ b/tensorflow/python/framework/ops_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function import gc +import threading import weakref from tensorflow.core.framework import attr_value_pb2 @@ -762,6 +763,7 @@ class CreateOpFromTFOperationTest(test_util.TensorFlowTestCase): self.assertEqual(g.get_operation_by_name("myop"), op) self.assertEqual(g.get_tensor_by_name("myop:0"), op.outputs[0]) + @test_util.enable_c_shapes def testShape(self): g = ops.Graph() with g.as_default(): @@ -916,7 +918,6 @@ class CreateOpFromTFOperationTest(test_util.TensorFlowTestCase): op = g.get_operation_by_name("myloop/myop") self.assertIsNotNone(op) - self.assertEqual(len(op.control_inputs), 1) # External control dep is removed and replaced with internal control dep self.assertNotEqual(op.control_inputs[0], c.op) self.assertIsNotNone(op.control_inputs[0]._get_control_flow_context()) @@ -1382,6 +1383,180 @@ class DeviceTest(test_util.TensorFlowTestCase): """, gd) +@test_util.with_c_api +class MultithreadedGraphStateTest(test_util.TensorFlowTestCase): + + class TestThread(threading.Thread): + + def __init__(self, graph, replica_id): + super(MultithreadedGraphStateTest.TestThread, self).__init__() + self._graph = graph + self._replica_id = replica_id + # This thread sets this event when it mutated the graph. The caller can + # wait for that. + self.has_mutated_graph = threading.Event() + # This thread waits for when it should continue. The caller can set this + # event. + self.should_continue = threading.Event() + + def run(self): + # Mutate a graph's stack, then set `has_mutated_graph`, then wait for + # `should_continue`, then add an op to the graph affected by the graph's + # stack. + raise NotImplementedError("must be implemented in descendants") + + def testDeviceFunctionStack(self): + + class DeviceSettingThread(self.TestThread): + + def run(self): + with g.device("/job:worker/replica:{}".format(self._replica_id)): + self.has_mutated_graph.set() + self.should_continue.wait() + self.should_continue.clear() + g.create_op( + "FloatOutput", [], [dtypes.float32], + name="FloatOutput_{}".format(self._replica_id)) + + g = ops.Graph() + # If `switch_to_thread` isn't called, then device placement of the ops + # below is not deterministic. + g.switch_to_thread_local() + threads = [DeviceSettingThread(g, i) for i in range(3)] + for t in threads: + t.start() + t.has_mutated_graph.wait() + t.has_mutated_graph.clear() + for t in threads: + t.should_continue.set() + t.join() + + gd = g.as_graph_def() + self.assertProtoEqualsVersion(""" + node { name: "FloatOutput_0" op: "FloatOutput" + device: "/job:worker/replica:0" } + node { name: "FloatOutput_1" op: "FloatOutput" + device: "/job:worker/replica:1" } + node { name: "FloatOutput_2" op: "FloatOutput" + device: "/job:worker/replica:2" } + """, gd) + + def testColocateWith(self): + + class ColocatingThread(self.TestThread): + + def __init__(self, graph, replica_id, op_to_colocate_with): + super(ColocatingThread, self).__init__(graph, replica_id) + self._op_to_colocate_with = op_to_colocate_with + + def run(self): + with g.colocate_with(self._op_to_colocate_with): + self.has_mutated_graph.set() + self.should_continue.wait() + self.should_continue.clear() + g.create_op( + "FloatOutput", [], [dtypes.float32], + name="FloatOutput_{}".format(self._replica_id)) + + g = ops.Graph() + ops_to_colocate_with = [] + for i in range(3): + with g.device("/job:worker/replica:{}".format(i)): + ops_to_colocate_with.append( + g.create_op( + "FloatOutput", [], [dtypes.float32], + name="ColocateWithMe_{}".format(i))) + + # If `switch_to_thread` isn't called, then `device` and `attr` values for + # the ops below are not deterministic. + g.switch_to_thread_local() + threads = [ + ColocatingThread(g, i, ops_to_colocate_with[i]) for i in range(3) + ] + for t in threads: + t.start() + t.has_mutated_graph.wait() + t.has_mutated_graph.clear() + for t in threads: + t.should_continue.set() + t.join() + + gd = g.as_graph_def() + self.assertProtoEqualsVersion(""" + node { name: "ColocateWithMe_0" op: "FloatOutput" + device: "/job:worker/replica:0" } + node { name: "ColocateWithMe_1" op: "FloatOutput" + device: "/job:worker/replica:1" } + node { name: "ColocateWithMe_2" op: "FloatOutput" + device: "/job:worker/replica:2" } + node { name: "FloatOutput_0" op: "FloatOutput" + device: "/job:worker/replica:0" + attr { key: "_class" + value { list { + s: "loc:@ColocateWithMe_0"}}}} + node { name: "FloatOutput_1" op: "FloatOutput" + device: "/job:worker/replica:1" + attr { key: "_class" + value { list { + s: "loc:@ColocateWithMe_1"}}}} + node { name: "FloatOutput_2" op: "FloatOutput" + device: "/job:worker/replica:2" + attr { key: "_class" + value { list { + s: "loc:@ColocateWithMe_2"}}}} + """, gd) + + def testControlDependencies(self): + + class DependingThread(self.TestThread): + + def __init__(self, graph, replica_id, dependency_op): + super(DependingThread, self).__init__(graph, replica_id) + self._dependency_op = dependency_op + + def run(self): + with g.control_dependencies([self._dependency_op]): + self.has_mutated_graph.set() + self.should_continue.wait() + self.should_continue.clear() + g.create_op( + "FloatOutput", [], [dtypes.float32], + name="FloatOutput_{}".format(self._replica_id)) + + g = ops.Graph() + dependency_ops = [] + for i in range(3): + dependency_ops.append( + g.create_op( + "FloatOutput", [], [dtypes.float32], + name="ColocateWithMe_{}".format(i))) + + # If `switch_to_thread` isn't called, then `input` values for the ops below + # are not deterministic. + g.switch_to_thread_local() + threads = [DependingThread(g, i, dependency_ops[i]) for i in range(3)] + for t in threads: + t.start() + t.has_mutated_graph.wait() + t.has_mutated_graph.clear() + for t in threads: + t.should_continue.set() + t.join() + + gd = g.as_graph_def() + self.assertProtoEqualsVersion(""" + node { name: "ColocateWithMe_0" op: "FloatOutput" } + node { name: "ColocateWithMe_1" op: "FloatOutput" } + node { name: "ColocateWithMe_2" op: "FloatOutput" } + node { name: "FloatOutput_0" op: "FloatOutput" + input: "^ColocateWithMe_0" } + node { name: "FloatOutput_1" op: "FloatOutput" + input: "^ColocateWithMe_1" } + node { name: "FloatOutput_2" op: "FloatOutput" + input: "^ColocateWithMe_2" } + """, gd) + + @test_util.with_c_api class ObjectWithName(object): @@ -1589,7 +1764,13 @@ class ControlDependenciesTest(test_util.TensorFlowTestCase): return constant_op.constant(2.0) future.calls = 0 - if context.in_graph_mode(): + if context.executing_eagerly(): + a = constant_op.constant(1.0) + b = future + with ops.control_dependencies([a, b]): + c = constant_op.constant(3.0) + self.assertEqual(future.calls, 1) + else: g = ops.Graph() with g.as_default(): a = constant_op.constant(1.0) @@ -1598,12 +1779,6 @@ class ControlDependenciesTest(test_util.TensorFlowTestCase): c = constant_op.constant(3.0) self.assertEqual(c.op.control_inputs, [a.op, b.op]) self.assertEqual(future.calls, 1) - else: - a = constant_op.constant(1.0) - b = future() - with ops.control_dependencies([a, b]): - c = constant_op.constant(3.0) - self.assertEqual(future.calls, 1) def testBasicWithConversion(self): g = ops.Graph() @@ -1976,19 +2151,11 @@ class InitScopeTest(test_util.TensorFlowTestCase): with ops.init_scope(): # Because g is building a function, init_scope should # escape out to the eager context. - self.assertTrue(context.in_eager_mode()) + self.assertTrue(context.executing_eagerly()) # g should be reinstated as the default graph, and the # graph context should be re-entered. self.assertIs(g, ops.get_default_graph()) - self.assertTrue(context.in_graph_mode()) - - def testAllGraphsBuildingFunctionsRaisesError(self): - g = ops.Graph() - g._building_function = True # pylint: disable=protected-access - with g.as_default(): - with self.assertRaises(AssertionError): - with ops.init_scope(): - pass + self.assertFalse(context.executing_eagerly()) def testStaysInEagerWhenOnlyEagerContextActive(self): with context.eager_mode(): @@ -2067,6 +2234,29 @@ class InitScopeTest(test_util.TensorFlowTestCase): self.assertEqual(4, int(compiled_outer(inner=compiled_inner))) self.assertEqual(7, int(compiled_outer(inner=compiled_inner))) + def testFallsBackToGlobalGraphWhenAllGraphsAreBuildingFunctions(self): + with context.graph_mode(): + ops.reset_default_graph() + # This doesn't push anything onto the graph stack, but it does + # set the stack's global graph. + global_graph = ops.get_default_graph() + fn_graph = ops.Graph() + + # pylint: disable=protected-access + fn_graph._building_function = True + self.assertEqual(len(ops._default_graph_stack.stack), 0) + with fn_graph.as_default(): + self.assertEqual(len(ops._default_graph_stack.stack), 1) + with ops.init_scope(): + self.assertGreater(len(ops._default_graph_stack.stack), 1) + dummy = constant_op.constant(1.0) + self.assertEqual(len(ops._default_graph_stack.stack), 1) + # Note that the global graph is _not_ on the graph stack. + self.assertEqual(len(ops._default_graph_stack.stack), 0) + # Ensure that `dummy` was added to the global graph. + self.assertEqual(global_graph, dummy.graph) + # pylint: enable=protected-access + def testInstallsDefaultGraphWhenGraphStackIsEmptyInGraphMode(self): with context.graph_mode(): # pylint: disable=protected-access @@ -2088,12 +2278,13 @@ class InitScopeTest(test_util.TensorFlowTestCase): with context.eager_mode(): def foo(): with ops.name_scope("inner"), ops.init_scope(): - if context.in_graph_mode(): - self.assertEqual(ops.get_name_scope(), "inner") - else: + if context.executing_eagerly(): # A trailing slash is always appended when eager execution is # enabled. self.assertEqual(context.context().scope_name, "inner/") + else: + self.assertEqual(ops.get_name_scope(), "inner") + foo() self.assertEqual(ops.get_name_scope(), "") foo_compiled = eager_function.defun(foo) @@ -2703,7 +2894,7 @@ class OutputTypesTest(test_util.TensorFlowTestCase): with g.as_default(): x = constant_op.constant([1, 1, 2, 4, 4, 4, 7, 8, 8], dtype=dtypes.double) - y, _ = gen_array_ops._unique(x) + y, _ = gen_array_ops.unique(x) self.assertEqual([types_pb2.DT_DOUBLE, types_pb2.DT_INT32], y.op._output_types) # pylint: disable=protected-access @@ -2728,6 +2919,9 @@ class EnableEagerExecutionTest(test_util.TensorFlowTestCase): with self.assertRaisesRegexp(ValueError, "device_policy must be one of"): c = config_pb2.ConfigProto() ops.enable_eager_execution(c, c) + with self.assertRaisesRegexp(ValueError, "execution_mode must be one of"): + c = config_pb2.ConfigProto() + ops.enable_eager_execution(c, execution_mode=c) if __name__ == "__main__": diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc index 65810fa7094409c7429dbaaa6c1e62efb263eafc..03721c9a680c285142eca8a399d26c5c376d62ee 100644 --- a/tensorflow/python/framework/python_op_gen.cc +++ b/tensorflow/python/framework/python_op_gen.cc @@ -75,6 +75,33 @@ bool IsPythonReserved(const string& s) { return kPythonReserved->count(s) > 0; } +bool IsOpWithUnderscorePrefix(const string& s) { + static const std::set* const kUnderscoreOps = new std::set( + {// Lowercase built-in functions and types in Python, from: + // [x for x in dir(__builtins__) if x[0].islower()] + // These need to be excluded so they don't conflict with actual built-in + // functions since we use '*' imports. + "abs", "all", "any", "apply", "bin", "bool", "buffer", "bytearray", + "bytes", "callable", "chr", "classmethod", "cmp", "coerce", "compile", + "complex", "copyright", "credits", "delattr", "dict", "dir", "divmod", + "enumerate", "eval", "execfile", "exit", "file", "filter", "float", + "format", "frozenset", "getattr", "globals", "hasattr", "hash", "help", + "hex", "id", "input", "int", "intern", "isinstance", "issubclass", + "iter", "len", "license", "list", "locals", "long", "map", "max", + "memoryview", "min", "next", "object", "oct", "open", "ord", "pow", + "print", "property", "quit", "range", "raw_input", "reduce", "reload", + "repr", "reversed", "round", "set", "setattr", "slice", "sorted", + "staticmethod", "str", "sum", "super", "tuple", "type", "unichr", + "unicode", "vars", "xrange", "zip", + // These have the same name as ops defined in Python and might be used + // incorrectly depending on order of '*' imports. + // TODO(annarev): reduce usage of '*' imports and remove these from the + // list. + "fused_batch_norm", "histogram_fixed_width", "stack", + "batch_norm_with_global_normalization"}); + return kUnderscoreOps->count(s) > 0; +} + string AvoidPythonReserved(const string& s) { if (IsPythonReserved(s)) return strings::StrCat(s, "_"); return s; @@ -476,9 +503,6 @@ GenPythonOp::GenPythonOp(const OpDef& op_def, const ApiDef& api_def, GenPythonOp::~GenPythonOp() {} string GenPythonOp::Code() { - if (api_def_.visibility() == ApiDef::SKIP) { - return ""; - } // This has all the input args followed by those attrs that don't have // defaults. std::vector params_no_default; @@ -583,8 +607,13 @@ void GenPythonOp::AddExport() { strings::StrAppend(&result_, ")\n"); } +void GenPythonOp::AddDefLine(const string& function_name, + const string& parameters) { + strings::StrAppend(&result_, "def ", function_name, "(", parameters, "):\n"); +} + void GenPythonOp::AddDefLine(const string& parameters) { - strings::StrAppend(&result_, "def ", function_name_, "(", parameters, "):\n"); + AddDefLine(function_name_, parameters); } void GenPythonOp::AddDocStringDescription() { @@ -805,28 +834,47 @@ from tensorflow.python.util.tf_export import tf_export auto out = cleaned_ops.mutable_op(); out->Reserve(ops.op_size()); for (const auto& op_def : ops.op()) { - bool is_hidden = false; - for (const string& hidden : hidden_ops) { - if (op_def.name() == hidden) { - is_hidden = true; - break; + const auto* api_def = api_defs.GetApiDef(op_def.name()); + + if (api_def->visibility() == ApiDef::SKIP) { + continue; + } + + // An op is hidden if either its ApiDef visibility is HIDDEN + // or it is in the hidden_ops list. + bool is_hidden = api_def->visibility() == ApiDef::HIDDEN; + bool hidden_by_api_def = is_hidden; + if (!is_hidden) { + for (const string& hidden : hidden_ops) { + if (op_def.name() == hidden) { + is_hidden = true; + break; + } } } string function_name; python_op_gen_internal::GenerateLowerCaseOpName(op_def.name(), &function_name); - if (is_hidden) function_name = strings::StrCat("_", function_name); - - // When users create custom python wrappers, they may link in the - // default op registry by accident, and because they can't - // enumerate all 'hidden' symbols, this guard is to prevent - // instantiating a python reserved word in their wrapper. - if (python_op_gen_internal::IsPythonReserved(function_name)) { + bool is_reserved = python_op_gen_internal::IsPythonReserved(function_name); + + // Prefix an op with underscore if the op is listed in hidden_ops or + // name is reserved or it is of the exceptions in IsOpWithUnderscorePrefix. + // Do not add underscores to ops set to HIDDEN in ApiDef otherwise. + // TODO(annarev): don't prefix with underscores even if op is in hidden_ops. + if (is_hidden) { + if (!hidden_by_api_def || is_reserved || + python_op_gen_internal::IsOpWithUnderscorePrefix(function_name)) { + function_name = strings::StrCat("_", function_name); + } + } else if (is_reserved) { + // When users create custom python wrappers, they may link in the + // default op registry by accident, and because they can't + // enumerate all 'hidden' symbols, this guard is to prevent + // instantiating a python reserved word in their wrapper. continue; } - const auto* api_def = api_defs.GetApiDef(op_def.name()); strings::StrAppend(&result, GetPythonOp(op_def, *api_def, function_name)); if (!require_shapes) { diff --git a/tensorflow/python/framework/python_op_gen_internal.h b/tensorflow/python/framework/python_op_gen_internal.h index d09b36a3e8247241420649c6a4a950be6edc3c00..e0cfb05f4bdf8afd09957c62a9ba3af1fd0882a6 100644 --- a/tensorflow/python/framework/python_op_gen_internal.h +++ b/tensorflow/python/framework/python_op_gen_internal.h @@ -29,6 +29,9 @@ namespace python_op_gen_internal { // Returns true if s is a Python keyword or built-in. bool IsPythonReserved(const string& s); +// Whether the op should be prefixed with underscore. +bool IsOpWithUnderscorePrefix(const string& s); + // Add a _ to the end of s if necessary to avoid a Python keyword or built-in. string AvoidPythonReserved(const string& s); @@ -73,6 +76,7 @@ class GenPythonOp { protected: // Print: def Function(parameters): + void AddDefLine(const string& function_name, const string& parameters); void AddDefLine(const string& parameters); // Format the Op's descriptions so that it can be a Python docstring. diff --git a/tensorflow/python/framework/random_seed.py b/tensorflow/python/framework/random_seed.py index 1e74a790a3fb0c72b7c0fb1127ffac95f386d85e..b724432e00b0d11de86a0fff9ff31758ad36479f 100644 --- a/tensorflow/python/framework/random_seed.py +++ b/tensorflow/python/framework/random_seed.py @@ -52,20 +52,20 @@ def get_seed(op_seed): A tuple of two integers that should be used for the local seed of this operation. """ - is_graph_mode = context.in_graph_mode() + eager = context.executing_eagerly() - if is_graph_mode: - global_seed = ops.get_default_graph().seed - else: + if eager: global_seed = context.global_seed() + else: + global_seed = ops.get_default_graph().seed if global_seed is not None: if op_seed is None: # pylint: disable=protected-access - if is_graph_mode: - op_seed = ops.get_default_graph()._last_id - else: + if eager: op_seed = context.internal_operation_seed() + else: + op_seed = ops.get_default_graph()._last_id seeds = _truncate_seed(global_seed), _truncate_seed(op_seed) else: @@ -176,7 +176,7 @@ def set_random_seed(seed): Args: seed: integer. """ - if context.in_graph_mode(): - ops.get_default_graph().seed = seed - else: + if context.executing_eagerly(): context.set_global_seed(seed) + else: + ops.get_default_graph().seed = seed diff --git a/tensorflow/python/framework/random_seed_test.py b/tensorflow/python/framework/random_seed_test.py index b4c98ab8b289c850c6171425167bb17606a4162d..194492268631abfa911bd45f13a302c09a2c8bda 100644 --- a/tensorflow/python/framework/random_seed_test.py +++ b/tensorflow/python/framework/random_seed_test.py @@ -40,13 +40,13 @@ class RandomSeedTest(test.TestCase): ((2**31 - 1, 0), (0, 2**31 - 1)), # Don't wrap to (0, 0) either ((0, 2**31 - 1), (0, 2**31 - 1)), # Wrapping for the other argument ] - if context.in_graph_mode(): - # 0 will be the default_graph._lastid. - test_cases.append(((1, None), (1, 0))) - else: + if context.executing_eagerly(): # operation seed is random number generated based on global seed. # it's not tested due to possibility of platform or version difference. pass + else: + # 0 will be the default_graph._lastid. + test_cases.append(((1, None), (1, 0))) for tc in test_cases: tinput, toutput = tc[0], tc[1] random_seed.set_random_seed(tinput[0]) diff --git a/tensorflow/python/framework/smart_cond.py b/tensorflow/python/framework/smart_cond.py new file mode 100644 index 0000000000000000000000000000000000000000..c7ff23e4ff809ed7bc57259fa3ec9feb921b5a71 --- /dev/null +++ b/tensorflow/python/framework/smart_cond.py @@ -0,0 +1,123 @@ +# 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. +# ============================================================================== +"""smart_cond and related utilties.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python import pywrap_tensorflow as c_api +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import control_flow_ops + + +def smart_cond(pred, true_fn=None, false_fn=None, name=None): + """Return either `true_fn()` if predicate `pred` is true else `false_fn()`. + + If `pred` is a bool or has a constant value, we return either `true_fn()` + or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both. + + Arguments: + pred: A scalar determining whether to return the result of `true_fn` or + `false_fn`. + true_fn: The callable to be performed if pred is true. + false_fn: The callable to be performed if pred is false. + name: Optional name prefix when using `tf.cond`. + + Returns: + Tensors returned by the call to either `true_fn` or `false_fn`. + + Raises: + TypeError: If `true_fn` or `false_fn` is not callable. + """ + if not callable(true_fn): + raise TypeError("`true_fn` must be callable.") + if not callable(false_fn): + raise TypeError("`false_fn` must be callable.") + + pred_value = smart_constant_value(pred) + if pred_value is not None: + if pred_value: + return true_fn() + else: + return false_fn() + else: + return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn, + name=name) + + +def smart_constant_value(pred): + """Return the bool value for `pred`, or None if `pred` had a dynamic value. + + Arguments: + pred: A scalar, either a Python bool or tensor. + + Returns: + True or False if `pred` has a constant boolean value, None otherwise. + + Raises: + TypeError: If `pred` is not a Tensor or bool. + """ + if pred in {0, 1}: # Accept 1/0 as valid boolean values + pred_value = bool(pred) + elif isinstance(pred, bool): + pred_value = pred + elif isinstance(pred, ops.Tensor): + pred_value = tensor_util.constant_value(pred) + # TODO(skyewm): consider folding this into tensor_util.constant_value when + # _USE_C_API is removed (there may be performance and correctness bugs, so I + # wanted to limit the change hidden behind _USE_C_API). + # pylint: disable=protected-access + if pred_value is None and ops._USE_C_API: + with errors.raise_exception_on_not_ok_status() as status: + pred_value = c_api.TF_TryEvaluateConstant_wrapper( + pred.graph._c_graph, pred._as_tf_output(), status) + # pylint: enable=protected-access + + else: + raise TypeError("`pred` must be a Tensor, or a Python bool, or 1 or 0. " + "Found instead: %s" % pred) + return pred_value + + +def smart_case(pred_fn_pairs, default=None, exclusive=False, name="smart_case"): + """Like tf.case, except attempts to statically evaluate predicates. + + If any predicate in `pred_fn_pairs` is a bool or has a constant value, the + associated callable will be called or omitted depending on its value. + Otherwise this functions like tf.case. + + Args: + pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor and a + callable which returns a list of tensors. + default: Optional callable that returns a list of tensors. + exclusive: True iff at most one predicate is allowed to evaluate to `True`. + name: A name for this operation (optional). + + Returns: + The tensors returned by the first pair whose predicate evaluated to True, or + those returned by `default` if none does. + + Raises: + TypeError: If `pred_fn_pairs` is not a list/dictionary. + TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples. + TypeError: If `fns[i]` is not callable for any i, or `default` is not + callable. + """ + return control_flow_ops._case_helper( # pylint: disable=protected-access + smart_cond, pred_fn_pairs, default, exclusive, name, + allow_python_preds=True) diff --git a/tensorflow/python/framework/smart_cond_test.py b/tensorflow/python/framework/smart_cond_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1170a41c99995ae875e58a2d5491e05bc1e40df6 --- /dev/null +++ b/tensorflow/python/framework/smart_cond_test.py @@ -0,0 +1,166 @@ +# 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 + +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.framework import smart_cond +from tensorflow.python.framework import test_util +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.platform import googletest + + +def raise_exception(): + raise RuntimeError("did not expect to be called") + + +@test_util.with_c_api +class SmartCondTest(test_util.TensorFlowTestCase): + + def testTrue(self): + with ops.Graph().as_default(): + with session.Session(): + x = constant_op.constant(2) + y = constant_op.constant(5) + z = smart_cond.smart_cond(True, lambda: math_ops.multiply(x, 16), + lambda: math_ops.multiply(y, 5)) + self.assertEqual(z.eval(), 32) + + def testFalse(self): + with ops.Graph().as_default(): + with session.Session(): + x = constant_op.constant(4) + y = constant_op.constant(3) + z = smart_cond.smart_cond(False, lambda: math_ops.multiply(x, 16), + lambda: math_ops.multiply(y, 3)) + self.assertEqual(z.eval(), 9) + + def testUnknown(self): + with ops.Graph().as_default(): + with session.Session(): + x = array_ops.placeholder(dtype=dtypes.int32) + y = smart_cond.smart_cond(x > 0, lambda: constant_op.constant(1), + lambda: constant_op.constant(2)) + self.assertEqual(y.eval(feed_dict={x: 1}), 1) + self.assertEqual(y.eval(feed_dict={x: -1}), 2) + + def testEval(self): + # Constant expression evaluation only works with the C API enabled. + if not ops._USE_C_API: return + + with ops.Graph().as_default(): + with session.Session(): + x = constant_op.constant(1) + y = constant_op.constant(2) + # x * y > 0 can be evaluated at graph construction time, so the false + # branch shouldn't be evaluated at all. + z = smart_cond.smart_cond(x * y > 0, lambda: constant_op.constant(1), + raise_exception) + self.assertEqual(z.eval(feed_dict={x: 1}), 1) + + def testPlaceholderWithDefault(self): + with ops.Graph().as_default(): + with session.Session(): + x = array_ops.placeholder_with_default(1, shape=()) + y = smart_cond.smart_cond(x > 0, lambda: constant_op.constant(1), + lambda: constant_op.constant(2)) + self.assertEqual(y.eval(), 1) + self.assertEqual(y.eval(feed_dict={x: -1}), 2) + + def testMissingArg1(self): + with ops.Graph().as_default(): + with session.Session(): + x = constant_op.constant(1) + with self.assertRaises(TypeError): + smart_cond.smart_cond(True, false_fn=lambda: x) + + def testMissingArg2(self): + with ops.Graph().as_default(): + with session.Session(): + x = constant_op.constant(1) + with self.assertRaises(TypeError): + smart_cond.smart_cond(True, lambda: x) + + +@test_util.with_c_api +class SmartCaseTest(test_util.TensorFlowTestCase): + + def testTrue(self): + x = array_ops.placeholder(dtype=dtypes.int32, shape=[]) + conditions = [(True, lambda: constant_op.constant(1)), + (x == 0, raise_exception)] + y = smart_cond.smart_case(conditions, default=raise_exception, + exclusive=False) + z = smart_cond.smart_case(conditions, default=raise_exception, + exclusive=True) + with session.Session() as sess: + # No feed_dict necessary + self.assertEqual(sess.run(y), 1) + self.assertEqual(sess.run(z), 1) + + def testFalse(self): + conditions = [(False, raise_exception)] + y = smart_cond.smart_case(conditions, + default=lambda: constant_op.constant(1), + exclusive=False) + z = smart_cond.smart_case(conditions, + default=lambda: constant_op.constant(1), + exclusive=True) + with session.Session() as sess: + self.assertEqual(sess.run(y), 1) + self.assertEqual(sess.run(z), 1) + + def testMix(self): + # Constant expression evaluation only works with the C API enabled. + if not ops._USE_C_API: return + + x = array_ops.placeholder(dtype=dtypes.int32, shape=[]) + y = constant_op.constant(10) + conditions = [(x > 1, lambda: constant_op.constant(1)), + (y < 1, raise_exception), + (False, raise_exception), + (True, lambda: constant_op.constant(3))] + z = smart_cond.smart_case(conditions, default=raise_exception) + with session.Session() as sess: + self.assertEqual(sess.run(z, feed_dict={x: 2}), 1) + self.assertEqual(sess.run(z, feed_dict={x: 0}), 3) + + +@test_util.with_c_api +class SmartConstantValueTest(test_util.TensorFlowTestCase): + + # TODO(skyewm): this is essentially a regression test for + # TF_TryEvaluateConstant, and is not really a valid smart_constant_value test + # (smart_constant_value is only supposed to return bools). Move the + # TF_TryEvaluateConstant call to tensor_util.constant_value and make this a + # constant_value test instead. + def testCond(self): + with ops.Graph().as_default(): + pred = array_ops.placeholder_with_default(True, shape=()) + x = control_flow_ops.cond(pred, + lambda: constant_op.constant(1), + lambda: constant_op.constant(2)) + self.assertIsNone(smart_cond.smart_constant_value(x)) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/python/framework/tensor_shape.py b/tensorflow/python/framework/tensor_shape.py index 222071cb9e87aa0fdd9788d1c72df4c66ea61547..af2a5b1a7ef9a70c0baf5d02257951803a7a76fa 100644 --- a/tensorflow/python/framework/tensor_shape.py +++ b/tensorflow/python/framework/tensor_shape.py @@ -156,7 +156,7 @@ class Dimension(object): ``` Args: - other: Another Dimension. + other: Another Dimension, or a value accepted by `as_dimension`. Returns: A Dimension whose value is the sum of `self` and `other`. @@ -167,6 +167,17 @@ class Dimension(object): else: return Dimension(self._value + other.value) + def __radd__(self, other): + """Returns the sum of `other` and `self`. + + Args: + other: Another Dimension, or a value accepted by `as_dimension`. + + Returns: + A Dimension whose value is the sum of `self` and `other`. + """ + return self + other + def __sub__(self, other): """Returns the subtraction of `other` from `self`. @@ -180,10 +191,10 @@ class Dimension(object): ``` Args: - other: Another Dimension. + other: Another Dimension, or a value accepted by `as_dimension`. Returns: - A Dimension whose value is the subtraction of sum of `other` from `self`. + A Dimension whose value is the subtraction of `other` from `self`. """ other = as_dimension(other) if self._value is None or other.value is None: @@ -191,6 +202,21 @@ class Dimension(object): else: return Dimension(self._value - other.value) + def __rsub__(self, other): + """Returns the subtraction of `self` from `other`. + + Args: + other: Another Dimension, or a value accepted by `as_dimension`. + + Returns: + A Dimension whose value is the subtraction of `self` from `other`. + """ + other = as_dimension(other) + if self._value is None or other.value is None: + return Dimension(None) + else: + return Dimension(other.value - self._value) + def __mul__(self, other): """Returns the product of `self` and `other`. @@ -204,17 +230,32 @@ class Dimension(object): ``` Args: - other: Another Dimension. + other: Another Dimension, or a value accepted by `as_dimension`. Returns: A Dimension whose value is the product of `self` and `other`. """ - other = as_dimension(other) + try: + other = as_dimension(other) + except (TypeError, ValueError): + return NotImplemented + if self._value is None or other.value is None: return Dimension(None) else: return Dimension(self._value * other.value) + def __rmul__(self, other): + """Returns the product of `self` and `other`. + + Args: + other: Another Dimension, or a value accepted by `as_dimension`. + + Returns: + A Dimension whose value is the product of `self` and `other`. + """ + return self * other + def __floordiv__(self, other): """Returns the quotient of `self` and `other` rounded down. @@ -228,17 +269,35 @@ class Dimension(object): ``` Args: - other: Another `Dimension`. + other: Another Dimension, or a value accepted by `as_dimension`. Returns: A `Dimension` whose value is the integer quotient of `self` and `other`. """ - other = as_dimension(other) + try: + other = as_dimension(other) + except (TypeError, ValueError): + return NotImplemented if self._value is None or other.value is None: return Dimension(None) else: return Dimension(self._value // other.value) + def __rfloordiv__(self, other): + """Returns the quotient of `other` and `self` rounded down. + + Args: + other: Another Dimension, or a value accepted by `as_dimension`. + + Returns: + A `Dimension` whose value is the integer quotient of `self` and `other`. + """ + other = as_dimension(other) + if self._value is None or other.value is None: + return Dimension(None) + else: + return Dimension(other.value // self._value) + def __div__(self, other): """DEPRECATED: Use `__floordiv__` via `x // y` instead. @@ -256,7 +315,7 @@ class Dimension(object): return self // other def __mod__(self, other): - """Returns `self` modulo `other. + """Returns `self` modulo `other`. Dimension moduli are computed as follows: @@ -268,17 +327,35 @@ class Dimension(object): ``` Args: - other: Another Dimension. + other: Another Dimension, or a value accepted by `as_dimension`. Returns: A Dimension whose value is `self` modulo `other`. """ - other = as_dimension(other) + try: + other = as_dimension(other) + except (TypeError, ValueError): + return NotImplemented if self._value is None or other.value is None: return Dimension(None) else: return Dimension(self._value % other.value) + def __rmod__(self, other): + """Returns `other` modulo `self`. + + Args: + other: Another Dimension, or a value accepted by `as_dimension`. + + Returns: + A Dimension whose value is `other` modulo `self`. + """ + try: + other = as_dimension(other) + except (TypeError, ValueError): + return NotImplemented + return other % self + def __lt__(self, other): """Returns True if `self` is known to be less than `other`. @@ -456,6 +533,7 @@ class TensorShape(object): else: # Got a list of dimensions self._dims = [as_dimension(d) for d in dims_iter] + self._ndims = None def __repr__(self): return "TensorShape(%r)" % self._dims @@ -473,19 +551,26 @@ class TensorShape(object): """Returns a list of Dimensions, or None if the shape is unspecified.""" return self._dims + @dims.setter + def dims(self, dims): + self._dims = dims + self._ndims = None + @property def ndims(self): """Returns the rank of this shape, or None if it is unspecified.""" if self._dims is None: return None else: - return len(self._dims) + if self._ndims is None: + self._ndims = len(self._dims) + return self._ndims def __len__(self): """Returns the rank of this shape, or raises ValueError if unspecified.""" if self._dims is None: raise ValueError("Cannot take the length of Shape with unknown rank.") - return len(self._dims) + return self.ndims def __bool__(self): """Returns True if this shape contains non-zero information.""" diff --git a/tensorflow/python/framework/tensor_shape_test.py b/tensorflow/python/framework/tensor_shape_test.py index fffd86c7a6241b8be92ad33852da244ab9b5284d..4e8ce4d889c4ef0c6e56806587a64e8f9be7e10a 100644 --- a/tensorflow/python/framework/tensor_shape_test.py +++ b/tensorflow/python/framework/tensor_shape_test.py @@ -34,12 +34,20 @@ class DimensionTest(test_util.TensorFlowTestCase): self.assertEqual(tensor_shape.Dimension(15), dim + tensor_shape.Dimension(3)) self.assertEqual(tensor_shape.Dimension(15), dim + 3) + self.assertEqual(tensor_shape.Dimension(15), 3 + dim) + self.assertEqual(tensor_shape.Dimension(9), dim - 3) + self.assertEqual(tensor_shape.Dimension(1), 13 - dim) self.assertEqual(tensor_shape.Dimension(24), dim * tensor_shape.Dimension(2)) self.assertEqual(tensor_shape.Dimension(24), dim * 2) + self.assertEqual(tensor_shape.Dimension(24), 2 * dim) + self.assertEqual([4] * 12, [4] * dim) + self.assertEqual(12 * [4], dim * [4]) + self.assertEqual(tensor_shape.Dimension(24), 2 * dim) self.assertEqual( tensor_shape.Dimension(6), dim // tensor_shape.Dimension(2)) self.assertEqual(tensor_shape.Dimension(6), dim // 2) + self.assertEqual(tensor_shape.Dimension(0), 2 // dim) self.assertEqual(tensor_shape.Dimension(12), dim.merge_with(tensor_shape.Dimension(12))) self.assertEqual(tensor_shape.Dimension(12), dim.merge_with(12)) @@ -176,6 +184,14 @@ class DimensionTest(test_util.TensorFlowTestCase): self.assertEqual(str(tensor_shape.Dimension(7)), "7") self.assertEqual(str(tensor_shape.Dimension(None)), "?") + def testMod(self): + four = tensor_shape.Dimension(4) + nine = tensor_shape.Dimension(9) + self.assertEqual(nine % four, 1) + # test both __mod__ and __rmod__. + self.assertEqual(nine % 4, 1) + self.assertEqual(4 % nine, 4) + class ShapeTest(test_util.TensorFlowTestCase): diff --git a/tensorflow/python/framework/tensor_spec.py b/tensorflow/python/framework/tensor_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..546c48adbae7a97b404b2507994dfff77dd69ec6 --- /dev/null +++ b/tensorflow/python/framework/tensor_spec.py @@ -0,0 +1,213 @@ +# 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. +# ============================================================================== +"""A TensorSpec class.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.python.framework import common_shapes +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape + + +class TensorSpec(object): + """Describes a tf.Tensor. + + A TensorSpec allows an API to describe the Tensors that it accepts or + returns, before that Tensor exists. This allows dynamic and flexible graph + construction and configuration. + """ + + __slots__ = ["_shape", "_dtype", "_name"] + + def __init__(self, shape, dtype, name=None): + """Creates a TensorSpec. + + Args: + shape: Value convertible to `tf.TensorShape`. The shape of the tensor. + dtype: Value convertible to `tf.DType`. The type of the tensor values. + name: Optional name for the Tensor. + + Raises: + TypeError: If shape is not convertible to a `tf.TensorShape`, or dtype is + not convertible to a `tf.DType`. + """ + self._shape = tensor_shape.TensorShape(shape) + self._dtype = dtypes.as_dtype(dtype) + self._name = name + + @classmethod + def from_spec(cls, spec, name=None): + return cls(spec.shape, spec.dtype, name or spec.name) + + @classmethod + def from_tensor(cls, tensor, name=None): + if isinstance(tensor, ops.EagerTensor): + return TensorSpec(tensor.shape, tensor.dtype, name) + elif isinstance(tensor, ops.Tensor): + return TensorSpec(tensor.shape, tensor.dtype, name or tensor.op.name) + else: + raise ValueError("`tensor` should be a tf.Tensor") + + @classmethod + def is_bounded(cls): + del cls + return False + + @property + def shape(self): + """Returns the `TensorShape` that represents the shape of the tensor.""" + return self._shape + + @property + def dtype(self): + """Returns the `dtype` of elements in the tensor.""" + return self._dtype + + @property + def name(self): + """Returns the name of the described tensor.""" + return self._name + + @property + def is_discrete(self): + """Whether spec is discrete.""" + return self.dtype.is_integer + + @property + def is_continuous(self): + """Whether spec is continuous.""" + return self.dtype.is_floating + + def is_compatible_with(self, spec_or_tensor): + """True if the shape and dtype of `spec_or_tensor` are compatible.""" + return (self._dtype.is_compatible_with(spec_or_tensor.dtype) and + self._shape.is_compatible_with(spec_or_tensor.shape)) + + def __repr__(self): + return "TensorSpec(shape={}, dtype={}, name={})".format( + self.shape, repr(self.dtype), repr(self.name)) + + def __eq__(self, other): + return self.shape == other.shape and self.dtype == other.dtype + + def __ne__(self, other): + return not self == other + + +class BoundedTensorSpec(TensorSpec): + """A `TensorSpec` that specifies minimum and maximum values. + + Example usage: + ```python + spec = tensor_spec.BoundedTensorSpec((1, 2, 3), tf.float32, 0, (5, 5, 5)) + tf_minimum = tf.convert_to_tensor(spec.minimum, dtype=spec.dtype) + tf_maximum = tf.convert_to_tensor(spec.maximum, dtype=spec.dtype) + ``` + + Bounds are meant to be inclusive. This is especially important for + integer types. The following spec will be satisfied by tensors + with values in the set {0, 1, 2}: + ```python + spec = tensor_spec.BoundedTensorSpec((3, 5), tf.int32, 0, 2) + ``` + """ + + __slots__ = ("_minimum", "_maximum") + + def __init__(self, shape, dtype, minimum, maximum, name=None): + """Initializes a new `BoundedTensorSpec`. + + Args: + shape: Value convertible to `tf.TensorShape`. The shape of the tensor. + dtype: Value convertible to `tf.DType`. The type of the tensor values. + minimum: Number or sequence specifying the minimum element bounds + (inclusive). Must be broadcastable to `shape`. + maximum: Number or sequence specifying the maximum element bounds + (inclusive). Must be broadcastable to `shape`. + name: Optional string containing a semantic name for the corresponding + array. Defaults to `None`. + + Raises: + ValueError: If `minimum` or `maximum` are not provided or not + broadcastable to `shape`. + TypeError: If the shape is not an iterable or if the `dtype` is an invalid + numpy dtype. + """ + super(BoundedTensorSpec, self).__init__(shape, dtype, name) + + if minimum is None or maximum is None: + raise ValueError("minimum and maximum must be provided; but saw " + "'%s' and '%s'" % (minimum, maximum)) + + try: + minimum_shape = np.shape(minimum) + common_shapes.broadcast_shape( + tensor_shape.TensorShape(minimum_shape), self.shape) + except ValueError as exception: + raise ValueError("minimum is not compatible with shape. " + "Message: {!r}.".format(exception)) + + try: + maximum_shape = np.shape(maximum) + common_shapes.broadcast_shape( + tensor_shape.TensorShape(maximum_shape), self.shape) + except ValueError as exception: + raise ValueError("maximum is not compatible with shape. " + "Message: {!r}.".format(exception)) + + self._minimum = np.array(minimum, dtype=self.dtype.as_numpy_dtype()) + self._minimum.setflags(write=False) + + self._maximum = np.array(maximum, dtype=self.dtype.as_numpy_dtype()) + self._maximum.setflags(write=False) + + @classmethod + def is_bounded(cls): + del cls + return True + + @classmethod + def from_spec(cls, spec): + dtype = dtypes.as_dtype(spec.dtype) + minimum = getattr(spec, "minimum", dtype.min) + maximum = getattr(spec, "maximum", dtype.max) + return BoundedTensorSpec(spec.shape, dtype, minimum, maximum, spec.name) + + @property + def minimum(self): + """Returns a NumPy array specifying the minimum bounds (inclusive).""" + return self._minimum + + @property + def maximum(self): + """Returns a NumPy array specifying the maximum bounds (inclusive).""" + return self._maximum + + def __repr__(self): + s = "BoundedTensorSpec(shape={}, dtype={}, name={}, minimum={}, maximum={})" + return s.format(self.shape, repr(self.dtype), repr(self.name), + repr(self.minimum), repr(self.maximum)) + + def __eq__(self, other): + tensor_spec_eq = super(BoundedTensorSpec, self).__eq__(other) + return (tensor_spec_eq and np.allclose(self.minimum, other.minimum) and + np.allclose(self.maximum, other.maximum)) + + diff --git a/tensorflow/python/framework/tensor_spec_test.py b/tensorflow/python/framework/tensor_spec_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b33d769d86d3ea03f6a59082a132b0a94d5921ea --- /dev/null +++ b/tensorflow/python/framework/tensor_spec_test.py @@ -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. +# ============================================================================== +"""Tests for tensor_spec.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import tensor_spec +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import googletest + + +class TensorSpecTest(test_util.TensorFlowTestCase): + + def testAcceptsNumpyDType(self): + desc = tensor_spec.TensorSpec([1], np.float32) + self.assertEqual(desc.dtype, dtypes.float32) + + def testAcceptsTensorShape(self): + desc = tensor_spec.TensorSpec(tensor_shape.TensorShape([1]), dtypes.float32) + self.assertEqual(desc.shape, tensor_shape.TensorShape([1])) + + def testUnknownShape(self): + desc = tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32) + self.assertEqual(desc.shape, tensor_shape.TensorShape(None)) + + def testShapeCompatibility(self): + unknown = array_ops.placeholder(dtypes.int64) + partial = array_ops.placeholder(dtypes.int64, shape=[None, 1]) + full = array_ops.placeholder(dtypes.int64, shape=[2, 3]) + rank3 = array_ops.placeholder(dtypes.int64, shape=[4, 5, 6]) + + desc_unknown = tensor_spec.TensorSpec(None, dtypes.int64) + self.assertTrue(desc_unknown.is_compatible_with(unknown)) + self.assertTrue(desc_unknown.is_compatible_with(partial)) + self.assertTrue(desc_unknown.is_compatible_with(full)) + self.assertTrue(desc_unknown.is_compatible_with(rank3)) + + desc_partial = tensor_spec.TensorSpec([2, None], dtypes.int64) + self.assertTrue(desc_partial.is_compatible_with(unknown)) + self.assertTrue(desc_partial.is_compatible_with(partial)) + self.assertTrue(desc_partial.is_compatible_with(full)) + self.assertFalse(desc_partial.is_compatible_with(rank3)) + + desc_full = tensor_spec.TensorSpec([2, 3], dtypes.int64) + self.assertTrue(desc_full.is_compatible_with(unknown)) + self.assertFalse(desc_full.is_compatible_with(partial)) + self.assertTrue(desc_full.is_compatible_with(full)) + self.assertFalse(desc_full.is_compatible_with(rank3)) + + desc_rank3 = tensor_spec.TensorSpec([4, 5, 6], dtypes.int64) + self.assertTrue(desc_rank3.is_compatible_with(unknown)) + self.assertFalse(desc_rank3.is_compatible_with(partial)) + self.assertFalse(desc_rank3.is_compatible_with(full)) + self.assertTrue(desc_rank3.is_compatible_with(rank3)) + + def testTypeCompatibility(self): + floats = array_ops.placeholder(dtypes.float32, shape=[10, 10]) + ints = array_ops.placeholder(dtypes.int32, shape=[10, 10]) + desc = tensor_spec.TensorSpec(shape=(10, 10), dtype=dtypes.float32) + self.assertTrue(desc.is_compatible_with(floats)) + self.assertFalse(desc.is_compatible_with(ints)) + + def testName(self): + desc = tensor_spec.TensorSpec([1], dtypes.float32, name="beep") + self.assertEqual(desc.name, "beep") + + def testRepr(self): + desc1 = tensor_spec.TensorSpec([1], dtypes.float32, name="beep") + self.assertEqual( + repr(desc1), + "TensorSpec(shape=(1,), dtype=tf.float32, name='beep')") + desc2 = tensor_spec.TensorSpec([1, None], dtypes.int32) + self.assertEqual( + repr(desc2), + "TensorSpec(shape=(1, ?), dtype=tf.int32, name=None)") + + def testFromTensorSpec(self): + spec_1 = tensor_spec.TensorSpec((1, 2), dtypes.int32) + spec_2 = tensor_spec.TensorSpec.from_spec(spec_1) + self.assertEqual(spec_1, spec_2) + + def testFromTensor(self): + zero = constant_op.constant(0) + spec = tensor_spec.TensorSpec.from_tensor(zero) + self.assertEqual(spec.dtype, dtypes.int32) + self.assertEqual(spec.shape, []) + self.assertEqual(spec.name, "Const") + + def testFromPlaceholder(self): + unknown = array_ops.placeholder(dtypes.int64, name="unknown") + partial = array_ops.placeholder(dtypes.float32, + shape=[None, 1], + name="partial") + spec_1 = tensor_spec.TensorSpec.from_tensor(unknown) + self.assertEqual(spec_1.dtype, dtypes.int64) + self.assertEqual(spec_1.shape, None) + self.assertEqual(spec_1.name, "unknown") + spec_2 = tensor_spec.TensorSpec.from_tensor(partial) + self.assertEqual(spec_2.dtype, dtypes.float32) + self.assertEqual(spec_2.shape.as_list(), [None, 1]) + self.assertEqual(spec_2.name, "partial") + + def testFromBoundedTensorSpec(self): + bounded_spec = tensor_spec.BoundedTensorSpec((1, 2), dtypes.int32, 0, 1) + spec = tensor_spec.TensorSpec.from_spec(bounded_spec) + self.assertEqual(bounded_spec.shape, spec.shape) + self.assertEqual(bounded_spec.dtype, spec.dtype) + self.assertEqual(bounded_spec.name, spec.name) + + def testIsDiscrete(self): + discrete_spec = tensor_spec.TensorSpec((1, 2), dtypes.int32) + continuous_spec = tensor_spec.TensorSpec((1, 2), dtypes.float32) + self.assertTrue(discrete_spec.is_discrete) + self.assertFalse(continuous_spec.is_discrete) + + def testIsContinuous(self): + discrete_spec = tensor_spec.TensorSpec((1, 2), dtypes.int32) + continuous_spec = tensor_spec.TensorSpec((1, 2), dtypes.float32) + self.assertFalse(discrete_spec.is_continuous) + self.assertTrue(continuous_spec.is_continuous) + + def testIsBounded(self): + unbounded_spec = tensor_spec.TensorSpec((1, 2), dtypes.int32) + self.assertFalse(unbounded_spec.is_bounded()) + + +class BoundedTensorSpecTest(test_util.TensorFlowTestCase): + + def testInvalidMinimum(self): + with self.assertRaisesRegexp(ValueError, "not compatible"): + tensor_spec.BoundedTensorSpec((3, 5), dtypes.uint8, (0, 0, 0), (1, 1)) + + def testInvalidMaximum(self): + with self.assertRaisesRegexp(ValueError, "not compatible"): + tensor_spec.BoundedTensorSpec((3, 5), dtypes.uint8, 0, (1, 1, 1)) + + def testIsBounded(self): + bounded_spec = tensor_spec.BoundedTensorSpec( + (1, 2), dtypes.int32, minimum=0, maximum=1) + self.assertTrue(bounded_spec.is_bounded()) + + def testMinimumMaximumAttributes(self): + spec = tensor_spec.BoundedTensorSpec( + (1, 2, 3), dtypes.float32, 0, (5, 5, 5)) + self.assertEqual(type(spec.minimum), np.ndarray) + self.assertEqual(type(spec.maximum), np.ndarray) + self.assertAllEqual(spec.minimum, np.array(0, dtype=np.float32)) + self.assertAllEqual(spec.maximum, np.array([5, 5, 5], dtype=np.float32)) + + def testNotWriteableNP(self): + spec = tensor_spec.BoundedTensorSpec( + (1, 2, 3), dtypes.float32, 0, (5, 5, 5)) + with self.assertRaisesRegexp(ValueError, "read-only"): + spec.minimum[0] = -1 + with self.assertRaisesRegexp(ValueError, "read-only"): + spec.maximum[0] = 100 + + def testReuseSpec(self): + spec_1 = tensor_spec.BoundedTensorSpec((1, 2), dtypes.int32, + minimum=0, maximum=1) + spec_2 = tensor_spec.BoundedTensorSpec( + spec_1.shape, spec_1.dtype, spec_1.minimum, spec_1.maximum) + self.assertEqual(spec_1, spec_2) + + def testScalarBounds(self): + spec = tensor_spec.BoundedTensorSpec( + (), dtypes.float32, minimum=0.0, maximum=1.0) + + self.assertIsInstance(spec.minimum, np.ndarray) + self.assertIsInstance(spec.maximum, np.ndarray) + + # Sanity check that numpy compares correctly to a scalar for an empty shape. + self.assertEqual(0.0, spec.minimum) + self.assertEqual(1.0, spec.maximum) + + # Check that the spec doesn't fail its own input validation. + _ = tensor_spec.BoundedTensorSpec( + spec.shape, spec.dtype, spec.minimum, spec.maximum) + + def testFromBoundedTensorSpec(self): + spec_1 = tensor_spec.BoundedTensorSpec((1, 2), dtypes.int32, + minimum=0, maximum=1) + spec_2 = tensor_spec.BoundedTensorSpec.from_spec(spec_1) + self.assertEqual(spec_1, spec_2) + + def testEquality(self): + spec_1_1 = tensor_spec.BoundedTensorSpec((1, 2, 3), dtypes.float32, + 0, (5, 5, 5)) + spec_1_2 = tensor_spec.BoundedTensorSpec((1, 2, 3), dtypes.float32, + 0.00000001, + (5, 5, 5.00000000000000001)) + spec_2_1 = tensor_spec.BoundedTensorSpec((1, 2, 3), dtypes.float32, + 1, (5, 5, 5)) + spec_2_2 = tensor_spec.BoundedTensorSpec((1, 2, 3), dtypes.float32, + (1, 1, 1), (5, 5, 5)) + spec_2_3 = tensor_spec.BoundedTensorSpec((1, 2, 3), dtypes.float32, + (1, 1, 1), 5) + spec_3_1 = tensor_spec.BoundedTensorSpec((1, 2, 3), dtypes.float32, + (2, 1, 1), (5, 5, 5)) + + self.assertEqual(spec_1_1, spec_1_2) + self.assertEqual(spec_1_2, spec_1_1) + + self.assertNotEqual(spec_1_1, spec_2_2) + self.assertNotEqual(spec_1_1, spec_2_1) + self.assertNotEqual(spec_2_2, spec_1_1) + self.assertNotEqual(spec_2_1, spec_1_1) + + self.assertEqual(spec_2_1, spec_2_2) + self.assertEqual(spec_2_2, spec_2_1) + self.assertEqual(spec_2_2, spec_2_3) + + self.assertNotEqual(spec_1_1, spec_3_1) + self.assertNotEqual(spec_2_1, spec_3_1) + self.assertNotEqual(spec_2_2, spec_3_1) + + def testFromTensorSpec(self): + spec = tensor_spec.TensorSpec((1, 2), dtypes.int32) + bounded_spec = tensor_spec.BoundedTensorSpec.from_spec(spec) + self.assertEqual(spec.shape, bounded_spec.shape) + self.assertEqual(spec.dtype, bounded_spec.dtype) + self.assertEqual(spec.dtype.min, bounded_spec.minimum) + self.assertEqual(spec.dtype.max, bounded_spec.maximum) + self.assertEqual(spec.name, bounded_spec.name) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/python/framework/tensor_util.py b/tensorflow/python/framework/tensor_util.py index 0e5f696111ae7f74b41f8af21a5190fc2617e51a..984bcecdfe05efd79bdf218197c410b14abe3516 100644 --- a/tensorflow/python/framework/tensor_util.py +++ b/tensorflow/python/framework/tensor_util.py @@ -557,17 +557,18 @@ def MakeNdarray(tensor): dtype = tensor_dtype.as_numpy_dtype if tensor.tensor_content: - return np.fromstring(tensor.tensor_content, dtype=dtype).reshape(shape) - elif tensor_dtype == dtypes.float16: + return (np.frombuffer(tensor.tensor_content, dtype=dtype).copy() + .reshape(shape)) + elif tensor_dtype == dtypes.float16 or tensor_dtype == dtypes.bfloat16: # the half_val field of the TensorProto stores the binary representation # of the fp16: we need to reinterpret this as a proper float16 if len(tensor.half_val) == 1: tmp = np.array(tensor.half_val[0], dtype=np.uint16) - tmp.dtype = np.float16 + tmp.dtype = tensor_dtype.as_numpy_dtype return np.repeat(tmp, num_elements).reshape(shape) else: tmp = np.fromiter(tensor.half_val, dtype=np.uint16) - tmp.dtype = np.float16 + tmp.dtype = tensor_dtype.as_numpy_dtype return tmp.reshape(shape) elif tensor_dtype == dtypes.float32: if len(tensor.float_val) == 1: @@ -585,8 +586,7 @@ def MakeNdarray(tensor): return np.fromiter(tensor.double_val, dtype=dtype).reshape(shape) elif tensor_dtype in [ dtypes.int32, dtypes.uint8, dtypes.uint16, dtypes.int16, dtypes.int8, - dtypes.qint32, dtypes.quint8, dtypes.qint8, dtypes.qint16, dtypes.quint16, - dtypes.bfloat16 + dtypes.qint32, dtypes.quint8, dtypes.qint8, dtypes.qint16, dtypes.quint16 ]: if len(tensor.int_val) == 1: return np.repeat(np.array(tensor.int_val[0], dtype=dtype), @@ -828,7 +828,7 @@ def constant_value_as_shape(tensor): # pylint: disable=invalid-name Returns: A `TensorShape` based on the constant value of the given `tensor`. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return tensor_shape.as_shape( [dim if dim != -1 else None for dim in tensor.numpy()]) diff --git a/tensorflow/python/framework/tensor_util_test.py b/tensorflow/python/framework/tensor_util_test.py index f2de69e159646b4a085645fa1bfef7782e78cd59..35fff80c61b98e7603d3b7b5df3cabdb59059a72 100644 --- a/tensorflow/python/framework/tensor_util_test.py +++ b/tensorflow/python/framework/tensor_util_test.py @@ -199,6 +199,25 @@ class TensorUtilTest(test.TestCase): dtype=nptype), a) + def testFloatMutateArray(self): + t = tensor_util.make_tensor_proto([10.0, 20.0, 30.0], dtype=dtypes.float32) + a = tensor_util.MakeNdarray(t) + a[0] = 5.0 + self.assertEquals(np.float32, a.dtype) + self.assertAllClose(np.array([5.0, 20.0, 30.0], dtype=np.float32), a) + if sys.byteorder == "big": + self.assertProtoEquals(""" + dtype: DT_FLOAT + tensor_shape { dim { size: 3 } } + tensor_content: "A \000\000A\240\000\000A\360\000\000" + """, t) + else: + self.assertProtoEquals(""" + dtype: DT_FLOAT + tensor_shape { dim { size: 3 } } + tensor_content: "\000\000 A\000\000\240A\000\000\360A" + """, t) + def testHalf(self): t = tensor_util.make_tensor_proto(np.array([10.0, 20.0], dtype=np.float16)) self.assertProtoEquals(""" @@ -216,6 +235,26 @@ class TensorUtilTest(test.TestCase): self.assertEquals(np.float16, a.dtype) self.assertAllClose(np.array([10.0, 20.0], dtype=np.float16), a) + def testBfloat16(self): + test_type = dtypes.bfloat16.as_numpy_dtype + t = tensor_util.make_tensor_proto(np.array([10.0, 20.0], dtype=test_type)) + # 10.0: 16672 = 010000010(130) 0100000: (1+0/2+1/4) * 2^(130-127) + # 20.0: 16800 = 010000011(131) 0100000: (1+0/2+1/4) * 2^(131-127) + self.assertProtoEquals(""" + dtype: DT_BFLOAT16 + tensor_shape { + dim { + size: 2 + } + } + half_val: 16672 + half_val: 16800 + """, t) + + a = tensor_util.MakeNdarray(t) + self.assertEquals(test_type, a.dtype) + self.assertAllClose(np.array([10.0, 20.0], dtype=test_type), a) + def testInt(self): t = tensor_util.make_tensor_proto(10) self.assertProtoEquals(""" @@ -749,7 +788,7 @@ class ConstantValueTest(test.TestCase): self.assertAllClose(np_val, tensor_util.constant_value(tf_val)) def testUnknown(self): - tf_val = gen_state_ops._variable( + tf_val = gen_state_ops.variable( shape=[3, 4, 7], dtype=dtypes.float32, name="tf_val", diff --git a/tensorflow/python/framework/test_ops.cc b/tensorflow/python/framework/test_ops.cc index c6c6c2233c9a81467f57abe2d42f0df9b7ce7106..070b5ac11f563443a97b304ddcdaabd2f4338445 100644 --- a/tensorflow/python/framework/test_ops.cc +++ b/tensorflow/python/framework/test_ops.cc @@ -76,6 +76,11 @@ REGISTER_OP("TestStringOutput") .Output("output2: string") .SetShapeFn(shape_inference::UnknownShape); +REGISTER_OP("TestAttr") + .Output("out: T") + .Attr("T: {float, double}") + .SetShapeFn(shape_inference::UnknownShape); + namespace { enum KernelLabel { DEFAULT_LABEL, OVERLOAD_1_LABEL, OVERLOAD_2_LABEL }; } // namespace @@ -188,6 +193,20 @@ class ResourceUsingOp : public OpKernel { REGISTER_KERNEL_BUILDER(Name("ResourceUsingOp").Device(DEVICE_CPU), ResourceUsingOp); +class TestAttrOp : public OpKernel { + public: + explicit TestAttrOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} + + void Compute(OpKernelContext* ctx) override { + Tensor* output; + OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &output)); + output->scalar()() = 1.0; + } +}; + +REGISTER_KERNEL_BUILDER( + Name("TestAttr").Device(DEVICE_CPU).TypeConstraint("T"), TestAttrOp); + // Various test ops without kernels. These are used to test graph construction. REGISTER_OP("A") diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py index 70f6a2acbafa43c57f1bab9ae36c2816ab4c78b5..e9e86e452b19fd23b4eddfb4f6a985458b503eff 100644 --- a/tensorflow/python/framework/test_util.py +++ b/tensorflow/python/framework/test_util.py @@ -49,13 +49,14 @@ from tensorflow.python.client import device_lib from tensorflow.python.client import session from tensorflow.python.eager import backprop from tensorflow.python.eager import context -from tensorflow.python.eager import tape +from tensorflow.python.eager import tape # pylint: disable=unused-import from tensorflow.python.framework import device as pydev from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import importer from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed +from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import versions from tensorflow.python.ops import array_ops from tensorflow.python.ops import resource_variable_ops @@ -200,11 +201,14 @@ def _strip_checkpoint_v2_randomized(graph_def): def IsGoogleCudaEnabled(): return pywrap_tensorflow.IsGoogleCudaEnabled() - def CudaSupportsHalfMatMulAndConv(): return pywrap_tensorflow.CudaSupportsHalfMatMulAndConv() +def IsMklEnabled(): + return pywrap_tensorflow.IsMklEnabled() + + def InstallStackTraceHandler(): pywrap_tensorflow.InstallStacktraceHandler() @@ -403,6 +407,31 @@ def enable_c_api(fn): return wrapper +def enable_c_shapes(fn): + """Decorator for enabling C shapes on a test. + + Note this enables the C shapes after running the test class's setup/teardown + methods. + + Args: + fn: the function to be wrapped + + Returns: + The wrapped function + """ + + def wrapper(*args, **kwargs): + prev_value = ops._USE_C_SHAPES + # Only use C shapes if the C API is already enabled. + ops._USE_C_SHAPES = ops._USE_C_API + try: + fn(*args, **kwargs) + finally: + ops._USE_C_SHAPES = prev_value + + return wrapper + + # This decorator is a hacky way to run all the test methods in a decorated # class with and without C API enabled. # TODO(iga): Remove this and its uses once we switch to using C API by default. @@ -419,71 +448,42 @@ def with_c_api(cls): Returns: cls with new test methods added """ + # If the C API is already enabled, don't do anything. Some tests break if the + # same test is run twice, so this allows us to turn on the C API by default + # without breaking these tests. + if ops._USE_C_API: return cls + for name, value in cls.__dict__.copy().items(): if callable(value) and name.startswith("test"): setattr(cls, name + "WithCApi", enable_c_api(value)) return cls -class IsolateTest(object): - """A context manager which isolates resources in its block. - - Provides an Eager-agnostic abstraction for preventing the sharing of - variables and other resources. - - In graph mode, resource handle ops are only executed in a particular Session, - isolating them from resources with the same name in other Graphs. In Eager, - separate Sessions do not exist, so resources (particularly ResourceVariables) - would be shared implicitly if a resource of the same name were created - anywhere in a Python process. Multiple handles to the same resource would - cause several issues, and so this type of sharing will raise an exception. +def assert_no_new_pyobjects_executing_eagerly(f): + """Decorator for asserting that no new Python objects persist after a test. - Using resources with the same name in a single Python process may be useful - (especially for unit tests), so this context manager provides an abstraction - for isolating resources. Using a resource created in one Isolation environment - in another is an error. + Runs the test multiple times executing eagerly, first as a warmup and then + several times to let objects accumulate. The warmup helps ignore caches which + do not grow as the test is run repeatedly. - Example usage in Eager mode: - - ```python - import tensorflow as tf - # Import subject to change - from tensorflow.contrib.eager.python import tfe - - tfe.enable_eager_execution() - - for hyperparameter in [1, 2, 3]: - with tfe.IsolateTest(): - v = tfe.Variable(name="v", initial_value=hyperparameter) - # train model, test results ... - ``` - - IsolateTest is currently exposed through contrib.eager, but it creates a new - default Graph and provides equivalent safety in graph mode. + Useful for checking that there are no missing Py_DECREFs in the C exercised by + a bit of Python. """ - - def __init__(self): - if context.in_eager_mode() and tape.could_possibly_record(): - raise ValueError("Cannot isolate Eager execution with an active tape.") - # In Eager, Graphs set a container which isolates resources, and maintain a - # VariableStore which caches ResourceVariable objects created through - # get_variable. So setting the default Graph has the side effect of - # isolating Eager resources. + def decorator(self, **kwargs): + """Warms up, gets an object count, runs the test, checks for new objects.""" with context.eager_mode(): - # Create the graph in Eager mode, as this provides stricter semantics - # (i.e. has a unique container prefix). This prevents implicit sharing - # when a Graph-mode graph is created and then Eager mode is enabled (an - # error through enable_eager_execution, but common with context managers - # in unit tests). - self._graph_as_default_context_manager = ops.Graph().as_default() - - def __enter__(self): - self._graph_as_default_context_manager.__enter__() - - def __exit__(self, type_arg, value_arg, traceback_arg): - return self._graph_as_default_context_manager.__exit__( - type_arg, value_arg, traceback_arg) - + gc.disable() + f(self, **kwargs) + gc.collect() + previous_count = len(gc.get_objects()) + for _ in range(3): + f(self, **kwargs) + gc.collect() + # There should be no new Python objects hanging around. + new_count = len(gc.get_objects()) + self.assertEqual(previous_count, new_count) + gc.enable() + return decorator def assert_no_new_tensors(f): """Decorator for asserting that no new Tensors persist after a test. @@ -506,31 +506,34 @@ def assert_no_new_tensors(f): def decorator(self, **kwargs): """Finds existing Tensors, runs the test, checks for new Tensors.""" - def _is_tensor(obj): + def _is_tensorflow_object(obj): try: - return (isinstance(obj, ops.Tensor) or - isinstance(obj, variables.Variable)) + return isinstance(obj, ( + ops.Tensor, + variables.Variable, + tensor_shape.Dimension, + tensor_shape.TensorShape)) except ReferenceError: # If the object no longer exists, we don't care about it. return False - tensors_before = set(id(obj) for obj in gc.get_objects() if _is_tensor(obj)) - outside_container_prefix = ops.get_default_graph()._container_prefix - with IsolateTest(): + tensors_before = set(id(obj) for obj in gc.get_objects() + if _is_tensorflow_object(obj)) + outside_graph_key = ops.get_default_graph()._graph_key + with ops.Graph().as_default(): # Run the test in a new graph so that collections get cleared when it's - # done, but inherit the container prefix so that we can print the values - # of variables which get leaked when executing eagerly. - ops.get_default_graph()._container_prefix = outside_container_prefix + # done, but inherit the graph key so optimizers behave. + ops.get_default_graph()._graph_key = outside_graph_key f(self, **kwargs) # Make an effort to clear caches, which would otherwise look like leaked # Tensors. - backprop._last_zero = [None] - backprop._shape_dtype = [None, None] + backprop._zeros_cache.flush() + context.get_default_context().ones_rank_cache().flush() context.get_default_context().scalar_cache().clear() gc.collect() tensors_after = [ obj for obj in gc.get_objects() - if _is_tensor(obj) and id(obj) not in tensors_before + if _is_tensorflow_object(obj) and id(obj) not in tensors_before ] if tensors_after: raise AssertionError(("%d Tensors not deallocated after test: %s" % ( @@ -563,6 +566,30 @@ def assert_no_garbage_created(f): previous_garbage = len(gc.garbage) f(self, **kwargs) gc.collect() + if len(gc.garbage) > previous_garbage: + logging.error( + "The decorated test created work for Python's garbage collector, " + "likely due to a reference cycle. New objects in cycle(s):") + for i, obj in enumerate(gc.garbage[previous_garbage:]): + try: + logging.error( + "Object %d of %d" % (i, len(gc.garbage) - previous_garbage)) + def _safe_object_str(obj): + return "<%s %d>" % (obj.__class__.__name__, id(obj)) + logging.error(" Object type: %s" % (_safe_object_str(obj),)) + logging.error(" Referrer types: %s" % ( + ', '.join([_safe_object_str(ref) + for ref in gc.get_referrers(obj)]),)) + logging.error(" Referent types: %s" % ( + ', '.join([_safe_object_str(ref) + for ref in gc.get_referents(obj)]),)) + logging.error(" Object attribute names: %s" % (dir(obj),)) + logging.error(" Object __str__:") + logging.error(obj) + logging.error(" Object __repr__:") + logging.error(repr(obj)) + except Exception: + logging.error("(Exception while printing object)") # This will fail if any garbage has been created, typically because of a # reference cycle. self.assertEqual(previous_garbage, len(gc.garbage)) @@ -621,6 +648,7 @@ def run_in_graph_and_eager_modes(__unused__=None, # This decorator runs the wrapped test twice. # Reset the test environment between runs. self.tearDown() + self._tempdir = None self.setUp() def run_eager_mode(self, **kwargs): @@ -642,7 +670,7 @@ def run_in_graph_and_eager_modes(__unused__=None, assert_no_garbage_created(run_eager_mode)) with context.eager_mode(): - with IsolateTest(): + with ops.Graph().as_default(): run_eager_mode(self, **kwargs) return decorated @@ -753,7 +781,7 @@ class TensorFlowTestCase(googletest.TestCase): self._tempdir = tempfile.mkdtemp(dir=googletest.GetTempDir()) return self._tempdir - def _AssertProtoEquals(self, a, b): + def _AssertProtoEquals(self, a, b, msg=None): """Asserts that a and b are the same proto. Uses ProtoEq() first, as it returns correct results @@ -763,11 +791,12 @@ class TensorFlowTestCase(googletest.TestCase): Args: a: a proto. b: another proto. + msg: Optional message to report on failure. """ if not compare.ProtoEq(a, b): - compare.assertProtoEqual(self, a, b, normalize_numbers=True) + compare.assertProtoEqual(self, a, b, normalize_numbers=True, msg=msg) - def assertProtoEquals(self, expected_message_maybe_ascii, message): + def assertProtoEquals(self, expected_message_maybe_ascii, message, msg=None): """Asserts that message is same as parsed expected_message_ascii. Creates another prototype of message, reads the ascii message into it and @@ -776,8 +805,9 @@ class TensorFlowTestCase(googletest.TestCase): Args: expected_message_maybe_ascii: proto message in original or ascii form. message: the message to validate. + msg: Optional message to report on failure. """ - + msg = msg if msg else "" if isinstance(expected_message_maybe_ascii, type(message)): expected_message = expected_message_maybe_ascii self._AssertProtoEquals(expected_message, message) @@ -787,20 +817,21 @@ class TensorFlowTestCase(googletest.TestCase): expected_message_maybe_ascii, expected_message, descriptor_pool=descriptor_pool.Default()) - self._AssertProtoEquals(expected_message, message) + self._AssertProtoEquals(expected_message, message, msg=msg) else: - assert False, ("Can't compare protos of type %s and %s" % - (type(expected_message_maybe_ascii), type(message))) + assert False, ("Can't compare protos of type %s and %s. %s" % + (type(expected_message_maybe_ascii), type(message), msg)) def assertProtoEqualsVersion( self, expected, actual, producer=versions.GRAPH_DEF_VERSION, - min_consumer=versions.GRAPH_DEF_VERSION_MIN_CONSUMER): + min_consumer=versions.GRAPH_DEF_VERSION_MIN_CONSUMER, + msg=None): expected = "versions { producer: %d min_consumer: %d };\n%s" % ( producer, min_consumer, expected) - self.assertProtoEquals(expected, actual) + self.assertProtoEquals(expected, actual, msg=msg) def assertStartsWith(self, actual, expected_start, msg=None): """Assert that actual.startswith(expected_start) is True. @@ -841,7 +872,7 @@ class TensorFlowTestCase(googletest.TestCase): Returns: tensors numpy values. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return self._eval_helper(tensors) else: sess = ops.get_default_session() @@ -1090,7 +1121,7 @@ class TensorFlowTestCase(googletest.TestCase): "%f != %f +/- %f%s" % (f1, f2, err, " (%s)" % msg if msg is not None else "")) - def assertArrayNear(self, farray1, farray2, err): + def assertArrayNear(self, farray1, farray2, err, msg=None): """Asserts that two float arrays are near each other. Checks that for all elements of farray1 and farray2 @@ -1100,23 +1131,25 @@ class TensorFlowTestCase(googletest.TestCase): farray1: a list of float values. farray2: a list of float values. err: a float value. + msg: Optional message to report on failure. """ - self.assertEqual(len(farray1), len(farray2)) + self.assertEqual(len(farray1), len(farray2), msg=msg) for f1, f2 in zip(farray1, farray2): - self.assertNear(float(f1), float(f2), err) + self.assertNear(float(f1), float(f2), err, msg=msg) def _NDArrayNear(self, ndarray1, ndarray2, err): return np.linalg.norm(ndarray1 - ndarray2) < err - def assertNDArrayNear(self, ndarray1, ndarray2, err): + def assertNDArrayNear(self, ndarray1, ndarray2, err, msg=None): """Asserts that two numpy arrays have near values. Args: ndarray1: a numpy ndarray. ndarray2: a numpy ndarray. err: a float. The maximum absolute difference allowed. + msg: Optional message to report on failure. """ - self.assertTrue(self._NDArrayNear(ndarray1, ndarray2, err)) + self.assertTrue(self._NDArrayNear(ndarray1, ndarray2, err), msg=msg) def _GetNdArray(self, a): if not isinstance(a, np.ndarray): @@ -1158,9 +1191,16 @@ class TensorFlowTestCase(googletest.TestCase): np.testing.assert_allclose( a, b, rtol=rtol, atol=atol, err_msg=msg, equal_nan=True) - def _assertAllCloseRecursive(self, a, b, rtol=1e-6, atol=1e-6, path=None): + def _assertAllCloseRecursive(self, + a, + b, + rtol=1e-6, + atol=1e-6, + path=None, + msg=None): path = path or [] path_str = (("[" + "][".join([str(p) for p in path]) + "]") if path else "") + msg = msg if msg else "" # Check if a and/or b are namedtuples. if hasattr(a, "_asdict"): @@ -1169,18 +1209,18 @@ class TensorFlowTestCase(googletest.TestCase): b = b._asdict() a_is_dict = isinstance(a, dict) if a_is_dict != isinstance(b, dict): - raise ValueError("Can't compare dict to non-dict, a%s vs b%s." % - (path_str, path_str)) + raise ValueError("Can't compare dict to non-dict, a%s vs b%s. %s" % + (path_str, path_str, msg)) if a_is_dict: self.assertItemsEqual( a.keys(), b.keys(), - msg="mismatched keys: a%s has keys %s, but b%s has keys %s" % - (path_str, a.keys(), path_str, b.keys())) + msg="mismatched keys: a%s has keys %s, but b%s has keys %s. %s" % + (path_str, a.keys(), path_str, b.keys(), msg)) for k in a: path.append(k) self._assertAllCloseRecursive( - a[k], b[k], rtol=rtol, atol=atol, path=path) + a[k], b[k], rtol=rtol, atol=atol, path=path, msg=msg) del path[-1] elif isinstance(a, (list, tuple)): # Try to directly compare a, b as ndarrays; if not work, then traverse @@ -1193,29 +1233,35 @@ class TensorFlowTestCase(googletest.TestCase): b_as_ndarray, rtol=rtol, atol=atol, - msg="Mismatched value: a%s is different from b%s." % (path_str, - path_str)) + msg="Mismatched value: a%s is different from b%s. %s" % + (path_str, path_str, msg)) except (ValueError, TypeError) as e: if len(a) != len(b): raise ValueError( - "Mismatched length: a%s has %d items, but b%s has %d items" % - (path_str, len(a), path_str, len(b))) + "Mismatched length: a%s has %d items, but b%s has %d items. %s" % + (path_str, len(a), path_str, len(b), msg)) for idx, (a_ele, b_ele) in enumerate(zip(a, b)): path.append(str(idx)) self._assertAllCloseRecursive( - a_ele, b_ele, rtol=rtol, atol=atol, path=path) + a_ele, b_ele, rtol=rtol, atol=atol, path=path, msg=msg) del path[-1] # a and b are ndarray like objects else: - self._assertArrayLikeAllClose( - a, - b, - rtol=rtol, - atol=atol, - msg="Mismatched value: a%s is different from b%s." % (path_str, - path_str)) - - def assertAllClose(self, a, b, rtol=1e-6, atol=1e-6): + try: + self._assertArrayLikeAllClose( + a, + b, + rtol=rtol, + atol=atol, + msg="Mismatched value: a%s is different from b%s." % (path_str, + path_str)) + except TypeError as e: + msg = "Error: a%s has %s, but b%s has %s" % ( + path_str, type(a), path_str, type(b)) + e.args = ((e.args[0] + ' : ' + msg,) + e.args[1:]) + raise + + def assertAllClose(self, a, b, rtol=1e-6, atol=1e-6, msg=None): """Asserts that two structures of numpy arrays, have near values. `a` and `b` can be arbitrarily nested structures. A layer of a nested @@ -1228,6 +1274,7 @@ class TensorFlowTestCase(googletest.TestCase): numpy `ndarray`, or any arbitrarily nested of structure of these. rtol: relative tolerance. atol: absolute tolerance. + msg: Optional message to report on failure. Raises: ValueError: if only one of `a[p]` and `b[p]` is a dict or @@ -1235,7 +1282,7 @@ class TensorFlowTestCase(googletest.TestCase): to the nested structure, e.g. given `a = [(1, 1), {'d': (6, 7)}]` and `[p] = [1]['d']`, then `a[p] = (6, 7)`. """ - self._assertAllCloseRecursive(a, b, rtol=rtol, atol=atol) + self._assertAllCloseRecursive(a, b, rtol=rtol, atol=atol, msg=msg) def assertAllCloseAccordingToType(self, a, @@ -1247,7 +1294,8 @@ class TensorFlowTestCase(googletest.TestCase): half_rtol=1e-3, half_atol=1e-3, bfloat16_rtol=1e-2, - bfloat16_atol=1e-2): + bfloat16_atol=1e-2, + msg=None): """Like assertAllClose, but also suitable for comparing fp16 arrays. In particular, the tolerance is reduced to 1e-3 if at least @@ -1264,6 +1312,7 @@ class TensorFlowTestCase(googletest.TestCase): half_atol: absolute tolerance for float16. bfloat16_rtol: relative tolerance for bfloat16. bfloat16_atol: absolute tolerance for bfloat16. + msg: Optional message to report on failure. """ a = self._GetNdArray(a) b = self._GetNdArray(b) @@ -1280,19 +1329,21 @@ class TensorFlowTestCase(googletest.TestCase): rtol = max(rtol, bfloat16_rtol) atol = max(atol, bfloat16_atol) - self.assertAllClose(a, b, rtol=rtol, atol=atol) + self.assertAllClose(a, b, rtol=rtol, atol=atol, msg=msg) - def assertAllEqual(self, a, b): + def assertAllEqual(self, a, b, msg=None): """Asserts that two numpy arrays have the same values. Args: a: the expected numpy ndarray or anything can be converted to one. b: the actual numpy ndarray or anything can be converted to one. + msg: Optional message to report on failure. """ + msg = msg if msg else "" a = self._GetNdArray(a) b = self._GetNdArray(b) - self.assertEqual(a.shape, b.shape, "Shape mismatch: expected %s, got %s." % - (a.shape, b.shape)) + self.assertEqual(a.shape, b.shape, "Shape mismatch: expected %s, got %s." + " %s" % (a.shape, b.shape, msg)) same = (a == b) if a.dtype == np.float32 or a.dtype == np.float64: @@ -1309,7 +1360,7 @@ class TensorFlowTestCase(googletest.TestCase): x, y = a, b print("not equal lhs = ", x) print("not equal rhs = ", y) - np.testing.assert_array_equal(a, b) + np.testing.assert_array_equal(a, b, err_msg=msg) # pylint: disable=g-doc-return-or-yield @contextlib.contextmanager @@ -1359,12 +1410,13 @@ class TensorFlowTestCase(googletest.TestCase): return self.assertRaisesWithPredicateMatch(errors.OpError, expected_err_re_or_predicate) - def assertShapeEqual(self, np_array, tf_tensor): + def assertShapeEqual(self, np_array, tf_tensor, msg=None): """Asserts that a Numpy ndarray and a TensorFlow tensor have the same shape. Args: np_array: A Numpy ndarray or Numpy scalar. tf_tensor: A Tensor. + msg: Optional message to report on failure. Raises: TypeError: If the arguments have the wrong type. @@ -1373,19 +1425,22 @@ class TensorFlowTestCase(googletest.TestCase): raise TypeError("np_array must be a Numpy ndarray or Numpy scalar") if not isinstance(tf_tensor, ops.Tensor): raise TypeError("tf_tensor must be a Tensor") - self.assertAllEqual(np_array.shape, tf_tensor.get_shape().as_list()) + self.assertAllEqual( + np_array.shape, tf_tensor.get_shape().as_list(), msg=msg) - def assertDeviceEqual(self, device1, device2): + def assertDeviceEqual(self, device1, device2, msg=None): """Asserts that the two given devices are the same. Args: device1: A string device name or TensorFlow `DeviceSpec` object. device2: A string device name or TensorFlow `DeviceSpec` object. + msg: Optional message to report on failure. """ device1 = pydev.canonical_name(device1) device2 = pydev.canonical_name(device2) self.assertEqual(device1, device2, - "Devices %s and %s are not equal" % (device1, device2)) + "Devices %s and %s are not equal. %s" % + (device1, device2, msg)) # Fix Python 3 compatibility issues if six.PY3: diff --git a/tensorflow/python/framework/test_util_test.py b/tensorflow/python/framework/test_util_test.py index 3594d125bf616917727bea4958eaabf159d0aee0..02ffa93baee5c643ebdceaa274710f9d58e6eecb 100644 --- a/tensorflow/python/framework/test_util_test.py +++ b/tensorflow/python/framework/test_util_test.py @@ -29,7 +29,6 @@ from google.protobuf import text_format from tensorflow.core.framework import graph_pb2 from tensorflow.core.protobuf import meta_graph_pb2 -from tensorflow.python.client import session from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import errors @@ -39,7 +38,6 @@ from tensorflow.python.framework import test_util from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import resource_variable_ops -from tensorflow.python.ops import variables from tensorflow.python.platform import googletest @@ -84,6 +82,14 @@ class TestUtilTest(test_util.TensorFlowTestCase): else: print("GoogleCuda is disabled") + def testIsMklEnabled(self): + # This test doesn't assert anything. + # It ensures the py wrapper function is generated correctly. + if test_util.IsMklEnabled(): + print("MKL is enabled") + else: + print("MKL is disabled") + def testAssertProtoEqualsStr(self): graph_str = "node { name: 'w1' op: 'params' }" @@ -442,72 +448,26 @@ class GarbageCollectionTest(test_util.TensorFlowTestCase): LeakedTensorTest().test_has_no_leak() + def test_no_new_objects_decorator(self): -@test_util.with_c_api -class IsolationTest(test_util.TensorFlowTestCase): + class LeakedObjectTest(object): - @test_util.run_in_graph_and_eager_modes() - def test_variable_reuse_exception(self): - with test_util.IsolateTest(), session.Session(): - first_container_variable = resource_variable_ops.ResourceVariable( - name="first_container_variable", - initial_value=1) - if context.in_graph_mode(): - self.evaluate([variables.global_variables_initializer()]) - with test_util.IsolateTest(): - if context.in_graph_mode(): - with self.assertRaises(RuntimeError): - self.evaluate(first_container_variable.read_value()) - else: - with self.assertRaises(ValueError): - first_container_variable.read_value() + def __init__(inner_self): # pylint: disable=no-self-argument + inner_self.assertEqual = self.assertEqual # pylint: disable=invalid-name + inner_self.accumulation = [] - @test_util.run_in_graph_and_eager_modes() - def test_variable_reuse_exception_nested(self): - with test_util.IsolateTest(), session.Session(): - first_container_variable = resource_variable_ops.ResourceVariable( - name="first_container_variable", - initial_value=1) - if context.in_graph_mode(): - self.evaluate([variables.global_variables_initializer()]) - with test_util.IsolateTest(), session.Session(): - if context.in_graph_mode(): - with self.assertRaises(RuntimeError): - self.evaluate(first_container_variable.read_value()) - else: - with self.assertRaises(ValueError): - first_container_variable.read_value() + @test_util.assert_no_new_pyobjects_executing_eagerly + def test_has_leak(self): + self.accumulation.append([1.]) - @test_util.run_in_graph_and_eager_modes() - def test_no_sharing(self): - with test_util.IsolateTest(), session.Session(): - first_container_variable = resource_variable_ops.ResourceVariable( - name="same_name", - initial_value=1) - if context.in_graph_mode(): - self.evaluate([variables.global_variables_initializer()]) - with test_util.IsolateTest(), session.Session(): - second_container_variable = resource_variable_ops.ResourceVariable( - name="same_name", - initial_value=2) - if context.in_graph_mode(): - self.evaluate([variables.global_variables_initializer()]) - self.assertEqual( - 2, self.evaluate(second_container_variable.read_value())) - self.assertEqual(1, self.evaluate(first_container_variable.read_value())) - - def test_graph_mode_isolation(self): - with context.graph_mode(): - # Even if we've (accidentally) called IsolateTest in Graph mode, it should - # provide Eager isolation. - with test_util.IsolateTest(): - with context.eager_mode(): - first_container_variable = resource_variable_ops.ResourceVariable( - name="first_container_variable", - initial_value=1) - with context.eager_mode(): - with self.assertRaises(ValueError): - first_container_variable.read_value() + @test_util.assert_no_new_pyobjects_executing_eagerly + def test_has_no_leak(self): + self.not_accumulating = [1.] + + with self.assertRaises(AssertionError): + LeakedObjectTest().test_has_leak() + + LeakedObjectTest().test_has_no_leak() if __name__ == "__main__": googletest.main() diff --git a/tensorflow/python/framework/versions.py b/tensorflow/python/framework/versions.py index bdcbc15af63c57d712abfac97537f86b3bbe1737..06955b885852a641bc814f88c99838effe03bfd4 100644 --- a/tensorflow/python/framework/versions.py +++ b/tensorflow/python/framework/versions.py @@ -35,7 +35,9 @@ tf_export("GIT_VERSION").export_constant(__name__, "GIT_VERSION") COMPILER_VERSION = __compiler_version__ tf_export("COMPILER_VERSION").export_constant(__name__, "COMPILER_VERSION") CXX11_ABI_FLAG = __cxx11_abi_flag__ +tf_export("CXX11_ABI_FLAG").export_constant(__name__, "CXX11_ABI_FLAG") MONOLITHIC_BUILD = __monolithic_build__ +tf_export("MONOLITHIC_BUILD").export_constant(__name__, "MONOLITHIC_BUILD") GRAPH_DEF_VERSION = pywrap_tensorflow.GRAPH_DEF_VERSION tf_export("GRAPH_DEF_VERSION").export_constant(__name__, "GRAPH_DEF_VERSION") diff --git a/tensorflow/python/grappler/cluster.i b/tensorflow/python/grappler/cluster.i index 0c8d04ff29518d587079a76e3fee3b2e327c6c5c..067c8213d4741936e4c28aaedf4f30639b8cdc41 100644 --- a/tensorflow/python/grappler/cluster.i +++ b/tensorflow/python/grappler/cluster.i @@ -140,6 +140,7 @@ static GCluster TF_NewCluster(bool allow_soft_placement, timeout_s, num_cpu_cores, num_gpus); cluster_->DisableDetailedStats(disable_detailed_stats); cluster_->AllowSoftPlacement(allow_soft_placement); + cluster_->SetNumWarmupSteps(10); tensorflow::Status status = cluster_->Provision(); tensorflow::Set_TF_Status_from_Status(out_status, status); return GCluster(cluster_); @@ -205,7 +206,7 @@ static PyObject* TF_ListDevices(GCluster cluster) { return result; } -static std::vector TF_ListAvailableOps() { +static PyObject* TF_ListAvailableOps() { tensorflow::OpRegistry* registry = tensorflow::OpRegistry::Global(); std::vector ops; registry->GetRegisteredOps(&ops); @@ -214,7 +215,14 @@ static std::vector TF_ListAvailableOps() { op_names.push_back(op.name()); } std::sort(op_names.begin(), op_names.end()); - return op_names; + + PyGILState_STATE gstate = PyGILState_Ensure(); + PyObject* result = PyList_New(op_names.size()); + for (int i = 0; i < op_names.size(); ++i) { + PyList_SetItem(result, i, PyString_FromString(op_names[i].c_str())); + } + PyGILState_Release(gstate); + return result; } static PyObject* TF_GetSupportedDevices(GCluster cluster, GItem item) { @@ -431,7 +439,7 @@ static GCluster TF_NewVirtualCluster( TF_Status* out_status); static void TF_ShutdownCluster(GCluster cluster); static PyObject* TF_ListDevices(GCluster cluster); -static std::vector TF_ListAvailableOps(); +static PyObject* TF_ListAvailableOps(); static PyObject* TF_GetSupportedDevices(GCluster cluster, GItem item); static float TF_EstimatePerformance(const tensorflow::NamedDevice& device); static PyObject* TF_MeasureCosts( diff --git a/tensorflow/python/grappler/cluster_test.py b/tensorflow/python/grappler/cluster_test.py index 2292b2c732b2d5d0d40b44d8ca831f4e72b057c6..a3c4c2bbeba7c4ee5d00268c0e475e11a31fa7eb 100644 --- a/tensorflow/python/grappler/cluster_test.py +++ b/tensorflow/python/grappler/cluster_test.py @@ -45,7 +45,7 @@ class ClusterTest(test.TestCase): op_perfs, run_time, step_stats = grappler_cluster.MeasureCosts( grappler_item) self.assertTrue(run_time > 0) - self.assertEqual(len(op_perfs), 9) + self.assertEqual(len(op_perfs), 8) self.assertTrue(step_stats.dev_stats) def testNoDetailedStats(self): @@ -125,14 +125,14 @@ class ClusterTest(test.TestCase): disable_detailed_stats=False, disable_timeline=False) as gcluster: op_perfs, run_time, step_stats = gcluster.MeasureCosts(grappler_item) self.assertTrue(run_time > 0) - self.assertEqual(len(op_perfs), 9) + self.assertEqual(len(op_perfs), 8) self.assertTrue(step_stats.dev_stats) def testAvailableOps(self): with cluster.Provision() as gcluster: op_names = gcluster.ListAvailableOps() - self.assertTrue(b'Add' in op_names) - self.assertTrue(b'MatMul' in op_names) + self.assertTrue('Add' in op_names) + self.assertTrue('MatMul' in op_names) self.assertEqual(op_names, sorted(op_names)) def testSupportDevices(self): diff --git a/tensorflow/python/grappler/controller.py b/tensorflow/python/grappler/controller.py new file mode 100644 index 0000000000000000000000000000000000000000..5677f4f52310dd68dc80c87275b50be95ba86b60 --- /dev/null +++ b/tensorflow/python/grappler/controller.py @@ -0,0 +1,142 @@ +# 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. +# ============================================================================== +"""Controller Class.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from collections import defaultdict + + +class Controller(object): + """Controller class.""" + + def __init__(self, item, cluster): + """Controller class initializer. + + Args: + item: The metagraph to place wrapped in a cluster. + cluster: A cluster of devices on which to place the item. + """ + self.item = item + + self._node = {} + for node in item.metagraph.graph_def.node: + self._node[node.name] = node + + self._fanout = defaultdict(lambda: []) + for node in item.metagraph.graph_def.node: + for fanin in self._get_node_fanin(node): + self._fanout[fanin.name].append(node) + + important_op_names = item.IdentifyImportantOps(sort_topologically=True) + + # List of important ops (these are the ops to place) sorted in topological + # order. The order of this collection is deterministic. + self.important_ops = [] + for name in important_op_names: + self.important_ops.append(self._node[name]) + + self.node_properties = item.GetOpProperties() + + self.cluster = cluster + self.devices = cluster.ListDevices() + + self.colocation_constraints = item.GetColocationGroups() + + self.placement_constraints = cluster.GetSupportedDevices(item) + for node_name, dev in self.placement_constraints.items(): + if len(dev) == 1: + # Place the node on the supported device + node = self._node[node_name] + node.device = dev[0] + fanout = self.get_node_fanout(node) + # Update the fanout of the fanin to bypass the node + for fanin in self._get_node_fanin(node): + fanout_of_fanin = self.get_node_fanout(fanin) + fanout_of_fanin += fanout + fanout_of_fanin.remove(node) + # Remove node from the list of important ops since we don't need to + # place the node. + if node in self.important_ops: + self.important_ops.remove(node) + important_op_names.remove(node.name) + + # List of important op names, in non deterministic order. + self.important_op_names = frozenset(important_op_names) + + @property + def input_graph_def(self): + return self.item.metagraph.graph_def + + @property + def num_devices(self): + return len(self.devices) + + def get_node_by_name(self, node_name): + return self._node[node_name] + + def get_node_fanout(self, node): + return self._fanout[node.name] + + def get_placements(self, *args, **kwargs): + """Returns: Two TF ops. + + Args: + *args: "". + **kwargs: "". + + Returns: + y_preds: tensor of size [batch_size, num_ops] + log_probs: python dict of at least two fields: "sample", "target" each + containing a tensor of size [batch_size], corresponding to the log_probs. + """ + raise NotImplementedError + + def eval_placement(self, sess, *args, **kwargs): + """At this time, this method evaluates ONLY ONE placement. + + Args: + sess: a tf.Session() object used to retrieve cached assignment info. + *args: "". + **kwargs: "". + + Returns: + run_time: scalar + """ + raise NotImplementedError + + def export_placement(self, metagraph): + """Annotate the placement onto the specified metagraph. + + Args: + metagraph: the metagraph to annotate with the placement. + """ + for node in metagraph.graph_def.node: + if node.name in self.important_op_names: + node.device = self.get_node_by_name(node.name).device + + # Get the nodes in the immediate fanin of node. + # Beware: this doesn't take into account the nodes that may be skipped + # since placement constraints force their placement. + def _get_node_fanin(self, node): + input_ops = [] + for fanin_name in node.input: + if fanin_name[0] == "^": + fanin_name = fanin_name[1:] + fanin_name = fanin_name.split(":")[0] + input_ops.append(self.get_node_by_name(fanin_name)) + return input_ops diff --git a/tensorflow/python/grappler/cost_analyzer.cc b/tensorflow/python/grappler/cost_analyzer.cc index 88bf900dca6d97773959eb309a4a3c5931fdcb88..b474e19894957d01c7c8978282c547df81a9b2b3 100644 --- a/tensorflow/python/grappler/cost_analyzer.cc +++ b/tensorflow/python/grappler/cost_analyzer.cc @@ -30,11 +30,12 @@ CostAnalyzer::CostAnalyzer(const GrapplerItem& item, Cluster* cluster, analytical_estimator_(cluster, false), suffix_(suffix) {} -Status CostAnalyzer::GenerateReport(std::ostream& os, bool per_node_report) { +Status CostAnalyzer::GenerateReport(std::ostream& os, bool per_node_report, + bool verbose) { GatherCosts(); PreprocessCosts(); AnalyzeCosts(); - PrintAnalysis(os, per_node_report); + PrintAnalysis(os, per_node_report, verbose); return Status::OK(); } @@ -158,7 +159,8 @@ void CostAnalyzer::AnalyzeCosts() { } } -void CostAnalyzer::PrintAnalysis(std::ostream& os, bool per_node_report) const { +void CostAnalyzer::PrintAnalysis(std::ostream& os, bool per_node_report, + bool verbose) const { os << std::endl; os << std::left << std::setw(50) << "Total time measured in ns (serialized): " << std::right @@ -227,10 +229,55 @@ void CostAnalyzer::PrintAnalysis(std::ostream& os, bool per_node_report) const { os << std::endl; if (per_node_report) { - os << "Below is the per-node report:" << std::endl; - os << op_perf_.DebugString(); + if (verbose) { + os << "Below is the full per-node report:" << std::endl; + os << op_perf_.DebugString(); + } else { + os << "Below is the per-node report summary:" << std::endl; + int width = 35; + int width_narrow = 15; + int width_wide = 20; + os << std::setw(width + 1) << "Op,"; + os << std::setw(width_wide + 1) << "Measured time (ns),"; + os << std::setw(width_wide + 1) << "Compute time (ns),"; + os << std::setw(width_wide + 1) << "Memory time (ns),"; + os << std::setw(width_narrow + 2) << "Compute eff,"; + os << std::setw(width_narrow + 2) << "Memory eff,"; + os << " Inputs" << std::endl; + for (int i = 0; i < op_perf_.op_performance_size(); i++) { + const auto& perf = op_perf_.op_performance(i); + string op_name = perf.op().op(); + os << std::setw(width) << op_name << ","; + os << std::setw(width_wide) << perf.compute_cost() << ","; + os << std::setw(width_wide) << perf.compute_time() << ","; + os << std::setw(width_wide) << perf.memory_time() << ","; + os << std::setw(width_narrow) << std::setprecision(2) + << perf.compute_efficiency() * 100 << "%,"; + os << std::setw(width_narrow) << std::setprecision(2) + << perf.memory_efficiency() * 100 << "%,"; + os << " ["; + for (int j = 0; j < perf.op().inputs_size(); j++) { + const auto& shape = perf.op().inputs(j).shape(); + if (shape.dim_size() > 0) { + os << "("; + std::vector dims; + for (int k = 0; k < shape.dim_size(); k++) { + os << shape.dim(k).size(); + if (k < shape.dim_size() - 1) { + os << ", "; + } + } + os << ")"; + if (j < perf.op().inputs_size() - 1) { + os << ", "; + } + } + } + os << "]" << std::endl; + } + os << std::endl; + } } } - } // end namespace grappler } // end namespace tensorflow diff --git a/tensorflow/python/grappler/cost_analyzer.h b/tensorflow/python/grappler/cost_analyzer.h index 0e860e0fee9923510292d3cf1a8069435787476f..b5364aa37ab2fbbeb0a33e6764539cca795f2fa6 100644 --- a/tensorflow/python/grappler/cost_analyzer.h +++ b/tensorflow/python/grappler/cost_analyzer.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "tensorflow/core/framework/cost_graph.pb.h" #include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/grappler/clusters/cluster.h" #include "tensorflow/core/grappler/costs/analytical_cost_estimator.h" #include "tensorflow/core/grappler/costs/cost_estimator.h" @@ -50,7 +51,7 @@ class CostAnalyzer { public: explicit CostAnalyzer(const GrapplerItem& item, Cluster* cluster, const string& suffix); - Status GenerateReport(std::ostream& os, bool per_node_report); + Status GenerateReport(std::ostream& os, bool per_node_report, bool verbose); private: void PredictCosts(CostEstimator* cost_estimator, CostGraphDef* cost_graph, @@ -59,7 +60,8 @@ class CostAnalyzer { void PreprocessCosts(); void AnalyzeCosts(); void SortOpsByTime(std::map ops); - void PrintAnalysis(std::ostream& os, bool per_node_report) const; + void PrintAnalysis(std::ostream& os, bool per_node_report, + bool verbose) const; const GrapplerItem* item_; MeasuringCostEstimator measure_estimator_; diff --git a/tensorflow/python/grappler/cost_analyzer.i b/tensorflow/python/grappler/cost_analyzer.i index 4c0953435ba3fa6423bbc869fcca909d0c2ccb25..8f7fdb47f267bea582e371eb9ea6982b6e9341ad 100644 --- a/tensorflow/python/grappler/cost_analyzer.i +++ b/tensorflow/python/grappler/cost_analyzer.i @@ -44,7 +44,7 @@ limitations under the License. %{ string GenerateCostReport(const tensorflow::MetaGraphDef& metagraph, bool per_node_report, - GCluster cluster) { + bool verbose, GCluster cluster) { tensorflow::grappler::ItemConfig cfg; cfg.apply_optimizations = false; std::unique_ptr item = @@ -57,11 +57,11 @@ string GenerateCostReport(const tensorflow::MetaGraphDef& metagraph, bool per_no tensorflow::grappler::CostAnalyzer analyzer(*item, cluster.get(), suffix); std::stringstream os; - analyzer.GenerateReport(os, per_node_report); + analyzer.GenerateReport(os, per_node_report, verbose); return os.str(); } %} string GenerateCostReport(const tensorflow::MetaGraphDef& metagraph, bool per_node_report, - GCluster cluster); + bool verbose, GCluster cluster); diff --git a/tensorflow/python/grappler/cost_analyzer.py b/tensorflow/python/grappler/cost_analyzer.py index a1ff915c61ba14d9a899d7f6c9a2c49855969b00..6a4690e91ba981706eed0d9fdfae2e64359d0416 100644 --- a/tensorflow/python/grappler/cost_analyzer.py +++ b/tensorflow/python/grappler/cost_analyzer.py @@ -24,7 +24,10 @@ from tensorflow.python.grappler import cluster as gcluster from tensorflow.python.grappler import item as gitem -def GenerateCostReport(metagraph, per_node_report=False, cluster=None): +def GenerateCostReport(metagraph, + per_node_report=False, + verbose=False, + cluster=None): """Analyze the cost of each TensorFlow op and node in the provided metagraph. Args: @@ -32,6 +35,7 @@ def GenerateCostReport(metagraph, per_node_report=False, cluster=None): per_node_report: by default the report contains stats aggregated on a per op type basis, setting per_node_report to True adds results for each individual node to the report. + verbose: Prints out the entire operation proto instead of a summary table. cluster: Analyze the costs using the specified cluster, or the local machine if no cluster was specified. @@ -42,8 +46,9 @@ def GenerateCostReport(metagraph, per_node_report=False, cluster=None): cluster = gcluster.Cluster(disable_detailed_stats=False) with errors.raise_exception_on_not_ok_status(): - ret_from_swig = tf_wrap.GenerateCostReport( - metagraph.SerializeToString(), per_node_report, cluster.tf_cluster) + ret_from_swig = tf_wrap.GenerateCostReport(metagraph.SerializeToString(), + per_node_report, verbose, + cluster.tf_cluster) return ret_from_swig diff --git a/tensorflow/python/grappler/cost_analyzer_test.py b/tensorflow/python/grappler/cost_analyzer_test.py index 511908c79ce47d6849bf97d11bc42f2f1bb13f18..b8225b81a52f1a2ee10663544d54f1c9bd7ee785 100644 --- a/tensorflow/python/grappler/cost_analyzer_test.py +++ b/tensorflow/python/grappler/cost_analyzer_test.py @@ -48,7 +48,7 @@ class CostAnalysisTest(test.TestCase): train_op.append(d) mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph()) - report = cost_analyzer.GenerateCostReport(mg) + report = cost_analyzer.GenerateCostReport(mg, per_node_report=True) # Check the report headers self.assertTrue(b"Total time measured in ns (serialized):" in report) @@ -57,6 +57,26 @@ class CostAnalysisTest(test.TestCase): self.assertTrue(b"Total time analytical in ns (lower bound):" in report) self.assertTrue(b"Overall efficiency (analytical upper/actual):" in report) self.assertTrue(b"Overall efficiency (analytical lower/actual):" in report) + self.assertTrue(b"Below is the per-node report summary:" in report) + + # Also print the report to make it easier to debug + print("{}".format(report)) + + def testVerbose(self): + """Make sure the full report is generated with verbose=True.""" + a = constant_op.constant(10, name="a") + b = constant_op.constant(20, name="b") + c = math_ops.add_n([a, b], name="c") + d = math_ops.add_n([b, c], name="d") + train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP) + train_op.append(d) + mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph()) + + report = cost_analyzer.GenerateCostReport( + mg, per_node_report=True, verbose=True) + + # Check the report headers + self.assertTrue(b"Below is the full per-node report:" in report) # Also print the report to make it easier to debug print("{}".format(report)) diff --git a/tensorflow/python/grappler/cost_analyzer_tool.py b/tensorflow/python/grappler/cost_analyzer_tool.py index 61dc4e2afb833414f875d66bb12b0aa010f9d62e..0853db252406966cec36b63efafec9ec755c7e87 100644 --- a/tensorflow/python/grappler/cost_analyzer_tool.py +++ b/tensorflow/python/grappler/cost_analyzer_tool.py @@ -22,7 +22,7 @@ import argparse import sys from google.protobuf import text_format - +from tensorflow.contrib.fused_conv.ops import gen_fused_conv2d_bias_activation_op # pylint: disable=unused-import from tensorflow.core.framework import graph_pb2 from tensorflow.core.protobuf import meta_graph_pb2 from tensorflow.core.protobuf import rewriter_config_pb2 @@ -35,30 +35,51 @@ from tensorflow.python.platform import gfile from tensorflow.python.training import saver -def main(_): +def get_metagraph(): + """Constructs and returns a MetaGraphDef from the input file.""" if FLAGS.metagraphdef: with gfile.GFile(FLAGS.metagraphdef) as meta_file: metagraph = meta_graph_pb2.MetaGraphDef() - metagraph.ParseFromString(meta_file.read()) + if FLAGS.metagraphdef.endswith(".pbtxt"): + text_format.Merge(meta_file.read(), metagraph) + else: + metagraph.ParseFromString(meta_file.read()) + if FLAGS.fetch is not None: + fetch_collection = meta_graph_pb2.CollectionDef() + for fetch in FLAGS.fetch.split(","): + fetch_collection.node_list.value.append(fetch) + metagraph.collection_def["train_op"].CopyFrom(fetch_collection) else: with gfile.GFile(FLAGS.graphdef) as graph_file: graph_def = graph_pb2.GraphDef() - graph_def.ParseFromString(graph_file.read()) + if FLAGS.graphdef.endswith(".pbtxt"): + text_format.Merge(graph_file.read(), graph_def) + else: + graph_def.ParseFromString(graph_file.read()) importer.import_graph_def(graph_def, name="") graph = ops.get_default_graph() - fetch = graph.get_operation_by_name(FLAGS.fetch) - graph.add_to_collection("train_op", fetch) + for fetch in FLAGS.fetch.split(","): + fetch_op = graph.get_operation_by_name(fetch) + graph.add_to_collection("train_op", fetch_op) metagraph = saver.export_meta_graph( graph_def=graph.as_graph_def(), graph=graph) + return metagraph + +def main(_): + metagraph = get_metagraph() + rewriter_config = rewriter_config_pb2.RewriterConfig() if FLAGS.rewriter_config is not None: - rewriter_config = rewriter_config_pb2.RewriterConfig() text_format.Merge(FLAGS.rewriter_config, rewriter_config) - optimized_graph = tf_optimizer.OptimizeGraph(rewriter_config, metagraph) - metagraph.graph_def.CopyFrom(optimized_graph) + optimized_graph = tf_optimizer.OptimizeGraph(rewriter_config, metagraph) + metagraph.graph_def.CopyFrom(optimized_graph) - report = cost_analyzer.GenerateCostReport(metagraph, FLAGS.per_node_report) + report = cost_analyzer.GenerateCostReport(metagraph, FLAGS.per_node_report, + FLAGS.verbose) print(report) + if FLAGS.memory_report: + report = cost_analyzer.GenerateMemoryReport(metagraph) + print(report) if __name__ == "__main__": @@ -73,16 +94,11 @@ if __name__ == "__main__": type=str, default=None, help="Input .pb GraphDef file path.") - # Consider making flag fetch work together with flag metagraphdef. As some - # MetaGraphDef files don't have collection train_op. parser.add_argument( "--fetch", type=str, default=None, - help= - "The name of the fetch node. This flag is ignored if flag " - "metagraphdef is used." - ) + help="The names of the fetch node delimited by comma.") parser.add_argument( "--rewriter_config", type=str, @@ -98,5 +114,13 @@ if __name__ == "__main__": help="Generate per-node report. By default the report contains stats " "aggregated on a per op type basis, per_node_report adds results " "for each individual node to the report.") + parser.add_argument( + "--memory_report", + action="store_true", + help="Generate memory usage report.") + parser.add_argument( + "--verbose", + action="store_true", + help="Generate verbose reports. By default, succinct reports are used.") FLAGS, unparsed = parser.parse_known_args() app.run(main=main, argv=[sys.argv[0]] + unparsed) diff --git a/tensorflow/python/grappler/graph_placer.py b/tensorflow/python/grappler/graph_placer.py new file mode 100644 index 0000000000000000000000000000000000000000..1cd51df4d962583555e08ae973ab43d15ba01997 --- /dev/null +++ b/tensorflow/python/grappler/graph_placer.py @@ -0,0 +1,120 @@ +# 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. +# ============================================================================== +"""Graph Placer.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import time +from tensorflow.core.protobuf import meta_graph_pb2 +from tensorflow.core.protobuf import rewriter_config_pb2 +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops as tf_ops +from tensorflow.python.grappler import cluster as gcluster +from tensorflow.python.grappler import hierarchical_controller +from tensorflow.python.grappler import item as gitem +from tensorflow.python.grappler import tf_optimizer +from tensorflow.python.training import training + + +def PlaceGraph(metagraph, + cluster=None, + allotted_time=3600, + hparams=None, + verbose=False): + """Place the provided metagraph. + + Args: + metagraph: the metagraph to place. + cluster: an optional set of hardware resource to optimize the placement for. + If none is specified, we'll optimize the placement for the hardware + available on the local machine. + allotted_time: the maximum amount to time in seconds to spend optimizing + the placement. + hparams: hyperparameters used to fine tune the placer. + verbose: prints debug information if True. + + Returns: + The placed metagraph. + """ + if cluster is None: + cluster = gcluster.Cluster() + + # Optimize the metagraph to speedup the placement + rewriter_config = rewriter_config_pb2.RewriterConfig() + rewriter_config.optimizers.append("pruning") + rewriter_config.optimizers.append("constfold") + rewriter_config.optimizers.append("arithmetic") + rewriter_config.optimizers.append("dependency") + rewriter_config.optimizers.append("pruning") + optimized_graph = tf_optimizer.OptimizeGraph( + rewriter_config, metagraph, verbose=verbose, cluster=cluster) + optimized_metagraph = meta_graph_pb2.MetaGraphDef() + optimized_metagraph.CopyFrom(metagraph) + optimized_metagraph.graph_def.CopyFrom(optimized_graph) + + item = gitem.Item(optimized_metagraph) + + # Measure the runtime achievable with the original placement. + try: + _, original_run_time, _ = cluster.MeasureCosts(item) + if verbose: + print("Runtime for original placement: " + str(original_run_time)) + except errors.OpError as e: + if verbose: + print("Original placement isn't feasible: " + str(e)) + original_run_time = hparams.failing_signal + + if hparams is None: + hparams = hierarchical_controller.hierarchical_controller_hparams() + # We run with a single child + hparams.num_children = 1 + + with tf_ops.Graph().as_default(): + # Place all the nodes of the controller on the CPU. We don't want them to + # fight for accelerator memory with the model to optimize. + with tf_ops.device("/device:CPU:0"): + model = hierarchical_controller.HierarchicalController( + hparams, item, cluster) + ops = model.build_controller() + session_creator = training.ChiefSessionCreator() + with training.MonitoredSession(session_creator=session_creator) as sess: + start_time = time.time() + current_time = start_time + while current_time - start_time < allotted_time: + grouping_actions = model.generate_grouping(sess) + input_to_seq2seq = model.create_group_embeddings( + grouping_actions, verbose=verbose) + model.generate_placement(input_to_seq2seq, sess) + try: + run_time = model.eval_placement( + sess, + verbose=verbose) + except errors.OpError as e: + if verbose: + print("Failed to run graph:" + str(e)) + run_time = hparams.failing_signal + updated = model.update_reward(sess, run_time, verbose=verbose) + if updated and run_time < original_run_time: + if verbose: + print("Found better placement, with runtime " + str(run_time)) + model.export_placement(metagraph) + + model.process_reward(sess) + + current_time = time.time() + + return metagraph diff --git a/tensorflow/python/grappler/graph_placer_test.py b/tensorflow/python/grappler/graph_placer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9eabe3cd5437022eb3b98010d0f384cc9f6bac2a --- /dev/null +++ b/tensorflow/python/grappler/graph_placer_test.py @@ -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. +# ============================================================================== +"""Tests the graph placer.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from tensorflow.core.protobuf import device_properties_pb2 +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import meta_graph +from tensorflow.python.framework import ops as tf_ops +from tensorflow.python.grappler import cluster +from tensorflow.python.grappler import graph_placer +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 variable_scope +from tensorflow.python.platform import test + + +class GraphPlacerTest(test.TestCase): + + @staticmethod + def _buildMnist(batch_size=128, + input_size=256, + num_classes=1024, + num_layers=10, + hidden_size=256, + name='mnist'): + g = tf_ops.get_default_graph() + with g.as_default(): + ops = {} + x = random_ops.random_uniform( + [batch_size, input_size], -0.1, 0.1, dtype=dtypes.float32) + for layer_id in range(num_layers): + with variable_scope.variable_scope('layer_{}'.format(layer_id)): + a = input_size if layer_id == 0 else hidden_size + b = hidden_size if layer_id < num_layers - 1 else num_classes + w = variable_scope.get_variable('w', [a, b]) + x = math_ops.matmul(x, w) + x = nn_ops.relu(x) + ops['y_preds'] = math_ops.argmax(x, axis=1) + + train_op = g.get_collection_ref(tf_ops.GraphKeys.TRAIN_OP) + train_op.append(ops['y_preds']) + return g + + @staticmethod + def _buildCluster(num_cpus=1, num_gpus=1): + devices = [] + if num_gpus > 0: + device_properties = device_properties_pb2.DeviceProperties( + type='GPU', + vendor='NVidia', + model='GeForce GTX TITAN X', + frequency=1076, + num_cores=24, + environment={'architecture': '5.2', + 'cuda': '8000', + 'cudnn': '6021'}, + num_registers=65536, + l1_cache_size=24576, + l2_cache_size=3145728, + shared_memory_size_per_multiprocessor=98304, + memory_size=12783648768, + bandwidth=336480000) + for i in range(num_gpus): + devices.append( + device_properties_pb2.NamedDevice( + properties=device_properties, name='/GPU:' + str(i))) + + assert num_cpus > 0 + device_properties = device_properties_pb2.DeviceProperties( + type='CPU', + frequency=2000, + num_cores=4, + l1_cache_size=32768, + l2_cache_size=262144, + l3_cache_size=12582912) + for i in range(num_cpus): + devices.append( + device_properties_pb2.NamedDevice( + properties=device_properties, name='/CPU:' + str(i))) + + return cluster.Cluster(devices=devices) + + def testBasic(self): + """Place a trivial graph.""" + a = constant_op.constant(10, name='a') + b = constant_op.constant(20, name='b') + c = math_ops.add_n([a, b], name='c') + d = math_ops.add_n([b, c], name='d') + train_op = tf_ops.get_collection_ref(tf_ops.GraphKeys.TRAIN_OP) + train_op.append(d) + mg = meta_graph.create_meta_graph_def(graph=tf_ops.get_default_graph()) + + gcluster = cluster.Cluster() + placed_mg = graph_placer.PlaceGraph(mg, allotted_time=15, cluster=gcluster) + + self.assertEqual(4, len(placed_mg.graph_def.node)) + self.assertItemsEqual([node.name for node in placed_mg.graph_def.node], + [node.name for node in mg.graph_def.node]) + + available_devices = [device.name for device in gcluster.ListDevices()] + for node in placed_mg.graph_def.node: + # The constant nodes are optimized away before the placer is run, and + # therefore won't be placed. + self.assertTrue(not node.device or node.device in available_devices) + + def testMNIST(self): + graph = GraphPlacerTest._buildMnist() + mg = meta_graph.create_meta_graph_def(graph=graph) + gcluster = GraphPlacerTest._buildCluster(num_gpus=1) + # Spend 15 seconds trying to optimize the placement of the model. This + # should give us enough time to exercise the code, but not enough to find + # a good placement, so we'll just check for legality. + placed_mg = graph_placer.PlaceGraph(mg, allotted_time=15, cluster=gcluster) + self.assertEqual(len(placed_mg.graph_def.node), len(mg.graph_def.node)) + self.assertItemsEqual([node.name for node in placed_mg.graph_def.node], + [node.name for node in mg.graph_def.node]) + available_devices = [device.name for device in gcluster.ListDevices()] + for node in placed_mg.graph_def.node: + self.assertTrue(not node.device or node.device in available_devices) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/grappler/hierarchical_controller.py b/tensorflow/python/grappler/hierarchical_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..c0866c1069ac7f7e25cbd12cb5a490e2ed5e4bec --- /dev/null +++ b/tensorflow/python/grappler/hierarchical_controller.py @@ -0,0 +1,1117 @@ +# 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. +# ============================================================================== +"""HierarchicalController Class. + +The HierarchicalController encompasses the entire lifecycle of training the +device placement policy, including generating op embeddings, getting groups for +each op, placing those groups and running the predicted placements. + +Different assignment models can inherit from this class. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import math +import numpy as np +import six +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 as tf_ops +from tensorflow.python.grappler.controller import Controller +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import clip_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import embedding_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 nn_ops +from tensorflow.python.ops import random_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.summary import summary +from tensorflow.python.training import adam +from tensorflow.python.training import gradient_descent +from tensorflow.python.training import learning_rate_decay +from tensorflow.python.training import training_util + + +class PlacerParams(object): + """Class to hold a set of placement parameters as name-value pairs. + + A typical usage is as follows: + + ```python + # Create a PlacerParams object specifying names and values of the model + # parameters: + params = PlacerParams(hidden_size=128, decay_steps=50) + + # The parameters are available as attributes of the PlacerParams object: + hparams.hidden_size ==> 128 + hparams.decay_steps ==> 50 + ``` + + """ + + def __init__(self, **kwargs): + """Create an instance of `PlacerParams` from keyword arguments. + + The keyword arguments specify name-values pairs for the parameters. + The parameter types are inferred from the type of the values passed. + + The parameter names are added as attributes of `PlacerParams` object, + and they can be accessed directly with the dot notation `params._name_`. + + Example: + + ```python + # Define 1 parameter: 'hidden_size' + params = PlacerParams(hidden_size=128) + params.hidden_size ==> 128 + ``` + + Args: + **kwargs: Key-value pairs where the key is the parameter name and + the value is the value for the parameter. + """ + for name, value in six.iteritems(kwargs): + self.add_param(name, value) + + def add_param(self, name, value): + """Adds {name, value} pair to hyperparameters. + + Args: + name: Name of the hyperparameter. + value: Value of the hyperparameter. Can be one of the following types: + int, float, string, int list, float list, or string list. + + Raises: + ValueError: if one of the arguments is invalid. + """ + # Keys in kwargs are unique, but 'name' could be the name of a pre-existing + # attribute of this object. In that case we refuse to use it as a + # parameter name. + if getattr(self, name, None) is not None: + raise ValueError("Parameter name is reserved: %s" % name) + setattr(self, name, value) + + +def hierarchical_controller_hparams(): + """Hyperparameters for hierarchical planner.""" + return PlacerParams( + hidden_size=512, + forget_bias_init=1.0, + temperature=1.0, + logits_std_noise=0.5, + stop_noise_step=750, + decay_steps=50, + max_num_outputs=5, + max_output_size=5, + tanh_constant=1.0, + adj_embed_dim=20, + grouping_hidden_size=64, + num_groups=None, + bi_lstm=True, + failing_signal=100, + stop_sampling=500, + start_with_failing_signal=True, + always_update_baseline=False, + bl_dec=0.9, + grad_bound=1.0, + lr=0.1, + lr_dec=0.95, + start_decay_step=400, + optimizer_type="adam", + stop_updating_after_steps=1000, + name="hierarchical_controller", + keep_prob=1.0, + reward_function="sqrt", + seed=1234, + # distributed training params + num_children=1) + + +class HierarchicalController(Controller): + """HierarchicalController class.""" + + def __init__(self, hparams, item, cluster, controller_id=0): + """HierarchicalController class initializer. + + Args: + hparams: All hyper-parameters. + item: The metagraph to place. + cluster: The cluster of hardware devices to optimize for. + controller_id: the id of the controller in a multi-controller setup. + """ + super(HierarchicalController, self).__init__(item, cluster) + self.ctrl_id = controller_id + self.hparams = hparams + + if self.hparams.num_groups is None: + self.num_groups = min(256, 20 * self.num_devices) + else: + self.num_groups = self.hparams.num_groups + + # creates self.op_embeddings and self.type_dict + self.create_op_embeddings(verbose=False) + # TODO(azalia) clean up embedding/group_embedding_size names + self.group_emb_size = ( + 2 * self.num_groups + len(self.type_dict) + + self.hparams.max_num_outputs * self.hparams.max_output_size) + self.embedding_size = self.group_emb_size + self.initializer = init_ops.glorot_uniform_initializer( + seed=self.hparams.seed) + + with variable_scope.variable_scope( + self.hparams.name, + initializer=self.initializer, + reuse=variable_scope.AUTO_REUSE): + # define parameters of feedforward + variable_scope.get_variable("w_grouping_ff", [ + 1 + self.hparams.max_num_outputs * self.hparams.max_output_size + + self.hparams.adj_embed_dim, self.hparams.grouping_hidden_size + ]) + variable_scope.get_variable( + "w_grouping_softmax", + [self.hparams.grouping_hidden_size, self.num_groups]) + if self.hparams.bi_lstm: + variable_scope.get_variable("encoder_lstm_forward", [ + self.embedding_size + self.hparams.hidden_size / 2, + 2 * self.hparams.hidden_size + ]) + variable_scope.get_variable("encoder_lstm_backward", [ + self.embedding_size + self.hparams.hidden_size / 2, + 2 * self.hparams.hidden_size + ]) + variable_scope.get_variable( + "device_embeddings", [self.num_devices, self.hparams.hidden_size]) + variable_scope.get_variable( + "decoder_lstm", + [2 * self.hparams.hidden_size, 4 * self.hparams.hidden_size]) + variable_scope.get_variable( + "device_softmax", [2 * self.hparams.hidden_size, self.num_devices]) + variable_scope.get_variable("device_go_embedding", + [1, self.hparams.hidden_size]) + variable_scope.get_variable( + "encoder_forget_bias", + shape=1, + dtype=dtypes.float32, + initializer=init_ops.constant_initializer( + self.hparams.forget_bias_init)) + variable_scope.get_variable( + "decoder_forget_bias", + shape=1, + dtype=dtypes.float32, + initializer=init_ops.constant_initializer( + self.hparams.forget_bias_init)) + variable_scope.get_variable( + "attn_w_1", [self.hparams.hidden_size, self.hparams.hidden_size]) + variable_scope.get_variable( + "attn_w_2", [self.hparams.hidden_size, self.hparams.hidden_size]) + variable_scope.get_variable("attn_v", [self.hparams.hidden_size, 1]) + + else: + variable_scope.get_variable("encoder_lstm", [ + self.embedding_size + self.hparams.hidden_size, + 4 * self.hparams.hidden_size + ]) + variable_scope.get_variable( + "device_embeddings", [self.num_devices, self.hparams.hidden_size]) + variable_scope.get_variable( + "decoder_lstm", + [2 * self.hparams.hidden_size, 4 * self.hparams.hidden_size]) + variable_scope.get_variable( + "device_softmax", [2 * self.hparams.hidden_size, self.num_devices]) + variable_scope.get_variable("device_go_embedding", + [1, self.hparams.hidden_size]) + variable_scope.get_variable( + "encoder_forget_bias", + shape=1, + dtype=dtypes.float32, + initializer=init_ops.constant_initializer( + self.hparams.forget_bias_init)) + variable_scope.get_variable( + "decoder_forget_bias", + shape=1, + dtype=dtypes.float32, + initializer=init_ops.constant_initializer( + self.hparams.forget_bias_init)) + variable_scope.get_variable( + "attn_w_1", [self.hparams.hidden_size, self.hparams.hidden_size]) + variable_scope.get_variable( + "attn_w_2", [self.hparams.hidden_size, self.hparams.hidden_size]) + variable_scope.get_variable("attn_v", [self.hparams.hidden_size, 1]) + seq2seq_input_layer = array_ops.placeholder_with_default( + array_ops.zeros([self.hparams.num_children, + self.num_groups, + self.group_emb_size], + dtypes.float32), + shape=(self.hparams.num_children, self.num_groups, self.group_emb_size)) + self.seq2seq_input_layer = seq2seq_input_layer + + def compute_reward(self, run_time): + if self.hparams.reward_function == "id": + reward = run_time + elif self.hparams.reward_function == "sqrt": + reward = math.sqrt(run_time) + elif self.hparams.reward_function == "log": + reward = math.log1p(run_time) + else: + raise NotImplementedError( + "Unrecognized reward function '%s', consider your " + "--reward_function flag value." % self.hparams.reward_function) + return reward + + def build_controller(self): + """RL optimization interface. + + Returns: + ops: A dictionary holding handles of the model used for training. + """ + + self._global_step = training_util.get_or_create_global_step() + ops = {} + ops["loss"] = 0 + + failing_signal = self.compute_reward(self.hparams.failing_signal) + + ctr = {} + + with tf_ops.name_scope("controller_{}".format(self.ctrl_id)): + with variable_scope.variable_scope("controller_{}".format(self.ctrl_id)): + ctr["reward"] = {"value": [], "ph": [], "update": []} + ctr["ready"] = {"value": [], "ph": [], "update": []} + ctr["best_reward"] = {"value": [], "update": []} + for i in range(self.hparams.num_children): + reward_value = variable_scope.get_local_variable( + "reward_{}".format(i), + initializer=0.0, + dtype=dtypes.float32, + trainable=False) + reward_ph = array_ops.placeholder( + dtypes.float32, shape=(), name="reward_ph_{}".format(i)) + reward_update = state_ops.assign( + reward_value, reward_ph, use_locking=True) + ctr["reward"]["value"].append(reward_value) + ctr["reward"]["ph"].append(reward_ph) + ctr["reward"]["update"].append(reward_update) + best_reward = variable_scope.get_local_variable( + "best_reward_{}".format(i), + initializer=failing_signal, + dtype=dtypes.float32, + trainable=False) + ctr["best_reward"]["value"].append(best_reward) + ctr["best_reward"]["update"].append( + state_ops.assign(best_reward, + math_ops.minimum(best_reward, reward_update))) + + ready_value = variable_scope.get_local_variable( + "ready_{}".format(i), + initializer=True, + dtype=dtypes.bool, + trainable=False) + ready_ph = array_ops.placeholder( + dtypes.bool, shape=(), name="ready_ph_{}".format(i)) + ready_update = state_ops.assign( + ready_value, ready_ph, use_locking=True) + ctr["ready"]["value"].append(ready_value) + ctr["ready"]["ph"].append(ready_ph) + ctr["ready"]["update"].append(ready_update) + + ctr["grouping_y_preds"], ctr["grouping_log_probs"] = self.get_groupings() + summary.histogram( + "grouping_actions", + array_ops.slice(ctr["grouping_y_preds"]["sample"], [0, 0], + [1, array_ops.shape(self.op_embeddings)[0]])) + + with variable_scope.variable_scope("controller_{}".format(self.ctrl_id)): + ctr["baseline"] = variable_scope.get_local_variable( + "baseline", + initializer=failing_signal + if self.hparams.start_with_failing_signal else 0.0, + dtype=dtypes.float32, + trainable=False) + + new_baseline = self.hparams.bl_dec * ctr["baseline"] + ( + 1 - self.hparams.bl_dec) * math_ops.reduce_mean( + ctr["reward"]["value"]) + if not self.hparams.always_update_baseline: + baseline_mask = math_ops.less(ctr["reward"]["value"], failing_signal) + selected_reward = array_ops.boolean_mask(ctr["reward"]["value"], + baseline_mask) + selected_baseline = control_flow_ops.cond( + math_ops.reduce_any(baseline_mask), + lambda: math_ops.reduce_mean(selected_reward), + lambda: constant_op.constant(0, dtype=dtypes.float32)) + ctr["pos_reward"] = selected_baseline + pos_ = math_ops.less( + constant_op.constant(0, dtype=dtypes.float32), selected_baseline) + selected_baseline = self.hparams.bl_dec * ctr["baseline"] + ( + 1 - self.hparams.bl_dec) * selected_baseline + selected_baseline = control_flow_ops.cond( + pos_, lambda: selected_baseline, lambda: ctr["baseline"]) + new_baseline = control_flow_ops.cond( + math_ops.less(self.global_step, + self.hparams.stop_updating_after_steps), + lambda: new_baseline, lambda: selected_baseline) + ctr["baseline_update"] = state_ops.assign( + ctr["baseline"], new_baseline, use_locking=True) + + ctr["y_preds"], ctr["log_probs"] = self.get_placements() + summary.histogram("actions", ctr["y_preds"]["sample"]) + mask = math_ops.less(ctr["reward"]["value"], failing_signal) + ctr["loss"] = ctr["reward"]["value"] - ctr["baseline"] + ctr["loss"] *= ( + ctr["log_probs"]["sample"] + ctr["grouping_log_probs"]["sample"]) + + selected_loss = array_ops.boolean_mask(ctr["loss"], mask) + selected_loss = control_flow_ops.cond( + math_ops.reduce_any(mask), + lambda: math_ops.reduce_mean(-selected_loss), + lambda: constant_op.constant(0, dtype=dtypes.float32)) + + ctr["loss"] = control_flow_ops.cond( + math_ops.less(self.global_step, + self.hparams.stop_updating_after_steps), + lambda: math_ops.reduce_mean(-ctr["loss"]), lambda: selected_loss) + + ctr["reward_s"] = math_ops.reduce_mean(ctr["reward"]["value"]) + summary.scalar("loss", ctr["loss"]) + summary.scalar("avg_reward", ctr["reward_s"]) + summary.scalar("best_reward_so_far", best_reward) + summary.scalar( + "advantage", + math_ops.reduce_mean(ctr["reward"]["value"] - ctr["baseline"])) + + with variable_scope.variable_scope( + "optimizer", reuse=variable_scope.AUTO_REUSE): + (ctr["train_op"], ctr["lr"], ctr["grad_norm"], + ctr["grad_norms"]) = self._get_train_ops( + ctr["loss"], + tf_ops.get_collection(tf_ops.GraphKeys.TRAINABLE_VARIABLES), + self.global_step, + grad_bound=self.hparams.grad_bound, + lr_init=self.hparams.lr, + lr_dec=self.hparams.lr_dec, + start_decay_step=self.hparams.start_decay_step, + decay_steps=self.hparams.decay_steps, + optimizer_type=self.hparams.optimizer_type) + + summary.scalar("gradnorm", ctr["grad_norm"]) + summary.scalar("lr", ctr["lr"]) + ctr["summary"] = summary.merge_all() + ops["controller"] = ctr + + self.ops = ops + return ops + + @property + def global_step(self): + return self._global_step + + def create_op_embeddings(self, verbose=False): + if verbose: + print("process input graph for op embeddings") + self.num_ops = len(self.important_ops) + # topological sort of important nodes + topo_order = [op.name for op in self.important_ops] + + # create index to name for topologicaly sorted important nodes + name_to_topo_order_index = {} + for idx, x in enumerate(topo_order): + name_to_topo_order_index[x] = idx + self.name_to_topo_order_index = name_to_topo_order_index + + # create adj matrix + adj_dict = {} + for idx, op in enumerate(self.important_ops): + for output_op in self.get_node_fanout(op): + output_op_name = output_op.name + if output_op_name in self.important_op_names: + if name_to_topo_order_index[op.name] not in adj_dict: + adj_dict[name_to_topo_order_index[op.name]] = [] + adj_dict[name_to_topo_order_index[op.name]].extend( + [name_to_topo_order_index[output_op_name], 1]) + if output_op_name not in adj_dict: + adj_dict[name_to_topo_order_index[output_op_name]] = [] + adj_dict[name_to_topo_order_index[output_op_name]].extend( + [name_to_topo_order_index[op.name], -1]) + + # get op_type op_output_shape, and adj info + output_embed_dim = (self.hparams.max_num_outputs * + self.hparams.max_output_size) + + # TODO(bsteiner): don't filter based on used ops so that we can generalize + # to models that use other types of ops. + used_ops = set() + for node in self.important_ops: + op_type = str(node.op) + used_ops.add(op_type) + + self.type_dict = {} + for op_type in self.cluster.ListAvailableOps(): + if op_type in used_ops: + self.type_dict[op_type] = len(self.type_dict) + + op_types = np.zeros([self.num_ops], dtype=np.int32) + op_output_shapes = np.full( + [self.num_ops, output_embed_dim], -1.0, dtype=np.float32) + for idx, node in enumerate(self.important_ops): + op_types[idx] = self.type_dict[node.op] + # output shape + op_name = node.name + for i, output_prop in enumerate(self.node_properties[op_name]): + if output_prop.shape.__str__() == "": + continue + shape = output_prop.shape + for j, dim in enumerate(shape.dim): + if dim.size >= 0: + if i * self.hparams.max_output_size + j >= output_embed_dim: + break + op_output_shapes[idx, + i * self.hparams.max_output_size + j] = dim.size + # adj for padding + op_adj = np.full( + [self.num_ops, self.hparams.adj_embed_dim], 0, dtype=np.float32) + for idx in adj_dict: + neighbors = adj_dict[int(idx)] + min_dim = min(self.hparams.adj_embed_dim, len(neighbors)) + padding_size = self.hparams.adj_embed_dim - min_dim + neighbors = neighbors[:min_dim] + [0] * padding_size + op_adj[int(idx)] = neighbors + + # op_embedding starts here + op_embeddings = np.zeros( + [ + self.num_ops, + 1 + self.hparams.max_num_outputs * self.hparams.max_output_size + + self.hparams.adj_embed_dim + ], + dtype=np.float32) + for idx, op_name in enumerate(topo_order): + op_embeddings[idx] = np.concatenate( + (np.array([op_types[idx]]), op_output_shapes[idx], op_adj[int(idx)])) + self.op_embeddings = constant_op.constant( + op_embeddings, dtype=dtypes.float32) + if verbose: + print("num_ops = {}".format(self.num_ops)) + print("num_types = {}".format(len(self.type_dict))) + + def get_groupings(self, *args, **kwargs): + num_children = self.hparams.num_children + with variable_scope.variable_scope("controller_{}".format(self.ctrl_id)): + grouping_actions_cache = variable_scope.get_local_variable( + "grouping_actions_cache", + initializer=init_ops.zeros_initializer, + dtype=dtypes.int32, + shape=[num_children, self.num_ops], + trainable=False) + input_layer = self.op_embeddings + input_layer = array_ops.expand_dims(input_layer, 0) + feed_ff_input_layer = array_ops.tile(input_layer, [num_children, 1, 1]) + grouping_actions, grouping_log_probs = {}, {} + grouping_actions["sample"], grouping_log_probs[ + "sample"] = self.make_grouping_predictions(feed_ff_input_layer) + + grouping_actions["sample"] = state_ops.assign(grouping_actions_cache, + grouping_actions["sample"]) + self.grouping_actions_cache = grouping_actions_cache + + return grouping_actions, grouping_log_probs + + def make_grouping_predictions(self, input_layer, reuse=None): + """model that predicts grouping (grouping_actions). + + Args: + input_layer: group_input_layer + reuse: reuse + + Returns: + grouping_actions: actions + grouping_log_probs: log probabilities corresponding to actions + """ + with variable_scope.variable_scope(self.hparams.name, reuse=True): + # input_layer: tensor of size [1, num_ops, hidden_size] + w_grouping_ff = variable_scope.get_variable("w_grouping_ff") + w_grouping_softmax = variable_scope.get_variable("w_grouping_softmax") + + batch_size = array_ops.shape(input_layer)[0] + embedding_dim = array_ops.shape(input_layer)[2] + + reshaped = array_ops.reshape(input_layer, + [batch_size * self.num_ops, embedding_dim]) + ff_output = math_ops.matmul(reshaped, w_grouping_ff) + logits = math_ops.matmul(ff_output, w_grouping_softmax) + if self.hparams.logits_std_noise > 0: + num_in_logits = math_ops.cast( + array_ops.size(logits), dtype=dtypes.float32) + avg_norm = math_ops.divide( + linalg_ops.norm(logits), math_ops.sqrt(num_in_logits)) + logits_noise = random_ops.random_normal( + array_ops.shape(logits), + stddev=self.hparams.logits_std_noise * avg_norm) + logits = control_flow_ops.cond( + self.global_step > self.hparams.stop_noise_step, lambda: logits, + lambda: logits + logits_noise) + logits = array_ops.reshape(logits, + [batch_size * self.num_ops, self.num_groups]) + actions = random_ops.multinomial(logits, 1, seed=self.hparams.seed) + actions = math_ops.to_int32(actions) + actions = array_ops.reshape(actions, [batch_size, self.num_ops]) + action_label = array_ops.reshape(actions, [-1]) + log_probs = nn_ops.sparse_softmax_cross_entropy_with_logits( + logits=logits, labels=action_label) + log_probs = array_ops.reshape(log_probs, [batch_size, -1]) + log_probs = math_ops.reduce_sum(log_probs, 1) + grouping_actions = actions + grouping_log_probs = log_probs + return grouping_actions, grouping_log_probs + + def create_group_embeddings(self, grouping_actions, verbose=False): + """Approximating the blocks of a TF graph from a graph_def. + + Args: + grouping_actions: grouping predictions. + verbose: print stuffs. + + Returns: + groups: list of groups. + """ + groups = [ + self._create_group_embeddings(grouping_actions, i, verbose) for + i in range(self.hparams.num_children) + ] + return np.stack(groups, axis=0) + + def _create_group_embeddings(self, grouping_actions, child_id, verbose=False): + """Approximating the blocks of a TF graph from a graph_def for each child. + + Args: + grouping_actions: grouping predictions. + child_id: child_id for the group. + verbose: print stuffs. + + Returns: + groups: group embedding for the child_id. + """ + if verbose: + print("Processing input_graph") + + # TODO(azalia): Build inter-adjacencies dag matrix. + # record dag_matrix + dag_matrix = np.zeros([self.num_groups, self.num_groups], dtype=np.float32) + for op in self.important_ops: + topo_op_index = self.name_to_topo_order_index[op.name] + group_index = grouping_actions[child_id][topo_op_index] + for output_op in self.get_node_fanout(op): + if output_op.name not in self.important_op_names: + continue + output_group_index = ( + grouping_actions[child_id][self.name_to_topo_order_index[ + output_op.name]]) + dag_matrix[group_index, output_group_index] += 1.0 + num_connections = np.sum(dag_matrix) + num_intra_group_connections = dag_matrix.trace() + num_inter_group_connections = num_connections - num_intra_group_connections + if verbose: + print("grouping evaluation metric") + print(("num_connections={} num_intra_group_connections={} " + "num_inter_group_connections={}").format( + num_connections, num_intra_group_connections, + num_inter_group_connections)) + self.dag_matrix = dag_matrix + + # output_shape + op_output_shapes = np.zeros( + [ + len(self.important_ops), + self.hparams.max_num_outputs * self.hparams.max_output_size + ], + dtype=np.float32) + + for idx, op in enumerate(self.important_ops): + for i, output_properties in enumerate(self.node_properties[op.name]): + if output_properties.shape.__str__() == "": + continue + if i > self.hparams.max_num_outputs: + break + shape = output_properties.shape + for j, dim in enumerate(shape.dim): + if dim.size > 0: + k = i * self.hparams.max_output_size + j + if k >= self.hparams.max_num_outputs * self.hparams.max_output_size: + break + op_output_shapes[idx, k] = dim.size + + # group_embedding + group_embedding = np.zeros( + [ + self.num_groups, len(self.type_dict) + + self.hparams.max_num_outputs * self.hparams.max_output_size + ], + dtype=np.float32) + for op_index, op in enumerate(self.important_ops): + group_index = grouping_actions[child_id][ + self.name_to_topo_order_index[op.name]] + type_name = str(op.op) + type_index = self.type_dict[type_name] + group_embedding[group_index, type_index] += 1 + group_embedding[group_index, :self.hparams.max_num_outputs * self.hparams. + max_output_size] += ( + op_output_shapes[op_index]) + grouping_adjacencies = np.concatenate( + [dag_matrix, np.transpose(dag_matrix)], axis=1) + group_embedding = np.concatenate( + [grouping_adjacencies, group_embedding], axis=1) + group_normalizer = np.amax(group_embedding, axis=1, keepdims=True) + group_embedding /= (group_normalizer + 1.0) + if verbose: + print("Finished Processing Input Graph") + return group_embedding + + def get_placements(self, *args, **kwargs): + num_children = self.hparams.num_children + with variable_scope.variable_scope("controller_{}".format(self.ctrl_id)): + actions_cache = variable_scope.get_local_variable( + "actions_cache", + initializer=init_ops.zeros_initializer, + dtype=dtypes.int32, + shape=[num_children, self.num_groups], + trainable=False) + + x = self.seq2seq_input_layer + last_c, last_h, attn_mem = self.encode(x) + actions, log_probs = {}, {} + actions["sample"], log_probs["sample"] = ( + self.decode( + x, last_c, last_h, attn_mem, mode="sample")) + actions["target"], log_probs["target"] = ( + self.decode( + x, + last_c, + last_h, + attn_mem, + mode="target", + y=actions_cache)) + actions["greedy"], log_probs["greedy"] = ( + self.decode( + x, last_c, last_h, attn_mem, mode="greedy")) + actions["sample"] = control_flow_ops.cond( + self.global_step < self.hparams.stop_sampling, + lambda: state_ops.assign(actions_cache, actions["sample"]), + lambda: state_ops.assign(actions_cache, actions["target"])) + self.actions_cache = actions_cache + + return actions, log_probs + + def encode(self, x): + """Encoder using LSTM. + + Args: + x: tensor of size [num_children, num_groups, embedding_size] + + Returns: + last_c, last_h: tensors of size [num_children, hidden_size], the final + LSTM states + attn_mem: tensor of size [num_children, num_groups, hidden_size], the + attention + memory, i.e. concatenation of all hidden states, linearly transformed by + an attention matrix attn_w_1 + """ + if self.hparams.bi_lstm: + with variable_scope.variable_scope(self.hparams.name, reuse=True): + w_lstm_forward = variable_scope.get_variable("encoder_lstm_forward") + w_lstm_backward = variable_scope.get_variable("encoder_lstm_backward") + forget_bias = variable_scope.get_variable("encoder_forget_bias") + attn_w_1 = variable_scope.get_variable("attn_w_1") + else: + with variable_scope.variable_scope(self.hparams.name, reuse=True): + w_lstm = variable_scope.get_variable("encoder_lstm") + forget_bias = variable_scope.get_variable("encoder_forget_bias") + attn_w_1 = variable_scope.get_variable("attn_w_1") + + embedding_size = array_ops.shape(x)[2] + + signals = array_ops.split(x, self.num_groups, axis=1) + for i in range(len(signals)): + signals[i] = array_ops.reshape( + signals[i], [self.hparams.num_children, embedding_size]) + + if self.hparams.bi_lstm: + + def body(i, prev_c_forward, prev_h_forward, prev_c_backward, + prev_h_backward): + """while loop for LSTM.""" + signal_forward = signals[i] + next_c_forward, next_h_forward = lstm(signal_forward, prev_c_forward, + prev_h_forward, w_lstm_forward, + forget_bias) + + signal_backward = signals[self.num_groups - 1 - i] + next_c_backward, next_h_backward = lstm( + signal_backward, prev_c_backward, prev_h_backward, w_lstm_backward, + forget_bias) + + next_h = array_ops.concat([next_h_forward, next_h_backward], axis=1) + all_h.append(next_h) + + return (next_c_forward, next_h_forward, next_c_backward, + next_h_backward) + + c_forward = array_ops.zeros( + [self.hparams.num_children, self.hparams.hidden_size / 2], + dtype=dtypes.float32) + h_forward = array_ops.zeros( + [self.hparams.num_children, self.hparams.hidden_size / 2], + dtype=dtypes.float32) + + c_backward = array_ops.zeros( + [self.hparams.num_children, self.hparams.hidden_size / 2], + dtype=dtypes.float32) + h_backward = array_ops.zeros( + [self.hparams.num_children, self.hparams.hidden_size / 2], + dtype=dtypes.float32) + all_h = [] + + for i in range(0, self.num_groups): + c_forward, h_forward, c_backward, h_backward = body( + i, c_forward, h_forward, c_backward, h_backward) + + last_c = array_ops.concat([c_forward, c_backward], axis=1) + last_h = array_ops.concat([h_forward, h_backward], axis=1) + attn_mem = array_ops.stack(all_h) + + else: + + def body(i, prev_c, prev_h): + signal = signals[i] + next_c, next_h = lstm(signal, prev_c, prev_h, w_lstm, forget_bias) + all_h.append(next_h) + return next_c, next_h + + c = array_ops.zeros( + [self.hparams.num_children, self.hparams.hidden_size], + dtype=dtypes.float32) + h = array_ops.zeros( + [self.hparams.num_children, self.hparams.hidden_size], + dtype=dtypes.float32) + all_h = [] + + for i in range(0, self.num_groups): + c, h = body(i, c, h) + + last_c = c + last_h = h + attn_mem = array_ops.stack(all_h) + + attn_mem = array_ops.transpose(attn_mem, [1, 0, 2]) + attn_mem = array_ops.reshape( + attn_mem, + [self.hparams.num_children * self.num_groups, self.hparams.hidden_size]) + attn_mem = math_ops.matmul(attn_mem, attn_w_1) + attn_mem = array_ops.reshape( + attn_mem, + [self.hparams.num_children, self.num_groups, self.hparams.hidden_size]) + + return last_c, last_h, attn_mem + + def decode(self, + x, + last_c, + last_h, + attn_mem, + mode="target", + y=None): + """Decoder using LSTM. + + Args: + x: tensor of size [num_children, num_groups, embedding_size]. + last_c: tensor of size [num_children, hidden_size], the final LSTM states + computed by self.encoder. + last_h: same as last_c. + attn_mem: tensor of size [num_children, num_groups, hidden_size]. + mode: "target" or "sample". + y: tensor of size [num_children, num_groups], the device placements. + + Returns: + actions: tensor of size [num_children, num_groups], the placements of + devices + """ + with variable_scope.variable_scope(self.hparams.name, reuse=True): + w_lstm = variable_scope.get_variable("decoder_lstm") + forget_bias = variable_scope.get_variable("decoder_forget_bias") + device_embeddings = variable_scope.get_variable("device_embeddings") + device_softmax = variable_scope.get_variable("device_softmax") + device_go_embedding = variable_scope.get_variable("device_go_embedding") + attn_w_2 = variable_scope.get_variable("attn_w_2") + attn_v = variable_scope.get_variable("attn_v") + + actions = tensor_array_ops.TensorArray( + dtypes.int32, + size=self.num_groups, + infer_shape=False, + clear_after_read=False) + + # pylint: disable=unused-argument + def condition(i, *args): + return math_ops.less(i, self.num_groups) + + # pylint: disable=missing-docstring + def body(i, prev_c, prev_h, actions, log_probs): + # pylint: disable=g-long-lambda + signal = control_flow_ops.cond( + math_ops.equal(i, 0), + lambda: array_ops.tile(device_go_embedding, + [self.hparams.num_children, 1]), + lambda: embedding_ops.embedding_lookup(device_embeddings, + actions.read(i - 1)) + ) + if self.hparams.keep_prob is not None: + signal = nn_ops.dropout(signal, self.hparams.keep_prob) + next_c, next_h = lstm(signal, prev_c, prev_h, w_lstm, forget_bias) + query = math_ops.matmul(next_h, attn_w_2) + query = array_ops.reshape( + query, [self.hparams.num_children, 1, self.hparams.hidden_size]) + query = math_ops.tanh(query + attn_mem) + query = array_ops.reshape(query, [ + self.hparams.num_children * self.num_groups, self.hparams.hidden_size + ]) + query = math_ops.matmul(query, attn_v) + query = array_ops.reshape(query, + [self.hparams.num_children, self.num_groups]) + query = nn_ops.softmax(query) + query = array_ops.reshape(query, + [self.hparams.num_children, self.num_groups, 1]) + query = math_ops.reduce_sum(attn_mem * query, axis=1) + query = array_ops.concat([next_h, query], axis=1) + logits = math_ops.matmul(query, device_softmax) + logits /= self.hparams.temperature + if self.hparams.tanh_constant > 0: + logits = math_ops.tanh(logits) * self.hparams.tanh_constant + if self.hparams.logits_std_noise > 0: + num_in_logits = math_ops.cast( + array_ops.size(logits), dtype=dtypes.float32) + avg_norm = math_ops.divide( + linalg_ops.norm(logits), math_ops.sqrt(num_in_logits)) + logits_noise = random_ops.random_normal( + array_ops.shape(logits), + stddev=self.hparams.logits_std_noise * avg_norm) + logits = control_flow_ops.cond( + self.global_step > self.hparams.stop_noise_step, lambda: logits, + lambda: logits + logits_noise) + + if mode == "sample": + next_y = random_ops.multinomial(logits, 1, seed=self.hparams.seed) + elif mode == "greedy": + next_y = math_ops.argmax(logits, 1) + elif mode == "target": + next_y = array_ops.slice(y, [0, i], [-1, 1]) + else: + raise NotImplementedError + next_y = math_ops.to_int32(next_y) + next_y = array_ops.reshape(next_y, [self.hparams.num_children]) + actions = actions.write(i, next_y) + log_probs += nn_ops.sparse_softmax_cross_entropy_with_logits( + logits=logits, labels=next_y) + return i + 1, next_c, next_h, actions, log_probs + + loop_vars = [ + constant_op.constant(0, dtype=dtypes.int32), last_c, last_h, actions, + array_ops.zeros([self.hparams.num_children], dtype=dtypes.float32) + ] + loop_outputs = control_flow_ops.while_loop(condition, body, loop_vars) + + last_c = loop_outputs[-4] + last_h = loop_outputs[-3] + actions = loop_outputs[-2].stack() + actions = array_ops.transpose(actions, [1, 0]) + log_probs = loop_outputs[-1] + return actions, log_probs + + def eval_placement(self, + sess, + child_id=0, + verbose=False): + grouping_actions, actions = sess.run([ + self.grouping_actions_cache, + self.actions_cache + ]) + grouping_actions = grouping_actions[child_id] + actions = actions[child_id] + if verbose: + global_step = sess.run(self.global_step) + if global_step % 100 == 0: + log_string = "op group assignments: " + for a in grouping_actions: + log_string += "{} ".format(a) + print(log_string[:-1]) + log_string = "group device assignments: " + for a in actions: + log_string += "{} ".format(a) + print(log_string[:-1]) + + for op in self.important_ops: + topo_order_index = self.name_to_topo_order_index[op.name] + group_index = grouping_actions[topo_order_index] + op.device = self.devices[actions[group_index]].name + try: + _, run_time, _ = self.cluster.MeasureCosts(self.item) + except errors.ResourceExhaustedError: + run_time = self.hparams.failing_signal + return run_time + + def update_reward(self, + sess, + run_time, + child_id=0, + verbose=False): + reward = self.compute_reward(run_time) + controller_ops = self.ops["controller"] + _, best_reward = sess.run( + [ + controller_ops["reward"]["update"][child_id], + controller_ops["best_reward"]["update"][child_id] + ], + feed_dict={ + controller_ops["reward"]["ph"][child_id]: reward, + }) + if verbose: + print(("run_time={:<.5f} reward={:<.5f} " + "best_reward={:<.5f}").format(run_time, reward, best_reward)) + + # Reward is a double, best_reward a float: allow for some slack in the + # comparison. + updated = abs(best_reward - reward) < 1e-6 + return updated + + def generate_grouping(self, sess): + controller_ops = self.ops["controller"] + grouping_actions = sess.run(controller_ops["grouping_y_preds"]["sample"]) + return grouping_actions + + def generate_placement(self, grouping, sess): + controller_ops = self.ops["controller"] + feed_seq2seq_input_dict = {} + feed_seq2seq_input_dict[self.seq2seq_input_layer] = grouping + sess.run( + controller_ops["y_preds"]["sample"], feed_dict=feed_seq2seq_input_dict) + + def process_reward(self, sess): + controller_ops = self.ops["controller"] + run_ops = [ + controller_ops["loss"], controller_ops["lr"], + controller_ops["grad_norm"], controller_ops["grad_norms"], + controller_ops["train_op"] + ] + sess.run(run_ops) + sess.run(controller_ops["baseline_update"]) + + def _get_train_ops(self, + loss, + tf_variables, + global_step, + grad_bound=1.25, + lr_init=1e-3, + lr_dec=0.9, + start_decay_step=10000, + decay_steps=100, + optimizer_type="adam"): + """Loss optimizer. + + Args: + loss: scalar tf tensor + tf_variables: list of training variables, typically + tf.trainable_variables() + global_step: global_step + grad_bound: max gradient norm + lr_init: initial learning rate + lr_dec: leaning rate decay coefficient + start_decay_step: start decaying learning rate after this many steps + decay_steps: apply decay rate factor at this step intervals + optimizer_type: optimizer type should be either adam or sgd + + Returns: + train_op: training op + learning_rate: scalar learning rate tensor + grad_norm: l2 norm of the gradient vector + all_grad_norms: l2 norm of each component + """ + lr_gstep = global_step - start_decay_step + + def f1(): + return constant_op.constant(lr_init) + + def f2(): + return learning_rate_decay.exponential_decay(lr_init, lr_gstep, + decay_steps, lr_dec, True) + + learning_rate = control_flow_ops.cond( + math_ops.less(global_step, start_decay_step), + f1, + f2, + name="learning_rate") + + if optimizer_type == "adam": + opt = adam.AdamOptimizer(learning_rate) + elif optimizer_type == "sgd": + opt = gradient_descent.GradientDescentOptimizer(learning_rate) + grads_and_vars = opt.compute_gradients(loss, tf_variables) + grad_norm = clip_ops.global_norm([g for g, v in grads_and_vars]) + all_grad_norms = {} + clipped_grads = [] + clipped_rate = math_ops.maximum(grad_norm / grad_bound, 1.0) + for g, v in grads_and_vars: + if g is not None: + if isinstance(g, tf_ops.IndexedSlices): + clipped = g.values / clipped_rate + norm_square = math_ops.reduce_sum(clipped * clipped) + clipped = tf_ops.IndexedSlices(clipped, g.indices) + else: + clipped = g / clipped_rate + norm_square = math_ops.reduce_sum(clipped * clipped) + all_grad_norms[v.name] = math_ops.sqrt(norm_square) + clipped_grads.append((clipped, v)) + + train_op = opt.apply_gradients(clipped_grads, global_step) + return train_op, learning_rate, grad_norm, all_grad_norms + + +def lstm(x, prev_c, prev_h, w_lstm, forget_bias): + """LSTM cell. + + Args: + x: tensors of size [num_children, hidden_size]. + prev_c: tensors of size [num_children, hidden_size]. + prev_h: same as prev_c. + w_lstm: . + forget_bias: . + + Returns: + next_c: + next_h: + """ + ifog = math_ops.matmul(array_ops.concat([x, prev_h], axis=1), w_lstm) + i, f, o, g = array_ops.split(ifog, 4, axis=1) + i = math_ops.sigmoid(i) + f = math_ops.sigmoid(f + forget_bias) + o = math_ops.sigmoid(o) + g = math_ops.tanh(g) + next_c = i * g + f * prev_c + next_h = o * math_ops.tanh(next_c) + return next_c, next_h diff --git a/tensorflow/python/grappler/item.i b/tensorflow/python/grappler/item.i index d0fc1a04f220e0a053257e0206bb07b25f3767c6..9a84c60b04029a64ed35a01f045a6eec5e492504 100644 --- a/tensorflow/python/grappler/item.i +++ b/tensorflow/python/grappler/item.i @@ -96,10 +96,10 @@ static GItem TF_NewItem( return GItem(item.release()); } -static std::vector TF_IdentifyImportantOps(GItem item, bool sort_topologically, +static PyObject* TF_IdentifyImportantOps(GItem item, bool sort_topologically, TF_Status* status) { if (item.is_none()) { - return {}; + Py_RETURN_NONE; } std::vector main_ops = item->MainOpsFanin(); @@ -132,7 +132,13 @@ static std::vector TF_IdentifyImportantOps(GItem item, bool sort_topolog } } - return ops; + PyGILState_STATE gstate = PyGILState_Ensure(); + PyObject* result = PyList_New(ops.size()); + for (int i = 0; i < ops.size(); ++i) { + PyList_SetItem(result, i, PyString_FromString(ops[i].c_str())); + } + PyGILState_Release(gstate); + return result; } static PyObject* TF_GetOpProperties(GItem item) { @@ -305,7 +311,7 @@ static PyObject* TF_GetColocationGroups(GItem item) { static GItem TF_NewItem( const tensorflow::MetaGraphDef& meta_graph, bool ignore_colocation, bool ignore_user_placement, TF_Status* out_status); -static std::vector TF_IdentifyImportantOps(GItem item, bool sort_topologically, - TF_Status* status); +static PyObject* TF_IdentifyImportantOps(GItem item, bool sort_topologically, + TF_Status* status); static PyObject* TF_GetOpProperties(GItem item); static PyObject* TF_GetColocationGroups(GItem item); diff --git a/tensorflow/python/grappler/item_test.py b/tensorflow/python/grappler/item_test.py index cd70e2fdecc74f9d99240ac566f3c28e900a06c2..c40de9da0abca3bb99a82a1456261f45b1c45c99 100644 --- a/tensorflow/python/grappler/item_test.py +++ b/tensorflow/python/grappler/item_test.py @@ -56,7 +56,7 @@ class ItemTest(test.TestCase): mg = meta_graph.create_meta_graph_def(graph=g) grappler_item = item.Item(mg) op_list = grappler_item.IdentifyImportantOps() - self.assertItemsEqual([b'Const', b'Const_1', b'add'], op_list) + self.assertItemsEqual(['Const', 'Const_1', 'add'], op_list) def testOpProperties(self): with ops.Graph().as_default() as g: @@ -111,7 +111,7 @@ class ItemTest(test.TestCase): with ops.Graph().as_default() as g: c = constant_op.constant([10]) v = variables.Variable([3], dtype=dtypes.int32) - i = gen_array_ops._ref_identity(v) + i = gen_array_ops.ref_identity(v) a = state_ops.assign(i, c) train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP) train_op.append(a) diff --git a/tensorflow/python/grappler/layout_optimizer_test.py b/tensorflow/python/grappler/layout_optimizer_test.py index 578f86ca5a0c1f2446dbf26ce412e34f3bdbd23a..5a84b16a23f567fba6d08aaefd3b816a76907735 100644 --- a/tensorflow/python/grappler/layout_optimizer_test.py +++ b/tensorflow/python/grappler/layout_optimizer_test.py @@ -157,6 +157,7 @@ def _get_config(layout_optimizer=True): graph_options = config_pb2.GraphOptions( rewrite_options=rewrite_options, build_cost_model=1) config = config_pb2.ConfigProto(graph_options=graph_options) + config.graph_options.optimizer_options.opt_level = -1 return config @@ -179,6 +180,8 @@ def _get_cluster(): named_device = device_properties_pb2.NamedDevice() named_device.name = '/GPU:0' named_device.properties.type = 'GPU' + named_device.properties.num_cores = 24 + named_device.properties.frequency = 1000 named_device.properties.environment['architecture'] = '4' cluster = gcluster.Cluster(devices=[named_device]) return cluster @@ -253,7 +256,7 @@ class LayoutOptimizerTest(test.TestCase): x = random_ops.truncated_normal([1, 784], seed=0) output = _two_layer_model(x) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -290,7 +293,7 @@ class LayoutOptimizerTest(test.TestCase): add = bn0[0] + bn1[0] output = array_ops.identity(add) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={dim: 3}) with session.Session(config=_get_config()) as sess: @@ -318,11 +321,11 @@ class LayoutOptimizerTest(test.TestCase): conv = _two_layer_model(x) dim = array_ops.placeholder(dtype='int32') sizes = constant_op.constant([50, 10, 4], shape=[3]) - split = gen_array_ops._split_v( + split = gen_array_ops.split_v( value=conv, size_splits=sizes, axis=dim, num_split=3) output = math_ops.reduce_sum(split[0]) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={dim: 3}) with session.Session(config=_get_config()) as sess: @@ -356,7 +359,7 @@ class LayoutOptimizerTest(test.TestCase): pad = array_ops.pad(conv, paddings) output = array_ops.identity(pad) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -387,7 +390,7 @@ class LayoutOptimizerTest(test.TestCase): reduce_sum = math_ops.reduce_sum(conv) output = array_ops.identity(reduce_sum) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -416,7 +419,7 @@ class LayoutOptimizerTest(test.TestCase): cast = math_ops.cast(conv, dtype='bool') output = array_ops.identity(cast) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -447,7 +450,67 @@ class LayoutOptimizerTest(test.TestCase): squeeze = array_ops.squeeze(reduce_sum) output = array_ops.identity(squeeze) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run(output, run_metadata=metadata) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + # Three transposes were initially added in the Expand phase of + # LayoutOptimizer; two of them are cancelled out in the Collapse phase. + expected_num_transposes = 1 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self.assertAllClose(output_val_ref, output_val, atol=1e-3) + + def testSqueezeAlongHW(self): + if test.is_gpu_available(cuda_only=True): + random_seed.set_random_seed(0) + x = random_ops.truncated_normal([1, 784], seed=0) + conv = _two_layer_model(x) + reduce_sum = math_ops.reduce_sum(conv, axis=[1, 2], keep_dims=True) + squeeze = array_ops.squeeze(reduce_sum, axis=[1, 2]) + output = array_ops.identity(squeeze) + + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run(output, run_metadata=metadata) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + # Three transposes were initially added in the Expand phase of + # LayoutOptimizer; two of them are cancelled out in the Collapse phase. + expected_num_transposes = 1 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self.assertAllClose(output_val_ref, output_val, atol=1e-3) + + def testSqueezeAlongNHW(self): + if test.is_gpu_available(cuda_only=True): + random_seed.set_random_seed(0) + x = random_ops.truncated_normal([1, 784], seed=0) + conv = _two_layer_model(x) + reduce_sum = math_ops.reduce_sum(conv, axis=[0, 1, 2], keep_dims=True) + squeeze = array_ops.squeeze(reduce_sum, axis=[0, 1, 2]) + output = array_ops.identity(squeeze) + + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -476,7 +539,7 @@ class LayoutOptimizerTest(test.TestCase): reduce_sum = math_ops.reduce_sum(conv, axis=[1, 2, 3]) output = array_ops.identity(reduce_sum) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -505,7 +568,7 @@ class LayoutOptimizerTest(test.TestCase): reduce_sum = math_ops.reduce_sum(conv, axis=[0, 1, 2]) output = array_ops.identity(reduce_sum) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -534,7 +597,7 @@ class LayoutOptimizerTest(test.TestCase): reduce_sum = math_ops.reduce_sum(conv, axis=[3]) output = array_ops.identity(reduce_sum) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -555,6 +618,94 @@ class LayoutOptimizerTest(test.TestCase): self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) self.assertAllClose(output_val_ref, output_val, atol=1e-3) + def testReduceSumAlongCKeepDims(self): + if test.is_gpu_available(cuda_only=True): + random_seed.set_random_seed(0) + x = random_ops.truncated_normal([1, 784], seed=0) + conv = _two_layer_model(x) + reduce_sum = math_ops.reduce_sum(conv, axis=[3], keep_dims=True) + output = array_ops.identity(reduce_sum) + + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run(output, run_metadata=metadata) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + # Four transposes were initially added in the Expand phase of + # LayoutOptimizer; two of them are cancelled out in the Collapse phase. + expected_num_transposes = 2 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self._assert_trans_nchw_to_nhwc('Sum-0-0', nodes) + self.assertAllClose(output_val_ref, output_val, atol=1e-3) + + def testReduceSumAlongHKeepDims(self): + if test.is_gpu_available(cuda_only=True): + random_seed.set_random_seed(0) + x = random_ops.truncated_normal([1, 784], seed=0) + conv = _two_layer_model(x) + reduce_sum = math_ops.reduce_sum(conv, axis=[2], keep_dims=True) + output = array_ops.identity(reduce_sum) + + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run(output, run_metadata=metadata) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + # Four transposes were initially added in the Expand phase of + # LayoutOptimizer; two of them are cancelled out in the Collapse phase. + expected_num_transposes = 2 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self.assertAllClose(output_val_ref, output_val, atol=1e-3) + + def testReduceSumAlongWCKeepDims(self): + if test.is_gpu_available(cuda_only=True): + random_seed.set_random_seed(0) + x = random_ops.truncated_normal([1, 784], seed=0) + conv = _two_layer_model(x) + reduce_sum = math_ops.reduce_sum(conv, axis=[2, 3], keep_dims=True) + output = array_ops.identity(reduce_sum) + + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run(output, run_metadata=metadata) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + # Four transposes were initially added in the Expand phase of + # LayoutOptimizer; two of them are cancelled out in the Collapse phase. + expected_num_transposes = 2 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self.assertAllClose(output_val_ref, output_val, atol=1e-3) + def testConcatWithControlDependency(self): if test.is_gpu_available(cuda_only=True): random_seed.set_random_seed(0) @@ -567,7 +718,7 @@ class LayoutOptimizerTest(test.TestCase): concat = array_ops.concat([conv, conv], axis) output = array_ops.identity(concat) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -601,7 +752,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(fill) x_val = [3.4] * 784 - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={x: x_val}) with session.Session(config=_get_config()) as sess: @@ -643,7 +794,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(tile) multiple_val = [2, 3, 4, 1] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={multiple: multiple_val}) with session.Session(config=_get_config()) as sess: @@ -678,7 +829,7 @@ class LayoutOptimizerTest(test.TestCase): reverse = array_ops.reverse(conv, dims) output = array_ops.identity(reverse) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -711,7 +862,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(reverse) dims_val = [2, 3] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={dims: dims_val}) with session.Session(config=_get_config()) as sess: @@ -745,10 +896,10 @@ class LayoutOptimizerTest(test.TestCase): add = math_ops.add(conv, conv) mean = math_ops.reduce_mean(conv) condition = math_ops.less(conv, mean) - select = gen_math_ops._select(condition, conv, add) + select = gen_math_ops.select(condition, conv, add) output = array_ops.identity(select) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -768,6 +919,37 @@ class LayoutOptimizerTest(test.TestCase): self._assert_trans_nchw_to_nhwc('Select-0-0', nodes) self.assertAllClose(output_val_ref, output_val, atol=1e-3) + def testSelectOpConditionUnknownShape(self): + if test.is_gpu_available(cuda_only=True): + random_seed.set_random_seed(0) + x = random_ops.truncated_normal([1, 784], seed=0) + conv = _two_layer_model(x) + add = math_ops.add(conv, conv) + condition = array_ops.placeholder(dtype='bool') + select = gen_math_ops.select(condition, conv, add) + output = array_ops.identity(select) + + condition_val = np.zeros((1, 7, 7, 64)) + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output, feed_dict={condition: condition_val}) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run( + output, run_metadata=metadata, feed_dict={condition: condition_val}) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + expected_num_transposes = 3 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self.assertAllClose(output_val_ref, output_val, atol=1e-3) + def testSelectOpScalarCondition(self): if test.is_gpu_available(cuda_only=True): random_seed.set_random_seed(0) @@ -775,10 +957,10 @@ class LayoutOptimizerTest(test.TestCase): conv = _two_layer_model(x) add = math_ops.add(conv, conv) condition = constant_op.constant(True) - select = gen_math_ops._select(condition, conv, add) + select = gen_math_ops.select(condition, conv, add) output = array_ops.identity(select) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -808,7 +990,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(pad) paddings_val = [[1, 2], [3, 4], [5, 6], [7, 8]] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={paddings: paddings_val}) with session.Session(config=_get_config()) as sess: @@ -841,11 +1023,11 @@ class LayoutOptimizerTest(test.TestCase): conv = _two_layer_model(x) ksize = constant_op.constant([1, 2, 3, 1], shape=[4]) strides = array_ops.placeholder(dtype='int32', shape=[4]) - max_pool = gen_nn_ops._max_pool_v2(conv, ksize, strides, 'VALID') + max_pool = gen_nn_ops.max_pool_v2(conv, ksize, strides, 'VALID') output = array_ops.identity(max_pool) strides_val = [1, 3, 2, 1] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={strides: strides_val}) with session.Session(config=_get_config()) as sess: @@ -882,7 +1064,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(max_pool_grad) strides_val = [1, 3, 2, 1] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={strides: strides_val}) with session.Session(config=_get_config()) as sess: @@ -917,7 +1099,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(s) size_val = [1, 2, 3, 4] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={size: size_val}) with session.Session(config=_get_config()) as sess: @@ -953,7 +1135,7 @@ class LayoutOptimizerTest(test.TestCase): output = array_ops.identity(s) end_val = [1, 2, 3, 4] - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={end: end_val}) with session.Session(config=_get_config()) as sess: @@ -991,7 +1173,7 @@ class LayoutOptimizerTest(test.TestCase): s = conv[:, :, 1:-1, :] output = array_ops.identity(s) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -1026,7 +1208,7 @@ class LayoutOptimizerTest(test.TestCase): s = conv[:, :, :, 1:-1] output = array_ops.identity(s) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -1065,7 +1247,7 @@ class LayoutOptimizerTest(test.TestCase): [1, 2, 3, 1], s) output = array_ops.identity(s_grad) - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={end: end_val}) with session.Session(config=_get_config()) as sess: @@ -1101,7 +1283,7 @@ class LayoutOptimizerTest(test.TestCase): output = math_ops.add(shapen[0], shapen[1]) x_val = [1.7] * 784 - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output, feed_dict={x: x_val}) with session.Session(config=_get_config()) as sess: @@ -1124,11 +1306,42 @@ class LayoutOptimizerTest(test.TestCase): self._assert_vec_nchw_to_nhwc('ShapeN-0-0', nodes) self.assertAllEqual(output_val_ref, output_val) + def testShapeNFollowedByNotConvertibleNodeReshape(self): + if test.is_gpu_available(cuda_only=True): + x = array_ops.placeholder(dtype='float32') + conv = _two_layer_model(x) + conv_reshape = array_ops.reshape(conv, [1, 1, 1, -1]) + shapen = array_ops.shape_n([conv, conv_reshape]) + shape = array_ops.identity(shapen[1]) + ones = array_ops.ones(shape) + output = math_ops.add_n([conv_reshape, ones]) + + x_val = [1.7] * 784 + with session.Session(config=_get_config(False)) as sess: + output_val_ref = sess.run(output, feed_dict={x: x_val}) + + with session.Session(config=_get_config()) as sess: + metadata = config_pb2.RunMetadata() + output_val = sess.run( + output, run_metadata=metadata, feed_dict={x: x_val}) + + nodes = [] + num_transposes = 0 + for node in metadata.cost_graph.node: + if _is_transpose(node.name): + num_transposes += 1 + nodes.append(node.name) + + expected_num_transposes = 2 + self.assertEqual(expected_num_transposes, num_transposes) + self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes) + self.assertAllEqual(output_val_ref, output_val) + def testLoop(self): if test.is_gpu_available(cuda_only=True): output = _loop() - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -1155,7 +1368,7 @@ class LayoutOptimizerTest(test.TestCase): if test.is_gpu_available(cuda_only=True): output = _loop_with_branch() - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -1169,7 +1382,7 @@ class LayoutOptimizerTest(test.TestCase): num_transposes += 1 nodes.append(node.name) - expected_num_transposes = 2 + expected_num_transposes = 3 self.assertEqual(expected_num_transposes, num_transposes) self._assert_trans_nhwc_to_nchw('map/while/Conv2D-0', nodes) self._assert_trans_nchw_to_nhwc('map/while/Add-0-2', nodes) @@ -1179,7 +1392,7 @@ class LayoutOptimizerTest(test.TestCase): if test.is_gpu_available(cuda_only=True): output = _loop_with_vec_and_4d() - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: @@ -1203,7 +1416,7 @@ class LayoutOptimizerTest(test.TestCase): if test.is_gpu_available(cuda_only=True): output = _model_with_second_port() - with session.Session() as sess: + with session.Session(config=_get_config(False)) as sess: output_val_ref = sess.run(output) with session.Session(config=_get_config()) as sess: diff --git a/tensorflow/python/grappler/memory_optimizer_test.py b/tensorflow/python/grappler/memory_optimizer_test.py index 948911f099674af4c6dd19bfdac75e5fc1f75c78..4df959ce04169395589aeebaef9e3e7839e2300c 100644 --- a/tensorflow/python/grappler/memory_optimizer_test.py +++ b/tensorflow/python/grappler/memory_optimizer_test.py @@ -162,7 +162,8 @@ class MemoryOptimizerRecomputeTest(test.TestCase): arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF, memory_optimization=rewriter_config_pb2.RewriterConfig. RECOMPUTATION_HEURISTICS, - memory_optimizer_target_node_name_prefix='optimizer/gradients/'), + # Checks that name scope "gradients/" also match sub-scope. + memory_optimizer_target_node_name_scope='gradients/'), original_metagraph) self.assertGreater( len(rewritten_graph_def.node), @@ -176,6 +177,35 @@ class MemoryOptimizerRecomputeTest(test.TestCase): len([node for node in rewritten_graph_def.node if 'Recomputed/' in node.name])) + def testRewritingNameScopedGradientNamesScope(self): + """Tests that rewriting occurs with non-standard gradient names.""" + (original_metagraph, _, _, + _) = self._GetMetaGraph(optimizer_scope_name='foo/bar') + rewritten_graph_def = tf_optimizer.OptimizeGraph( + rewriter_config_pb2.RewriterConfig( + disable_model_pruning=True, + constant_folding=rewriter_config_pb2.RewriterConfig.OFF, + dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF, + layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF, + arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF, + memory_optimization=rewriter_config_pb2.RewriterConfig. + RECOMPUTATION_HEURISTICS, + # This should not match anything. + memory_optimizer_target_node_name_scope='r/gradients/'), + original_metagraph) + self.assertEqual( + len(rewritten_graph_def.node), len(original_metagraph.graph_def.node)) + self.assertEqual(0, + len([ + node for node in original_metagraph.graph_def.node + if 'Recomputed/' in node.name + ])) + self.assertEqual(0, + len([ + node for node in rewritten_graph_def.node + if 'Recomputed/' in node.name + ])) + def _GetMemoryOptimizerSessionConfig(self): rewrite_options = rewriter_config_pb2.RewriterConfig( disable_model_pruning=True, diff --git a/tensorflow/python/grappler/tf_optimizer.i b/tensorflow/python/grappler/tf_optimizer.i index 1b657983a4690dd0ddb7f569ce514b08cb10400a..de9326ccfc1653c2afd0833dcdca2cc4bfdabed5 100644 --- a/tensorflow/python/grappler/tf_optimizer.i +++ b/tensorflow/python/grappler/tf_optimizer.i @@ -100,6 +100,7 @@ PyObject* TF_OptimizeGraph( tensorflow::grappler::ItemConfig item_config; item_config.inline_functions = false; item_config.apply_optimizations = false; + item_config.ignore_user_placement = false; std::unique_ptr grappler_item = tensorflow::grappler::GrapplerItemFromMetaGraphDef(graph_id, metagraph, item_config); diff --git a/tensorflow/python/grappler/tf_optimizer_test.py b/tensorflow/python/grappler/tf_optimizer_test.py index 55dcbe2071f74204e0bbdd141560f33cefdf174d..3ee4d7807ea5677a742514eb56267b94c6b92bba 100644 --- a/tensorflow/python/grappler/tf_optimizer_test.py +++ b/tensorflow/python/grappler/tf_optimizer_test.py @@ -24,6 +24,7 @@ from tensorflow.python.framework import meta_graph from tensorflow.python.framework import ops from tensorflow.python.grappler import tf_optimizer from tensorflow.python.ops import math_ops +from tensorflow.python.ops import variables from tensorflow.python.platform import test @@ -48,6 +49,31 @@ class PyWrapOptimizeGraphTest(test.TestCase): self.assertEqual(len(graph.node), 1) self.assertItemsEqual([node.name for node in graph.node], ['d']) + def testKeepNodes(self): + g = ops.Graph() + with g.as_default(): + a1 = variables.Variable( + 1.0) # Must be preserved since it's in the collection 'variables'. + a2 = constant_op.constant(0, shape=[50, 50], name='keep') + ops.add_to_collection('a2', a2) # Explicitly add to collection. + b = constant_op.constant(1, shape=[100, 10]) + c = constant_op.constant(0, shape=[10, 30]) + d = math_ops.matmul(b, c) + ops.add_to_collection('train_op', d) # d is the fetch node. + + # Optimize the graph. + mg = meta_graph.create_meta_graph_def(graph=g) + rewriter_config = rewriter_config_pb2.RewriterConfig() + optimized_graph = tf_optimizer.OptimizeGraph(rewriter_config, mg) + + # Check that the nodes referenced in various collections have been preserved + self.assertEqual(len(optimized_graph.node), 5) + self.assertEqual(d.op.name, optimized_graph.node[0].name) + self.assertEqual(a1.op.name, optimized_graph.node[1].name) + self.assertEqual('Variable/initial_value', optimized_graph.node[2].name) + self.assertEqual(a2.op.name, optimized_graph.node[3].name) + self.assertEqual('Variable/Assign', optimized_graph.node[4].name) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/BUILD b/tensorflow/python/keras/BUILD index 61257557751359a45e3ef9f74ee6307b4c6d21dc..f74881f179115ed64ca0a177cdc7711e2f2ae3bf 100755 --- a/tensorflow/python/keras/BUILD +++ b/tensorflow/python/keras/BUILD @@ -3,10 +3,17 @@ licenses(["notice"]) # Apache 2.0 +exports_files(["LICENSE"]) + package(default_visibility = ["//visibility:public"]) load("//tensorflow:tensorflow.bzl", "py_test") +config_setting( + name = "empty_condition", + values = {"define": "UNUSED=unused"}, +) + py_library( name = "keras", srcs = [ @@ -37,8 +44,16 @@ py_library( "_impl/keras/datasets/mnist.py", "_impl/keras/datasets/reuters.py", "_impl/keras/engine/__init__.py", - "_impl/keras/engine/topology.py", + "_impl/keras/engine/base_layer.py", + "_impl/keras/engine/input_layer.py", + "_impl/keras/engine/network.py", + "_impl/keras/engine/saving.py", + "_impl/keras/engine/sequential.py", "_impl/keras/engine/training.py", + "_impl/keras/engine/training_arrays.py", + "_impl/keras/engine/training_eager.py", + "_impl/keras/engine/training_generator.py", + "_impl/keras/engine/training_utils.py", "_impl/keras/estimator.py", "_impl/keras/initializers.py", "_impl/keras/layers/__init__.py", @@ -71,8 +86,8 @@ py_library( "_impl/keras/utils/generic_utils.py", "_impl/keras/utils/io_utils.py", "_impl/keras/utils/layer_utils.py", + "_impl/keras/utils/multi_gpu_utils.py", "_impl/keras/utils/np_utils.py", - "_impl/keras/utils/training_utils.py", "_impl/keras/utils/vis_utils.py", "_impl/keras/wrappers/__init__.py", "_impl/keras/wrappers/scikit_learn.py", @@ -116,7 +131,11 @@ py_library( ], srcs_version = "PY2AND3", visibility = ["//visibility:public"], - deps = [ + deps = select({ + ":empty_condition": [], + "//conditions:default": [], + }) + [ + "@six_archive//:six", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", @@ -155,7 +174,6 @@ py_library( "//tensorflow/python/estimator", "//tensorflow/python/estimator:model_fn", "//tensorflow/python/saved_model", - "@six_archive//:six", ], ) @@ -253,6 +271,11 @@ py_test( size = "small", srcs = ["_impl/keras/metrics_test.py"], srcs_version = "PY2AND3", + tags = [ + "manual", + "no_oss", + "notap", + ], deps = [ ":keras", "//tensorflow/python:client_testlib", @@ -392,7 +415,7 @@ py_test( py_test( name = "convolutional_test", - size = "medium", + size = "large", srcs = ["_impl/keras/layers/convolutional_test.py"], srcs_version = "PY2AND3", tags = [ @@ -481,6 +504,7 @@ py_test( size = "small", srcs = ["_impl/keras/layers/normalization_test.py"], srcs_version = "PY2AND3", + tags = ["notsan"], deps = [ ":keras", "//tensorflow/python:client_testlib", @@ -612,7 +636,10 @@ py_test( size = "small", srcs = ["_impl/keras/utils/io_utils_test.py"], srcs_version = "PY2AND3", - tags = ["notsan"], + tags = [ + "no_windows", # TODO: needs investigation on Windows + "notsan", + ], deps = [ ":keras", "//tensorflow/python:client_testlib", @@ -633,9 +660,9 @@ py_test( ) py_test( - name = "training_utils_test", + name = "multi_gpu_utils_test", size = "medium", - srcs = ["_impl/keras/utils/training_utils_test.py"], + srcs = ["_impl/keras/utils/multi_gpu_utils_test.py"], srcs_version = "PY2AND3", tags = ["multi_gpu"], deps = [ @@ -719,16 +746,67 @@ py_test( ], ) +py_test( + name = "training_eager_test", + size = "medium", + srcs = ["_impl/keras/engine/training_eager_test.py"], + srcs_version = "PY2AND3", + tags = ["notsan"], + deps = [ + ":keras", + "//tensorflow/python:client_testlib", + "//third_party/py/numpy", + ], +) + +py_test( + name = "model_subclassing_test", + size = "medium", + srcs = ["_impl/keras/model_subclassing_test.py"], + srcs_version = "PY2AND3", + tags = ["notsan"], + deps = [ + ":keras", + "//tensorflow/python:client_testlib", + "//third_party/py/numpy", + ], +) + py_test( name = "topology_test", size = "small", srcs = ["_impl/keras/engine/topology_test.py"], srcs_version = "PY2AND3", + tags = [ + "no-internal-py3", + ], + deps = [ + ":keras", + "//tensorflow/python:client_testlib", + "//third_party/py/numpy", + ], +) + +py_test( + name = "saving_test", + size = "small", + srcs = ["_impl/keras/engine/saving_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":keras", + "//tensorflow/python:client_testlib", + "//third_party/py/numpy", + ], +) + +py_test( + name = "sequential_test", + size = "small", + srcs = ["_impl/keras/engine/sequential_test.py"], + srcs_version = "PY2AND3", deps = [ ":keras", - "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", "//third_party/py/numpy", ], ) @@ -749,7 +827,7 @@ py_test( py_test( name = "estimator_test", - size = "medium", + size = "large", srcs = ["_impl/keras/estimator_test.py"], srcs_version = "PY2AND3", tags = ["notsan"], diff --git a/tensorflow/python/keras/_impl/keras/__init__.py b/tensorflow/python/keras/_impl/keras/__init__.py index 73113539329c5493141db243b85254062f7b8f88..53f5d31e9c5b861c551a7a9ca3700c383ea679d7 100644 --- a/tensorflow/python/keras/_impl/keras/__init__.py +++ b/tensorflow/python/keras/_impl/keras/__init__.py @@ -40,4 +40,4 @@ from tensorflow.python.keras._impl.keras.layers import Input from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.models import Sequential -__version__ = '2.1.3-tf' +__version__ = '2.1.5-tf' diff --git a/tensorflow/python/keras/_impl/keras/activations.py b/tensorflow/python/keras/_impl/keras/activations.py index 4852b8c36ae5f475b33f12b7c7f21ae424ba242e..236e17653e1b762e1e6962f453b714d1bf7bcbf7 100644 --- a/tensorflow/python/keras/_impl/keras/activations.py +++ b/tensorflow/python/keras/_impl/keras/activations.py @@ -24,8 +24,10 @@ from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.generic_utils import deserialize_keras_object from tensorflow.python.layers.base import Layer from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.activations.softmax') def softmax(x, axis=-1): """Softmax activation function. @@ -50,10 +52,12 @@ def softmax(x, axis=-1): raise ValueError('Cannot apply softmax to a tensor that is 1D') +@tf_export('keras.activations.elu') def elu(x, alpha=1.0): return K.elu(x, alpha) +@tf_export('keras.activations.selu') def selu(x): """Scaled Exponential Linear Unit. (Klambauer et al., 2017). @@ -73,38 +77,47 @@ def selu(x): return scale * K.elu(x, alpha) +@tf_export('keras.activations.softplus') def softplus(x): return K.softplus(x) +@tf_export('keras.activations.softsign') def softsign(x): return K.softsign(x) +@tf_export('keras.activations.relu') def relu(x, alpha=0., max_value=None): return K.relu(x, alpha=alpha, max_value=max_value) +@tf_export('keras.activations.tanh') def tanh(x): return K.tanh(x) +@tf_export('keras.activations.sigmoid') def sigmoid(x): return K.sigmoid(x) +@tf_export('keras.activations.hard_sigmoid') def hard_sigmoid(x): return K.hard_sigmoid(x) +@tf_export('keras.activations.linear') def linear(x): return x +@tf_export('keras.activations.serialize') def serialize(activation): return activation.__name__ +@tf_export('keras.activations.deserialize') def deserialize(name, custom_objects=None): return deserialize_keras_object( name, @@ -113,6 +126,7 @@ def deserialize(name, custom_objects=None): printable_module_name='activation function') +@tf_export('keras.activations.get') def get(identifier): if identifier is None: return linear diff --git a/tensorflow/python/keras/_impl/keras/applications/densenet.py b/tensorflow/python/keras/_impl/keras/applications/densenet.py index 9e40d3493024c5472afca636ffc1510fad9b52d1..ca83e8691237216e799f2ca738dcb6822506e2cb 100644 --- a/tensorflow/python/keras/_impl/keras/applications/densenet.py +++ b/tensorflow/python/keras/_impl/keras/applications/densenet.py @@ -31,7 +31,7 @@ from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.applications import imagenet_utils from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import AveragePooling2D from tensorflow.python.keras._impl.keras.layers import BatchNormalization @@ -45,6 +45,7 @@ from tensorflow.python.keras._impl.keras.layers import MaxPooling2D from tensorflow.python.keras._impl.keras.layers import ZeroPadding2D from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils.data_utils import get_file +from tensorflow.python.util.tf_export import tf_export DENSENET121_WEIGHT_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.8/densenet121_weights_tf_dim_ordering_tf_kernels.h5' @@ -298,6 +299,8 @@ def DenseNet(blocks, return model +@tf_export('keras.applications.DenseNet121', + 'keras.applications.densenet.DenseNet121') def DenseNet121(include_top=True, weights='imagenet', input_tensor=None, @@ -308,6 +311,8 @@ def DenseNet121(include_top=True, input_shape, pooling, classes) +@tf_export('keras.applications.DenseNet169', + 'keras.applications.densenet.DenseNet169') def DenseNet169(include_top=True, weights='imagenet', input_tensor=None, @@ -318,6 +323,8 @@ def DenseNet169(include_top=True, input_shape, pooling, classes) +@tf_export('keras.applications.DenseNet201', + 'keras.applications.densenet.DenseNet201') def DenseNet201(include_top=True, weights='imagenet', input_tensor=None, @@ -328,6 +335,7 @@ def DenseNet201(include_top=True, input_shape, pooling, classes) +@tf_export('keras.applications.densenet.preprocess_input') def preprocess_input(x, data_format=None): """Preprocesses a numpy array encoding a batch of images. diff --git a/tensorflow/python/keras/_impl/keras/applications/imagenet_utils.py b/tensorflow/python/keras/_impl/keras/applications/imagenet_utils.py index f1f20f12a8da2cb9dcb918838a4df55cdcd04602..c26a28ed4087e30968585ec8ac0b64b51513bcae 100644 --- a/tensorflow/python/keras/_impl/keras/applications/imagenet_utils.py +++ b/tensorflow/python/keras/_impl/keras/applications/imagenet_utils.py @@ -25,6 +25,7 @@ import numpy as np from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export CLASS_INDEX = None @@ -162,6 +163,9 @@ def _preprocess_symbolic_input(x, data_format, mode): return x +@tf_export('keras.applications.resnet50.preprocess_input', + 'keras.applications.vgg19.preprocess_input', + 'keras.applications.vgg16.preprocess_input') def preprocess_input(x, data_format=None, mode='caffe'): """Preprocesses a tensor or Numpy array encoding a batch of images. @@ -193,6 +197,15 @@ def preprocess_input(x, data_format=None, mode='caffe'): return _preprocess_symbolic_input(x, data_format=data_format, mode=mode) +@tf_export('keras.applications.nasnet.decode_predictions', + 'keras.applications.resnet50.decode_predictions', + 'keras.applications.vgg19.decode_predictions', + 'keras.applications.vgg16.decode_predictions', + 'keras.applications.inception_resnet_v2.decode_predictions', + 'keras.applications.inception_v3.decode_predictions', + 'keras.applications.densenet.decode_predictions', + 'keras.applications.mobilenet.decode_predictions', + 'keras.applications.xception.decode_predictions') def decode_predictions(preds, top=5): """Decodes the prediction of an ImageNet model. @@ -221,7 +234,8 @@ def decode_predictions(preds, top=5): CLASS_INDEX_PATH, cache_subdir='models', file_hash='c2c37ea517e94d9795004a39431a14cb') - CLASS_INDEX = json.load(open(fpath)) + with open(fpath) as f: + CLASS_INDEX = json.load(f) results = [] for pred in preds: top_indices = pred.argsort()[-top:][::-1] diff --git a/tensorflow/python/keras/_impl/keras/applications/inception_resnet_v2.py b/tensorflow/python/keras/_impl/keras/applications/inception_resnet_v2.py index 1dc15b5b3437718e285d694eb058ef124ee52c0b..17e407dd58460e6d6802a3e137a96faf38a6f576 100644 --- a/tensorflow/python/keras/_impl/keras/applications/inception_resnet_v2.py +++ b/tensorflow/python/keras/_impl/keras/applications/inception_resnet_v2.py @@ -31,7 +31,7 @@ from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.applications import imagenet_utils from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import AveragePooling2D from tensorflow.python.keras._impl.keras.layers import BatchNormalization @@ -46,11 +46,13 @@ from tensorflow.python.keras._impl.keras.layers import MaxPooling2D from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export BASE_WEIGHT_URL = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.7/' +@tf_export('keras.applications.inception_resnet_v2.preprocess_input') def preprocess_input(x): """Preprocesses a numpy array encoding a batch of images. @@ -190,6 +192,8 @@ def inception_resnet_block(x, scale, block_type, block_idx, activation='relu'): return x +@tf_export('keras.applications.InceptionResNetV2', + 'keras.applications.inception_resnet_v2.InceptionResNetV2') def InceptionResNetV2(include_top=True, weights='imagenet', input_tensor=None, diff --git a/tensorflow/python/keras/_impl/keras/applications/inception_v3.py b/tensorflow/python/keras/_impl/keras/applications/inception_v3.py index ff57116f2dabed58e6993320425384ed55fde65b..2897c6058eb445ceacc34084b53dc89f556e3e9c 100644 --- a/tensorflow/python/keras/_impl/keras/applications/inception_v3.py +++ b/tensorflow/python/keras/_impl/keras/applications/inception_v3.py @@ -37,7 +37,7 @@ from tensorflow.python.keras._impl.keras import layers from tensorflow.python.keras._impl.keras.applications import imagenet_utils from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import AveragePooling2D from tensorflow.python.keras._impl.keras.layers import BatchNormalization @@ -50,6 +50,7 @@ from tensorflow.python.keras._impl.keras.layers import MaxPooling2D from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export WEIGHTS_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.5/inception_v3_weights_tf_dim_ordering_tf_kernels.h5' @@ -101,6 +102,8 @@ def conv2d_bn(x, return x +@tf_export('keras.applications.InceptionV3', + 'keras.applications.inception_v3.InceptionV3') def InceptionV3(include_top=True, weights='imagenet', input_tensor=None, @@ -399,6 +402,8 @@ def InceptionV3(include_top=True, return model +@tf_export('keras.applications.nasnet.preprocess_input', + 'keras.applications.inception_v3.preprocess_input') def preprocess_input(x): """Preprocesses a numpy array encoding a batch of images. diff --git a/tensorflow/python/keras/_impl/keras/applications/mobilenet.py b/tensorflow/python/keras/_impl/keras/applications/mobilenet.py index 790bf8ceade6039bbd651fce1960c04f9c51c63e..ad96b53a4528d99a014a0214b52a78d6a60076f8 100644 --- a/tensorflow/python/keras/_impl/keras/applications/mobilenet.py +++ b/tensorflow/python/keras/_impl/keras/applications/mobilenet.py @@ -79,8 +79,8 @@ from tensorflow.python.keras._impl.keras.applications import imagenet_utils from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions from tensorflow.python.keras._impl.keras.engine import InputSpec -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import BatchNormalization from tensorflow.python.keras._impl.keras.layers import Conv2D @@ -93,6 +93,7 @@ from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils import conv_utils from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export BASE_WEIGHT_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.6/' @@ -102,6 +103,7 @@ def relu6(x): return K.relu(x, max_value=6) +@tf_export('keras.applications.mobilenet.preprocess_input') def preprocess_input(x): """Preprocesses a numpy array encoding a batch of images. @@ -303,6 +305,8 @@ class DepthwiseConv2D(Conv2D): return config +@tf_export('keras.applications.MobileNet', + 'keras.applications.mobilenet.MobileNet') def MobileNet(input_shape=None, alpha=1.0, depth_multiplier=1, @@ -557,7 +561,7 @@ def _conv_block(inputs, filters, alpha, kernel=(3, 3), strides=(1, 1)): and width and height should be no smaller than 32. E.g. `(224, 224, 3)` would be one valid value. filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output filters in the convolution). alpha: controls the width of the network. - If `alpha` < 1.0, proportionally decreases the number of filters in each layer. @@ -623,7 +627,7 @@ def _depthwise_conv_block(inputs, (with `channels_last` data format) or (channels, rows, cols) (with `channels_first` data format). pointwise_conv_filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the pointwise convolution). + (i.e. the number of output filters in the pointwise convolution). alpha: controls the width of the network. - If `alpha` < 1.0, proportionally decreases the number of filters in each layer. diff --git a/tensorflow/python/keras/_impl/keras/applications/nasnet.py b/tensorflow/python/keras/_impl/keras/applications/nasnet.py index 5dd038c096f09422e381ba96006c3ebcd88481d9..dd33230a7eb9272f8fc60daee63e1f92574cf5e3 100644 --- a/tensorflow/python/keras/_impl/keras/applications/nasnet.py +++ b/tensorflow/python/keras/_impl/keras/applications/nasnet.py @@ -49,7 +49,7 @@ from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions from tensorflow.python.keras._impl.keras.applications.inception_v3 import preprocess_input -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import add from tensorflow.python.keras._impl.keras.layers import AveragePooling2D @@ -67,6 +67,7 @@ from tensorflow.python.keras._impl.keras.layers import ZeroPadding2D from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export NASNET_MOBILE_WEIGHT_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.8/NASNet-mobile.h5' @@ -323,6 +324,8 @@ def NASNet(input_shape=None, return model +@tf_export('keras.applications.NASNetLarge', + 'keras.applications.nasnet.NASNetLarge') def NASNetLarge(input_shape=None, include_top=True, weights='imagenet', @@ -390,6 +393,8 @@ def NASNetLarge(input_shape=None, default_size=331) +@tf_export('keras.applications.NASNetMobile', + 'keras.applications.nasnet.NASNetMobile') def NASNetMobile(input_shape=None, include_top=True, weights='imagenet', diff --git a/tensorflow/python/keras/_impl/keras/applications/resnet50.py b/tensorflow/python/keras/_impl/keras/applications/resnet50.py index 5705b3481afac4fe8f50d91c3099db5d4fb63878..46c0e635578c7f4707b027247943d75b16d703ad 100644 --- a/tensorflow/python/keras/_impl/keras/applications/resnet50.py +++ b/tensorflow/python/keras/_impl/keras/applications/resnet50.py @@ -34,7 +34,7 @@ from tensorflow.python.keras._impl.keras import layers from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions from tensorflow.python.keras._impl.keras.applications.imagenet_utils import preprocess_input -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import AveragePooling2D from tensorflow.python.keras._impl.keras.layers import BatchNormalization @@ -49,6 +49,7 @@ from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils import layer_utils from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export WEIGHTS_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.2/resnet50_weights_tf_dim_ordering_tf_kernels.h5' @@ -146,6 +147,8 @@ def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, return x +@tf_export('keras.applications.ResNet50', + 'keras.applications.resnet50.ResNet50') def ResNet50(include_top=True, weights='imagenet', input_tensor=None, diff --git a/tensorflow/python/keras/_impl/keras/applications/vgg16.py b/tensorflow/python/keras/_impl/keras/applications/vgg16.py index c91c24e6fbed100c55d38432b2338a51213f64d2..cefb25063e30505c9c34b49fd2df6eb7210d7ca8 100644 --- a/tensorflow/python/keras/_impl/keras/applications/vgg16.py +++ b/tensorflow/python/keras/_impl/keras/applications/vgg16.py @@ -32,7 +32,7 @@ from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions from tensorflow.python.keras._impl.keras.applications.imagenet_utils import preprocess_input -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Conv2D from tensorflow.python.keras._impl.keras.layers import Dense from tensorflow.python.keras._impl.keras.layers import Flatten @@ -44,12 +44,14 @@ from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils import layer_utils from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export WEIGHTS_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels.h5' WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5' +@tf_export('keras.applications.VGG16', 'keras.applications.vgg16.VGG16') def VGG16(include_top=True, weights='imagenet', input_tensor=None, diff --git a/tensorflow/python/keras/_impl/keras/applications/vgg19.py b/tensorflow/python/keras/_impl/keras/applications/vgg19.py index 223cd79d7bfb660098273e46444d9588ae10f7a3..dadaf4fdf0cc5922752c6867720c5d8cdbcab19a 100644 --- a/tensorflow/python/keras/_impl/keras/applications/vgg19.py +++ b/tensorflow/python/keras/_impl/keras/applications/vgg19.py @@ -32,7 +32,7 @@ from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions from tensorflow.python.keras._impl.keras.applications.imagenet_utils import preprocess_input -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Conv2D from tensorflow.python.keras._impl.keras.layers import Dense from tensorflow.python.keras._impl.keras.layers import Flatten @@ -44,12 +44,14 @@ from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils import layer_utils from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export WEIGHTS_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg19_weights_tf_dim_ordering_tf_kernels.h5' WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5' +@tf_export('keras.applications.VGG19', 'keras.applications.vgg19.VGG19') def VGG19(include_top=True, weights='imagenet', input_tensor=None, diff --git a/tensorflow/python/keras/_impl/keras/applications/xception.py b/tensorflow/python/keras/_impl/keras/applications/xception.py index 0a6eb4953a7ab82635f5b99ca021898876c02052..971063a16d1f5ba0e25189f1ef2f6c24eb5f8d61 100644 --- a/tensorflow/python/keras/_impl/keras/applications/xception.py +++ b/tensorflow/python/keras/_impl/keras/applications/xception.py @@ -44,7 +44,7 @@ from tensorflow.python.keras._impl.keras import layers from tensorflow.python.keras._impl.keras.applications import imagenet_utils from tensorflow.python.keras._impl.keras.applications.imagenet_utils import _obtain_input_shape from tensorflow.python.keras._impl.keras.applications.imagenet_utils import decode_predictions -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs from tensorflow.python.keras._impl.keras.layers import Activation from tensorflow.python.keras._impl.keras.layers import BatchNormalization from tensorflow.python.keras._impl.keras.layers import Conv2D @@ -57,12 +57,15 @@ from tensorflow.python.keras._impl.keras.layers import SeparableConv2D from tensorflow.python.keras._impl.keras.models import Model from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export TF_WEIGHTS_PATH = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.4/xception_weights_tf_dim_ordering_tf_kernels.h5' TF_WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.4/xception_weights_tf_dim_ordering_tf_kernels_notop.h5' +@tf_export('keras.applications.Xception', + 'keras.applications.xception.Xception') def Xception(include_top=True, weights='imagenet', input_tensor=None, @@ -328,6 +331,7 @@ def Xception(include_top=True, return model +@tf_export('keras.applications.xception.preprocess_input') def preprocess_input(x): """Preprocesses a numpy array encoding a batch of images. diff --git a/tensorflow/python/keras/_impl/keras/backend.py b/tensorflow/python/keras/_impl/keras/backend.py index 460c0dc5f39baac7b171568e6014c22eac23ccfc..04866fbe0f48dabd57fd0336258420765b24a793 100644 --- a/tensorflow/python/keras/_impl/keras/backend.py +++ b/tensorflow/python/keras/_impl/keras/backend.py @@ -29,6 +29,7 @@ import numpy as np from tensorflow.core.protobuf import config_pb2 from tensorflow.python.client import session as session_module +from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes as dtypes_module from tensorflow.python.framework import ops @@ -47,14 +48,16 @@ from tensorflow.python.ops import logging_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 resource_variable_ops from tensorflow.python.ops import sparse_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import from tensorflow.python.ops import tensor_array_ops from tensorflow.python.ops import variables as variables_module from tensorflow.python.training import moving_averages +from tensorflow.python.util import tf_contextlib from tensorflow.python.util import tf_inspect - +from tensorflow.python.util.tf_export import tf_export py_all = all py_sum = sum @@ -96,6 +99,7 @@ _IMAGE_DATA_FORMAT = 'channels_last' _LOCAL_DEVICES = None +@tf_export('keras.backend.backend') def backend(): """Publicly accessible method for determining the current backend. @@ -107,6 +111,7 @@ def backend(): return 'tensorflow' +@tf_export('keras.backend.epsilon') def epsilon(): """Returns the value of the fuzz factor used in numeric expressions. @@ -122,6 +127,7 @@ def epsilon(): return _EPSILON +@tf_export('keras.backend.set_epsilon') def set_epsilon(value): """Sets the value of the fuzz factor used in numeric expressions. @@ -142,6 +148,7 @@ def set_epsilon(value): _EPSILON = value +@tf_export('keras.backend.floatx') def floatx(): """Returns the default float type, as a string. @@ -159,6 +166,7 @@ def floatx(): return _FLOATX +@tf_export('keras.backend.set_floatx') def set_floatx(value): """Sets the default float type. @@ -184,6 +192,7 @@ def set_floatx(value): _FLOATX = str(value) +@tf_export('keras.backend.cast_to_floatx') def cast_to_floatx(x): """Cast a Numpy array to the default Keras float type. @@ -211,6 +220,7 @@ def cast_to_floatx(x): return np.asarray(x, dtype=_FLOATX) +@tf_export('keras.backend.image_data_format') def image_data_format(): """Returns the default image data format convention. @@ -226,6 +236,7 @@ def image_data_format(): return _IMAGE_DATA_FORMAT +@tf_export('keras.backend.set_image_data_format') def set_image_data_format(data_format): """Sets the value of the image data format convention. @@ -247,10 +258,11 @@ def set_image_data_format(data_format): """ global _IMAGE_DATA_FORMAT if data_format not in {'channels_last', 'channels_first'}: - raise ValueError('Unknown data_format:', data_format) + raise ValueError('Unknown data_format: ' + str(data_format)) _IMAGE_DATA_FORMAT = str(data_format) +@tf_export('keras.backend.get_uid') def get_uid(prefix=''): """Associates a string prefix with an integer counter in a TensorFlow graph. @@ -278,6 +290,7 @@ def get_uid(prefix=''): return layer_name_uids[prefix] +@tf_export('keras.backend.reset_uids') def reset_uids(): per_graph_layer_name_uids = tf_base_layers.PER_GRAPH_LAYER_NAME_UIDS keys = list(per_graph_layer_name_uids.keys()) @@ -285,6 +298,7 @@ def reset_uids(): del per_graph_layer_name_uids[key] +@tf_export('keras.backend.clear_session') def clear_session(): """Destroys the current TF graph and creates a new one. @@ -301,6 +315,7 @@ def clear_session(): _GRAPH_LEARNING_PHASES[ops.get_default_graph()] = phase +@tf_export('keras.backend.manual_variable_initialization') def manual_variable_initialization(value): """Sets the manual variable initialization flag. @@ -317,6 +332,7 @@ def manual_variable_initialization(value): _MANUAL_VAR_INIT = value +@tf_export('keras.backend.learning_phase') def learning_phase(): """Returns the learning phase flag. @@ -327,6 +343,12 @@ def learning_phase(): Returns: Learning phase (scalar integer tensor or Python integer). """ + if context.executing_eagerly(): + if 'eager' not in _GRAPH_LEARNING_PHASES: + # Fallback to inference mode as default. + return 0 + return _GRAPH_LEARNING_PHASES['eager'] + graph = ops.get_default_graph() if graph not in _GRAPH_LEARNING_PHASES: phase = array_ops.placeholder_with_default( @@ -335,6 +357,7 @@ def learning_phase(): return _GRAPH_LEARNING_PHASES[graph] +@tf_export('keras.backend.set_learning_phase') def set_learning_phase(value): """Sets the learning phase to a fixed value. @@ -346,10 +369,43 @@ def set_learning_phase(value): """ global _GRAPH_LEARNING_PHASES # pylint: disable=global-variable-not-assigned if value not in {0, 1}: - raise ValueError('Expected learning phase to be ' '0 or 1.') - _GRAPH_LEARNING_PHASES[ops.get_default_graph()] = value + raise ValueError('Expected learning phase to be 0 or 1.') + if context.executing_eagerly(): + _GRAPH_LEARNING_PHASES['eager'] = value + else: + _GRAPH_LEARNING_PHASES[ops.get_default_graph()] = value + + +@tf_contextlib.contextmanager +def learning_phase_scope(value): + """Provides a scope within which the learning phase is equal to `value`. + + The learning phase gets restored to its original value upon exiting the scope. + + Arguments: + value: Learning phase value, either 0 or 1 (integers). + + Yields: + The provided value. + + Raises: + ValueError: if `value` is neither `0` nor `1`. + """ + if value not in {0, 1}: + raise ValueError('Expected learning phase to be 0 or 1.') + previous_value = learning_phase() + try: + set_learning_phase(value) + yield value + finally: + # Restore learning phase to initial value. + if context.executing_eagerly(): + _GRAPH_LEARNING_PHASES['eager'] = previous_value + else: + _GRAPH_LEARNING_PHASES[ops.get_default_graph()] = previous_value +@tf_export('keras.backend.get_session') def get_session(): """Returns the TF session to be used by the backend. @@ -367,8 +423,9 @@ def get_session(): A TensorFlow session. """ global _SESSION - if ops.get_default_session() is not None: - session = ops.get_default_session() + default_session = ops.get_default_session() + if default_session is not None: + session = default_session else: if _SESSION is None: if not os.environ.get('OMP_NUM_THREADS'): @@ -385,6 +442,7 @@ def get_session(): return session +@tf_export('keras.backend.set_session') def set_session(session): """Sets the global TensorFlow session. @@ -438,7 +496,7 @@ def _is_current_explicit_device(device_type): """ device_type = device_type.upper() if device_type not in ['CPU', 'GPU']: - raise ValueError('device_type should be either "CPU" or "GPU".') + raise ValueError('`device_type` should be either "CPU" or "GPU".') device = _get_current_tf_device() return device is not None and device.device_type == device_type.upper() @@ -458,7 +516,7 @@ def _get_available_gpus(): def _has_nchw_support(): """Check whether the current scope supports NCHW ops. - Tensorflow does not support NCHW on CPU. Therefore we check if we are not + TensorFlow does not support NCHW on CPU. Therefore we check if we are not explicitly put on CPU, and have GPUs available. In this case there will be soft-placing on the GPU device. @@ -487,6 +545,7 @@ def _to_tensor(x, dtype): return ops.convert_to_tensor(x, dtype=dtype) +@tf_export('keras.backend.is_sparse') def is_sparse(tensor): """Returns whether a tensor is a sparse tensor. @@ -510,6 +569,7 @@ def is_sparse(tensor): return isinstance(tensor, sparse_tensor.SparseTensor) +@tf_export('keras.backend.to_dense') def to_dense(tensor): """Converts a sparse tensor into a dense tensor and returns it. @@ -539,6 +599,7 @@ def to_dense(tensor): name_scope = ops.name_scope +@tf_export('keras.backend.variable') def variable(value, dtype=None, name=None, constraint=None): """Instantiates a variable and returns it. @@ -577,7 +638,7 @@ def variable(value, dtype=None, name=None, constraint=None): v._keras_shape = sparse_coo.shape v._uses_learning_phase = False return v - v = variables_module.Variable( + v = resource_variable_ops.ResourceVariable( value, dtype=dtypes_module.as_dtype(dtype), name=name, @@ -611,6 +672,7 @@ def _initialize_variables(session): session.run(variables_module.variables_initializer(uninitialized_vars)) +@tf_export('keras.backend.constant') def constant(value, dtype=None, shape=None, name=None): """Creates a constant tensor. @@ -679,6 +741,7 @@ def is_keras_tensor(x): return hasattr(x, '_keras_history') +@tf_export('keras.backend.placeholder') def placeholder(shape=None, ndim=None, dtype=None, sparse=False, name=None): """Instantiates a placeholder tensor and returns it. @@ -731,6 +794,7 @@ def is_placeholder(x): return False +@tf_export('keras.backend.shape') def shape(x): """Returns the symbolic shape of a tensor or variable. @@ -763,6 +827,7 @@ def shape(x): return array_ops.shape(x) +@tf_export('keras.backend.int_shape') def int_shape(x): """Returns the shape of tensor or variable as a tuple of int or None entries. @@ -790,6 +855,7 @@ def int_shape(x): return None +@tf_export('keras.backend.ndim') def ndim(x): """Returns the number of axes in a tensor, as an integer. @@ -817,6 +883,7 @@ def ndim(x): return None +@tf_export('keras.backend.dtype') def dtype(x): """Returns the dtype of a Keras tensor or variable, as a string. @@ -847,6 +914,7 @@ def dtype(x): return x.dtype.base_dtype.name +@tf_export('keras.backend.eval') def eval(x): """Evaluates the value of a variable. @@ -868,6 +936,7 @@ def eval(x): return to_dense(x).eval(session=get_session()) +@tf_export('keras.backend.zeros') def zeros(shape, dtype=None, name=None): """Instantiates an all-zeros variable and returns it. @@ -900,6 +969,7 @@ def zeros(shape, dtype=None, name=None): return v +@tf_export('keras.backend.ones') def ones(shape, dtype=None, name=None): """Instantiates an all-ones variable and returns it. @@ -932,6 +1002,7 @@ def ones(shape, dtype=None, name=None): return v +@tf_export('keras.backend.eye') def eye(size, dtype=None, name=None): """Instantiate an identity matrix and returns it. @@ -960,6 +1031,7 @@ def eye(size, dtype=None, name=None): return variable(linalg_ops.eye(size, dtype=tf_dtype), dtype, name) +@tf_export('keras.backend.zeros_like') def zeros_like(x, dtype=None, name=None): """Instantiates an all-zeros variable of the same shape as another tensor. @@ -985,6 +1057,7 @@ def zeros_like(x, dtype=None, name=None): return array_ops.zeros_like(x, dtype=dtype, name=name) +@tf_export('keras.backend.ones_like') def ones_like(x, dtype=None, name=None): """Instantiates an all-ones variable of the same shape as another tensor. @@ -1023,6 +1096,7 @@ def identity(x, name=None): return array_ops.identity(x, name=name) +@tf_export('keras.backend.random_uniform_variable') def random_uniform_variable(shape, low, high, dtype=None, name=None, seed=None): """Instantiates a variable with values drawn from a uniform distribution. @@ -1059,6 +1133,7 @@ def random_uniform_variable(shape, low, high, dtype=None, name=None, seed=None): return variable(value, dtype=dtype, name=name) +@tf_export('keras.backend.random_normal_variable') def random_normal_variable(shape, mean, scale, dtype=None, name=None, seed=None): """Instantiates a variable with values drawn from a normal distribution. @@ -1096,6 +1171,7 @@ def random_normal_variable(shape, mean, scale, dtype=None, name=None, return variable(value, dtype=dtype, name=name) +@tf_export('keras.backend.count_params') def count_params(x): """Returns the static number of elements in a variable or tensor. @@ -1118,6 +1194,7 @@ def count_params(x): return np.prod(x.get_shape().as_list()) +@tf_export('keras.backend.cast') def cast(x, dtype): """Casts a tensor to a different dtype and returns it. @@ -1153,10 +1230,12 @@ def cast(x, dtype): # UPDATES OPS +@tf_export('keras.backend.update') def update(x, new_x): return state_ops.assign(x, new_x) +@tf_export('keras.backend.update_add') def update_add(x, increment): """Update the value of `x` by adding `increment`. @@ -1170,6 +1249,7 @@ def update_add(x, increment): return state_ops.assign_add(x, increment) +@tf_export('keras.backend.update_sub') def update_sub(x, decrement): """Update the value of `x` by subtracting `decrement`. @@ -1183,6 +1263,7 @@ def update_sub(x, decrement): return state_ops.assign_sub(x, decrement) +@tf_export('keras.backend.moving_average_update') def moving_average_update(x, value, momentum): """Compute the moving average of a variable. @@ -1201,6 +1282,7 @@ def moving_average_update(x, value, momentum): # LINEAR ALGEBRA +@tf_export('keras.backend.dot') def dot(x, y): """Multiplies 2 tensors (and/or variables) and returns a *tensor*. @@ -1272,6 +1354,7 @@ def dot(x, y): return out +@tf_export('keras.backend.batch_dot') def batch_dot(x, y, axes=None): """Batchwise dot product. @@ -1364,6 +1447,7 @@ def batch_dot(x, y, axes=None): return out +@tf_export('keras.backend.transpose') def transpose(x): """Transposes a tensor and returns it. @@ -1399,6 +1483,7 @@ def transpose(x): return array_ops.transpose(x) +@tf_export('keras.backend.gather') def gather(reference, indices): """Retrieves the elements of indices `indices` in the tensor `reference`. @@ -1415,6 +1500,7 @@ def gather(reference, indices): # ELEMENT-WISE OPERATIONS +@tf_export('keras.backend.max') def max(x, axis=None, keepdims=False): """Maximum value in a tensor. @@ -1432,6 +1518,7 @@ def max(x, axis=None, keepdims=False): return math_ops.reduce_max(x, axis, keepdims) +@tf_export('keras.backend.min') def min(x, axis=None, keepdims=False): """Minimum value in a tensor. @@ -1449,6 +1536,7 @@ def min(x, axis=None, keepdims=False): return math_ops.reduce_min(x, axis, keepdims) +@tf_export('keras.backend.sum') def sum(x, axis=None, keepdims=False): """Sum of the values in a tensor, alongside the specified axis. @@ -1466,6 +1554,7 @@ def sum(x, axis=None, keepdims=False): return math_ops.reduce_sum(x, axis, keepdims) +@tf_export('keras.backend.prod') def prod(x, axis=None, keepdims=False): """Multiplies the values in a tensor, alongside the specified axis. @@ -1509,6 +1598,7 @@ def cumprod(x, axis=0): return math_ops.cumprod(x, axis=axis) +@tf_export('keras.backend.var') def var(x, axis=None, keepdims=False): """Variance of a tensor, alongside the specified axis. @@ -1531,6 +1621,7 @@ def var(x, axis=None, keepdims=False): devs_squared, axis, keepdims) +@tf_export('keras.backend.std') def std(x, axis=None, keepdims=False): """Standard deviation of a tensor, alongside the specified axis. @@ -1548,6 +1639,7 @@ def std(x, axis=None, keepdims=False): return math_ops.sqrt(var(x, axis=axis, keepdims=keepdims)) +@tf_export('keras.backend.mean') def mean(x, axis=None, keepdims=False): """Mean of a tensor, alongside the specified axis. @@ -1567,6 +1659,7 @@ def mean(x, axis=None, keepdims=False): return math_ops.reduce_mean(x, axis, keepdims) +@tf_export('keras.backend.any') def any(x, axis=None, keepdims=False): """Bitwise reduction (logical OR). @@ -1582,6 +1675,7 @@ def any(x, axis=None, keepdims=False): return math_ops.reduce_any(x, axis, keepdims) +@tf_export('keras.backend.all') def all(x, axis=None, keepdims=False): """Bitwise reduction (logical AND). @@ -1597,6 +1691,7 @@ def all(x, axis=None, keepdims=False): return math_ops.reduce_all(x, axis, keepdims) +@tf_export('keras.backend.argmax') def argmax(x, axis=-1): """Returns the index of the maximum value along an axis. @@ -1610,6 +1705,7 @@ def argmax(x, axis=-1): return math_ops.argmax(x, axis) +@tf_export('keras.backend.argmin') def argmin(x, axis=-1): """Returns the index of the minimum value along an axis. @@ -1623,6 +1719,7 @@ def argmin(x, axis=-1): return math_ops.argmin(x, axis) +@tf_export('keras.backend.square') def square(x): """Element-wise square. @@ -1635,6 +1732,7 @@ def square(x): return math_ops.square(x) +@tf_export('keras.backend.abs') def abs(x): """Element-wise absolute value. @@ -1647,6 +1745,7 @@ def abs(x): return math_ops.abs(x) +@tf_export('keras.backend.sqrt') def sqrt(x): """Element-wise square root. @@ -1662,6 +1761,7 @@ def sqrt(x): return math_ops.sqrt(x) +@tf_export('keras.backend.exp') def exp(x): """Element-wise exponential. @@ -1674,6 +1774,7 @@ def exp(x): return math_ops.exp(x) +@tf_export('keras.backend.log') def log(x): """Element-wise log. @@ -1707,6 +1808,7 @@ def logsumexp(x, axis=None, keepdims=False): return math_ops.reduce_logsumexp(x, axis, keepdims) +@tf_export('keras.backend.round') def round(x): """Element-wise rounding to the closest integer. @@ -1721,6 +1823,7 @@ def round(x): return math_ops.round(x) +@tf_export('keras.backend.sign') def sign(x): """Element-wise sign. @@ -1733,6 +1836,7 @@ def sign(x): return math_ops.sign(x) +@tf_export('keras.backend.pow') def pow(x, a): """Element-wise exponentiation. @@ -1746,6 +1850,7 @@ def pow(x, a): return math_ops.pow(x, a) +@tf_export('keras.backend.clip') def clip(x, min_value, max_value): """Element-wise value clipping. @@ -1766,6 +1871,7 @@ def clip(x, min_value, max_value): return clip_ops.clip_by_value(x, min_value, max_value) +@tf_export('keras.backend.equal') def equal(x, y): """Element-wise equality between two tensors. @@ -1779,6 +1885,7 @@ def equal(x, y): return math_ops.equal(x, y) +@tf_export('keras.backend.not_equal') def not_equal(x, y): """Element-wise inequality between two tensors. @@ -1792,6 +1899,7 @@ def not_equal(x, y): return math_ops.not_equal(x, y) +@tf_export('keras.backend.greater') def greater(x, y): """Element-wise truth value of (x > y). @@ -1805,6 +1913,7 @@ def greater(x, y): return math_ops.greater(x, y) +@tf_export('keras.backend.greater_equal') def greater_equal(x, y): """Element-wise truth value of (x >= y). @@ -1818,6 +1927,7 @@ def greater_equal(x, y): return math_ops.greater_equal(x, y) +@tf_export('keras.backend.less') def less(x, y): """Element-wise truth value of (x < y). @@ -1831,6 +1941,7 @@ def less(x, y): return math_ops.less(x, y) +@tf_export('keras.backend.less_equal') def less_equal(x, y): """Element-wise truth value of (x <= y). @@ -1844,6 +1955,7 @@ def less_equal(x, y): return math_ops.less_equal(x, y) +@tf_export('keras.backend.maximum') def maximum(x, y): """Element-wise maximum of two tensors. @@ -1857,6 +1969,7 @@ def maximum(x, y): return math_ops.maximum(x, y) +@tf_export('keras.backend.minimum') def minimum(x, y): """Element-wise minimum of two tensors. @@ -1870,6 +1983,7 @@ def minimum(x, y): return math_ops.minimum(x, y) +@tf_export('keras.backend.sin') def sin(x): """Computes sin of x element-wise. @@ -1882,6 +1996,7 @@ def sin(x): return math_ops.sin(x) +@tf_export('keras.backend.cos') def cos(x): """Computes cos of x element-wise. @@ -1996,6 +2111,7 @@ def _fused_normalize_batch_in_training(x, x, gamma, beta, epsilon=epsilon, data_format=tf_data_format) +@tf_export('keras.backend.normalize_batch_in_training') def normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=1e-3): """Computes mean and std for batch then apply batch_normalization on batch. @@ -2025,6 +2141,7 @@ def normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=1e-3): x, gamma, beta, reduction_axes, epsilon=epsilon) +@tf_export('keras.backend.batch_normalization') def batch_normalization(x, mean, var, beta, gamma, epsilon=1e-3): """Applies batch normalization on x given mean, var, beta and gamma. @@ -2048,6 +2165,7 @@ def batch_normalization(x, mean, var, beta, gamma, epsilon=1e-3): # SHAPE OPERATIONS +@tf_export('keras.backend.concatenate') def concatenate(tensors, axis=-1): """Concatenates a list of tensors alongside the specified axis. @@ -2071,6 +2189,7 @@ def concatenate(tensors, axis=-1): return array_ops.concat([to_dense(x) for x in tensors], axis) +@tf_export('keras.backend.reshape') def reshape(x, shape): """Reshapes a tensor to the specified shape. @@ -2084,6 +2203,7 @@ def reshape(x, shape): return array_ops.reshape(x, shape) +@tf_export('keras.backend.permute_dimensions') def permute_dimensions(x, pattern): """Permutes axes in a tensor. @@ -2098,6 +2218,7 @@ def permute_dimensions(x, pattern): return array_ops.transpose(x, perm=pattern) +@tf_export('keras.backend.resize_images') def resize_images(x, height_factor, width_factor, data_format): """Resizes the images contained in a 4D tensor. @@ -2139,9 +2260,10 @@ def resize_images(x, height_factor, width_factor, data_format): if original_shape[2] is not None else None, None)) return x else: - raise ValueError('Invalid data_format:', data_format) + raise ValueError('Invalid data_format: ' + str(data_format)) +@tf_export('keras.backend.resize_volumes') def resize_volumes(x, depth_factor, height_factor, width_factor, data_format): """Resizes the volume contained in a 5D tensor. @@ -2170,9 +2292,10 @@ def resize_volumes(x, depth_factor, height_factor, width_factor, data_format): output = repeat_elements(output, width_factor, axis=3) return output else: - raise ValueError('Invalid data_format:', data_format) + raise ValueError('Invalid data_format: ' + str(data_format)) +@tf_export('keras.backend.repeat_elements') def repeat_elements(x, rep, axis): """Repeats the elements of a tensor along an axis, like `np.repeat`. @@ -2225,6 +2348,7 @@ def repeat_elements(x, rep, axis): return x_rep +@tf_export('keras.backend.repeat') def repeat(x, n): """Repeats a 2D tensor. @@ -2244,12 +2368,13 @@ def repeat(x, n): return array_ops.tile(x, pattern) +@tf_export('keras.backend.arange') def arange(start, stop=None, step=1, dtype='int32'): """Creates a 1D tensor containing a sequence of integers. The function arguments use the same convention as Theano's arange: if only one argument is provided, - it is in fact the "stop" argument. + it is in fact the "stop" argument and "start" is 0. The default type of the returned tensor is `'int32'` to match TensorFlow's default. @@ -2264,7 +2389,7 @@ def arange(start, stop=None, step=1, dtype='int32'): An integer tensor. """ - # Match the behavior of numpy and Theano by returning an empty seqence. + # Match the behavior of numpy and Theano by returning an empty sequence. if stop is None and start < 0: start = 0 result = math_ops.range(start, limit=stop, delta=step, name='arange') @@ -2289,6 +2414,7 @@ def tile(x, n): return array_ops.tile(x, n) +@tf_export('keras.backend.flatten') def flatten(x): """Flatten a tensor. @@ -2301,6 +2427,7 @@ def flatten(x): return array_ops.reshape(x, [-1]) +@tf_export('keras.backend.batch_flatten') def batch_flatten(x): """Turn a nD tensor into a 2D tensor with same 0th dimension. @@ -2316,6 +2443,7 @@ def batch_flatten(x): return x +@tf_export('keras.backend.expand_dims') def expand_dims(x, axis=-1): """Adds a 1-sized dimension at index "axis". @@ -2329,6 +2457,7 @@ def expand_dims(x, axis=-1): return array_ops.expand_dims(x, axis) +@tf_export('keras.backend.squeeze') def squeeze(x, axis): """Removes a 1-dimension from the tensor at index "axis". @@ -2342,6 +2471,7 @@ def squeeze(x, axis): return array_ops.squeeze(x, [axis]) +@tf_export('keras.backend.temporal_padding') def temporal_padding(x, padding=(1, 1)): """Pads the middle dimension of a 3D tensor. @@ -2358,6 +2488,7 @@ def temporal_padding(x, padding=(1, 1)): return array_ops.pad(x, pattern) +@tf_export('keras.backend.spatial_2d_padding') def spatial_2d_padding(x, padding=((1, 1), (1, 1)), data_format=None): """Pads the 2nd and 3rd dimensions of a 4D tensor. @@ -2379,7 +2510,7 @@ def spatial_2d_padding(x, padding=((1, 1), (1, 1)), data_format=None): if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) if data_format == 'channels_first': pattern = [[0, 0], [0, 0], list(padding[0]), list(padding[1])] @@ -2388,6 +2519,7 @@ def spatial_2d_padding(x, padding=((1, 1), (1, 1)), data_format=None): return array_ops.pad(x, pattern) +@tf_export('keras.backend.spatial_3d_padding') def spatial_3d_padding(x, padding=((1, 1), (1, 1), (1, 1)), data_format=None): """Pads 5D tensor with zeros along the depth, height, width dimensions. @@ -2419,7 +2551,7 @@ def spatial_3d_padding(x, padding=((1, 1), (1, 1), (1, 1)), data_format=None): if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) if data_format == 'channels_first': pattern = [[0, 0], [0, 0], [padding[0][0], padding[0][1]], @@ -2431,6 +2563,7 @@ def spatial_3d_padding(x, padding=((1, 1), (1, 1), (1, 1)), data_format=None): return array_ops.pad(x, pattern) +@tf_export('keras.backend.stack') def stack(x, axis=0): """Stacks a list of rank `R` tensors into a rank `R+1` tensor. @@ -2444,6 +2577,7 @@ def stack(x, axis=0): return array_ops.stack(x, axis=axis) +@tf_export('keras.backend.one_hot') def one_hot(indices, num_classes): """Computes the one-hot representation of an integer tensor. @@ -2462,6 +2596,7 @@ def one_hot(indices, num_classes): return array_ops.one_hot(indices, depth=num_classes, axis=-1) +@tf_export('keras.backend.reverse') def reverse(x, axes): """Reverse a tensor along the specified axes. @@ -2481,6 +2616,7 @@ def reverse(x, axes): # VALUE MANIPULATION +@tf_export('keras.backend.get_value') def get_value(x): """Returns the value of a variable. @@ -2490,9 +2626,12 @@ def get_value(x): Returns: A Numpy array. """ + if context.executing_eagerly(): + return x.numpy() return x.eval(session=get_session()) +@tf_export('keras.backend.batch_get_value') def batch_get_value(tensors): """Returns the value of more than one tensor variable. @@ -2502,12 +2641,15 @@ def batch_get_value(tensors): Returns: A list of Numpy arrays. """ + if context.executing_eagerly(): + return [x.numpy() for x in tensors] if tensors: return get_session().run(tensors) else: return [] +@tf_export('keras.backend.set_value') def set_value(x, value): """Sets the value of a variable, from a Numpy array. @@ -2517,18 +2659,22 @@ def set_value(x, value): (of the same shape). """ value = np.asarray(value, dtype=dtype(x)) - tf_dtype = dtypes_module.as_dtype(x.dtype.name.split('_')[0]) - if hasattr(x, '_assign_placeholder'): - assign_placeholder = x._assign_placeholder - assign_op = x._assign_op + if context.executing_eagerly(): + x.assign(value) else: - assign_placeholder = array_ops.placeholder(tf_dtype, shape=value.shape) - assign_op = x.assign(assign_placeholder) - x._assign_placeholder = assign_placeholder - x._assign_op = assign_op - get_session().run(assign_op, feed_dict={assign_placeholder: value}) + tf_dtype = dtypes_module.as_dtype(x.dtype.name.split('_')[0]) + if hasattr(x, '_assign_placeholder'): + assign_placeholder = x._assign_placeholder + assign_op = x._assign_op + else: + assign_placeholder = array_ops.placeholder(tf_dtype, shape=value.shape) + assign_op = x.assign(assign_placeholder) + x._assign_placeholder = assign_placeholder + x._assign_op = assign_op + get_session().run(assign_op, feed_dict={assign_placeholder: value}) +@tf_export('keras.backend.batch_set_value') def batch_set_value(tuples): """Sets the values of many tensor variables at once. @@ -2536,25 +2682,31 @@ def batch_set_value(tuples): tuples: a list of tuples `(tensor, value)`. `value` should be a Numpy array. """ - if tuples: - assign_ops = [] - feed_dict = {} + if context.executing_eagerly(): for x, value in tuples: - value = np.asarray(value, dtype=dtype(x)) - tf_dtype = dtypes_module.as_dtype(x.dtype.name.split('_')[0]) - if hasattr(x, '_assign_placeholder'): - assign_placeholder = x._assign_placeholder - assign_op = x._assign_op - else: - assign_placeholder = array_ops.placeholder(tf_dtype, shape=value.shape) - assign_op = x.assign(assign_placeholder) - x._assign_placeholder = assign_placeholder - x._assign_op = assign_op - assign_ops.append(assign_op) - feed_dict[assign_placeholder] = value - get_session().run(assign_ops, feed_dict=feed_dict) + x.assign(np.asarray(value, dtype=dtype(x))) + else: + if tuples: + assign_ops = [] + feed_dict = {} + for x, value in tuples: + value = np.asarray(value, dtype=dtype(x)) + tf_dtype = dtypes_module.as_dtype(x.dtype.name.split('_')[0]) + if hasattr(x, '_assign_placeholder'): + assign_placeholder = x._assign_placeholder + assign_op = x._assign_op + else: + assign_placeholder = array_ops.placeholder(tf_dtype, + shape=value.shape) + assign_op = x.assign(assign_placeholder) + x._assign_placeholder = assign_placeholder + x._assign_op = assign_op + assign_ops.append(assign_op) + feed_dict[assign_placeholder] = value + get_session().run(assign_ops, feed_dict=feed_dict) +@tf_export('keras.backend.print_tensor') def print_tensor(x, message=''): """Prints `message` and the tensor value when evaluated. @@ -2627,7 +2779,7 @@ class Function(object): self.updates_op = control_flow_ops.group(*updates_ops) self.name = name # additional tensor substitutions - self.feed_dict = session_kwargs.pop('feed_dict', {}) + self.feed_dict = session_kwargs.pop('feed_dict', None) # additional operations self.fetches = session_kwargs.pop('fetches', []) if not isinstance(self.fetches, list): @@ -2637,8 +2789,15 @@ class Function(object): def __call__(self, inputs): if not isinstance(inputs, (list, tuple)): raise TypeError('`inputs` should be a list or tuple.') - feed_dict = self.feed_dict.copy() + + if self.feed_dict: + feed_dict = self.feed_dict.copy() + else: + feed_dict = {} + for tensor, value in zip(self.inputs, inputs): + if value is None: + continue if is_sparse(tensor): sparse_coo = value.tocoo() indices = np.concatenate((np.expand_dims(sparse_coo.row, 1), @@ -2652,6 +2811,7 @@ class Function(object): return updated[:len(self.outputs)] +@tf_export('keras.backend.function') def function(inputs, outputs, updates=None, **kwargs): """Instantiates a Keras function. @@ -2671,12 +2831,13 @@ def function(inputs, outputs, updates=None, **kwargs): for key in kwargs: if (key not in tf_inspect.getargspec(session_module.Session.run)[0] and key not in tf_inspect.getargspec(Function.__init__)[0]): - msg = ('Invalid argument "%s" passed to K.function with Tensorflow ' + msg = ('Invalid argument "%s" passed to K.function with TensorFlow ' 'backend') % key raise ValueError(msg) return Function(inputs, outputs, updates=updates, **kwargs) +@tf_export('keras.backend.gradients') def gradients(loss, variables): """Returns the gradients of `variables` w.r.t. `loss`. @@ -2691,6 +2852,7 @@ def gradients(loss, variables): loss, variables, colocate_gradients_with_ops=True) +@tf_export('keras.backend.stop_gradient') def stop_gradient(variables): """Returns `variables` but with zero gradient w.r.t. every other variable. @@ -2711,6 +2873,7 @@ def stop_gradient(variables): # CONTROL FLOW +@tf_export('keras.backend.rnn') def rnn(step_function, inputs, initial_states, @@ -2769,6 +2932,7 @@ def rnn(step_function, ndim = len(inputs.get_shape()) if ndim < 3: raise ValueError('Input should be at least 3D.') + inputs_shape = inputs.get_shape() axes = [1, 0] + list(range(2, ndim)) inputs = array_ops.transpose(inputs, (axes)) @@ -2787,7 +2951,7 @@ def rnn(step_function, if unroll: if not inputs.get_shape()[0]: - raise ValueError('Unrolling requires a ' 'fixed number of timesteps.') + raise ValueError('Unrolling requires a fixed number of timesteps.') states = initial_states successive_states = [] successive_outputs = [] @@ -2953,10 +3117,19 @@ def rnn(step_function, axes = [1, 0] + list(range(2, len(outputs.get_shape()))) outputs = array_ops.transpose(outputs, axes) - last_output._uses_learning_phase = uses_learning_phase + + # Static shape inference: (samples, time, ...) + outputs_shape = outputs.get_shape().as_list() + outputs_shape[0] = inputs_shape[0] + outputs_shape[1] = inputs_shape[1] + outputs.set_shape(outputs_shape) + + if not context.executing_eagerly(): + last_output._uses_learning_phase = uses_learning_phase return last_output, outputs, new_states +@tf_export('keras.backend.switch') def switch(condition, then_expression, else_expression): """Switches between two operations depending on a scalar value. @@ -3020,6 +3193,7 @@ def switch(condition, then_expression, else_expression): return x +@tf_export('keras.backend.in_train_phase') def in_train_phase(x, alt, training=None): """Selects `x` in train phase, and `alt` otherwise. @@ -3063,6 +3237,7 @@ def in_train_phase(x, alt, training=None): return x +@tf_export('keras.backend.in_test_phase') def in_test_phase(x, alt, training=None): """Selects `x` in test phase, and `alt` otherwise. @@ -3086,6 +3261,7 @@ def in_test_phase(x, alt, training=None): # NN OPERATIONS +@tf_export('keras.backend.relu') def relu(x, alpha=0., max_value=None): """Rectified linear unit. @@ -3112,6 +3288,7 @@ def relu(x, alpha=0., max_value=None): return x +@tf_export('keras.backend.elu') def elu(x, alpha=1.): """Exponential linear unit. @@ -3129,6 +3306,7 @@ def elu(x, alpha=1.): return array_ops.where(x > 0, res, alpha * res) +@tf_export('keras.backend.softmax') def softmax(x): """Softmax of a tensor. @@ -3141,6 +3319,7 @@ def softmax(x): return nn.softmax(x) +@tf_export('keras.backend.softplus') def softplus(x): """Softplus of a tensor. @@ -3153,6 +3332,7 @@ def softplus(x): return nn.softplus(x) +@tf_export('keras.backend.softsign') def softsign(x): """Softsign of a tensor. @@ -3165,6 +3345,7 @@ def softsign(x): return nn.softsign(x) +@tf_export('keras.backend.categorical_crossentropy') def categorical_crossentropy(target, output, from_logits=False): """Categorical crossentropy between an output tensor and a target tensor. @@ -3183,7 +3364,7 @@ def categorical_crossentropy(target, output, from_logits=False): # expects logits, Keras expects probabilities. if not from_logits: # scale preds so that the class probas of each sample sum to 1 - output /= math_ops.reduce_sum( + output = output / math_ops.reduce_sum( # pylint: disable=g-no-augmented-assignment output, len(output.get_shape()) - 1, True) # manual computation of crossentropy epsilon_ = _to_tensor(epsilon(), output.dtype.base_dtype) @@ -3195,6 +3376,7 @@ def categorical_crossentropy(target, output, from_logits=False): return nn.softmax_cross_entropy_with_logits(labels=target, logits=output) +@tf_export('keras.backend.sparse_categorical_crossentropy') def sparse_categorical_crossentropy(target, output, from_logits=False): """Categorical crossentropy with integer targets. @@ -3228,6 +3410,7 @@ def sparse_categorical_crossentropy(target, output, from_logits=False): return res +@tf_export('keras.backend.binary_crossentropy') def binary_crossentropy(target, output, from_logits=False): """Binary crossentropy between an output tensor and a target tensor. @@ -3251,6 +3434,7 @@ def binary_crossentropy(target, output, from_logits=False): return nn.sigmoid_cross_entropy_with_logits(labels=target, logits=output) +@tf_export('keras.backend.sigmoid') def sigmoid(x): """Element-wise sigmoid. @@ -3263,6 +3447,7 @@ def sigmoid(x): return nn.sigmoid(x) +@tf_export('keras.backend.hard_sigmoid') def hard_sigmoid(x): """Segment-wise linear approximation of sigmoid. @@ -3283,6 +3468,7 @@ def hard_sigmoid(x): return x +@tf_export('keras.backend.tanh') def tanh(x): """Element-wise tanh. @@ -3295,6 +3481,7 @@ def tanh(x): return nn.tanh(x) +@tf_export('keras.backend.dropout') def dropout(x, level, noise_shape=None, seed=None): """Sets entries in `x` to zero at random, while scaling the entire tensor. @@ -3317,6 +3504,7 @@ def dropout(x, level, noise_shape=None, seed=None): return nn.dropout(x * 1., retain_prob, noise_shape, seed=seed) +@tf_export('keras.backend.l2_normalize') def l2_normalize(x, axis=None): """Normalizes a tensor wrt the L2 norm alongside the specified axis. @@ -3327,9 +3515,10 @@ def l2_normalize(x, axis=None): Returns: A tensor. """ - return nn.l2_normalize(x, dim=axis) + return nn.l2_normalize(x, axis=axis) +@tf_export('keras.backend.in_top_k') def in_top_k(predictions, targets, k): """Returns whether the `targets` are in the top `k` `predictions`. @@ -3407,7 +3596,7 @@ def _preprocess_conv3d_input(x, data_format): def _preprocess_padding(padding): - """Convert keras' padding to tensorflow's padding. + """Convert keras' padding to TensorFlow's padding. Arguments: padding: string, one of 'same' , 'valid' @@ -3423,10 +3612,11 @@ def _preprocess_padding(padding): elif padding == 'valid': padding = 'VALID' else: - raise ValueError('Invalid padding:', padding) + raise ValueError('Invalid padding: ' + str(padding)) return padding +@tf_export('keras.backend.conv1d') def conv1d(x, kernel, strides=1, @@ -3453,7 +3643,7 @@ def conv1d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) kernel_shape = kernel.get_shape().as_list() if padding == 'causal': @@ -3476,6 +3666,7 @@ def conv1d(x, return x +@tf_export('keras.backend.conv2d') def conv2d(x, kernel, strides=(1, 1), @@ -3504,7 +3695,7 @@ def conv2d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv2d_input(x, data_format) padding = _preprocess_padding(padding) @@ -3520,6 +3711,7 @@ def conv2d(x, return x +@tf_export('keras.backend.conv2d_transpose') def conv2d_transpose(x, kernel, output_shape, @@ -3550,7 +3742,7 @@ def conv2d_transpose(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) if isinstance(output_shape, (tuple, list)): output_shape = array_ops.stack(output_shape) @@ -3609,16 +3801,18 @@ def separable_conv1d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv1d_input(x, data_format) padding = _preprocess_padding(padding) + if not isinstance(strides, tuple): + strides = tuple(strides) if tf_data_format == 'NHWC': spatial_start_dim = 1 - strides = (1, 1) + strides + (1,) + strides = (1,) + strides * 2 + (1,) else: spatial_start_dim = 2 - strides = (1, 1, 1) + strides + strides = (1, 1) + strides * 2 x = array_ops.expand_dims(x, spatial_start_dim) depthwise_kernel = array_ops.expand_dims(depthwise_kernel, 0) pointwise_kernel = array_ops.expand_dims(pointwise_kernel, 0) @@ -3641,6 +3835,7 @@ def separable_conv1d(x, return x +@tf_export('keras.backend.separable_conv2d') def separable_conv2d(x, depthwise_kernel, pointwise_kernel, @@ -3670,10 +3865,12 @@ def separable_conv2d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv2d_input(x, data_format) padding = _preprocess_padding(padding) + if not isinstance(strides, tuple): + strides = tuple(strides) if tf_data_format == 'NHWC': strides = (1,) + strides + (1,) else: @@ -3719,7 +3916,7 @@ def depthwise_conv2d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv2d_input(x, data_format) padding = _preprocess_padding(padding) @@ -3740,6 +3937,7 @@ def depthwise_conv2d(x, return x +@tf_export('keras.backend.conv3d') def conv3d(x, kernel, strides=(1, 1, 1), @@ -3768,7 +3966,7 @@ def conv3d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv3d_input(x, data_format) padding = _preprocess_padding(padding) @@ -3814,7 +4012,7 @@ def conv3d_transpose(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) if isinstance(output_shape, (tuple, list)): output_shape = array_ops.stack(output_shape) @@ -3845,6 +4043,7 @@ def conv3d_transpose(x, return x +@tf_export('keras.backend.pool2d') def pool2d(x, pool_size, strides=(1, 1), @@ -3872,7 +4071,7 @@ def pool2d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv2d_input(x, data_format) padding = _preprocess_padding(padding) @@ -3890,13 +4089,14 @@ def pool2d(x, x = nn.avg_pool( x, pool_size, strides, padding=padding, data_format=tf_data_format) else: - raise ValueError('Invalid pooling mode:', pool_mode) + raise ValueError('Invalid pooling mode: ' + str(pool_mode)) if data_format == 'channels_first' and tf_data_format == 'NHWC': x = array_ops.transpose(x, (0, 3, 1, 2)) # NHWC -> NCHW return x +@tf_export('keras.backend.pool3d') def pool3d(x, pool_size, strides=(1, 1, 1), @@ -3924,7 +4124,7 @@ def pool3d(x, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) x, tf_data_format = _preprocess_conv3d_input(x, data_format) padding = _preprocess_padding(padding) @@ -3942,7 +4142,7 @@ def pool3d(x, x = nn.avg_pool3d( x, pool_size, strides, padding=padding, data_format=tf_data_format) else: - raise ValueError('Invalid pooling mode:', pool_mode) + raise ValueError('Invalid pooling mode: ' + str(pool_mode)) if data_format == 'channels_first' and tf_data_format == 'NDHWC': x = array_ops.transpose(x, (0, 4, 1, 2, 3)) @@ -3973,7 +4173,7 @@ def local_conv1d(inputs, kernel, kernel_size, strides, data_format=None): if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) stride = strides[0] kernel_shape = int_shape(kernel) @@ -4029,7 +4229,7 @@ def local_conv2d(inputs, if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) stride_row, stride_col = strides output_row, output_col = output_shape @@ -4060,6 +4260,7 @@ def local_conv2d(inputs, return output +@tf_export('keras.backend.bias_add') def bias_add(x, bias, data_format=None): """Adds a bias vector to a tensor. @@ -4081,56 +4282,59 @@ def bias_add(x, bias, data_format=None): if data_format is None: data_format = image_data_format() if data_format not in {'channels_first', 'channels_last'}: - raise ValueError('Unknown data_format ' + str(data_format)) + raise ValueError('Unknown data_format: ' + str(data_format)) bias_shape = int_shape(bias) if len(bias_shape) != 1 and len(bias_shape) != ndim(x) - 1: raise ValueError( 'Unexpected bias dimensions %d, expect to be 1 or %d dimensions' % (len(bias_shape), ndim(x))) + # pylint: disable=g-no-augmented-assignment if ndim(x) == 5: if data_format == 'channels_first': if len(bias_shape) == 1: - x += reshape(bias, (1, bias_shape[0], 1, 1, 1)) + x = x + reshape(bias, (1, bias_shape[0], 1, 1, 1)) else: - x += reshape(bias, (1, bias_shape[3]) + bias_shape[:3]) + x = x + reshape(bias, (1, bias_shape[3]) + bias_shape[:3]) elif data_format == 'channels_last': if len(bias_shape) == 1: - x += reshape(bias, (1, 1, 1, bias_shape[0])) + x = x + reshape(bias, (1, 1, 1, bias_shape[0])) else: - x += reshape(bias, (1,) + bias_shape) + x = x + reshape(bias, (1,) + bias_shape) elif ndim(x) == 4: if data_format == 'channels_first': if len(bias_shape) == 1: if _has_nchw_support(): x = nn.bias_add(x, bias, data_format='NCHW') else: - x += reshape(bias, (1, bias_shape[0], 1, 1)) + x = x + reshape(bias, (1, bias_shape[0], 1, 1)) else: - x += reshape(bias, (1, bias_shape[2]) + bias_shape[:2]) + x = x + reshape(bias, (1, bias_shape[2]) + bias_shape[:2]) elif data_format == 'channels_last': if len(bias_shape) == 1: x = nn.bias_add(x, bias, data_format='NHWC') else: - x += reshape(bias, (1,) + bias_shape) + x = x + reshape(bias, (1,) + bias_shape) elif ndim(x) == 3: if data_format == 'channels_first': if len(bias_shape) == 1: - x += reshape(bias, (1, bias_shape[0], 1)) + x = x + reshape(bias, (1, bias_shape[0], 1)) else: - x += reshape(bias, (1, bias_shape[1], bias_shape[0])) + x = x + reshape(bias, (1, bias_shape[1], bias_shape[0])) elif data_format == 'channels_last': if len(bias_shape) == 1: - x += reshape(bias, (1, 1, bias_shape[0])) + x = x + reshape(bias, (1, 1, bias_shape[0])) else: - x += reshape(bias, (1,) + bias_shape) + x = x + reshape(bias, (1,) + bias_shape) else: x = nn.bias_add(x, bias) + # pylint: enable=g-no-augmented-assignment return x # RANDOMNESS +@tf_export('keras.backend.random_normal') def random_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None): """Returns a tensor with normal distribution of values. @@ -4153,6 +4357,7 @@ def random_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None): shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed) +@tf_export('keras.backend.random_uniform') def random_uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None): """Returns a tensor with uniform distribution of values. @@ -4176,6 +4381,7 @@ def random_uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None): shape, minval=minval, maxval=maxval, dtype=dtype, seed=seed) +@tf_export('keras.backend.random_binomial') def random_binomial(shape, p=0.0, dtype=None, seed=None): """Returns a tensor with random binomial distribution of values. @@ -4197,6 +4403,7 @@ def random_binomial(shape, p=0.0, dtype=None, seed=None): array_ops.ones(shape, dtype=dtype), array_ops.zeros(shape, dtype=dtype)) +@tf_export('keras.backend.truncated_normal') def truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None): """Returns a tensor with truncated random normal distribution of values. @@ -4230,6 +4437,7 @@ def truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None): # in TensorFlow's CTC implementation +@tf_export('keras.backend.ctc_label_dense_to_sparse') def ctc_label_dense_to_sparse(labels, label_lengths): """Converts CTC labels from dense to sparse. @@ -4274,6 +4482,7 @@ def ctc_label_dense_to_sparse(labels, label_lengths): math_ops.to_int64(indices), vals_sparse, math_ops.to_int64(label_shape)) +@tf_export('keras.backend.ctc_batch_cost') def ctc_batch_cost(y_true, y_pred, input_length, label_length): """Runs CTC loss algorithm on each batch element. @@ -4303,6 +4512,7 @@ def ctc_batch_cost(y_true, y_pred, input_length, label_length): inputs=y_pred, labels=sparse_labels, sequence_length=input_length), 1) +@tf_export('keras.backend.ctc_decode') def ctc_decode(y_pred, input_length, greedy=True, beam_width=100, top_paths=1): """Decodes the output of a softmax. @@ -4354,6 +4564,7 @@ def ctc_decode(y_pred, input_length, greedy=True, beam_width=100, top_paths=1): # HIGH ORDER FUNCTIONS +@tf_export('keras.backend.map_fn') def map_fn(fn, elems, name=None, dtype=None): """Map the function fn over the elements elems and return the outputs. @@ -4369,6 +4580,7 @@ def map_fn(fn, elems, name=None, dtype=None): return functional_ops.map_fn(fn, elems, name=name, dtype=dtype) +@tf_export('keras.backend.foldl') def foldl(fn, elems, initializer=None, name=None): """Reduce elems using fn to combine them from left to right. @@ -4385,6 +4597,7 @@ def foldl(fn, elems, initializer=None, name=None): return functional_ops.foldl(fn, elems, initializer=initializer, name=name) +@tf_export('keras.backend.foldr') def foldr(fn, elems, initializer=None, name=None): """Reduce elems using fn to combine them from right to left. diff --git a/tensorflow/python/keras/_impl/keras/backend_test.py b/tensorflow/python/keras/_impl/keras/backend_test.py index 27833e368d1e1c6938b83d08553fca9f4c3669a2..fb4b2a0e1dc06c904d4b93038840dbf688d42ed4 100644 --- a/tensorflow/python/keras/_impl/keras/backend_test.py +++ b/tensorflow/python/keras/_impl/keras/backend_test.py @@ -128,6 +128,22 @@ class BackendUtilsTest(test.TestCase): sess.run(variables.global_variables_initializer()) sess.run(y, feed_dict={x: np.random.random((2, 3))}) + def test_learning_phase_scope(self): + with self.test_session(): + initial_learning_phase = keras.backend.learning_phase() + with keras.backend.learning_phase_scope(1) as lp: + self.assertEqual(lp, 1) + self.assertEqual(keras.backend.learning_phase(), 1) + self.assertEqual(keras.backend.learning_phase(), initial_learning_phase) + with keras.backend.learning_phase_scope(0) as lp: + self.assertEqual(lp, 0) + self.assertEqual(keras.backend.learning_phase(), 0) + self.assertEqual(keras.backend.learning_phase(), initial_learning_phase) + with self.assertRaises(ValueError): + with keras.backend.learning_phase_scope(None): + pass + self.assertEqual(keras.backend.learning_phase(), initial_learning_phase) + def test_int_shape(self): x = keras.backend.placeholder(shape=(3, 4)) self.assertEqual(keras.backend.int_shape(x), (3, 4)) @@ -915,6 +931,15 @@ class BackendNNOpsTest(test.TestCase): last_output, outputs, new_states = keras.backend.rnn(rnn_fn, inputs, initial_states, **kwargs) + # check static shape inference + self.assertEquals(last_output.get_shape().as_list(), + [num_samples, output_dim]) + self.assertEquals(outputs.get_shape().as_list(), + [num_samples, timesteps, output_dim]) + for state in new_states: + self.assertEquals(state.get_shape().as_list(), + [num_samples, output_dim]) + last_output_list[i].append(keras.backend.eval(last_output)) outputs_list[i].append(keras.backend.eval(outputs)) self.assertEqual(len(new_states), 1) diff --git a/tensorflow/python/keras/_impl/keras/callbacks.py b/tensorflow/python/keras/_impl/keras/callbacks.py index f0d9e0b0f522d344d9e04365f04655f00eea787f..deb1e8867dba3d52816ebda02bd9a3bf2ec7bc09 100644 --- a/tensorflow/python/keras/_impl/keras/callbacks.py +++ b/tensorflow/python/keras/_impl/keras/callbacks.py @@ -35,6 +35,7 @@ from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar from tensorflow.python.ops import array_ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.summary import summary as tf_summary +from tensorflow.python.util.tf_export import tf_export try: @@ -159,10 +160,11 @@ class CallbackList(object): return iter(self.callbacks) +@tf_export('keras.callbacks.Callback') class Callback(object): """Abstract base class used to build new callbacks. - # Properties + Attributes: params: dict. Training parameters (eg. verbosity, batch size, number of epochs...). model: instance of `keras.models.Model`. @@ -215,12 +217,23 @@ class Callback(object): pass +@tf_export('keras.callbacks.BaseLogger') class BaseLogger(Callback): """Callback that accumulates epoch averages of metrics. This callback is automatically applied to every Keras model. + + Arguments: + stateful_metrics: Iterable of string names of metrics that + should *not* be averaged over an epoch. + Metrics in this list will be logged as-is in `on_epoch_end`. + All others will be averaged in `on_epoch_end`. """ + def __init__(self, stateful_metrics=None): + super(BaseLogger, self).__init__() + self.stateful_metrics = set(stateful_metrics or []) + def on_epoch_begin(self, epoch, logs=None): self.seen = 0 self.totals = {} @@ -231,19 +244,26 @@ class BaseLogger(Callback): self.seen += batch_size for k, v in logs.items(): - if k in self.totals: - self.totals[k] += v * batch_size + if k in self.stateful_metrics: + self.totals[k] = v else: - self.totals[k] = v * batch_size + if k in self.totals: + self.totals[k] += v * batch_size + else: + self.totals[k] = v * batch_size def on_epoch_end(self, epoch, logs=None): if logs is not None: for k in self.params['metrics']: if k in self.totals: # Make value available to next callbacks. - logs[k] = self.totals[k] / self.seen + if k in self.stateful_metrics: + logs[k] = self.totals[k] + else: + logs[k] = self.totals[k] / self.seen +@tf_export('keras.callbacks.TerminateOnNaN') class TerminateOnNaN(Callback): """Callback that terminates training when a NaN loss is encountered. """ @@ -260,6 +280,7 @@ class TerminateOnNaN(Callback): self.model.stop_training = True +@tf_export('keras.callbacks.ProgbarLogger') class ProgbarLogger(Callback): """Callback that prints metrics to stdout. @@ -267,12 +288,16 @@ class ProgbarLogger(Callback): count_mode: One of "steps" or "samples". Whether the progress bar should count samples seen or steps (batches) seen. + stateful_metrics: Iterable of string names of metrics that + should *not* be averaged over an epoch. + Metrics in this list will be logged as-is. + All others will be averaged over time (e.g. loss, etc). Raises: ValueError: In case of invalid `count_mode`. """ - def __init__(self, count_mode='samples'): + def __init__(self, count_mode='samples', stateful_metrics=None): super(ProgbarLogger, self).__init__() if count_mode == 'samples': self.use_steps = False @@ -280,6 +305,7 @@ class ProgbarLogger(Callback): self.use_steps = True else: raise ValueError('Unknown `count_mode`: ' + str(count_mode)) + self.stateful_metrics = set(stateful_metrics or []) def on_train_begin(self, logs=None): self.verbose = self.params['verbose'] @@ -293,7 +319,10 @@ class ProgbarLogger(Callback): else: target = self.params['samples'] self.target = target - self.progbar = Progbar(target=self.target, verbose=self.verbose) + self.progbar = Progbar( + target=self.target, + verbose=self.verbose, + stateful_metrics=self.stateful_metrics) self.seen = 0 def on_batch_begin(self, batch, logs=None): @@ -323,9 +352,10 @@ class ProgbarLogger(Callback): if k in logs: self.log_values.append((k, logs[k])) if self.verbose: - self.progbar.update(self.seen, self.log_values, force=True) + self.progbar.update(self.seen, self.log_values) +@tf_export('keras.callbacks.History') class History(Callback): """Callback that records events into a `History` object. @@ -345,6 +375,7 @@ class History(Callback): self.history.setdefault(k, []).append(v) +@tf_export('keras.callbacks.ModelCheckpoint') class ModelCheckpoint(Callback): """Save the model after every epoch. @@ -448,6 +479,7 @@ class ModelCheckpoint(Callback): self.model.save(filepath, overwrite=True) +@tf_export('keras.callbacks.EarlyStopping') class EarlyStopping(Callback): """Stop training when a monitored quantity has stopped improving. @@ -531,6 +563,7 @@ class EarlyStopping(Callback): print('Epoch %05d: early stopping' % (self.stopped_epoch + 1)) +@tf_export('keras.callbacks.RemoteMonitor') class RemoteMonitor(Callback): """Callback used to stream events to a server. @@ -575,6 +608,7 @@ class RemoteMonitor(Callback): 'root server at ' + str(self.root)) +@tf_export('keras.callbacks.LearningRateScheduler') class LearningRateScheduler(Callback): """Learning rate scheduler. @@ -603,6 +637,7 @@ class LearningRateScheduler(Callback): 'rate to %s.' % (epoch + 1, lr)) +@tf_export('keras.callbacks.TensorBoard') class TensorBoard(Callback): # pylint: disable=line-too-long """Tensorboard basic visualizations. @@ -743,16 +778,24 @@ class TensorBoard(Callback): while i < val_size: step = min(self.batch_size, val_size - i) batch_val = [] - batch_val.append(val_data[0][i:i + step]) - batch_val.append(val_data[1][i:i + step]) - batch_val.append(val_data[2][i:i + step]) + batch_val.append(val_data[0][i:i + step] + if val_data[0] is not None else None) + batch_val.append(val_data[1][i:i + step] + if val_data[1] is not None else None) + batch_val.append(val_data[2][i:i + step] + if val_data[2] is not None else None) if self.model.uses_learning_phase: # do not slice the learning phase - batch_val = [x[i:i + step] for x in val_data[:-1]] + batch_val = [x[i:i + step] if x is not None else None + for x in val_data[:-1]] batch_val.append(val_data[-1]) else: - batch_val = [x[i:i + step] for x in val_data] - feed_dict = dict(zip(tensors, batch_val)) + batch_val = [x[i:i + step] if x is not None else None + for x in val_data] + feed_dict = {} + for key, val in zip(tensors, batch_val): + if val is not None: + feed_dict[key] = val result = self.sess.run([self.merged], feed_dict=feed_dict) summary_str = result[0] self.writer.add_summary(summary_str, epoch) @@ -772,6 +815,7 @@ class TensorBoard(Callback): self.writer.close() +@tf_export('keras.callbacks.ReduceLROnPlateau') class ReduceLROnPlateau(Callback): """Reduce learning rate when a metric has stopped improving. @@ -891,6 +935,7 @@ class ReduceLROnPlateau(Callback): return self.cooldown_counter > 0 +@tf_export('keras.callbacks.CSVLogger') class CSVLogger(Callback): """Callback that streams epoch results to a csv file. @@ -971,6 +1016,7 @@ class CSVLogger(Callback): self.writer = None +@tf_export('keras.callbacks.LambdaCallback') class LambdaCallback(Callback): r"""Callback for creating simple, custom callbacks on-the-fly. diff --git a/tensorflow/python/keras/_impl/keras/constraints.py b/tensorflow/python/keras/_impl/keras/constraints.py index 4b051c93f3aaca18e10df666966c66b4dce28111..271fbbb63d3dfd50507837e190860d48315a14f2 100644 --- a/tensorflow/python/keras/_impl/keras/constraints.py +++ b/tensorflow/python/keras/_impl/keras/constraints.py @@ -24,8 +24,10 @@ import six from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.generic_utils import deserialize_keras_object from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_keras_object +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.constraints.Constraint') class Constraint(object): def __call__(self, w): @@ -35,6 +37,7 @@ class Constraint(object): return {} +@tf_export('keras.constraints.MaxNorm', 'keras.constraints.max_norm') class MaxNorm(Constraint): """MaxNorm weight constraint. @@ -64,22 +67,22 @@ class MaxNorm(Constraint): def __call__(self, w): norms = K.sqrt(K.sum(K.square(w), axis=self.axis, keepdims=True)) desired = K.clip(norms, 0, self.max_value) - w *= (desired / (K.epsilon() + norms)) - return w + return w * (desired / (K.epsilon() + norms)) def get_config(self): return {'max_value': self.max_value, 'axis': self.axis} +@tf_export('keras.constraints.NonNeg', 'keras.constraints.non_neg') class NonNeg(Constraint): """Constrains the weights to be non-negative. """ def __call__(self, w): - w *= K.cast(K.greater_equal(w, 0.), K.floatx()) - return w + return w * K.cast(K.greater_equal(w, 0.), K.floatx()) +@tf_export('keras.constraints.UnitNorm', 'keras.constraints.unit_norm') class UnitNorm(Constraint): """Constrains the weights incident to each hidden unit to have unit norm. @@ -108,6 +111,7 @@ class UnitNorm(Constraint): return {'axis': self.axis} +@tf_export('keras.constraints.MinMaxNorm', 'keras.constraints.min_max_norm') class MinMaxNorm(Constraint): """MinMaxNorm weight constraint. @@ -148,8 +152,7 @@ class MinMaxNorm(Constraint): desired = ( self.rate * K.clip(norms, self.min_value, self.max_value) + (1 - self.rate) * norms) - w *= (desired / (K.epsilon() + norms)) - return w + return w * (desired / (K.epsilon() + norms)) def get_config(self): return { @@ -173,10 +176,12 @@ nonneg = non_neg unitnorm = unit_norm +@tf_export('keras.constraints.serialize') def serialize(constraint): return serialize_keras_object(constraint) +@tf_export('keras.constraints.deserialize') def deserialize(config, custom_objects=None): return deserialize_keras_object( config, @@ -185,6 +190,7 @@ def deserialize(config, custom_objects=None): printable_module_name='constraint') +@tf_export('keras.constraints.get') def get(identifier): if identifier is None: return None @@ -196,4 +202,5 @@ def get(identifier): elif callable(identifier): return identifier else: - raise ValueError('Could not interpret constraint identifier:', identifier) + raise ValueError('Could not interpret constraint identifier: ' + + str(identifier)) diff --git a/tensorflow/python/keras/_impl/keras/datasets/boston_housing.py b/tensorflow/python/keras/_impl/keras/datasets/boston_housing.py index cfd7df61d5ea47b810776ac8da1bdfcff77d177f..13fa9aed2b8da124af4e9f68c779e08d3094cb5d 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/boston_housing.py +++ b/tensorflow/python/keras/_impl/keras/datasets/boston_housing.py @@ -21,8 +21,10 @@ from __future__ import print_function import numpy as np from tensorflow.python.keras._impl.keras.utils.data_utils import get_file +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.boston_housing.load_data') def load_data(path='boston_housing.npz', test_split=0.2, seed=113): """Loads the Boston Housing dataset. diff --git a/tensorflow/python/keras/_impl/keras/datasets/cifar.py b/tensorflow/python/keras/_impl/keras/datasets/cifar.py index 7ada3340a59e114d73095068ec476da5973b67fb..02344897f774723d0ad690ae641952cb63022bdf 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/cifar.py +++ b/tensorflow/python/keras/_impl/keras/datasets/cifar.py @@ -34,17 +34,16 @@ def load_batch(fpath, label_key='labels'): Returns: A tuple `(data, labels)`. """ - f = open(fpath, 'rb') - if sys.version_info < (3,): - d = cPickle.load(f) - else: - d = cPickle.load(f, encoding='bytes') - # decode utf8 - d_decoded = {} - for k, v in d.items(): - d_decoded[k.decode('utf8')] = v - d = d_decoded - f.close() + with open(fpath, 'rb') as f: + if sys.version_info < (3,): + d = cPickle.load(f) + else: + d = cPickle.load(f, encoding='bytes') + # decode utf8 + d_decoded = {} + for k, v in d.items(): + d_decoded[k.decode('utf8')] = v + d = d_decoded data = d['data'] labels = d[label_key] diff --git a/tensorflow/python/keras/_impl/keras/datasets/cifar10.py b/tensorflow/python/keras/_impl/keras/datasets/cifar10.py index fb9d98d42cca9c98d6e9ea46782d1e3d31a4d7dc..6b772433822474c06efcce1701226a4a67abe361 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/cifar10.py +++ b/tensorflow/python/keras/_impl/keras/datasets/cifar10.py @@ -25,8 +25,10 @@ import numpy as np from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.datasets.cifar import load_batch from tensorflow.python.keras._impl.keras.utils.data_utils import get_file +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.cifar10.load_data') def load_data(): """Loads CIFAR10 dataset. diff --git a/tensorflow/python/keras/_impl/keras/datasets/cifar100.py b/tensorflow/python/keras/_impl/keras/datasets/cifar100.py index 95aace599a721618eaae51c89f05de01653c682d..28d74116a50979abab207dbec88e384210dfc070 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/cifar100.py +++ b/tensorflow/python/keras/_impl/keras/datasets/cifar100.py @@ -25,8 +25,10 @@ import numpy as np from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.datasets.cifar import load_batch from tensorflow.python.keras._impl.keras.utils.data_utils import get_file +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.cifar100.load_data') def load_data(label_mode='fine'): """Loads CIFAR100 dataset. diff --git a/tensorflow/python/keras/_impl/keras/datasets/fashion_mnist.py b/tensorflow/python/keras/_impl/keras/datasets/fashion_mnist.py index b9ae41a0d4d0e8d9df70e3fc1952e81c5f57e8d9..508e95f719a02977960b80c283495ced642293c5 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/fashion_mnist.py +++ b/tensorflow/python/keras/_impl/keras/datasets/fashion_mnist.py @@ -24,8 +24,10 @@ import os import numpy as np from tensorflow.python.keras._impl.keras.utils.data_utils import get_file +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.fashion_mnist.load_data') def load_data(): """Loads the Fashion-MNIST dataset. diff --git a/tensorflow/python/keras/_impl/keras/datasets/imdb.py b/tensorflow/python/keras/_impl/keras/datasets/imdb.py index 880c9c821b832caed4a3072c094d72a3171f7a63..7467bb24646227705972262381aa5cf1de809f1c 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/imdb.py +++ b/tensorflow/python/keras/_impl/keras/datasets/imdb.py @@ -25,8 +25,10 @@ import numpy as np from tensorflow.python.keras._impl.keras.preprocessing.sequence import _remove_long_seq from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.imdb.load_data') def load_data(path='imdb.npz', num_words=None, skip_top=0, @@ -128,6 +130,7 @@ def load_data(path='imdb.npz', return (x_train, y_train), (x_test, y_test) +@tf_export('keras.datasets.imdb.get_word_index') def get_word_index(path='imdb_word_index.json'): """Retrieves the dictionary mapping word indices back to words. @@ -141,7 +144,5 @@ def get_word_index(path='imdb_word_index.json'): path, origin='https://s3.amazonaws.com/text-datasets/imdb_word_index.json', file_hash='bfafd718b763782e994055a2d397834f') - f = open(path) - data = json.load(f) - f.close() - return data + with open(path) as f: + return json.load(f) diff --git a/tensorflow/python/keras/_impl/keras/datasets/mnist.py b/tensorflow/python/keras/_impl/keras/datasets/mnist.py index ec12a31dcf07cad6ec076aad9bd1f671236cf457..e30691373e9aafad61b101476e21d6860527ce98 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/mnist.py +++ b/tensorflow/python/keras/_impl/keras/datasets/mnist.py @@ -21,8 +21,10 @@ from __future__ import print_function import numpy as np from tensorflow.python.keras._impl.keras.utils.data_utils import get_file +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.mnist.load_data') def load_data(path='mnist.npz'): """Loads the MNIST dataset. diff --git a/tensorflow/python/keras/_impl/keras/datasets/reuters.py b/tensorflow/python/keras/_impl/keras/datasets/reuters.py index 95cf8852a9c6b1866573231a833f9a95af7b0d55..b711696b5eecf9ba07a66cef25c1811c182b3b60 100644 --- a/tensorflow/python/keras/_impl/keras/datasets/reuters.py +++ b/tensorflow/python/keras/_impl/keras/datasets/reuters.py @@ -25,8 +25,10 @@ import numpy as np from tensorflow.python.keras._impl.keras.preprocessing.sequence import _remove_long_seq from tensorflow.python.keras._impl.keras.utils.data_utils import get_file from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.datasets.reuters.load_data') def load_data(path='reuters.npz', num_words=None, skip_top=0, @@ -112,6 +114,7 @@ def load_data(path='reuters.npz', return (x_train, y_train), (x_test, y_test) +@tf_export('keras.datasets.reuters.get_word_index') def get_word_index(path='reuters_word_index.json'): """Retrieves the dictionary mapping word indices back to words. diff --git a/tensorflow/python/keras/_impl/keras/engine/__init__.py b/tensorflow/python/keras/_impl/keras/engine/__init__.py index 31f624f9af65cac60b6466d4eb5753cbdee984c6..1bc533ab8f7ba37948d82bc69fe1c9bfe00d6834 100644 --- a/tensorflow/python/keras/_impl/keras/engine/__init__.py +++ b/tensorflow/python/keras/_impl/keras/engine/__init__.py @@ -18,13 +18,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.keras._impl.keras.engine.topology import get_source_inputs -from tensorflow.python.keras._impl.keras.engine.topology import Input -from tensorflow.python.keras._impl.keras.engine.topology import InputLayer -from tensorflow.python.keras._impl.keras.engine.topology import InputSpec -from tensorflow.python.keras._impl.keras.engine.topology import Layer +from tensorflow.python.keras._impl.keras.engine.base_layer import InputSpec +from tensorflow.python.keras._impl.keras.engine.base_layer import Layer +from tensorflow.python.keras._impl.keras.engine.input_layer import Input +from tensorflow.python.keras._impl.keras.engine.input_layer import InputLayer +from tensorflow.python.keras._impl.keras.engine.network import get_source_inputs +from tensorflow.python.keras._impl.keras.engine.network import Network from tensorflow.python.keras._impl.keras.engine.training import Model - - -# Note: topology.Node is an internal class, -# it isn't meant to be used by Keras users. diff --git a/tensorflow/python/keras/_impl/keras/engine/base_layer.py b/tensorflow/python/keras/_impl/keras/engine/base_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..5615241ae3077102ef40f9c0619161964a62a335 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/base_layer.py @@ -0,0 +1,505 @@ +# 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. +# ============================================================================== +# pylint: disable=protected-access +"""Base layer code (`Layer`). +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from six.moves import zip # pylint: disable=redefined-builtin + +from tensorflow.python.eager import context +from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import constraints +from tensorflow.python.keras._impl.keras import initializers +from tensorflow.python.keras._impl.keras import regularizers +from tensorflow.python.keras._impl.keras.utils import generic_utils +from tensorflow.python.layers import base as tf_base_layers +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export + + +# pylint: disable=invalid-name +InputSpec = tf_base_layers.InputSpec +Node = tf_base_layers.Node +TFBaseLayer = tf_base_layers.Layer +# pylint: enable=invalid-name + + +@tf_export('keras.layers.Layer') +class Layer(tf_base_layers.Layer): + """Abstract base layer class. + + # Properties + name: String, must be unique within a model. + input_spec: List of InputSpec class instances + each entry describes one required input: + - ndim + - dtype + A layer with `n` input tensors must have + an `input_spec` of length `n`. + trainable: Boolean, whether the layer weights + will be updated during training. + uses_learning_phase: Whether any operation + of the layer uses `K.in_training_phase()` + or `K.in_test_phase()`. + input_shape: Shape tuple. Provided for convenience, + but note that there may be cases in which this + attribute is ill-defined (e.g. a shared layer + with multiple input shapes), in which case + requesting `input_shape` will raise an Exception. + Prefer using `layer.get_input_shape_for(input_shape)`, + or `layer.get_input_shape_at(node_index)`. + output_shape: Shape tuple. See above. + inbound_nodes: List of nodes. + outbound_nodes: List of nodes. + input, output: Input/output tensor(s). Note that if the layer is used + more than once (shared layer), this is ill-defined + and will raise an exception. In such cases, use + `layer.get_input_at(node_index)`. + input_mask, output_mask: Same as above, for masks. + trainable_weights: List of variables. + non_trainable_weights: List of variables. + weights: The concatenation of the lists trainable_weights and + non_trainable_weights (in this order). + + # Methods + call(x, mask=None): Where the layer's logic lives. + __call__(x, mask=None): Wrapper around the layer logic (`call`). + If x is a Keras tensor: + - Connect current layer with last layer from tensor: + `self._add_inbound_node(last_layer)` + - Add layer to tensor history + If layer is not built: + - Build from inputs shape + get_weights() + set_weights(weights) + get_config() + count_params() + compute_output_shape(input_shape) + compute_mask(x, mask) + get_input_at(node_index) + get_output_at(node_index) + get_input_shape_at(node_index) + get_output_shape_at(node_index) + get_input_mask_at(node_index) + get_output_mask_at(node_index) + + # Class Methods + from_config(config) + + # Internal methods: + build(input_shape) + _add_inbound_node(layer, index=0) + """ + + def __init__(self, **kwargs): + # These properties should be set by the user via keyword arguments. + # note that 'dtype', 'input_shape' and 'batch_input_shape' + # are only applicable to input layers: do not pass these keywords + # to non-input layers. + allowed_kwargs = { + 'activity_regularizer', + 'input_shape', + 'batch_input_shape', + 'batch_size', + 'dtype', + 'name', + 'trainable', + 'weights', + } + # Validate optional keyword arguments. + for kwarg in kwargs: + if kwarg not in allowed_kwargs: + raise TypeError('Keyword argument not understood:', kwarg) + + # Get layer name. + name = kwargs.get('name') + + # Get `trainable` status. + trainable = kwargs.get('trainable', True) + + # Get `dtype`. + dtype = kwargs.get('dtype') + if dtype is None: + dtype = K.floatx() + + # Call super, which will set all properties common to Keras layers + # and core TF layers. + super(Layer, self).__init__( + name=name, dtype=dtype, trainable=trainable, + activity_regularizer=kwargs.get('activity_regularizer')) + + # Add properties that are Keras-only for now. + self.supports_masking = False + + # Manage input shape information if passed. + if 'input_shape' in kwargs or 'batch_input_shape' in kwargs: + # In this case we will later create an input layer + # to insert before the current layer + if 'batch_input_shape' in kwargs: + batch_input_shape = tuple(kwargs['batch_input_shape']) + elif 'input_shape' in kwargs: + if 'batch_size' in kwargs: + batch_size = kwargs['batch_size'] + else: + batch_size = None + batch_input_shape = (batch_size,) + tuple(kwargs['input_shape']) + self._batch_input_shape = batch_input_shape + + # Manage initial weight values if passed. + if 'weights' in kwargs: + self._initial_weights = kwargs['weights'] + else: + self._initial_weights = None + + def add_weight(self, + name, + shape, + dtype=None, + initializer=None, + regularizer=None, + trainable=True, + constraint=None): + """Adds a weight variable to the layer. + + Arguments: + name: String, the name for the weight variable. + shape: The shape tuple of the weight. + dtype: The dtype of the weight. + initializer: An Initializer instance (callable). + regularizer: An optional Regularizer instance. + trainable: A boolean, whether the weight should + be trained via backprop or not (assuming + that the layer itself is also trainable). + constraint: An optional Constraint instance. + + Returns: + The created weight variable. + """ + if dtype is None: + dtype = K.floatx() + weight = self.add_variable(name, shape, + dtype=dtype, + initializer=initializers.get(initializer), + regularizer=regularizers.get(regularizer), + constraint=constraints.get(constraint), + trainable=trainable) + return weight + + def call(self, inputs, **kwargs): # pylint: disable=unused-argument + """This is where the layer's logic lives. + + Arguments: + inputs: Input tensor, or list/tuple of input tensors. + **kwargs: Additional keyword arguments. + + Returns: + A tensor or list/tuple of tensors. + """ + return inputs + + def __call__(self, inputs, **kwargs): + """Wrapper around self.call(), for handling internal references. + + If a Keras tensor is passed: + - We call self._add_inbound_node(). + - If necessary, we `build` the layer to match + the shape of the input(s). + - We update the _keras_history of the output tensor(s) + with the current layer. + This is done as part of _add_inbound_node(). + + Arguments: + inputs: Can be a tensor or list/tuple of tensors. + **kwargs: Additional keyword arguments to be passed to `call()`. + + Returns: + Output of the layer's `call` method. + + Raises: + ValueError: in case the layer is missing shape information + for its `build` call. + """ + # Actually call the layer (optionally building it). + output = super(Layer, self).__call__(inputs, **kwargs) + if context.executing_eagerly(): + return output + + if hasattr(self, '_symbolic_set_inputs') and not self.inputs: + # Subclassed network: explicitly set metadata normally set by a call to + # self._set_inputs(). + self._symbolic_set_inputs(inputs, output) + + # Update learning phase info. + output_tensors = generic_utils.to_list(output) + uses_lp = any( + [getattr(x, '_uses_learning_phase', False) + for x in generic_utils.to_list(inputs)]) + uses_lp = getattr(self, 'uses_learning_phase', False) or uses_lp + for i in range(len(output_tensors)): + output_tensors[i]._uses_learning_phase = getattr( + output_tensors[i], '_uses_learning_phase', False) or uses_lp + + # Optionally load weight values that were specified at layer instantiation. + if hasattr(self, '_initial_weights') and self._initial_weights is not None: + self.set_weights(self._initial_weights) + del self._initial_weights + return output + + def compute_output_shape(self, input_shape): + """Computes the output shape of the layer. + + Assumes that the layer will be built + to match that input shape provided. + + Arguments: + input_shape: Shape tuple (tuple of integers) + or list of shape tuples (one per output tensor of the layer). + Shape tuples can include None for free dimensions, + instead of an integer. + + Returns: + An input shape tuple. + """ + logging.warning( + 'All custom layers should implement the ' + '`compute_output_shape` method. This layer (' + self.name + ') ' + 'is relying on the base `Layer.compute_output_shape` implementation, ' + 'which will start raising a `NotImplementedError` ' + 'as of July 1st, 2018.') + return input_shape + + def compute_mask(self, inputs, mask=None): # pylint: disable=unused-argument + """Computes an output mask tensor. + + Arguments: + inputs: Tensor or list of tensors. + mask: Tensor or list of tensors. + + Returns: + None or a tensor (or list of tensors, + one per output tensor of the layer). + """ + if not self.supports_masking: + if mask is not None: + if isinstance(mask, list): + if any(m is not None for m in mask): + raise TypeError('Layer ' + self.name + ' does not support masking, ' + 'but was passed an input_mask: ' + str(mask)) + else: + raise TypeError('Layer ' + self.name + ' does not support masking, ' + 'but was passed an input_mask: ' + str(mask)) + # masking not explicitly supported: return None as mask + return None + # if masking is explicitly supported, by default + # carry over the input mask + return mask + + def get_input_mask_at(self, node_index): + """Retrieves the input mask tensor(s) of a layer at a given node. + + Arguments: + node_index: Integer, index of the node + from which to retrieve the attribute. + E.g. `node_index=0` will correspond to the + first time the layer was called. + + Returns: + A mask tensor + (or list of tensors if the layer has multiple inputs). + """ + inputs = self.get_input_at(node_index) + if isinstance(inputs, list): + return [getattr(x, '_keras_mask', None) for x in inputs] + else: + return getattr(inputs, '_keras_mask', None) + + def get_output_mask_at(self, node_index): + """Retrieves the output mask tensor(s) of a layer at a given node. + + Arguments: + node_index: Integer, index of the node + from which to retrieve the attribute. + E.g. `node_index=0` will correspond to the + first time the layer was called. + + Returns: + A mask tensor + (or list of tensors if the layer has multiple outputs). + """ + output = self.get_output_at(node_index) + if isinstance(output, list): + return [getattr(x, '_keras_mask', None) for x in output] + else: + return getattr(output, '_keras_mask', None) + + @property + def input_mask(self): + """Retrieves the input mask tensor(s) of a layer. + + Only applicable if the layer has exactly one inbound node, + i.e. if it is connected to one incoming layer. + + Returns: + Input mask tensor (potentially None) or list of input + mask tensors. + + Raises: + AttributeError: if the layer is connected to + more than one incoming layers. + """ + inputs = self.input + if isinstance(inputs, list): + return [getattr(x, '_keras_mask', None) for x in inputs] + else: + return getattr(inputs, '_keras_mask', None) + + @property + def output_mask(self): + """Retrieves the output mask tensor(s) of a layer. + + Only applicable if the layer has exactly one inbound node, + i.e. if it is connected to one incoming layer. + + Returns: + Output mask tensor (potentially None) or list of output + mask tensors. + + Raises: + AttributeError: if the layer is connected to + more than one incoming layers. + """ + output = self.output + if isinstance(output, list): + return [getattr(x, '_keras_mask', None) for x in output] + else: + return getattr(output, '_keras_mask', None) + + def set_weights(self, weights): + """Sets the weights of the layer, from Numpy arrays. + + Arguments: + weights: a list of Numpy arrays. The number + of arrays and their shape must match + number of the dimensions of the weights + of the layer (i.e. it should match the + output of `get_weights`). + + Raises: + ValueError: If the provided weights list does not match the + layer's specifications. + """ + params = self.weights + if len(params) != len(weights): + raise ValueError('You called `set_weights(weights)` on layer "' + + self.name + '" with a weight list of length ' + + str(len(weights)) + ', but the layer was expecting ' + + str(len(params)) + ' weights. Provided weights: ' + + str(weights)[:50] + '...') + if not params: + return + weight_value_tuples = [] + param_values = K.batch_get_value(params) + for pv, p, w in zip(param_values, params, weights): + if pv.shape != w.shape: + raise ValueError('Layer weight shape ' + str(pv.shape) + + ' not compatible with ' + 'provided weight shape ' + str(w.shape)) + weight_value_tuples.append((p, w)) + K.batch_set_value(weight_value_tuples) + + def get_weights(self): + """Returns the current weights of the layer. + + Returns: + Weights values as a list of numpy arrays. + """ + params = self.weights + return K.batch_get_value(params) + + def get_config(self): + """Returns the config of the layer. + + A layer config is a Python dictionary (serializable) + containing the configuration of a layer. + The same layer can be reinstantiated later + (without its trained weights) from this configuration. + + The config of a layer does not include connectivity + information, nor the layer class name. These are handled + by `Network` (one layer of abstraction above). + + Returns: + Python dictionary. + """ + config = {'name': self.name, 'trainable': self.trainable} + if hasattr(self, '_batch_input_shape'): + config['batch_input_shape'] = self._batch_input_shape + if hasattr(self, 'dtype'): + config['dtype'] = self.dtype + return config + + @classmethod + def from_config(cls, config): + """Creates a layer from its config. + + This method is the reverse of `get_config`, + capable of instantiating the same layer from the config + dictionary. It does not handle layer connectivity + (handled by Network), nor weights (handled by `set_weights`). + + Arguments: + config: A Python dictionary, typically the + output of get_config. + + Returns: + A layer instance. + """ + return cls(**config) + + @tf_base_layers.Layer.activity_regularizer.setter + def activity_regularizer(self, activity_regularizer): + self._activity_regularizer = activity_regularizer + + +def shape_type_conversion(fn): + """Decorator that handles tuple/TensorShape conversion. + + Used in `compute_output_shape` and `build`. + + Arguments: + fn: function to wrap. + + Returns: + Wrapped function. + """ + + def wrapper(instance, input_shape): + if input_shape is not None: + if isinstance(input_shape, list): + input_shape = [ + tuple(tensor_shape.TensorShape(x).as_list()) for x in input_shape] + else: + input_shape = tuple(tensor_shape.TensorShape(input_shape).as_list()) + output_shape = fn(instance, input_shape) + if output_shape is not None: + if isinstance(output_shape, list): + return [tensor_shape.TensorShape(x) for x in output_shape] + return tensor_shape.TensorShape(output_shape) + + return wrapper diff --git a/tensorflow/python/keras/_impl/keras/engine/input_layer.py b/tensorflow/python/keras/_impl/keras/engine/input_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..b51dd8a2189d0c8542c84dfeac9be0d72b96ff1b --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/input_layer.py @@ -0,0 +1,231 @@ +# 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. +# ============================================================================== +# pylint: disable=protected-access +"""Input layer code (`Input` and `InputLayer`). +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.eager import context +from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras.engine import base_layer +from tensorflow.python.layers import base as tf_base_layers +from tensorflow.python.ops import array_ops +from tensorflow.python.util.tf_export import tf_export + + +@tf_export('keras.layers.InputLayer') +class InputLayer(base_layer.Layer): + """Layer to be used as an entry point into a Network (a graph of layers). + + It can either wrap an existing tensor (pass an `input_tensor` argument) + or create its a placeholder tensor (pass arguments `input_shape`, and + optionally, `dtype`). + + It is generally recommend to use the functional layer API via `Input`, + (which creates an `InputLayer`) without directly using `InputLayer`. + + Arguments: + input_shape: Shape tuple (not including the batch axis), or `TensorShape` + instance (not including the batch axis). + batch_size: Optional input batch size (integer or None). + dtype: Datatype of the input. + input_tensor: Optional tensor to use as layer input + instead of creating a placeholder. + sparse: Boolean, whether the placeholder created + is meant to be sparse. + name: Name of the layer (string). + """ + + def __init__(self, + input_shape=None, + batch_size=None, + dtype=None, + input_tensor=None, + sparse=False, + name=None, + **kwargs): + if 'batch_input_shape' in kwargs: + batch_input_shape = kwargs.pop('batch_input_shape') + if input_shape and batch_input_shape: + raise ValueError('Only provide the input_shape OR ' + 'batch_input_shape argument to ' + 'InputLayer, not both at the same time.') + batch_size = batch_input_shape[0] + input_shape = batch_input_shape[1:] + if kwargs: + raise ValueError('Unrecognized keyword arguments:', kwargs.keys()) + + if not name: + prefix = 'input' + name = prefix + '_' + str(K.get_uid(prefix)) + + if not dtype: + if input_tensor is None: + dtype = K.floatx() + else: + dtype = K.dtype(input_tensor) + super(InputLayer, self).__init__(dtype=dtype, name=name) + self.built = True + self.sparse = sparse + self.batch_size = batch_size + + if isinstance(input_shape, tensor_shape.TensorShape): + input_shape = tuple(input_shape.as_list()) + + if input_tensor is None: + if input_shape is not None: + batch_input_shape = (batch_size,) + tuple(input_shape) + else: + batch_input_shape = None + + if context.executing_eagerly(): + # In eager mode, create a temporary placeholder to call the layer on. + input_tensor = tf_base_layers._DeferredTensor( # pylint: disable=protected-access + shape=batch_input_shape, + dtype=dtype, + name=self.name) + else: + # In graph mode, create a graph placeholder to call the layer on. + if sparse: + input_tensor = array_ops.sparse_placeholder( + shape=batch_input_shape, + dtype=dtype, + name=self.name) + else: + input_tensor = array_ops.placeholder( + shape=batch_input_shape, + dtype=dtype, + name=self.name) + + # For compatibility with Keras API. + self.is_placeholder = True + self._batch_input_shape = batch_input_shape + else: + # For compatibility with Keras API. + self.is_placeholder = False + self._batch_input_shape = tuple(input_tensor.get_shape().as_list()) + + # Create an input node to add to self.outbound_node + # and set output_tensors' _keras_history. + input_tensor._keras_history = (self, 0, 0) # pylint: disable=protected-access + tf_base_layers.Node( + self, + inbound_layers=[], + node_indices=[], + tensor_indices=[], + input_tensors=[input_tensor], + output_tensors=[input_tensor]) + + def get_config(self): + config = { + 'batch_input_shape': self._batch_input_shape, + 'dtype': self.dtype, + 'sparse': self.sparse, + 'name': self.name + } + return config + + +@tf_export('keras.layers.Input', 'keras.Input') +def Input( # pylint: disable=invalid-name + shape=None, + batch_size=None, + name=None, + dtype=None, + sparse=False, + tensor=None, + **kwargs): + """`Input()` is used to instantiate a Keras tensor. + + A Keras tensor is a tensor object from the underlying backend + (Theano or TensorFlow), which we augment with certain + attributes that allow us to build a Keras model + just by knowing the inputs and outputs of the model. + + For instance, if a, b and c are Keras tensors, + it becomes possible to do: + `model = Model(input=[a, b], output=c)` + + The added Keras attribute is: + `_keras_history`: Last layer applied to the tensor. + the entire layer graph is retrievable from that layer, + recursively. + + Arguments: + shape: A shape tuple (integers), not including the batch size. + For instance, `shape=(32,)` indicates that the expected input + will be batches of 32-dimensional vectors. + batch_size: optional static batch size (integer). + name: An optional name string for the layer. + Should be unique in a model (do not reuse the same name twice). + It will be autogenerated if it isn't provided. + dtype: The data type expected by the input, as a string + (`float32`, `float64`, `int32`...) + sparse: A boolean specifying whether the placeholder + to be created is sparse. + tensor: Optional existing tensor to wrap into the `Input` layer. + If set, the layer will not create a placeholder tensor. + **kwargs: deprecated arguments support. + + Returns: + A tensor. + + Example: + + ```python + # this is a logistic regression in Keras + x = Input(shape=(32,)) + y = Dense(16, activation='softmax')(x) + model = Model(x, y) + ``` + + Raises: + ValueError: in case of invalid arguments. + """ + if 'batch_shape' in kwargs: + batch_shape = kwargs.pop('batch_shape') + if shape and batch_shape: + raise ValueError('Only provide the shape OR ' + 'batch_shape argument to ' + 'Input, not both at the same time.') + batch_size = batch_shape[0] + shape = batch_shape[1:] + if kwargs: + raise ValueError('Unrecognized keyword arguments:', kwargs.keys()) + + if dtype is None: + dtype = K.floatx() + if not shape and tensor is None: + raise ValueError('Please provide to Input either a `shape`' + ' or a `tensor` argument. Note that ' + '`shape` does not include the batch ' + 'dimension.') + input_layer = InputLayer( + input_shape=shape, + batch_size=batch_size, + name=name, + dtype=dtype, + sparse=sparse, + input_tensor=tensor) + # Return tensor including `_keras_history`. + # Note that in this case train_output and test_output are the same pointer. + outputs = input_layer._inbound_nodes[0].output_tensors + if len(outputs) == 1: + return outputs[0] + else: + return outputs diff --git a/tensorflow/python/keras/_impl/keras/engine/network.py b/tensorflow/python/keras/_impl/keras/engine/network.py new file mode 100644 index 0000000000000000000000000000000000000000..bf8239043874794c6617937cfc9c619d743502a9 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/network.py @@ -0,0 +1,1516 @@ +# 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. +# ============================================================================== +# pylint: disable=protected-access +"""A `Network` is way to compose layers: the topological form of a `Model`. +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import copy +import json +import os + +import numpy as np +from six.moves import zip # pylint: disable=redefined-builtin + +from tensorflow.python.eager import context +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras.engine import base_layer +from tensorflow.python.keras._impl.keras.engine import saving +from tensorflow.python.keras._impl.keras.utils import generic_utils +from tensorflow.python.keras._impl.keras.utils.io_utils import ask_to_proceed_with_overwrite +from tensorflow.python.keras._impl.keras.utils.layer_utils import print_summary as print_layer_summary +from tensorflow.python.layers import base as tf_base_layers +from tensorflow.python.layers import utils as tf_layers_util +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpointable +from tensorflow.python.util import nest +from tensorflow.python.util import tf_inspect + + +# pylint: disable=g-import-not-at-top +try: + import h5py +except ImportError: + h5py = None + +try: + import yaml +except ImportError: + yaml = None +# pylint: enable=g-import-not-at-top + + +class Network(base_layer.Layer): + """A `Network` is a composition of layers. + + It is the topological form of a "model". A `Model` + is simply a `Network` with added training routines. + """ + + def __init__(self, *args, **kwargs): # pylint: disable=super-init-not-called + # Signature detection + if (len(args) == 2 or + len(args) == 1 and 'outputs' in kwargs or + 'inputs' in kwargs and 'outputs' in kwargs): + # Graph network + self._init_graph_network(*args, **kwargs) + else: + # Subclassed network + self._init_subclassed_network(**kwargs) + + def _base_init(self, name=None): + # The following are implemented as property functions: + # self.trainable_weights + # self.non_trainable_weights + # self.input_spec + # self.losses + # self.updates + + self._init_set_name(name) + self._activity_regularizer = None + # This acts just like the `trainable` attribute of any layer instance. + # It does not affect users of the underlying layers, only users of the + # Network instance. + self.trainable = True + self._is_compiled = False + self._expects_training_arg = False + + self.supports_masking = False + self.optimizer = None + + # Private attributes to implement compatibility with Layer. + self._updates = [] # Used in symbolic mode only. + self._losses = [] # Used in symbolic mode only. + self._scope = None # Never used. + self._reuse = None # Never used. + if context.executing_eagerly(): + self._graph = None + else: + self._graph = ops.get_default_graph() # Used in symbolic mode only. + # A Network does not create weights of its own, thus has no dtype. + self._dtype = None + + # All layers in order of horizontal graph traversal. + # Entries are unique. Includes input and output layers. + self._layers = [] + + # Used in symbolic mode only, only in conjonction with graph-networks + self._outbound_nodes = [] + self._inbound_nodes = [] + + def _init_graph_network(self, inputs, outputs, name=None): + # Normalize and set self.inputs, self.outputs. + if isinstance(inputs, (list, tuple)): + self.inputs = list(inputs) # Tensor or list of tensors. + else: + self.inputs = [inputs] + if isinstance(outputs, (list, tuple)): + self.outputs = list(outputs) + else: + self.outputs = [outputs] + + # User-prodived argument validation. + if context.executing_eagerly(): + # Check that all inputs/outputs are DeferredTensors. + for tensor in self.inputs: + if not isinstance(tensor, tf_base_layers._DeferredTensor): # pylint: disable=protected-access + raise TypeError('When eager execution is enabled, ' + 'inputs must come from a call to ' + '`tf.keras.Input` (called after ' + 'tfe.enable_eager_execution()). ' + 'Received invalid input: ' + str(tensor)) + for tensor in self.outputs: + if not isinstance(tensor, tf_base_layers._DeferredTensor): # pylint: disable=protected-access + raise TypeError('When eager execution is enabled, ' + 'outputs must come from a call to ' + 'a layer (called after ' + 'tfe.enable_eager_execution()). ' + 'Received invalid output: ' + str(tensor)) + # Check for redundancy in inputs. + if len(set(self.inputs)) != len(self.inputs): + raise ValueError('The list of inputs passed to the model ' + 'is redundant. ' + 'All inputs should only appear once.' + ' Found: ' + str(self.inputs)) + for x in self.inputs: + # Check that x has appropriate `_keras_history` metadata. + if not hasattr(x, '_keras_history'): + cls_name = self.__class__.__name__ + raise ValueError('Input tensors to a ' + cls_name + ' ' + + 'must come from `tf.layers.Input`. ' + 'Received: ' + str(x) + + ' (missing previous layer metadata).') + # Check that x is an input tensor. + # pylint: disable=protected-access + layer, node_index, tensor_index = x._keras_history + if len(layer._inbound_nodes) > 1 or ( + layer._inbound_nodes and layer._inbound_nodes[0].inbound_layers): + cls_name = self.__class__.__name__ + logging.warning(cls_name + ' inputs must come from ' + '`tf.layers.Input` (thus holding past layer metadata), ' + 'they cannot be the output of ' + 'a previous non-Input layer. ' + 'Here, a tensor specified as ' + 'input to "' + self.name + '" was not an Input tensor, ' + 'it was generated by layer ' + layer.name + '.\n' + 'Note that input tensors are ' + 'instantiated via `tensor = tf.layers.Input(shape)`.\n' + 'The tensor that caused the issue was: ' + str(x.name)) + for x in self.outputs: + if not hasattr(x, '_keras_history'): + cls_name = self.__class__.__name__ + raise ValueError('Output tensors to a ' + cls_name + ' must be ' + 'the output of a TensorFlow `Layer` ' + '(thus holding past layer metadata). Found: ' + str(x)) + + self._base_init(name=name) + self._compute_previous_mask = ( + 'mask' in tf_inspect.getargspec(self.call).args or + hasattr(self, 'compute_mask')) + # A Network does not create weights of its own, thus it is already + # built. + self.built = True + self._is_graph_network = True + + # # List of initial layers (1 to 1 mapping with self.inputs, + # # hence the same layer might appear twice) + # self._input_layers = [] + # self._input_layers_node_indices = [] + # self._input_layers_tensor_indices = [] + # # list of layers (1 to 1 mapping with self.inputs, + # # hence the same layer might appear twice) + # self._output_layers = [] + # self._output_layers_node_indices = [] + # self._output_layers_tensor_indices = [] + + self._input_layers = [] + self._output_layers = [] + self._input_coordinates = [] + self._output_coordinates = [] + + # This is for performance optimization when calling the Network on new + # inputs. Every time the Network is called on a set on input tensors, + # we compute the output tensors, output masks and output shapes in one pass, + # then cache them here. When any of these outputs is queried later, we + # retrieve it from there instead of recomputing it. + self._output_mask_cache = {} + self._output_tensor_cache = {} + self._output_shape_cache = {} + + # Build self._output_layers: + for x in self.outputs: + layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access + self._output_layers.append(layer) + self._output_coordinates.append((layer, node_index, tensor_index)) + + # Build self._input_layers: + for x in self.inputs: + layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access + # It's supposed to be an input layer, so only one node + # and one tensor output. + assert node_index == 0 + assert tensor_index == 0 + self._input_layers.append(layer) + self._input_coordinates.append((layer, node_index, tensor_index)) + + # Keep track of the network's nodes and layers. + nodes, nodes_by_depth, layers, layers_by_depth = _map_graph_network( + self.inputs, self.outputs) + self._network_nodes = nodes + self._nodes_by_depth = nodes_by_depth + self._layers = layers + self._layers_by_depth = layers_by_depth + + # Create the node linking internal inputs to internal outputs. + tf_base_layers.Node( + outbound_layer=self, + inbound_layers=[], + node_indices=[], + tensor_indices=[], + input_tensors=self.inputs, + output_tensors=self.outputs) + + # Fill in the output mask cache. + masks = [] + for x in self.inputs: + mask = x._keras_mask if hasattr(x, '_keras_mask') else None # pylint: disable=protected-access + masks.append(mask) + mask_cache_key = (tf_layers_util.object_list_uid(self.inputs) + '_' + + tf_layers_util.object_list_uid(masks)) + masks = [] + for x in self.outputs: + mask = x._keras_mask if hasattr(x, '_keras_mask') else None # pylint: disable=protected-access + masks.append(mask) + if len(masks) == 1: + mask = masks[0] + else: + mask = masks + self._output_mask_cache[mask_cache_key] = mask + + # Build self.input_names and self.output_names. + self.input_names = [] + self.output_names = [] + self._feed_input_names = [] + self._feed_inputs = [] + self._feed_input_shapes = [] + for i, layer in enumerate(self._input_layers): + self.input_names.append(layer.name) + if layer.is_placeholder: + self._feed_input_names.append(layer.name) + self._feed_input_shapes.append(K.int_shape(self.inputs[i])) + # layer.input gives an error in eager mode + if not context.executing_eagerly(): + self._feed_inputs.append(layer.input) + for layer in self._output_layers: + self.output_names.append(layer.name) + + def _init_subclassed_network(self, name=None): + self._base_init(name=name) + self._is_graph_network = False + if 'training' in tf_inspect.getargspec(self.call).args: + self._expects_training_arg = True + else: + self._expects_training_arg = False + + self.outputs = None + self.inputs = None + self.built = False + + def __setattr__(self, name, value): + if isinstance(value, (tf_base_layers.Layer, Network)): + try: + is_graph_network = self._is_graph_network + except AttributeError: + raise RuntimeError('It looks like you are subclassing `Model` and you ' + 'forgot to call `super(YourClass, self).__init__()`.' + ' Always start with this line.') + if not is_graph_network: + if value not in self._layers: + self._layers.append(value) + if isinstance(value, checkpointable.CheckpointableBase): + # Layer (and therefore Network/Model) inherit from CheckpointableBase + # rather than Checkpointable, which means there is no Checkpointable + # __setattr__ override (it would be a performance issue for functional + # layers). Therefore Model tracks Checkpointable objects itself. + self._track_checkpointable( + checkpointable=value, name=name, overwrite=True) + super(Network, self).__setattr__(name, value) + + def add_variable(self, name, shape, dtype=None, initializer=None, + regularizer=None, trainable=True, constraint=None): + raise NotImplementedError('`add_variable` is not supported on Networks.') + + def add_loss(self, *args, **kwargs): + if context.executing_eagerly(): + raise NotImplementedError('`add_loss` is not supported on Networks ' + 'when eager execution is enabled.') + super(Network, self).add_loss(*args, **kwargs) + + @property + def uses_learning_phase(self): + return any( + [getattr(x, '_uses_learning_phase', False) for x in self.outputs]) + + @property + def stateful(self): + return any([(hasattr(layer, 'stateful') and layer.stateful) + for layer in self.layers]) + + def reset_states(self): + for layer in self.layers: + if hasattr(layer, 'reset_states') and getattr(layer, 'stateful', False): + layer.reset_states() + + @property + def state_updates(self): + """Returns the `updates` from all layers that are stateful. + + This is useful for separating training updates and + state updates, e.g. when we need to update a layer's internal state + during prediction. + + Returns: + A list of update ops. + """ + state_updates = [] + for layer in self.layers: + if getattr(layer, 'stateful', False): + if hasattr(layer, 'updates'): + state_updates += layer.updates + return state_updates + + def get_weights(self): + """Retrieves the weights of the model. + + Returns: + A flat list of Numpy arrays. + """ + weights = [] + for layer in self.layers: + weights += layer.weights + return K.batch_get_value(weights) + + def set_weights(self, weights): + """Sets the weights of the model. + + Arguments: + weights: A list of Numpy arrays with shapes and types matching + the output of `model.get_weights()`. + """ + tuples = [] + for layer in self.layers: + num_param = len(layer.weights) + layer_weights = weights[:num_param] + for sw, w in zip(layer.weights, layer_weights): + tuples.append((sw, w)) + weights = weights[num_param:] + K.batch_set_value(tuples) + + def compute_mask(self, inputs, mask): + if not self._is_graph_network: + return None + + inputs = generic_utils.to_list(inputs) + if mask is None: + masks = [None for _ in range(len(inputs))] + else: + masks = generic_utils.to_list(mask) + cache_key = (tf_layers_util.object_list_uid(inputs) + + '_' + tf_layers_util.object_list_uid(masks)) + if cache_key in self._output_mask_cache: + return self._output_mask_cache[cache_key] + else: + _, output_masks = self._run_internal_graph(inputs, mask=masks) + return output_masks + + @property + def layers(self): + return self._layers + + def get_layer(self, name=None, index=None): + """Retrieves a layer based on either its name (unique) or index. + + If `name` and `index` are both provided, `index` will take precedence. + Indices are based on order of horizontal graph traversal (bottom-up). + + Arguments: + name: String, name of layer. + index: Integer, index of layer. + + Returns: + A layer instance. + + Raises: + ValueError: In case of invalid layer name or index. + """ + # TODO(fchollet): We could build a dictionary based on layer names + # since they are constant, but we have not done that yet. + if index is not None: + if len(self.layers) <= index: + raise ValueError('Was asked to retrieve layer at index ' + str(index) + + ' but model only has ' + str(len(self.layers)) + + ' layers.') + else: + return self.layers[index] + else: + if not name: + raise ValueError('Provide either a layer name or layer index.') + for layer in self.layers: + if layer.name == name: + return layer + raise ValueError('No such layer: ' + name) + + @property + def updates(self): + """Retrieves the network's updates. + + Will only include updates that are either + unconditional, or conditional on inputs to this model + (e.g. will not include updates that were created by layers of this model + outside of the model). + + Effectively, `network.updates` behaves like `layer.updates`. + + Concrete example: + + ```python + bn = keras.layers.BatchNormalization() + x1 = keras.layers.Input(shape=(10,)) + _ = bn(x1) # This creates 2 updates. + + x2 = keras.layers.Input(shape=(10,)) + y2 = bn(x2) # This creates 2 more updates. + + # The BN layer has now 4 updates. + self.assertEqual(len(bn.updates), 4) + + # Let's create a model from x2 to y2. + model = keras.models.Model(x2, y2) + + # The model does not list all updates from its underlying layers, + # but only the updates that are relevant to it. Updates created by layers + # outside of the model are discarded. + self.assertEqual(len(model.updates), 2) + + # If you keep calling the model, you append to its updates, just like + # what happens for a layer. + x3 = keras.layers.Input(shape=(10,)) + y3 = model(x3) + self.assertEqual(len(model.updates), 4) + + # But if you call the inner BN layer independently, you don't affect + # the model's updates. + x4 = keras.layers.Input(shape=(10,)) + _ = bn(x4) + self.assertEqual(len(model.updates), 4) + ``` + + Returns: + A list of update ops. + """ + if context.executing_eagerly(): + return [] + + if not self.trainable and not self.stateful: + return [] + + updates = [] + for layer in self.layers: + updates += layer.updates + + # `updates` might contain irrelevant updates, so it needs to be filtered + # with respect to inputs the model has been called on. + if self.inputs: + relevant_inputs = self.inputs[:] + else: + relevant_inputs = [] + for i in range(1, len(self._inbound_nodes)): + inputs = self.get_input_at(i) + if isinstance(inputs, list): + relevant_inputs += inputs + else: + relevant_inputs.append(inputs) + reachable = tf_layers_util.get_reachable_from_inputs(relevant_inputs, + updates) + relevant_conditional_updates = [x for x in updates if x in reachable] + unconditional_updates = [ + x for x in updates if x._unconditional_update] # pylint: disable=protected-access + # A layer could be used multiple times in a nested structure, + # so the updates list must be de-duped. + return list(set( + relevant_conditional_updates + unconditional_updates + self._updates)) + + @property + def losses(self): + """Retrieves the network's losses. + + Will only include losses that are either + unconditional, or conditional on inputs to this model + (e.g. will not include losses that depend on tensors + that aren't inputs to this model). + + Returns: + A list of loss tensors. + """ + losses = [] + for layer in self.layers: + losses += layer.losses + if context.executing_eagerly(): + return losses + + if self.inputs: + relevant_inputs = self.inputs[:] + else: + relevant_inputs = [] + for i in range(1, len(self._inbound_nodes)): + inputs = self.get_input_at(i) + if isinstance(inputs, list): + relevant_inputs += inputs + else: + relevant_inputs.append(inputs) + reachable = tf_layers_util.get_reachable_from_inputs(relevant_inputs, + losses) + relevant_conditional_losses = [x for x in losses if x in reachable] + unconditional_losses = [ + x for x in losses if x._unconditional_loss] # pylint: disable=protected-access + return list(set( + relevant_conditional_losses + unconditional_losses + self._losses)) + + @property + def trainable_weights(self): + if not self.trainable: + return [] + weights = [] + for layer in self.layers: + weights += layer.trainable_weights + return weights + + @property + def non_trainable_weights(self): + weights = [] + for layer in self.layers: + weights += layer.non_trainable_weights + if not self.trainable: + trainable_weights = [] + for layer in self.layers: + trainable_weights += layer.trainable_weights + return trainable_weights + weights + return weights + + @property + def input_spec(self): + """Gets the network's input specs. + + Returns: + A list of `InputSpec` instances (one per input to the model) + or a single instance if the model has only one input. + """ + # If not a graph network, can't assume anything. + if not self._is_graph_network: + return None + + specs = [] + for layer in self._input_layers: + if layer.input_spec is None: + specs.append(None) + else: + if not isinstance(layer.input_spec, list): + raise TypeError('Layer ' + layer.name + + ' has an input_spec attribute that ' + 'is not a list. We expect a list. ' + 'Found input_spec = ' + str(layer.input_spec)) + specs += layer.input_spec + if len(specs) == 1: + return specs[0] + return specs + + def call(self, inputs, training=None, mask=None): + """Calls the model on new inputs. + + In this case `call` just reapplies + all ops in the graph to the new inputs + (e.g. build a new computational graph from the provided inputs). + + Arguments: + inputs: A tensor or list of tensors. + training: Boolean or boolean scalar tensor, indicating whether to run + the `Network` in training mode or inference mode. + mask: A mask or list of masks. A mask can be + either a tensor or None (no mask). + + Returns: + A tensor if there is a single output, or + a list of tensors if there are more than one outputs. + """ + inputs = nest.flatten(inputs) + if mask is None: + masks = [None for _ in range(len(inputs))] + else: + masks = nest.flatten(mask) + + if not context.executing_eagerly(): + # Try to retrieve cached outputs if the layer has already been called + # on these exact inputs. + cache_key = (tf_layers_util.object_list_uid(inputs) + + '_' + tf_layers_util.object_list_uid(masks)) + if cache_key in self._output_tensor_cache: + # Cache hit. + return self._output_tensor_cache[cache_key] + # Actually apply the network graph to the new inputs. + outputs, _ = self._run_internal_graph(inputs, + training=training, + mask=masks) + return outputs + + def compute_output_shape(self, input_shape): + if not self._is_graph_network: + raise NotImplementedError + + if isinstance(input_shape, list): + input_shapes = [] + for shape in input_shape: + if shape is not None: + input_shapes.append(tuple(tensor_shape.TensorShape(shape).as_list())) + else: + input_shapes.append(None) + else: + if input_shape is not None: + input_shapes = [tuple(tensor_shape.TensorShape(input_shape).as_list())] + else: + input_shapes = [None] + + if len(input_shapes) != len(self._input_layers): + raise ValueError('Invalid input_shape argument ' + str(input_shape) + + ': model has ' + str(len(self._input_layers)) + + ' tensor inputs.') + + cache_key = tf_layers_util.object_list_uid(input_shapes) + if cache_key not in self._output_shape_cache: + # Cache miss. We have to run the network graph manually (recursive calls + # to `compute_output_shape`). + layers_to_output_shapes = {} + for i in range(len(input_shapes)): + layer = self._input_layers[i] + input_shape = input_shapes[i] + # It's an input layer: then `compute_output_shape` is identity, + # and there is only one node and one tensor output. + shape_key = layer.name + '_0_0' + layers_to_output_shapes[shape_key] = input_shape + + depth_keys = list(self._nodes_by_depth.keys()) + depth_keys.sort(reverse=True) + # Iterate over nodes, by depth level. + if len(depth_keys) > 1: + for depth in depth_keys: + nodes = self._nodes_by_depth[depth] + for node in nodes: + # This is always a single layer, never a list. + layer = node.outbound_layer + if layer in self._input_layers: + # We've already covered the input layers + # a few lines above. + continue + # Potentially redundant list, + # same size as node.input_tensors. + input_shapes = [] + for j in range(len(node.inbound_layers)): + inbound_layer = node.inbound_layers[j] + node_index = node.node_indices[j] + tensor_index = node.tensor_indices[j] + shape_key = inbound_layer.name + '_%s_%s' % (node_index, + tensor_index) + input_shape = layers_to_output_shapes[shape_key] + input_shapes.append(input_shape) + + if len(input_shapes) == 1: + output_shape = layer.compute_output_shape(input_shapes[0]) + else: + output_shape = layer.compute_output_shape(input_shapes) + if isinstance(output_shape, list): + output_shapes = [ + tuple(tensor_shape.TensorShape(shape).as_list()) + for shape in output_shape + ] + else: + output_shapes = [ + tuple(tensor_shape.TensorShape(output_shape).as_list()) + ] + + node_index = layer._inbound_nodes.index(node) # pylint: disable=protected-access + for j in range(len(output_shapes)): + shape_key = layer.name + '_%s_%s' % (node_index, j) + layers_to_output_shapes[shape_key] = output_shapes[j] + + # Read final output shapes from layers_to_output_shapes. + output_shapes = [] + for i in range(len(self._output_layers)): + layer, node_index, tensor_index = self._output_coordinates[i] + shape_key = layer.name + '_%s_%s' % (node_index, tensor_index) + output_shapes.append(layers_to_output_shapes[shape_key]) + # Store in cache. + self._output_shape_cache[cache_key] = output_shapes + else: + # Cache hit. + output_shapes = self._output_shape_cache[cache_key] + + if isinstance(output_shapes, list): + if len(output_shapes) == 1: + return tensor_shape.TensorShape(output_shapes[0]) + else: + return [tensor_shape.TensorShape(shape) for shape in output_shapes] + else: + return tensor_shape.TensorShape(output_shapes) + + def _run_internal_graph(self, inputs, training=None, mask=None): + """Computes output tensors for new inputs. + + # Note: + - Expects `inputs` to be a list (potentially with 1 element). + - Can be run on non-Keras tensors. + + Arguments: + inputs: List of tensors + training: Boolean learning phase. + mask: List of masks (tensors or None). + + Returns: + Three lists: output_tensors, output_masks, output_shapes + """ + # Note: masking support is relevant mainly for Keras. + # It cannot be factored out without having the fully reimplement the network + # calling logic on the Keras side. We choose to incorporate it in + # Network because 1) it may be useful to fully support in tf.layers in + # the future and 2) Keras is a major user of Network. If you don't + # use masking, it does not interfere with regular behavior at all and you + # can ignore it. + if mask is None: + masks = [None for _ in range(len(inputs))] + else: + masks = mask + + # Dictionary mapping reference tensors to tuples + # (computed tensor, compute mask) + # we assume a 1:1 mapping from tensor to mask + # TODO(fchollet): raise exception when a `.compute_mask()` call + # does not return a list the same size as `call` + tensor_map = {} + for x, y, mask in zip(self.inputs, inputs, masks): + tensor_map[str(id(x))] = (y, mask) + + depth_keys = list(self._nodes_by_depth.keys()) + depth_keys.sort(reverse=True) + for depth in depth_keys: + nodes = self._nodes_by_depth[depth] + for node in nodes: + # This is always a single layer, never a list. + layer = node.outbound_layer + reference_input_tensors = node.input_tensors + reference_output_tensors = node.output_tensors + + # If all previous input tensors are available in tensor_map, + # then call node.inbound_layer on them. + computed_data = [] # List of tuples (input, mask). + for x in reference_input_tensors: + if str(id(x)) in tensor_map: + computed_data.append(tensor_map[str(id(x))]) + + if len(computed_data) == len(reference_input_tensors): + # Call layer (reapplying ops to new inputs). + with ops.name_scope(layer.name): + if node.arguments: + kwargs = node.arguments + else: + kwargs = {} + if len(computed_data) == 1: + computed_tensor, computed_mask = computed_data[0] + # Ensure mask propagation if applicable. + if 'mask' in tf_inspect.getargspec(layer.call).args: + kwargs.setdefault('mask', computed_mask) + if 'training' in tf_inspect.getargspec(layer.call).args: + kwargs.setdefault('training', training) + + output_tensors = nest.flatten( + layer.call(computed_tensor, **kwargs)) + if hasattr(layer, 'compute_mask'): + output_masks = nest.flatten( + layer.compute_mask(computed_tensor, computed_mask)) + else: + output_masks = [None for _ in range(len(output_tensors))] + computed_tensors = [computed_tensor] + computed_masks = [computed_mask] + else: + computed_tensors = [x[0] for x in computed_data] + computed_masks = [x[1] for x in computed_data] + if 'mask' in tf_inspect.getargspec(layer.call).args: + kwargs.setdefault('mask', computed_masks) + if 'training' in tf_inspect.getargspec(layer.call).args: + kwargs.setdefault('training', training) + + output_tensors = nest.flatten( + layer.call(computed_tensors, **kwargs)) + if hasattr(layer, 'compute_mask'): + output_masks = nest.flatten( + layer.compute_mask(computed_tensors, computed_masks)) + else: + output_masks = [None for _ in range(len(output_tensors))] + + if not context.executing_eagerly(): + if layer.activity_regularizer is not None: + regularization_losses = [ + layer.activity_regularizer(x) for x in output_tensors + ] + # Apply activity regularizer if any: + layer.add_loss(regularization_losses, computed_tensors) + + # Update tensor_map. + for x, y, mask in zip(reference_output_tensors, output_tensors, + output_masks): + tensor_map[str(id(x))] = (y, mask) + + output_tensors = [] + output_masks = [] + output_shapes = [] + for x in self.outputs: + assert str(id(x)) in tensor_map, 'Could not compute output ' + str(x) + tensor, mask = tensor_map[str(id(x))] + output_shapes.append(tf_layers_util.static_shape(x)) + output_tensors.append(tensor) + output_masks.append(mask) + + if len(output_tensors) == 1: + output_tensors = output_tensors[0] + if output_shapes is not None: + output_shapes = output_shapes[0] + if output_masks is not None: + output_masks = output_masks[0] + + if not context.executing_eagerly(): + # Update cache; + # keys are based on ids on input tensors and inputs masks. + cache_key = (tf_layers_util.object_list_uid(inputs) + + '_' + tf_layers_util.object_list_uid(masks)) + self._output_tensor_cache[cache_key] = output_tensors + self._output_mask_cache[cache_key] = output_masks + + if output_shapes is not None: + input_shapes = [tf_layers_util.static_shape(x) for x in inputs] + cache_key = tf_layers_util.object_list_uid(input_shapes) + self._output_shape_cache[cache_key] = output_shapes + + return output_tensors, output_masks + + def get_config(self): + if not self._is_graph_network: + raise NotImplementedError + + config = { + 'name': self.name, + } + node_conversion_map = {} + for layer in self.layers: + if issubclass(layer.__class__, Network): + # Networks start with a pre-existing node + # linking their input to output. + kept_nodes = 1 + else: + kept_nodes = 0 + for original_node_index, node in enumerate(layer._inbound_nodes): + node_key = _make_node_key(layer.name, original_node_index) + if node_key in self._network_nodes: + node_conversion_map[node_key] = kept_nodes + kept_nodes += 1 + layer_configs = [] + for layer in self.layers: # From the earliest layers on. + layer_class_name = layer.__class__.__name__ + layer_config = layer.get_config() + filtered_inbound_nodes = [] + for original_node_index, node in enumerate(layer._inbound_nodes): + node_key = _make_node_key(layer.name, original_node_index) + if node_key in self._network_nodes: + # The node is relevant to the model: + # add to filtered_inbound_nodes. + if node.arguments: + try: + json.dumps(node.arguments) + kwargs = node.arguments + except TypeError: + logging.warning( + 'Layer ' + layer.name + + ' was passed non-serializable keyword arguments: ' + + str(node.arguments) + '. They will not be included ' + 'in the serialized model (and thus will be missing ' + 'at deserialization time).') + kwargs = {} + else: + kwargs = {} + if node.inbound_layers: + node_data = [] + for i in range(len(node.inbound_layers)): + inbound_layer = node.inbound_layers[i] + node_index = node.node_indices[i] + tensor_index = node.tensor_indices[i] + node_key = _make_node_key(inbound_layer.name, node_index) + new_node_index = node_conversion_map.get(node_key, 0) + node_data.append( + [inbound_layer.name, new_node_index, tensor_index, kwargs]) + filtered_inbound_nodes.append(node_data) + layer_configs.append({ + 'name': layer.name, + 'class_name': layer_class_name, + 'config': layer_config, + 'inbound_nodes': filtered_inbound_nodes, + }) + config['layers'] = layer_configs + + # Gather info about inputs and outputs. + model_inputs = [] + for i in range(len(self._input_layers)): + layer, node_index, tensor_index = self._input_coordinates[i] + node_key = _make_node_key(layer.name, node_index) + if node_key not in self._network_nodes: + continue + new_node_index = node_conversion_map[node_key] + model_inputs.append([layer.name, new_node_index, tensor_index]) + config['input_layers'] = model_inputs + model_outputs = [] + for i in range(len(self._output_layers)): + layer, node_index, tensor_index = self._output_coordinates[i] + node_key = _make_node_key(layer.name, node_index) + if node_key not in self._network_nodes: + continue + new_node_index = node_conversion_map[node_key] + model_outputs.append([layer.name, new_node_index, tensor_index]) + config['output_layers'] = model_outputs + return copy.deepcopy(config) + + @classmethod + def from_config(cls, config, custom_objects=None): + """Instantiates a Model from its config (output of `get_config()`). + + Arguments: + config: Model config dictionary. + custom_objects: Optional dictionary mapping names + (strings) to custom classes or functions to be + considered during deserialization. + + Returns: + A model instance. + + Raises: + ValueError: In case of improperly formatted config dict. + """ + # Layer instances created during + # the graph reconstruction process + created_layers = {} + + # Dictionary mapping layer instances to + # node data that specifies a layer call. + # It acts as a queue that maintains any unprocessed + # layer call until it becomes possible to process it + # (i.e. until the input tensors to the call all exist). + unprocessed_nodes = {} + + def add_unprocessed_node(layer, node_data): + if layer not in unprocessed_nodes: + unprocessed_nodes[layer] = [node_data] + else: + unprocessed_nodes[layer].append(node_data) + + def process_node(layer, node_data): + """Deserialize a node. + + Arguments: + layer: layer instance. + node_data: node config dict. + + Raises: + ValueError: In case of improperly formatted `node_data` dict. + """ + input_tensors = [] + for input_data in node_data: + inbound_layer_name = input_data[0] + inbound_node_index = input_data[1] + inbound_tensor_index = input_data[2] + if len(input_data) == 3: + kwargs = {} + elif len(input_data) == 4: + kwargs = input_data[3] + else: + raise ValueError('Improperly formatted model config.') + if inbound_layer_name not in created_layers: + add_unprocessed_node(layer, node_data) + return + inbound_layer = created_layers[inbound_layer_name] + if len(inbound_layer._inbound_nodes) <= inbound_node_index: + add_unprocessed_node(layer, node_data) + return + inbound_node = inbound_layer._inbound_nodes[inbound_node_index] + input_tensors.append(inbound_node.output_tensors[inbound_tensor_index]) + # Call layer on its inputs, thus creating the node + # and building the layer if needed. + if input_tensors: + if len(input_tensors) == 1: + layer(input_tensors[0], **kwargs) + else: + layer(input_tensors, **kwargs) + + def process_layer(layer_data): + """Deserializes a layer, then call it on appropriate inputs. + + Arguments: + layer_data: layer config dict. + + Raises: + ValueError: In case of improperly formatted `layer_data` dict. + """ + layer_name = layer_data['name'] + + # Instantiate layer. + from tensorflow.python.keras._impl.keras.layers import deserialize as deserialize_layer # pylint: disable=g-import-not-at-top + + layer = deserialize_layer(layer_data, custom_objects=custom_objects) + created_layers[layer_name] = layer + + # Gather layer inputs. + inbound_nodes_data = layer_data['inbound_nodes'] + for node_data in inbound_nodes_data: + # We don't process nodes (i.e. make layer calls) + # on the fly because the inbound node may not yet exist, + # in case of layer shared at different topological depths + # (e.g. a model such as A(B(A(B(x))))) + add_unprocessed_node(layer, node_data) + + # First, we create all layers and enqueue nodes to be processed + for layer_data in config['layers']: + process_layer(layer_data) + # Then we process nodes in order of layer depth. + # Nodes that cannot yet be processed (if the inbound node + # does not yet exist) are re-enqueued, and the process + # is repeated until all nodes are processed. + while unprocessed_nodes: + for layer_data in config['layers']: + layer = created_layers[layer_data['name']] + if layer in unprocessed_nodes: + for node_data in unprocessed_nodes.pop(layer): + process_node(layer, node_data) + + name = config.get('name') + input_tensors = [] + output_tensors = [] + for layer_data in config['input_layers']: + layer_name, node_index, tensor_index = layer_data + assert layer_name in created_layers + layer = created_layers[layer_name] + layer_output_tensors = layer._inbound_nodes[node_index].output_tensors + input_tensors.append(layer_output_tensors[tensor_index]) + for layer_data in config['output_layers']: + layer_name, node_index, tensor_index = layer_data + assert layer_name in created_layers + layer = created_layers[layer_name] + layer_output_tensors = layer._inbound_nodes[node_index].output_tensors + output_tensors.append(layer_output_tensors[tensor_index]) + return cls(inputs=input_tensors, outputs=output_tensors, name=name) + + def save(self, filepath, overwrite=True, include_optimizer=True): + """Saves the model to a single HDF5 file. + + The savefile includes: + - The model architecture, allowing to re-instantiate the model. + - The model weights. + - The state of the optimizer, allowing to resume training + exactly where you left off. + + This allows you to save the entirety of the state of a model + in a single file. + + Saved models can be reinstantiated via `keras.models.load_model`. + The model returned by `load_model` + is a compiled model ready to be used (unless the saved model + was never compiled in the first place). + + Arguments: + filepath: String, path to the file to save the weights to. + overwrite: Whether to silently overwrite any existing file at the + target location, or provide the user with a manual prompt. + include_optimizer: If True, save optimizer's state together. + + Example: + + ```python + from keras.models import load_model + + model.save('my_model.h5') # creates a HDF5 file 'my_model.h5' + del model # deletes the existing model + + # returns a compiled model + # identical to the previous one + model = load_model('my_model.h5') + ``` + """ + if not self._is_graph_network: + raise NotImplementedError + + from tensorflow.python.keras._impl.keras.models import save_model # pylint: disable=g-import-not-at-top + save_model(self, filepath, overwrite, include_optimizer) + + def save_weights(self, filepath, overwrite=True): + """Dumps all layer weights to a HDF5 file. + + The weight file has: + - `layer_names` (attribute), a list of strings + (ordered names of model layers). + - For every layer, a `group` named `layer.name` + - For every such layer group, a group attribute `weight_names`, + a list of strings + (ordered names of weights tensor of the layer). + - For every weight in the layer, a dataset + storing the weight value, named after the weight tensor. + + Arguments: + filepath: String, path to the file to save the weights to. + overwrite: Whether to silently overwrite any existing file at the + target location, or provide the user with a manual prompt. + + Raises: + ImportError: If h5py is not available. + """ + if h5py is None: + raise ImportError('`save_weights` requires h5py.') + # If file exists and should not be overwritten: + if not overwrite and os.path.isfile(filepath): + proceed = ask_to_proceed_with_overwrite(filepath) + if not proceed: + return + with h5py.File(filepath, 'w') as f: + saving.save_weights_to_hdf5_group(f, self.layers) + + def load_weights(self, filepath, by_name=False): + """Loads all layer weights from a HDF5 save file. + + If `by_name` is False (default) weights are loaded + based on the network's topology, meaning the architecture + should be the same as when the weights were saved. + Note that layers that don't have weights are not taken + into account in the topological ordering, so adding or + removing layers is fine as long as they don't have weights. + + If `by_name` is True, weights are loaded into layers + only if they share the same name. This is useful + for fine-tuning or transfer-learning models where + some of the layers have changed. + + Arguments: + filepath: String, path to the weights file to load. + by_name: Boolean, whether to load weights by name + or by topological order. + + Raises: + ImportError: If h5py is not available. + """ + if h5py is None: + raise ImportError('`load_weights` requires h5py.') + with h5py.File(filepath, 'r') as f: + if 'layer_names' not in f.attrs and 'model_weights' in f: + f = f['model_weights'] + if by_name: + saving.load_weights_from_hdf5_group_by_name(f, self.layers) + else: + saving.load_weights_from_hdf5_group(f, self.layers) + + def _updated_config(self): + """Util shared between different serialization methods. + + Returns: + Model config with Keras version information added. + """ + from tensorflow.python.keras._impl.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top + + config = self.get_config() + model_config = { + 'class_name': self.__class__.__name__, + 'config': config, + 'keras_version': keras_version, + 'backend': K.backend() + } + return model_config + + def to_json(self, **kwargs): + """Returns a JSON string containing the network configuration. + + To load a network from a JSON save file, use + `keras.models.model_from_json(json_string, custom_objects={})`. + + Arguments: + **kwargs: Additional keyword arguments + to be passed to `json.dumps()`. + + Returns: + A JSON string. + """ + if not self._is_graph_network: + raise NotImplementedError + + def get_json_type(obj): + # If obj is any numpy type + if type(obj).__module__ == np.__name__: + return obj.item() + + # If obj is a python 'type' + if type(obj).__name__ == type.__name__: + return obj.__name__ + + raise TypeError('Not JSON Serializable:', obj) + + model_config = self._updated_config() + return json.dumps(model_config, default=get_json_type, **kwargs) + + def to_yaml(self, **kwargs): + """Returns a yaml string containing the network configuration. + + To load a network from a yaml save file, use + `keras.models.model_from_yaml(yaml_string, custom_objects={})`. + + `custom_objects` should be a dictionary mapping + the names of custom losses / layers / etc to the corresponding + functions / classes. + + Arguments: + **kwargs: Additional keyword arguments + to be passed to `yaml.dump()`. + + Returns: + A YAML string. + + Raises: + ImportError: if yaml module is not found. + """ + if not self._is_graph_network: + raise NotImplementedError + + if yaml is None: + raise ImportError('Requires yaml module installed.') + return yaml.dump(self._updated_config(), **kwargs) + + def summary(self, line_length=None, positions=None, print_fn=None): + """Prints a string summary of the network. + + Arguments: + line_length: Total length of printed lines + (e.g. set this to adapt the display to different + terminal window sizes). + positions: Relative or absolute positions of log elements + in each line. If not provided, + defaults to `[.33, .55, .67, 1.]`. + print_fn: Print function to use. Defaults to `print`. + It will be called on each line of the summary. + You can set it to a custom function + in order to capture the string summary. + """ + print_layer_summary(self, + line_length=line_length, + positions=positions, + print_fn=print_fn) + + +def get_source_inputs(tensor, layer=None, node_index=None): + """Returns the list of input tensors necessary to compute `tensor`. + + Output will always be a list of tensors + (potentially with 1 element). + + Arguments: + tensor: The tensor to start from. + layer: Origin layer of the tensor. Will be + determined via tensor._keras_history if not provided. + node_index: Origin node index of the tensor. + + Returns: + List of input tensors. + """ + if not hasattr(tensor, '_keras_history'): + return tensor + + if layer is None or node_index: + layer, node_index, _ = tensor._keras_history + if not layer._inbound_nodes: + return [tensor] + else: + node = layer._inbound_nodes[node_index] + if not node.inbound_layers: + # Reached an Input layer, stop recursion. + return node.input_tensors + else: + source_tensors = [] + for i in range(len(node.inbound_layers)): + x = node.input_tensors[i] + layer = node.inbound_layers[i] + node_index = node.node_indices[i] + previous_sources = get_source_inputs(x, layer, node_index) + # Avoid input redundancy. + for x in previous_sources: + if x not in source_tensors: + source_tensors.append(x) + return source_tensors + + +def _make_node_key(layer_name, node_index): + return layer_name + '_ib-' + str(node_index) + + +def _map_graph_network(inputs, outputs): + """Validates a network's topology and gather its layers and nodes. + + Arguments: + inputs: List of input tensors. + outputs: List of outputs tensors. + + Returns: + A tuple `(nodes, nodes_by_depth, layers, layers_by_depth)`. + - nodes: list of Node instances. + - nodes_by_depth: dict mapping ints (depth) to lists of node instances. + - layers: list of Layer instances. + - layers_by_depth: dict mapping ints (depth) to lists of layer instances. + + Raises: + ValueError: In case the network is not valid (e.g. disconnected graph). + """ + # Network_nodes: set of nodes included in the graph of layers + # (not all nodes included in the layers are relevant to the current graph). + network_nodes = set() # ids of all nodes relevant to the Network + nodes_depths = {} # dict {node: depth value} + layers_depths = {} # dict {layer: depth value} + layer_indices = {} # dict {layer: index in traversal} + nodes_in_decreasing_depth = [] + + def build_map(tensor, + finished_nodes, + nodes_in_progress, + layer, + node_index, + tensor_index): + """Builds a map of the graph of layers. + + This recursively updates the map `layer_indices`, + the list `nodes_in_decreasing_depth` and the set `network_nodes`. + + Arguments: + tensor: Some tensor in a graph. + finished_nodes: Set of nodes whose subgraphs have been traversed + completely. Useful to prevent duplicated work. + nodes_in_progress: Set of nodes that are currently active on the + recursion stack. Useful to detect cycles. + layer: Layer from which `tensor` comes from. If not provided, + will be obtained from `tensor._keras_history`. + node_index: Node index from which `tensor` comes from. + tensor_index: Tensor_index from which `tensor` comes from. + + Raises: + ValueError: if a cycle is detected. + """ + node = layer._inbound_nodes[node_index] # pylint: disable=protected-access + + # Prevent cycles. + if node in nodes_in_progress: + raise ValueError('The tensor ' + str(tensor) + ' at layer "' + + layer.name + '" is part of a cycle.') + + # Don't repeat work for shared subgraphs + if node in finished_nodes: + return + + node_key = _make_node_key(layer.name, node_index) + # Update network_nodes. + network_nodes.add(node_key) + + # Store the traversal order for layer sorting. + if layer not in layer_indices: + layer_indices[layer] = len(layer_indices) + + nodes_in_progress.add(node) + + # Propagate to all previous tensors connected to this node. + for i in range(len(node.inbound_layers)): + x = node.input_tensors[i] + layer = node.inbound_layers[i] + node_index = node.node_indices[i] + tensor_index = node.tensor_indices[i] + build_map(x, finished_nodes, nodes_in_progress, layer, + node_index, tensor_index) + + finished_nodes.add(node) + nodes_in_progress.remove(node) + nodes_in_decreasing_depth.append(node) + + finished_nodes = set() + nodes_in_progress = set() + for x in outputs: + layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access + build_map(x, finished_nodes, nodes_in_progress, + layer=layer, + node_index=node_index, + tensor_index=tensor_index) + + for node in reversed(nodes_in_decreasing_depth): + # If the depth is not set, the node has no outbound nodes (depth 0). + depth = nodes_depths.setdefault(node, 0) + + # Update the depth of the corresponding layer + previous_depth = layers_depths.get(node.outbound_layer, 0) + # If we've seen this layer before at a higher depth, + # we should use that depth instead of the node depth. + # This is necessary for shared layers that have inputs at different + # depth levels in the graph. + depth = max(depth, previous_depth) + layers_depths[node.outbound_layer] = depth + nodes_depths[node] = depth + + # Update the depth of inbound nodes. + # The "depth" of a node is the max of the depths + # of all layers it is connected to. + for i in range(len(node.inbound_layers)): + inbound_layer = node.inbound_layers[i] + node_index = node.node_indices[i] + inbound_node = inbound_layer._inbound_nodes[node_index] # pylint: disable=protected-access + previous_depth = nodes_depths.get(inbound_node, 0) + nodes_depths[inbound_node] = max(depth + 1, previous_depth) + + # Build a dict {depth: list of nodes with this depth} + nodes_by_depth = {} + for node, depth in nodes_depths.items(): + if depth not in nodes_by_depth: + nodes_by_depth[depth] = [] + nodes_by_depth[depth].append(node) + + # Build a dict {depth: list of layers with this depth} + layers_by_depth = {} + for layer, depth in layers_depths.items(): + if depth not in layers_by_depth: + layers_by_depth[depth] = [] + layers_by_depth[depth].append(layer) + + # Get sorted list of layer depths. + depth_keys = list(layers_by_depth.keys()) + depth_keys.sort(reverse=True) + + # Set self.layers and self._layers_by_depth. + layers = [] + for depth in depth_keys: + layers_for_depth = layers_by_depth[depth] + # Network.layers needs to have a deterministic order: + # here we order them by traversal order. + layers_for_depth.sort(key=lambda x: layer_indices[x]) + layers.extend(layers_for_depth) + + # Get sorted list of node depths. + depth_keys = list(nodes_by_depth.keys()) + depth_keys.sort(reverse=True) + + # Check that all tensors required are computable. + # computable_tensors: all tensors in the graph + # that can be computed from the inputs provided. + computable_tensors = [] + for x in inputs: + computable_tensors.append(x) + + layers_with_complete_input = [] # To provide a better error msg. + for depth in depth_keys: + for node in nodes_by_depth[depth]: + layer = node.outbound_layer + if layer: + for x in node.input_tensors: + if x not in computable_tensors: + raise ValueError('Graph disconnected: ' + 'cannot obtain value for tensor ' + str(x) + + ' at layer "' + layer.name + '". ' + 'The following previous layers ' + 'were accessed without issue: ' + + str(layers_with_complete_input)) + for x in node.output_tensors: + computable_tensors.append(x) + layers_with_complete_input.append(layer.name) + + # Ensure name unicity, which will be crucial for serialization + # (since serialized nodes refer to layers by their name). + all_names = [layer.name for layer in layers] + for name in all_names: + if all_names.count(name) != 1: + raise ValueError('The name "' + name + '" is used ' + + str(all_names.count(name)) + ' times in the model. ' + 'All layer names should be unique.') + return network_nodes, nodes_by_depth, layers, layers_by_depth diff --git a/tensorflow/python/keras/_impl/keras/engine/saving.py b/tensorflow/python/keras/_impl/keras/engine/saving.py new file mode 100644 index 0000000000000000000000000000000000000000..2ad06ca4fdcd55c12ba3ba192751f2f05aacc7ec --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/saving.py @@ -0,0 +1,844 @@ +# 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. +# ============================================================================== +# pylint: disable=protected-access +"""Model saving utilities. +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import json +import os + +import numpy as np +from six.moves import zip # pylint: disable=redefined-builtin + +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import optimizers +from tensorflow.python.keras._impl.keras.utils import conv_utils +from tensorflow.python.keras._impl.keras.utils.io_utils import ask_to_proceed_with_overwrite +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export + +# pylint: disable=g-import-not-at-top +try: + import h5py + HDF5_OBJECT_HEADER_LIMIT = 64512 +except ImportError: + h5py = None + +try: + import yaml +except ImportError: + yaml = None +# pylint: enable=g-import-not-at-top + + +@tf_export('keras.models.save_model') +def save_model(model, filepath, overwrite=True, include_optimizer=True): + """Saves a model to a HDF5 file. + + The saved model contains: + - the model's configuration (topology) + - the model's weights + - the model's optimizer's state (if any) + + Thus the saved model can be reinstantiated in + the exact same state, without any of the code + used for model definition or training. + + Arguments: + model: Keras model instance to be saved. + filepath: String, path where to save the model. + overwrite: Whether we should overwrite any existing + model at the target location, or instead + ask the user with a manual prompt. + include_optimizer: If True, save optimizer's state together. + + Raises: + ImportError: if h5py is not available. + """ + + if h5py is None: + raise ImportError('`save_model` requires h5py.') + + def get_json_type(obj): + """Serializes any object to a JSON-serializable structure. + + Arguments: + obj: the object to serialize + + Returns: + JSON-serializable structure representing `obj`. + + Raises: + TypeError: if `obj` cannot be serialized. + """ + # if obj is a serializable Keras class instance + # e.g. optimizer, layer + if hasattr(obj, 'get_config'): + return {'class_name': obj.__class__.__name__, 'config': obj.get_config()} + + # if obj is any numpy type + if type(obj).__module__ == np.__name__: + if isinstance(obj, np.ndarray): + return {'type': type(obj), 'value': obj.tolist()} + else: + return obj.item() + + # misc functions (e.g. loss function) + if callable(obj): + return obj.__name__ + + # if obj is a python 'type' + if type(obj).__name__ == type.__name__: + return obj.__name__ + + raise TypeError('Not JSON Serializable:', obj) + + from tensorflow.python.keras._impl.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top + + # If file exists and should not be overwritten. + if not overwrite and os.path.isfile(filepath): + proceed = ask_to_proceed_with_overwrite(filepath) + if not proceed: + return + + with h5py.File(filepath, mode='w') as f: + f.attrs['keras_version'] = str(keras_version).encode('utf8') + f.attrs['backend'] = K.backend().encode('utf8') + f.attrs['model_config'] = json.dumps( + { + 'class_name': model.__class__.__name__, + 'config': model.get_config() + }, + default=get_json_type).encode('utf8') + + model_weights_group = f.create_group('model_weights') + model_layers = model.layers + save_weights_to_hdf5_group(model_weights_group, model_layers) + + if include_optimizer and hasattr(model, 'optimizer'): + if isinstance(model.optimizer, optimizers.TFOptimizer): + logging.warning( + 'TensorFlow optimizers do not ' + 'make it possible to access ' + 'optimizer attributes or optimizer state ' + 'after instantiation. ' + 'As a result, we cannot save the optimizer ' + 'as part of the model save file.' + 'You will have to compile your model again after loading it. ' + 'Prefer using a Keras optimizer instead ' + '(see keras.io/optimizers).') + else: + f.attrs['training_config'] = json.dumps( + { + 'optimizer_config': { + 'class_name': model.optimizer.__class__.__name__, + 'config': model.optimizer.get_config() + }, + 'loss': model.loss, + 'metrics': model.metrics, + 'sample_weight_mode': model.sample_weight_mode, + 'loss_weights': model.loss_weights, + }, + default=get_json_type).encode('utf8') + + # Save optimizer weights. + symbolic_weights = getattr(model.optimizer, 'weights') + if symbolic_weights: + optimizer_weights_group = f.create_group('optimizer_weights') + weight_values = K.batch_get_value(symbolic_weights) + weight_names = [] + for w, val in zip(symbolic_weights, weight_values): + name = str(w.name) + weight_names.append(name.encode('utf8')) + optimizer_weights_group.attrs['weight_names'] = weight_names + for name, val in zip(weight_names, weight_values): + param_dset = optimizer_weights_group.create_dataset( + name, val.shape, dtype=val.dtype) + if not val.shape: + # scalar + param_dset[()] = val + else: + param_dset[:] = val + f.flush() + + +@tf_export('keras.models.load_model') +def load_model(filepath, custom_objects=None, compile=True): # pylint: disable=redefined-builtin + """Loads a model saved via `save_model`. + + Arguments: + filepath: String, path to the saved model. + custom_objects: Optional dictionary mapping names + (strings) to custom classes or functions to be + considered during deserialization. + compile: Boolean, whether to compile the model + after loading. + + Returns: + A Keras model instance. If an optimizer was found + as part of the saved model, the model is already + compiled. Otherwise, the model is uncompiled and + a warning will be displayed. When `compile` is set + to False, the compilation is omitted without any + warning. + + Raises: + ImportError: if h5py is not available. + ValueError: In case of an invalid savefile. + """ + if h5py is None: + raise ImportError('`load_model` requires h5py.') + + if not custom_objects: + custom_objects = {} + + def convert_custom_objects(obj): + """Handles custom object lookup. + + Arguments: + obj: object, dict, or list. + + Returns: + The same structure, where occurrences + of a custom object name have been replaced + with the custom object. + """ + if isinstance(obj, list): + deserialized = [] + for value in obj: + deserialized.append(convert_custom_objects(value)) + return deserialized + if isinstance(obj, dict): + deserialized = {} + for key, value in obj.items(): + deserialized[key] = convert_custom_objects(value) + return deserialized + if obj in custom_objects: + return custom_objects[obj] + return obj + + with h5py.File(filepath, mode='r') as f: + # instantiate model + model_config = f.attrs.get('model_config') + if model_config is None: + raise ValueError('No model found in config file.') + model_config = json.loads(model_config.decode('utf-8')) + model = model_from_config(model_config, custom_objects=custom_objects) + + # set weights + load_weights_from_hdf5_group(f['model_weights'], model.layers) + + # Early return if compilation is not required. + if not compile: + return model + + # instantiate optimizer + training_config = f.attrs.get('training_config') + if training_config is None: + logging.warning('No training configuration found in save file: ' + 'the model was *not* compiled. Compile it manually.') + return model + training_config = json.loads(training_config.decode('utf-8')) + optimizer_config = training_config['optimizer_config'] + optimizer = optimizers.deserialize( + optimizer_config, custom_objects=custom_objects) + + # Recover loss functions and metrics. + loss = convert_custom_objects(training_config['loss']) + metrics = convert_custom_objects(training_config['metrics']) + sample_weight_mode = training_config['sample_weight_mode'] + loss_weights = training_config['loss_weights'] + + # Compile model. + model.compile( + optimizer=optimizer, + loss=loss, + metrics=metrics, + loss_weights=loss_weights, + sample_weight_mode=sample_weight_mode) + + # Set optimizer weights. + if 'optimizer_weights' in f: + # Build train function (to get weight updates). + model._make_train_function() + optimizer_weights_group = f['optimizer_weights'] + optimizer_weight_names = [ + n.decode('utf8') + for n in optimizer_weights_group.attrs['weight_names'] + ] + optimizer_weight_values = [ + optimizer_weights_group[n] for n in optimizer_weight_names + ] + try: + model.optimizer.set_weights(optimizer_weight_values) + except ValueError: + logging.warning('Error in loading the saved optimizer ' + 'state. As a result, your model is ' + 'starting with a freshly initialized ' + 'optimizer.') + return model + + +@tf_export('keras.models.model_from_config') +def model_from_config(config, custom_objects=None): + """Instantiates a Keras model from its config. + + Arguments: + config: Configuration dictionary. + custom_objects: Optional dictionary mapping names + (strings) to custom classes or functions to be + considered during deserialization. + + Returns: + A Keras model instance (uncompiled). + + Raises: + TypeError: if `config` is not a dictionary. + """ + if isinstance(config, list): + raise TypeError('`model_from_config` expects a dictionary, not a list. ' + 'Maybe you meant to use ' + '`Sequential.from_config(config)`?') + from tensorflow.python.keras._impl.keras.layers import deserialize # pylint: disable=g-import-not-at-top + return deserialize(config, custom_objects=custom_objects) + + +@tf_export('keras.models.model_from_yaml') +def model_from_yaml(yaml_string, custom_objects=None): + """Parses a yaml model configuration file and returns a model instance. + + Arguments: + yaml_string: YAML string encoding a model configuration. + custom_objects: Optional dictionary mapping names + (strings) to custom classes or functions to be + considered during deserialization. + + Returns: + A Keras model instance (uncompiled). + + Raises: + ImportError: if yaml module is not found. + """ + if yaml is None: + raise ImportError('Requires yaml module installed.') + config = yaml.load(yaml_string) + from tensorflow.python.keras._impl.keras.layers import deserialize # pylint: disable=g-import-not-at-top + return deserialize(config, custom_objects=custom_objects) + + +@tf_export('keras.models.model_from_json') +def model_from_json(json_string, custom_objects=None): + """Parses a JSON model configuration file and returns a model instance. + + Arguments: + json_string: JSON string encoding a model configuration. + custom_objects: Optional dictionary mapping names + (strings) to custom classes or functions to be + considered during deserialization. + + Returns: + A Keras model instance (uncompiled). + """ + config = json.loads(json_string) + from tensorflow.python.keras._impl.keras.layers import deserialize # pylint: disable=g-import-not-at-top + return deserialize(config, custom_objects=custom_objects) + + +def preprocess_weights_for_loading(layer, + weights, + original_keras_version=None, + original_backend=None): + """Converts layers weights from Keras 1 format to Keras 2. + + Arguments: + layer: Layer instance. + weights: List of weights values (Numpy arrays). + original_keras_version: Keras version for the weights, as a string. + original_backend: Keras backend the weights were trained with, + as a string. + + Returns: + A list of weights values (Numpy arrays). + """ + if layer.__class__.__name__ == 'Bidirectional': + num_weights_per_layer = len(weights) // 2 + forward_weights = preprocess_weights_for_loading( + layer.forward_layer, weights[:num_weights_per_layer], + original_keras_version, original_backend) + backward_weights = preprocess_weights_for_loading( + layer.backward_layer, weights[num_weights_per_layer:], + original_keras_version, original_backend) + weights = forward_weights + backward_weights + + if original_keras_version == '1': + if layer.__class__.__name__ == 'TimeDistributed': + weights = preprocess_weights_for_loading( + layer.layer, weights, original_keras_version, original_backend) + + if layer.__class__.__name__ == 'Conv1D': + shape = weights[0].shape + # Handle Keras 1.1 format + if shape[:2] != (layer.kernel_size[0], 1) or shape[3] != layer.filters: + # Legacy shape: + # (filters, input_dim, filter_length, 1) + assert shape[0] == layer.filters and shape[2:] == (layer.kernel_size[0], + 1) + weights[0] = np.transpose(weights[0], (2, 3, 1, 0)) + weights[0] = weights[0][:, 0, :, :] + + if layer.__class__.__name__ == 'Conv2D': + if layer.data_format == 'channels_first': + # old: (filters, stack_size, kernel_rows, kernel_cols) + # new: (kernel_rows, kernel_cols, stack_size, filters) + weights[0] = np.transpose(weights[0], (2, 3, 1, 0)) + + if layer.__class__.__name__ == 'Conv2DTranspose': + if layer.data_format == 'channels_last': + # old: (kernel_rows, kernel_cols, stack_size, filters) + # new: (kernel_rows, kernel_cols, filters, stack_size) + weights[0] = np.transpose(weights[0], (0, 1, 3, 2)) + if layer.data_format == 'channels_first': + # old: (filters, stack_size, kernel_rows, kernel_cols) + # new: (kernel_rows, kernel_cols, filters, stack_size) + weights[0] = np.transpose(weights[0], (2, 3, 0, 1)) + + if layer.__class__.__name__ == 'Conv3D': + if layer.data_format == 'channels_first': + # old: (filters, stack_size, ...) + # new: (..., stack_size, filters) + weights[0] = np.transpose(weights[0], (2, 3, 4, 1, 0)) + + if layer.__class__.__name__ == 'GRU': + if len(weights) == 9: + kernel = np.concatenate([weights[0], weights[3], weights[6]], axis=-1) + recurrent_kernel = np.concatenate( + [weights[1], weights[4], weights[7]], axis=-1) + bias = np.concatenate([weights[2], weights[5], weights[8]], axis=-1) + weights = [kernel, recurrent_kernel, bias] + + if layer.__class__.__name__ == 'LSTM': + if len(weights) == 12: + # old: i, c, f, o + # new: i, f, c, o + kernel = np.concatenate( + [weights[0], weights[6], weights[3], weights[9]], axis=-1) + recurrent_kernel = np.concatenate( + [weights[1], weights[7], weights[4], weights[10]], axis=-1) + bias = np.concatenate( + [weights[2], weights[8], weights[5], weights[11]], axis=-1) + weights = [kernel, recurrent_kernel, bias] + + if layer.__class__.__name__ == 'ConvLSTM2D': + if len(weights) == 12: + kernel = np.concatenate( + [weights[0], weights[6], weights[3], weights[9]], axis=-1) + recurrent_kernel = np.concatenate( + [weights[1], weights[7], weights[4], weights[10]], axis=-1) + bias = np.concatenate( + [weights[2], weights[8], weights[5], weights[11]], axis=-1) + if layer.data_format == 'channels_first': + # old: (filters, stack_size, kernel_rows, kernel_cols) + # new: (kernel_rows, kernel_cols, stack_size, filters) + kernel = np.transpose(kernel, (2, 3, 1, 0)) + recurrent_kernel = np.transpose(recurrent_kernel, (2, 3, 1, 0)) + weights = [kernel, recurrent_kernel, bias] + + if layer.__class__.__name__ in ['Model', 'Sequential']: + new_weights = [] + # trainable weights + for sublayer in layer.layers: + num_weights = len(sublayer.trainable_weights) + if num_weights > 0: + new_weights.extend( + preprocess_weights_for_loading( + layer=sublayer, + weights=weights[:num_weights], + original_keras_version=original_keras_version, + original_backend=original_backend)) + weights = weights[num_weights:] + + # non-trainable weights + for sublayer in layer.layers: + num_weights = len([ + l for l in sublayer.weights if l not in sublayer.trainable_weights + ]) + if num_weights > 0: + new_weights.extend( + preprocess_weights_for_loading( + layer=sublayer, + weights=weights[:num_weights], + original_keras_version=original_keras_version, + original_backend=original_backend)) + weights = weights[num_weights:] + weights = new_weights + + conv_layers = ['Conv1D', 'Conv2D', 'Conv3D', 'Conv2DTranspose', 'ConvLSTM2D'] + if layer.__class__.__name__ in conv_layers: + if original_backend == 'theano': + weights[0] = conv_utils.convert_kernel(weights[0]) + if layer.__class__.__name__ == 'ConvLSTM2D': + weights[1] = conv_utils.convert_kernel(weights[1]) + if K.int_shape(layer.weights[0]) != weights[0].shape: + weights[0] = np.transpose(weights[0], (3, 2, 0, 1)) + if layer.__class__.__name__ == 'ConvLSTM2D': + weights[1] = np.transpose(weights[1], (3, 2, 0, 1)) + + # Convert the weights of CuDNNLSTM so that they could be loaded into LSTM + if layer.__class__.__name__ == 'LSTM' and len(weights) == 3: + # Determine if loading a CuDNNLSTM layer from the number of bias weights: + # CuDNNLSTM has (units * 8) weights; while LSTM has (units * 4) + # if there's no bias weight in the file, skip this conversion + units = weights[1].shape[0] + bias = weights[2] + if len(bias) == units * 8: + # reshape the kernels + kernels = np.split(weights[0], 4, axis=1) + kernels = [ + kernel.reshape(-1).reshape(kernel.shape, order='F') + for kernel in kernels + ] + weights[0] = np.concatenate(kernels, axis=1) + + # transpose the recurrent kernels + recurrent_kernels = np.split(weights[1], 4, axis=1) + recurrent_kernels = [kernel.T for kernel in recurrent_kernels] + weights[1] = np.concatenate(recurrent_kernels, axis=1) + + # split the bias into half and merge + weights[2] = bias[:units * 4] + bias[units * 4:] + + return convert_rnn_weights(layer, weights) + + +def convert_rnn_weights(layer, weights): + """Converts weights for RNN layers between native and CuDNN format. + + Input kernels for each gate are transposed and converted between Fortran + and C layout, recurrent kernels are transposed. For LSTM biases are summed/ + split in half, for GRU biases are reshaped. + + Weights can be converted in both directions between `LSTM` and`CuDNNSLTM` + and between `CuDNNGRU` and `GRU(reset_after=True)`. Default `GRU` is not + compatible with `CuDNNGRU`. + + For missing biases in `LSTM`/`GRU` (`use_bias=False`) no conversion is made. + + Arguments: + layer: Target layer instance. + weights: List of source weights values (input kernels, recurrent + kernels, [biases]) (Numpy arrays). + + Returns: + A list of converted weights values (Numpy arrays). + + Raises: + ValueError: for incompatible GRU layer/weights or incompatible biases + """ + + def transform_kernels(kernels, func, n_gates): + """Transforms kernel for each gate separately using given function. + + Arguments: + kernels: Stacked array of kernels for individual gates. + func: Function applied to kernel of each gate. + n_gates: Number of gates (4 for LSTM, 3 for GRU). + Returns: + Stacked array of transformed kernels. + """ + return np.hstack([func(k) for k in np.hsplit(kernels, n_gates)]) + + def transpose_input(from_cudnn): + """Makes a function that transforms input kernels from/to CuDNN format. + + It keeps the shape, but changes between the layout (Fortran/C). Eg.: + + ``` + Keras CuDNN + [[0, 1, 2], <---> [[0, 2, 4], + [3, 4, 5]] [1, 3, 5]] + ``` + + It can be passed to `transform_kernels()`. + + Arguments: + from_cudnn: `True` if source weights are in CuDNN format, `False` + if they're in plain Keras format. + Returns: + Function that converts input kernel to the other format. + """ + order = 'F' if from_cudnn else 'C' + + def transform(kernel): + return kernel.T.reshape(kernel.shape, order=order) + + return transform + + target_class = layer.__class__.__name__ + + # convert the weights between CuDNNLSTM and LSTM + if target_class in ['LSTM', 'CuDNNLSTM'] and len(weights) == 3: + # determine if we're loading a CuDNNLSTM layer + # from the number of bias weights: + # CuDNNLSTM has (units * 8) weights; while LSTM has (units * 4) + # if there's no bias weight in the file, skip this conversion + units = weights[1].shape[0] + bias_shape = weights[2].shape + n_gates = 4 + + if bias_shape == (2 * units * n_gates,): + source = 'CuDNNLSTM' + elif bias_shape == (units * n_gates,): + source = 'LSTM' + else: + raise ValueError('Invalid bias shape: ' + str(bias_shape)) + + def convert_lstm_weights(weights, from_cudnn=True): + # Transpose (and reshape) input and recurrent kernels. + kernels = transform_kernels(weights[0], transpose_input(from_cudnn), + n_gates) + recurrent_kernels = transform_kernels(weights[1], lambda k: k.T, n_gates) + if from_cudnn: # Merge input and recurrent biases into a single set. + biases = np.sum(np.split(weights[2], 2, axis=0), axis=0) + else: + # Split single set of biases evenly to two sets. + biases = np.tile(0.5 * weights[2], 2) + return [kernels, recurrent_kernels, biases] + + if source != target_class: + weights = convert_lstm_weights(weights, from_cudnn=source == 'CuDNNLSTM') + + # TODO(fchollet): add feature after GRU is refactored: + # convert the weights between `CuDNNGRU` and `GRU(reset_after=True)` + return weights + + +def save_weights_to_hdf5_group(f, layers): + from tensorflow.python.keras._impl.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top + + save_attributes_to_hdf5_group( + f, 'layer_names', [layer.name.encode('utf8') for layer in layers]) + f.attrs['backend'] = K.backend().encode('utf8') + f.attrs['keras_version'] = str(keras_version).encode('utf8') + + for layer in layers: + g = f.create_group(layer.name) + symbolic_weights = layer.weights + weight_values = K.batch_get_value(symbolic_weights) + weight_names = [] + for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)): + if hasattr(w, 'name') and w.name: + name = str(w.name) + else: + name = 'param_' + str(i) + weight_names.append(name.encode('utf8')) + save_attributes_to_hdf5_group(g, 'weight_names', weight_names) + for name, val in zip(weight_names, weight_values): + param_dset = g.create_dataset(name, val.shape, dtype=val.dtype) + if not val.shape: + # scalar + param_dset[()] = val + else: + param_dset[:] = val + + +def load_weights_from_hdf5_group(f, layers): + """Implements topological (order-based) weight loading. + + Arguments: + f: A pointer to a HDF5 group. + layers: a list of target layers. + + Raises: + ValueError: in case of mismatch between provided layers + and weights file. + """ + if 'keras_version' in f.attrs: + original_keras_version = f.attrs['keras_version'].decode('utf8') + else: + original_keras_version = '1' + if 'backend' in f.attrs: + original_backend = f.attrs['backend'].decode('utf8') + else: + original_backend = None + + filtered_layers = [] + for layer in layers: + weights = layer.weights + if weights: + filtered_layers.append(layer) + + layer_names = load_attributes_from_hdf5_group(f, 'layer_names') + filtered_layer_names = [] + for name in layer_names: + g = f[name] + weight_names = load_attributes_from_hdf5_group(g, 'weight_names') + if weight_names: + filtered_layer_names.append(name) + layer_names = filtered_layer_names + if len(layer_names) != len(filtered_layers): + raise ValueError('You are trying to load a weight file ' + 'containing ' + str(len(layer_names)) + + ' layers into a model with ' + str(len(filtered_layers)) + + ' layers.') + + # We batch weight value assignments in a single backend call + # which provides a speedup in TensorFlow. + weight_value_tuples = [] + for k, name in enumerate(layer_names): + g = f[name] + weight_names = load_attributes_from_hdf5_group(g, 'weight_names') + weight_values = [g[weight_name] for weight_name in weight_names] + layer = filtered_layers[k] + symbolic_weights = layer.weights + weight_values = preprocess_weights_for_loading( + layer, weight_values, original_keras_version, original_backend) + if len(weight_values) != len(symbolic_weights): + raise ValueError('Layer #' + str(k) + ' (named "' + layer.name + + '" in the current model) was found to ' + 'correspond to layer ' + name + ' in the save file. ' + 'However the new layer ' + layer.name + ' expects ' + + str(len(symbolic_weights)) + + ' weights, but the saved weights have ' + + str(len(weight_values)) + ' elements.') + weight_value_tuples += zip(symbolic_weights, weight_values) + K.batch_set_value(weight_value_tuples) + + +def load_weights_from_hdf5_group_by_name(f, layers): + """Implements name-based weight loading. + + (instead of topological weight loading). + + Layers that have no matching name are skipped. + + Arguments: + f: A pointer to a HDF5 group. + layers: a list of target layers. + + Raises: + ValueError: in case of mismatch between provided layers + and weights file. + """ + if 'keras_version' in f.attrs: + original_keras_version = f.attrs['keras_version'].decode('utf8') + else: + original_keras_version = '1' + if 'backend' in f.attrs: + original_backend = f.attrs['backend'].decode('utf8') + else: + original_backend = None + + # New file format. + layer_names = load_attributes_from_hdf5_group(f, 'layer_names') + + # Reverse index of layer name to list of layers with name. + index = {} + for layer in layers: + if layer.name: + index.setdefault(layer.name, []).append(layer) + + # We batch weight value assignments in a single backend call + # which provides a speedup in TensorFlow. + weight_value_tuples = [] + for k, name in enumerate(layer_names): + g = f[name] + weight_names = load_attributes_from_hdf5_group(g, 'weight_names') + weight_values = [g[weight_name] for weight_name in weight_names] + + for layer in index.get(name, []): + symbolic_weights = layer.weights + weight_values = preprocess_weights_for_loading( + layer, weight_values, original_keras_version, original_backend) + if len(weight_values) != len(symbolic_weights): + raise ValueError('Layer #' + str(k) + ' (named "' + layer.name + + '") expects ' + str(len(symbolic_weights)) + + ' weight(s), but the saved weights' + ' have ' + + str(len(weight_values)) + ' element(s).') + # Set values. + for i in range(len(weight_values)): + weight_value_tuples.append((symbolic_weights[i], weight_values[i])) + K.batch_set_value(weight_value_tuples) + + +def save_attributes_to_hdf5_group(group, name, data): + """Saves attributes (data) of the specified name into the HDF5 group. + + This method deals with an inherent problem of HDF5 file which is not + able to store data larger than HDF5_OBJECT_HEADER_LIMIT bytes. + + Arguments: + group: A pointer to a HDF5 group. + name: A name of the attributes to save. + data: Attributes data to store. + + Raises: + RuntimeError: If any single attribute is too large to be saved. + """ + # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT` + # because in that case even chunking the array would not make the saving + # possible. + bad_attributes = [x for x in data if len(x) > HDF5_OBJECT_HEADER_LIMIT] + + # Expecting this to never be true. + if bad_attributes: + raise RuntimeError('The following attributes cannot be saved to HDF5 ' + 'file because they are larger than %d bytes: %s' % + (HDF5_OBJECT_HEADER_LIMIT, + ', '.join([x for x in bad_attributes]))) + + data_npy = np.asarray(data) + + num_chunks = 1 + chunked_data = np.array_split(data_npy, num_chunks) + + # This will never loop forever thanks to the test above. + while any([x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data]): + num_chunks += 1 + chunked_data = np.array_split(data_npy, num_chunks) + + if num_chunks > 1: + for chunk_id, chunk_data in enumerate(chunked_data): + group.attrs['%s%d' % (name, chunk_id)] = chunk_data + else: + group.attrs[name] = data + + +def load_attributes_from_hdf5_group(group, name): + """Loads attributes of the specified name from the HDF5 group. + + This method deals with an inherent problem + of HDF5 file which is not able to store + data larger than HDF5_OBJECT_HEADER_LIMIT bytes. + + Arguments: + group: A pointer to a HDF5 group. + name: A name of the attributes to load. + + Returns: + data: Attributes data. + """ + if name in group.attrs: + data = [n.decode('utf8') for n in group.attrs[name]] + else: + data = [] + chunk_id = 0 + while '%s%d' % (name, chunk_id) in group.attrs: + data.extend( + [n.decode('utf8') for n in group.attrs['%s%d' % (name, chunk_id)]]) + chunk_id += 1 + return data diff --git a/tensorflow/python/keras/_impl/keras/engine/saving_test.py b/tensorflow/python/keras/_impl/keras/engine/saving_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4a18cc2e119d7cfb3f15da593d4944abd445905b --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/saving_test.py @@ -0,0 +1,461 @@ +# 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 model saving.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import shutil +import tempfile + +import numpy as np + +from tensorflow.python.keras._impl import keras +from tensorflow.python.platform import test +from tensorflow.python.training import training as training_module + +try: + import h5py # pylint:disable=g-import-not-at-top +except ImportError: + h5py = None + + +class TestWeightSavingAndLoading(test.TestCase): + + def test_weight_loading(self): + with self.test_session(): + a = keras.layers.Input(shape=(2,)) + x = keras.layers.Dense(3)(a) + b = keras.layers.Dense(1)(x) + model = keras.models.Model(a, b) + + x = np.random.random((3, 2)) + ref_y = model.predict(x) + weights = model.get_weights() + model.set_weights(weights) + y = model.predict(x) + self.assertAllClose(ref_y, y) + + with self.assertRaises(ValueError): + model.set_weights(weights[1:]) + with self.assertRaises(ValueError): + model.set_weights(weights[::-1]) + + if h5py is None: + return # Skip rest of test if H5py isn't available. + + temp_dir = self.get_temp_dir() + self.addCleanup(shutil.rmtree, temp_dir) + + h5_path = os.path.join(temp_dir, 'test.h5') + model.save_weights(h5_path) + model.load_weights(h5_path) + y = model.predict(x) + self.assertAllClose(ref_y, y) + + model.load_weights(h5_path, by_name=True) + y = model.predict(x) + self.assertAllClose(ref_y, y) + + def test_weight_preprocessing(self): + input_dim = 3 + output_dim = 3 + size = 2 + cases = [ + [ + (keras.layers.Bidirectional(keras.layers.SimpleRNN(2))), + [np.random.random((2, 1)), np.random.random((2, 1))], + (None, 3, 2), + ], + [ + (keras.layers.TimeDistributed(keras.layers.Dense(1))), + [np.random.random((2, 1)), np.random.random((1,))], + (None, 3, 2), + ], + [ + (keras.layers.Conv1D(output_dim, size, use_bias=False)), + [np.random.random((output_dim, input_dim, size, 1))], + (None, 4, input_dim), + ], + [ + (keras.layers.Conv2D(output_dim, size, + use_bias=False, data_format='channels_first')), + [np.random.random((output_dim, input_dim, size, size))], + (None, input_dim, 4, 4), + ], + [ + (keras.layers.Conv2DTranspose(output_dim, size, + use_bias=False, + data_format='channels_first')), + [np.random.random((output_dim, input_dim, size, size))], + (None, input_dim, 4, 4), + ], + [ + (keras.layers.Conv2DTranspose(output_dim, size, + use_bias=False, + data_format='channels_last')), + [np.random.random((size, size, input_dim, output_dim))], + (None, 4, 4, input_dim), + ], + [ + (keras.layers.Conv3D(output_dim, size, + use_bias=False, data_format='channels_first')), + [np.random.random((output_dim, input_dim, size, size, size))], + (None, input_dim, 4, 4, 4), + ], + [ + (keras.layers.GRU(output_dim)), + [np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,)), + np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,)), + np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,))], + (None, 4, input_dim), + ], + [ + (keras.layers.LSTM(output_dim)), + [np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,)), + np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,)), + np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,)), + np.random.random((input_dim, output_dim)), + np.random.random((output_dim, output_dim)), + np.random.random((output_dim,))], + (None, 4, input_dim), + ], + ] + for layer, weights, input_shape in cases: + layer.build(input_shape) + _ = keras.engine.saving.preprocess_weights_for_loading( + layer, weights, original_keras_version='1') + + model = keras.models.Sequential([keras.layers.Dense(2, input_dim=2)]) + _ = keras.engine.saving.preprocess_weights_for_loading( + model, model.weights, original_keras_version='1') + + x = keras.Input((2,)) + y = keras.layers.Dense(2)(x) + model = keras.models.Model(x, y) + _ = keras.engine.saving.preprocess_weights_for_loading( + model, model.weights, original_keras_version='1') + + def test_sequential_weight_loading(self): + if h5py is None: + return + + temp_dir = self.get_temp_dir() + self.addCleanup(shutil.rmtree, temp_dir) + h5_path = os.path.join(temp_dir, 'test.h5') + + num_hidden = 5 + input_dim = 3 + batch_size = 5 + num_classes = 2 + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) + model.add(keras.layers.Dense(num_classes)) + + x = np.random.random((batch_size, input_dim)) + ref_y = model.predict(x) + + model.save_weights(h5_path) + + model = keras.models.Sequential() + model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) + model.add(keras.layers.Dense(num_classes)) + model.load_weights(h5_path) + y = model.predict(x) + + self.assertAllClose(y, ref_y) + + +class TestWholeModelSaving(test.TestCase): + + def test_sequential_model_saving(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(2, input_shape=(3,))) + model.add(keras.layers.RepeatVector(3)) + model.add(keras.layers.TimeDistributed(keras.layers.Dense(3))) + model.compile(loss=keras.losses.MSE, + optimizer=keras.optimizers.RMSprop(lr=0.0001), + metrics=[keras.metrics.categorical_accuracy], + sample_weight_mode='temporal') + x = np.random.random((1, 3)) + y = np.random.random((1, 3, 3)) + model.train_on_batch(x, y) + + out = model.predict(x) + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + + new_model = keras.models.load_model(fname) + os.close(fd) + os.remove(fname) + + out2 = new_model.predict(x) + self.assertAllClose(out, out2, atol=1e-05) + + # test that new updates are the same with both models + x = np.random.random((1, 3)) + y = np.random.random((1, 3, 3)) + model.train_on_batch(x, y) + new_model.train_on_batch(x, y) + out = model.predict(x) + out2 = new_model.predict(x) + self.assertAllClose(out, out2, atol=1e-05) + + def test_sequential_model_saving_2(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + # test with custom optimizer, loss + + class CustomOp(keras.optimizers.RMSprop): + pass + + def custom_loss(y_true, y_pred): + return keras.losses.mse(y_true, y_pred) + + model = keras.models.Sequential() + model.add(keras.layers.Dense(2, input_shape=(3,))) + model.add(keras.layers.Dense(3)) + model.compile(loss=custom_loss, optimizer=CustomOp(), metrics=['acc']) + + x = np.random.random((1, 3)) + y = np.random.random((1, 3)) + model.train_on_batch(x, y) + + out = model.predict(x) + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + + model = keras.models.load_model( + fname, + custom_objects={'CustomOp': CustomOp, + 'custom_loss': custom_loss}) + os.close(fd) + os.remove(fname) + + out2 = model.predict(x) + self.assertAllClose(out, out2, atol=1e-05) + + def test_functional_model_saving(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + inputs = keras.layers.Input(shape=(3,)) + x = keras.layers.Dense(2)(inputs) + output = keras.layers.Dense(3)(x) + + model = keras.models.Model(inputs, output) + model.compile(loss=keras.losses.MSE, + optimizer=keras.optimizers.RMSprop(lr=0.0001), + metrics=[keras.metrics.categorical_accuracy]) + x = np.random.random((1, 3)) + y = np.random.random((1, 3)) + model.train_on_batch(x, y) + + out = model.predict(x) + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + + model = keras.models.load_model(fname) + os.close(fd) + os.remove(fname) + + out2 = model.predict(x) + self.assertAllClose(out, out2, atol=1e-05) + + def test_saving_without_compilation(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(2, input_shape=(3,))) + model.add(keras.layers.Dense(3)) + model.compile(loss='mse', optimizer='sgd', metrics=['acc']) + + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + model = keras.models.load_model(fname) + os.close(fd) + os.remove(fname) + + def test_saving_with_tf_optimizer(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(2, input_shape=(3,))) + model.add(keras.layers.Dense(3)) + model.compile(loss='mse', + optimizer=training_module.AdadeltaOptimizer(0.1), + metrics=['acc']) + + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + model = keras.models.load_model(fname) + os.close(fd) + os.remove(fname) + + def test_saving_right_after_compilation(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(2, input_shape=(3,))) + model.add(keras.layers.Dense(3)) + model.compile(loss='mse', optimizer='sgd', metrics=['acc']) + model.model._make_train_function() + + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + model = keras.models.load_model(fname) + os.close(fd) + os.remove(fname) + + def test_saving_lambda_numpy_array_arguments(self): + if h5py is None: + return # Skip test if models cannot be saved. + + mean = np.random.random((4, 2, 3)) + std = np.abs(np.random.random((4, 2, 3))) + 1e-5 + inputs = keras.layers.Input(shape=(4, 2, 3)) + output = keras.layers.Lambda(lambda image, mu, std: (image - mu) / std, + arguments={'mu': mean, 'std': std})(inputs) + model = keras.models.Model(inputs, output) + model.compile(loss='mse', optimizer='sgd', metrics=['acc']) + + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + + model = keras.models.load_model(fname) + os.close(fd) + os.remove(fname) + + self.assertAllClose(mean, model.layers[1].arguments['mu']) + self.assertAllClose(std, model.layers[1].arguments['std']) + + def test_saving_model_with_long_layer_names(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + # This layer name will make the `layers_name` HDF5 attribute blow + # out of proportion. Note that it fits into the internal HDF5 + # attribute memory limit on its own but because h5py converts + # the list of layer names into numpy array, which uses the same + # amout of memory for every item, it increases the memory + # requirements substantially. + x = keras.Input(shape=(2,), name='input_' + ('x' * (2**15))) + f = x + for i in range(4): + f = keras.layers.Dense(2, name='dense_%d' % (i,))(f) + model = keras.Model(inputs=[x], outputs=[f]) + model.compile(loss='mse', optimizer='adam', metrics=['acc']) + + x = np.random.random((1, 2)) + y = np.random.random((1, 2)) + model.train_on_batch(x, y) + out = model.predict(x) + + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + model = keras.models.load_model(fname) + + # Check that the HDF5 files contains chunked array + # of layer names. + with h5py.File(fname, 'r') as h5file: + num_names_arrays = len([attr for attr in h5file['model_weights'].attrs + if attr.startswith('layer_names')]) + # The chunking of layer names array should have happend. + self.assertGreater(num_names_arrays, 0) + out2 = model.predict(x) + self.assertAllClose(out, out2, atol=1e-05) + + # Cleanup + os.close(fd) + os.remove(fname) + + def test_saving_model_with_long_weights_names(self): + if h5py is None: + return # Skip test if models cannot be saved. + + with self.test_session(): + x = keras.Input(shape=(2,), name='nested_model_input') + f = x + for i in range(4): + f = keras.layers.Dense(2, name='nested_model_dense_%d' % (i,))(f) + # This layer name will make the `weights_name` + # HDF5 attribute blow out of proportion. + f = keras.layers.Dense(2, name='nested_model_output' + ('x' * (2**15)))(f) + nested_model = keras.Model(inputs=[x], outputs=[f], name='nested_model') + + x = keras.Input(shape=(2,), name='outer_model_input') + f = nested_model(x) + f = keras.layers.Dense(2, name='outer_model_output')(f) + + model = keras.Model(inputs=[x], outputs=[f]) + model.compile(loss='mse', optimizer='adam', metrics=['acc']) + + x = np.random.random((1, 2)) + y = np.random.random((1, 2)) + model.train_on_batch(x, y) + out = model.predict(x) + + fd, fname = tempfile.mkstemp('.h5') + keras.models.save_model(model, fname) + model = keras.models.load_model(fname) + + # Check that the HDF5 files contains chunked array + # of weight names. + with h5py.File(fname, 'r') as h5file: + num_weight_arrays = len( + [attr for attr in h5file['model_weights']['nested_model'].attrs + if attr.startswith('weight_names')]) + # The chunking of layer names array should have happend. + self.assertGreater(num_weight_arrays, 0) + out2 = model.predict(x) + self.assertAllClose(out, out2, atol=1e-05) + + # Cleanup + os.close(fd) + os.remove(fname) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/keras/_impl/keras/engine/sequential.py b/tensorflow/python/keras/_impl/keras/engine/sequential.py new file mode 100644 index 0000000000000000000000000000000000000000..db5e7754bc22ba360dbf635f1bd80334f58e8509 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/sequential.py @@ -0,0 +1,997 @@ +# 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. +# ============================================================================== +# pylint: disable=protected-access +"""Home of the `Sequential` model. +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import copy +import os + +from tensorflow.python.framework import ops +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import layers as layer_module +from tensorflow.python.keras._impl.keras.engine import base_layer +from tensorflow.python.keras._impl.keras.engine import network +from tensorflow.python.keras._impl.keras.engine import saving +from tensorflow.python.keras._impl.keras.engine.input_layer import Input +from tensorflow.python.keras._impl.keras.engine.input_layer import InputLayer +from tensorflow.python.keras._impl.keras.engine.training import Model +from tensorflow.python.keras._impl.keras.utils.io_utils import ask_to_proceed_with_overwrite +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export + +try: + import h5py # pylint: disable=g-import-not-at-top +except ImportError: + h5py = None + + +@tf_export('keras.models.Sequential', 'keras.Sequential') +class Sequential(Model): + """Linear stack of layers. + + Arguments: + layers: list of layers to add to the model. + + # Note + The first layer passed to a Sequential model + should have a defined input shape. What that + means is that it should have received an `input_shape` + or `batch_input_shape` argument, + or for some type of layers (recurrent, Dense...) + an `input_dim` argument. + + Example: + + ```python + model = Sequential() + # first layer must have a defined input shape + model.add(Dense(32, input_dim=500)) + # afterwards, Keras does automatic shape inference + model.add(Dense(32)) + + # also possible (equivalent to the above): + model = Sequential() + model.add(Dense(32, input_shape=(500,))) + model.add(Dense(32)) + + # also possible (equivalent to the above): + model = Sequential() + # here the batch dimension is None, + # which means any batch size will be accepted by the model. + model.add(Dense(32, batch_input_shape=(None, 500))) + model.add(Dense(32)) + ``` + """ + + def __init__(self, layers=None, name=None): + self._is_graph_network = True + self._is_compiled = False + self._layers = [] # Stack of layers. + self.model = None # Internal Model instance. + self.inputs = [] # List of input tensors + self.outputs = [] # List of length 1: the output tensor (unique). + self._trainable = True + self._initial_weights = None + self._input_layers = [] + + # Model attributes. + self._inbound_nodes = [] + self._outbound_nodes = [] + self.built = False + + # Set model name. + if not name: + prefix = 'sequential_' + name = prefix + str(K.get_uid(prefix)) + self._name = name + + # Used by Layer base class. + self._dtype = None + self._activity_regularizer = None + + # The following properties are not actually used by Keras; + # they exist for compatibility with TF's variable scoping mechanism. + self._updates = [] + self._losses = [] + self._scope = None + self._reuse = None + self._base_name = name + self._graph = ops.get_default_graph() + + # Add to the model any layers passed to the constructor. + if layers: + for layer in layers: + self.add(layer) + + def add(self, layer): + """Adds a layer instance on top of the layer stack. + + Arguments: + layer: layer instance. + + Raises: + TypeError: If `layer` is not a layer instance. + ValueError: In case the `layer` argument does not + know its input shape. + ValueError: In case the `layer` argument has + multiple output tensors, or is already connected + somewhere else (forbidden in `Sequential` models). + """ + if not isinstance(layer, (base_layer.Layer, base_layer.TFBaseLayer)): + raise TypeError('The added layer must be ' + 'an instance of class Layer. ' + 'Found: ' + str(layer)) + if not self.outputs: + # First layer in model: check that it is an input layer. + if not isinstance(layer, InputLayer): + # Create an input layer. + # First, we need to infer its expected input shape and dtype. + if isinstance(layer, (Model, Sequential)): + # We were passed a model as first layer. + # This requires a specific way to figure out the + # input shape and dtype. + if not layer.layers: + raise ValueError('Cannot add an empty model ' + 'to a `Sequential` model.') + # In case of nested models: recover the first layer + # of the deepest model to infer input shape and dtype. + first_layer = layer.layers[0] + while isinstance(first_layer, (Model, Sequential)): + first_layer = first_layer.layers[0] + batch_shape = first_layer._batch_input_shape + dtype = first_layer.dtype + else: + # We were passed a regular layer, and it should + # know about its input shape. Otherwise, that's an error. + if not hasattr(layer, '_batch_input_shape'): + raise ValueError('The first layer in a ' + 'Sequential model must ' + 'get an `input_shape` argument.') + batch_shape = layer._batch_input_shape + dtype = layer.dtype + # Instantiate the input layer. + x = Input( + batch_shape=batch_shape, dtype=dtype, name=layer.name + '_input') + # This will build the current layer + # and create the node connecting the current layer + # to the input layer we just created. + layer(x) + + if len(layer._inbound_nodes[-1].output_tensors) != 1: + raise ValueError('All layers in a Sequential model ' + 'should have a single output tensor. ' + 'For multi-output layers, ' + 'use the functional API.') + + self.outputs = [layer._inbound_nodes[-1].output_tensors[0]] + self.inputs = network.get_source_inputs(self.outputs[0]) + + # We create an input node, which we will keep updated + # as we add more layers + base_layer.Node( + outbound_layer=self, + inbound_layers=[], + node_indices=[], + tensor_indices=[], + input_tensors=self.inputs, + output_tensors=self.outputs) + else: + output_tensor = layer(self.outputs[0]) + if isinstance(output_tensor, list): + raise TypeError('All layers in a Sequential model ' + 'should have a single output tensor. ' + 'For multi-output layers, ' + 'use the functional API.') + self.outputs = [output_tensor] + # update self._inbound_nodes + self._inbound_nodes[0].output_tensors = self.outputs + self._inbound_nodes[0].output_shapes = [K.int_shape(self.outputs[0])] + + self._layers.append(layer) + self.built = False + + def pop(self): + """Removes the last layer in the model. + + Raises: + TypeError: if there are no layers in the model. + """ + if not self.layers: + raise TypeError('There are no layers in the model.') + + self.layers.pop() + if not self.layers: + self.outputs = [] + self._inbound_nodes = [] + self._outbound_nodes = [] + else: + self.layers[-1]._outbound_nodes = [] + self.outputs = [self.layers[-1].output] + # update self._inbound_nodes + self._inbound_nodes[0].output_tensors = self.outputs + self._inbound_nodes[0].output_shapes = [K.int_shape(self.outputs[0])] + self.built = False + + def get_layer(self, name=None, index=None): + """Retrieve a layer that is part of the model. + + Returns a layer based on either its name (unique) + or its index in the graph. Indices are based on + order of horizontal graph traversal (bottom-up). + + Arguments: + name: string, name of layer. + index: integer, index of layer. + + Returns: + A layer instance. + """ + if not self.built: + self.build() + return self.model.get_layer(name, index) + + def call(self, inputs, **kwargs): + if not self.built: + self.build() + return self.model.call(inputs, **kwargs) + + def build(self, input_shape=None): + if not self.inputs or not self.outputs: + raise TypeError('Sequential model cannot be built: model is empty.' + ' Add some layers first.') + # actually create the model + self.model = Model(self.inputs, self.outputs[0], name=self.name + '_model') + self.model.trainable = self.trainable + + # mirror model attributes + self.supports_masking = self.model.supports_masking + self._output_mask_cache = self.model._output_mask_cache + self._output_tensor_cache = self.model._output_tensor_cache + self._output_shape_cache = self.model._output_shape_cache + self._input_layers = self.model._input_layers + self._output_layers = self.model._output_layers + self._input_coordinates = self.model._input_coordinates + self._output_coordinates = self.model._output_coordinates + self._nodes_by_depth = self.model._nodes_by_depth + self._network_nodes = self.model._network_nodes + self.output_names = self.model.output_names + self.input_names = self.model.input_names + self._feed_input_names = self.model._feed_input_names + self._feed_inputs = self.model._feed_inputs + + # Make sure child model callbacks + # will call the parent Sequential model. + self.model.callback_model = self + + self.built = True + + @property + def uses_learning_phase(self): + if not self.built: + self.build() + return self.model.uses_learning_phase + + def _gather_list_attr(self, attr): + all_attrs = [] + for layer in self.layers: + all_attrs += getattr(layer, attr, []) + return all_attrs + + def _make_train_function(self): + self.model._make_train_function() + + def _make_test_function(self): + self.model._make_test_function() + + def _make_predict_function(self): + self.model._make_predict_function() + + @property + def trainable(self): + return self._trainable + + @trainable.setter + def trainable(self, value): + if self.model: + self.model.trainable = value + self._trainable = value + + @property + def trainable_weights(self): + if not self.trainable: + return [] + return self._gather_list_attr('trainable_weights') + + @property + def non_trainable_weights(self): + weights = self._gather_list_attr('non_trainable_weights') + if not self.trainable: + trainable_weights = self._gather_list_attr('trainable_weights') + return trainable_weights + weights + return weights + + @property + def regularizers(self): + if not self.built: + self.build() + return self.model.regularizers + + def get_weights(self): + """Retrieves the weights of the model. + + Returns: + A flat list of Numpy arrays + (one array per model weight). + """ + if not self.built: + self.build() + return self.model.get_weights() + + def set_weights(self, weights): + """Sets the weights of the model. + + Arguments: + weights: Should be a list + of Numpy arrays with shapes and types matching + the output of `model.get_weights()`. + """ + if not self.built: + self.build() + self.model.set_weights(weights) + + def load_weights(self, filepath, by_name=False): + if h5py is None: + raise ImportError('`load_weights` requires h5py.') + f = h5py.File(filepath, mode='r') + if 'layer_names' not in f.attrs and 'model_weights' in f: + f = f['model_weights'] + layers = self.layers + if by_name: + saving.load_weights_from_hdf5_group_by_name(f, layers) + else: + saving.load_weights_from_hdf5_group(f, layers) + if hasattr(f, 'close'): + f.close() + + def save_weights(self, filepath, overwrite=True): + if h5py is None: + raise ImportError('`save_weights` requires h5py.') + # If file exists and should not be overwritten: + if not overwrite and os.path.isfile(filepath): + proceed = ask_to_proceed_with_overwrite(filepath) + if not proceed: + return + layers = self.layers + f = h5py.File(filepath, 'w') + saving.save_weights_to_hdf5_group(f, layers) + f.flush() + f.close() + + def compile(self, + optimizer, + loss, + metrics=None, + sample_weight_mode=None, + weighted_metrics=None, + target_tensors=None, + **kwargs): + """Configures the model for training. + + Arguments: + optimizer: String (name of optimizer) or optimizer object. + See [optimizers](/optimizers). + loss: String (name of objective function) or objective function. + See [losses](/losses). + If the model has multiple outputs, you can use a different loss + on each output by passing a dictionary or a list of losses. + The loss value that will be minimized by the model + will then be the sum of all individual losses. + metrics: List of metrics to be evaluated by the model + during training and testing. + Typically you will use `metrics=['accuracy']`. + To specify different metrics for different outputs of a + multi-output model, you could also pass a dictionary, + such as `metrics={'output_a': 'accuracy'}`. + sample_weight_mode: If you need to do timestep-wise + sample weighting (2D weights), set this to `"temporal"`. + `None` defaults to sample-wise weights (1D). + If the model has multiple outputs, you can use a different + `sample_weight_mode` on each output by passing a + dictionary or a list of modes. + weighted_metrics: list of metrics to be evaluated and weighted + by `sample_weight` or `class_weight` during training and testing. + target_tensors: By default, Keras will create a placeholder for the + model's target, which will be fed with the target data during + training. If instead you would like to use your own + target tensor (in turn, Keras will not expect external + Numpy data for these targets at training time), you + can specify them via the `target_tensors` argument. + It should be a single tensor + (for a single-output `Sequential` model). + **kwargs: These arguments are passed into `tf.Session.run`. + + Example: + ```python + model = Sequential() + model.add(Dense(32, input_shape=(500,))) + model.add(Dense(10, activation='softmax')) + model.compile(optimizer='rmsprop', + loss='categorical_crossentropy', + metrics=['accuracy']) + ``` + """ + # create the underlying model + self.build() + # call compile method of Model class + self.model.compile( + optimizer, + loss, + metrics=metrics, + sample_weight_mode=sample_weight_mode, + weighted_metrics=weighted_metrics, + target_tensors=target_tensors, + **kwargs) + self.optimizer = self.model.optimizer + self.loss = self.model.loss + self.metrics = self.model.metrics + self.loss_weights = self.model.loss_weights + self.sample_weight_mode = self.model.sample_weight_mode + self.weighted_metrics = self.model.weighted_metrics + self.targets = self.model.targets + self.metrics_tensors = self.model.metrics_tensors + self.metrics_names = self.model.metrics_names + self.sample_weights = self.model.sample_weights + self.total_loss = self.model.total_loss + + def fit(self, + x=None, + y=None, + batch_size=None, + epochs=1, + verbose=1, + callbacks=None, + validation_split=0., + validation_data=None, + shuffle=True, + class_weight=None, + sample_weight=None, + initial_epoch=0, + steps_per_epoch=None, + validation_steps=None, + **kwargs): + """Trains the model for a fixed number of epochs. + + Arguments: + x: Numpy array of training data. + If the input layer in the model is named, you can also pass a + dictionary mapping the input name to a Numpy array. + `x` can be `None` (default) if feeding from + TensorFlow data tensors. + y: Numpy array of target (label) data. + If the output layer in the model is named, you can also pass a + dictionary mapping the output name to a Numpy array. + `y` can be `None` (default) if feeding from + TensorFlow data tensors. + batch_size: Integer or `None`. + Number of samples per gradient update. + If unspecified, it will default to 32. + epochs: Integer. Number of epochs to train the model. + An epoch is an iteration over the entire `x` and `y` + data provided. + Note that in conjunction with `initial_epoch`, + `epochs` is to be understood as "final epoch". + The model is not trained for a number of iterations + given by `epochs`, but merely until the epoch + of index `epochs` is reached. + verbose: 0, 1, or 2. Verbosity mode. + 0 = silent, 1 = progress bar, 2 = one line per epoch. + callbacks: List of `keras.callbacks.Callback` instances. + List of callbacks to apply during training. + See [callbacks](/callbacks). + validation_split: Float between 0 and 1: + Fraction of the training data to be used as validation data. + The model will set apart this fraction of the training data, + will not train on it, and will evaluate + the loss and any model metrics + on this data at the end of each epoch. + The validation data is selected from the last samples + in the `x` and `y` data provided, before shuffling. + validation_data: tuple `(x_val, y_val)` or tuple + `(x_val, y_val, val_sample_weights)` on which to evaluate + the loss and any model metrics at the end of each epoch. + The model will not be trained on this data. + This will override `validation_split`. + shuffle: Boolean (whether to shuffle the training data + before each epoch) or str (for 'batch'). + 'batch' is a special option for dealing with the + limitations of HDF5 data; it shuffles in batch-sized chunks. + Has no effect when `steps_per_epoch` is not `None`. + class_weight: Optional dictionary mapping class indices (integers) + to a weight (float) value, used for weighting the loss function + (during training only). + This can be useful to tell the model to + "pay more attention" to samples from + an under-represented class. + sample_weight: Optional Numpy array of weights for + the training samples, used for weighting the loss function + (during training only). You can either pass a flat (1D) + Numpy array with the same length as the input samples + (1:1 mapping between weights and samples), + or in the case of temporal data, + you can pass a 2D array with shape + `(samples, sequence_length)`, + to apply a different weight to every timestep of every sample. + In this case you should make sure to specify + `sample_weight_mode="temporal"` in `compile()`. + initial_epoch: Epoch at which to start training + (useful for resuming a previous training run). + steps_per_epoch: Total number of steps (batches of samples) + before declaring one epoch finished and starting the + next epoch. When training with input tensors such as + TensorFlow data tensors, the default `None` is equal to + the number of unique samples in your dataset divided by + the batch size, or 1 if that cannot be determined. + validation_steps: Only relevant if `steps_per_epoch` + is specified. Total number of steps (batches of samples) + to validate before stopping. + **kwargs: Used for backwards compatibility support. + + Returns: + A `History` object. Its `History.history` attribute is + a record of training loss values and metrics values + at successive epochs, as well as validation loss values + and validation metrics values (if applicable). + + Raises: + RuntimeError: If the model was never compiled. + ValueError: In case of mismatch between the provided input data + and what the model expects. + """ + if not self.built: + raise RuntimeError('The model needs to be compiled before being used.') + return self.model.fit( + x, + y, + batch_size=batch_size, + epochs=epochs, + verbose=verbose, + callbacks=callbacks, + validation_split=validation_split, + validation_data=validation_data, + shuffle=shuffle, + class_weight=class_weight, + sample_weight=sample_weight, + initial_epoch=initial_epoch, + steps_per_epoch=steps_per_epoch, + validation_steps=validation_steps) + + def evaluate(self, x, y, batch_size=32, verbose=1, sample_weight=None): + """Computes the loss on some input data, batch by batch. + + Arguments: + x: input data, as a Numpy array or list of Numpy arrays + (if the model has multiple inputs). + y: labels, as a Numpy array. + batch_size: integer. Number of samples per gradient update. + verbose: verbosity mode, 0 or 1. + sample_weight: sample weights, as a Numpy array. + + Returns: + Scalar test loss (if the model has no metrics) + or list of scalars (if the model computes other metrics). + The attribute `model.metrics_names` will give you + the display labels for the scalar outputs. + + Raises: + RuntimeError: if the model was never compiled. + """ + if not self.built: + raise RuntimeError('The model needs to be compiled before being used.') + return self.model.evaluate( + x, + y, + batch_size=batch_size, + verbose=verbose, + sample_weight=sample_weight) + + def predict(self, x, batch_size=32, verbose=0): + """Generates output predictions for the input samples. + + The input samples are processed batch by batch. + + Arguments: + x: the input data, as a Numpy array. + batch_size: integer. + verbose: verbosity mode, 0 or 1. + + Returns: + A Numpy array of predictions. + """ + if not self.built: + self.build() + return self.model.predict(x, batch_size=batch_size, verbose=verbose) + + def predict_on_batch(self, x): + """Returns predictions for a single batch of samples. + + Arguments: + x: input data, as a Numpy array or list of Numpy arrays + (if the model has multiple inputs). + + Returns: + A Numpy array of predictions. + """ + if not self.built: + self.build() + return self.model.predict_on_batch(x) + + def train_on_batch(self, x, y, class_weight=None, sample_weight=None): + """Single gradient update over one batch of samples. + + Arguments: + x: input data, as a Numpy array or list of Numpy arrays + (if the model has multiple inputs). + y: labels, as a Numpy array. + class_weight: dictionary mapping classes to a weight value, + used for scaling the loss function (during training only). + sample_weight: sample weights, as a Numpy array. + + Returns: + Scalar training loss (if the model has no metrics) + or list of scalars (if the model computes other metrics). + The attribute `model.metrics_names` will give you + the display labels for the scalar outputs. + + Raises: + RuntimeError: if the model was never compiled. + """ + if not self.built: + raise RuntimeError('The model needs to be compiled before being used.') + return self.model.train_on_batch( + x, y, sample_weight=sample_weight, class_weight=class_weight) + + def test_on_batch(self, x, y, sample_weight=None): + """Evaluates the model over a single batch of samples. + + Arguments: + x: input data, as a Numpy array or list of Numpy arrays + (if the model has multiple inputs). + y: labels, as a Numpy array. + sample_weight: sample weights, as a Numpy array. + + Returns: + Scalar test loss (if the model has no metrics) + or list of scalars (if the model computes other metrics). + The attribute `model.metrics_names` will give you + the display labels for the scalar outputs. + + Raises: + RuntimeError: if the model was never compiled. + """ + if not self.built: + raise RuntimeError('The model needs to be compiled before being used.') + return self.model.test_on_batch(x, y, sample_weight=sample_weight) + + def predict_proba(self, x, batch_size=32, verbose=0): + """Generates class probability predictions for the input samples. + + The input samples are processed batch by batch. + + Arguments: + x: input data, as a Numpy array or list of Numpy arrays + (if the model has multiple inputs). + batch_size: integer. + verbose: verbosity mode, 0 or 1. + + Returns: + A Numpy array of probability predictions. + """ + preds = self.predict(x, batch_size, verbose) + if preds.min() < 0. or preds.max() > 1.: + logging.warning('Network returning invalid probability values. ' + 'The last layer might not normalize predictions ' + 'into probabilities ' + '(like softmax or sigmoid would).') + return preds + + def predict_classes(self, x, batch_size=32, verbose=0): + """Generate class predictions for the input samples. + + The input samples are processed batch by batch. + + Arguments: + x: input data, as a Numpy array or list of Numpy arrays + (if the model has multiple inputs). + batch_size: integer. + verbose: verbosity mode, 0 or 1. + + Returns: + A numpy array of class predictions. + """ + proba = self.predict(x, batch_size=batch_size, verbose=verbose) + if proba.shape[-1] > 1: + return proba.argmax(axis=-1) + else: + return (proba > 0.5).astype('int32') + + def fit_generator(self, + generator, + steps_per_epoch=None, + epochs=1, + verbose=1, + callbacks=None, + validation_data=None, + validation_steps=None, + class_weight=None, + max_queue_size=10, + workers=1, + use_multiprocessing=False, + shuffle=True, + initial_epoch=0, + **kwargs): + """Fits the model on data generated batch-by-batch by a Python generator. + + The generator is run in parallel to the model, for efficiency. + For instance, this allows you to do real-time data augmentation + on images on CPU in parallel to training your model on GPU. + + Arguments: + generator: A generator. + The output of the generator must be either + - a tuple (inputs, targets) + - a tuple (inputs, targets, sample_weights). + All arrays should contain the same number of samples. + The generator is expected to loop over its data + indefinitely. An epoch finishes when `steps_per_epoch` + batches have been seen by the model. + steps_per_epoch: Total number of steps (batches of samples) + to yield from `generator` before declaring one epoch + finished and starting the next epoch. It should typically + be equal to the number of samples of your dataset + divided by the batch size. + Optional for `Sequence`: if unspecified, will use + the `len(generator)` as a number of steps. + epochs: Integer, total number of iterations on the data. + Note that in conjunction with initial_epoch, the parameter + epochs is to be understood as "final epoch". The model is + not trained for n steps given by epochs, but until the + epoch epochs is reached. + verbose: Verbosity mode, 0, 1, or 2. + callbacks: List of callbacks to be called during training. + validation_data: This can be either + - A generator for the validation data + - A tuple (inputs, targets) + - A tuple (inputs, targets, sample_weights). + validation_steps: Only relevant if `validation_data` + is a generator. + Number of steps to yield from validation generator + at the end of every epoch. It should typically + be equal to the number of samples of your + validation dataset divided by the batch size. + Optional for `Sequence`: if unspecified, will use + the `len(validation_data)` as a number of steps. + class_weight: Dictionary mapping class indices to a weight + for the class. + max_queue_size: Maximum size for the generator queue + workers: Maximum number of processes to spin up + use_multiprocessing: If True, use process based threading. + Note that because + this implementation relies on multiprocessing, + you should not pass + non picklable arguments to the generator + as they can't be passed + easily to children processes. + shuffle: Whether to shuffle the order of the batches at + the beginning of each epoch. Only used with instances + of `Sequence` (keras.utils.Sequence). + initial_epoch: Epoch at which to start training + (useful for resuming a previous training run) + **kwargs: support for legacy arguments. + + Returns: + A `History` object. + + Raises: + RuntimeError: if the model was never compiled. + ValueError: In case the generator yields + data in an invalid format. + + Example: + + ```python + def generate_arrays_from_file(path): + while 1: + f = open(path) + for line in f: + # create Numpy arrays of input data + # and labels, from each line in the file + x, y = process_line(line) + yield (x, y) + f.close() + + model.fit_generator(generate_arrays_from_file('/my_file.txt'), + steps_per_epoch=1000, epochs=10) + ``` + """ + # Legacy support + if 'max_q_size' in kwargs: + max_queue_size = kwargs.pop('max_q_size') + logging.warning('The argument `max_q_size` has been renamed ' + '`max_queue_size`. Update your method calls accordingly.') + if 'pickle_safe' in kwargs: + use_multiprocessing = kwargs.pop('pickle_safe') + logging.warning('The argument `pickle_safe` has been renamed ' + '`use_multiprocessing`. ' + 'Update your method calls accordingly.') + if kwargs: + raise ValueError('Unrecognized keyword arguments: ' + str(kwargs)) + + if not self.built: + raise RuntimeError('The model needs to be compiled before being used.') + return self.model.fit_generator( + generator, + steps_per_epoch, + epochs, + verbose=verbose, + callbacks=callbacks, + validation_data=validation_data, + validation_steps=validation_steps, + class_weight=class_weight, + max_queue_size=max_queue_size, + workers=workers, + use_multiprocessing=use_multiprocessing, + shuffle=shuffle, + initial_epoch=initial_epoch) + + def evaluate_generator(self, + generator, + steps=None, + max_queue_size=10, + workers=1, + use_multiprocessing=False, + **kwargs): + """Evaluates the model on a data generator. + + The generator should return the same kind of data + as accepted by `test_on_batch`. + + Arguments: + generator: Generator yielding tuples (inputs, targets) + or (inputs, targets, sample_weights) + steps: Total number of steps (batches of samples) + to yield from `generator` before stopping. + Optional for `Sequence`: if unspecified, will use + the `len(generator)` as a number of steps. + max_queue_size: maximum size for the generator queue + workers: maximum number of processes to spin up + use_multiprocessing: if True, use process based threading. + Note that because this implementation + relies on multiprocessing, you should not pass + non picklable arguments to the generator + as they can't be passed easily to children processes. + **kwargs: support for legacy arguments. + + Returns: + Scalar test loss (if the model has no metrics) + or list of scalars (if the model computes other metrics). + The attribute `model.metrics_names` will give you + the display labels for the scalar outputs. + + Raises: + RuntimeError: if the model was never compiled. + ValueError: In case the generator yields + data in an invalid format. + """ + # Legacy support + if 'max_q_size' in kwargs: + max_queue_size = kwargs.pop('max_q_size') + logging.warning('The argument `max_q_size` has been renamed ' + '`max_queue_size`. Update your method calls accordingly.') + if 'pickle_safe' in kwargs: + use_multiprocessing = kwargs.pop('pickle_safe') + logging.warning('The argument `pickle_safe` has been renamed ' + '`use_multiprocessing`. ' + 'Update your method calls accordingly.') + if kwargs: + raise ValueError('Unrecognized keyword arguments: ' + str(kwargs)) + + if not self.built: + raise RuntimeError('The model needs to be compiled before being used.') + return self.model.evaluate_generator( + generator, + steps, + max_queue_size=max_queue_size, + workers=workers, + use_multiprocessing=use_multiprocessing) + + def predict_generator(self, + generator, + steps=None, + max_queue_size=10, + workers=1, + use_multiprocessing=False, + verbose=0, + **kwargs): + """Generates predictions for the input samples from a data generator. + + The generator should return the same kind of data as accepted by + `predict_on_batch`. + + Arguments: + generator: generator yielding batches of input samples. + steps: Total number of steps (batches of samples) + to yield from `generator` before stopping. + Optional for `Sequence`: if unspecified, will use + the `len(generator)` as a number of steps. + max_queue_size: maximum size for the generator queue + workers: maximum number of processes to spin up + use_multiprocessing: if True, use process based threading. + Note that because this implementation + relies on multiprocessing, you should not pass + non picklable arguments to the generator + as they can't be passed easily to children processes. + verbose: verbosity mode, 0 or 1. + **kwargs: support for legacy arguments. + + Returns: + A Numpy array of predictions. + + Raises: + ValueError: In case the generator yields + data in an invalid format. + """ + # Legacy support + if 'max_q_size' in kwargs: + max_queue_size = kwargs.pop('max_q_size') + logging.warning('The argument `max_q_size` has been renamed ' + '`max_queue_size`. Update your method calls accordingly.') + if 'pickle_safe' in kwargs: + use_multiprocessing = kwargs.pop('pickle_safe') + logging.warning('The argument `pickle_safe` has been renamed ' + '`use_multiprocessing`. ' + 'Update your method calls accordingly.') + if kwargs: + raise ValueError('Unrecognized keyword arguments: ' + str(kwargs)) + + if not self.built: + self.build() + return self.model.predict_generator( + generator, + steps, + max_queue_size=max_queue_size, + workers=workers, + use_multiprocessing=use_multiprocessing, + verbose=verbose) + + def get_config(self): + config = [] + for layer in self.layers: + config.append({ + 'class_name': layer.__class__.__name__, + 'config': layer.get_config() + }) + return copy.deepcopy(config) + + @classmethod + def from_config(cls, config, custom_objects=None): + model = cls() + for conf in config: + layer = layer_module.deserialize(conf, custom_objects=custom_objects) + model.add(layer) + return model diff --git a/tensorflow/python/keras/_impl/keras/engine/sequential_test.py b/tensorflow/python/keras/_impl/keras/engine/sequential_test.py new file mode 100644 index 0000000000000000000000000000000000000000..166634bd8219b831ce212ba983a4ab695b00c3b7 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/sequential_test.py @@ -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. +# ============================================================================== +"""Tests specific to `Sequential` model.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.python.keras._impl import keras +from tensorflow.python.platform import test + + +class TestSequential(test.TestCase): + """Most Sequential model API tests are covered in `training_test.py`. + """ + + def test_basic_methods(self): + model = keras.models.Sequential() + model.add(keras.layers.Dense(1, input_dim=2)) + model.add(keras.layers.Dropout(0.3, name='dp')) + model.add(keras.layers.Dense(2, kernel_regularizer='l2', + kernel_constraint='max_norm')) + model.build() + self.assertEqual(model.state_updates, model.model.state_updates) + self.assertEqual(model.get_layer(name='dp').name, 'dp') + + def test_sequential_pop(self): + num_hidden = 5 + input_dim = 3 + batch_size = 5 + num_classes = 2 + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) + model.add(keras.layers.Dense(num_classes)) + model.compile(loss='mse', optimizer='sgd') + x = np.random.random((batch_size, input_dim)) + y = np.random.random((batch_size, num_classes)) + model.fit(x, y, epochs=1) + model.pop() + self.assertEqual(len(model.layers), 1) + self.assertEqual(model.output_shape, (None, num_hidden)) + model.compile(loss='mse', optimizer='sgd') + y = np.random.random((batch_size, num_hidden)) + model.fit(x, y, epochs=1) + + # Test popping single-layer model + model = keras.models.Sequential() + model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) + model.pop() + self.assertEqual(len(model.layers), 0) + self.assertEqual(len(model.outputs), 0) + + # Invalid use case + model = keras.models.Sequential() + with self.assertRaises(TypeError): + model.pop() + + def test_invalid_use_cases(self): + with self.test_session(): + # Added objects must be layer instances + with self.assertRaises(TypeError): + model = keras.models.Sequential() + model.add(None) + + # Added layers must have an inputs shape + with self.assertRaises(ValueError): + model = keras.models.Sequential() + model.add(keras.layers.Dense(1)) + + # Added layers cannot have multiple outputs + class MyLayer(keras.layers.Layer): + + def call(self, inputs): + return [3 * inputs, 2 * inputs] + + def compute_output_shape(self, input_shape): + return [input_shape, input_shape] + + with self.assertRaises(ValueError): + model = keras.models.Sequential() + model.add(MyLayer(input_shape=(3,))) + with self.assertRaises(TypeError): + model = keras.models.Sequential() + model.add(keras.layers.Dense(1, input_dim=1)) + model.add(MyLayer()) + + # Building empty model + model = keras.models.Sequential() + with self.assertRaises(TypeError): + model.build() + + def test_nested_sequential_trainability(self): + input_dim = 20 + num_units = 10 + num_classes = 2 + + inner_model = keras.models.Sequential() + inner_model.add(keras.layers.Dense(num_units, input_shape=(input_dim,))) + + model = keras.models.Sequential() + model.add(inner_model) + model.add(keras.layers.Dense(num_classes)) + + self.assertEqual(len(model.trainable_weights), 4) + inner_model.trainable = False + self.assertEqual(len(model.trainable_weights), 2) + inner_model.trainable = True + self.assertEqual(len(model.trainable_weights), 4) + + def test_sequential_update_disabling(self): + val_a = np.random.random((10, 4)) + val_out = np.random.random((10, 4)) + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.BatchNormalization(input_shape=(4,))) + + model.trainable = False + assert not model.updates + + model.compile('sgd', 'mse') + assert not model.updates + assert not model.model.updates + + x1 = model.predict(val_a) + model.train_on_batch(val_a, val_out) + x2 = model.predict(val_a) + self.assertAllClose(x1, x2, atol=1e-7) + + model.trainable = True + model.compile('sgd', 'mse') + assert model.updates + assert model.model.updates + + model.train_on_batch(val_a, val_out) + x2 = model.predict(val_a) + assert np.abs(np.sum(x1 - x2)) > 1e-5 diff --git a/tensorflow/python/keras/_impl/keras/engine/topology.py b/tensorflow/python/keras/_impl/keras/engine/topology.py deleted file mode 100644 index 64aa868f3822c4dfcfbe8ae1764d617a00ffff4d..0000000000000000000000000000000000000000 --- a/tensorflow/python/keras/_impl/keras/engine/topology.py +++ /dev/null @@ -1,1601 +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. -# ============================================================================== -# pylint: disable=protected-access -"""Base layer code and base model (Network) code. -""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import copy -import json -import os - -import numpy as np -from six.moves import zip # pylint: disable=redefined-builtin - -from tensorflow.python.eager import context -from tensorflow.python.framework import tensor_shape -from tensorflow.python.keras._impl.keras import backend as K -from tensorflow.python.keras._impl.keras import constraints -from tensorflow.python.keras._impl.keras import initializers -from tensorflow.python.keras._impl.keras import regularizers -from tensorflow.python.keras._impl.keras.utils import conv_utils -from tensorflow.python.keras._impl.keras.utils.io_utils import ask_to_proceed_with_overwrite -from tensorflow.python.keras._impl.keras.utils.layer_utils import print_summary as print_layer_summary -from tensorflow.python.layers import base as tf_base_layers -from tensorflow.python.layers import network as tf_network -from tensorflow.python.layers import utils as tf_layers_util -from tensorflow.python.platform import tf_logging as logging - - -# pylint: disable=g-import-not-at-top -try: - import h5py -except ImportError: - h5py = None - -try: - import yaml -except ImportError: - yaml = None -# pylint: enable=g-import-not-at-top - -# pylint: disable=invalid-name -InputSpec = tf_base_layers.InputSpec -Node = tf_base_layers.Node -TFBaseLayer = tf_base_layers.Layer -# pylint: enable=invalid-name - - -class Layer(tf_base_layers.Layer): - """Abstract base layer class. - - # Properties - name: String, must be unique within a model. - input_spec: List of InputSpec class instances - each entry describes one required input: - - ndim - - dtype - A layer with `n` input tensors must have - an `input_spec` of length `n`. - trainable: Boolean, whether the layer weights - will be updated during training. - uses_learning_phase: Whether any operation - of the layer uses `K.in_training_phase()` - or `K.in_test_phase()`. - input_shape: Shape tuple. Provided for convenience, - but note that there may be cases in which this - attribute is ill-defined (e.g. a shared layer - with multiple input shapes), in which case - requesting `input_shape` will raise an Exception. - Prefer using `layer.get_input_shape_for(input_shape)`, - or `layer.get_input_shape_at(node_index)`. - output_shape: Shape tuple. See above. - inbound_nodes: List of nodes. - outbound_nodes: List of nodes. - input, output: Input/output tensor(s). Note that if the layer is used - more than once (shared layer), this is ill-defined - and will raise an exception. In such cases, use - `layer.get_input_at(node_index)`. - input_mask, output_mask: Same as above, for masks. - trainable_weights: List of variables. - non_trainable_weights: List of variables. - weights: The concatenation of the lists trainable_weights and - non_trainable_weights (in this order). - - # Methods - call(x, mask=None): Where the layer's logic lives. - __call__(x, mask=None): Wrapper around the layer logic (`call`). - If x is a Keras tensor: - - Connect current layer with last layer from tensor: - `self._add_inbound_node(last_layer)` - - Add layer to tensor history - If layer is not built: - - Build from inputs shape - get_weights() - set_weights(weights) - get_config() - count_params() - compute_output_shape(input_shape) - compute_mask(x, mask) - get_input_at(node_index) - get_output_at(node_index) - get_input_shape_at(node_index) - get_output_shape_at(node_index) - get_input_mask_at(node_index) - get_output_mask_at(node_index) - - # Class Methods - from_config(config) - - # Internal methods: - build(input_shape) - _add_inbound_node(layer, index=0) - """ - - def __init__(self, **kwargs): - # These properties should be set by the user via keyword arguments. - # note that 'dtype', 'input_shape' and 'batch_input_shape' - # are only applicable to input layers: do not pass these keywords - # to non-input layers. - allowed_kwargs = { - 'activity_regularizer', - 'input_shape', - 'batch_input_shape', - 'batch_size', - 'dtype', - 'name', - 'trainable', - 'weights', - } - # Validate optional keyword arguments. - for kwarg in kwargs: - if kwarg not in allowed_kwargs: - raise TypeError('Keyword argument not understood:', kwarg) - - # Get layer name. - name = kwargs.get('name') - - # Get `trainable` status. - trainable = kwargs.get('trainable', True) - - # Get `dtype`. - dtype = kwargs.get('dtype') - if dtype is None: - dtype = K.floatx() - - # Call super, which will set all properties common to Keras layers - # and core TF layers. - super(Layer, self).__init__( - name=name, dtype=dtype, trainable=trainable, - activity_regularizer=kwargs.get('activity_regularizer')) - - # Add properties that are Keras-only for now. - self.supports_masking = False - - # Manage input shape information if passed. - if 'input_shape' in kwargs or 'batch_input_shape' in kwargs: - # In this case we will later create an input layer - # to insert before the current layer - if 'batch_input_shape' in kwargs: - batch_input_shape = tuple(kwargs['batch_input_shape']) - elif 'input_shape' in kwargs: - if 'batch_size' in kwargs: - batch_size = kwargs['batch_size'] - else: - batch_size = None - batch_input_shape = (batch_size,) + tuple(kwargs['input_shape']) - self._batch_input_shape = batch_input_shape - - # Manage initial weight values if passed. - if 'weights' in kwargs: - self._initial_weights = kwargs['weights'] - else: - self._initial_weights = None - - def add_weight(self, - name, - shape, - dtype=None, - initializer=None, - regularizer=None, - trainable=True, - constraint=None): - """Adds a weight variable to the layer. - - Arguments: - name: String, the name for the weight variable. - shape: The shape tuple of the weight. - dtype: The dtype of the weight. - initializer: An Initializer instance (callable). - regularizer: An optional Regularizer instance. - trainable: A boolean, whether the weight should - be trained via backprop or not (assuming - that the layer itself is also trainable). - constraint: An optional Constraint instance. - - Returns: - The created weight variable. - """ - if dtype is None: - dtype = K.floatx() - weight = self.add_variable(name, shape, - dtype=dtype, - initializer=initializers.get(initializer), - regularizer=regularizers.get(regularizer), - constraint=constraints.get(constraint), - trainable=trainable) - return weight - - def call(self, inputs, **kwargs): # pylint: disable=unused-argument - """This is where the layer's logic lives. - - Arguments: - inputs: Input tensor, or list/tuple of input tensors. - **kwargs: Additional keyword arguments. - - Returns: - A tensor or list/tuple of tensors. - """ - return inputs - - def __call__(self, inputs, **kwargs): - """Wrapper around self.call(), for handling internal references. - - If a Keras tensor is passed: - - We call self._add_inbound_node(). - - If necessary, we `build` the layer to match - the shape of the input(s). - - We update the _keras_history of the output tensor(s) - with the current layer. - This is done as part of _add_inbound_node(). - - Arguments: - inputs: Can be a tensor or list/tuple of tensors. - **kwargs: Additional keyword arguments to be passed to `call()`. - - Returns: - Output of the layer's `call` method. - - Raises: - ValueError: in case the layer is missing shape information - for its `build` call. - """ - # Actually call the layer (optionally building it). - output = super(Layer, self).__call__(inputs, **kwargs) - if context.in_eager_mode(): - return output - - # Update learning phase info. - output_tensors = _to_list(output) - uses_lp = any( - [getattr(x, '_uses_learning_phase', False) for x in _to_list(inputs)]) - uses_lp = getattr(self, 'uses_learning_phase', False) or uses_lp - for i in range(len(output_tensors)): - output_tensors[i]._uses_learning_phase = getattr( - output_tensors[i], '_uses_learning_phase', False) or uses_lp - - # Optionally load weight values that were specified at layer instantiation. - if hasattr(self, '_initial_weights') and self._initial_weights is not None: - self.set_weights(self._initial_weights) - del self._initial_weights - return output - - def compute_output_shape(self, input_shape): - """Computes the output shape of the layer. - - Assumes that the layer will be built - to match that input shape provided. - - Arguments: - input_shape: Shape tuple (tuple of integers) - or list of shape tuples (one per output tensor of the layer). - Shape tuples can include None for free dimensions, - instead of an integer. - - Returns: - An input shape tuple. - """ - logging.warning( - 'All custom layers should implement the ' - '`compute_output_shape` method. This layer (' + self.name + ') ' - 'is relying on the base `Layer.compute_output_shape` implementation, ' - 'which will start raising a `NotImplementedError` ' - 'as of July 1st, 2018.') - return input_shape - - def compute_mask(self, inputs, mask=None): # pylint: disable=unused-argument - """Computes an output mask tensor. - - Arguments: - inputs: Tensor or list of tensors. - mask: Tensor or list of tensors. - - Returns: - None or a tensor (or list of tensors, - one per output tensor of the layer). - """ - if not self.supports_masking: - if mask is not None: - if isinstance(mask, list): - if any(m is not None for m in mask): - raise TypeError('Layer ' + self.name + ' does not support masking, ' - 'but was passed an input_mask: ' + str(mask)) - else: - raise TypeError('Layer ' + self.name + ' does not support masking, ' - 'but was passed an input_mask: ' + str(mask)) - # masking not explicitly supported: return None as mask - return None - # if masking is explicitly supported, by default - # carry over the input mask - return mask - - def get_input_mask_at(self, node_index): - """Retrieves the input mask tensor(s) of a layer at a given node. - - Arguments: - node_index: Integer, index of the node - from which to retrieve the attribute. - E.g. `node_index=0` will correspond to the - first time the layer was called. - - Returns: - A mask tensor - (or list of tensors if the layer has multiple inputs). - """ - inputs = self.get_input_at(node_index) - if isinstance(inputs, list): - return [getattr(x, '_keras_mask', None) for x in inputs] - else: - return getattr(inputs, '_keras_mask', None) - - def get_output_mask_at(self, node_index): - """Retrieves the output mask tensor(s) of a layer at a given node. - - Arguments: - node_index: Integer, index of the node - from which to retrieve the attribute. - E.g. `node_index=0` will correspond to the - first time the layer was called. - - Returns: - A mask tensor - (or list of tensors if the layer has multiple outputs). - """ - output = self.get_output_at(node_index) - if isinstance(output, list): - return [getattr(x, '_keras_mask', None) for x in output] - else: - return getattr(output, '_keras_mask', None) - - @property - def input_mask(self): - """Retrieves the input mask tensor(s) of a layer. - - Only applicable if the layer has exactly one inbound node, - i.e. if it is connected to one incoming layer. - - Returns: - Input mask tensor (potentially None) or list of input - mask tensors. - - Raises: - AttributeError: if the layer is connected to - more than one incoming layers. - """ - inputs = self.input - if isinstance(inputs, list): - return [getattr(x, '_keras_mask', None) for x in inputs] - else: - return getattr(inputs, '_keras_mask', None) - - @property - def output_mask(self): - """Retrieves the output mask tensor(s) of a layer. - - Only applicable if the layer has exactly one inbound node, - i.e. if it is connected to one incoming layer. - - Returns: - Output mask tensor (potentially None) or list of output - mask tensors. - - Raises: - AttributeError: if the layer is connected to - more than one incoming layers. - """ - output = self.output - if isinstance(output, list): - return [getattr(x, '_keras_mask', None) for x in output] - else: - return getattr(output, '_keras_mask', None) - - def set_weights(self, weights): - """Sets the weights of the layer, from Numpy arrays. - - Arguments: - weights: a list of Numpy arrays. The number - of arrays and their shape must match - number of the dimensions of the weights - of the layer (i.e. it should match the - output of `get_weights`). - - Raises: - ValueError: If the provided weights list does not match the - layer's specifications. - """ - params = self.weights - if len(params) != len(weights): - raise ValueError('You called `set_weights(weights)` on layer "' + - self.name + '" with a weight list of length ' + - str(len(weights)) + ', but the layer was expecting ' + - str(len(params)) + ' weights. Provided weights: ' + - str(weights)[:50] + '...') - if not params: - return - weight_value_tuples = [] - param_values = K.batch_get_value(params) - for pv, p, w in zip(param_values, params, weights): - if pv.shape != w.shape: - raise ValueError('Layer weight shape ' + str(pv.shape) + - ' not compatible with ' - 'provided weight shape ' + str(w.shape)) - weight_value_tuples.append((p, w)) - K.batch_set_value(weight_value_tuples) - - def get_weights(self): - """Returns the current weights of the layer. - - Returns: - Weights values as a list of numpy arrays. - """ - params = self.weights - return K.batch_get_value(params) - - def get_config(self): - """Returns the config of the layer. - - A layer config is a Python dictionary (serializable) - containing the configuration of a layer. - The same layer can be reinstantiated later - (without its trained weights) from this configuration. - - The config of a layer does not include connectivity - information, nor the layer class name. These are handled - by `Network` (one layer of abstraction above). - - Returns: - Python dictionary. - """ - config = {'name': self.name, 'trainable': self.trainable} - if hasattr(self, '_batch_input_shape'): - config['batch_input_shape'] = self._batch_input_shape - if hasattr(self, 'dtype'): - config['dtype'] = self.dtype - return config - - @classmethod - def from_config(cls, config): - """Creates a layer from its config. - - This method is the reverse of `get_config`, - capable of instantiating the same layer from the config - dictionary. It does not handle layer connectivity - (handled by Network), nor weights (handled by `set_weights`). - - Arguments: - config: A Python dictionary, typically the - output of get_config. - - Returns: - A layer instance. - """ - return cls(**config) - - @tf_base_layers.Layer.activity_regularizer.setter - def activity_regularizer(self, activity_regularizer): - self._activity_regularizer = activity_regularizer - - -class InputLayer(tf_network.InputLayer, Layer): - """Layer to be used as an entry point into a graph. - - It can either wrap an existing tensor (pass an `input_tensor` argument) - or create its a placeholder tensor (pass argument `input_shape`. - - Arguments: - input_shape: Shape tuple, not including the batch axis. - batch_size: Optional input batch size (integer or None). - dtype: Datatype of the input. - input_tensor: Optional tensor to use as layer input - instead of creating a placeholder. - sparse: Boolean, whether the placeholder created - is meant to be sparse. - name: Name of the layer (string). - """ - - def __init__(self, - input_shape=None, - batch_size=None, - dtype=None, - input_tensor=None, - sparse=False, - name=None, - **kwargs): - if 'batch_input_shape' in kwargs: - batch_input_shape = kwargs.pop('batch_input_shape') - if input_shape and batch_input_shape: - raise ValueError('Only provide the input_shape OR ' - 'batch_input_shape argument to ' - 'InputLayer, not both at the same time.') - batch_size = batch_input_shape[0] - input_shape = batch_input_shape[1:] - if kwargs: - raise ValueError('Unrecognized keyword arguments:', kwargs.keys()) - - if not name: - prefix = 'input' - name = prefix + '_' + str(K.get_uid(prefix)) - - if not dtype: - if input_tensor is None: - dtype = K.floatx() - else: - dtype = K.dtype(input_tensor) - super(InputLayer, self).__init__(input_shape=input_shape, - batch_size=batch_size, - dtype=dtype, - input_tensor=input_tensor, - sparse=sparse, - name=name) - - def get_config(self): - config = { - 'batch_input_shape': self._batch_input_shape, - 'dtype': self.dtype, - 'sparse': self.sparse, - 'name': self.name - } - return config - - -def Input( # pylint: disable=invalid-name - shape=None, - batch_size=None, - name=None, - dtype=None, - sparse=False, - tensor=None, - **kwargs): - """`Input()` is used to instantiate a Keras tensor. - - A Keras tensor is a tensor object from the underlying backend - (Theano or TensorFlow), which we augment with certain - attributes that allow us to build a Keras model - just by knowing the inputs and outputs of the model. - - For instance, if a, b and c are Keras tensors, - it becomes possible to do: - `model = Model(input=[a, b], output=c)` - - The added Keras attribute is: - `_keras_history`: Last layer applied to the tensor. - the entire layer graph is retrievable from that layer, - recursively. - - Arguments: - shape: A shape tuple (integers), not including the batch size. - For instance, `shape=(32,)` indicates that the expected input - will be batches of 32-dimensional vectors. - batch_size: optional static batch size (integer). - name: An optional name string for the layer. - Should be unique in a model (do not reuse the same name twice). - It will be autogenerated if it isn't provided. - dtype: The data type expected by the input, as a string - (`float32`, `float64`, `int32`...) - sparse: A boolean specifying whether the placeholder - to be created is sparse. - tensor: Optional existing tensor to wrap into the `Input` layer. - If set, the layer will not create a placeholder tensor. - **kwargs: deprecated arguments support. - - Returns: - A tensor. - - Example: - - ```python - # this is a logistic regression in Keras - x = Input(shape=(32,)) - y = Dense(16, activation='softmax')(x) - model = Model(x, y) - ``` - - Raises: - ValueError: in case of invalid arguments. - """ - if 'batch_shape' in kwargs: - batch_shape = kwargs.pop('batch_shape') - if shape and batch_shape: - raise ValueError('Only provide the shape OR ' - 'batch_shape argument to ' - 'Input, not both at the same time.') - batch_size = batch_shape[0] - shape = batch_shape[1:] - if kwargs: - raise ValueError('Unrecognized keyword arguments:', kwargs.keys()) - - if dtype is None: - dtype = K.floatx() - if not shape and tensor is None: - raise ValueError('Please provide to Input either a `shape`' - ' or a `tensor` argument. Note that ' - '`shape` does not include the batch ' - 'dimension.') - input_layer = InputLayer( - input_shape=shape, - batch_size=batch_size, - name=name, - dtype=dtype, - sparse=sparse, - input_tensor=tensor) - # Return tensor including `_keras_history`. - # Note that in this case train_output and test_output are the same pointer. - outputs = input_layer._inbound_nodes[0].output_tensors - if len(outputs) == 1: - return outputs[0] - else: - return outputs - - -class Network(tf_network.GraphNetwork, Layer): - """A Network is a directed acyclic graph of layers. - - It is the topological form of a "model". A Model - is simply a Network with added training routines. - - # Properties - name - inputs - outputs - input_layers - output_layers - input_spec (list of class instances) - each entry describes one required input: - - ndim - - dtype - trainable (boolean) - input_shape - output_shape - inbound_nodes: list of nodes - outbound_nodes: list of nodes - trainable_weights (list of variables) - non_trainable_weights (list of variables) - - # Methods - summary - get_layer - get_weights - set_weights - get_config - compute_output_shape - - # Class Methods - from_config - """ - - def __init__(self, inputs, outputs, name=None): - super(Network, self).__init__(inputs, outputs, name=name) - - self.supports_masking = False - # Fill in the output mask cache. - masks = [] - for x in self.inputs: - mask = x._keras_mask if hasattr(x, '_keras_mask') else None - masks.append(mask) - mask_cache_key = (tf_layers_util.object_list_uid(self.inputs) + '_' + - tf_layers_util.object_list_uid(masks)) - masks = [] - for x in self.outputs: - mask = x._keras_mask if hasattr(x, '_keras_mask') else None - masks.append(mask) - if len(masks) == 1: - mask = masks[0] - else: - mask = masks - self._output_mask_cache[mask_cache_key] = mask - - # Build self.input_names and self.output_names. - self.input_names = [] - self.output_names = [] - self._feed_input_names = [] - self._feed_inputs = [] - self._feed_input_shapes = [] - for i, layer in enumerate(self._input_layers): - self.input_names.append(layer.name) - if layer.is_placeholder: - self._feed_input_names.append(layer.name) - self._feed_inputs.append(layer.input) - self._feed_input_shapes.append(K.int_shape(self.inputs[i])) - for layer in self._output_layers: - self.output_names.append(layer.name) - - self._internal_input_shapes = [K.int_shape(x) for x in self.inputs] - self._internal_output_shapes = [K.int_shape(x) for x in self.outputs] - - @property - def uses_learning_phase(self): - return any( - [getattr(x, '_uses_learning_phase', False) for x in self.outputs]) - - @property - def stateful(self): - return any([(hasattr(layer, 'stateful') and layer.stateful) - for layer in self.layers]) - - def reset_states(self): - for layer in self.layers: - if hasattr(layer, 'reset_states') and getattr(layer, 'stateful', False): - layer.reset_states() - - @property - def state_updates(self): - """Returns the `updates` from all layers that are stateful. - - This is useful for separating training updates and - state updates, e.g. when we need to update a layer's internal state - during prediction. - - Returns: - A list of update ops. - """ - state_updates = [] - for layer in self.layers: - if getattr(layer, 'stateful', False): - if hasattr(layer, 'updates'): - state_updates += layer.updates - return state_updates - - def get_weights(self): - """Retrieves the weights of the model. - - Returns: - A flat list of Numpy arrays. - """ - weights = [] - for layer in self.layers: - weights += layer.weights - return K.batch_get_value(weights) - - def set_weights(self, weights): - """Sets the weights of the model. - - Arguments: - weights: A list of Numpy arrays with shapes and types matching - the output of `model.get_weights()`. - """ - tuples = [] - for layer in self.layers: - num_param = len(layer.weights) - layer_weights = weights[:num_param] - for sw, w in zip(layer.weights, layer_weights): - tuples.append((sw, w)) - weights = weights[num_param:] - K.batch_set_value(tuples) - - def compute_mask(self, inputs, mask): - inputs = _to_list(inputs) - if mask is None: - masks = [None for _ in range(len(inputs))] - else: - masks = _to_list(mask) - cache_key = ','.join([str(id(x)) for x in inputs]) - cache_key += '_' + ','.join([str(id(x)) for x in masks]) - if cache_key in self._output_mask_cache: - return self._output_mask_cache[cache_key] - else: - _, output_masks = self._run_internal_graph(inputs, masks) - return output_masks - - def get_config(self): - config = { - 'name': self.name, - } - node_conversion_map = {} - for layer in self.layers: - if issubclass(layer.__class__, Network): - # Networks start with a pre-existing node - # linking their input to output. - kept_nodes = 1 - else: - kept_nodes = 0 - for original_node_index, node in enumerate(layer._inbound_nodes): - node_key = tf_network._make_node_key(layer.name, - original_node_index) - if node_key in self._network_nodes: - node_conversion_map[node_key] = kept_nodes - kept_nodes += 1 - layer_configs = [] - for layer in self.layers: # From the earliest layers on. - layer_class_name = layer.__class__.__name__ - layer_config = layer.get_config() - filtered_inbound_nodes = [] - for original_node_index, node in enumerate(layer._inbound_nodes): - node_key = tf_network._make_node_key(layer.name, - original_node_index) - if node_key in self._network_nodes: - # The node is relevant to the model: - # add to filtered_inbound_nodes. - if node.arguments: - try: - json.dumps(node.arguments) - kwargs = node.arguments - except TypeError: - logging.warning( - 'Layer ' + layer.name + - ' was passed non-serializable keyword arguments: ' + - str(node.arguments) + '. They will not be included ' - 'in the serialized model (and thus will be missing ' - 'at deserialization time).') - kwargs = {} - else: - kwargs = {} - if node.inbound_layers: - node_data = [] - for i in range(len(node.inbound_layers)): - inbound_layer = node.inbound_layers[i] - node_index = node.node_indices[i] - tensor_index = node.tensor_indices[i] - node_key = tf_network._make_node_key(inbound_layer.name, - node_index) - new_node_index = node_conversion_map.get(node_key, 0) - node_data.append( - [inbound_layer.name, new_node_index, tensor_index, kwargs]) - filtered_inbound_nodes.append(node_data) - layer_configs.append({ - 'name': layer.name, - 'class_name': layer_class_name, - 'config': layer_config, - 'inbound_nodes': filtered_inbound_nodes, - }) - config['layers'] = layer_configs - - # Gather info about inputs and outputs. - model_inputs = [] - for i in range(len(self._input_layers)): - layer, node_index, tensor_index = self._input_coordinates[i] - node_key = tf_network._make_node_key(layer.name, - node_index) - if node_key not in self._network_nodes: - continue - new_node_index = node_conversion_map[node_key] - model_inputs.append([layer.name, new_node_index, tensor_index]) - config['input_layers'] = model_inputs - model_outputs = [] - for i in range(len(self._output_layers)): - layer, node_index, tensor_index = self._output_coordinates[i] - node_key = tf_network._make_node_key(layer.name, - node_index) - if node_key not in self._network_nodes: - continue - new_node_index = node_conversion_map[node_key] - model_outputs.append([layer.name, new_node_index, tensor_index]) - config['output_layers'] = model_outputs - return copy.deepcopy(config) - - @classmethod - def from_config(cls, config, custom_objects=None): - """Instantiates a Model from its config (output of `get_config()`). - - Arguments: - config: Model config dictionary. - custom_objects: Optional dictionary mapping names - (strings) to custom classes or functions to be - considered during deserialization. - - Returns: - A model instance. - - Raises: - ValueError: In case of improperly formatted config dict. - """ - # Layer instances created during - # the graph reconstruction process - created_layers = {} - - # Dictionary mapping layer instances to - # node data that specifies a layer call. - # It acts as a queue that maintains any unprocessed - # layer call until it becomes possible to process it - # (i.e. until the input tensors to the call all exist). - unprocessed_nodes = {} - - def add_unprocessed_node(layer, node_data): - if layer not in unprocessed_nodes: - unprocessed_nodes[layer] = [node_data] - else: - unprocessed_nodes[layer].append(node_data) - - def process_node(layer, node_data): - """Deserialize a node. - - Arguments: - layer: layer instance. - node_data: node config dict. - - Raises: - ValueError: In case of improperly formatted `node_data` dict. - """ - input_tensors = [] - for input_data in node_data: - inbound_layer_name = input_data[0] - inbound_node_index = input_data[1] - inbound_tensor_index = input_data[2] - if len(input_data) == 3: - kwargs = {} - elif len(input_data) == 4: - kwargs = input_data[3] - else: - raise ValueError('Improperly formatted model config.') - if inbound_layer_name not in created_layers: - add_unprocessed_node(layer, node_data) - return - inbound_layer = created_layers[inbound_layer_name] - if len(inbound_layer._inbound_nodes) <= inbound_node_index: - add_unprocessed_node(layer, node_data) - return - inbound_node = inbound_layer._inbound_nodes[inbound_node_index] - input_tensors.append(inbound_node.output_tensors[inbound_tensor_index]) - # Call layer on its inputs, thus creating the node - # and building the layer if needed. - if input_tensors: - if len(input_tensors) == 1: - layer(input_tensors[0], **kwargs) - else: - layer(input_tensors, **kwargs) - - def process_layer(layer_data): - """Deserialize a layer, then call it on appropriate inputs. - - Arguments: - layer_data: layer config dict. - - Raises: - ValueError: In case of improperly formatted `layer_data` dict. - """ - layer_name = layer_data['name'] - - # Instantiate layer. - from tensorflow.python.keras._impl.keras.layers import deserialize as deserialize_layer # pylint: disable=g-import-not-at-top - - layer = deserialize_layer(layer_data, custom_objects=custom_objects) - created_layers[layer_name] = layer - - # Gather layer inputs. - inbound_nodes_data = layer_data['inbound_nodes'] - for node_data in inbound_nodes_data: - # We don't process nodes (i.e. make layer calls) - # on the fly because the inbound node may not yet exist, - # in case of layer shared at different topological depths - # (e.g. a model such as A(B(A(B(x))))) - add_unprocessed_node(layer, node_data) - - # First, we create all layers and enqueue nodes to be processed - for layer_data in config['layers']: - process_layer(layer_data) - # Then we process nodes in order of layer depth. - # Nodes that cannot yet be processed (if the inbound node - # does not yet exist) are re-enqueued, and the process - # is repeated until all nodes are processed. - while unprocessed_nodes: - for layer_data in config['layers']: - layer = created_layers[layer_data['name']] - if layer in unprocessed_nodes: - for node_data in unprocessed_nodes.pop(layer): - process_node(layer, node_data) - - name = config.get('name') - input_tensors = [] - output_tensors = [] - for layer_data in config['input_layers']: - layer_name, node_index, tensor_index = layer_data - assert layer_name in created_layers - layer = created_layers[layer_name] - layer_output_tensors = layer._inbound_nodes[node_index].output_tensors - input_tensors.append(layer_output_tensors[tensor_index]) - for layer_data in config['output_layers']: - layer_name, node_index, tensor_index = layer_data - assert layer_name in created_layers - layer = created_layers[layer_name] - layer_output_tensors = layer._inbound_nodes[node_index].output_tensors - output_tensors.append(layer_output_tensors[tensor_index]) - return cls(inputs=input_tensors, outputs=output_tensors, name=name) - - def save(self, filepath, overwrite=True, include_optimizer=True): - """Save the model to a single HDF5 file. - - The savefile includes: - - The model architecture, allowing to re-instantiate the model. - - The model weights. - - The state of the optimizer, allowing to resume training - exactly where you left off. - - This allows you to save the entirety of the state of a model - in a single file. - - Saved models can be reinstantiated via `keras.models.load_model`. - The model returned by `load_model` - is a compiled model ready to be used (unless the saved model - was never compiled in the first place). - - Arguments: - filepath: String, path to the file to save the weights to. - overwrite: Whether to silently overwrite any existing file at the - target location, or provide the user with a manual prompt. - include_optimizer: If True, save optimizer's state together. - - Example: - - ```python - from keras.models import load_model - - model.save('my_model.h5') # creates a HDF5 file 'my_model.h5' - del model # deletes the existing model - - # returns a compiled model - # identical to the previous one - model = load_model('my_model.h5') - ``` - """ - from tensorflow.python.keras._impl.keras.models import save_model # pylint: disable=g-import-not-at-top - save_model(self, filepath, overwrite, include_optimizer) - - def save_weights(self, filepath, overwrite=True): - """Dumps all layer weights to a HDF5 file. - - The weight file has: - - `layer_names` (attribute), a list of strings - (ordered names of model layers). - - For every layer, a `group` named `layer.name` - - For every such layer group, a group attribute `weight_names`, - a list of strings - (ordered names of weights tensor of the layer). - - For every weight in the layer, a dataset - storing the weight value, named after the weight tensor. - - Arguments: - filepath: String, path to the file to save the weights to. - overwrite: Whether to silently overwrite any existing file at the - target location, or provide the user with a manual prompt. - - Raises: - ImportError: If h5py is not available. - """ - if h5py is None: - raise ImportError('`save_weights` requires h5py.') - # If file exists and should not be overwritten: - if not overwrite and os.path.isfile(filepath): - proceed = ask_to_proceed_with_overwrite(filepath) - if not proceed: - return - f = h5py.File(filepath, 'w') - save_weights_to_hdf5_group(f, self.layers) - f.flush() - f.close() - - def load_weights(self, filepath, by_name=False): - """Loads all layer weights from a HDF5 save file. - - If `by_name` is False (default) weights are loaded - based on the network's topology, meaning the architecture - should be the same as when the weights were saved. - Note that layers that don't have weights are not taken - into account in the topological ordering, so adding or - removing layers is fine as long as they don't have weights. - - If `by_name` is True, weights are loaded into layers - only if they share the same name. This is useful - for fine-tuning or transfer-learning models where - some of the layers have changed. - - Arguments: - filepath: String, path to the weights file to load. - by_name: Boolean, whether to load weights by name - or by topological order. - - Raises: - ImportError: If h5py is not available. - """ - if h5py is None: - raise ImportError('`load_weights` requires h5py.') - f = h5py.File(filepath, mode='r') - if 'layer_names' not in f.attrs and 'model_weights' in f: - f = f['model_weights'] - if by_name: - load_weights_from_hdf5_group_by_name(f, self.layers) - else: - load_weights_from_hdf5_group(f, self.layers) - - if hasattr(f, 'close'): - f.close() - - def _updated_config(self): - """Util hared between different serialization methods. - - Returns: - Model config with Keras version information added. - """ - from tensorflow.python.keras._impl.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top - - config = self.get_config() - model_config = { - 'class_name': self.__class__.__name__, - 'config': config, - 'keras_version': keras_version, - 'backend': K.backend() - } - return model_config - - def to_json(self, **kwargs): - """Returns a JSON string containing the network configuration. - - To load a network from a JSON save file, use - `keras.models.model_from_json(json_string, custom_objects={})`. - - Arguments: - **kwargs: Additional keyword arguments - to be passed to `json.dumps()`. - - Returns: - A JSON string. - """ - - def get_json_type(obj): - # If obj is any numpy type - if type(obj).__module__ == np.__name__: - return obj.item() - - # If obj is a python 'type' - if type(obj).__name__ == type.__name__: - return obj.__name__ - - raise TypeError('Not JSON Serializable:', obj) - - model_config = self._updated_config() - return json.dumps(model_config, default=get_json_type, **kwargs) - - def to_yaml(self, **kwargs): - """Returns a yaml string containing the network configuration. - - To load a network from a yaml save file, use - `keras.models.model_from_yaml(yaml_string, custom_objects={})`. - - `custom_objects` should be a dictionary mapping - the names of custom losses / layers / etc to the corresponding - functions / classes. - - Arguments: - **kwargs: Additional keyword arguments - to be passed to `yaml.dump()`. - - Returns: - A YAML string. - - Raises: - ImportError: if yaml module is not found. - """ - if yaml is None: - raise ImportError('Requires yaml module installed.') - return yaml.dump(self._updated_config(), **kwargs) - - def summary(self, line_length=None, positions=None, print_fn=None): - """Prints a string summary of the network. - - Arguments: - line_length: Total length of printed lines - (e.g. set this to adapt the display to different - terminal window sizes). - positions: Relative or absolute positions of log elements - in each line. If not provided, - defaults to `[.33, .55, .67, 1.]`. - print_fn: Print function to use. Defaults to `print`. - It will be called on each line of the summary. - You can set it to a custom function - in order to capture the string summary. - """ - print_layer_summary(self, - line_length=line_length, - positions=positions, - print_fn=print_fn) - - -def get_source_inputs(tensor, layer=None, node_index=None): - """Returns the list of input tensors necessary to compute `tensor`. - - Output will always be a list of tensors - (potentially with 1 element). - - Arguments: - tensor: The tensor to start from. - layer: Origin layer of the tensor. Will be - determined via tensor._keras_history if not provided. - node_index: Origin node index of the tensor. - - Returns: - List of input tensors. - """ - if not hasattr(tensor, '_keras_history'): - return tensor - - if layer is None or node_index: - layer, node_index, _ = tensor._keras_history - if not layer._inbound_nodes: - return [tensor] - else: - node = layer._inbound_nodes[node_index] - if not node.inbound_layers: - # Reached an Input layer, stop recursion. - return node.input_tensors - else: - source_tensors = [] - for i in range(len(node.inbound_layers)): - x = node.input_tensors[i] - layer = node.inbound_layers[i] - node_index = node.node_indices[i] - previous_sources = get_source_inputs(x, layer, node_index) - # Avoid input redundancy. - for x in previous_sources: - if x not in source_tensors: - source_tensors.append(x) - return source_tensors - - -def _to_list(x): - """Normalizes a list/tensor into a list. - - If a tensor is passed, we return - a list of size 1 containing the tensor. - - Arguments: - x: target object to be normalized. - - Returns: - A list. - """ - if isinstance(x, list): - return x - return [x] - - -def save_weights_to_hdf5_group(f, layers): - from tensorflow.python.keras._impl.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top - - f.attrs['layer_names'] = [layer.name.encode('utf8') for layer in layers] - f.attrs['backend'] = K.backend().encode('utf8') - f.attrs['keras_version'] = str(keras_version).encode('utf8') - - for layer in layers: - g = f.create_group(layer.name) - symbolic_weights = layer.weights - weight_values = K.batch_get_value(symbolic_weights) - weight_names = [] - for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)): - if hasattr(w, 'name') and w.name: - name = str(w.name) - else: - name = 'param_' + str(i) - weight_names.append(name.encode('utf8')) - g.attrs['weight_names'] = weight_names - for name, val in zip(weight_names, weight_values): - param_dset = g.create_dataset(name, val.shape, dtype=val.dtype) - if not val.shape: - # scalar - param_dset[()] = val - else: - param_dset[:] = val - - -def preprocess_weights_for_loading(layer, - weights, - original_keras_version=None, - original_backend=None): - """Converts layers weights from Keras 1 format to Keras 2. - - Arguments: - layer: Layer instance. - weights: List of weights values (Numpy arrays). - original_keras_version: Keras version for the weights, as a string. - original_backend: Keras backend the weights were trained with, - as a string. - - Returns: - A list of weights values (Numpy arrays). - """ - if layer.__class__.__name__ == 'Bidirectional': - num_weights_per_layer = len(weights) // 2 - forward_weights = preprocess_weights_for_loading( - layer.forward_layer, weights[:num_weights_per_layer], - original_keras_version, original_backend) - backward_weights = preprocess_weights_for_loading( - layer.backward_layer, weights[num_weights_per_layer:], - original_keras_version, original_backend) - weights = forward_weights + backward_weights - - if original_keras_version == '1': - if layer.__class__.__name__ == 'TimeDistributed': - weights = preprocess_weights_for_loading( - layer.layer, weights, original_keras_version, original_backend) - - if layer.__class__.__name__ == 'Conv1D': - shape = weights[0].shape - # Handle Keras 1.1 format - if shape[:2] != (layer.kernel_size[0], 1) or shape[3] != layer.filters: - # Legacy shape: - # (filters, input_dim, filter_length, 1) - assert shape[0] == layer.filters and shape[2:] == (layer.kernel_size[0], - 1) - weights[0] = np.transpose(weights[0], (2, 3, 1, 0)) - weights[0] = weights[0][:, 0, :, :] - - if layer.__class__.__name__ == 'Conv2D': - if layer.data_format == 'channels_first': - # old: (filters, stack_size, kernel_rows, kernel_cols) - # new: (kernel_rows, kernel_cols, stack_size, filters) - weights[0] = np.transpose(weights[0], (2, 3, 1, 0)) - - if layer.__class__.__name__ == 'Conv2DTranspose': - if layer.data_format == 'channels_last': - # old: (kernel_rows, kernel_cols, stack_size, filters) - # new: (kernel_rows, kernel_cols, filters, stack_size) - weights[0] = np.transpose(weights[0], (0, 1, 3, 2)) - if layer.data_format == 'channels_first': - # old: (filters, stack_size, kernel_rows, kernel_cols) - # new: (kernel_rows, kernel_cols, filters, stack_size) - weights[0] = np.transpose(weights[0], (2, 3, 0, 1)) - - if layer.__class__.__name__ == 'Conv3D': - if layer.data_format == 'channels_first': - # old: (filters, stack_size, ...) - # new: (..., stack_size, filters) - weights[0] = np.transpose(weights[0], (2, 3, 4, 1, 0)) - - if layer.__class__.__name__ == 'GRU': - if len(weights) == 9: - kernel = np.concatenate([weights[0], weights[3], weights[6]], axis=-1) - recurrent_kernel = np.concatenate( - [weights[1], weights[4], weights[7]], axis=-1) - bias = np.concatenate([weights[2], weights[5], weights[8]], axis=-1) - weights = [kernel, recurrent_kernel, bias] - - if layer.__class__.__name__ == 'LSTM': - if len(weights) == 12: - # old: i, c, f, o - # new: i, f, c, o - kernel = np.concatenate( - [weights[0], weights[6], weights[3], weights[9]], axis=-1) - recurrent_kernel = np.concatenate( - [weights[1], weights[7], weights[4], weights[10]], axis=-1) - bias = np.concatenate( - [weights[2], weights[8], weights[5], weights[11]], axis=-1) - weights = [kernel, recurrent_kernel, bias] - - if layer.__class__.__name__ == 'ConvLSTM2D': - if len(weights) == 12: - kernel = np.concatenate( - [weights[0], weights[6], weights[3], weights[9]], axis=-1) - recurrent_kernel = np.concatenate( - [weights[1], weights[7], weights[4], weights[10]], axis=-1) - bias = np.concatenate( - [weights[2], weights[8], weights[5], weights[11]], axis=-1) - if layer.data_format == 'channels_first': - # old: (filters, stack_size, kernel_rows, kernel_cols) - # new: (kernel_rows, kernel_cols, stack_size, filters) - kernel = np.transpose(kernel, (2, 3, 1, 0)) - recurrent_kernel = np.transpose(recurrent_kernel, (2, 3, 1, 0)) - weights = [kernel, recurrent_kernel, bias] - - if layer.__class__.__name__ in ['Model', 'Sequential']: - new_weights = [] - # trainable weights - for sublayer in layer.layers: - num_weights = len(sublayer.trainable_weights) - if num_weights > 0: - new_weights.extend( - preprocess_weights_for_loading( - layer=sublayer, - weights=weights[:num_weights], - original_keras_version=original_keras_version, - original_backend=original_backend)) - weights = weights[num_weights:] - - # non-trainable weights - for sublayer in layer.layers: - num_weights = len([ - l for l in sublayer.weights if l not in sublayer.trainable_weights - ]) - if num_weights > 0: - new_weights.extend( - preprocess_weights_for_loading( - layer=sublayer, - weights=weights[:num_weights], - original_keras_version=original_keras_version, - original_backend=original_backend)) - weights = weights[num_weights:] - weights = new_weights - - conv_layers = ['Conv1D', 'Conv2D', 'Conv3D', 'Conv2DTranspose', 'ConvLSTM2D'] - if layer.__class__.__name__ in conv_layers: - if original_backend == 'theano': - weights[0] = conv_utils.convert_kernel(weights[0]) - if layer.__class__.__name__ == 'ConvLSTM2D': - weights[1] = conv_utils.convert_kernel(weights[1]) - if K.int_shape(layer.weights[0]) != weights[0].shape: - weights[0] = np.transpose(weights[0], (3, 2, 0, 1)) - if layer.__class__.__name__ == 'ConvLSTM2D': - weights[1] = np.transpose(weights[1], (3, 2, 0, 1)) - - # Convert the weights of CuDNNLSTM so that they could be loaded into LSTM - if layer.__class__.__name__ == 'LSTM' and len(weights) == 3: - # Determine if loading a CuDNNLSTM layer from the number of bias weights: - # CuDNNLSTM has (units * 8) weights; while LSTM has (units * 4) - # if there's no bias weight in the file, skip this conversion - units = weights[1].shape[0] - bias = weights[2] - if len(bias) == units * 8: - # reshape the kernels - kernels = np.split(weights[0], 4, axis=1) - kernels = [ - kernel.reshape(-1).reshape(kernel.shape, order='F') - for kernel in kernels - ] - weights[0] = np.concatenate(kernels, axis=1) - - # transpose the recurrent kernels - recurrent_kernels = np.split(weights[1], 4, axis=1) - recurrent_kernels = [kernel.T for kernel in recurrent_kernels] - weights[1] = np.concatenate(recurrent_kernels, axis=1) - - # split the bias into half and merge - weights[2] = bias[:units * 4] + bias[units * 4:] - - return weights - - -def load_weights_from_hdf5_group(f, layers): - """Implements topological (order-based) weight loading. - - Arguments: - f: A pointer to a HDF5 group. - layers: a list of target layers. - - Raises: - ValueError: in case of mismatch between provided layers - and weights file. - """ - if 'keras_version' in f.attrs: - original_keras_version = f.attrs['keras_version'].decode('utf8') - else: - original_keras_version = '1' - if 'backend' in f.attrs: - original_backend = f.attrs['backend'].decode('utf8') - else: - original_backend = None - - filtered_layers = [] - for layer in layers: - weights = layer.weights - if weights: - filtered_layers.append(layer) - - layer_names = [n.decode('utf8') for n in f.attrs['layer_names']] - filtered_layer_names = [] - for name in layer_names: - g = f[name] - weight_names = [n.decode('utf8') for n in g.attrs['weight_names']] - if weight_names: - filtered_layer_names.append(name) - layer_names = filtered_layer_names - if len(layer_names) != len(filtered_layers): - raise ValueError('You are trying to load a weight file ' - 'containing ' + str(len(layer_names)) + - ' layers into a model with ' + str(len(filtered_layers)) + - ' layers.') - - # We batch weight value assignments in a single backend call - # which provides a speedup in TensorFlow. - weight_value_tuples = [] - for k, name in enumerate(layer_names): - g = f[name] - weight_names = [n.decode('utf8') for n in g.attrs['weight_names']] - weight_values = [g[weight_name] for weight_name in weight_names] - layer = filtered_layers[k] - symbolic_weights = layer.weights - weight_values = preprocess_weights_for_loading( - layer, weight_values, original_keras_version, original_backend) - if len(weight_values) != len(symbolic_weights): - raise ValueError('Layer #' + str(k) + ' (named "' + layer.name + - '" in the current model) was found to ' - 'correspond to layer ' + name + ' in the save file. ' - 'However the new layer ' + layer.name + ' expects ' + - str(len(symbolic_weights)) + - ' weights, but the saved weights have ' + - str(len(weight_values)) + ' elements.') - weight_value_tuples += zip(symbolic_weights, weight_values) - K.batch_set_value(weight_value_tuples) - - -def load_weights_from_hdf5_group_by_name(f, layers): - """Implements name-based weight loading. - - (instead of topological weight loading). - - Layers that have no matching name are skipped. - - Arguments: - f: A pointer to a HDF5 group. - layers: a list of target layers. - - Raises: - ValueError: in case of mismatch between provided layers - and weights file. - """ - if 'keras_version' in f.attrs: - original_keras_version = f.attrs['keras_version'].decode('utf8') - else: - original_keras_version = '1' - if 'backend' in f.attrs: - original_backend = f.attrs['backend'].decode('utf8') - else: - original_backend = None - - # New file format. - layer_names = [n.decode('utf8') for n in f.attrs['layer_names']] - - # Reverse index of layer name to list of layers with name. - index = {} - for layer in layers: - if layer.name: - index.setdefault(layer.name, []).append(layer) - - # We batch weight value assignments in a single backend call - # which provides a speedup in TensorFlow. - weight_value_tuples = [] - for k, name in enumerate(layer_names): - g = f[name] - weight_names = [n.decode('utf8') for n in g.attrs['weight_names']] - weight_values = [g[weight_name] for weight_name in weight_names] - - for layer in index.get(name, []): - symbolic_weights = layer.weights - weight_values = preprocess_weights_for_loading( - layer, weight_values, original_keras_version, original_backend) - if len(weight_values) != len(symbolic_weights): - raise ValueError('Layer #' + str(k) + ' (named "' + layer.name + - '") expects ' + str(len(symbolic_weights)) + - ' weight(s), but the saved weights' + ' have ' + - str(len(weight_values)) + ' element(s).') - # Set values. - for i in range(len(weight_values)): - weight_value_tuples.append((symbolic_weights[i], weight_values[i])) - K.batch_set_value(weight_value_tuples) - - -def shape_type_conversion(fn): - """Decorator that handles tuple/TensorShape conversion. - - Used in `compute_output_shape` and `build`. - - Arguments: - fn: function to wrap. - - Returns: - Wrapped function. - """ - - def wrapper(instance, input_shape): - if input_shape is not None: - if isinstance(input_shape, list): - input_shape = [ - tuple(tensor_shape.TensorShape(x).as_list()) for x in input_shape] - else: - input_shape = tuple(tensor_shape.TensorShape(input_shape).as_list()) - output_shape = fn(instance, input_shape) - if output_shape is not None: - if isinstance(output_shape, list): - return [tensor_shape.TensorShape(x) for x in output_shape] - return tensor_shape.TensorShape(output_shape) - - return wrapper diff --git a/tensorflow/python/keras/_impl/keras/engine/topology_test.py b/tensorflow/python/keras/_impl/keras/engine/topology_test.py index 479ee877fd2471a67b5b5b81e8fbf338ce755a7b..b50277c8fff917d77694903c989fd02ea98b1711 100644 --- a/tensorflow/python/keras/_impl/keras/engine/topology_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/topology_test.py @@ -18,14 +18,17 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os -import shutil - import numpy as np +from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import test_util from tensorflow.python.keras._impl import keras +from tensorflow.python.layers import base as tf_base_layers from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import state_ops from tensorflow.python.platform import test try: @@ -33,31 +36,255 @@ try: except ImportError: yaml = None -try: - import h5py # pylint:disable=g-import-not-at-top -except ImportError: - h5py = None - class TopologyConstructionTest(test.TestCase): - def test_get_updates_for(self): - a = keras.layers.Input(shape=(2,)) - dense_layer = keras.layers.Dense(1) - dense_layer.add_update(0, inputs=a) - dense_layer.add_update(1, inputs=None) + def test_get_updates(self): + + class MyLayer(keras.layers.Layer): + + def build(self, input_shape): + self.a = self.add_variable('a', + (1, 1), + 'float32', + trainable=False) + self.b = self.add_variable('b', + (1, 1), + 'float32', + trainable=False) + self.add_update(state_ops.assign_add(self.a, [[1.]])) + self.built = True + + def call(self, inputs): + self.add_update(state_ops.assign_add(self.a, inputs), + inputs=True) + return inputs + 1 + + x1 = keras.Input(shape=(1,)) + layer = MyLayer() + _ = layer.apply(x1) + + self.assertEqual(len(layer.updates), 2) + self.assertEqual(len(layer.get_updates_for(x1)), 1) + self.assertEqual(len(layer.get_updates_for(None)), 1) + + x2 = keras.Input(shape=(1,)) + y2 = layer.apply(x2) + + self.assertEqual(len(layer.updates), 3) + self.assertEqual(len(layer.get_updates_for(x1)), 1) + self.assertEqual(len(layer.get_updates_for(x2)), 1) + self.assertEqual(len(layer.get_updates_for(None)), 1) + + network = keras.engine.Network(x2, y2) + self.assertEqual(len(network.updates), 2) + self.assertEqual(len(network.get_updates_for(x1)), 0) + self.assertEqual(len(network.get_updates_for(x2)), 1) + self.assertEqual(len(network.get_updates_for(None)), 1) + + x3 = keras.Input(shape=(1,)) + _ = layer.apply(x3) + self.assertEqual(len(network.updates), 2) + + x4 = keras.Input(shape=(1,)) + _ = network(x4) + self.assertEqual(len(network.updates), 3) + self.assertEqual(len(network.get_updates_for(x2)), 1) + self.assertEqual(len(network.get_updates_for(x4)), 1) + self.assertEqual(len(network.get_updates_for(None)), 1) + + network.add_update(state_ops.assign_add(layer.a, [[1]])) + self.assertEqual(len(network.updates), 4) + self.assertEqual(len(network.get_updates_for(None)), 2) + + network.add_update(state_ops.assign_add(layer.a, x4), inputs=True) + self.assertEqual(len(network.updates), 5) + self.assertEqual(len(network.get_updates_for(x4)), 2) + + def test_get_losses(self): + + class MyLayer(keras.layers.Layer): + + def build(self, input_shape): + self.a = self.add_variable('a', + (1, 1), + 'float32', + trainable=False) + self.b = self.add_variable('b', + (1, 1), + 'float32', + trainable=False) + self.add_loss(math_ops.reduce_sum(self.a)) + self.built = True + + def call(self, inputs): + self.add_loss(math_ops.reduce_sum(inputs), + inputs=True) + return inputs + 1 + + x1 = keras.Input(shape=(1,)) + layer = MyLayer() + _ = layer.apply(x1) + + self.assertEqual(len(layer.losses), 2) + self.assertEqual(len(layer.get_losses_for(x1)), 1) + self.assertEqual(len(layer.get_losses_for(None)), 1) + + x2 = keras.Input(shape=(1,)) + y2 = layer.apply(x2) + + self.assertEqual(len(layer.losses), 3) + self.assertEqual(len(layer.get_losses_for(x1)), 1) + self.assertEqual(len(layer.get_losses_for(x2)), 1) + self.assertEqual(len(layer.get_losses_for(None)), 1) + + network = keras.engine.Network(x2, y2) + self.assertEqual(len(network.losses), 2) + self.assertEqual(len(network.get_losses_for(x1)), 0) + self.assertEqual(len(network.get_losses_for(x2)), 1) + self.assertEqual(len(network.get_losses_for(None)), 1) + + x3 = keras.Input(shape=(1,)) + _ = layer.apply(x3) + self.assertEqual(len(network.losses), 2) + + x4 = keras.Input(shape=(1,)) + _ = network(x4) + self.assertEqual(len(network.losses), 3) + self.assertEqual(len(network.get_losses_for(x2)), 1) + self.assertEqual(len(network.get_losses_for(x4)), 1) + self.assertEqual(len(network.get_losses_for(None)), 1) + + network.add_loss(math_ops.reduce_sum(layer.a)) + self.assertEqual(len(network.losses), 4) + self.assertEqual(len(network.get_losses_for(None)), 2) + + network.add_loss(math_ops.reduce_sum(x4), inputs=True) + self.assertEqual(len(network.losses), 5) + self.assertEqual(len(network.get_losses_for(x4)), 2) + + def testTopologicalAttributes(self): + # test layer attributes / methods related to cross-layer connectivity. + a = keras.Input(shape=(32,), name='input_a') + b = keras.Input(shape=(32,), name='input_b') + + # test input, output, input_shape, output_shape + test_layer = keras.layers.Dense(16, name='test_layer') + a_test = test_layer(a) + self.assertEqual(test_layer.input, a) + self.assertEqual(test_layer.output, a_test) + self.assertEqual(test_layer.input_shape, (None, 32)) + self.assertEqual(test_layer.output_shape, (None, 16)) - self.assertListEqual(dense_layer.get_updates_for(a), [0]) - self.assertListEqual(dense_layer.get_updates_for(None), [1]) + # test `get_*_at` methods + dense = keras.layers.Dense(16, name='dense_1') + a_2 = dense(a) + b_2 = dense(b) - def test_get_losses_for(self): - a = keras.layers.Input(shape=(2,)) - dense_layer = keras.layers.Dense(1) - dense_layer.add_loss(0, inputs=a) - dense_layer.add_loss(1, inputs=None) + self.assertEqual(dense.get_input_at(0), a) + self.assertEqual(dense.get_input_at(1), b) + self.assertEqual(dense.get_output_at(0), a_2) + self.assertEqual(dense.get_output_at(1), b_2) + self.assertEqual(dense.get_input_shape_at(0), (None, 32)) + self.assertEqual(dense.get_input_shape_at(1), (None, 32)) + self.assertEqual(dense.get_output_shape_at(0), (None, 16)) + self.assertEqual(dense.get_output_shape_at(1), (None, 16)) + + # Test invalid value for attribute retrieval. + with self.assertRaises(ValueError): + dense.get_input_at(2) + with self.assertRaises(AttributeError): + new_dense = keras.layers.Dense(16) + _ = new_dense.input + with self.assertRaises(AttributeError): + new_dense = keras.layers.Dense(16) + _ = new_dense.output + with self.assertRaises(AttributeError): + new_dense = keras.layers.Dense(16) + _ = new_dense.output_shape + with self.assertRaises(AttributeError): + new_dense = keras.layers.Dense(16) + _ = new_dense.input_shape + with self.assertRaises(AttributeError): + new_dense = keras.layers.Dense(16) + a = keras.Input(shape=(3, 32)) + a = keras.Input(shape=(5, 32)) + a_2 = dense(a) + b_2 = dense(b) + _ = new_dense.input_shape + with self.assertRaises(AttributeError): + new_dense = keras.layers.Dense(16) + a = keras.Input(shape=(3, 32)) + a = keras.Input(shape=(5, 32)) + a_2 = dense(a) + b_2 = dense(b) + _ = new_dense.output_shape + + def testTopologicalAttributesMultiOutputLayer(self): + + class PowersLayer(keras.layers.Layer): - self.assertListEqual(dense_layer.get_losses_for(a), [0]) - self.assertListEqual(dense_layer.get_losses_for(None), [1]) + def call(self, inputs): + return [inputs**2, inputs**3] + + x = keras.Input(shape=(32,)) + test_layer = PowersLayer() + p1, p2 = test_layer(x) # pylint: disable=not-callable + + self.assertEqual(test_layer.input, x) + self.assertEqual(test_layer.output, [p1, p2]) + self.assertEqual(test_layer.input_shape, (None, 32)) + self.assertEqual(test_layer.output_shape, [(None, 32), (None, 32)]) + + def testTopologicalAttributesMultiInputLayer(self): + + class AddLayer(keras.layers.Layer): + + def call(self, inputs): + assert len(inputs) == 2 + return inputs[0] + inputs[1] + + a = keras.Input(shape=(32,)) + b = keras.Input(shape=(32,)) + test_layer = AddLayer() + y = test_layer([a, b]) # pylint: disable=not-callable + + self.assertEqual(test_layer.input, [a, b]) + self.assertEqual(test_layer.output, y) + self.assertEqual(test_layer.input_shape, [(None, 32), (None, 32)]) + self.assertEqual(test_layer.output_shape, (None, 32)) + + def testBasicNetwork(self): + # minimum viable network + x = keras.Input(shape=(32,)) + dense = keras.layers.Dense(2) + y = dense(x) + network = keras.engine.Network(x, y, name='dense_network') + + # test basic attributes + self.assertEqual(network.name, 'dense_network') + self.assertEqual(len(network.layers), 2) # InputLayer + Dense + self.assertEqual(network.layers[1], dense) + self.assertEqual(network.weights, dense.weights) + self.assertEqual(network.trainable_weights, dense.trainable_weights) + self.assertEqual(network.non_trainable_weights, dense.non_trainable_weights) + + # test callability on Input + x_2 = keras.Input(shape=(32,)) + y_2 = network(x_2) + self.assertEqual(y_2.get_shape().as_list(), [None, 2]) + + # test callability on regular tensor + x_2 = array_ops.placeholder(dtype='float32', shape=(None, 32)) + y_2 = network(x_2) + self.assertEqual(y_2.get_shape().as_list(), [None, 2]) + + # test network `trainable` attribute + network.trainable = False + self.assertEqual(network.weights, dense.weights) + self.assertEqual(network.trainable_weights, []) + self.assertEqual(network.non_trainable_weights, + dense.trainable_weights + dense.non_trainable_weights) def test_trainable_weights(self): a = keras.layers.Input(shape=(2,)) @@ -100,41 +327,6 @@ class TopologyConstructionTest(test.TestCase): self.assertListEqual(model.trainable_weights, []) self.assertListEqual(model.non_trainable_weights, weights) - def test_weight_loading(self): - with self.test_session(): - a = keras.layers.Input(shape=(2,)) - x = keras.layers.Dense(3)(a) - b = keras.layers.Dense(1)(x) - model = keras.models.Model(a, b) - - x = np.random.random((3, 2)) - ref_y = model.predict(x) - weights = model.get_weights() - model.set_weights(weights) - y = model.predict(x) - self.assertAllClose(ref_y, y) - - with self.assertRaises(ValueError): - model.set_weights(weights[1:]) - with self.assertRaises(ValueError): - model.set_weights(weights[::-1]) - - if h5py is None: - return # Skip rest of test if H5py isn't available. - - temp_dir = self.get_temp_dir() - self.addCleanup(shutil.rmtree, temp_dir) - - h5_path = os.path.join(temp_dir, 'test.h5') - model.save_weights(h5_path) - model.load_weights(h5_path) - y = model.predict(x) - self.assertAllClose(ref_y, y) - - model.load_weights(h5_path, by_name=True) - y = model.predict(x) - self.assertAllClose(ref_y, y) - def test_learning_phase(self): with self.test_session(): a = keras.layers.Input(shape=(32,), name='input_a') @@ -302,7 +494,7 @@ class TopologyConstructionTest(test.TestCase): self.assertListEqual([x.shape for x in fn_outputs], [(10, 64), (10, 5)]) # test get_source_inputs - self.assertListEqual(keras.engine.topology.get_source_inputs(c), [a, b]) + self.assertListEqual(keras.engine.network.get_source_inputs(c), [a, b]) # serialization / deserialization json_config = model.to_json() @@ -339,7 +531,10 @@ class TopologyConstructionTest(test.TestCase): e = keras.layers.Input(shape=(32,), name='input_e') f = keras.layers.Input(shape=(32,), name='input_f') + self.assertEqual(len(model.inputs), 2) g, h = model([e, f]) + self.assertEqual(len(model.inputs), 2) + self.assertEqual(g.name, 'model/dense_2/BiasAdd:0') self.assertListEqual(g.get_shape().as_list(), c.get_shape().as_list()) self.assertListEqual(h.get_shape().as_list(), d.get_shape().as_list()) @@ -520,7 +715,9 @@ class TopologyConstructionTest(test.TestCase): j = keras.layers.Input(shape=(32,), name='input_j') k = keras.layers.Input(shape=(32,), name='input_k') + self.assertEqual(len(model.inputs), 2) m, n = model([j, k]) + self.assertEqual(len(model.inputs), 2) tf_model = keras.models.Model([j, k], [m, n]) j_tf = array_ops.placeholder(dtype=dtypes.float32, shape=(None, 32)) @@ -546,96 +743,63 @@ class TopologyConstructionTest(test.TestCase): model = keras.models.Model(a, b) self.assertEqual(model.output_mask.get_shape().as_list(), [None, 10]) - def test_weight_preprocessing(self): - input_dim = 3 - output_dim = 3 - size = 2 - cases = [ - [ - (keras.layers.Bidirectional(keras.layers.SimpleRNN(2))), - [np.random.random((2, 1)), np.random.random((2, 1))], - (None, 3, 2), - ], - [ - (keras.layers.TimeDistributed(keras.layers.Dense(1))), - [np.random.random((2, 1)), np.random.random((1,))], - (None, 3, 2), - ], - [ - (keras.layers.Conv1D(output_dim, size, use_bias=False)), - [np.random.random((output_dim, input_dim, size, 1))], - (None, 4, input_dim), - ], - [ - (keras.layers.Conv2D(output_dim, size, - use_bias=False, data_format='channels_first')), - [np.random.random((output_dim, input_dim, size, size))], - (None, input_dim, 4, 4), - ], - [ - (keras.layers.Conv2DTranspose(output_dim, size, - use_bias=False, - data_format='channels_first')), - [np.random.random((output_dim, input_dim, size, size))], - (None, input_dim, 4, 4), - ], - [ - (keras.layers.Conv2DTranspose(output_dim, size, - use_bias=False, - data_format='channels_last')), - [np.random.random((size, size, input_dim, output_dim))], - (None, 4, 4, input_dim), - ], - [ - (keras.layers.Conv3D(output_dim, size, - use_bias=False, data_format='channels_first')), - [np.random.random((output_dim, input_dim, size, size, size))], - (None, input_dim, 4, 4, 4), - ], - [ - (keras.layers.GRU(output_dim)), - [np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,)), - np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,)), - np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,))], - (None, 4, input_dim), - ], - [ - (keras.layers.LSTM(output_dim)), - [np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,)), - np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,)), - np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,)), - np.random.random((input_dim, output_dim)), - np.random.random((output_dim, output_dim)), - np.random.random((output_dim,))], - (None, 4, input_dim), - ], - ] - for layer, weights, input_shape in cases: - layer.build(input_shape) - _ = keras.engine.topology.preprocess_weights_for_loading( - layer, weights, original_keras_version='1') - - model = keras.models.Sequential([keras.layers.Dense(2, input_dim=2)]) - _ = keras.engine.topology.preprocess_weights_for_loading( - model, model.weights, original_keras_version='1') - - x = keras.Input((2,)) - y = keras.layers.Dense(2)(x) - model = keras.models.Model(x, y) - _ = keras.engine.topology.preprocess_weights_for_loading( - model, model.weights, original_keras_version='1') + def testMaskingSingleInput(self): + + class MaskedLayer(keras.layers.Layer): + + def call(self, inputs, mask=None): + if mask is not None: + return inputs * mask + return inputs + + def compute_mask(self, inputs, mask=None): + return array_ops.ones_like(inputs) + + if context.executing_eagerly(): + a = constant_op.constant([2] * 32) + mask = constant_op.constant([0, 1] * 16) + a._keras_mask = mask + b = MaskedLayer().apply(a) + self.assertTrue(hasattr(b, '_keras_mask')) + self.assertAllEqual( + self.evaluate(array_ops.ones_like(mask)), + self.evaluate(getattr(b, '_keras_mask'))) + self.assertAllEqual(self.evaluate(a * mask), self.evaluate(b)) + else: + x = keras.Input(shape=(32,)) + y = MaskedLayer()(x) # pylint: disable=not-callable + network = keras.engine.Network(x, y) + + # test callability on Input + x_2 = keras.Input(shape=(32,)) + y_2 = network(x_2) + self.assertEqual(y_2.get_shape().as_list(), [None, 32]) + + # test callability on regular tensor + x_2 = array_ops.placeholder(dtype='float32', shape=(None, 32)) + y_2 = network(x_2) + self.assertEqual(y_2.get_shape().as_list(), [None, 32]) + + def test_activity_regularization_with_model_composition(self): + + def reg(x): + return keras.backend.sum(x) + + net_a_input = keras.Input((2,)) + net_a = net_a_input + net_a = keras.layers.Dense(2, kernel_initializer='ones', + use_bias=False, + activity_regularizer=reg)(net_a) + model_a = keras.Model([net_a_input], [net_a]) + + net_b_input = keras.Input((2,)) + net_b = model_a(net_b_input) + model_b = keras.Model([net_b_input], [net_b]) + + model_b.compile(optimizer='sgd', loss=None) + x = np.ones((1, 2)) + loss = model_b.evaluate(x) + self.assertEqual(loss, 4.) def test_layer_sharing_at_heterogenous_depth(self): with self.test_session(): @@ -685,5 +849,92 @@ class TopologyConstructionTest(test.TestCase): output_val_2 = m2.predict(x_val) self.assertAllClose(output_val, output_val_2, atol=1e-6) + def test_explicit_training_argument(self): + with self.test_session(): + a = keras.layers.Input(shape=(2,)) + b = keras.layers.Dropout(0.5)(a) + base_model = keras.models.Model(a, b) + + a = keras.layers.Input(shape=(2,)) + b = base_model(a, training=False) + model = keras.models.Model(a, b) + + x = np.ones((100, 2)) + y = np.ones((100, 2)) + model.compile(optimizer='sgd', loss='mse') + loss = model.train_on_batch(x, y) + self.assertEqual(loss, 0) # In inference mode, output is equal to input. + + a = keras.layers.Input(shape=(2,)) + b = base_model(a, training=True) + model = keras.models.Model(a, b) + preds = model.predict(x) + self.assertEqual(np.min(preds), 0.) # At least one unit was dropped. + + +class DeferredModeTest(test.TestCase): + + def testDeferredTensorAttributes(self): + x = tf_base_layers._DeferredTensor(shape=(None, 2), + dtype='float32', + name='x') + self.assertEqual(str(x), + 'DeferredTensor(\'x\', shape=(?, 2), dtype=float32)') + self.assertEqual(repr(x), + '<_DeferredTensor \'x\' shape=(?, 2) dtype=float32>') + + @test_util.run_in_graph_and_eager_modes() + def testSimpleNetworkBuilding(self): + inputs = keras.engine.Input(shape=(32,)) + if context.executing_eagerly(): + self.assertIsInstance(inputs, tf_base_layers._DeferredTensor) + self.assertEqual(inputs.dtype.name, 'float32') + self.assertEqual(inputs.shape.as_list(), [None, 32]) + + x = keras.layers.Dense(2)(inputs) + if context.executing_eagerly(): + self.assertIsInstance(x, tf_base_layers._DeferredTensor) + self.assertEqual(x.dtype.name, 'float32') + self.assertEqual(x.shape.as_list(), [None, 2]) + + outputs = keras.layers.Dense(4)(x) + network = keras.engine.Network(inputs, outputs) + self.assertIsInstance(network, keras.engine.Network) + + if context.executing_eagerly(): + # It should be possible to call such a network on EagerTensors. + inputs = constant_op.constant( + np.random.random((10, 32)).astype('float32')) + outputs = network(inputs) + self.assertEqual(outputs.shape.as_list(), [10, 4]) + + @test_util.run_in_graph_and_eager_modes() + def testMultiIONetworkbuilding(self): + input_a = keras.engine.Input(shape=(32,)) + input_b = keras.engine.Input(shape=(16,)) + a = keras.layers.Dense(16)(input_a) + + class AddLayer(keras.layers.Layer): + + def call(self, inputs): + return inputs[0] + inputs[1] + + def compute_output_shape(self, input_shape): + return input_shape[0] + + c = AddLayer()([a, input_b]) # pylint: disable=not-callable + c = keras.layers.Dense(2)(c) + + network = keras.engine.Network([input_a, input_b], [a, c]) + if context.executing_eagerly(): + a_val = constant_op.constant( + np.random.random((10, 32)).astype('float32')) + b_val = constant_op.constant( + np.random.random((10, 16)).astype('float32')) + outputs = network([a_val, b_val]) + self.assertEqual(len(outputs), 2) + self.assertEqual(outputs[0].shape.as_list(), [10, 16]) + self.assertEqual(outputs[1].shape.as_list(), [10, 2]) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/_impl/keras/engine/training.py b/tensorflow/python/keras/_impl/keras/engine/training.py index 699ae2edf0db1cdcc73763607d04329c76888565..4acb41553eda2e07962f6ac510f08988a5adb90c 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training.py +++ b/tensorflow/python/keras/_impl/keras/engine/training.py @@ -18,551 +18,95 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import copy - import numpy as np +from tensorflow.python.eager import context +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util from tensorflow.python.keras._impl.keras import backend as K -from tensorflow.python.keras._impl.keras import callbacks as cbks from tensorflow.python.keras._impl.keras import losses from tensorflow.python.keras._impl.keras import metrics as metrics_module from tensorflow.python.keras._impl.keras import optimizers -from tensorflow.python.keras._impl.keras.engine.topology import Network -from tensorflow.python.keras._impl.keras.utils.data_utils import GeneratorEnqueuer -from tensorflow.python.keras._impl.keras.utils.data_utils import OrderedEnqueuer -from tensorflow.python.keras._impl.keras.utils.data_utils import Sequence -from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar +from tensorflow.python.keras._impl.keras.engine import training_arrays +from tensorflow.python.keras._impl.keras.engine import training_eager +from tensorflow.python.keras._impl.keras.engine import training_generator +from tensorflow.python.keras._impl.keras.engine import training_utils +from tensorflow.python.keras._impl.keras.engine.base_layer import Layer +from tensorflow.python.keras._impl.keras.engine.network import Network +from tensorflow.python.keras._impl.keras.utils.generic_utils import slice_arrays +from tensorflow.python.layers.base import _DeferredTensor +from tensorflow.python.ops import array_ops from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import optimizer as tf_optimizer_module +from tensorflow.python.util.tf_export import tf_export -try: - from scipy.sparse import issparse # pylint: disable=g-import-not-at-top -except ImportError: - issparse = None - - -def _standardize_input_data(data, - names, - shapes=None, - check_batch_axis=True, - exception_prefix=''): - """Normalizes inputs and targets provided by users. - - Users may pass data as a list of arrays, dictionary of arrays, - or as a single array. We normalize this to an ordered list of - arrays (same order as `names`), while checking that the provided - arrays have shapes that match the network's expectations. - - Arguments: - data: User-provided input data (polymorphic). - names: List of expected array names. - shapes: Optional list of expected array shapes. - check_batch_axis: Boolean; whether to check that - the batch axis of the arrays matches the expected - value found in `shapes`. - exception_prefix: String prefix used for exception formatting. - - Returns: - List of standardized input arrays (one array per model input). - - Raises: - ValueError: in case of improperly formatted user-provided data. - """ - if not names: - if data is not None and hasattr(data, '__len__') and len(data): - raise ValueError('Error when checking model ' + exception_prefix + ': ' - 'expected no data, but got:', data) - return [] - if data is None: - return [None for _ in range(len(names))] - - if isinstance(data, dict): - try: - data = [ - data[x].values - if data[x].__class__.__name__ == 'DataFrame' else data[x] - for x in names - ] - data = [np.expand_dims(x, 1) if x.ndim == 1 else x for x in data] - except KeyError as e: - raise ValueError('No data provided for "' + e.args[0] + '". Need data ' - 'for each key in: ' + str(names)) - elif isinstance(data, list): - data = [ - x.values if x.__class__.__name__ == 'DataFrame' else x for x in data - ] - data = [ - np.expand_dims(x, 1) if x is not None and x.ndim == 1 else x - for x in data - ] - else: - data = data.values if data.__class__.__name__ == 'DataFrame' else data - data = [np.expand_dims(data, 1)] if data.ndim == 1 else [data] - - if len(data) != len(names): - if data and hasattr(data[0], 'shape'): - raise ValueError('Error when checking model ' + exception_prefix + - ': the list of Numpy arrays that you are passing to ' - 'your model is not the size the model expected. ' - 'Expected to see ' + str(len(names)) + ' array(s), ' - 'but instead got the following list of ' + - str(len(data)) + ' arrays: ' + str(data)[:200] + '...') - elif len(names) > 1: - raise ValueError( - 'Error when checking model ' + exception_prefix + - ': you are passing a list as input to your model, ' - 'but the model expects a list of ' + str(len(names)) + - ' Numpy arrays instead. The list you passed was: ' + str(data)[:200]) - elif len(data) == 1 and not hasattr(data[0], 'shape'): - raise TypeError('Error when checking model ' + exception_prefix + - ': data should be a Numpy array, or list/dict of ' - 'Numpy arrays. Found: ' + str(data)[:200] + '...') - elif len(names) == 1: - data = [np.asarray(data)] - - # Check shapes compatibility. - if shapes: - for i in range(len(names)): - if shapes[i] is not None: - data_shape = data[i].shape - shape = shapes[i] - if data[i].ndim != len(shape): - raise ValueError('Error when checking ' + exception_prefix + - ': expected ' + names[i] + ' to have ' + - str(len(shape)) + ' dimensions, but got array ' - 'with shape ' + str(data_shape)) - if not check_batch_axis: - data_shape = data_shape[1:] - shape = shape[1:] - for dim, ref_dim in zip(data_shape, shape): - if ref_dim != dim and ref_dim: - raise ValueError( - 'Error when checking ' + exception_prefix + ': expected ' + - names[i] + ' to have shape ' + str(shape) + - ' but got array with shape ' + str(data_shape)) - return data - - -def _standardize_sample_or_class_weights(x_weight, output_names, weight_type): - """Maps `sample_weight` or `class_weight` to model outputs. - - Arguments: - x_weight: User-provided `sample_weight` or `class_weight` argument. - output_names: List of output names (strings) in the model. - weight_type: A string used purely for exception printing. - - Returns: - A list of `sample_weight` or `class_weight` where there are exactly - one element per model output. - Raises: - ValueError: In case of invalid user-provided argument. - """ - if x_weight is None or len(x_weight) == 0: # pylint: disable=g-explicit-length-test - return [None for _ in output_names] - if len(output_names) == 1: - if isinstance(x_weight, list) and len(x_weight) == 1: - return x_weight - if isinstance(x_weight, dict) and output_names[0] in x_weight: - return [x_weight[output_names[0]]] - else: - return [x_weight] - if isinstance(x_weight, list): - if len(x_weight) != len(output_names): - raise ValueError('Provided `' + weight_type + '` was a list of ' + - str(len(x_weight)) + ' elements, but the model has ' + - str(len(output_names)) + ' outputs. ' - 'You should provide one `' + weight_type + '`' - 'array per model output.') - return x_weight - if isinstance(x_weight, dict): - x_weights = [] - for name in output_names: - x_weights.append(x_weight.get(name)) - return x_weights - else: - raise TypeError( - 'The model has multiple outputs, so `' + weight_type + '` ' - 'should be either a list or a dict. ' - 'Provided `' + weight_type + '` type not understood: ' + str(x_weight)) - - -def _standardize_class_weights(class_weight, output_names): - return _standardize_sample_or_class_weights(class_weight, output_names, - 'class_weight') - - -def _standardize_sample_weights(sample_weight, output_names): - return _standardize_sample_or_class_weights(sample_weight, output_names, - 'sample_weight') - - -def _check_array_lengths(inputs, targets, weights=None): - """Does user input validation for numpy arrays. - - Arguments: - inputs: list of Numpy arrays of inputs. - targets: list of Numpy arrays of targets. - weights: list of Numpy arrays of sample weights. - - Raises: - ValueError: in case of incorrectly formatted data. - """ - - def set_of_lengths(x): - # return a set with the variation between - # different shapes, with None => 0 - if x is None: - return {0} - else: - return set([0 if y is None else y.shape[0] for y in x]) - - set_x = set_of_lengths(inputs) - set_y = set_of_lengths(targets) - set_w = set_of_lengths(weights) - if len(set_x) > 1: - raise ValueError('All input arrays (x) should have ' - 'the same number of samples. Got array shapes: ' + - str([x.shape for x in inputs])) - if len(set_y) > 1: - raise ValueError('All target arrays (y) should have ' - 'the same number of samples. Got array shapes: ' + - str([y.shape for y in targets])) - if set_x and set_y and list(set_x)[0] != list(set_y)[0]: - raise ValueError('Input arrays should have ' - 'the same number of samples as target arrays. ' - 'Found ' + str(list(set_x)[0]) + ' input samples ' - 'and ' + str(list(set_y)[0]) + ' target samples.') - if len(set_w) > 1: - raise ValueError('All sample_weight arrays should have ' - 'the same number of samples. Got array shapes: ' + - str([w.shape for w in weights])) - if set_y and set_w and list(set_y)[0] != list(set_w)[0]: - raise ValueError('Sample_weight arrays should have ' - 'the same number of samples as target arrays. Got ' + - str(list(set_y)[0]) + ' input samples and ' + - str(list(set_w)[0]) + ' target samples.') - - -def _check_loss_and_target_compatibility(targets, loss_fns, output_shapes): - """Does validation on the compatibility of targets and loss functions. - - This helps prevent users from using loss functions incorrectly. - - Arguments: - targets: list of Numpy arrays of targets. - loss_fns: list of loss functions. - output_shapes: list of shapes of model outputs. - - Raises: - ValueError: if a loss function or target array - is incompatible with an output. - """ - key_losses = { - losses.mean_squared_error, losses.binary_crossentropy, - losses.categorical_crossentropy - } - for y, loss, shape in zip(targets, loss_fns, output_shapes): - if loss is None: - continue - if loss is losses.categorical_crossentropy: - if y.shape[-1] == 1: - raise ValueError('You are passing a target array of shape ' + str( - y.shape) + ' while using as loss `categorical_crossentropy`. ' - '`categorical_crossentropy` expects ' - 'targets to be binary matrices (1s and 0s) ' - 'of shape (samples, classes). ' - 'If your targets are integer classes, ' - 'you can convert them to the expected format via:\n' - '```\n' - 'from keras.utils import to_categorical\n' - 'y_binary = to_categorical(y_int)\n' - '```\n' - '\n' - 'Alternatively, you can use the loss function ' - '`sparse_categorical_crossentropy` instead, ' - 'which does expect integer targets.') - if loss in key_losses: - for target_dim, out_dim in zip(y.shape[1:], shape[1:]): - if out_dim is not None and target_dim != out_dim: - raise ValueError('A target array with shape ' + str(y.shape) + - ' was passed for an output of shape ' + str(shape) + - ' while using as loss `' + loss.__name__ + '`. ' - 'This loss expects ' - 'targets to have the same shape ' - 'as the output.') - - -def _collect_metrics(metrics, output_names): - """Maps metric functions to model outputs. - - Arguments: - metrics: a list or dict of metric functions. - output_names: a list of the names (strings) of model outputs. - - Returns: - A list (one entry per model output) of lists of metric functions. - For instance, if the model has 2 outputs, and for the first output - we want to compute "binary_accuracy" and "binary_crossentropy", - and just "binary_accuracy" for the second output, - the list would look like: - `[[binary_accuracy, binary_crossentropy], [binary_accuracy]]` - - Raises: - TypeError: if an incorrect type is passed for the `metrics` argument. - """ - if not metrics: - return [[] for _ in output_names] - if isinstance(metrics, list): - # we then apply all metrics to all outputs. - return [copy.copy(metrics) for _ in output_names] - elif isinstance(metrics, dict): - nested_metrics = [] - for name in output_names: - output_metrics = metrics.get(name, []) - if not isinstance(output_metrics, list): - output_metrics = [output_metrics] - nested_metrics.append(output_metrics) - return nested_metrics - else: - raise TypeError('Type of `metrics` argument not understood. ' - 'Expected a list or dictionary, found: ' + str(metrics)) - - -def _batch_shuffle(index_array, batch_size): - """Shuffles an array in a batch-wise fashion. - - Useful for shuffling HDF5 arrays - (where one cannot access arbitrary indices). - - Arguments: - index_array: array of indices to be shuffled. - batch_size: integer. - - Returns: - The `index_array` array, shuffled in a batch-wise fashion. - """ - batch_count = int(len(index_array) / batch_size) - # to reshape we need to be cleanly divisible by batch size - # we stash extra items and reappend them after shuffling - last_batch = index_array[batch_count * batch_size:] - index_array = index_array[:batch_count * batch_size] - index_array = index_array.reshape((batch_count, batch_size)) - np.random.shuffle(index_array) - index_array = index_array.flatten() - return np.append(index_array, last_batch) - - -def _make_batches(size, batch_size): - """Returns a list of batch indices (tuples of indices). - - Arguments: - size: Integer, total size of the data to slice into batches. - batch_size: Integer, batch size. - - Returns: - A list of tuples of array indices. - """ - num_batches = (size + batch_size - 1) // batch_size # round up - return [(i * batch_size, min(size, (i + 1) * batch_size)) - for i in range(num_batches)] +@tf_export('keras.models.Model', 'keras.Model') +class Model(Network): + """`Model` groups layers into an object with training and inference features. + There are two ways to instantiate a `Model`: -def _slice_arrays(arrays, start=None, stop=None): - """Slice an array or list of arrays. + 1 - With the "functional API", where you start from `Input`, + you chain layer calls to specify the model's forward pass, + and finally you create your model from inputs and outputs: - This takes an array-like, or a list of - array-likes, and outputs: - - arrays[start:stop] if `arrays` is an array-like - - [x[start:stop] for x in arrays] if `arrays` is a list + ```python + import tensorflow as tf - Can also work on list/array of indices: `_slice_arrays(x, indices)` + inputs = tf.keras.Input(shape=(3,)) + x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs) + outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(x) + model = tf.keras.Model(inputs=inputs, outputs=outputs) + ``` - Arguments: - arrays: Single array or list of arrays. - start: can be an integer index (start index) - or a list/array of indices - stop: integer (stop index); should be None if - `start` was a list. + 2 - By subclassing the `Model` class: in that case, you should define your + layers in `__init__` and you should implement the model's forward pass + in `call`. - Returns: - A slice of the array(s). - """ - if arrays is None: - return [None] - elif isinstance(arrays, list): - if hasattr(start, '__len__'): - # hdf5 datasets only support list objects as indices - if hasattr(start, 'shape'): - start = start.tolist() - return [None if x is None else x[start] for x in arrays] - else: - return [None if x is None else x[start:stop] for x in arrays] - else: - if hasattr(start, '__len__'): - if hasattr(start, 'shape'): - start = start.tolist() - return arrays[start] - elif hasattr(start, '__getitem__'): - return arrays[start:stop] - else: - return [None] + ```python + import tensorflow as tf + class MyModel(tf.keras.Model): -def _weighted_masked_objective(fn): - """Adds support for masking and sample-weighting to an objective function. + def __init__(self): + self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu) + self.dense2 = tf.keras.layers.Dense(5, activation=tf.nn.softmax) - It transforms an objective function `fn(y_true, y_pred)` - into a sample-weighted, cost-masked objective function - `fn(y_true, y_pred, weights, mask)`. + def call(self, inputs): + x = self.dense1(inputs) + return self.dense2(x) - Arguments: - fn: The objective function to wrap, - with signature `fn(y_true, y_pred)`. + model = MyModel() + ``` - Returns: - A function with signature `fn(y_true, y_pred, weights, mask)`. - """ - if fn is None: - return None + If you subclass `Model`, you can optionally have + a `training` argument (boolean) in `call`, which you can use to specify + a different behavior in training and inference: - def weighted(y_true, y_pred, weights, mask=None): - """Wrapper function. + ```python + import tensorflow as tf - Arguments: - y_true: `y_true` argument of `fn`. - y_pred: `y_pred` argument of `fn`. - weights: Weights tensor. - mask: Mask tensor. + class MyModel(tf.keras.Model): - Returns: - Scalar tensor. - """ - # score_array has ndim >= 2 - score_array = fn(y_true, y_pred) - if mask is not None: - # Cast the mask to floatX to avoid float64 upcasting in theano - mask = K.cast(mask, K.floatx()) - # mask should have the same shape as score_array - score_array *= mask - # the loss per batch should be proportional - # to the number of unmasked samples. - score_array /= K.mean(mask) - - # apply sample weighting - if weights is not None: - # reduce score_array to same ndim as weight array - ndim = K.ndim(score_array) - weight_ndim = K.ndim(weights) - score_array = K.mean(score_array, axis=list(range(weight_ndim, ndim))) - score_array *= weights - score_array /= K.mean(K.cast(K.not_equal(weights, 0), K.floatx())) - return K.mean(score_array) - - return weighted - - -def _standardize_weights(y, - sample_weight=None, - class_weight=None, - sample_weight_mode=None): - """Performs sample weight validation and standardization. - - Everything gets normalized to a single sample-wise (or timestep-wise) - weight array. - - Arguments: - y: Numpy array of model targets to be weighted. - sample_weight: User-provided `sample_weight` argument. - class_weight: User-provided `class_weight` argument. - sample_weight_mode: One of `None` or `"temporal"`. - `"temporal"` indicated that we expect 2D weight data - that will be applied to the last 2 dimensions of - the targets (i.e. we are weighting timesteps, not samples). - - Returns: - A numpy array of target weights, one entry per sample to weight. - - Raises: - ValueError: In case of invalid user-provided arguments. - """ - if sample_weight_mode is not None: - if sample_weight_mode != 'temporal': - raise ValueError('"sample_weight_mode ' - 'should be None or "temporal". ' - 'Found: ' + str(sample_weight_mode)) - if len(y.shape) < 3: - raise ValueError('Found a sample_weight array for ' - 'an input with shape ' + str(y.shape) + '. ' - 'Timestep-wise sample weighting (use of ' - 'sample_weight_mode="temporal") is restricted to ' - 'outputs that are at least 3D, i.e. that have ' - 'a time dimension.') - if sample_weight is not None and len(sample_weight.shape) != 2: - raise ValueError('Found a sample_weight array with shape ' + - str(sample_weight.shape) + '. ' - 'In order to use timestep-wise sample weighting, ' - 'you should pass a 2D sample_weight array.') - else: - if sample_weight is not None and len(sample_weight.shape) != 1: - raise ValueError('Found a sample_weight array with shape ' + - str(sample_weight.shape) + '. ' - 'In order to use timestep-wise sample weights, ' - 'you should specify ' - 'sample_weight_mode="temporal" ' - 'in compile(). If you just mean to use ' - 'sample-wise weights, make sure your ' - 'sample_weight array is 1D.') - - if sample_weight is not None: - if len(sample_weight.shape) > len(y.shape): - raise ValueError( - 'Found a sample_weight with shape' + str(sample_weight.shape) + '.' - 'Expected sample_weight with rank ' - 'less than or equal to ' + str(len(y.shape))) - - if y.shape[:sample_weight.ndim] != sample_weight.shape: - raise ValueError( - 'Found a sample_weight array with shape ' + str(sample_weight.shape) + - ' for an input with shape ' + str(y.shape) + '. ' - 'sample_weight cannot be broadcast.') - return sample_weight - elif isinstance(class_weight, dict): - if len(y.shape) > 2: - raise ValueError('`class_weight` not supported for ' - '3+ dimensional targets.') - if y.shape[1] > 1: - y_classes = y.argmax(axis=1) - elif y.shape[1] == 1: - y_classes = np.reshape(y, y.shape[0]) - else: - y_classes = y - - weights = np.asarray( - [class_weight[cls] for cls in y_classes if cls in class_weight]) - - if len(weights) != len(y_classes): - # subtract the sets to pick all missing classes - existing_classes = set(y_classes) - existing_class_weight = set(class_weight.keys()) - raise ValueError('`class_weight` must contain all classes in the data.' - ' The classes %s exist in the data but not in ' - '`class_weight`.' % - (existing_classes - existing_class_weight)) - return weights - else: - if sample_weight_mode is None: - return np.ones((y.shape[0],), dtype=K.floatx()) - else: - return np.ones((y.shape[0], y.shape[1]), dtype=K.floatx()) + def __init__(self): + self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu) + self.dense2 = tf.keras.layers.Dense(5, activation=tf.nn.softmax) + self.dropout = tf.keras.layers.Dropout(0.5) + def call(self, inputs, training=False): + x = self.dense1(inputs) + if training: + x = self.dropout(x, training=training) + return self.dense2(x) -class Model(Network): - """The `Model` class adds training & evaluation routines to a `Network`. + model = MyModel() + ``` """ def compile(self, optimizer, - loss, + loss=None, metrics=None, loss_weights=None, sample_weight_mode=None, @@ -618,10 +162,30 @@ class Model(Network): `optimizer`, `loss`, `metrics` or `sample_weight_mode`. """ loss = loss or {} + if context.executing_eagerly() and not isinstance( + optimizer, (tf_optimizer_module.Optimizer, optimizers.TFOptimizer)): + raise ValueError('Only TF native optimizers are supported in Eager mode.') + self.optimizer = optimizers.get(optimizer) self.loss = loss + self.metrics = metrics or [] self.loss_weights = loss_weights + if context.executing_eagerly() and sample_weight_mode is not None: + raise ValueError('sample_weight_mode is not supported in Eager mode.') self.sample_weight_mode = sample_weight_mode + if context.executing_eagerly() and weighted_metrics is not None: + raise ValueError('weighted_metrics is not supported in Eager mode.') + self.weighted_metrics = weighted_metrics + if context.executing_eagerly() and target_tensors is not None: + raise ValueError('target_tensors is not supported in Eager mode.') + self.target_tensors = target_tensors + + if not self.built: + # Model is not compilable because it does not know its number of inputs + # and outputs, nor their shapes and names. We will compile after the first + # time the model gets called on training data. + return + self._is_compiled = True # Prepare loss functions. if isinstance(loss, dict): @@ -651,7 +215,9 @@ class Model(Network): loss_function = losses.get(loss) loss_functions = [loss_function for _ in range(len(self.outputs))] self.loss_functions = loss_functions - weighted_losses = [_weighted_masked_objective(fn) for fn in loss_functions] + + weighted_losses = [training_utils.weighted_masked_objective(fn) + for fn in loss_functions] skip_target_indices = [] skip_target_weighing_indices = [] self._feed_outputs = [] @@ -664,11 +230,12 @@ class Model(Network): skip_target_weighing_indices.append(i) # Prepare output masks. - masks = self.compute_mask(self.inputs, mask=None) - if masks is None: - masks = [None for _ in self.outputs] - if not isinstance(masks, list): - masks = [masks] + if not context.executing_eagerly(): + masks = self.compute_mask(self.inputs, mask=None) + if masks is None: + masks = [None for _ in self.outputs] + if not isinstance(masks, list): + masks = [masks] # Prepare loss weights. if loss_weights is None: @@ -694,11 +261,35 @@ class Model(Network): else: raise TypeError('Could not interpret loss_weights argument: ' + str(loss_weights) + ' - expected a list of dicts.') + self.loss_weights_list = loss_weights_list + + # initialization for Eager mode execution + if context.executing_eagerly(): + if target_tensors is not None: + raise ValueError('target_tensors are not currently supported in Eager ' + 'mode.') + self.total_loss = None + self.metrics_tensors = [] + self.metrics_names = ['loss'] + for i in range(len(self.outputs)): + if len(self.outputs) > 1: + self.metrics_names.append(self.output_names[i] + '_loss') + self.nested_metrics = training_utils.collect_metrics(metrics, + self.output_names) + self._feed_sample_weight_modes = [] + for i in range(len(self.outputs)): + self._feed_sample_weight_modes.append(None) + self.sample_weights = [] + self.targets = [] + for i in range(len(self.outputs)): + self._feed_output_names.append(self.output_names[i]) + self._collected_trainable_weights = self.trainable_weights + return # Prepare targets of model. self.targets = [] self._feed_targets = [] - if target_tensors is not None: + if target_tensors not in (None, []): if isinstance(target_tensors, list): if len(target_tensors) != len(self.outputs): raise ValueError( @@ -720,13 +311,14 @@ class Model(Network): else: raise TypeError('Expected `target_tensors` to be ' 'a list or dict, but got:', target_tensors) + for i in range(len(self.outputs)): if i in skip_target_indices: self.targets.append(None) else: - shape = self._internal_output_shapes[i] + shape = K.int_shape(self.outputs[i]) name = self.output_names[i] - if target_tensors is not None: + if target_tensors not in (None, []): target = target_tensors[i] else: target = None @@ -769,7 +361,7 @@ class Model(Network): weight = K.placeholder(ndim=2, name=name + '_sample_weights') sample_weight_modes.append('temporal') else: - weight = K.placeholder(ndim=1, name=name + '_sample_weights') + weight = K.placeholder(ndim=1, name=name + 'sample_weights') sample_weight_modes.append(None) sample_weights.append(weight) elif isinstance(sample_weight_mode, list): @@ -800,12 +392,12 @@ class Model(Network): sample_weights.append(None) else: if sample_weight_mode == 'temporal': - sample_weights.append( - K.placeholder(ndim=2, name=name + '_sample_weights')) + sample_weights.append(array_ops.placeholder_with_default( + [[1.]], shape=[None, None], name=name + '_sample_weights')) sample_weight_modes.append('temporal') else: - sample_weights.append( - K.placeholder(ndim=1, name=name + '_sample_weights')) + sample_weights.append(array_ops.placeholder_with_default( + [1.], shape=[None], name=name + '_sample_weights')) sample_weight_modes.append(None) self.sample_weight_modes = sample_weight_modes self._feed_sample_weight_modes = [] @@ -814,7 +406,6 @@ class Model(Network): self._feed_sample_weight_modes.append(self.sample_weight_modes[i]) # Prepare metrics. - self.metrics = metrics self.weighted_metrics = weighted_metrics self.metrics_names = ['loss'] self.metrics_tensors = [] @@ -854,17 +445,11 @@ class Model(Network): # List of same size as output_names. # contains tuples (metrics for output, names of metrics). - nested_metrics = _collect_metrics(metrics, self.output_names) - nested_weighted_metrics = _collect_metrics(weighted_metrics, - self.output_names) - - def append_metric(layer_index, metric_name, metric_tensor): - """Helper function used in loop below.""" - if len(self.output_names) > 1: - metric_name = self.output_names[layer_index] + '_' + metric_name - self.metrics_names.append(metric_name) - self.metrics_tensors.append(metric_tensor) - + nested_metrics = training_utils.collect_metrics(metrics, self.output_names) + nested_weighted_metrics = training_utils.collect_metrics(weighted_metrics, + self.output_names) + self.metrics_updates = [] + self.stateful_metric_names = [] with K.name_scope('metrics'): for i in range(len(self.outputs)): if i in skip_target_indices: @@ -883,42 +468,67 @@ class Model(Network): if metric in ('accuracy', 'acc', 'crossentropy', 'ce'): # custom handling of accuracy/crossentropy # (because of class mode duality) - output_shape = self._internal_output_shapes[i] + output_shape = self.outputs[i].get_shape().as_list() if (output_shape[-1] == 1 or self.loss_functions[i] == losses.binary_crossentropy): # case: binary accuracy/crossentropy if metric in ('accuracy', 'acc'): - acc_fn = metrics_module.binary_accuracy + metric_fn = metrics_module.binary_accuracy elif metric in ('crossentropy', 'ce'): - acc_fn = metrics_module.binary_crossentropy + metric_fn = metrics_module.binary_crossentropy elif self.loss_functions[ i] == losses.sparse_categorical_crossentropy: # case: categorical accuracy/crossentropy with sparse targets if metric in ('accuracy', 'acc'): - acc_fn = metrics_module.sparse_categorical_accuracy + metric_fn = metrics_module.sparse_categorical_accuracy elif metric in ('crossentropy', 'ce'): - acc_fn = metrics_module.sparse_categorical_crossentropy + metric_fn = metrics_module.sparse_categorical_crossentropy else: # case: categorical accuracy/crossentropy if metric in ('accuracy', 'acc'): - acc_fn = metrics_module.categorical_accuracy + metric_fn = metrics_module.categorical_accuracy elif metric in ('crossentropy', 'ce'): - acc_fn = metrics_module.categorical_crossentropy + metric_fn = metrics_module.categorical_crossentropy if metric in ('accuracy', 'acc'): suffix = 'acc' elif metric in ('crossentropy', 'ce'): suffix = 'ce' - weighted_metric_fn = _weighted_masked_objective(acc_fn) + weighted_metric_fn = training_utils.weighted_masked_objective( + metric_fn) metric_name = metric_name_prefix + suffix else: metric_fn = metrics_module.get(metric) - weighted_metric_fn = _weighted_masked_objective(metric_fn) - metric_name = metric_name_prefix + metric_fn.__name__ + weighted_metric_fn = training_utils.weighted_masked_objective( + metric_fn) + # Get metric name as string + if hasattr(metric_fn, 'name'): + metric_name = metric_fn.name + else: + metric_name = metric_fn.__name__ + metric_name = metric_name_prefix + metric_name with K.name_scope(metric_name): metric_result = weighted_metric_fn( y_true, y_pred, weights=weights, mask=masks[i]) - append_metric(i, metric_name, metric_result) + + # Append to self.metrics_names, self.metric_tensors, + # self.stateful_metric_names + if len(self.output_names) > 1: + metric_name = '%s_%s' % (self.output_names[i], metric_name) + # Dedupe name + j = 1 + base_metric_name = metric_name + while metric_name in self.metrics_names: + metric_name = '%s_%d' % (base_metric_name, j) + j += 1 + self.metrics_names.append(metric_name) + self.metrics_tensors.append(metric_result) + + # Keep track of state updates created by + # stateful metrics (i.e. metrics layers). + if isinstance(metric_fn, Layer): + self.stateful_metric_names.append(metric_name) + self.metrics_updates += metric_fn.updates handle_metrics(output_metrics) handle_metrics(output_weighted_metrics, weights=weights) @@ -929,7 +539,7 @@ class Model(Network): self._feed_sample_weights = [] for i in range(len(self.sample_weights)): if i not in skip_target_weighing_indices: - self._feed_sample_weights.append(sample_weights[i]) + self._feed_sample_weights.append(self.sample_weights[i]) # Functions for train, test and predict will # be compiled lazily when required. @@ -976,9 +586,15 @@ class Model(Network): with K.name_scope('training'): with K.name_scope(self.optimizer.__class__.__name__): - training_updates = self.optimizer.get_updates( + # Training updates + updates = self.optimizer.get_updates( params=self._collected_trainable_weights, loss=self.total_loss) - updates = self.updates + training_updates + # Unconditional updates + updates += self.get_updates_for(None) + # Conditional updates relevant to this model + updates += self.get_updates_for(self._feed_inputs) + # Stateful metrics updates + updates += self.metrics_updates # Gets loss and metrics. Updates weights at each call. self.train_function = K.function( inputs, [self.total_loss] + self.metrics_tensors, @@ -999,7 +615,7 @@ class Model(Network): # Does update the network states. self.test_function = K.function( inputs, [self.total_loss] + self.metrics_tensors, - updates=self.state_updates, + updates=self.state_updates + self.metrics_updates, name='test_function', **self._function_kwargs) @@ -1021,465 +637,215 @@ class Model(Network): name='predict_function', **kwargs) - def _check_num_samples(self, - ins, - batch_size=None, - steps=None, - steps_name='steps'): - """Determine the number of samples provided for training and evaluation. - - The number of samples is not defined when running with `steps`, - in which case the number of samples is set to `None`. - - Arguments: - ins: List of tensors to be fed to the Keras function. - batch_size: Integer batch size or `None` if not defined. - steps: Total number of steps (batches of samples) - before declaring `_predict_loop` finished. - Ignored with the default value of `None`. - steps_name: The public API's parameter name for `steps`. - - Raises: - ValueError: when `steps` is `None` and the attribute `ins.shape` - does not exist. Also raises ValueError when `steps` is not `None` - and `batch_size` is not `None` because they are mutually - exclusive. - - Returns: - When steps is `None`, returns the number of samples to be - processed based on the size of the first dimension of the - first input numpy array. When steps is not `None` and - `batch_size` is `None`, returns `None`. - - Raises: - ValueError: In case of invalid arguments. - """ - if steps is not None: - num_samples = None - if batch_size is not None: - raise ValueError( - 'If ' + steps_name + ' is set, the `batch_size` must be None.') - elif ins and hasattr(ins[0], 'shape'): - num_samples = ins[0].shape[0] - else: - raise ValueError( - 'Either the input data should have ' - 'a defined shape, or ' + steps_name + ' should be specified.') - return num_samples - - def _fit_loop(self, - f, - ins, - out_labels=None, - batch_size=None, - epochs=100, - verbose=1, - callbacks=None, - val_f=None, - val_ins=None, - shuffle=True, - callback_metrics=None, - initial_epoch=0, - steps_per_epoch=None, - validation_steps=None): - """Abstract fit function for `f(ins)`. - - Assume that f returns a list, labeled by out_labels. - - Arguments: - f: Keras function returning a list of tensors - ins: List of tensors to be fed to `f` - out_labels: List of strings, display names of - the outputs of `f` - batch_size: Integer batch size or None if unknown. - epochs: Number of times to iterate over the data - verbose: Verbosity mode, 0, 1 or 2 - callbacks: List of callbacks to be called during training - val_f: Keras function to call for validation - val_ins: List of tensors to be fed to `val_f` - shuffle: Whether to shuffle the data at the beginning of each epoch - callback_metrics: List of strings, the display names of the metrics - passed to the callbacks. They should be the - concatenation of list the display names of the outputs of - `f` and the list of display names of the outputs of `f_val`. - initial_epoch: Epoch at which to start training - (useful for resuming a previous training run) - steps_per_epoch: Total number of steps (batches of samples) - before declaring one epoch finished and starting the - next epoch. Ignored with the default value of `None`. - validation_steps: Number of steps to run validation for - (only if doing validation from data tensors). - Ignored with the default value of `None`. + def _standardize_user_data(self, + x, + y=None, + sample_weight=None, + class_weight=None, + batch_size=None): + """Runs validation checks on input and target data passed by the user. + + Also standardizes the data to lists of arrays, in order. + + Also builds and compiles the model on the fly if it is a subclassed model + that has never been called before (and thus has no inputs/outputs). + + This is a purely internal method, subject to refactoring at any time. + + Args: + x: An array or list of arrays, to be used as input data. If the model + has known, named inputs, this could also be a dict mapping input names + to the corresponding array. + y: An array or list of arrays, to be used as target data. If the model + has known, named outputs, this could also be a dict mapping output names + to the corresponding array. + sample_weight: An optional sample-weight array passed by the user to + weight the importance of each sample in `x`. + class_weight: An optional class-weight array by the user to + weight the importance of samples in `x` based on the class they belong + to, as conveyed by `y`. + batch_size: Integer batch size. If provided, it is used to run additional + validation checks on stateful models. Returns: - `History` object. + A tuple of 3 lists: input arrays, target arrays, sample-weight arrays. + If the model's input and targets are symbolic, these lists are empty + (since the model takes no user-provided data, instead the data comes + from the symbolic inputs/targets). Raises: - ValueError: in case of invalid arguments. + ValueError: In case of invalid user-provided data. + RuntimeError: If the model was never compiled. """ - do_validation = False - if val_f and val_ins: - do_validation = True - if verbose and ins and hasattr(ins[0], 'shape') and hasattr( - val_ins[0], 'shape'): - print('Train on %d samples, validate on %d samples' % - (ins[0].shape[0], val_ins[0].shape[0])) - if validation_steps: - do_validation = True - if steps_per_epoch is None: - raise ValueError('Can only use `validation_steps` ' - 'when doing step-wise ' - 'training, i.e. `steps_per_epoch` ' - 'must be set.') - - num_train_samples = self._check_num_samples( - ins, batch_size, steps_per_epoch, 'steps_per_epoch') - if num_train_samples is not None: - index_array = np.arange(num_train_samples) - - self.history = cbks.History() - callbacks = [cbks.BaseLogger()] + (callbacks or []) + [self.history] - if verbose: - if steps_per_epoch is not None: - count_mode = 'steps' + # First, we build/compile the model on the fly if necessary. + all_inputs = [] + if not self.built: + # We need to use `x` to set the model inputs. + # We type-check that `x` and `y` are either single arrays + # or lists of arrays. + if isinstance(x, (list, tuple)): + if not all(isinstance(v, np.ndarray) or + tensor_util.is_tensor(v) for v in x): + raise ValueError('Please provide as model inputs either a single ' + 'array or a list of arrays. You passed: x=' + str(x)) + all_inputs += list(x) + elif isinstance(x, dict): + raise ValueError('Please do not pass a dictionary as model inputs.') else: - count_mode = 'samples' - callbacks += [cbks.ProgbarLogger(count_mode)] - callbacks = cbks.CallbackList(callbacks) - out_labels = out_labels or [] - - # it's possible to callback a different model than self - # (used by Sequential models) - if hasattr(self, 'callback_model') and self.callback_model: - callback_model = self.callback_model + if not isinstance(x, np.ndarray) and not tensor_util.is_tensor(x): + raise ValueError('Please provide as model inputs either a single ' + 'array or a list of arrays. You passed: x=' + str(x)) + all_inputs.append(x) + + # Build the model using the retrieved inputs (value or symbolic). + # If values, then in symbolic-mode placeholders will be created + # to match the value shapes. + if not self.inputs: + self._set_inputs(x) + + if y is not None: + if not self.optimizer: + raise RuntimeError('You must compile a model before ' + 'training/testing. ' + 'Use `model.compile(optimizer, loss)`.') + if not self._is_compiled: + # On-the-fly compilation of the model. + # We need to use `y` to set the model targets. + if isinstance(y, (list, tuple)): + if not all(isinstance(v, np.ndarray) or + tensor_util.is_tensor(v) for v in y): + raise ValueError('Please provide as model targets either a single ' + 'array or a list of arrays. ' + 'You passed: y=' + str(y)) + elif isinstance(y, dict): + raise ValueError('Please do not pass a dictionary as model targets.') + else: + if not isinstance(y, np.ndarray) and not tensor_util.is_tensor(y): + raise ValueError('Please provide as model targets either a single ' + 'array or a list of arrays. ' + 'You passed: y=' + str(y)) + + # Typecheck that all inputs are *either* value *or* symbolic. + # TODO(fchollet): this check could be removed in Eager mode? + if y is not None: + if isinstance(y, (list, tuple)): + all_inputs += list(y) + else: + all_inputs.append(y) + if any(tensor_util.is_tensor(v) for v in all_inputs): + if not all(tensor_util.is_tensor(v) for v in all_inputs): + raise ValueError('Do not pass inputs that mix Numpy arrays and ' + 'TensorFlow tensors. ' + 'You passed: x=' + str(x) + '; y=' + str(y)) + + if context.executing_eagerly(): + target_tensors = None + else: + # Handle target tensors if any passed. + if not isinstance(y, (list, tuple)): + y = [y] + target_tensors = [v for v in y if tensor_util.is_tensor(v)] + self.compile(optimizer=self.optimizer, + loss=self.loss, + metrics=self.metrics, + loss_weights=self.loss_weights, + target_tensors=target_tensors) + + # If `x` and `y` were all symbolic, then no model should not be fed any + # inputs and targets. + # Note: in this case, `any` and `all` are equivalent since we disallow + # mixed symbolic/value inputs. + if any(tensor_util.is_tensor(v) for v in all_inputs): + return [], [], [] + + # What follows is input validation and standardization to list format, + # in the case where all inputs are value arrays. + + if context.executing_eagerly(): + # In eager mode, do not do shape validation. + feed_input_names = self.input_names + feed_input_shapes = None + elif not self._is_graph_network: + # Case: symbolic-mode subclassed network. Do not do shape validation. + feed_input_names = self._feed_input_names + feed_input_shapes = None else: - callback_model = self - - callbacks.set_model(callback_model) - callbacks.set_params({ - 'batch_size': batch_size, - 'epochs': epochs, - 'steps': steps_per_epoch, - 'samples': num_train_samples, - 'verbose': verbose, - 'do_validation': do_validation, - 'metrics': callback_metrics or [], - }) - callbacks.on_train_begin() - callback_model.stop_training = False - for cbk in callbacks: - cbk.validation_data = val_ins - - # To prevent a slowdown, we find beforehand the arrays that need conversion. - feed = self._feed_inputs + self._feed_targets + self._feed_sample_weights - indices_for_conversion_to_dense = [] - for i in range(len(feed)): - if issparse is not None and issparse(ins[i]) and not K.is_sparse(feed[i]): - indices_for_conversion_to_dense.append(i) - - for epoch in range(initial_epoch, epochs): - callbacks.on_epoch_begin(epoch) - epoch_logs = {} - if steps_per_epoch is not None: - for step_index in range(steps_per_epoch): - batch_logs = {} - batch_logs['batch'] = step_index - batch_logs['size'] = 1 - callbacks.on_batch_begin(step_index, batch_logs) - outs = f(ins) - - if not isinstance(outs, list): - outs = [outs] - for l, o in zip(out_labels, outs): - batch_logs[l] = o - - callbacks.on_batch_end(step_index, batch_logs) - if callback_model.stop_training: - break - - if do_validation: - val_outs = self._test_loop( - val_f, - val_ins, - batch_size=batch_size, - steps=validation_steps, - verbose=0) - if not isinstance(val_outs, list): - val_outs = [val_outs] - # Same labels assumed. - for l, o in zip(out_labels, val_outs): - epoch_logs['val_' + l] = o - else: - if shuffle == 'batch': - index_array = _batch_shuffle(index_array, batch_size) - elif shuffle: - np.random.shuffle(index_array) - - batches = _make_batches(num_train_samples, batch_size) - for batch_index, (batch_start, batch_end) in enumerate(batches): - batch_ids = index_array[batch_start:batch_end] - try: - if isinstance(ins[-1], float): - # Do not slice the training phase flag. - ins_batch = _slice_arrays(ins[:-1], batch_ids) + [ins[-1]] - else: - ins_batch = _slice_arrays(ins, batch_ids) - except TypeError: - raise TypeError('TypeError while preparing batch. ' - 'If using HDF5 input data, ' - 'pass shuffle="batch".') - batch_logs = {} - batch_logs['batch'] = batch_index - batch_logs['size'] = len(batch_ids) - callbacks.on_batch_begin(batch_index, batch_logs) - for i in indices_for_conversion_to_dense: - ins_batch[i] = ins_batch[i].toarray() - - outs = f(ins_batch) - if not isinstance(outs, list): - outs = [outs] - for l, o in zip(out_labels, outs): - batch_logs[l] = o - - callbacks.on_batch_end(batch_index, batch_logs) - if callback_model.stop_training: - break - - if batch_index == len(batches) - 1: # Last batch. - if do_validation: - val_outs = self._test_loop( - val_f, val_ins, batch_size=batch_size, verbose=0) - if not isinstance(val_outs, list): - val_outs = [val_outs] - # Same labels assumed. - for l, o in zip(out_labels, val_outs): - epoch_logs['val_' + l] = o - callbacks.on_epoch_end(epoch, epoch_logs) - if callback_model.stop_training: - break - callbacks.on_train_end() - return self.history - - def _predict_loop(self, f, ins, batch_size=32, verbose=0, steps=None): - """Abstract method to loop over some data in batches. + # Case: symbolic-mode graph network. + # In this case, we run extensive shape validation checks. + feed_input_names = self._feed_input_names + feed_input_shapes = self._feed_input_shapes - Arguments: - f: Keras function returning a list of tensors. - ins: list of tensors to be fed to `f`. - batch_size: integer batch size. - verbose: verbosity mode. - steps: Total number of steps (batches of samples) - before declaring `_predict_loop` finished. - Ignored with the default value of `None`. + # Standardize the inputs. + x = training_utils.standardize_input_data( + x, + feed_input_names, + feed_input_shapes, + check_batch_axis=False, # Don't enforce the batch size. + exception_prefix='input') - Returns: - Array of predictions (if the model has a single output) - or list of arrays of predictions - (if the model has multiple outputs). - """ - num_samples = self._check_num_samples(ins, batch_size, steps, 'steps') - if verbose == 1: - if steps is not None: - progbar = Progbar(target=steps) + if y is not None: + if context.executing_eagerly(): + feed_output_names = self.output_names + feed_output_shapes = None + # Sample weighting not supported in this case. + # TODO(fchollet): consider supporting it. + feed_sample_weight_modes = [None for _ in self.outputs] + elif not self._is_graph_network: + feed_output_names = self._feed_output_names + feed_output_shapes = None + # Sample weighting not supported in this case. + # TODO(fchollet): consider supporting it. + feed_sample_weight_modes = [None for _ in self.outputs] else: - progbar = Progbar(target=num_samples) - - indices_for_conversion_to_dense = [] - for i in range(len(self._feed_inputs)): - if (issparse is not None and issparse(ins[i]) and - not K.is_sparse(self._feed_inputs[i])): - indices_for_conversion_to_dense.append(i) - - if steps is not None: - # Step-based predictions. - # Since we do not know how many samples - # we will see, we cannot pre-allocate - # the returned Numpy arrays. - # Instead, we store one array per batch seen - # and concatenate them upon returning. - unconcatenated_outs = [] - for step in range(steps): - batch_outs = f(ins) - if not isinstance(batch_outs, list): - batch_outs = [batch_outs] - if step == 0: - for batch_out in batch_outs: - unconcatenated_outs.append([]) - for i, batch_out in enumerate(batch_outs): - unconcatenated_outs[i].append(batch_out) - if verbose == 1: - progbar.update(step + 1) - if len(unconcatenated_outs) == 1: - return np.concatenate(unconcatenated_outs[0], axis=0) - return [ - np.concatenate(unconcatenated_outs[i], axis=0) - for i in range(len(unconcatenated_outs)) + feed_output_names = self._feed_output_names + feed_sample_weight_modes = self._feed_sample_weight_modes + feed_output_shapes = [] + for output_shape, loss_fn in zip(self._feed_output_shapes, + self._feed_loss_fns): + if loss_fn is losses.sparse_categorical_crossentropy: + feed_output_shapes.append(output_shape[:-1] + (1,)) + elif (not hasattr(loss_fn, '__name__') or + getattr(losses, loss_fn.__name__, None) is None): + # If `loss_fn` is not a function (e.g. callable class) + # or if it not in the `losses` module, then + # it is a user-defined loss and we make no assumptions + # about it. + feed_output_shapes.append(None) + else: + feed_output_shapes.append(output_shape) + + # Standardize the outputs. + y = training_utils.standardize_input_data( + y, + feed_output_names, + feed_output_shapes, + check_batch_axis=False, # Don't enforce the batch size. + exception_prefix='target') + + # Generate sample-wise weight values given the `sample_weight` and + # `class_weight` arguments. + sample_weights = training_utils.standardize_sample_weights( + sample_weight, feed_output_names) + class_weights = training_utils.standardize_class_weights( + class_weight, feed_output_names) + sample_weights = [ + training_utils.standardize_weights(ref, sw, cw, mode) + for (ref, sw, cw, mode) in zip(y, sample_weights, class_weights, + feed_sample_weight_modes) ] + # Check that all arrays have the same length. + training_utils.check_array_lengths(x, y, sample_weights) + if self._is_graph_network and not context.executing_eagerly(): + # Additional checks to avoid users mistakenly using improper loss fns. + training_utils.check_loss_and_target_compatibility( + y, self._feed_loss_fns, feed_output_shapes) else: - # Sample-based predictions. - outs = [] - batches = _make_batches(num_samples, batch_size) - index_array = np.arange(num_samples) - for batch_index, (batch_start, batch_end) in enumerate(batches): - batch_ids = index_array[batch_start:batch_end] - if ins and isinstance(ins[-1], float): - # Do not slice the training phase flag. - ins_batch = _slice_arrays(ins[:-1], batch_ids) + [ins[-1]] - else: - ins_batch = _slice_arrays(ins, batch_ids) - for i in indices_for_conversion_to_dense: - ins_batch[i] = ins_batch[i].toarray() - - batch_outs = f(ins_batch) - if not isinstance(batch_outs, list): - batch_outs = [batch_outs] - if batch_index == 0: - # Pre-allocate the results arrays. - for batch_out in batch_outs: - shape = (num_samples,) + batch_out.shape[1:] - outs.append(np.zeros(shape, dtype=batch_out.dtype)) - for i, batch_out in enumerate(batch_outs): - outs[i][batch_start:batch_end] = batch_out - if verbose == 1: - progbar.update(batch_end) - if len(outs) == 1: - return outs[0] - return outs - - def _test_loop(self, f, ins, batch_size=None, verbose=0, steps=None): - """Abstract method to loop over some data in batches. - - Arguments: - f: Keras function returning a list of tensors. - ins: list of tensors to be fed to `f`. - batch_size: integer batch size or `None`. - verbose: verbosity mode. - steps: Total number of steps (batches of samples) - before declaring predictions finished. - Ignored with the default value of `None`. + y = [] + sample_weights = [] - Returns: - Scalar loss (if the model has a single output and no metrics) - or list of scalars (if the model has multiple outputs - and/or metrics). The attribute `model.metrics_names` will give you - the display labels for the scalar outputs. - """ - num_samples = self._check_num_samples(ins, batch_size, steps, 'steps') - outs = [] - if verbose == 1: - if steps is not None: - progbar = Progbar(target=steps) - else: - progbar = Progbar(target=num_samples) - - # To prevent a slowdown, we find beforehand the arrays that need conversion. - feed = self._feed_inputs + self._feed_targets + self._feed_sample_weights - indices_for_conversion_to_dense = [] - for i in range(len(feed)): - if issparse is not None and issparse(ins[i]) and not K.is_sparse(feed[i]): - indices_for_conversion_to_dense.append(i) - - if steps is not None: - for step in range(steps): - batch_outs = f(ins) - if isinstance(batch_outs, list): - if step == 0: - for _ in enumerate(batch_outs): - outs.append(0.) - for i, batch_out in enumerate(batch_outs): - outs[i] += batch_out - else: - if step == 0: - outs.append(0.) - outs[0] += batch_outs - if verbose == 1: - progbar.update(step + 1) - for i in range(len(outs)): - outs[i] /= steps - else: - batches = _make_batches(num_samples, batch_size) - index_array = np.arange(num_samples) - for batch_index, (batch_start, batch_end) in enumerate(batches): - batch_ids = index_array[batch_start:batch_end] - if isinstance(ins[-1], float): - # Do not slice the training phase flag. - ins_batch = _slice_arrays(ins[:-1], batch_ids) + [ins[-1]] - else: - ins_batch = _slice_arrays(ins, batch_ids) - for i in indices_for_conversion_to_dense: - ins_batch[i] = ins_batch[i].toarray() - - batch_outs = f(ins_batch) - if isinstance(batch_outs, list): - if batch_index == 0: - for batch_out in enumerate(batch_outs): - outs.append(0.) - for i, batch_out in enumerate(batch_outs): - outs[i] += batch_out * len(batch_ids) - else: - if batch_index == 0: - outs.append(0.) - outs[0] += batch_outs * len(batch_ids) - - if verbose == 1: - progbar.update(batch_end) - for i in range(len(outs)): - outs[i] /= num_samples - if len(outs) == 1: - return outs[0] - return outs - - def _standardize_user_data(self, - x, - y, - sample_weight=None, - class_weight=None, - check_batch_axis=True, - batch_size=None): - if not hasattr(self, 'optimizer'): - raise RuntimeError('You must compile a model before ' - 'training/testing. ' - 'Use `model.compile(optimizer, loss)`.') - - output_shapes = [] - for output_shape, loss_fn in zip(self._feed_output_shapes, - self._feed_loss_fns): - if loss_fn is losses.sparse_categorical_crossentropy: - output_shapes.append(output_shape[:-1] + (1,)) - elif (not hasattr(loss_fn, '__name__') or - getattr(losses, loss_fn.__name__, None) is None): - # If `loss_fn` is not a function (e.g. callable class) - # or if it not in the `losses` module, then - # it is a user-defined loss and we make no assumptions - # about it. - output_shapes.append(None) - else: - output_shapes.append(output_shape) - x = _standardize_input_data( - x, - self._feed_input_names, - self._feed_input_shapes, - check_batch_axis=False, - exception_prefix='input') - y = _standardize_input_data( - y, - self._feed_output_names, - output_shapes, - check_batch_axis=False, - exception_prefix='target') - sample_weights = _standardize_sample_weights(sample_weight, - self._feed_output_names) - class_weights = _standardize_class_weights(class_weight, - self._feed_output_names) - sample_weights = [ - _standardize_weights(ref, sw, cw, mode) - for (ref, sw, cw, mode) in zip(y, sample_weights, class_weights, - self._feed_sample_weight_modes) - ] - _check_array_lengths(x, y, sample_weights) - _check_loss_and_target_compatibility(y, self._feed_loss_fns, - self._feed_output_shapes) if self.stateful and batch_size: + # Check that for stateful networks, number of samples is a multiple + # of the static batch size. if x[0].shape[0] % batch_size != 0: raise ValueError('In a stateful network, ' 'you should only pass inputs with ' @@ -1488,19 +854,164 @@ class Model(Network): str(x[0].shape[0]) + ' samples') return x, y, sample_weights - def _get_deduped_metrics_names(self): - out_labels = self.metrics_names + def _set_inputs(self, inputs, training=None): + """Set model's input and output specs based on the input data received. + + This is to be used for Model subclasses, which do not know at instantiation + time what their inputs look like. + + Args: + inputs: Single array, or list of arrays. The arrays could be placeholders, + Numpy arrays, or data tensors. + - if placeholders: the model is built on top of these placeholders, + and we expect Numpy data to be fed for them when calling `fit`/etc. + - if Numpy data: we create placeholders matching the shape of the Numpy + arrays. We expect Numpy data to be fed for these placeholders + when calling `fit`/etc. + - if data tensors: the model is built on top of these tensors. + We do not expect any Numpy data to be provided when calling `fit`/etc. + training: Boolean or None. Only relevant in symbolic mode. Specifies + whether to build the model's graph in inference mode (False), training + mode (True), or using the Keras learning phase (None). + """ + if context.executing_eagerly(): + self._eager_set_inputs(inputs) + else: + self._symbolic_set_inputs(inputs, training=training) + + def _set_scope(self, scope=None): + """Modify the Layer scope creation logic to create ResourceVariables.""" + super(Model, self)._set_scope(scope=scope) + # Subclassed Models create ResourceVariables by default. This makes it + # easier to use Models in an eager/graph agnostic way (since eager execution + # always uses ResourceVariables). + if not self._is_graph_network: + self._scope.set_use_resource(True) - # Rename duplicated metrics name - # (can happen with an output layer shared among multiple dataflows). - deduped_out_labels = [] - for i, label in enumerate(out_labels): - new_label = label - if out_labels.count(label) > 1: - dup_idx = out_labels[:i].count(label) - new_label += '_' + str(dup_idx + 1) - deduped_out_labels.append(new_label) - return deduped_out_labels + def _eager_set_inputs(self, inputs): + """Set model's input and output specs based on the input data received. + + This is to be used for Model subclasses, which do not know at instantiation + time what their inputs look like. + + We assume the number and ndim of outputs + does not change over different calls. + + Args: + inputs: Argument `x` (input data) passed by the user upon first model use. + + Raises: + ValueError: If the model's inputs are already set. + """ + assert context.executing_eagerly() + if self.inputs: + raise ValueError('Model inputs are already set.') + # On-the-fly setting of model inputs/outputs as DeferredTensors, + # to keep track of number of inputs and outputs and their ndim. + if isinstance(inputs, (list, tuple)): + dummy_output_values = self.call( + [ops.convert_to_tensor(v, dtype=K.floatx()) for v in inputs]) + dummy_input_values = list(inputs) + else: + dummy_output_values = self.call( + ops.convert_to_tensor(inputs, dtype=K.floatx())) + dummy_input_values = [inputs] + if isinstance(dummy_output_values, (list, tuple)): + dummy_output_values = list(dummy_output_values) + else: + dummy_output_values = [dummy_output_values] + self.outputs = [ + _DeferredTensor(shape=(None for _ in v.shape), + dtype=v.dtype) for v in dummy_output_values] + self.inputs = [ + _DeferredTensor(shape=(None for _ in v.shape), + dtype=v.dtype) for v in dummy_input_values] + self.input_names = [ + 'input_%d' % (i + 1) for i in range(len(dummy_input_values))] + self.output_names = [ + 'output_%d' % (i + 1) for i in range(len(dummy_output_values))] + self.built = True + + def _symbolic_set_inputs(self, inputs, outputs=None, training=None): + """Set model's inputs and output specs based. + + This is to be used for Model subclasses, which do not know at instantiation + time what their inputs look like. + + Args: + inputs: Argument `x` (input data) passed by the user upon first model use. + outputs: None, a data tensor, or a list of data tensors. If None, the + outputs will be determined by invoking self.call(), otherwise the + provided value will be used. + training: Boolean or None. Only relevant in symbolic mode. Specifies + whether to build the model's graph in inference mode (False), training + mode (True), or using the Keras learning phase (None). + + Raises: + ValueError: If the model's inputs are already set. + """ + assert not context.executing_eagerly() + if self.inputs: + raise ValueError('Model inputs are already set.') + + # On-the-fly setting of symbolic model inputs (either by using the tensor + # provided, or by creating a placeholder if Numpy data was provided). + self.inputs = [] + self.input_names = [] + self._feed_inputs = [] + self._feed_input_names = [] + self._feed_input_shapes = [] + if isinstance(inputs, (list, tuple)): + inputs = list(inputs) + else: + inputs = [inputs] + + for i, v in enumerate(inputs): + name = 'input_%d' % (i + 1) + self.input_names.append(name) + if isinstance(v, list): + v = np.asarray(v) + if v.ndim == 1: + v = np.expand_dims(v, 1) + if isinstance(v, (np.ndarray)): + # We fix the placeholder shape except the batch size. + # This is suboptimal, but it is the best we can do with the info + # we have. The user should call `model._set_inputs(placeholders)` + # to specify custom placeholders if the need arises. + shape = (None,) + v.shape[1:] + placeholder = K.placeholder(shape=shape, name=name) + self.inputs.append(placeholder) + self._feed_inputs.append(placeholder) + self._feed_input_names.append(name) + self._feed_input_shapes.append(shape) + else: + # Assumed tensor - TODO(fchollet) additional type check? + self.inputs.append(v) + if K.is_placeholder(v): + self._feed_inputs.append(v) + self._feed_input_names.append(name) + self._feed_input_shapes.append(K.int_shape(v)) + + if outputs is None: + # Obtain symbolic outputs by calling the model. + if len(self.inputs) == 1: + if self._expects_training_arg: + outputs = self.call(self.inputs[0], training=training) + else: + outputs = self.call(self.inputs[0]) + else: + if self._expects_training_arg: + outputs = self.call(self.inputs, training=training) + else: + outputs = self.call(self.inputs) + if isinstance(outputs, (list, tuple)): + outputs = list(outputs) + else: + outputs = [outputs] + self.outputs = outputs + self.output_names = [ + 'output_%d' % (i + 1) for i in range(len(self.outputs))] + self.built = True def fit(self, x=None, @@ -1611,6 +1122,9 @@ class Model(Network): ValueError: In case of mismatch between the provided input data and what the model expects. """ + # TODO(fchollet): this method may be creating reference cycles, which would + # lead to accumulating garbage in memory when called in a loop. Investigate. + # Backwards compatibility if batch_size is None and steps_per_epoch is None: batch_size = 32 @@ -1626,18 +1140,16 @@ class Model(Network): raise ValueError('If fitting from data tensors, ' 'you should specify the `steps_per_epoch` ' 'argument.') + # Validate user data. x, y, sample_weights = self._standardize_user_data( x, y, sample_weight=sample_weight, class_weight=class_weight, - check_batch_axis=False, batch_size=batch_size) # Prepare validation data. - do_validation = False if validation_data: - do_validation = True if len(validation_data) == 2: val_x, val_y = validation_data # pylint: disable=unpacking-non-sequence val_sample_weight = None @@ -1654,71 +1166,58 @@ class Model(Network): val_x, val_y, sample_weight=val_sample_weight, - check_batch_axis=False, batch_size=batch_size) - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - val_ins = val_x + val_y + val_sample_weights + [0.] - else: - val_ins = val_x + val_y + val_sample_weights elif validation_split and 0. < validation_split < 1.: - do_validation = True if hasattr(x[0], 'shape'): split_at = int(x[0].shape[0] * (1. - validation_split)) else: split_at = int(len(x[0]) * (1. - validation_split)) - x, val_x = (_slice_arrays(x, 0, split_at), _slice_arrays(x, split_at)) - y, val_y = (_slice_arrays(y, 0, split_at), _slice_arrays(y, split_at)) - sample_weights, val_sample_weights = (_slice_arrays( - sample_weights, 0, split_at), _slice_arrays(sample_weights, split_at)) - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - val_ins = val_x + val_y + val_sample_weights + [0.] - else: - val_ins = val_x + val_y + val_sample_weights - + x, val_x = (slice_arrays(x, 0, split_at), slice_arrays(x, split_at)) + y, val_y = (slice_arrays(y, 0, split_at), slice_arrays(y, split_at)) + sample_weights, val_sample_weights = (slice_arrays( + sample_weights, 0, split_at), slice_arrays(sample_weights, split_at)) elif validation_steps: - do_validation = True - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - val_ins = [0.] - - # Prepare input arrays and training function. - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - ins = x + y + sample_weights + [1.] + val_x = [] + val_y = [] + val_sample_weights = [] else: - ins = x + y + sample_weights - self._make_train_function() - f = self.train_function - - # Prepare display labels. - out_labels = self._get_deduped_metrics_names() - - if do_validation: - self._make_test_function() - val_f = self.test_function - callback_metrics = copy.copy(out_labels) + [ - 'val_' + n for n in out_labels - ] + val_x = None + val_y = None + val_sample_weights = None + + if context.executing_eagerly(): + return training_eager.fit_loop( + self, + inputs=x, + targets=y, + sample_weights=sample_weights, + batch_size=batch_size, + epochs=epochs, + verbose=verbose, + callbacks=callbacks, + val_inputs=val_x, + val_targets=val_y, + val_sample_weights=val_sample_weights, + shuffle=shuffle, + initial_epoch=initial_epoch, + steps_per_epoch=steps_per_epoch, + validation_steps=validation_steps) else: - callback_metrics = copy.copy(out_labels) - val_f = None - val_ins = [] - - # Delegate logic to `_fit_loop`. - return self._fit_loop( - f, - ins, - out_labels=out_labels, - batch_size=batch_size, - epochs=epochs, - verbose=verbose, - callbacks=callbacks, - val_f=val_f, - val_ins=val_ins, - shuffle=shuffle, - callback_metrics=callback_metrics, - initial_epoch=initial_epoch, - steps_per_epoch=steps_per_epoch, - validation_steps=validation_steps) + return training_arrays.fit_loop( + self, x, y, + sample_weights=sample_weights, + batch_size=batch_size, + epochs=epochs, + verbose=verbose, + callbacks=callbacks, + val_inputs=val_x, + val_targets=val_y, + val_sample_weights=val_sample_weights, + shuffle=shuffle, + initial_epoch=initial_epoch, + steps_per_epoch=steps_per_epoch, + validation_steps=validation_steps) def evaluate(self, x=None, @@ -1782,22 +1281,22 @@ class Model(Network): raise ValueError('If evaluating from data tensors, ' 'you should specify the `steps` ' 'argument.') + # Validate user data. x, y, sample_weights = self._standardize_user_data( x, y, sample_weight=sample_weight, - check_batch_axis=False, batch_size=batch_size) - # Prepare inputs, delegate logic to `_test_loop`. - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - ins = x + y + sample_weights + [0.] + + if context.executing_eagerly(): + return training_eager.test_loop( + self, inputs=x, targets=y, sample_weights=sample_weights, + batch_size=batch_size, verbose=verbose, steps=steps) else: - ins = x + y + sample_weights - self._make_test_function() - f = self.test_function - return self._test_loop( - f, ins, batch_size=batch_size, verbose=verbose, steps=steps) + return training_arrays.test_loop( + self, inputs=x, targets=y, sample_weights=sample_weights, + batch_size=batch_size, verbose=verbose, steps=steps) def predict(self, x, batch_size=None, verbose=0, steps=None): """Generates output predictions for the input samples. @@ -1829,30 +1328,14 @@ class Model(Network): raise ValueError('If predicting from data tensors, ' 'you should specify the `steps` ' 'argument.') - # Validate user data. - x = _standardize_input_data( - x, - self._feed_input_names, - self._feed_input_shapes, - check_batch_axis=False) - if self.stateful: - if x[0].shape[0] > batch_size and x[0].shape[0] % batch_size != 0: - raise ValueError('In a stateful network, ' - 'you should only pass inputs with ' - 'a number of samples that can be ' - 'divided by the batch size. Found: ' + - str(x[0].shape[0]) + ' samples. ' - 'Batch size: ' + str(batch_size) + '.') + x, _, _ = self._standardize_user_data(x) - # Prepare inputs, delegate logic to `_predict_loop`. - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - ins = x + [0.] + if context.executing_eagerly(): + return training_eager.predict_loop( + self, x, batch_size=batch_size, verbose=verbose, steps=steps) else: - ins = x - self._make_predict_function() - f = self.predict_function - return self._predict_loop( - f, ins, batch_size=batch_size, verbose=verbose, steps=steps) + return training_arrays.predict_loop( + self, x, batch_size=batch_size, verbose=verbose, steps=steps) def train_on_batch(self, x, y, sample_weight=None, class_weight=None): """Runs a single gradient update on a single batch of data. @@ -1888,19 +1371,28 @@ class Model(Network): or list of scalars (if the model has multiple outputs and/or metrics). The attribute `model.metrics_names` will give you the display labels for the scalar outputs. + + Raises: + ValueError: In case of invalid user-provided arguments. """ x, y, sample_weights = self._standardize_user_data( x, y, sample_weight=sample_weight, - class_weight=class_weight, - check_batch_axis=True) - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - ins = x + y + sample_weights + [1.] + class_weight=class_weight) + + if context.executing_eagerly(): + outputs = training_eager.train_on_batch( + self, x, y, sample_weights=sample_weights) else: - ins = x + y + sample_weights - self._make_train_function() - outputs = self.train_function(ins) + if self.uses_learning_phase and not isinstance(K.learning_phase(), int): + ins = x + y + sample_weights + [1] + else: + ins = x + y + sample_weights + + self._make_train_function() + outputs = self.train_function(ins) + if len(outputs) == 1: return outputs[0] return outputs @@ -1934,16 +1426,22 @@ class Model(Network): the display labels for the scalar outputs. Raises: - ValueError: in case of invalid arguments. + ValueError: In case of invalid user-provided arguments. """ x, y, sample_weights = self._standardize_user_data( - x, y, sample_weight=sample_weight, check_batch_axis=True) - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - ins = x + y + sample_weights + [0.] + x, y, sample_weight=sample_weight) + + if context.executing_eagerly(): + outputs = training_eager.test_on_batch( + self, x, y, sample_weights=sample_weights) else: - ins = x + y + sample_weights - self._make_test_function() - outputs = self.test_function(ins) + if self.uses_learning_phase and not isinstance(K.learning_phase(), int): + ins = x + y + sample_weights + [0] + else: + ins = x + y + sample_weights + self._make_test_function() + outputs = self.test_function(ins) + if len(outputs) == 1: return outputs[0] return outputs @@ -1956,18 +1454,26 @@ class Model(Network): Returns: Numpy array(s) of predictions. + """ - x = _standardize_input_data(x, self._feed_input_names, - self._feed_input_shapes) - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): - ins = x + [0.] - else: - ins = x - self._make_predict_function() - outputs = self.predict_function(ins) - if len(outputs) == 1: - return outputs[0] - return outputs + x, _, _ = self._standardize_user_data(x) + + if context.executing_eagerly(): + inputs = [ops.convert_to_tensor(val, dtype=K.floatx()) for val in x] + return self(inputs) # pylint: disable=not-callable + + if not context.executing_eagerly(): + if self.uses_learning_phase and not isinstance(K.learning_phase(), int): + ins = x + [0] + else: + ins = x + + self._make_predict_function() + outputs = self.predict_function(ins) + + if len(outputs) == 1: + return outputs[0] + return outputs def fit_generator(self, generator, @@ -2036,20 +1542,19 @@ class Model(Network): max_queue_size: Integer. Maximum size for the generator queue. If unspecified, `max_queue_size` will default to 10. workers: Integer. Maximum number of processes to spin up - when using process based threading. + when using process-based threading. If unspecified, `workers` will default to 1. If 0, will execute the generator on the main thread. - use_multiprocessing: Boolean. If True, use process based threading. - If unspecified, `workers` will default to False. - Note that because - this implementation relies on multiprocessing, - you should not pass - non picklable arguments to the generator - as they can't be passed - easily to children processes. - shuffle: Whether to shuffle the order of the batches at + use_multiprocessing: Boolean. + If `True`, use process-based threading. + If unspecified, `use_multiprocessing` will default to `False`. + Note that because this implementation relies on multiprocessing, + you should not pass non-picklable arguments to the generator + as they can't be passed easily to children processes. + shuffle: Boolean. Whether to shuffle the order of the batches at the beginning of each epoch. Only used with instances - of `Sequence` (keras.utils.Sequence). + of `Sequence` (`keras.utils.Sequence`). + Has no effect when `steps_per_epoch` is not `None`. initial_epoch: Epoch at which to start training (useful for resuming a previous training run) @@ -2072,218 +1577,29 @@ class Model(Network): model.fit_generator(generate_arrays_from_file('/my_file.txt'), steps_per_epoch=10000, epochs=10) ``` - Raises: ValueError: In case the generator yields data in an invalid format. """ - wait_time = 0.01 # in seconds - epoch = initial_epoch - - do_validation = bool(validation_data) - self._make_train_function() - if do_validation: - self._make_test_function() - - is_sequence = isinstance(generator, Sequence) - if not is_sequence and use_multiprocessing and workers > 1: - logging.warning( - UserWarning('Using a generator with `use_multiprocessing=True`' - ' and multiple workers may duplicate your data.' - ' Please consider using the`keras.utils.Sequence' - ' class.')) - if steps_per_epoch is None: - if is_sequence: - steps_per_epoch = len(generator) - else: - raise ValueError('`steps_per_epoch=None` is only valid for a' - ' generator based on the `keras.utils.Sequence`' - ' class. Please specify `steps_per_epoch` or use' - ' the `keras.utils.Sequence` class.') - - # python 2 has 'next', 3 has '__next__' - # avoid any explicit version checks - val_gen = ( - hasattr(validation_data, 'next') or - hasattr(validation_data, '__next__') or - isinstance(validation_data, Sequence)) - if (val_gen and not isinstance(validation_data, Sequence) and - not validation_steps): - raise ValueError('`validation_steps=None` is only valid for a' - ' generator based on the `keras.utils.Sequence`' - ' class. Please specify `validation_steps` or use' - ' the `keras.utils.Sequence` class.') - - # Prepare display labels. - out_labels = self._get_deduped_metrics_names() - callback_metrics = out_labels + ['val_' + n for n in out_labels] - - # prepare callbacks - self.history = cbks.History() - callbacks = [cbks.BaseLogger()] + (callbacks or []) + [self.history] - if verbose: - callbacks += [cbks.ProgbarLogger(count_mode='steps')] - callbacks = cbks.CallbackList(callbacks) - - # it's possible to callback a different model than self: - if hasattr(self, 'callback_model') and self.callback_model: - callback_model = self.callback_model - else: - callback_model = self - callbacks.set_model(callback_model) - callbacks.set_params({ - 'epochs': epochs, - 'steps': steps_per_epoch, - 'verbose': verbose, - 'do_validation': do_validation, - 'metrics': callback_metrics, - }) - callbacks.on_train_begin() - - enqueuer = None - val_enqueuer = None - - try: - if do_validation: - if val_gen: - if workers > 0: - if isinstance(validation_data, Sequence): - val_enqueuer = OrderedEnqueuer( - validation_data, use_multiprocessing=use_multiprocessing) - if validation_steps is None: - validation_steps = len(validation_data) - else: - val_enqueuer = GeneratorEnqueuer( - validation_data, - use_multiprocessing=use_multiprocessing, - wait_time=wait_time) - val_enqueuer.start(workers=workers, max_queue_size=max_queue_size) - validation_generator = val_enqueuer.get() - else: - validation_generator = validation_data - else: - if len(validation_data) == 2: - val_x, val_y = validation_data # pylint: disable=unpacking-non-sequence - val_sample_weight = None - elif len(validation_data) == 3: - val_x, val_y, val_sample_weight = validation_data # pylint: disable=unpacking-non-sequence - else: - raise ValueError( - '`validation_data` should be a tuple ' - '`(val_x, val_y, val_sample_weight)` ' - 'or `(val_x, val_y)`. Found: ' + str(validation_data)) - val_x, val_y, val_sample_weights = self._standardize_user_data( - val_x, val_y, val_sample_weight) - val_data = val_x + val_y + val_sample_weights - if self.uses_learning_phase and not isinstance( - K.learning_phase(), int): - val_data += [0.] - for cbk in callbacks: - cbk.validation_data = val_data - - if workers > 0: - if is_sequence: - enqueuer = OrderedEnqueuer( - generator, - use_multiprocessing=use_multiprocessing, - shuffle=shuffle) - else: - enqueuer = GeneratorEnqueuer( - generator, - use_multiprocessing=use_multiprocessing, - wait_time=wait_time) - enqueuer.start(workers=workers, max_queue_size=max_queue_size) - output_generator = enqueuer.get() - else: - output_generator = generator - - callback_model.stop_training = False - # Construct epoch logs. - epoch_logs = {} - while epoch < epochs: - callbacks.on_epoch_begin(epoch) - steps_done = 0 - batch_index = 0 - while steps_done < steps_per_epoch: - generator_output = next(output_generator) - - if not hasattr(generator_output, '__len__'): - raise ValueError('Output of generator should be ' - 'a tuple `(x, y, sample_weight)` ' - 'or `(x, y)`. Found: ' + str(generator_output)) - - if len(generator_output) == 2: - x, y = generator_output - sample_weight = None - elif len(generator_output) == 3: - x, y, sample_weight = generator_output - else: - raise ValueError('Output of generator should be ' - 'a tuple `(x, y, sample_weight)` ' - 'or `(x, y)`. Found: ' + str(generator_output)) - # build batch logs - batch_logs = {} - if isinstance(x, list): - batch_size = x[0].shape[0] - elif isinstance(x, dict): - batch_size = list(x.values())[0].shape[0] - else: - batch_size = x.shape[0] - batch_logs['batch'] = batch_index - batch_logs['size'] = batch_size - callbacks.on_batch_begin(batch_index, batch_logs) - - outs = self.train_on_batch( - x, y, sample_weight=sample_weight, class_weight=class_weight) + if not self._is_graph_network: + raise NotImplementedError( + '`fit_generator` is not yet enabled for Model subclasses') - if not isinstance(outs, list): - outs = [outs] - for l, o in zip(out_labels, outs): - batch_logs[l] = o - - callbacks.on_batch_end(batch_index, batch_logs) - - batch_index += 1 - steps_done += 1 - - # Epoch finished. - if steps_done >= steps_per_epoch and do_validation: - if val_gen: - val_outs = self.evaluate_generator( - validation_generator, validation_steps, workers=0) - else: - # No need for try/except because - # data has already been validated. - val_outs = self.evaluate( - val_x, - val_y, - batch_size=batch_size, - sample_weight=val_sample_weights, - verbose=0) - if not isinstance(val_outs, list): - val_outs = [val_outs] - # Same labels assumed. - for l, o in zip(out_labels, val_outs): - epoch_logs['val_' + l] = o - - if callback_model.stop_training: - break - - callbacks.on_epoch_end(epoch, epoch_logs) - epoch += 1 - if callback_model.stop_training: - break - - finally: - try: - if enqueuer is not None: - enqueuer.stop() - finally: - if val_enqueuer is not None: - val_enqueuer.stop() - - callbacks.on_train_end() - return self.history + return training_generator.fit_generator( + self, + generator, + steps_per_epoch=steps_per_epoch, + epochs=epochs, + verbose=verbose, + callbacks=callbacks, + validation_data=validation_data, + validation_steps=validation_steps, + class_weight=class_weight, + max_queue_size=max_queue_size, + workers=workers, + use_multiprocessing=use_multiprocessing, + shuffle=shuffle, + initial_epoch=initial_epoch) def evaluate_generator(self, generator, @@ -2308,16 +1624,15 @@ class Model(Network): the `len(generator)` as a number of steps. max_queue_size: maximum size for the generator queue workers: Integer. Maximum number of processes to spin up - when using process based threading. + when using process-based threading. If unspecified, `workers` will default to 1. If 0, will execute the generator on the main thread. - use_multiprocessing: if True, use process based threading. - Note that because - this implementation relies on multiprocessing, - you should not pass - non picklable arguments to the generator - as they can't be passed - easily to children processes. + use_multiprocessing: Boolean. + If `True`, use process-based threading. + If unspecified, `use_multiprocessing` will default to `False`. + Note that because this implementation relies on multiprocessing, + you should not pass non-picklable arguments to the generator + as they can't be passed easily to children processes. Returns: Scalar test loss (if the model has a single output and no metrics) @@ -2332,87 +1647,17 @@ class Model(Network): ValueError: In case the generator yields data in an invalid format. """ - self._make_test_function() - - steps_done = 0 - wait_time = 0.01 - all_outs = [] - batch_sizes = [] - is_sequence = isinstance(generator, Sequence) - if not is_sequence and use_multiprocessing and workers > 1: - logging.warning( - UserWarning('Using a generator with `use_multiprocessing=True`' - ' and multiple workers may duplicate your data.' - ' Please consider using the`keras.utils.Sequence' - ' class.')) - if steps is None: - if is_sequence: - steps = len(generator) - else: - raise ValueError('`steps=None` is only valid for a generator' - ' based on the `keras.utils.Sequence` class.' - ' Please specify `steps` or use the' - ' `keras.utils.Sequence` class.') - enqueuer = None - - try: - if workers > 0: - if is_sequence: - enqueuer = OrderedEnqueuer( - generator, use_multiprocessing=use_multiprocessing) - else: - enqueuer = GeneratorEnqueuer( - generator, - use_multiprocessing=use_multiprocessing, - wait_time=wait_time) - enqueuer.start(workers=workers, max_queue_size=max_queue_size) - output_generator = enqueuer.get() - else: - output_generator = generator - - while steps_done < steps: - generator_output = next(output_generator) - if not hasattr(generator_output, '__len__'): - raise ValueError('Output of generator should be a tuple ' - '(x, y, sample_weight) ' - 'or (x, y). Found: ' + str(generator_output)) - if len(generator_output) == 2: - x, y = generator_output - sample_weight = None - elif len(generator_output) == 3: - x, y, sample_weight = generator_output - else: - raise ValueError('Output of generator should be a tuple ' - '(x, y, sample_weight) ' - 'or (x, y). Found: ' + str(generator_output)) - outs = self.test_on_batch(x, y, sample_weight=sample_weight) - - if isinstance(x, list): - batch_size = x[0].shape[0] - elif isinstance(x, dict): - batch_size = list(x.values())[0].shape[0] - else: - batch_size = x.shape[0] - if batch_size == 0: - raise ValueError('Received an empty batch. ' - 'Batches should at least contain one item.') - all_outs.append(outs) - - steps_done += 1 - batch_sizes.append(batch_size) - - finally: - if enqueuer is not None: - enqueuer.stop() - - if not isinstance(outs, list): - return np.average(np.asarray(all_outs), weights=batch_sizes) - else: - averages = [] - for i in range(len(outs)): - averages.append( - np.average([out[i] for out in all_outs], weights=batch_sizes)) - return averages + if not self._is_graph_network: + raise NotImplementedError( + '`evaluate_generator` is not yet enabled for Model subclasses') + + return training_generator.evaluate_generator( + self, + generator, + steps=steps, + max_queue_size=max_queue_size, + workers=workers, + use_multiprocessing=use_multiprocessing) def predict_generator(self, generator, @@ -2437,16 +1682,15 @@ class Model(Network): the `len(generator)` as a number of steps. max_queue_size: Maximum size for the generator queue. workers: Integer. Maximum number of processes to spin up - when using process based threading. + when using process-based threading. If unspecified, `workers` will default to 1. If 0, will execute the generator on the main thread. - use_multiprocessing: If `True`, use process based threading. - Note that because - this implementation relies on multiprocessing, - you should not pass - non picklable arguments to the generator - as they can't be passed - easily to children processes. + use_multiprocessing: Boolean. + If `True`, use process-based threading. + If unspecified, `use_multiprocessing` will default to `False`. + Note that because this implementation relies on multiprocessing, + you should not pass non-picklable arguments to the generator + as they can't be passed easily to children processes. verbose: verbosity mode, 0 or 1. Returns: @@ -2456,88 +1700,15 @@ class Model(Network): ValueError: In case the generator yields data in an invalid format. """ - self._make_predict_function() - - steps_done = 0 - wait_time = 0.01 - all_outs = [] - is_sequence = isinstance(generator, Sequence) - if not is_sequence and use_multiprocessing and workers > 1: - logging.warning( - UserWarning('Using a generator with `use_multiprocessing=True`' - ' and multiple workers may duplicate your data.' - ' Please consider using the`keras.utils.Sequence' - ' class.')) - if steps is None: - if is_sequence: - steps = len(generator) - else: - raise ValueError('`steps=None` is only valid for a generator' - ' based on the `keras.utils.Sequence` class.' - ' Please specify `steps` or use the' - ' `keras.utils.Sequence` class.') - enqueuer = None - - try: - if workers > 0: - if is_sequence: - enqueuer = OrderedEnqueuer( - generator, use_multiprocessing=use_multiprocessing) - else: - enqueuer = GeneratorEnqueuer( - generator, - use_multiprocessing=use_multiprocessing, - wait_time=wait_time) - enqueuer.start(workers=workers, max_queue_size=max_queue_size) - output_generator = enqueuer.get() - else: - output_generator = generator - - if verbose == 1: - progbar = Progbar(target=steps) - - while steps_done < steps: - generator_output = next(output_generator) - if isinstance(generator_output, tuple): - # Compatibility with the generators - # used for training. - if len(generator_output) == 2: - x, _ = generator_output - elif len(generator_output) == 3: - x, _, _ = generator_output - else: - raise ValueError('Output of generator should be ' - 'a tuple `(x, y, sample_weight)` ' - 'or `(x, y)`. Found: ' + str(generator_output)) - else: - # Assumes a generator that only - # yields inputs (not targets and sample weights). - x = generator_output - - outs = self.predict_on_batch(x) - if not isinstance(outs, list): - outs = [outs] - - if not all_outs: - for out in outs: - all_outs.append([]) - - for i, out in enumerate(outs): - all_outs[i].append(out) - steps_done += 1 - if verbose == 1: - progbar.update(steps_done) - - finally: - if enqueuer is not None: - enqueuer.stop() - - if len(all_outs) == 1: - if steps_done == 1: - return all_outs[0][0] - else: - return np.concatenate(all_outs[0]) - if steps_done == 1: - return [out[0] for out in all_outs] - else: - return [np.concatenate(out) for out in all_outs] + if not self._is_graph_network: + raise NotImplementedError( + '`predict_generator` is not yet enabled for Model subclasses') + + return training_generator.predict_generator( + self, + generator, + steps=steps, + max_queue_size=max_queue_size, + workers=workers, + use_multiprocessing=use_multiprocessing, + verbose=verbose) diff --git a/tensorflow/python/keras/_impl/keras/engine/training_arrays.py b/tensorflow/python/keras/_impl/keras/engine/training_arrays.py new file mode 100644 index 0000000000000000000000000000000000000000..18116e3a14d6b1365f1a9db1a23243cd07763a62 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/training_arrays.py @@ -0,0 +1,488 @@ +# 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. +# ============================================================================== +"""Part of the Keras training engine related to plain array data. +""" +# pylint: disable=protected-access +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import copy + +import numpy as np + +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import callbacks as cbks +from tensorflow.python.keras._impl.keras.engine import training_utils +from tensorflow.python.keras._impl.keras.engine.base_layer import Layer +from tensorflow.python.keras._impl.keras.utils.generic_utils import make_batches +from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar +from tensorflow.python.keras._impl.keras.utils.generic_utils import slice_arrays + +try: + from scipy.sparse import issparse # pylint: disable=g-import-not-at-top +except ImportError: + issparse = None + + +def fit_loop(model, + inputs, + targets, + sample_weights=None, + batch_size=None, + epochs=100, + verbose=1, + callbacks=None, + val_inputs=None, + val_targets=None, + val_sample_weights=None, + shuffle=True, + callback_metrics=None, + initial_epoch=0, + steps_per_epoch=None, + validation_steps=None): + """Abstract fit function for arrays of data. + + Arguments: + model: Keras Model instance. + inputs: List of input arrays. + targets: List of target arrays. + sample_weights: Optional list of sample weight arrays. + batch_size: Integer batch size or None if unknown. + epochs: Number of times to iterate over the data + verbose: Verbosity mode, 0, 1 or 2 + callbacks: List of callbacks to be called during training + val_inputs: List of input arrays. + val_targets: List of target arrays. + val_sample_weights: Optional list of sample weight arrays. + shuffle: Whether to shuffle the data at the beginning of each epoch + callback_metrics: List of strings, the display names of the metrics + passed to the callbacks. They should be the + concatenation of list the display names of the outputs of + `f` and the list of display names of the outputs of `f_val`. + initial_epoch: Epoch at which to start training + (useful for resuming a previous training run) + steps_per_epoch: Total number of steps (batches of samples) + before declaring one epoch finished and starting the + next epoch. Ignored with the default value of `None`. + validation_steps: Number of steps to run validation for + (only if doing validation from data tensors). + Ignored with the default value of `None`. + + Returns: + `History` object. + + Raises: + ValueError: in case of invalid arguments. + """ + model._make_train_function() + f = model.train_function + + sample_weights = sample_weights or [] + val_sample_weights = val_sample_weights or [] + if model.uses_learning_phase and not isinstance(K.learning_phase(), int): + ins = inputs + targets + sample_weights + [1] + if val_inputs: + val_ins = val_inputs + val_targets + val_sample_weights + [1] + else: + ins = inputs + targets + sample_weights + if val_inputs: + val_ins = val_inputs + val_targets + val_sample_weights + if not val_inputs: + val_ins = [] + + do_validation = False + if val_inputs: + do_validation = True + if verbose and inputs and hasattr(inputs[0], 'shape') and hasattr( + val_inputs[0], 'shape'): + print('Train on %d samples, validate on %d samples' % + (inputs[0].shape[0], val_inputs[0].shape[0])) + if validation_steps: + do_validation = True + if steps_per_epoch is None: + raise ValueError('Can only use `validation_steps` ' + 'when doing step-wise ' + 'training, i.e. `steps_per_epoch` ' + 'must be set.') + + out_labels = model.metrics_names + if do_validation: + callback_metrics = copy.copy(out_labels) + [ + 'val_' + n for n in out_labels + ] + else: + callback_metrics = copy.copy(out_labels) + + num_train_samples = training_utils.check_num_samples( + ins, batch_size, steps_per_epoch, 'steps_per_epoch') + if num_train_samples is not None: + index_array = np.arange(num_train_samples) + + model.history = cbks.History() + all_callbacks = [cbks.BaseLogger( + stateful_metrics=model.stateful_metric_names)] + if verbose: + if steps_per_epoch is not None: + count_mode = 'steps' + else: + count_mode = 'samples' + all_callbacks.append( + cbks.ProgbarLogger( + count_mode, stateful_metrics=model.stateful_metric_names)) + all_callbacks += (callbacks or []) + [model.history] + callbacks = cbks.CallbackList(all_callbacks) + out_labels = out_labels or [] + + # it's possible to callback a different model than self + # (used by Sequential models) + if hasattr(model, 'callback_model') and model.callback_model: + callback_model = model.callback_model + else: + callback_model = model + + callbacks.set_model(callback_model) + + callbacks.set_params({ + 'batch_size': batch_size, + 'epochs': epochs, + 'steps': steps_per_epoch, + 'samples': num_train_samples, + 'verbose': verbose, + 'do_validation': do_validation, + 'metrics': callback_metrics or [], + }) + callbacks.on_train_begin() + callback_model.stop_training = False + for cbk in callbacks: + cbk.validation_data = val_ins + + # To prevent a slowdown, we find beforehand the arrays that need conversion. + feed = model._feed_inputs + model._feed_targets + model._feed_sample_weights + indices_for_conversion_to_dense = [] + for i in range(len(feed)): + if issparse is not None and issparse(ins[i]) and not K.is_sparse(feed[i]): + indices_for_conversion_to_dense.append(i) + + for epoch in range(initial_epoch, epochs): + # Reset stateful metrics + for m in model.metrics: + if isinstance(m, Layer): + m.reset_states() + # Update callbacks + callbacks.on_epoch_begin(epoch) + epoch_logs = {} + if steps_per_epoch is not None: + for step_index in range(steps_per_epoch): + batch_logs = {} + batch_logs['batch'] = step_index + batch_logs['size'] = 1 + callbacks.on_batch_begin(step_index, batch_logs) + outs = f(ins) + + if not isinstance(outs, list): + outs = [outs] + for l, o in zip(out_labels, outs): + batch_logs[l] = o + + callbacks.on_batch_end(step_index, batch_logs) + if callback_model.stop_training: + break + + if do_validation: + val_outs = test_loop( + model, + val_inputs, + val_targets, + sample_weights=val_sample_weights, + batch_size=batch_size, + steps=validation_steps, + verbose=0) + if not isinstance(val_outs, list): + val_outs = [val_outs] + # Same labels assumed. + for l, o in zip(out_labels, val_outs): + epoch_logs['val_' + l] = o + else: + if shuffle == 'batch': + index_array = training_utils.batch_shuffle(index_array, batch_size) + elif shuffle: + np.random.shuffle(index_array) + + batches = make_batches(num_train_samples, batch_size) + + for batch_index, (batch_start, batch_end) in enumerate(batches): + batch_ids = index_array[batch_start:batch_end] + try: + if isinstance(ins[-1], int): + # Do not slice the training phase flag. + ins_batch = slice_arrays(ins[:-1], batch_ids) + [ins[-1]] + else: + ins_batch = slice_arrays(ins, batch_ids) + except TypeError: + raise TypeError('TypeError while preparing batch. ' + 'If using HDF5 input data, ' + 'pass shuffle="batch".') + batch_logs = {} + batch_logs['batch'] = batch_index + batch_logs['size'] = len(batch_ids) + callbacks.on_batch_begin(batch_index, batch_logs) + for i in indices_for_conversion_to_dense: + ins_batch[i] = ins_batch[i].toarray() + + outs = f(ins_batch) + if not isinstance(outs, list): + outs = [outs] + for l, o in zip(out_labels, outs): + batch_logs[l] = o + + callbacks.on_batch_end(batch_index, batch_logs) + if callback_model.stop_training: + break + + if batch_index == len(batches) - 1: # Last batch. + if do_validation: + val_outs = test_loop( + model, + val_inputs, + val_targets, + sample_weights=val_sample_weights, + batch_size=batch_size, + verbose=0) + if not isinstance(val_outs, list): + val_outs = [val_outs] + # Same labels assumed. + for l, o in zip(out_labels, val_outs): + epoch_logs['val_' + l] = o + callbacks.on_epoch_end(epoch, epoch_logs) + if callback_model.stop_training: + break + callbacks.on_train_end() + return model.history + + +def predict_loop(model, inputs, batch_size=32, verbose=0, steps=None): + """Abstract method to loop over some data in batches. + + Arguments: + model: Keras Model instance. + inputs: list of tensors to be fed to `f`. + batch_size: integer batch size. + verbose: verbosity mode. + steps: Total number of steps (batches of samples) + before declaring `_predict_loop` finished. + Ignored with the default value of `None`. + + Returns: + Array of predictions (if the model has a single output) + or list of arrays of predictions + (if the model has multiple outputs). + """ + model._make_predict_function() + f = model.predict_function + + if model.uses_learning_phase and not isinstance(K.learning_phase(), int): + ins = inputs + [0] + else: + ins = inputs + + num_samples = training_utils.check_num_samples( + inputs, batch_size, steps, 'steps') + if verbose == 1: + if steps is not None: + progbar = Progbar(target=steps) + else: + progbar = Progbar(target=num_samples) + + indices_for_conversion_to_dense = [] + for i in range(len(model._feed_inputs)): + if (issparse is not None and issparse(inputs[i]) and + not K.is_sparse(model._feed_inputs[i])): + indices_for_conversion_to_dense.append(i) + + if steps is not None: + # Step-based predictions. + # Since we do not know how many samples + # we will see, we cannot pre-allocate + # the returned Numpy arrays. + # Instead, we store one array per batch seen + # and concatenate them upon returning. + unconcatenated_outs = [] + for step in range(steps): + batch_outs = f(ins) + if not isinstance(batch_outs, list): + batch_outs = [batch_outs] + if step == 0: + for batch_out in batch_outs: + unconcatenated_outs.append([]) + for i, batch_out in enumerate(batch_outs): + unconcatenated_outs[i].append(batch_out) + if verbose == 1: + progbar.update(step + 1) + if len(unconcatenated_outs) == 1: + return np.concatenate(unconcatenated_outs[0], axis=0) + return [ + np.concatenate(unconcatenated_outs[i], axis=0) + for i in range(len(unconcatenated_outs)) + ] + else: + # Sample-based predictions. + outs = [] + batches = make_batches(num_samples, batch_size) + index_array = np.arange(num_samples) + for batch_index, (batch_start, batch_end) in enumerate(batches): + batch_ids = index_array[batch_start:batch_end] + if ins and isinstance(ins[-1], int): + # Do not slice the training phase flag. + ins_batch = slice_arrays(ins[:-1], batch_ids) + [ins[-1]] + else: + ins_batch = slice_arrays(ins, batch_ids) + for i in indices_for_conversion_to_dense: + ins_batch[i] = ins_batch[i].toarray() + + batch_outs = f(ins_batch) + if not isinstance(batch_outs, list): + batch_outs = [batch_outs] + if batch_index == 0: + # Pre-allocate the results arrays. + for batch_out in batch_outs: + shape = (num_samples,) + batch_out.shape[1:] + outs.append(np.zeros(shape, dtype=batch_out.dtype)) + for i, batch_out in enumerate(batch_outs): + outs[i][batch_start:batch_end] = batch_out + if verbose == 1: + progbar.update(batch_end) + if len(outs) == 1: + return outs[0] + return outs + + +def test_loop(model, inputs, targets, + sample_weights=None, + batch_size=None, + verbose=0, + steps=None): + """Abstract method to loop over some data in batches. + + Arguments: + model: Keras Model instance. + inputs: List of input arrays. + targets: List of target arrays. + sample_weights: Optional list of sample weight arrays. + batch_size: integer batch size or `None`. + verbose: verbosity mode. + steps: Total number of steps (batches of samples) + before declaring predictions finished. + Ignored with the default value of `None`. + + Returns: + Scalar loss (if the model has a single output and no metrics) + or list of scalars (if the model has multiple outputs + and/or metrics). The attribute `model.metrics_names` will give you + the display labels for the scalar outputs. + """ + model._make_test_function() + f = model.test_function + + sample_weights = sample_weights or [] + if model.uses_learning_phase and not isinstance(K.learning_phase(), int): + ins = inputs + targets + sample_weights + [0] + else: + ins = inputs + targets + sample_weights + + if hasattr(model, 'metrics'): + for m in model.metrics: + if isinstance(m, Layer): + m.reset_states() + stateful_metric_indices = [ + i for i, name in enumerate(model.metrics_names) + if str(name) in model.stateful_metric_names + ] + else: + stateful_metric_indices = [] + + num_samples = training_utils.check_num_samples( + ins, batch_size, steps, 'steps') + outs = [] + if verbose == 1: + if steps is not None: + progbar = Progbar(target=steps) + else: + progbar = Progbar(target=num_samples) + + # To prevent a slowdown, we find beforehand the arrays that need conversion. + feed = model._feed_inputs + model._feed_targets + model._feed_sample_weights + indices_for_conversion_to_dense = [] + for i in range(len(feed)): + if issparse is not None and issparse(ins[i]) and not K.is_sparse(feed[i]): + indices_for_conversion_to_dense.append(i) + + if steps is not None: + for step in range(steps): + batch_outs = f(ins) + if isinstance(batch_outs, list): + if step == 0: + for _ in enumerate(batch_outs): + outs.append(0.) + for i, batch_out in enumerate(batch_outs): + if i in stateful_metric_indices: + outs[i] = batch_out + else: + outs[i] += batch_out + else: + if step == 0: + outs.append(0.) + outs[0] += batch_outs + if verbose == 1: + progbar.update(step + 1) + for i in range(len(outs)): + if i not in stateful_metric_indices: + outs[i] /= steps + else: + batches = make_batches(num_samples, batch_size) + index_array = np.arange(num_samples) + for batch_index, (batch_start, batch_end) in enumerate(batches): + batch_ids = index_array[batch_start:batch_end] + if isinstance(ins[-1], int): + # Do not slice the training phase flag. + ins_batch = slice_arrays(ins[:-1], batch_ids) + [ins[-1]] + else: + ins_batch = slice_arrays(ins, batch_ids) + for i in indices_for_conversion_to_dense: + ins_batch[i] = ins_batch[i].toarray() + + batch_outs = f(ins_batch) + + if isinstance(batch_outs, list): + if batch_index == 0: + for batch_out in enumerate(batch_outs): + outs.append(0.) + for i, batch_out in enumerate(batch_outs): + if i in stateful_metric_indices: + outs[i] = batch_out + else: + outs[i] += batch_out * len(batch_ids) + else: + if batch_index == 0: + outs.append(0.) + outs[0] += batch_outs * len(batch_ids) + if verbose == 1: + progbar.update(batch_end) + for i in range(len(outs)): + if i not in stateful_metric_indices: + outs[i] /= num_samples + if len(outs) == 1: + return outs[0] + return outs diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py new file mode 100644 index 0000000000000000000000000000000000000000..67858a578c5c95b3099e1e6713f3287748fc861f --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py @@ -0,0 +1,664 @@ +# 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. +# ============================================================================== +"""Keras training and evaluation routines for eager execution. +""" +# pylint: disable=protected-access +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import copy + +import numpy as np + +from tensorflow.python.eager.backprop import GradientTape +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util +from tensorflow.python.keras._impl.keras import backend +from tensorflow.python.keras._impl.keras import callbacks as cbks +from tensorflow.python.keras._impl.keras import losses +from tensorflow.python.keras._impl.keras import metrics as metrics_module +from tensorflow.python.keras._impl.keras.engine import training_utils +from tensorflow.python.keras._impl.keras.utils.generic_utils import make_batches +from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar +from tensorflow.python.keras._impl.keras.utils.generic_utils import slice_arrays +from tensorflow.python.platform import tf_logging as logging + + +def _get_metrics_info(metric, internal_output_shapes=None, loss_func=None): + if metric == 'accuracy' or metric == 'acc': + # custom handling of accuracy + # (because of class mode duality) + output_shape = internal_output_shapes + if output_shape[-1] == 1 or loss_func == losses.binary_crossentropy: + # case: binary accuracy + acc_fn = metrics_module.binary_accuracy + elif loss_func == losses.sparse_categorical_crossentropy: + # case: categorical accuracy with sparse targets + acc_fn = metrics_module.sparse_categorical_accuracy + else: + acc_fn = metrics_module.categorical_accuracy + + metric_name = 'acc' + return metric_name, acc_fn + else: + metric_fn = metrics_module.get(metric) + metric_name = metric_fn.__name__ + return metric_name, metric_fn + + +def _eager_loss_fn(outputs, targets, loss_fn, output_name): + with backend.name_scope(output_name + '_loss'): + loss = loss_fn(targets, outputs) + return loss + + +def _eager_metrics_fn(model, outputs, targets): + """Calculates the metrics for each output of the given model. + + Arguments: + model: The model on which metrics are being calculated. + outputs: The outputs of the given model. + targets: The predictions or targets of the given model. + + Returns: + Returns the metric names and metric results for each output of the model. + """ + metric_names = [] + metric_results = [] + if not isinstance(outputs, list): + outputs = [outputs] + + if not isinstance(targets, list): + targets = [targets] + + for i in range(len(model.outputs)): + output_metrics = model.nested_metrics[i] + for nested_output_metric in output_metrics: + metric_name, metric_fn = _get_metrics_info( + nested_output_metric, backend.int_shape(model.outputs[i]), + model.loss_functions[i]) + + if len(model.output_names) > 1: + metric_name = model.output_names[i] + '_' + metric_name + if metric_name not in model.metrics_names: + model.metrics_names.append(metric_name) + + with backend.name_scope(metric_name): + metric_result = metric_fn(outputs[i], targets[i]) + metric_names.append(metric_name) + metric_results.append(backend.mean(metric_result)) + + return metric_names, metric_results + + +def _model_loss(model, inputs, targets, sample_weights=None, training=False): + """Calculates the loss for a given model. + + Arguments: + model: The model on which metrics are being calculated. + inputs: List of input arrays. + targets: List of target arrays. + sample_weights: Optional list of sample weight arrays. + training: Whether the model should be run in inference or training mode. + + Returns: + Returns the model output, total loss and loss value calculated using the + specified loss function. The total loss includes regularization losses and + applies masking and sample weighting to the loss value. + """ + total_loss = 0 + if len(inputs) == 1: + if model._expects_training_arg: + outs = model.call(inputs[0], training=training) + else: + outs = model.call(inputs[0]) + else: + if model._expects_training_arg: + outs = model.call(inputs, training=training) + else: + outs = model.call(inputs) + if not isinstance(outs, list): + outs = [outs] + + if not isinstance(targets, list): + targets = [targets] + + loss_metrics = [] + with backend.name_scope('loss'): + for i, loss_fn in enumerate(model.loss_functions): + if sample_weights: + weights = sample_weights[i] + else: + weights = None + + # TODO(fchollet): support masking; in practice `_keras_mask` is never + # set in this context currently. + mask = outs[i]._keras_mask + + weighted_masked_fn = training_utils.weighted_masked_objective(loss_fn) + with backend.name_scope(model.output_names[i] + '_loss'): + output_loss = weighted_masked_fn( + outs[i], targets[i], weights, mask=mask) + loss_metrics.append(backend.mean(output_loss)) + + loss_weight = model.loss_weights_list[i] + if total_loss is None: + total_loss = loss_weight * output_loss + else: + total_loss += loss_weight * output_loss + + total_loss = backend.mean(total_loss) + # Add regularization losses + custom_losses = [] + for layer in model.layers: + if layer.losses: + custom_losses += layer.losses + + if custom_losses: + total_loss += sum(custom_losses) + + return outs, total_loss, loss_metrics + + +def _process_single_batch(model, + inputs, + targets, + sample_weights=None, + training=False): + """Calculate the loss and gradient for one input batch. + + The model weights are updated if training is set to True. + + Arguments: + model: Model whose loss has to be calculated. + inputs: List of input arrays. + targets: List of target arrays. + sample_weights: Optional list of sample weight arrays. + training: The boolean represents if the weights of the model are updated. + 'fit' methods will set this to True while 'evaluate' methods will + set this to False. + + Returns: + output of the model, total loss and the loss associated with each output. + + Raises: + ValueError: If the model has no loss to optimize. + """ + with backend.learning_phase_scope(1 if training else 0): + with GradientTape() as tape: + outs, loss, loss_metrics = _model_loss(model, inputs, targets, + sample_weights=sample_weights, + training=training) + if loss is None: + raise ValueError('The model cannot be run ' + 'because it has no loss to optimize.') + if training: + if not model._collected_trainable_weights: + logging.warning('The list of trainable weights is empty. Make sure that' + ' you are not setting model.trainable to False before ' + 'compiling the model.') + else: + grads = tape.gradient(loss, model._collected_trainable_weights) + model.optimizer.apply_gradients(zip(grads, + model._collected_trainable_weights)) + return outs, loss, loss_metrics + + +def train_on_batch(model, inputs, targets, sample_weights=None): + """Calculates the loss and gradient updates for one input batch. + + Arguments: + model: Model whose loss has to be calculated. + inputs: Input batch data. + targets: Target batch data. + sample_weights: Sample weight batch data. + + Returns: + total loss and the loss associated with each output. + """ + inputs = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) for val in inputs] + targets = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) for val in targets] + sample_weights = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + if val is not None else None for val in sample_weights] + outs, loss, _ = _process_single_batch( + model, inputs, targets, sample_weights=sample_weights, training=True) + if not isinstance(outs, list): + outs = [outs] + _, metrics_results = _eager_metrics_fn( + model, outs, targets) + if not isinstance(loss, list): + loss = [loss] + return loss + metrics_results + + +def test_on_batch(model, inputs, targets, sample_weights=None): + """Calculates the loss for one input batch. + + Arguments: + model: Model whose loss has to be calculated. + inputs: Input batch data. + targets: Target batch data. + sample_weights: Sample weight batch data. + + Returns: + total loss, loss and metrics associated with each output. + """ + inputs = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) for val in inputs] + targets = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) for val in targets] + sample_weights = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + if val is not None else None for val in sample_weights] + outs, loss, loss_metrics = _process_single_batch( + model, inputs, targets, sample_weights=sample_weights, training=False) + if not isinstance(outs, list): + outs = [outs] + metric_names, metrics_results = _eager_metrics_fn( + model, outs, targets) + model.metrics_names.append(metric_names) + if not isinstance(loss, list): + loss = [loss] + return loss + loss_metrics + metrics_results + + +def fit_loop( + model, + inputs, + targets, + sample_weights=None, + val_inputs=None, + val_targets=None, + val_sample_weights=None, + batch_size=None, + epochs=100, + verbose=1, + callbacks=None, + shuffle=True, + callback_metrics=None, + initial_epoch=0, + steps_per_epoch=None, + validation_steps=None): + """Abstract fit function for eager execution. + + Arguments: + model: Instance of the model that is being executed in Eager mode. + inputs: List of input arrays. + targets: List of target arrays. + sample_weights: Optional list of sample weight arrays. + val_inputs: Input data for validation. + val_targets: Target data for validation. + val_sample_weights: Sample weight data for validation. + batch_size: Integer batch size or None if unknown. + epochs: Number of times to iterate over the data + verbose: Verbosity mode, 0, 1 or 2 + callbacks: List of callbacks to be called during training + shuffle: Whether to shuffle the data at the beginning of each epoch + callback_metrics: List of strings, the display names of the metrics + passed to the callbacks. They should be the + concatenation of list the display names of the outputs of + `f` and the list of display names of the outputs of `f_val`. + initial_epoch: Epoch at which to start training + (useful for resuming a previous training run) + steps_per_epoch: Total number of steps (batches of samples) + before declaring one epoch finished and starting the + next epoch. Ignored with the default value of `None`. + validation_steps: Number of steps to run validation for (only if doing + validation from data tensors). Ignored with default value of `None`. + + Returns: + `History` object. + + Raises: + ValueError: In case of invalid argument values. + """ + # Required for Eager mode + with backend.learning_phase_scope(1): + do_validation = False + if val_inputs: + do_validation = True + if (verbose and inputs and hasattr(inputs[0], 'shape') and + hasattr(val_inputs[0], 'shape')): + print('Train on %d samples, validate on %d samples' % + (inputs[0].shape[0], val_inputs[0].shape[0])) + if validation_steps: + if steps_per_epoch is None: + raise ValueError('Can only use `validation_steps` when doing step-wise ' + 'training, i.e. `steps_per_epoch` must be set.') + do_validation = True + + out_labels = model.metrics_names + if do_validation: + callback_metrics = copy.copy(out_labels) + [ + 'val_' + n for n in out_labels + ] + else: + callback_metrics = copy.copy(out_labels) + + if sample_weights: + feed_data = inputs + targets + sample_weights + else: + feed_data = inputs + targets + num_train_samples = training_utils.check_num_samples( + feed_data, + batch_size=batch_size, + steps=steps_per_epoch, + steps_name='steps_per_epoch') + + if num_train_samples is not None: + index_array = np.arange(num_train_samples) + + model.history = cbks.History() + callbacks = [cbks.BaseLogger()] + (callbacks or []) + [model.history] + if verbose: + if steps_per_epoch is not None: + count_mode = 'steps' + else: + count_mode = 'samples' + callbacks += [cbks.ProgbarLogger(count_mode)] + callbacks = cbks.CallbackList(callbacks) + + # it's possible to callback a different model than self + # (used by Sequential models) + if hasattr(model, 'callback_model') and model.callback_model: + callback_model = model.callback_model + else: + callback_model = model + + callbacks.set_model(callback_model) + + callbacks.set_params({ + 'batch_size': batch_size, + 'epochs': epochs, + 'steps': steps_per_epoch, + 'samples': num_train_samples, + 'verbose': verbose, + 'do_validation': do_validation, + 'metrics': callback_metrics or [], + }) + callbacks.on_train_begin() + callback_model.stop_training = False + for cbk in callbacks: + if not val_inputs: + cbk.validation_data = [] + elif val_sample_weights: + cbk.validation_data = val_inputs + val_targets + val_sample_weights + else: + cbk.validation_data = val_inputs + val_targets + + for epoch in range(initial_epoch, epochs): + callbacks.on_epoch_begin(epoch) + epoch_logs = {} + if shuffle == 'batch': + index_array = model._batch_shuffle(index_array, batch_size) + elif shuffle: + np.random.shuffle(index_array) + + batches = make_batches(num_train_samples, batch_size) + + for batch_index, (batch_start, batch_end) in enumerate(batches): + batch_ids = index_array[batch_start:batch_end] + try: + inputs_batch = slice_arrays(inputs, batch_ids) + targets_batch = slice_arrays(targets, batch_ids) + if sample_weights: + sample_weights_batch = slice_arrays(sample_weights, batch_ids) + else: + sample_weights_batch = None + except TypeError: + raise TypeError('TypeError while preparing batch. ' + 'If using HDF5 input data, ' + 'pass shuffle="batch".') + batch_logs = {} + batch_logs['batch'] = batch_index + batch_logs['size'] = len(batch_ids) + + callbacks.on_batch_begin(batch_index, batch_logs) + + inputs_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + for val in inputs_batch] + targets_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + for val in targets_batch] + if sample_weights: + sample_weights_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + if val is not None else None + for val in sample_weights_batch] + + outs, loss, loss_metrics = _process_single_batch( + model, + inputs_batch, + targets_batch, + sample_weights=sample_weights_batch, + training=True) + + if not isinstance(outs, list): + outs = [outs] + + for l, o in zip(out_labels, outs): + batch_logs[l] = o + # Required for Eager mode + metrics_names, metrics_results = _eager_metrics_fn( + model, outs, targets_batch) + batch_logs['loss'] = tensor_util.constant_value(backend.mean(loss)) + + # TODO(anjalisridhar): Move this to compile to avoid duplicate code. + # In graph mode we set the metric names in compile. However in + # Eager mode we calculate the metrics for each batch in fit_loop. + # We could calculate the metric names and functions in compile. + # This would avoid setting the callback parameters separately. + # We need to do this for the first iteration alone + for m in metrics_names: + if m not in callback_metrics: + callback_metrics.append(m) + + callbacks.set_params({ + 'batch_size': batch_size, + 'epochs': epochs, + 'steps': steps_per_epoch, + 'samples': num_train_samples, + 'verbose': verbose, + 'do_validation': do_validation, + 'metrics': callback_metrics or [], + }) + + for k, v in zip(model.metrics_names, + [backend.mean(loss)] + loss_metrics + metrics_results): + batch_logs[k] = tensor_util.constant_value(v) + + callbacks.on_batch_end(batch_index, batch_logs) + if callback_model.stop_training: + break + + if batch_index == len(batches) - 1: # Last batch. + if do_validation: + val_outs = test_loop( + model, val_inputs, val_targets, + sample_weights=val_sample_weights, + batch_size=batch_size, + verbose=0) + if not isinstance(val_outs, list): + val_outs = [val_outs] + # Same labels assumed. + for l, o in zip(out_labels, val_outs): + epoch_logs['val_' + l] = o + callbacks.on_epoch_end(epoch, epoch_logs) + if callback_model.stop_training: + break + callbacks.on_train_end() + return model.history + + +def test_loop(model, inputs, targets, + sample_weights=None, + batch_size=None, + verbose=0, + steps=None): + """Abstract method to loop over some data in batches. + + Arguments: + model: Model instance that is being evaluated in Eager mode. + inputs: List of input arrays. + targets: List of target arrays. + sample_weights: Optional list of sample weight arrays. + batch_size: integer batch size or `None`. + verbose: verbosity mode. + steps: Total number of steps (batches of samples) + before declaring predictions finished. + Ignored with the default value of `None`. + + Returns: + Scalar loss (if the model has a single output and no metrics) + or list of scalars (if the model has multiple outputs + and/or metrics). The attribute `model.metrics_names` will give you + the display labels for the scalar outputs. + """ + with backend.learning_phase_scope(0): + feed_data = inputs + targets + if sample_weights: + feed_data += sample_weights + num_samples = training_utils.check_num_samples( + feed_data, batch_size=batch_size, steps=steps, steps_name='steps') + outs = [] + if verbose == 1: + progbar = Progbar(target=num_samples) + batches = make_batches(num_samples, batch_size) + index_array = np.arange(num_samples) + for batch_index, (batch_start, batch_end) in enumerate(batches): + batch_ids = index_array[batch_start:batch_end] + inputs_batch = slice_arrays(inputs, batch_ids) + targets_batch = slice_arrays(targets, batch_ids) + if sample_weights: + sample_weights_batch = slice_arrays(sample_weights, batch_ids) + else: + sample_weights_batch = None + + inputs_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + for val in inputs_batch] + targets_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + for val in targets_batch] + if sample_weights: + sample_weights_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + if val is not None else None + for val in sample_weights_batch] + + loss_outs, loss, loss_metrics = _model_loss( + model, + inputs_batch, + targets_batch, + sample_weights=sample_weights_batch, + training=False) + _, metrics_results = _eager_metrics_fn(model, loss_outs, targets_batch) + batch_outs = [] + for _, v in zip(model.metrics_names, + [backend.mean(loss)] + loss_metrics + metrics_results): + batch_outs.append(tensor_util.constant_value(v)) + + if isinstance(batch_outs, list): + if batch_index == 0: + for batch_out in enumerate(batch_outs): + outs.append(0.) + for i, batch_out in enumerate(batch_outs): + outs[i] += batch_out * len(batch_ids) + else: + if batch_index == 0: + outs.append(0.) + outs[0] += batch_outs * len(batch_ids) + + if verbose == 1: + progbar.update(batch_end) + for i in range(len(outs)): + outs[i] /= num_samples + if len(outs) == 1: + return outs[0] + return outs + + +def predict_loop(model, inputs, + batch_size=32, + verbose=0, + steps=None): + """Abstract method to loop over some data in batches. + + Arguments: + model: + inputs: List of input arrays. + batch_size: integer batch size. + verbose: verbosity mode. + steps: Total number of steps (batches of samples) + before declaring `_predict_loop` finished. + Ignored with the default value of `None`. + + Returns: + Array of predictions (if the model has a single output) + or list of arrays of predictions + (if the model has multiple outputs). + """ + with backend.learning_phase_scope(0): + num_samples = training_utils.check_num_samples( + inputs, batch_size, steps, 'steps') + if verbose == 1: + if steps is not None: + progbar = Progbar(target=steps) + else: + progbar = Progbar(target=num_samples) + + outs = [] + batches = make_batches(num_samples, batch_size) + index_array = np.arange(num_samples) + for batch_index, (batch_start, batch_end) in enumerate(batches): + batch_ids = index_array[batch_start:batch_end] + inputs_batch = slice_arrays(inputs, batch_ids) + + inputs_batch = [ + ops.convert_to_tensor(val, dtype=backend.floatx()) + for val in inputs_batch] + + if len(inputs_batch) == 1: + if model._expects_training_arg: + batch_outs = model.call(inputs_batch[0], training=False) + else: + batch_outs = model.call(inputs_batch[0]) + else: + if model._expects_training_arg: + batch_outs = model.call(inputs_batch, training=False) + else: + batch_outs = model.call(inputs_batch) + + if not isinstance(batch_outs, list): + batch_outs = [batch_outs] + if batch_index == 0: + # Pre-allocate the results arrays. + for batch_out in batch_outs: + dims = batch_out.shape[1:].dims + dims_list = [d.value for d in dims] + shape = (num_samples,) + tuple(dims_list) + outs.append(np.zeros(shape, dtype=batch_out.dtype.as_numpy_dtype)) + for i, batch_out in enumerate(batch_outs): + outs[i][batch_start:batch_end] = batch_out + if verbose == 1: + progbar.update(batch_end) + if len(outs) == 1: + return outs[0] + return outs diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py new file mode 100644 index 0000000000000000000000000000000000000000..8848b393d5e602e564cb357c32a937eaabd68203 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py @@ -0,0 +1,546 @@ +# 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 training routines.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import numpy as np + +from tensorflow.python.framework import ops +from tensorflow.python.keras._impl import keras +from tensorflow.python.keras._impl.keras import testing_utils +from tensorflow.python.platform import test +from tensorflow.python.training.rmsprop import RMSPropOptimizer + + +class TrainingTest(test.TestCase): + + def test_fit_on_arrays(self): + a = keras.layers.Input(shape=(3,), name='input_a') + b = keras.layers.Input(shape=(3,), name='input_b') + + dense = keras.layers.Dense(4, name='dense') + c = dense(a) + d = dense(b) + e = keras.layers.Dropout(0.5, name='dropout')(c) + + model = keras.models.Model([a, b], [d, e]) + + optimizer = RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + loss_weights = [1., 0.5] + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics, loss_weights=loss_weights) + + input_a_np = np.random.random((10, 3)) + input_b_np = np.random.random((10, 3)) + + output_d_np = np.random.random((10, 4)) + output_e_np = np.random.random((10, 4)) + + # Test fit at different verbosity + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=1, + batch_size=5, + verbose=0) + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=1, + batch_size=5, + verbose=1) + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=2, + batch_size=5, + verbose=2) + + # Test with validation data + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + validation_data=([input_a_np, input_b_np], [output_d_np, + output_e_np]), + epochs=1, + batch_size=5, + verbose=0) + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + validation_data=([input_a_np, input_b_np], [output_d_np, + output_e_np]), + epochs=2, + batch_size=5, + verbose=1) + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + validation_data=([input_a_np, input_b_np], [output_d_np, + output_e_np]), + epochs=2, + batch_size=5, + verbose=2) + model.train_on_batch([input_a_np, input_b_np], [output_d_np, output_e_np]) + + # Test with validation split + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=2, + batch_size=5, + verbose=0, + validation_split=0.2) + + # Test with dictionary inputs + model.fit( + { + 'input_a': input_a_np, + 'input_b': input_b_np + }, {'dense': output_d_np, + 'dropout': output_e_np}, + epochs=1, + batch_size=5, + verbose=0) + model.fit( + { + 'input_a': input_a_np, + 'input_b': input_b_np + }, {'dense': output_d_np, + 'dropout': output_e_np}, + epochs=1, + batch_size=5, + verbose=1) + model.fit( + { + 'input_a': input_a_np, + 'input_b': input_b_np + }, {'dense': output_d_np, + 'dropout': output_e_np}, + validation_data=({'input_a': input_a_np, + 'input_b': input_b_np + }, + { + 'dense': output_d_np, + 'dropout': output_e_np + }), + epochs=1, + batch_size=5, + verbose=0) + model.train_on_batch({ + 'input_a': input_a_np, + 'input_b': input_b_np + }, {'dense': output_d_np, + 'dropout': output_e_np}) + # Test with lists for loss, metrics + loss = ['mae', 'mse'] + metrics = ['acc', 'mae'] + model.compile(optimizer, loss, metrics=metrics) + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=1, + batch_size=5, + verbose=0) + + # Test with dictionaries for loss, metrics, loss weights + loss = {'dense': 'mse', 'dropout': 'mae'} + loss_weights = {'dense': 1., 'dropout': 0.5} + metrics = {'dense': 'mse', 'dropout': 'mae'} + model.compile(optimizer, loss, metrics=metrics, loss_weights=loss_weights) + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=1, + batch_size=5, + verbose=0) + + # Invalid use cases + with self.assertRaises(AttributeError): + model.fit( + [input_a_np, input_b_np], [output_d_np, output_e_np], + epochs=1, + validation_data=([input_a_np, input_b_np], 0, 0), + verbose=0) + with self.assertRaises(ValueError): + model.train_on_batch({'input_a': input_a_np}, + [output_d_np, output_e_np]) + with self.assertRaises(ValueError): + model.train_on_batch([input_a_np], [output_d_np, output_e_np]) + with self.assertRaises(AttributeError): + model.train_on_batch(1, [output_d_np, output_e_np]) + with self.assertRaises(ValueError): + model.train_on_batch(input_a_np, [output_d_np, output_e_np]) + with self.assertRaises(ValueError): + bad_input = np.random.random((11, 3)) + model.train_on_batch([bad_input, input_b_np], + [output_d_np, output_e_np]) + with self.assertRaises(ValueError): + bad_target = np.random.random((11, 4)) + model.train_on_batch([input_a_np, input_b_np], + [bad_target, output_e_np]) + + # Build single-input model + x = keras.layers.Input(shape=(3,), name='input_a') + y = keras.layers.Dense(4)(x) + model = keras.models.Model(x, y) + model.compile(optimizer=RMSPropOptimizer(learning_rate=0.001), loss='mse') + # This will work + model.fit([input_a_np], output_d_np, epochs=1) + with self.assertRaises(ValueError): + model.fit([input_a_np, input_a_np], output_d_np, epochs=1) + + def test_evaluate_predict_on_arrays(self): + a = keras.layers.Input(shape=(3,), name='input_a') + b = keras.layers.Input(shape=(3,), name='input_b') + + dense = keras.layers.Dense(4, name='dense') + c = dense(a) + d = dense(b) + e = keras.layers.Dropout(0.5, name='dropout')(c) + + model = keras.models.Model([a, b], [d, e]) + + optimizer = RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + loss_weights = [1., 0.5] + metrics = ['mae'] + model.compile( + optimizer, + loss, + metrics=metrics, + loss_weights=loss_weights, + sample_weight_mode=None) + + input_a_np = np.random.random((10, 3)) + input_b_np = np.random.random((10, 3)) + + output_d_np = np.random.random((10, 4)) + output_e_np = np.random.random((10, 4)) + + # Test evaluate at different verbosity + out = model.evaluate( + [input_a_np, input_b_np], [output_d_np, output_e_np], + batch_size=5, + verbose=0) + self.assertEqual(len(out), 5) + out = model.evaluate( + [input_a_np, input_b_np], [output_d_np, output_e_np], + batch_size=5, + verbose=1) + self.assertEqual(len(out), 5) + out = model.evaluate( + [input_a_np, input_b_np], [output_d_np, output_e_np], + batch_size=5, + verbose=2) + self.assertEqual(len(out), 5) + out = model.test_on_batch([input_a_np, input_b_np], + [output_d_np, output_e_np]) + self.assertEqual(len(out), 5) + + # Test evaluate with dictionary inputs + model.evaluate( + { + 'input_a': input_a_np, + 'input_b': input_b_np + }, {'dense': output_d_np, + 'dropout': output_e_np}, + batch_size=5, + verbose=0) + model.evaluate( + { + 'input_a': input_a_np, + 'input_b': input_b_np + }, {'dense': output_d_np, + 'dropout': output_e_np}, + batch_size=5, + verbose=1) + + # Test predict + out = model.predict([input_a_np, input_b_np], batch_size=5) + self.assertEqual(len(out), 2) + out = model.predict({'input_a': input_a_np, 'input_b': input_b_np}) + self.assertEqual(len(out), 2) + out = model.predict_on_batch({ + 'input_a': input_a_np, + 'input_b': input_b_np + }) + self.assertEqual(len(out), 2) + + def test_invalid_loss_or_metrics(self): + num_classes = 5 + train_samples = 1000 + test_samples = 1000 + input_dim = 5 + + model = keras.models.Sequential() + model.add(keras.layers.Dense(10, input_shape=(input_dim,))) + model.add(keras.layers.Activation('relu')) + model.add(keras.layers.Dense(num_classes)) + model.add(keras.layers.Activation('softmax')) + model.compile(loss='categorical_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001)) + np.random.seed(1337) + + (x_train, y_train), (_, _) = testing_utils.get_test_data( + train_samples=train_samples, + test_samples=test_samples, + input_shape=(input_dim,), + num_classes=num_classes) + + with self.assertRaises(ValueError): + model.fit(x_train, np.concatenate([y_train, y_train], axis=-1)) + + with self.assertRaises(TypeError): + model.compile(loss='categorical_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001), + metrics=set(0)) + + with self.assertRaises(ValueError): + model.compile(loss=None, + optimizer='rms') + + +class LossWeightingTest(test.TestCase): + + def test_class_weights(self): + num_classes = 5 + batch_size = 5 + weighted_class = 3 + train_samples = 300 + test_samples = 300 + input_dim = 5 + + model = keras.models.Sequential() + model.add(keras.layers.Dense(10, input_shape=(input_dim,))) + model.add(keras.layers.Activation('relu')) + model.add(keras.layers.Dense(num_classes)) + model.add(keras.layers.Activation('softmax')) + model.compile(loss='categorical_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001)) + + np.random.seed(1337) + (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( + train_samples=train_samples, + test_samples=test_samples, + input_shape=(input_dim,), + num_classes=num_classes) + int_y_test = y_test.copy() + int_y_train = y_train.copy() + # convert class vectors to binary class matrices + y_train = keras.utils.to_categorical(y_train, num_classes) + y_test = keras.utils.to_categorical(y_test, num_classes) + test_ids = np.where(int_y_test == np.array(weighted_class))[0] + + class_weight = dict([(i, 1.) for i in range(num_classes)]) + class_weight[weighted_class] = 4. + + sample_weight = np.ones((y_train.shape[0])) + sample_weight[int_y_train == weighted_class] = 4. + + model.fit( + x_train, + y_train, + batch_size=batch_size, + epochs=2, + verbose=0, + class_weight=class_weight, + validation_data=(x_train, y_train, sample_weight)) + model.fit( + x_train, + y_train, + batch_size=batch_size, + epochs=2, + verbose=0, + class_weight=class_weight) + model.fit( + x_train, + y_train, + batch_size=batch_size, + epochs=2, + verbose=0, + class_weight=class_weight, + validation_split=0.1) + + model.train_on_batch( + x_train[:batch_size], y_train[:batch_size], class_weight=class_weight) + ref_score = model.evaluate(x_test, y_test, verbose=0) + score = model.evaluate( + x_test[test_ids, :], y_test[test_ids, :], verbose=0) + self.assertLess(score, ref_score) + + def test_sample_weights(self): + num_classes = 5 + batch_size = 5 + weighted_class = 3 + train_samples = 300 + test_samples = 300 + input_dim = 5 + + model = keras.models.Sequential() + model.add(keras.layers.Dense(10, input_shape=(input_dim,))) + model.add(keras.layers.Activation('relu')) + model.add(keras.layers.Dense(num_classes)) + model.add(keras.layers.Activation('softmax')) + model.compile(loss='categorical_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001)) + + np.random.seed(43) + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=train_samples, + test_samples=test_samples, + input_shape=(input_dim,), + num_classes=num_classes) + int_y_train = y_train.copy() + y_train = keras.utils.to_categorical(y_train, num_classes) + + class_weight = dict([(i, 1.) for i in range(num_classes)]) + class_weight[weighted_class] = 4. + + sample_weight = np.ones((y_train.shape[0])) + sample_weight[int_y_train == weighted_class] = 4. + + model.fit( + x_train, + y_train, + batch_size=batch_size, + epochs=2, + verbose=0, + sample_weight=sample_weight) + model.fit( + x_train, + y_train, + batch_size=batch_size, + epochs=2, + verbose=0, + sample_weight=sample_weight, + validation_split=0.1) + model.train_on_batch( + x_train[:batch_size], + y_train[:batch_size], + sample_weight=sample_weight[:batch_size]) + model.test_on_batch( + x_train[:batch_size], + y_train[:batch_size], + sample_weight=sample_weight[:batch_size]) + + def test_temporal_sample_weights(self): + num_classes = 5 + weighted_class = 3 + train_samples = 1000 + test_samples = 1000 + input_dim = 5 + timesteps = 3 + + model = keras.models.Sequential() + model.add( + keras.layers.TimeDistributed( + keras.layers.Dense(num_classes), + input_shape=(timesteps, input_dim))) + model.add(keras.layers.Activation('softmax')) + + np.random.seed(1337) + (_, y_train), _ = testing_utils.get_test_data( + train_samples=train_samples, + test_samples=test_samples, + input_shape=(input_dim,), + num_classes=num_classes) + int_y_train = y_train.copy() + # convert class vectors to binary class matrices + y_train = keras.utils.to_categorical(y_train, num_classes) + + class_weight = dict([(i, 1.) for i in range(num_classes)]) + class_weight[weighted_class] = 2. + + sample_weight = np.ones((y_train.shape[0])) + sample_weight[int_y_train == weighted_class] = 2. + with self.assertRaises(ValueError): + model.compile( + loss='binary_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001), + sample_weight_mode='temporal') + + def test_class_weight_invalid_use_case(self): + num_classes = 5 + train_samples = 1000 + test_samples = 1000 + input_dim = 5 + timesteps = 3 + + model = keras.models.Sequential() + model.add( + keras.layers.TimeDistributed( + keras.layers.Dense(num_classes), + input_shape=(timesteps, input_dim))) + model.add(keras.layers.Activation('softmax')) + model.compile( + loss='binary_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001)) + + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=train_samples, + test_samples=test_samples, + input_shape=(input_dim,), + num_classes=num_classes) + # convert class vectors to binary class matrices + y_train = keras.utils.to_categorical(y_train, num_classes) + class_weight = dict([(i, 1.) for i in range(num_classes)]) + + del class_weight[1] + with self.assertRaises(ValueError): + model.fit(x_train, y_train, + epochs=0, verbose=0, class_weight=class_weight) + + with self.assertRaises(ValueError): + model.compile( + loss='binary_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001), + sample_weight_mode=[]) + + # Build multi-output model + x = keras.Input((3,)) + y1 = keras.layers.Dense(4, name='1')(x) + y2 = keras.layers.Dense(4, name='2')(x) + model = keras.models.Model(x, [y1, y2]) + model.compile(optimizer=RMSPropOptimizer(learning_rate=0.001), loss='mse') + x_np = np.random.random((10, 3)) + y_np = np.random.random((10, 4)) + w_np = np.random.random((10,)) + # This will work + model.fit(x_np, [y_np, y_np], epochs=1, sample_weight={'1': w_np}) + # These will not + with self.assertRaises(ValueError): + model.fit(x_np, [y_np, y_np], epochs=1, sample_weight=[w_np]) + with self.assertRaises(TypeError): + model.fit(x_np, [y_np, y_np], epochs=1, sample_weight=w_np) + with self.assertRaises(ValueError): + bad_w_np = np.random.random((11,)) + model.fit(x_np, [y_np, y_np], epochs=1, sample_weight={'1': bad_w_np}) + with self.assertRaises(ValueError): + bad_w_np = np.random.random((10, 2)) + model.fit(x_np, [y_np, y_np], epochs=1, sample_weight={'1': bad_w_np}) + with self.assertRaises(ValueError): + bad_w_np = np.random.random((10, 2, 2)) + model.fit(x_np, [y_np, y_np], epochs=1, sample_weight={'1': bad_w_np}) + + +if __name__ == '__main__': + # Bazel sets these environment variables to very long paths. + # Tempfile uses them to create long paths, and in turn multiprocessing + # library tries to create sockets named after paths. Delete whatever bazel + # writes to these to avoid tests failing due to socket addresses being too + # long. + for var in ('TMPDIR', 'TMP', 'TEMP'): + if var in os.environ: + del os.environ[var] + + ops.enable_eager_execution() + test.main() diff --git a/tensorflow/python/keras/_impl/keras/engine/training_generator.py b/tensorflow/python/keras/_impl/keras/engine/training_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..58b5bc39c10ea06f680eb030e14ecd19a3888588 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/training_generator.py @@ -0,0 +1,436 @@ +# 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. +# ============================================================================== +"""Part of the Keras training engine related to Python generators of array data. +""" +# pylint: disable=protected-access +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import callbacks as cbks +from tensorflow.python.keras._impl.keras.utils.data_utils import GeneratorEnqueuer +from tensorflow.python.keras._impl.keras.utils.data_utils import OrderedEnqueuer +from tensorflow.python.keras._impl.keras.utils.data_utils import Sequence +from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar +from tensorflow.python.platform import tf_logging as logging + + +def fit_generator(model, + generator, + steps_per_epoch=None, + epochs=1, + verbose=1, + callbacks=None, + validation_data=None, + validation_steps=None, + class_weight=None, + max_queue_size=10, + workers=1, + use_multiprocessing=False, + shuffle=True, + initial_epoch=0): + """See docstring for `Model.fit_generator`.""" + wait_time = 0.01 # in seconds + epoch = initial_epoch + + do_validation = bool(validation_data) + model._make_train_function() + if do_validation: + model._make_test_function() + + is_sequence = isinstance(generator, Sequence) + if not is_sequence and use_multiprocessing and workers > 1: + logging.warning( + UserWarning('Using a generator with `use_multiprocessing=True`' + ' and multiple workers may duplicate your data.' + ' Please consider using the`keras.utils.Sequence' + ' class.')) + if steps_per_epoch is None: + if is_sequence: + steps_per_epoch = len(generator) + else: + raise ValueError('`steps_per_epoch=None` is only valid for a' + ' generator based on the `keras.utils.Sequence`' + ' class. Please specify `steps_per_epoch` or use' + ' the `keras.utils.Sequence` class.') + + # python 2 has 'next', 3 has '__next__' + # avoid any explicit version checks + val_gen = ( + hasattr(validation_data, 'next') or + hasattr(validation_data, '__next__') or + isinstance(validation_data, Sequence)) + if (val_gen and not isinstance(validation_data, Sequence) and + not validation_steps): + raise ValueError('`validation_steps=None` is only valid for a' + ' generator based on the `keras.utils.Sequence`' + ' class. Please specify `validation_steps` or use' + ' the `keras.utils.Sequence` class.') + + # Prepare display labels. + out_labels = model.metrics_names + callback_metrics = out_labels + ['val_%s' % n for n in out_labels] + + # prepare callbacks + model.history = cbks.History() + callbacks = [cbks.BaseLogger()] + (callbacks or []) + [model.history] + if verbose: + callbacks += [cbks.ProgbarLogger(count_mode='steps')] + callbacks = cbks.CallbackList(callbacks) + + # it's possible to callback a different model than self: + if hasattr(model, 'callback_model') and model.callback_model: + callback_model = model.callback_model + else: + callback_model = model + callbacks.set_model(callback_model) + callbacks.set_params({ + 'epochs': epochs, + 'steps': steps_per_epoch, + 'verbose': verbose, + 'do_validation': do_validation, + 'metrics': callback_metrics, + }) + callbacks.on_train_begin() + + enqueuer = None + val_enqueuer = None + + try: + if do_validation and not val_gen: + # Prepare data for validation + if len(validation_data) == 2: + val_x, val_y = validation_data # pylint: disable=unpacking-non-sequence + val_sample_weight = None + elif len(validation_data) == 3: + val_x, val_y, val_sample_weight = validation_data # pylint: disable=unpacking-non-sequence + else: + raise ValueError( + '`validation_data` should be a tuple ' + '`(val_x, val_y, val_sample_weight)` ' + 'or `(val_x, val_y)`. Found: ' + str(validation_data)) + val_x, val_y, val_sample_weights = model._standardize_user_data( + val_x, val_y, val_sample_weight) + val_data = val_x + val_y + val_sample_weights + if model.uses_learning_phase and not isinstance(K.learning_phase(), int): + val_data += [0.] + for cbk in callbacks: + cbk.validation_data = val_data + + if workers > 0: + if is_sequence: + enqueuer = OrderedEnqueuer( + generator, + use_multiprocessing=use_multiprocessing, + shuffle=shuffle) + else: + enqueuer = GeneratorEnqueuer( + generator, + use_multiprocessing=use_multiprocessing, + wait_time=wait_time) + enqueuer.start(workers=workers, max_queue_size=max_queue_size) + output_generator = enqueuer.get() + else: + if is_sequence: + output_generator = iter(generator) + else: + output_generator = generator + + callback_model.stop_training = False + # Construct epoch logs. + epoch_logs = {} + while epoch < epochs: + callbacks.on_epoch_begin(epoch) + steps_done = 0 + batch_index = 0 + while steps_done < steps_per_epoch: + generator_output = next(output_generator) + + if not hasattr(generator_output, '__len__'): + raise ValueError('Output of generator should be ' + 'a tuple `(x, y, sample_weight)` ' + 'or `(x, y)`. Found: ' + str(generator_output)) + + if len(generator_output) == 2: + x, y = generator_output + sample_weight = None + elif len(generator_output) == 3: + x, y, sample_weight = generator_output + else: + raise ValueError('Output of generator should be ' + 'a tuple `(x, y, sample_weight)` ' + 'or `(x, y)`. Found: ' + str(generator_output)) + # build batch logs + batch_logs = {} + if isinstance(x, list): + batch_size = x[0].shape[0] + elif isinstance(x, dict): + batch_size = list(x.values())[0].shape[0] + else: + batch_size = x.shape[0] + batch_logs['batch'] = batch_index + batch_logs['size'] = batch_size + callbacks.on_batch_begin(batch_index, batch_logs) + + outs = model.train_on_batch( + x, y, sample_weight=sample_weight, class_weight=class_weight) + + if not isinstance(outs, list): + outs = [outs] + for l, o in zip(out_labels, outs): + batch_logs[l] = o + + callbacks.on_batch_end(batch_index, batch_logs) + + batch_index += 1 + steps_done += 1 + + # Epoch finished. + if steps_done >= steps_per_epoch and do_validation: + if val_gen: + val_outs = evaluate_generator( + model, + validation_data, + validation_steps, + workers=workers, + use_multiprocessing=use_multiprocessing, + max_queue_size=max_queue_size) + else: + # No need for try/except because + # data has already been validated. + val_outs = model.evaluate( + val_x, + val_y, + batch_size=batch_size, + sample_weight=val_sample_weights, + verbose=0) + if not isinstance(val_outs, list): + val_outs = [val_outs] + # Same labels assumed. + for l, o in zip(out_labels, val_outs): + epoch_logs['val_' + l] = o + + if callback_model.stop_training: + break + + callbacks.on_epoch_end(epoch, epoch_logs) + epoch += 1 + if callback_model.stop_training: + break + + finally: + try: + if enqueuer is not None: + enqueuer.stop() + finally: + if val_enqueuer is not None: + val_enqueuer.stop() + + callbacks.on_train_end() + return model.history + + +def evaluate_generator(model, + generator, + steps=None, + max_queue_size=10, + workers=1, + use_multiprocessing=False): + """See docstring for `Model.evaluate_generator`.""" + model._make_test_function() + + steps_done = 0 + wait_time = 0.01 + all_outs = [] + batch_sizes = [] + is_sequence = isinstance(generator, Sequence) + if not is_sequence and use_multiprocessing and workers > 1: + logging.warning( + UserWarning('Using a generator with `use_multiprocessing=True`' + ' and multiple workers may duplicate your data.' + ' Please consider using the`keras.utils.Sequence' + ' class.')) + if steps is None: + if is_sequence: + steps = len(generator) + else: + raise ValueError('`steps=None` is only valid for a generator' + ' based on the `keras.utils.Sequence` class.' + ' Please specify `steps` or use the' + ' `keras.utils.Sequence` class.') + enqueuer = None + + try: + if workers > 0: + if is_sequence: + enqueuer = OrderedEnqueuer( + generator, use_multiprocessing=use_multiprocessing) + else: + enqueuer = GeneratorEnqueuer( + generator, + use_multiprocessing=use_multiprocessing, + wait_time=wait_time) + enqueuer.start(workers=workers, max_queue_size=max_queue_size) + output_generator = enqueuer.get() + else: + if is_sequence: + output_generator = iter(generator) + else: + output_generator = generator + + while steps_done < steps: + generator_output = next(output_generator) + if not hasattr(generator_output, '__len__'): + raise ValueError('Output of generator should be a tuple ' + '(x, y, sample_weight) ' + 'or (x, y). Found: ' + str(generator_output)) + if len(generator_output) == 2: + x, y = generator_output + sample_weight = None + elif len(generator_output) == 3: + x, y, sample_weight = generator_output + else: + raise ValueError('Output of generator should be a tuple ' + '(x, y, sample_weight) ' + 'or (x, y). Found: ' + str(generator_output)) + outs = model.test_on_batch(x, y, sample_weight=sample_weight) + + if isinstance(x, list): + batch_size = x[0].shape[0] + elif isinstance(x, dict): + batch_size = list(x.values())[0].shape[0] + else: + batch_size = x.shape[0] + if batch_size == 0: + raise ValueError('Received an empty batch. ' + 'Batches should at least contain one item.') + all_outs.append(outs) + + steps_done += 1 + batch_sizes.append(batch_size) + + finally: + if enqueuer is not None: + enqueuer.stop() + + if not isinstance(outs, list): + return np.average(np.asarray(all_outs), weights=batch_sizes) + else: + averages = [] + for i in range(len(outs)): + averages.append( + np.average([out[i] for out in all_outs], weights=batch_sizes)) + return averages + + +def predict_generator(model, + generator, + steps=None, + max_queue_size=10, + workers=1, + use_multiprocessing=False, + verbose=0): + """See docstring for `Model.predict_generator`.""" + model._make_predict_function() + + steps_done = 0 + wait_time = 0.01 + all_outs = [] + is_sequence = isinstance(generator, Sequence) + if not is_sequence and use_multiprocessing and workers > 1: + logging.warning( + UserWarning('Using a generator with `use_multiprocessing=True`' + ' and multiple workers may duplicate your data.' + ' Please consider using the`keras.utils.Sequence' + ' class.')) + if steps is None: + if is_sequence: + steps = len(generator) + else: + raise ValueError('`steps=None` is only valid for a generator' + ' based on the `keras.utils.Sequence` class.' + ' Please specify `steps` or use the' + ' `keras.utils.Sequence` class.') + enqueuer = None + + try: + if workers > 0: + if is_sequence: + enqueuer = OrderedEnqueuer( + generator, use_multiprocessing=use_multiprocessing) + else: + enqueuer = GeneratorEnqueuer( + generator, + use_multiprocessing=use_multiprocessing, + wait_time=wait_time) + enqueuer.start(workers=workers, max_queue_size=max_queue_size) + output_generator = enqueuer.get() + else: + if is_sequence: + output_generator = iter(generator) + else: + output_generator = generator + + if verbose == 1: + progbar = Progbar(target=steps) + + while steps_done < steps: + generator_output = next(output_generator) + if isinstance(generator_output, tuple): + # Compatibility with the generators + # used for training. + if len(generator_output) == 2: + x, _ = generator_output + elif len(generator_output) == 3: + x, _, _ = generator_output + else: + raise ValueError('Output of generator should be ' + 'a tuple `(x, y, sample_weight)` ' + 'or `(x, y)`. Found: ' + str(generator_output)) + else: + # Assumes a generator that only + # yields inputs (not targets and sample weights). + x = generator_output + + outs = model.predict_on_batch(x) + if not isinstance(outs, list): + outs = [outs] + + if not all_outs: + for out in outs: + all_outs.append([]) + + for i, out in enumerate(outs): + all_outs[i].append(out) + steps_done += 1 + if verbose == 1: + progbar.update(steps_done) + + finally: + if enqueuer is not None: + enqueuer.stop() + + if len(all_outs) == 1: + if steps_done == 1: + return all_outs[0][0] + else: + return np.concatenate(all_outs[0]) + if steps_done == 1: + return [out[0] for out in all_outs] + else: + return [np.concatenate(out) for out in all_outs] diff --git a/tensorflow/python/keras/_impl/keras/engine/training_test.py b/tensorflow/python/keras/_impl/keras/engine/training_test.py index 5a033a04ade6c9b93ab32fb45f31d3efec85cd3f..fd91dbba52ff7d152335514085ef3b057ae5eec4 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/training_test.py @@ -25,7 +25,8 @@ import numpy as np from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils -from tensorflow.python.keras._impl.keras.engine.training import _weighted_masked_objective +from tensorflow.python.keras._impl.keras.engine.training_utils import weighted_masked_objective +from tensorflow.python.keras._impl.keras.utils.generic_utils import slice_arrays from tensorflow.python.platform import test try: @@ -78,6 +79,14 @@ class TrainingTest(test.TestCase): verbose=2) model.train_on_batch([input_a_np, input_b_np], [output_d_np, output_e_np]) + # Test model with input data as a list of lists + model.fit( + [np.ndarray.tolist(input_a_np), np.ndarray.tolist(input_b_np)], + [output_d_np, output_e_np], + epochs=2, + batch_size=5, + verbose=2) + # Test with validation data model.fit( [input_a_np, input_b_np], [output_d_np, output_e_np], @@ -205,6 +214,16 @@ class TrainingTest(test.TestCase): with self.assertRaises(ValueError): model.fit([input_a_np, input_a_np], output_d_np, epochs=1) + # Test model on a list of floats + input_a_np = np.random.random((10, 3)) + input_b_np = np.random.random((10, 4)) + + model.fit([np.ndarray.tolist(input_a_np)], + [np.ndarray.tolist(input_b_np)], + epochs=2, + batch_size=5, + verbose=2) + def test_evaluate_predict_on_arrays(self): with self.test_session(): a = keras.layers.Input(shape=(3,), name='input_a') @@ -321,20 +340,21 @@ class TrainingTest(test.TestCase): if scipy_sparse is None: return - test_inputs = [ - scipy_sparse.random(6, 3, density=0.25).tocsr() for _ in range(2)] - test_outputs = [ - scipy_sparse.random(6, i, density=0.25).tocsr() for i in range(3, 5)] - in1 = keras.layers.Input(shape=(3,)) - in2 = keras.layers.Input(shape=(3,)) - out1 = keras.layers.Dropout(0.5, name='dropout')(in1) - out2 = keras.layers.Dense(4, name='dense_1')(in2) - model = keras.Model([in1, in2], [out1, out2]) - model.predict(test_inputs, batch_size=2) - model.compile('rmsprop', 'mse') - model.fit(test_inputs, test_outputs, - epochs=1, batch_size=2, validation_split=0.5) - model.evaluate(test_inputs, test_outputs, batch_size=2) + with self.test_session(): + test_inputs = [ + scipy_sparse.random(6, 3, density=0.25).tocsr() for _ in range(2)] + test_outputs = [ + scipy_sparse.random(6, i, density=0.25).tocsr() for i in range(3, 5)] + in1 = keras.layers.Input(shape=(3,)) + in2 = keras.layers.Input(shape=(3,)) + out1 = keras.layers.Dropout(0.5, name='dropout')(in1) + out2 = keras.layers.Dense(4, name='dense_1')(in2) + model = keras.Model([in1, in2], [out1, out2]) + model.predict(test_inputs, batch_size=2) + model.compile('rmsprop', 'mse') + model.fit(test_inputs, test_outputs, + epochs=1, batch_size=2, validation_split=0.5) + model.evaluate(test_inputs, test_outputs, batch_size=2) def test_that_trainable_disables_updates(self): val_a = np.random.random((10, 4)) @@ -686,7 +706,7 @@ class LossMaskingTest(test.TestCase): def test_loss_masking(self): with self.test_session(): - weighted_loss = _weighted_masked_objective(keras.losses.get('mae')) + weighted_loss = weighted_masked_objective(keras.losses.get('mae')) shape = (3, 4, 2) x = np.arange(24).reshape(shape) y = 2 * x @@ -857,9 +877,9 @@ class TestGeneratorMethods(test.TestCase): def custom_generator(): batch_size = 10 - n_samples = 50 + num_samples = 50 while True: - batch_index = np.random.randint(0, n_samples - batch_size) + batch_index = np.random.randint(0, num_samples - batch_size) start = batch_index end = start + batch_size x = arr_data[start: end] @@ -938,9 +958,9 @@ class TestGeneratorMethods(test.TestCase): def custom_generator(): batch_size = 10 - n_samples = 50 + num_samples = 50 while True: - batch_index = np.random.randint(0, n_samples - batch_size) + batch_index = np.random.randint(0, num_samples - batch_size) start = batch_index end = start + batch_size x = arr_data[start: end] @@ -1014,47 +1034,85 @@ class TestGeneratorMethods(test.TestCase): max_queue_size=10, use_multiprocessing=False) + def test_training_with_sequences(self): + + class DummySequence(keras.utils.Sequence): + + def __getitem__(self, idx): + return np.zeros([10, 2]), np.ones([10]) + + def __len__(self): + return 10 + + arr_data = np.random.random((50, 2)) + arr_labels = np.random.random((50,)) + arr_sample_weights = np.random.random((50,)) + + def custom_generator(): + batch_size = 10 + num_samples = 50 + while True: + batch_index = np.random.randint(0, num_samples - batch_size) + start = batch_index + end = start + batch_size + x = arr_data[start: end] + y = arr_labels[start: end] + w = arr_sample_weights[start: end] + yield x, y, w + + with self.test_session(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(1, input_shape=(2,))) + model.compile(loss='mse', optimizer='sgd') + + model.fit_generator(DummySequence(), + steps_per_epoch=10, + validation_data=custom_generator(), + validation_steps=1, + max_queue_size=10, + workers=0, + use_multiprocessing=True) + model.fit_generator(DummySequence(), + steps_per_epoch=10, + validation_data=custom_generator(), + validation_steps=1, + max_queue_size=10, + workers=0, + use_multiprocessing=False) + class TestTrainingUtils(test.TestCase): def test_check_array_lengths(self): - keras.engine.training._check_array_lengths(None, None, None) + keras.engine.training_utils.check_array_lengths(None, None, None) a_np = np.random.random((4, 3, 3)) - keras.engine.training._check_array_lengths(a_np, a_np, a_np) - keras.engine.training._check_array_lengths( + keras.engine.training_utils.check_array_lengths(a_np, a_np, a_np) + keras.engine.training_utils.check_array_lengths( [a_np, a_np], [a_np, a_np], [a_np, a_np]) - keras.engine.training._check_array_lengths([None], [None], [None]) + keras.engine.training_utils.check_array_lengths([None], [None], [None]) b_np = np.random.random((3, 4)) with self.assertRaises(ValueError): - keras.engine.training._check_array_lengths(a_np, None, None) - with self.assertRaises(ValueError): - keras.engine.training._check_array_lengths(a_np, a_np, None) - with self.assertRaises(ValueError): - keras.engine.training._check_array_lengths([a_np], [None], None) - with self.assertRaises(ValueError): - keras.engine.training._check_array_lengths([a_np], [b_np], None) - with self.assertRaises(ValueError): - keras.engine.training._check_array_lengths([a_np], None, [b_np]) + keras.engine.training_utils.check_array_lengths([a_np], [b_np], None) def test_slice_arrays(self): input_a = np.random.random((10, 3)) - keras.engine.training._slice_arrays(None) - keras.engine.training._slice_arrays(input_a, 0) - keras.engine.training._slice_arrays(input_a, 0, 1) - keras.engine.training._slice_arrays(input_a, stop=2) + slice_arrays(input_a, 0) + slice_arrays(None) + slice_arrays(input_a, 0, 1) + slice_arrays(input_a, stop=2) input_a = [None, [1, 1], None, [1, 1]] - keras.engine.training._slice_arrays(input_a, 0) - keras.engine.training._slice_arrays(input_a, 0, 1) - keras.engine.training._slice_arrays(input_a, stop=2) + slice_arrays(input_a, 0) + slice_arrays(input_a, 0, 1) + slice_arrays(input_a, stop=2) input_a = [None] - keras.engine.training._slice_arrays(input_a, 0) - keras.engine.training._slice_arrays(input_a, 0, 1) - keras.engine.training._slice_arrays(input_a, stop=2) + slice_arrays(input_a, 0) + slice_arrays(input_a, 0, 1) + slice_arrays(input_a, stop=2) input_a = None - keras.engine.training._slice_arrays(input_a, 0) - keras.engine.training._slice_arrays(input_a, 0, 1) - keras.engine.training._slice_arrays(input_a, stop=2) + slice_arrays(input_a, 0) + slice_arrays(input_a, 0, 1) + slice_arrays(input_a, stop=2) class TestTrainingWithDataTensors(test.TestCase): diff --git a/tensorflow/python/keras/_impl/keras/engine/training_utils.py b/tensorflow/python/keras/_impl/keras/engine/training_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..105638ce1087e8668b49b6653a847667e8f9157d --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/engine/training_utils.py @@ -0,0 +1,534 @@ +# 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-related utilities. +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import copy + +import numpy as np + +from tensorflow.python.framework import tensor_util +from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import losses + + +def check_num_samples(ins, + batch_size=None, + steps=None, + steps_name='steps'): + """Determine the number of samples provided for training and evaluation. + + The number of samples is not defined when running with `steps`, + in which case the number of samples is set to `None`. + + Arguments: + ins: List of tensors to be fed to the Keras function. + batch_size: Integer batch size or `None` if not defined. + steps: Total number of steps (batches of samples) + before declaring `_predict_loop` finished. + Ignored with the default value of `None`. + steps_name: The public API's parameter name for `steps`. + + Raises: + ValueError: when `steps` is `None` and the attribute `ins.shape` + does not exist. Also raises ValueError when `steps` is not `None` + and `batch_size` is not `None` because they are mutually + exclusive. + + Returns: + When steps is `None`, returns the number of samples to be + processed based on the size of the first dimension of the + first input numpy array. When steps is not `None` and + `batch_size` is `None`, returns `None`. + + Raises: + ValueError: In case of invalid arguments. + """ + if steps is not None: + num_samples = None + if batch_size is not None: + raise ValueError( + 'If ' + steps_name + ' is set, the `batch_size` must be None.') + elif ins and hasattr(ins[0], 'shape'): + num_samples = ins[0].shape[0] + else: + raise ValueError( + 'Either the input data should have ' + 'a defined shape, or ' + steps_name + ' should be specified.') + return num_samples + + +def standardize_input_data(data, + names, + shapes=None, + check_batch_axis=True, + exception_prefix=''): + """Normalizes inputs and targets provided by users. + + Users may pass data as a list of arrays, dictionary of arrays, + or as a single array. We normalize this to an ordered list of + arrays (same order as `names`), while checking that the provided + arrays have shapes that match the network's expectations. + + Arguments: + data: User-provided input data (polymorphic). + names: List of expected array names. + shapes: Optional list of expected array shapes. + check_batch_axis: Boolean; whether to check that + the batch axis of the arrays matches the expected + value found in `shapes`. + exception_prefix: String prefix used for exception formatting. + + Returns: + List of standardized input arrays (one array per model input). + + Raises: + ValueError: in case of improperly formatted user-provided data. + """ + if not names: + if data is not None and hasattr(data, '__len__') and len(data): + raise ValueError('Error when checking model ' + exception_prefix + ': ' + 'expected no data, but got:', data) + return [] + if data is None: + return [None for _ in range(len(names))] + + if isinstance(data, dict): + try: + data = [ + data[x].values + if data[x].__class__.__name__ == 'DataFrame' else data[x] + for x in names + ] + except KeyError as e: + raise ValueError('No data provided for "' + e.args[0] + '". Need data ' + 'for each key in: ' + str(names)) + elif isinstance(data, list): + if isinstance(data[0], list): + data = [np.asarray(d) for d in data] + elif len(names) == 1 and isinstance(data[0], (float, int)): + data = [np.asarray(data)] + else: + data = [ + x.values if x.__class__.__name__ == 'DataFrame' else x for x in data + ] + else: + data = data.values if data.__class__.__name__ == 'DataFrame' else data + data = [data] + data = [ + np.expand_dims(x, 1) if x is not None and x.ndim == 1 else x for x in data + ] + + if len(data) != len(names): + if data and hasattr(data[0], 'shape'): + raise ValueError('Error when checking model ' + exception_prefix + + ': the list of Numpy arrays that you are passing to ' + 'your model is not the size the model expected. ' + 'Expected to see ' + str(len(names)) + ' array(s), ' + 'but instead got the following list of ' + + str(len(data)) + ' arrays: ' + str(data)[:200] + '...') + elif len(names) > 1: + raise ValueError( + 'Error when checking model ' + exception_prefix + + ': you are passing a list as input to your model, ' + 'but the model expects a list of ' + str(len(names)) + + ' Numpy arrays instead. The list you passed was: ' + str(data)[:200]) + elif len(data) == 1 and not hasattr(data[0], 'shape'): + raise TypeError('Error when checking model ' + exception_prefix + + ': data should be a Numpy array, or list/dict of ' + 'Numpy arrays. Found: ' + str(data)[:200] + '...') + elif len(names) == 1: + data = [np.asarray(data)] + + # Check shapes compatibility. + if shapes: + for i in range(len(names)): + if shapes[i] is not None: + data_shape = data[i].shape + shape = shapes[i] + if data[i].ndim != len(shape): + raise ValueError('Error when checking ' + exception_prefix + + ': expected ' + names[i] + ' to have ' + + str(len(shape)) + ' dimensions, but got array ' + 'with shape ' + str(data_shape)) + if not check_batch_axis: + data_shape = data_shape[1:] + shape = shape[1:] + for dim, ref_dim in zip(data_shape, shape): + if ref_dim != dim and ref_dim: + raise ValueError( + 'Error when checking ' + exception_prefix + ': expected ' + + names[i] + ' to have shape ' + str(shape) + + ' but got array with shape ' + str(data_shape)) + return data + + +def standardize_sample_or_class_weights(x_weight, output_names, weight_type): + """Maps `sample_weight` or `class_weight` to model outputs. + + Arguments: + x_weight: User-provided `sample_weight` or `class_weight` argument. + output_names: List of output names (strings) in the model. + weight_type: A string used purely for exception printing. + + Returns: + A list of `sample_weight` or `class_weight` where there are exactly + one element per model output. + + Raises: + ValueError: In case of invalid user-provided argument. + """ + if x_weight is None or len(x_weight) == 0: # pylint: disable=g-explicit-length-test + return [None for _ in output_names] + if len(output_names) == 1: + if isinstance(x_weight, list) and len(x_weight) == 1: + return x_weight + if isinstance(x_weight, dict) and output_names[0] in x_weight: + return [x_weight[output_names[0]]] + else: + return [x_weight] + if isinstance(x_weight, list): + if len(x_weight) != len(output_names): + raise ValueError('Provided `' + weight_type + '` was a list of ' + + str(len(x_weight)) + ' elements, but the model has ' + + str(len(output_names)) + ' outputs. ' + 'You should provide one `' + weight_type + '`' + 'array per model output.') + return x_weight + if isinstance(x_weight, dict): + x_weights = [] + for name in output_names: + x_weights.append(x_weight.get(name)) + return x_weights + else: + raise TypeError( + 'The model has multiple outputs, so `' + weight_type + '` ' + 'should be either a list or a dict. ' + 'Provided `' + weight_type + '` type not understood: ' + str(x_weight)) + + +def standardize_class_weights(class_weight, output_names): + return standardize_sample_or_class_weights(class_weight, output_names, + 'class_weight') + + +def standardize_sample_weights(sample_weight, output_names): + return standardize_sample_or_class_weights(sample_weight, output_names, + 'sample_weight') + + +def check_array_lengths(inputs, targets, weights=None): + """Does user input validation for numpy arrays. + + Arguments: + inputs: list of Numpy arrays of inputs. + targets: list of Numpy arrays of targets. + weights: list of Numpy arrays of sample weights. + + Raises: + ValueError: in case of incorrectly formatted data. + """ + + def set_of_lengths(x): + # return a set with the variation between + # different shapes, with None => 0 + if x is None: + return {} + else: + return set([y.shape[0] for y in x if y is not None]) + + set_x = set_of_lengths(inputs) + set_y = set_of_lengths(targets) + set_w = set_of_lengths(weights) + if len(set_x) > 1: + raise ValueError('All input arrays (x) should have ' + 'the same number of samples. Got array shapes: ' + + str([x.shape for x in inputs])) + if len(set_y) > 1: + raise ValueError('All target arrays (y) should have ' + 'the same number of samples. Got array shapes: ' + + str([y.shape for y in targets])) + if set_x and set_y and list(set_x)[0] != list(set_y)[0]: + raise ValueError('Input arrays should have ' + 'the same number of samples as target arrays. ' + 'Found ' + str(list(set_x)[0]) + ' input samples ' + 'and ' + str(list(set_y)[0]) + ' target samples.') + if len(set_w) > 1: + raise ValueError('All sample_weight arrays should have ' + 'the same number of samples. Got array shapes: ' + + str([w.shape for w in weights])) + if set_y and set_w and list(set_y)[0] != list(set_w)[0]: + raise ValueError('Sample_weight arrays should have ' + 'the same number of samples as target arrays. Got ' + + str(list(set_y)[0]) + ' input samples and ' + + str(list(set_w)[0]) + ' target samples.') + + +def check_loss_and_target_compatibility(targets, loss_fns, output_shapes): + """Does validation on the compatibility of targets and loss functions. + + This helps prevent users from using loss functions incorrectly. This check + is purely for UX purposes. + + Arguments: + targets: list of Numpy arrays of targets. + loss_fns: list of loss functions. + output_shapes: list of shapes of model outputs. + + Raises: + ValueError: if a loss function or target array + is incompatible with an output. + """ + key_losses = { + losses.mean_squared_error, losses.binary_crossentropy, + losses.categorical_crossentropy + } + for y, loss, shape in zip(targets, loss_fns, output_shapes): + if y is None or loss is None or tensor_util.is_tensor(y): + continue + if loss is losses.categorical_crossentropy: + if y.shape[-1] == 1: + raise ValueError('You are passing a target array of shape ' + str( + y.shape) + ' while using as loss `categorical_crossentropy`. ' + '`categorical_crossentropy` expects ' + 'targets to be binary matrices (1s and 0s) ' + 'of shape (samples, classes). ' + 'If your targets are integer classes, ' + 'you can convert them to the expected format via:\n' + '```\n' + 'from keras.utils import to_categorical\n' + 'y_binary = to_categorical(y_int)\n' + '```\n' + '\n' + 'Alternatively, you can use the loss function ' + '`sparse_categorical_crossentropy` instead, ' + 'which does expect integer targets.') + if loss in key_losses: + for target_dim, out_dim in zip(y.shape[1:], shape[1:]): + if out_dim is not None and target_dim != out_dim: + raise ValueError('A target array with shape ' + str(y.shape) + + ' was passed for an output of shape ' + str(shape) + + ' while using as loss `' + loss.__name__ + '`. ' + 'This loss expects ' + 'targets to have the same shape ' + 'as the output.') + + +def collect_metrics(metrics, output_names): + """Maps metric functions to model outputs. + + Arguments: + metrics: a list or dict of metric functions. + output_names: a list of the names (strings) of model outputs. + + Returns: + A list (one entry per model output) of lists of metric functions. + For instance, if the model has 2 outputs, and for the first output + we want to compute "binary_accuracy" and "binary_crossentropy", + and just "binary_accuracy" for the second output, + the list would look like: + `[[binary_accuracy, binary_crossentropy], [binary_accuracy]]` + + Raises: + TypeError: if an incorrect type is passed for the `metrics` argument. + """ + if not metrics: + return [[] for _ in output_names] + if isinstance(metrics, list): + # we then apply all metrics to all outputs. + return [copy.copy(metrics) for _ in output_names] + elif isinstance(metrics, dict): + nested_metrics = [] + for name in output_names: + output_metrics = metrics.get(name, []) + if not isinstance(output_metrics, list): + output_metrics = [output_metrics] + nested_metrics.append(output_metrics) + return nested_metrics + else: + raise TypeError('Type of `metrics` argument not understood. ' + 'Expected a list or dictionary, found: ' + str(metrics)) + + +def batch_shuffle(index_array, batch_size): + """Shuffles an array in a batch-wise fashion. + + Useful for shuffling HDF5 arrays + (where one cannot access arbitrary indices). + + Arguments: + index_array: array of indices to be shuffled. + batch_size: integer. + + Returns: + The `index_array` array, shuffled in a batch-wise fashion. + """ + batch_count = int(len(index_array) / batch_size) + # to reshape we need to be cleanly divisible by batch size + # we stash extra items and reappend them after shuffling + last_batch = index_array[batch_count * batch_size:] + index_array = index_array[:batch_count * batch_size] + index_array = index_array.reshape((batch_count, batch_size)) + np.random.shuffle(index_array) + index_array = index_array.flatten() + return np.append(index_array, last_batch) + + +def weighted_masked_objective(fn): + """Adds support for masking and sample-weighting to an objective function. + + It transforms an objective function `fn(y_true, y_pred)` + into a sample-weighted, cost-masked objective function + `fn(y_true, y_pred, weights, mask)`. + + Arguments: + fn: The objective function to wrap, + with signature `fn(y_true, y_pred)`. + + Returns: + A function with signature `fn(y_true, y_pred, weights, mask)`. + """ + if fn is None: + return None + + def weighted(y_true, y_pred, weights, mask=None): + """Wrapper function. + + Arguments: + y_true: `y_true` argument of `fn`. + y_pred: `y_pred` argument of `fn`. + weights: Weights tensor. + mask: Mask tensor. + + Returns: + Scalar tensor. + """ + # score_array has ndim >= 2 + score_array = fn(y_true, y_pred) + if mask is not None: + # Cast the mask to floatX to avoid float64 upcasting in theano + mask = K.cast(mask, K.floatx()) + # mask should have the same shape as score_array + score_array *= mask + # the loss per batch should be proportional + # to the number of unmasked samples. + score_array /= K.mean(mask) + + # apply sample weighting + if weights is not None: + # reduce score_array to same ndim as weight array + ndim = K.ndim(score_array) + weight_ndim = K.ndim(weights) + score_array = K.mean(score_array, axis=list(range(weight_ndim, ndim))) + score_array *= weights + score_array /= K.mean(K.cast(K.not_equal(weights, 0), K.floatx())) + return K.mean(score_array) + + return weighted + + +def standardize_weights(y, + sample_weight=None, + class_weight=None, + sample_weight_mode=None): + """Performs sample weight validation and standardization. + + Everything gets normalized to a single sample-wise (or timestep-wise) + weight array. + + Arguments: + y: Numpy array of model targets to be weighted. + sample_weight: User-provided `sample_weight` argument. + class_weight: User-provided `class_weight` argument. + sample_weight_mode: One of `None` or `"temporal"`. + `"temporal"` indicated that we expect 2D weight data + that will be applied to the last 2 dimensions of + the targets (i.e. we are weighting timesteps, not samples). + + Returns: + A numpy array of target weights, one entry per sample to weight. + + Raises: + ValueError: In case of invalid user-provided arguments. + """ + if sample_weight_mode is not None: + if sample_weight_mode != 'temporal': + raise ValueError('"sample_weight_mode ' + 'should be None or "temporal". ' + 'Found: ' + str(sample_weight_mode)) + if len(y.shape) < 3: + raise ValueError('Found a sample_weight array for ' + 'an input with shape ' + str(y.shape) + '. ' + 'Timestep-wise sample weighting (use of ' + 'sample_weight_mode="temporal") is restricted to ' + 'outputs that are at least 3D, i.e. that have ' + 'a time dimension.') + if sample_weight is not None and len(sample_weight.shape) != 2: + raise ValueError('Found a sample_weight array with shape ' + + str(sample_weight.shape) + '. ' + 'In order to use timestep-wise sample weighting, ' + 'you should pass a 2D sample_weight array.') + else: + if sample_weight is not None and len(sample_weight.shape) != 1: + raise ValueError('Found a sample_weight array with shape ' + + str(sample_weight.shape) + '. ' + 'In order to use timestep-wise sample weights, ' + 'you should specify ' + 'sample_weight_mode="temporal" ' + 'in compile(). If you just mean to use ' + 'sample-wise weights, make sure your ' + 'sample_weight array is 1D.') + + if sample_weight is not None: + if len(sample_weight.shape) > len(y.shape): + raise ValueError( + 'Found a sample_weight with shape' + str(sample_weight.shape) + '.' + 'Expected sample_weight with rank ' + 'less than or equal to ' + str(len(y.shape))) + + if y.shape[:sample_weight.ndim] != sample_weight.shape: + raise ValueError( + 'Found a sample_weight array with shape ' + str(sample_weight.shape) + + ' for an input with shape ' + str(y.shape) + '. ' + 'sample_weight cannot be broadcast.') + return sample_weight + elif isinstance(class_weight, dict): + if len(y.shape) > 2: + raise ValueError('`class_weight` not supported for ' + '3+ dimensional targets.') + if y.shape[1] > 1: + y_classes = np.argmax(y, axis=1) + elif y.shape[1] == 1: + y_classes = np.reshape(y, y.shape[0]) + else: + y_classes = y + + weights = np.asarray( + [class_weight[cls] for cls in y_classes if cls in class_weight]) + + if len(weights) != len(y_classes): + # subtract the sets to pick all missing classes + existing_classes = set(y_classes) + existing_class_weight = set(class_weight.keys()) + raise ValueError('`class_weight` must contain all classes in the data.' + ' The classes %s exist in the data but not in ' + '`class_weight`.' % + (existing_classes - existing_class_weight)) + return weights + else: + return None diff --git a/tensorflow/python/keras/_impl/keras/estimator.py b/tensorflow/python/keras/_impl/keras/estimator.py index 624e92a04b8860d9a3974f2edb4a443482958259..081f25e91452582afabee1e43f6e15f7e5d0013d 100644 --- a/tensorflow/python/keras/_impl/keras/estimator.py +++ b/tensorflow/python/keras/_impl/keras/estimator.py @@ -25,11 +25,15 @@ from tensorflow.python.client import session from tensorflow.python.estimator import estimator as estimator_lib from tensorflow.python.estimator import export as export_lib from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.estimator import run_config as run_config_lib from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras import models +from tensorflow.python.keras._impl.keras import optimizers +from tensorflow.python.keras._impl.keras.engine.base_layer import Layer +from tensorflow.python.keras._impl.keras.engine.network import Network from tensorflow.python.keras._impl.keras.utils.generic_utils import CustomObjectScope from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics as metrics_module @@ -37,6 +41,7 @@ from tensorflow.python.platform import tf_logging as logging from tensorflow.python.saved_model import signature_constants from tensorflow.python.training import saver as saver_lib from tensorflow.python.training import training_util +from tensorflow.python.util.tf_export import tf_export _DEFAULT_SERVING_KEY = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY @@ -49,36 +54,174 @@ def _cast_tensor_to_floatx(x): return math_ops.cast(x, K.floatx()) -def _create_ordered_io(keras_model, estimator_io_dict, is_input=True): +def _create_ordered_io(keras_model, estimator_io, is_input=True): """Create a list of tensors from IO dictionary based on Keras IO order. Args: - keras_model: an instance of compiled keras model. - estimator_io_dict: features or labels dictionary from model_fn. + keras_model: An instance of compiled keras model. + estimator_io: The features or labels (dict or plain array) from model_fn. is_input: True if dictionary is for inputs. Returns: - a list of tensors based on Keras IO order. + A list of tensors based on Keras IO order. Raises: ValueError: if dictionary keys cannot be found in Keras model input_names or output_names. """ - if is_input: - keras_io_names = keras_model.input_names + if isinstance(estimator_io, (list, tuple)): + # Case currently not supported by most built-in input_fn, + # but it's good to have for sanity + return [_cast_tensor_to_floatx(x) for x in estimator_io] + elif isinstance(estimator_io, dict): + if is_input: + if keras_model._is_graph_network: + keras_io_names = keras_model.input_names + else: + keras_io_names = [ + 'input_%d' % i for i in range(1, len(estimator_io) + 1)] + else: + if keras_model._is_graph_network: + keras_io_names = keras_model.output_names + else: + keras_io_names = [ + 'output_%d' % i for i in range(1, len(estimator_io) + 1)] + + for key in estimator_io: + if key not in keras_io_names: + raise ValueError( + 'Cannot find %s with name "%s" in Keras Model. ' + 'It needs to match one ' + 'of the following: %s' % ('input' if is_input else 'output', key, + ', '.join(keras_io_names))) + tensors = [_cast_tensor_to_floatx(estimator_io[io_name]) + for io_name in keras_io_names] + return tensors else: - keras_io_names = keras_model.output_names + # Plain array. + return _cast_tensor_to_floatx(estimator_io) - for key in estimator_io_dict: - if key not in keras_io_names: - raise ValueError( - 'Cannot find %s with name "%s" in Keras Model. It needs to match ' - 'one of the following: %s' % ('input' if is_input else 'output', key, - ', '.join(keras_io_names))) - tensors = [] - for io_name in keras_io_names: - tensors.append(_cast_tensor_to_floatx(estimator_io_dict[io_name])) - return tensors + +def _in_place_subclassed_model_reset(model): + """Substitute for model cloning that works for subclassed models. + + Subclassed models cannot be cloned because their topology is not serializable. + To "instantiate" an identical model in a new TF graph, we reuse the original + model object, but we clear its state. + + After calling this function on a model intance, you can use the model instance + as if it were a model clone (in particular you can use it in a new graph). + + This method clears the state of the input model. It is thus destructive. + However the original state can be restored fully by calling + `_in_place_subclassed_model_state_restoration`. + + Args: + model: Instance of a Keras model created via subclassing. + + Raises: + ValueError: In case the model uses a subclassed model as inner layer. + """ + assert not model._is_graph_network # Only makes sense for subclassed networks + # Retrieve all layers tracked by the model as well as their attribute names + attributes_cache = {} + for name in dir(model): + try: + value = getattr(model, name) + except (AttributeError, ValueError, TypeError): + continue + if isinstance(value, Layer): + attributes_cache[name] = value + assert value in model._layers + elif isinstance(value, (list, tuple)) and name not in ('layers', '_layers'): + # Handle case: list/tuple of layers (also tracked by the Network API). + if value and all(isinstance(val, Layer) for val in value): + raise ValueError('We do not support the use of list-of-layers ' + 'attributes in subclassed models used with ' + '`model_to_estimator` at this time. Found list ' + 'model: %s' % name) + + # Replace layers on the model with fresh layers + layers_to_names = {value: key for key, value in attributes_cache.items()} + original_layers = model._layers[:] + model._layers = [] + for layer in original_layers: # We preserve layer order. + config = layer.get_config() + # This will not work for nested subclassed models used as layers. + # This would be theoretically possible to support, but would add complexity. + # Only do it if users complain. + if isinstance(layer, Network) and not layer._is_graph_network: + raise ValueError('We do not support the use of nested subclassed models ' + 'in `model_to_estimator` at this time. Found nested ' + 'model: %s' % layer) + fresh_layer = layer.__class__.from_config(config) + name = layers_to_names[layer] + setattr(model, name, fresh_layer) + + # Cache original model build attributes (in addition to layers) + if (not hasattr(model, '_original_attributes_cache') or + model._original_attributes_cache is None): + if model.built: + attributes_to_cache = [ + 'inputs', + 'outputs', + '_feed_outputs', + '_feed_output_names', + '_feed_output_shapes', + '_feed_loss_fns', + 'loss_weights_list', + 'targets', + '_feed_targets', + 'sample_weight_modes', + 'weighted_metrics', + 'metrics_names', + 'metrics_tensors', + 'metrics_updates', + 'stateful_metric_names', + 'total_loss', + 'sample_weights', + '_feed_sample_weights', + 'train_function', + 'test_function', + 'predict_function', + '_collected_trainable_weights', + '_feed_inputs', + '_feed_input_names', + '_feed_input_shapes', + 'optimizer', + ] + for name in attributes_to_cache: + attributes_cache[name] = getattr(model, name) + model._original_attributes_cache = attributes_cache + + # Reset built state + model.built = False + model.inputs = None + model.outputs = None + + +def _in_place_subclassed_model_state_restoration(model): + """Restores the original state of a model after it was "reset". + + This undoes this action of `_in_place_subclassed_model_reset`. + + Args: + model: Instance of a Keras model created via subclassing, on which + `_in_place_subclassed_model_reset` was previously called. + """ + assert not model._is_graph_network + # Restore layers and build attributes + if (hasattr(model, '_original_attributes_cache') and + model._original_attributes_cache is not None): + model._layers = [] + for name, value in model._original_attributes_cache.items(): + setattr(model, name, value) + model._original_attributes_cache = None + else: + # Restore to the state of a never-called model. + model.built = False + model.inputs = None + model.outputs = None def _clone_and_build_model(mode, @@ -92,8 +235,8 @@ def _clone_and_build_model(mode, mode: training mode. keras_model: an instance of compiled keras model. custom_objects: Dictionary for custom objects. - features: - labels: + features: Dict of tensors. + labels: Dict of tensors, or single tensor instance. Returns: The newly built model. @@ -101,33 +244,49 @@ def _clone_and_build_model(mode, # Set to True during training, False for inference. K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN) - # Clone keras model. - input_tensors = None if features is None else _create_ordered_io( - keras_model, features) - if custom_objects: - with CustomObjectScope(custom_objects): + # Get list of inputs. + if features is None: + input_tensors = None + else: + input_tensors = _create_ordered_io(keras_model, + estimator_io=features, + is_input=True) + # Get list of outputs. + if labels is None: + target_tensors = None + elif isinstance(labels, dict): + target_tensors = _create_ordered_io(keras_model, + estimator_io=labels, + is_input=False) + else: + target_tensors = [ + _cast_tensor_to_floatx( + sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(labels)) + ] + + if keras_model._is_graph_network: + if custom_objects: + with CustomObjectScope(custom_objects): + model = models.clone_model(keras_model, input_tensors=input_tensors) + else: model = models.clone_model(keras_model, input_tensors=input_tensors) else: - model = models.clone_model(keras_model, input_tensors=input_tensors) + model = keras_model + _in_place_subclassed_model_reset(model) + if input_tensors is not None: + model._set_inputs(input_tensors) # Compile/Build model - if mode is model_fn_lib.ModeKeys.PREDICT and not model.built: - model.build() + if mode is model_fn_lib.ModeKeys.PREDICT: + if isinstance(model, models.Sequential): + model.build() else: - optimizer_config = keras_model.optimizer.get_config() - optimizer = keras_model.optimizer.__class__.from_config(optimizer_config) - optimizer.iterations = training_util.get_or_create_global_step() - - # Get list of outputs. - if labels is None: - target_tensors = None - elif isinstance(labels, dict): - target_tensors = _create_ordered_io(keras_model, labels, is_input=False) + if isinstance(keras_model.optimizer, optimizers.TFOptimizer): + optimizer = keras_model.optimizer else: - target_tensors = [ - _cast_tensor_to_floatx( - sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(labels)) - ] + optimizer_config = keras_model.optimizer.get_config() + optimizer = keras_model.optimizer.__class__.from_config(optimizer_config) + optimizer.iterations = training_util.get_or_create_global_step() model.compile( optimizer, @@ -167,10 +326,14 @@ def _create_keras_model_fn(keras_model, custom_objects=None): # Set loss and metric only during train and evaluate. if mode is not model_fn_lib.ModeKeys.PREDICT: - model._make_train_function() # pylint: disable=protected-access + if mode is model_fn_lib.ModeKeys.TRAIN: + model._make_train_function() # pylint: disable=protected-access + else: + model._make_test_function() # pylint: disable=protected-access loss = model.total_loss if model.metrics: + # TODO(fchollet): support stateful metrics eval_metric_ops = {} # When each metric maps to an output if isinstance(model.metrics, dict): @@ -194,6 +357,10 @@ def _create_keras_model_fn(keras_model, custom_objects=None): if mode is model_fn_lib.ModeKeys.TRAIN: train_op = model.train_function.updates_op + if not model._is_graph_network: + # Reset model state to original state, + # to avoid `model_fn` being destructive for the initial model argument. + _in_place_subclassed_model_state_restoration(keras_model) return model_fn_lib.EstimatorSpec( mode=mode, predictions=predictions, @@ -221,18 +388,18 @@ def _save_first_checkpoint(keras_model, estimator, custom_objects, Returns: The model_fn for a keras Estimator. """ - with ops.Graph().as_default() as g, g.device(estimator._device_fn): - random_seed.set_random_seed(estimator.config.tf_random_seed) - training_util.create_global_step() - model = _clone_and_build_model(model_fn_lib.ModeKeys.TRAIN, keras_model, - custom_objects) - - if isinstance(model, models.Sequential): - model = model.model - # Load weights and save to checkpoint if there is no checkpoint - latest_path = saver_lib.latest_checkpoint(estimator.model_dir) - if not latest_path: - with session.Session() as sess: + # Load weights and save to checkpoint if there is no checkpoint + latest_path = saver_lib.latest_checkpoint(estimator.model_dir) + if not latest_path: + with ops.Graph().as_default(): + random_seed.set_random_seed(estimator.config.tf_random_seed) + training_util.create_global_step() + model = _clone_and_build_model(model_fn_lib.ModeKeys.TRAIN, keras_model, + custom_objects) + if isinstance(model, models.Sequential): + model = model.model + # save to checkpoint + with session.Session(config=estimator._session_config) as sess: model.set_weights(keras_weights) # Make update ops and initialize all variables. if not model.train_function: @@ -244,6 +411,7 @@ def _save_first_checkpoint(keras_model, estimator, custom_objects, saver.save(sess, os.path.join(estimator.model_dir, 'keras_model.ckpt')) +@tf_export('keras.estimator.model_to_estimator') def model_to_estimator(keras_model=None, keras_model_path=None, custom_objects=None, @@ -272,10 +440,11 @@ def model_to_estimator(keras_model=None, """ if (not keras_model) and (not keras_model_path): raise ValueError( - 'Either keras_model or keras_model_path needs to be provided.') + 'Either `keras_model` or `keras_model_path` needs to be provided.') if keras_model and keras_model_path: raise ValueError( - 'Please specity either keras_model or keras_model_path but not both.') + 'Please specity either `keras_model` or `keras_model_path`, ' + 'but not both.') if not keras_model: if keras_model_path.startswith( @@ -286,18 +455,42 @@ def model_to_estimator(keras_model=None, logging.info('Loading models from %s', keras_model_path) keras_model = models.load_model(keras_model_path) else: - logging.info('Using the Keras model from memory.') + logging.info('Using the Keras model provided.') keras_model = keras_model - if not hasattr(keras_model, 'optimizer'): + if not hasattr(keras_model, 'optimizer') or not keras_model.optimizer: raise ValueError( - 'Given keras model has not been compiled yet. Please compile first ' - 'before creating the estimator.') + 'The given keras model has not been compiled yet. Please compile first ' + 'before calling `model_to_estimator`.') + + if isinstance(config, dict): + config = run_config_lib.RunConfig(**config) - keras_weights = keras_model.get_weights() keras_model_fn = _create_keras_model_fn(keras_model, custom_objects) - est = estimator_lib.Estimator( + estimator = estimator_lib.Estimator( keras_model_fn, model_dir=model_dir, config=config) - # TODO(yifeif): move checkpoint initialization to scaffold.init_fn - _save_first_checkpoint(keras_model, est, custom_objects, keras_weights) - return est + + # Pass the config into keras backend's default session. + with session.Session(config=estimator._session_config) as sess: + K.set_session(sess) + + keras_weights = keras_model.get_weights() + if keras_model._is_graph_network: + # TODO(yifeif): move checkpoint initialization to scaffold.init_fn + _save_first_checkpoint(keras_model, + estimator, + custom_objects, + keras_weights) + elif keras_model.built: + logging.warning('You are creating an Estimator from a Keras model ' + 'manually subclassed from `Model`, that was ' + 'already called on some inputs (and thus already had ' + 'weights). We are currently unable to preserve ' + 'the model\'s state (its weights) ' + 'as part of the estimator ' + 'in this case. Be warned that the estimator ' + 'has been created using ' + 'a freshly initialized version of your model.\n' + 'Note that this doesn\'t affect the state of the ' + 'model instance you passed as `keras_model` argument.') + return estimator diff --git a/tensorflow/python/keras/_impl/keras/estimator_test.py b/tensorflow/python/keras/_impl/keras/estimator_test.py index 9fc48b4117e7ee2c717d5418754254aa02b82869..e076dc25b16900636313f0ddd85a61b8d917fc91 100644 --- a/tensorflow/python/keras/_impl/keras/estimator_test.py +++ b/tensorflow/python/keras/_impl/keras/estimator_test.py @@ -17,12 +17,14 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import json from math import log10 import os import tempfile import numpy as np +from tensorflow.core.protobuf import config_pb2 from tensorflow.python.estimator import run_config as run_config_lib from tensorflow.python.estimator.inputs import numpy_io from tensorflow.python.framework import test_util @@ -32,6 +34,7 @@ from tensorflow.python.keras._impl.keras.applications import mobilenet from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.summary.writer import writer_cache +from tensorflow.python.training import rmsprop try: @@ -62,12 +65,42 @@ def simple_functional_model(): return model -def get_resource_for_simple_model(is_sequential, is_evaluate): - model = simple_sequential_model( - ) if is_sequential else simple_functional_model() - if is_sequential: +def simple_subclassed_model(): + + class SimpleModel(keras.Model): + + def __init__(self): + super(SimpleModel, self).__init__() + self.dense1 = keras.layers.Dense(16, activation='relu') + self.dp = keras.layers.Dropout(0.1) + self.dense2 = keras.layers.Dense(_NUM_CLASS, activation='softmax') + + def call(self, inputs): + x = self.dense1(inputs) + x = self.dp(x) + return self.dense2(x) + + return SimpleModel() + + +def get_resource_for_simple_model(model_type='sequential', + is_evaluate=False,): + if model_type == 'sequential': + model = simple_sequential_model() model.build() - input_name = model.input_names[0] + elif model_type == 'subclass': + model = simple_subclassed_model() + else: + assert model_type == 'functional' + model = simple_functional_model() + + if model_type == 'subclass': + input_name = 'input_1' + output_name = 'output_1' + else: + input_name = model.input_names[0] + output_name = model.output_names[0] + np.random.seed(_RANDOM_SEED) (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( train_samples=_TRAIN_SIZE, @@ -78,17 +111,19 @@ def get_resource_for_simple_model(is_sequential, is_evaluate): y_test = keras.utils.to_categorical(y_test) train_input_fn = numpy_io.numpy_input_fn( - x={input_name: x_train}, - y=y_train, + x=randomize_io_type(x_train, input_name), + y=randomize_io_type(y_train, output_name), shuffle=False, num_epochs=None, batch_size=16) evaluate_input_fn = numpy_io.numpy_input_fn( - x={input_name: x_test}, y=y_test, num_epochs=1, shuffle=False) + x=randomize_io_type(x_test, input_name), + y=randomize_io_type(y_test, output_name), + num_epochs=1, shuffle=False) predict_input_fn = numpy_io.numpy_input_fn( - x={input_name: x_test}, num_epochs=1, shuffle=False) + x=randomize_io_type(x_test, input_name), num_epochs=1, shuffle=False) inference_input_fn = evaluate_input_fn if is_evaluate else predict_input_fn @@ -96,6 +131,14 @@ def get_resource_for_simple_model(is_sequential, is_evaluate): y_test), train_input_fn, inference_input_fn +def randomize_io_type(array, name): + switch = np.random.random() + if switch > 0.5: + return array + else: + return {name: array} + + def multi_inputs_multi_outputs_model(): # test multi-input layer a = keras.layers.Input(shape=(16,), name='input_a') @@ -132,10 +175,10 @@ class TestKerasEstimator(test_util.TensorFlowTestCase): gfile.DeleteRecursively(self._base_dir) def test_train(self): - for is_sequential in [True, False]: + for model_type in ['sequential', 'functional']: keras_model, (_, _), ( _, _), train_input_fn, eval_input_fn = get_resource_for_simple_model( - is_sequential=is_sequential, is_evaluate=True) + model_type=model_type, is_evaluate=True) keras_model.compile( loss='categorical_crossentropy', optimizer='rmsprop', @@ -153,10 +196,87 @@ class TestKerasEstimator(test_util.TensorFlowTestCase): writer_cache.FileWriterCache.clear() gfile.DeleteRecursively(self._config.model_dir) + def test_train_with_tf_optimizer(self): + for model_type in ['sequential', 'functional']: + keras_model, (_, _), ( + _, _), train_input_fn, eval_input_fn = get_resource_for_simple_model( + model_type=model_type, is_evaluate=True) + keras_model.compile( + loss='categorical_crossentropy', + optimizer=rmsprop.RMSPropOptimizer(1e-3), + metrics=['mse', keras.metrics.categorical_accuracy]) + + with self.test_session(): + est_keras = keras.estimator.model_to_estimator( + keras_model=keras_model, + # Also use dict config argument to get test coverage for that line. + config={ + 'tf_random_seed': _RANDOM_SEED, + 'model_dir': self._base_dir, + }) + before_eval_results = est_keras.evaluate( + input_fn=eval_input_fn, steps=1) + est_keras.train(input_fn=train_input_fn, steps=_TRAIN_SIZE / 16) + after_eval_results = est_keras.evaluate(input_fn=eval_input_fn, steps=1) + self.assertLess(after_eval_results['loss'], before_eval_results['loss']) + + writer_cache.FileWriterCache.clear() + gfile.DeleteRecursively(self._config.model_dir) + + def test_train_with_subclassed_model(self): + keras_model, (_, _), ( + _, _), train_input_fn, eval_input_fn = get_resource_for_simple_model( + model_type='subclass', is_evaluate=True) + keras_model.compile( + loss='categorical_crossentropy', + optimizer=rmsprop.RMSPropOptimizer(1e-3), + metrics=['mse', keras.metrics.categorical_accuracy]) + + with self.test_session(): + est_keras = keras.estimator.model_to_estimator( + keras_model=keras_model, config=self._config) + est_keras.train(input_fn=train_input_fn, steps=_TRAIN_SIZE / 16) + before_eval_results = est_keras.evaluate( + input_fn=eval_input_fn, steps=1) + est_keras.train(input_fn=train_input_fn, steps=_TRAIN_SIZE / 16) + after_eval_results = est_keras.evaluate(input_fn=eval_input_fn, steps=1) + self.assertLess(after_eval_results['loss'], before_eval_results['loss']) + + def test_train_with_subclassed_model_with_existing_state(self): + keras_model, (_, _), ( + _, _), train_input_fn, eval_input_fn = get_resource_for_simple_model( + model_type='subclass', is_evaluate=True) + keras_model.compile( + loss='categorical_crossentropy', + optimizer=rmsprop.RMSPropOptimizer(1e-3), + metrics=['mse', keras.metrics.categorical_accuracy]) + + with self.test_session(): + # Create state + keras_model.train_on_batch(np.random.random((10,) + _INPUT_SIZE), + np.random.random((10, _NUM_CLASS))) + original_preds = keras_model.predict(np.ones((10,) + _INPUT_SIZE)) + + est_keras = keras.estimator.model_to_estimator( + keras_model=keras_model, config=self._config) + est_keras.train(input_fn=train_input_fn, steps=_TRAIN_SIZE / 16) + before_eval_results = est_keras.evaluate( + input_fn=eval_input_fn, steps=1) + est_keras.train(input_fn=train_input_fn, steps=_TRAIN_SIZE / 16) + after_eval_results = est_keras.evaluate(input_fn=eval_input_fn, steps=1) + self.assertLess(after_eval_results['loss'], before_eval_results['loss']) + + # Check that original model state was not altered + preds = keras_model.predict(np.ones((10,) + _INPUT_SIZE)) + self.assertAllClose(original_preds, preds, atol=1e-5) + # Check that the original model compilation did not break + keras_model.train_on_batch(np.random.random((10,) + _INPUT_SIZE), + np.random.random((10, _NUM_CLASS))) + def test_evaluate(self): keras_model, (x_train, y_train), ( x_test, y_test), _, eval_input_fn = get_resource_for_simple_model( - is_sequential=False, is_evaluate=True) + model_type='functional', is_evaluate=True) with self.test_session(): metrics = [ @@ -198,7 +318,7 @@ class TestKerasEstimator(test_util.TensorFlowTestCase): # Check that predict on a pretrained model yield the same result. keras_model, (x_train, y_train), ( x_test, _), _, pred_input_fn = get_resource_for_simple_model( - is_sequential=True, is_evaluate=False) + model_type='sequential', is_evaluate=False) with self.test_session(): keras_model.compile( @@ -260,7 +380,7 @@ class TestKerasEstimator(test_util.TensorFlowTestCase): keras_model, (x_train, y_train), ( x_test, _), _, pred_input_fn = get_resource_for_simple_model( - is_sequential=False, is_evaluate=False) + model_type='functional', is_evaluate=False) with self.test_session(): keras_model.compile( @@ -352,6 +472,46 @@ class TestKerasEstimator(test_util.TensorFlowTestCase): model_dir=tempfile.mkdtemp(dir=self._base_dir), custom_objects=custom_objects) + def test_tf_config(self): + keras_model, (_, _), (_, _), _, _ = get_resource_for_simple_model() + keras_model.compile( + loss='categorical_crossentropy', + optimizer='rmsprop', + metrics=['mse', keras.metrics.categorical_accuracy]) + + tf_config = json.dumps({ + 'cluster': { + run_config_lib.TaskType.PS: ['localhost:1234'], + run_config_lib.TaskType.WORKER: ['localhost:1236'], + run_config_lib.TaskType.MASTER: ['localhost:1238'] + }, + 'task': { + 'type': run_config_lib.TaskType.MASTER, + 'index': 0 + } + }) + with test.mock.patch.dict('os.environ', {'TF_CONFIG': tf_config}): + with self.test_session(): + keras.estimator.model_to_estimator( + keras_model=keras_model, + model_dir=tempfile.mkdtemp(dir=self._base_dir)) + + def test_gpu_config(self): + keras_model, (_, _), (_, _), _, _ = get_resource_for_simple_model() + keras_model.compile( + loss='categorical_crossentropy', + optimizer='rmsprop', + metrics=['mse', keras.metrics.categorical_accuracy]) + + gpu_options = config_pb2.GPUOptions(per_process_gpu_memory_fraction=0.3) + sess_config = config_pb2.ConfigProto(gpu_options=gpu_options) + self._config._session_config = sess_config + keras.estimator.model_to_estimator( + keras_model=keras_model, config=self._config) + self.assertEqual(keras.backend.get_session() + ._config.gpu_options.per_process_gpu_memory_fraction, + gpu_options.per_process_gpu_memory_fraction) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/_impl/keras/initializers.py b/tensorflow/python/keras/_impl/keras/initializers.py index 8752faa534a3d6094ce530e490571ff939f86dbb..300bed5e1437074d010760c427c14f68e58ac363 100644 --- a/tensorflow/python/keras/_impl/keras/initializers.py +++ b/tensorflow/python/keras/_impl/keras/initializers.py @@ -32,8 +32,10 @@ from tensorflow.python.ops.init_ops import RandomUniform from tensorflow.python.ops.init_ops import TruncatedNormal from tensorflow.python.ops.init_ops import VarianceScaling from tensorflow.python.ops.init_ops import Zeros +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.initializers.lecun_normal') def lecun_normal(seed=None): """LeCun normal initializer. @@ -56,6 +58,7 @@ def lecun_normal(seed=None): scale=1., mode='fan_in', distribution='normal', seed=seed) +@tf_export('keras.initializers.lecun_uniform') def lecun_uniform(seed=None): """LeCun uniform initializer. @@ -77,6 +80,7 @@ def lecun_uniform(seed=None): scale=1., mode='fan_in', distribution='uniform', seed=seed) +@tf_export('keras.initializers.glorot_normal') def glorot_normal(seed=None): """Glorot normal initializer, also called Xavier normal initializer. @@ -99,6 +103,7 @@ def glorot_normal(seed=None): scale=1., mode='fan_avg', distribution='normal', seed=seed) +@tf_export('keras.initializers.glorot_uniform') def glorot_uniform(seed=None): """Glorot uniform initializer, also called Xavier uniform initializer. @@ -121,6 +126,7 @@ def glorot_uniform(seed=None): scale=1., mode='fan_avg', distribution='uniform', seed=seed) +@tf_export('keras.initializers.he_normal') def he_normal(seed=None): """He normal initializer. @@ -141,6 +147,7 @@ def he_normal(seed=None): scale=2., mode='fan_in', distribution='normal', seed=seed) +@tf_export('keras.initializers.he_uniform') def he_uniform(seed=None): """He uniform variance scaling initializer. @@ -178,10 +185,12 @@ orthogonal = Orthogonal # Utility functions +@tf_export('keras.initializers.serialize') def serialize(initializer): return serialize_keras_object(initializer) +@tf_export('keras.initializers.deserialize') def deserialize(config, custom_objects=None): return deserialize_keras_object( config, @@ -190,6 +199,7 @@ def deserialize(config, custom_objects=None): printable_module_name='initializer') +@tf_export('keras.initializers.get') def get(identifier): if isinstance(identifier, dict): return deserialize(identifier) @@ -199,4 +209,5 @@ def get(identifier): elif callable(identifier): return identifier else: - raise ValueError('Could not interpret initializer identifier:', identifier) + raise ValueError('Could not interpret initializer identifier: ' + + str(identifier)) diff --git a/tensorflow/python/keras/_impl/keras/integration_test.py b/tensorflow/python/keras/_impl/keras/integration_test.py index 15c3d14727a44c9726a1c2c86f47640bcc490e70..280f7ed1b11e2026ac196eb319f7d5da8301f060 100644 --- a/tensorflow/python/keras/_impl/keras/integration_test.py +++ b/tensorflow/python/keras/_impl/keras/integration_test.py @@ -23,7 +23,6 @@ import numpy as np from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.layers import core as tf_core_layers -from tensorflow.python.layers import network as tf_network_layers from tensorflow.python.ops import nn from tensorflow.python.platform import test @@ -275,10 +274,10 @@ class KerasIntegrationTest(test.TestCase): y_train = keras.utils.to_categorical(y_train) y_test = keras.utils.to_categorical(y_test) - inputs = tf_network_layers.Input(shape=(10,)) + inputs = keras.Input(shape=(10,)) x = tf_core_layers.Dense(32, activation=nn.relu)(inputs) outputs = tf_core_layers.Dense(2, activation=nn.softmax)(x) - model = keras.models.Model(inputs, outputs) + model = keras.Model(inputs, outputs) model.summary() model.compile(loss='categorical_crossentropy', diff --git a/tensorflow/python/keras/_impl/keras/layers/advanced_activations.py b/tensorflow/python/keras/_impl/keras/layers/advanced_activations.py index ffbf77c4b8e4fa4454bfa82e473522ee7a316222..c40ee109aaea7dacea72e095b1d8cea3ed2e9bf8 100644 --- a/tensorflow/python/keras/_impl/keras/layers/advanced_activations.py +++ b/tensorflow/python/keras/_impl/keras/layers/advanced_activations.py @@ -25,9 +25,11 @@ from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers from tensorflow.python.keras._impl.keras.engine import InputSpec from tensorflow.python.keras._impl.keras.engine import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.LeakyReLU') class LeakyReLU(Layer): """Leaky version of a Rectified Linear Unit. @@ -66,6 +68,7 @@ class LeakyReLU(Layer): return input_shape +@tf_export('keras.layers.PReLU') class PReLU(Layer): """Parametric Rectified Linear Unit. @@ -163,6 +166,7 @@ class PReLU(Layer): return input_shape +@tf_export('keras.layers.ELU') class ELU(Layer): """Exponential Linear Unit. @@ -201,6 +205,7 @@ class ELU(Layer): return input_shape +@tf_export('keras.layers.ThresholdedReLU') class ThresholdedReLU(Layer): """Thresholded Rectified Linear Unit. @@ -239,6 +244,7 @@ class ThresholdedReLU(Layer): return input_shape +@tf_export('keras.layers.Softmax') class Softmax(Layer): """Softmax activation function. diff --git a/tensorflow/python/keras/_impl/keras/layers/convolutional.py b/tensorflow/python/keras/_impl/keras/layers/convolutional.py index 2ee07327751f9f7dcc87aaa83e18bcc1b5991d5c..162ae6c28f1afae1dd8aaf70213b808d9ad9598f 100644 --- a/tensorflow/python/keras/_impl/keras/layers/convolutional.py +++ b/tensorflow/python/keras/_impl/keras/layers/convolutional.py @@ -38,8 +38,10 @@ from tensorflow.python.keras._impl.keras.layers.pooling import MaxPooling3D # pylint: enable=unused-import from tensorflow.python.keras._impl.keras.utils import conv_utils from tensorflow.python.layers import convolutional as tf_convolutional_layers +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.Conv1D', 'keras.layers.Convolution1D') class Conv1D(tf_convolutional_layers.Conv1D, Layer): """1D convolution layer (e.g. temporal convolution). @@ -58,7 +60,7 @@ class Conv1D(tf_convolutional_layers.Conv1D, Layer): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output 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, @@ -153,6 +155,7 @@ class Conv1D(tf_convolutional_layers.Conv1D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Conv2D', 'keras.layers.Convolution2D') class Conv2D(tf_convolutional_layers.Conv2D, Layer): """2D convolution layer (e.g. spatial convolution over images). @@ -170,7 +173,7 @@ class Conv2D(tf_convolutional_layers.Conv2D, Layer): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of 2 integers, specifying the width and height of the 2D convolution window. Can be a single integer to specify the same value for @@ -286,6 +289,7 @@ class Conv2D(tf_convolutional_layers.Conv2D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Conv3D', 'keras.layers.Convolution3D') class Conv3D(tf_convolutional_layers.Conv3D, Layer): """3D convolution layer (e.g. spatial convolution over volumes). @@ -304,7 +308,7 @@ class Conv3D(tf_convolutional_layers.Conv3D, Layer): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output 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 @@ -426,6 +430,8 @@ class Conv3D(tf_convolutional_layers.Conv3D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Conv2DTranspose', + 'keras.layers.Convolution2DTranspose') class Conv2DTranspose(tf_convolutional_layers.Conv2DTranspose, Layer): """Transposed convolution layer (sometimes called Deconvolution). @@ -563,6 +569,8 @@ class Conv2DTranspose(tf_convolutional_layers.Conv2DTranspose, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Conv3DTranspose', + 'keras.layers.Convolution3DTranspose') class Conv3DTranspose(tf_convolutional_layers.Conv3DTranspose, Layer): """Transposed convolution layer (sometimes called Deconvolution). @@ -711,6 +719,8 @@ class Conv3DTranspose(tf_convolutional_layers.Conv3DTranspose, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.SeparableConv1D', + 'keras.layers.SeparableConvolution1D') class SeparableConv1D(tf_convolutional_layers.SeparableConv1D, Layer): """Depthwise separable 1D convolution. @@ -849,6 +859,8 @@ class SeparableConv1D(tf_convolutional_layers.SeparableConv1D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.SeparableConv2D', + 'keras.layers.SeparableConvolution2D') class SeparableConv2D(tf_convolutional_layers.SeparableConv2D, Layer): """Depthwise separable 2D convolution. @@ -865,7 +877,7 @@ class SeparableConv2D(tf_convolutional_layers.SeparableConv2D, Layer): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of 2 integers, specifying the width and height of the 2D convolution window. Can be a single integer to specify the same value for @@ -1012,6 +1024,7 @@ class SeparableConv2D(tf_convolutional_layers.SeparableConv2D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.UpSampling1D') class UpSampling1D(Layer): """Upsampling layer for 1D inputs. @@ -1047,6 +1060,7 @@ class UpSampling1D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.UpSampling2D') class UpSampling2D(Layer): """Upsampling layer for 2D inputs. @@ -1114,6 +1128,7 @@ class UpSampling2D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.UpSampling3D') class UpSampling3D(Layer): """Upsampling layer for 3D inputs. @@ -1186,6 +1201,7 @@ class UpSampling3D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.ZeroPadding1D') class ZeroPadding1D(Layer): """Zero-padding layer for 1D input (e.g. temporal sequence). @@ -1226,6 +1242,7 @@ class ZeroPadding1D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.ZeroPadding2D') class ZeroPadding2D(Layer): """Zero-padding layer for 2D input (e.g. picture). @@ -1327,6 +1344,7 @@ class ZeroPadding2D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.ZeroPadding3D') class ZeroPadding3D(Layer): """Zero-padding layer for 3D data (spatial or spatio-temporal). @@ -1444,6 +1462,7 @@ class ZeroPadding3D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Cropping1D') class Cropping1D(Layer): """Cropping layer for 1D input (e.g. temporal sequence). @@ -1488,6 +1507,7 @@ class Cropping1D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Cropping2D') class Cropping2D(Layer): """Cropping layer for 2D input (e.g. picture). @@ -1619,6 +1639,7 @@ class Cropping2D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Cropping3D') class Cropping3D(Layer): """Cropping layer for 3D data (e.g. diff --git a/tensorflow/python/keras/_impl/keras/layers/convolutional_recurrent.py b/tensorflow/python/keras/_impl/keras/layers/convolutional_recurrent.py index 565db19e41fd4ebf334ea950566c8edf6896deaf..d95a0942452afa82e277c358be5c3b2ba061ac98 100644 --- a/tensorflow/python/keras/_impl/keras/layers/convolutional_recurrent.py +++ b/tensorflow/python/keras/_impl/keras/layers/convolutional_recurrent.py @@ -26,9 +26,10 @@ from tensorflow.python.keras._impl.keras import constraints from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers from tensorflow.python.keras._impl.keras.engine import InputSpec -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion from tensorflow.python.keras._impl.keras.layers.recurrent import Recurrent from tensorflow.python.keras._impl.keras.utils import conv_utils +from tensorflow.python.util.tf_export import tf_export class ConvRecurrent2D(Recurrent): @@ -38,7 +39,7 @@ class ConvRecurrent2D(Recurrent): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of n integers, specifying the dimensions of the convolution window. strides: An integer or tuple/list of n integers, @@ -190,6 +191,7 @@ class ConvRecurrent2D(Recurrent): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.ConvLSTM2D') class ConvLSTM2D(ConvRecurrent2D): """Convolutional LSTM. @@ -198,7 +200,7 @@ class ConvLSTM2D(ConvRecurrent2D): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of n integers, specifying the dimensions of the convolution window. strides: An integer or tuple/list of n integers, diff --git a/tensorflow/python/keras/_impl/keras/layers/convolutional_test.py b/tensorflow/python/keras/_impl/keras/layers/convolutional_test.py index 39c9d4f0fb2751b0eef3b28f6d5b8cb0a93e22e5..f4a134b96cec0385cb24a208f3403db944b68edc 100644 --- a/tensorflow/python/keras/_impl/keras/layers/convolutional_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/convolutional_test.py @@ -18,8 +18,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import copy + import numpy as np +from tensorflow.python.eager import context +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test @@ -27,45 +31,40 @@ from tensorflow.python.platform import test class Convolution1DTest(test.TestCase): - def test_dilated_conv1d(self): - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.Conv1D, - input_data=np.reshape(np.arange(4, dtype='float32'), (1, 4, 1)), - kwargs={ - 'filters': 1, - 'kernel_size': 2, - 'dilation_rate': 1, - 'padding': 'valid', - 'kernel_initializer': 'ones', - 'use_bias': False, - }, - expected_output=[[[1], [3], [5]]]) - - def test_conv_1d(self): - batch_size = 2 - steps = 8 - input_dim = 2 - kernel_size = 3 - filters = 3 + def _run_test(self, kwargs, arg, values): + num_samples = 2 + stack_size = 3 + length = 7 - for padding in ['valid', 'same']: - for strides in [1, 2]: - if padding == 'same' and strides != 1: - continue + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.Conv1D, + kwargs=test_kwargs, + input_shape=(num_samples, length, stack_size)) + + @tf_test_util.run_in_graph_and_eager_modes() + def test_conv1d(self): + kwargs = { + 'filters': 2, + 'kernel_size': 3, + } + + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [2]) + self._run_test(kwargs, 'dilation_rate', [2]) - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.Conv1D, - kwargs={ - 'filters': filters, - 'kernel_size': kernel_size, - 'padding': padding, - 'strides': strides - }, - input_shape=(batch_size, steps, input_dim)) - - def test_conv_1d_regularizers(self): + kwargs = { + 'filters': 2, + 'kernel_size': 3, + 'padding': 'same', + } + self._run_test(kwargs, 'dilation_rate', [2]) + self._run_test(kwargs, 'dilation_rate', [3]) + + def test_conv1d_regularizers(self): kwargs = { 'filters': 3, 'kernel_size': 3, @@ -82,7 +81,7 @@ class Convolution1DTest(test.TestCase): layer(keras.backend.variable(np.ones((1, 5, 2)))) self.assertEqual(len(layer.losses), 3) - def test_conv_1d_constraints(self): + def test_conv1d_constraints(self): k_constraint = lambda x: x b_constraint = lambda x: x @@ -103,35 +102,44 @@ class Convolution1DTest(test.TestCase): class Conv2DTest(test.TestCase): - def test_convolution_2d(self): + def _run_test(self, kwargs, arg, values): num_samples = 2 - filters = 2 stack_size = 3 - kernel_size = (3, 2) num_row = 7 num_col = 6 - for padding in ['valid', 'same']: - for strides in [(1, 1), (2, 2)]: - if padding == 'same' and strides != (1, 1): - continue + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.SeparableConv2D, + kwargs=test_kwargs, + input_shape=(num_samples, num_row, num_col, stack_size)) + + @tf_test_util.run_in_graph_and_eager_modes() + def test_conv2d(self): + kwargs = { + 'filters': 2, + 'kernel_size': (3, 3), + } + + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [(2, 2)]) + if test.is_gpu_available(cuda_only=True): + # Only runs on GPU with CUDA, channels_first is not supported on CPU. + # TODO(b/62340061): Support channels_first on CPU. + self._run_test(kwargs, 'data_format', ['channels_first']) + self._run_test(kwargs, 'dilation_rate', [(2, 2)]) - with self.test_session(use_gpu=True): - # Only runs on GPU with CUDA, channels_first is not supported on CPU. - # TODO(b/62340061): Support channels_first on CPU. - if test.is_gpu_available(cuda_only=True): - testing_utils.layer_test( - keras.layers.Conv2D, - kwargs={ - 'filters': filters, - 'kernel_size': kernel_size, - 'padding': padding, - 'strides': strides, - 'data_format': 'channels_first' - }, - input_shape=(num_samples, stack_size, num_row, num_col)) - - def test_convolution_2d_regularizers(self): + kwargs = { + 'filters': 2, + 'kernel_size': 3, + 'padding': 'same', + } + self._run_test(kwargs, 'dilation_rate', [2]) + + def test_conv2d_regularizers(self): kwargs = { 'filters': 3, 'kernel_size': 3, @@ -148,7 +156,7 @@ class Conv2DTest(test.TestCase): layer(keras.backend.variable(np.ones((1, 5, 5, 2)))) self.assertEqual(len(layer.losses), 3) - def test_convolution_2d_constraints(self): + def test_conv2d_constraints(self): k_constraint = lambda x: x b_constraint = lambda x: x @@ -166,51 +174,35 @@ class Conv2DTest(test.TestCase): self.assertEqual(layer.kernel.constraint, k_constraint) self.assertEqual(layer.bias.constraint, b_constraint) - def test_dilated_conv_2d(self): - num_samples = 2 - filters = 2 - stack_size = 3 - kernel_size = (3, 2) - num_row = 7 - num_col = 6 - - # Test dilation - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.Conv2D, - kwargs={ - 'filters': filters, - 'kernel_size': kernel_size, - 'dilation_rate': (2, 2) - }, - input_shape=(num_samples, num_row, num_col, stack_size)) - class Conv2DTransposeTest(test.TestCase): - def test_conv2d_transpose(self): + def _run_test(self, kwargs, arg, values): num_samples = 2 - filters = 2 stack_size = 3 - num_row = 5 + num_row = 7 num_col = 6 - for padding in ['valid', 'same']: - for strides in [(1, 1), (2, 2)]: - if padding == 'same' and strides != (1, 1): - continue + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.Conv2DTranspose, + kwargs=test_kwargs, + input_shape=(num_samples, num_row, num_col, stack_size)) - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.Conv2DTranspose, - kwargs={ - 'filters': filters, - 'kernel_size': 3, - 'padding': padding, - 'strides': strides, - 'data_format': 'channels_last' - }, - input_shape=(num_samples, num_row, num_col, stack_size)) + @tf_test_util.run_in_graph_and_eager_modes() + def test_conv2dtranspose(self): + kwargs = { + 'filters': 2, + 'kernel_size': (3, 3), + } + + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [(2, 2)]) + if test.is_gpu_available(cuda_only=True): + self._run_test(kwargs, 'data_format', ['channels_first']) def test_conv2dtranspose_regularizers(self): kwargs = { @@ -250,30 +242,33 @@ class Conv2DTransposeTest(test.TestCase): class Conv3DTransposeTest(test.TestCase): - def test_conv3d_transpose(self): + def _run_test(self, kwargs, arg, values): num_samples = 2 - filters = 2 stack_size = 3 - num_row = 5 + num_row = 7 num_col = 6 - depth = 4 + depth = 5 - for padding in ['valid', 'same']: - for strides in [(1, 1, 1), (2, 2, 2)]: - if padding == 'same' and strides != (1, 1, 1): - continue + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.Conv3DTranspose, + kwargs=test_kwargs, + input_shape=(num_samples, depth, num_row, num_col, stack_size)) - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.Conv3DTranspose, - kwargs={ - 'filters': filters, - 'kernel_size': 3, - 'padding': padding, - 'strides': strides, - 'data_format': 'channels_last' - }, - input_shape=(num_samples, depth, num_row, num_col, stack_size)) + @tf_test_util.run_in_graph_and_eager_modes() + def test_conv3dtranspose(self): + kwargs = { + 'filters': 2, + 'kernel_size': (3, 3, 3), + } + + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [(2, 2, 2)]) + if test.is_gpu_available(cuda_only=True): + self._run_test(kwargs, 'data_format', ['channels_first']) def test_conv3dtranspose_regularizers(self): kwargs = { @@ -313,29 +308,38 @@ class Conv3DTransposeTest(test.TestCase): class SeparableConv1DTest(test.TestCase): - def test_separable_conv_1d(self): + def _run_test(self, kwargs, arg, values): num_samples = 2 - filters = 6 stack_size = 3 length = 7 - strides = 1 - for padding in ['valid', 'same']: - for multiplier in [1, 2]: - if padding == 'same' and strides != 1: - continue + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.SeparableConv1D, + kwargs=test_kwargs, + input_shape=(num_samples, length, stack_size)) - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.SeparableConv1D, - kwargs={ - 'filters': filters, - 'kernel_size': 3, - 'padding': padding, - 'strides': strides, - 'depth_multiplier': multiplier - }, - input_shape=(num_samples, length, stack_size)) + @tf_test_util.run_in_graph_and_eager_modes() + def test_separable_conv1d(self): + kwargs = { + 'filters': 2, + 'kernel_size': 3, + } + + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [2]) + self._run_test(kwargs, 'dilation_rate', [2]) + self._run_test(kwargs, 'depth_multiplier', [2]) + + kwargs = { + 'filters': 2, + 'kernel_size': 3, + 'padding': 'same', + } + self._run_test(kwargs, 'dilation_rate', [2]) def test_separable_conv1d_regularizers(self): kwargs = { @@ -379,30 +383,41 @@ class SeparableConv1DTest(test.TestCase): class SeparableConv2DTest(test.TestCase): - def test_separable_conv_2d(self): + def _run_test(self, kwargs, arg, values): num_samples = 2 - filters = 6 stack_size = 3 num_row = 7 num_col = 6 - for padding in ['valid', 'same']: - for strides in [(1, 1), (2, 2)]: - for multiplier in [1, 2]: - if padding == 'same' and strides != (1, 1): - continue + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.SeparableConv2D, + kwargs=test_kwargs, + input_shape=(num_samples, num_row, num_col, stack_size)) + + @tf_test_util.run_in_graph_and_eager_modes() + def test_separable_conv2d(self): + kwargs = { + 'filters': 2, + 'kernel_size': 3, + } - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.SeparableConv2D, - kwargs={ - 'filters': filters, - 'kernel_size': (3, 3), - 'padding': padding, - 'strides': strides, - 'depth_multiplier': multiplier - }, - input_shape=(num_samples, num_row, num_col, stack_size)) + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [2]) + if test.is_gpu_available(cuda_only=True): + self._run_test(kwargs, 'data_format', ['channels_first']) + self._run_test(kwargs, 'dilation_rate', [2]) + self._run_test(kwargs, 'depth_multiplier', [2]) + + kwargs = { + 'filters': 2, + 'kernel_size': 3, + 'padding': 'same', + } + self._run_test(kwargs, 'dilation_rate', [2]) def test_separable_conv2d_regularizers(self): kwargs = { @@ -446,33 +461,36 @@ class SeparableConv2DTest(test.TestCase): class Conv3DTest(test.TestCase): - def test_convolution_3d(self): + def _run_test(self, kwargs, arg, values): num_samples = 2 - filters = 2 stack_size = 3 + num_row = 7 + num_col = 6 + depth = 5 - input_len_dim1 = 9 - input_len_dim2 = 8 - input_len_dim3 = 8 + test_kwargs = copy.copy(kwargs) + for value in values: + test_kwargs[arg] = value + with self.test_session(use_gpu=True): + testing_utils.layer_test( + keras.layers.Conv3D, + kwargs=test_kwargs, + input_shape=(num_samples, depth, num_row, num_col, stack_size)) - for padding in ['valid', 'same']: - for strides in [(1, 1, 1), (2, 2, 2)]: - if padding == 'same' and strides != (1, 1, 1): - continue + @tf_test_util.run_in_graph_and_eager_modes() + def test_conv3d(self): + kwargs = { + 'filters': 2, + 'kernel_size': (3, 3, 3), + } - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.Convolution3D, - kwargs={ - 'filters': filters, - 'kernel_size': 3, - 'padding': padding, - 'strides': strides - }, - input_shape=(num_samples, input_len_dim1, input_len_dim2, - input_len_dim3, stack_size)) - - def test_convolution_3d_regularizers(self): + self._run_test(kwargs, 'padding', ['valid', 'same']) + self._run_test(kwargs, 'strides', [(2, 2, 2)]) + self._run_test(kwargs, 'dilation_rate', [(2, 2, 2)]) + if test.is_gpu_available(cuda_only=True): + self._run_test(kwargs, 'data_format', ['channels_first']) + + def test_conv3d_regularizers(self): kwargs = { 'filters': 3, 'kernel_size': 3, @@ -490,7 +508,7 @@ class Conv3DTest(test.TestCase): layer(keras.backend.variable(np.ones((1, 5, 5, 5, 2)))) self.assertEqual(len(layer.losses), 3) - def test_convolution_3d_constraints(self): + def test_conv3d_constraints(self): k_constraint = lambda x: x b_constraint = lambda x: x @@ -511,6 +529,7 @@ class Conv3DTest(test.TestCase): class ZeroPaddingTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_zero_padding_1d(self): num_samples = 2 input_dim = 2 @@ -534,7 +553,10 @@ class ZeroPaddingTest(test.TestCase): layer = keras.layers.ZeroPadding1D(padding=2) layer.build(shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) for offset in [0, 1, -1, -2]: np.testing.assert_allclose(np_output[:, offset, :], 0.) np.testing.assert_allclose(np_output[:, 2:-2, :], 1.) @@ -542,7 +564,10 @@ class ZeroPaddingTest(test.TestCase): layer = keras.layers.ZeroPadding1D(padding=(1, 2)) layer.build(shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) for left_offset in [0]: np.testing.assert_allclose(np_output[:, left_offset, :], 0.) for right_offset in [-1, -2]: @@ -556,6 +581,7 @@ class ZeroPaddingTest(test.TestCase): with self.assertRaises(ValueError): keras.layers.ZeroPadding1D(padding=None) + @tf_test_util.run_in_graph_and_eager_modes() def test_zero_padding_2d(self): num_samples = 2 stack_size = 2 @@ -584,7 +610,10 @@ class ZeroPaddingTest(test.TestCase): padding=(2, 2), data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) if data_format == 'channels_last': for offset in [0, 1, -1, -2]: np.testing.assert_allclose(np_output[:, offset, :, :], 0.) @@ -600,7 +629,10 @@ class ZeroPaddingTest(test.TestCase): padding=((1, 2), (3, 4)), data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) if data_format == 'channels_last': for top_offset in [0]: np.testing.assert_allclose(np_output[:, top_offset, :, :], 0.) @@ -628,6 +660,7 @@ class ZeroPaddingTest(test.TestCase): with self.assertRaises(ValueError): keras.layers.ZeroPadding2D(padding=None) + @tf_test_util.run_in_graph_and_eager_modes() def test_zero_padding_3d(self): num_samples = 2 stack_size = 2 @@ -650,7 +683,10 @@ class ZeroPaddingTest(test.TestCase): layer = keras.layers.ZeroPadding3D(padding=(2, 2, 2)) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) for offset in [0, 1, -1, -2]: np.testing.assert_allclose(np_output[:, offset, :, :, :], 0.) np.testing.assert_allclose(np_output[:, :, offset, :, :], 0.) @@ -666,11 +702,13 @@ class ZeroPaddingTest(test.TestCase): class UpSamplingTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_upsampling_1d(self): with self.test_session(use_gpu=True): testing_utils.layer_test( keras.layers.UpSampling1D, kwargs={'size': 2}, input_shape=(3, 5, 4)) + @tf_test_util.run_in_graph_and_eager_modes() def test_upsampling_2d(self): num_samples = 2 stack_size = 2 @@ -699,7 +737,10 @@ class UpSamplingTest(test.TestCase): size=(length_row, length_col), data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) if data_format == 'channels_first': assert np_output.shape[2] == length_row * input_num_row assert np_output.shape[3] == length_col * input_num_col @@ -717,6 +758,7 @@ class UpSamplingTest(test.TestCase): np.testing.assert_allclose(np_output, expected_out) + @tf_test_util.run_in_graph_and_eager_modes() def test_upsampling_3d(self): num_samples = 2 stack_size = 2 @@ -748,7 +790,10 @@ class UpSamplingTest(test.TestCase): data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) if data_format == 'channels_first': assert np_output.shape[2] == length_dim1 * input_len_dim1 assert np_output.shape[3] == length_dim2 * input_len_dim2 @@ -773,6 +818,7 @@ class UpSamplingTest(test.TestCase): class CroppingTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_cropping_1d(self): num_samples = 2 time_length = 4 @@ -791,6 +837,7 @@ class CroppingTest(test.TestCase): with self.assertRaises(ValueError): keras.layers.Cropping1D(cropping=None) + @tf_test_util.run_in_graph_and_eager_modes() def test_cropping_2d(self): num_samples = 2 stack_size = 2 @@ -818,7 +865,10 @@ class CroppingTest(test.TestCase): cropping=cropping, data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) # compare with numpy if data_format == 'channels_first': expected_out = inputs[:, :, cropping[0][0]:-cropping[0][1], cropping[ @@ -842,7 +892,10 @@ class CroppingTest(test.TestCase): cropping=cropping, data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) # compare with input np.testing.assert_allclose(np_output, inputs) @@ -852,6 +905,7 @@ class CroppingTest(test.TestCase): with self.assertRaises(ValueError): keras.layers.Cropping2D(cropping=None) + @tf_test_util.run_in_graph_and_eager_modes() def test_cropping_3d(self): num_samples = 2 stack_size = 2 @@ -883,7 +937,10 @@ class CroppingTest(test.TestCase): cropping=cropping, data_format=data_format) layer.build(inputs.shape) output = layer(keras.backend.variable(inputs)) - np_output = keras.backend.eval(output) + if context.executing_eagerly(): + np_output = output.numpy() + else: + np_output = keras.backend.eval(output) # compare with numpy if data_format == 'channels_first': expected_out = inputs[:, :, @@ -897,7 +954,7 @@ class CroppingTest(test.TestCase): cropping[2][0]:-cropping[2][1], :] np.testing.assert_allclose(np_output, expected_out) - # test incorrect use + # test incorrect use with self.assertRaises(ValueError): keras.layers.Cropping3D(cropping=(1, 1)) with self.assertRaises(ValueError): diff --git a/tensorflow/python/keras/_impl/keras/layers/core.py b/tensorflow/python/keras/_impl/keras/layers/core.py index 6ee3fb48b2f1426b87c5c1947e90d0797e9b9ff7..73e4f15f7e259211c892fdc663e14dcb14aec58d 100644 --- a/tensorflow/python/keras/_impl/keras/layers/core.py +++ b/tensorflow/python/keras/_impl/keras/layers/core.py @@ -23,6 +23,7 @@ import types as python_types import numpy as np +from tensorflow.python.eager import context from tensorflow.python.framework import tensor_shape from tensorflow.python.keras._impl.keras import activations from tensorflow.python.keras._impl.keras import backend as K @@ -36,8 +37,10 @@ from tensorflow.python.keras._impl.keras.utils.generic_utils import func_dump from tensorflow.python.keras._impl.keras.utils.generic_utils import func_load from tensorflow.python.keras._impl.keras.utils.generic_utils import has_arg from tensorflow.python.layers import core as tf_core_layers +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.Masking') class Masking(Layer): """Masks a sequence by using a mask value to skip timesteps. @@ -88,6 +91,7 @@ class Masking(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Dropout') class Dropout(tf_core_layers.Dropout, Layer): """Applies Dropout to the input. @@ -119,7 +123,8 @@ class Dropout(tf_core_layers.Dropout, Layer): if training is None: training = K.learning_phase() output = super(Dropout, self).call(inputs, training=training) - if training is K.learning_phase(): + # EagerTensor object has no attribute _uses_learning_phase + if not context.executing_eagerly() and training is K.learning_phase(): output._uses_learning_phase = True # pylint: disable=protected-access return output @@ -133,6 +138,7 @@ class Dropout(tf_core_layers.Dropout, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.SpatialDropout1D') class SpatialDropout1D(Dropout): """Spatial 1D version of Dropout. @@ -169,6 +175,7 @@ class SpatialDropout1D(Dropout): return noise_shape +@tf_export('keras.layers.SpatialDropout2D') class SpatialDropout2D(Dropout): """Spatial 2D version of Dropout. @@ -222,6 +229,7 @@ class SpatialDropout2D(Dropout): return (input_shape[0], 1, 1, input_shape[3]) +@tf_export('keras.layers.SpatialDropout3D') class SpatialDropout3D(Dropout): """Spatial 3D version of Dropout. @@ -274,6 +282,7 @@ class SpatialDropout3D(Dropout): return (input_shape[0], 1, 1, 1, input_shape[4]) +@tf_export('keras.layers.Activation') class Activation(Layer): """Applies an activation function to an output. @@ -307,6 +316,7 @@ class Activation(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Reshape') class Reshape(Layer): """Reshapes an output to a certain shape. @@ -412,6 +422,7 @@ class Reshape(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Permute') class Permute(Layer): """Permutes the dimensions of the input according to a given pattern. @@ -464,6 +475,7 @@ class Permute(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Flatten') class Flatten(tf_core_layers.Flatten, Layer): """Flattens the input. Does not affect the batch size. @@ -483,6 +495,7 @@ class Flatten(tf_core_layers.Flatten, Layer): pass +@tf_export('keras.layers.RepeatVector') class RepeatVector(Layer): """Repeats the input n times. @@ -526,6 +539,7 @@ class RepeatVector(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Lambda') class Lambda(Layer): """Wraps arbitrary expression as a `Layer` object. @@ -707,6 +721,7 @@ class Lambda(Layer): return cls(**config) +@tf_export('keras.layers.Dense') class Dense(tf_core_layers.Dense, Layer): """Just your regular densely-connected NN layer. @@ -811,6 +826,7 @@ class Dense(tf_core_layers.Dense, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.ActivityRegularization') class ActivityRegularization(Layer): """Layer that applies an update to the cost function based input activity. diff --git a/tensorflow/python/keras/_impl/keras/layers/core_test.py b/tensorflow/python/keras/_impl/keras/layers/core_test.py index bdb99c91c289cf808fec7b891376dbfcf5504aca..2ca816adbdcecaf371776d99f3da60d0d8790832 100644 --- a/tensorflow/python/keras/_impl/keras/layers/core_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/core_test.py @@ -20,11 +20,9 @@ from __future__ import print_function import numpy as np -from tensorflow.python.eager import context -from tensorflow.python.framework import constant_op +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils -from tensorflow.python.ops import init_ops from tensorflow.python.platform import test @@ -52,146 +50,134 @@ class CoreLayersTest(test.TestCase): dropout = keras.layers.Dropout(0.5) self.assertEqual(True, dropout.supports_masking) - with self.test_session(): - testing_utils.layer_test( - keras.layers.SpatialDropout1D, - kwargs={'rate': 0.5}, - input_shape=(2, 3, 4)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.SpatialDropout2D, - kwargs={'rate': 0.5}, - input_shape=(2, 3, 4, 5)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.SpatialDropout2D, - kwargs={'rate': 0.5, 'data_format': 'channels_first'}, - input_shape=(2, 3, 4, 5)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.SpatialDropout3D, - kwargs={'rate': 0.5}, - input_shape=(2, 3, 4, 4, 5)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.SpatialDropout3D, - kwargs={'rate': 0.5, 'data_format': 'channels_first'}, - input_shape=(2, 3, 4, 4, 5)) - + @tf_test_util.run_in_graph_and_eager_modes() + def test_spatial_dropout(self): + testing_utils.layer_test( + keras.layers.SpatialDropout1D, + kwargs={'rate': 0.5}, + input_shape=(2, 3, 4)) + + testing_utils.layer_test( + keras.layers.SpatialDropout2D, + kwargs={'rate': 0.5}, + input_shape=(2, 3, 4, 5)) + + testing_utils.layer_test( + keras.layers.SpatialDropout2D, + kwargs={'rate': 0.5, 'data_format': 'channels_first'}, + input_shape=(2, 3, 4, 5)) + + testing_utils.layer_test( + keras.layers.SpatialDropout3D, + kwargs={'rate': 0.5}, + input_shape=(2, 3, 4, 4, 5)) + + testing_utils.layer_test( + keras.layers.SpatialDropout3D, + kwargs={'rate': 0.5, 'data_format': 'channels_first'}, + input_shape=(2, 3, 4, 4, 5)) + + @tf_test_util.run_in_graph_and_eager_modes() def test_activation(self): # with string argument - with self.test_session(): - testing_utils.layer_test( - keras.layers.Activation, - kwargs={'activation': 'relu'}, - input_shape=(3, 2)) + testing_utils.layer_test( + keras.layers.Activation, + kwargs={'activation': 'relu'}, + input_shape=(3, 2)) # with function argument - with self.test_session(): - testing_utils.layer_test( - keras.layers.Activation, - kwargs={'activation': keras.backend.relu}, - input_shape=(3, 2)) + testing_utils.layer_test( + keras.layers.Activation, + kwargs={'activation': keras.backend.relu}, + input_shape=(3, 2)) + @tf_test_util.run_in_graph_and_eager_modes() def test_reshape(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.Reshape, - kwargs={'target_shape': (8, 1)}, - input_shape=(3, 2, 4)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.Reshape, - kwargs={'target_shape': (-1, 1)}, - input_shape=(3, 2, 4)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.Reshape, - kwargs={'target_shape': (1, -1)}, - input_shape=(3, 2, 4)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.Reshape, - kwargs={'target_shape': (-1, 1)}, - input_shape=(None, None, 2)) - + testing_utils.layer_test( + keras.layers.Reshape, + kwargs={'target_shape': (8, 1)}, + input_shape=(3, 2, 4)) + + testing_utils.layer_test( + keras.layers.Reshape, + kwargs={'target_shape': (-1, 1)}, + input_shape=(3, 2, 4)) + + testing_utils.layer_test( + keras.layers.Reshape, + kwargs={'target_shape': (1, -1)}, + input_shape=(3, 2, 4)) + + testing_utils.layer_test( + keras.layers.Reshape, + kwargs={'target_shape': (-1, 1)}, + input_shape=(None, None, 2)) + + @tf_test_util.run_in_graph_and_eager_modes() def test_permute(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.Permute, kwargs={'dims': (2, 1)}, input_shape=(3, 2, 4)) + testing_utils.layer_test( + keras.layers.Permute, kwargs={'dims': (2, 1)}, input_shape=(3, 2, 4)) + @tf_test_util.run_in_graph_and_eager_modes() def test_flatten(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.Flatten, kwargs={}, input_shape=(3, 2, 4)) + testing_utils.layer_test( + keras.layers.Flatten, kwargs={}, input_shape=(3, 2, 4)) + @tf_test_util.run_in_graph_and_eager_modes() def test_repeat_vector(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.RepeatVector, kwargs={'n': 3}, input_shape=(3, 2)) + testing_utils.layer_test( + keras.layers.RepeatVector, kwargs={'n': 3}, input_shape=(3, 2)) + @tf_test_util.run_in_graph_and_eager_modes() def test_lambda(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.Lambda, - kwargs={'function': lambda x: x + 1}, - input_shape=(3, 2)) - - with self.test_session(): - testing_utils.layer_test( - keras.layers.Lambda, - kwargs={ - 'function': lambda x, a, b: x * a + b, - 'arguments': { - 'a': 0.6, - 'b': 0.4 - } - }, - input_shape=(3, 2)) - - with self.test_session(): - # test serialization with function - def f(x): - return x + 1 - - ld = keras.layers.Lambda(f) - config = ld.get_config() - ld = keras.layers.deserialize({ - 'class_name': 'Lambda', - 'config': config - }) - - # test with lambda - ld = keras.layers.Lambda( - lambda x: keras.backend.concatenate([keras.backend.square(x), x])) - config = ld.get_config() - ld = keras.layers.Lambda.from_config(config) - + testing_utils.layer_test( + keras.layers.Lambda, + kwargs={'function': lambda x: x + 1}, + input_shape=(3, 2)) + + testing_utils.layer_test( + keras.layers.Lambda, + kwargs={ + 'function': lambda x, a, b: x * a + b, + 'arguments': { + 'a': 0.6, + 'b': 0.4 + } + }, + input_shape=(3, 2)) + + # test serialization with function + def f(x): + return x + 1 + + ld = keras.layers.Lambda(f) + config = ld.get_config() + ld = keras.layers.deserialize({ + 'class_name': 'Lambda', + 'config': config + }) + + # test with lambda + ld = keras.layers.Lambda( + lambda x: keras.backend.concatenate([keras.backend.square(x), x])) + config = ld.get_config() + ld = keras.layers.Lambda.from_config(config) + + @tf_test_util.run_in_graph_and_eager_modes() def test_dense(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 2)) + testing_utils.layer_test( + keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 2)) - with self.test_session(): - testing_utils.layer_test( - keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 4, 2)) + testing_utils.layer_test( + keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 4, 2)) - with self.test_session(): - testing_utils.layer_test( - keras.layers.Dense, kwargs={'units': 3}, input_shape=(None, None, 2)) + testing_utils.layer_test( + keras.layers.Dense, kwargs={'units': 3}, input_shape=(None, None, 2)) - with self.test_session(): - testing_utils.layer_test( - keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 4, 5, 2)) + testing_utils.layer_test( + keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 4, 5, 2)) - # Test regularization + def test_dense_regularization(self): with self.test_session(): layer = keras.layers.Dense( 3, @@ -202,7 +188,7 @@ class CoreLayersTest(test.TestCase): layer(keras.backend.variable(np.ones((2, 4)))) self.assertEqual(3, len(layer.losses)) - # Test constraints + def test_dense_constraints(self): with self.test_session(): k_constraint = keras.constraints.max_norm(0.01) b_constraint = keras.constraints.max_norm(0.01) @@ -212,12 +198,6 @@ class CoreLayersTest(test.TestCase): self.assertEqual(layer.kernel.constraint, k_constraint) self.assertEqual(layer.bias.constraint, b_constraint) - def test_eager_dense(self): - with context.eager_mode(): - l = keras.layers.Dense(units=3, - kernel_initializer=init_ops.zeros_initializer()) - self.assertAllEqual(l(constant_op.constant([[1.0]])), [[0., 0., 0.]]) - def test_activity_regularization(self): with self.test_session(): layer = keras.layers.ActivityRegularization(l1=0.1) diff --git a/tensorflow/python/keras/_impl/keras/layers/embeddings.py b/tensorflow/python/keras/_impl/keras/layers/embeddings.py index f8e31068f8910ff44889963c8c301c39844d85cb..006ecd3135be25d43133daed1603734ecd1be955 100644 --- a/tensorflow/python/keras/_impl/keras/layers/embeddings.py +++ b/tensorflow/python/keras/_impl/keras/layers/embeddings.py @@ -23,9 +23,11 @@ from tensorflow.python.keras._impl.keras import constraints from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers from tensorflow.python.keras._impl.keras.engine import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.Embedding') class Embedding(Layer): """Turns positive integers (indexes) into dense vectors of fixed size. diff --git a/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py b/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py index 1712111b877cf1fee4353c5542f33a973a26de95..26fd1f1c114587c2f1b3e0155f1259dd5f0dcf60 100644 --- a/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test @@ -25,47 +26,44 @@ from tensorflow.python.platform import test class EmbeddingTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_embedding(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.Embedding, - kwargs={'output_dim': 4, - 'input_dim': 10, - 'input_length': 2}, - input_shape=(3, 2), - input_dtype='int32', - expected_output_dtype='float32') + testing_utils.layer_test( + keras.layers.Embedding, + kwargs={'output_dim': 4, + 'input_dim': 10, + 'input_length': 2}, + input_shape=(3, 2), + input_dtype='int32', + expected_output_dtype='float32') - with self.test_session(): - testing_utils.layer_test( - keras.layers.Embedding, - kwargs={'output_dim': 4, - 'input_dim': 10, - 'mask_zero': True}, - input_shape=(3, 2), - input_dtype='int32', - expected_output_dtype='float32') + testing_utils.layer_test( + keras.layers.Embedding, + kwargs={'output_dim': 4, + 'input_dim': 10, + 'mask_zero': True}, + input_shape=(3, 2), + input_dtype='int32', + expected_output_dtype='float32') - with self.test_session(): - testing_utils.layer_test( - keras.layers.Embedding, - kwargs={'output_dim': 4, - 'input_dim': 10, - 'mask_zero': True}, - input_shape=(3, 4, 2), - input_dtype='int32', - expected_output_dtype='float32') + testing_utils.layer_test( + keras.layers.Embedding, + kwargs={'output_dim': 4, + 'input_dim': 10, + 'mask_zero': True}, + input_shape=(3, 4, 2), + input_dtype='int32', + expected_output_dtype='float32') - with self.test_session(): - testing_utils.layer_test( - keras.layers.Embedding, - kwargs={'output_dim': 4, - 'input_dim': 10, - 'mask_zero': True, - 'input_length': (None, 2)}, - input_shape=(3, 4, 2), - input_dtype='int32', - expected_output_dtype='float32') + testing_utils.layer_test( + keras.layers.Embedding, + kwargs={'output_dim': 4, + 'input_dim': 10, + 'mask_zero': True, + 'input_length': (None, 2)}, + input_shape=(3, 4, 2), + input_dtype='int32', + expected_output_dtype='float32') if __name__ == '__main__': diff --git a/tensorflow/python/keras/_impl/keras/layers/gru_test.py b/tensorflow/python/keras/_impl/keras/layers/gru_test.py index c57fbac41cc43995ef3249414ed03928e7ffd044..48e7e14f5ab73b534ab0d1c765ad2572b2930b2b 100644 --- a/tensorflow/python/keras/_impl/keras/layers/gru_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/gru_test.py @@ -20,64 +20,66 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test +from tensorflow.python.training.rmsprop import RMSPropOptimizer class GRULayerTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_return_sequences_GRU(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - testing_utils.layer_test( - keras.layers.GRU, - kwargs={'units': units, - 'return_sequences': True}, - input_shape=(num_samples, timesteps, embedding_dim)) + testing_utils.layer_test( + keras.layers.GRU, + kwargs={'units': units, + 'return_sequences': True}, + input_shape=(num_samples, timesteps, embedding_dim)) + @tf_test_util.run_in_graph_and_eager_modes() def test_dynamic_behavior_GRU(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - layer = keras.layers.GRU(units, input_shape=(None, embedding_dim)) - model = keras.models.Sequential() - model.add(layer) - model.compile('sgd', 'mse') - x = np.random.random((num_samples, timesteps, embedding_dim)) - y = np.random.random((num_samples, units)) - model.train_on_batch(x, y) - + layer = keras.layers.GRU(units, input_shape=(None, embedding_dim)) + model = keras.models.Sequential() + model.add(layer) + model.compile(RMSPropOptimizer(0.01), 'mse') + x = np.random.random((num_samples, timesteps, embedding_dim)) + y = np.random.random((num_samples, units)) + model.train_on_batch(x, y) + + @tf_test_util.run_in_graph_and_eager_modes() def test_dropout_GRU(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - testing_utils.layer_test( - keras.layers.GRU, - kwargs={'units': units, - 'dropout': 0.1, - 'recurrent_dropout': 0.1}, - input_shape=(num_samples, timesteps, embedding_dim)) - + testing_utils.layer_test( + keras.layers.GRU, + kwargs={'units': units, + 'dropout': 0.1, + 'recurrent_dropout': 0.1}, + input_shape=(num_samples, timesteps, embedding_dim)) + + @tf_test_util.run_in_graph_and_eager_modes() def test_implementation_mode_GRU(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - for mode in [0, 1, 2]: - testing_utils.layer_test( - keras.layers.GRU, - kwargs={'units': units, - 'implementation': mode}, - input_shape=(num_samples, timesteps, embedding_dim)) + for mode in [0, 1, 2]: + testing_utils.layer_test( + keras.layers.GRU, + kwargs={'units': units, + 'implementation': mode}, + input_shape=(num_samples, timesteps, embedding_dim)) def test_statefulness_GRU(self): num_samples = 2 diff --git a/tensorflow/python/keras/_impl/keras/layers/local.py b/tensorflow/python/keras/_impl/keras/layers/local.py index b844b071e02b4f8f217a09a7d412609f6e2cadeb..13d96e939220c11a4090cf535e3efa4365fe8b62 100644 --- a/tensorflow/python/keras/_impl/keras/layers/local.py +++ b/tensorflow/python/keras/_impl/keras/layers/local.py @@ -25,10 +25,12 @@ from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers from tensorflow.python.keras._impl.keras.engine import InputSpec from tensorflow.python.keras._impl.keras.engine import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion from tensorflow.python.keras._impl.keras.utils import conv_utils +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.LocallyConnected1D') class LocallyConnected1D(Layer): """Locally-connected layer for 1D inputs. @@ -51,7 +53,7 @@ class LocallyConnected1D(Layer): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output 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, @@ -193,6 +195,7 @@ class LocallyConnected1D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.LocallyConnected2D') class LocallyConnected2D(Layer): """Locally-connected layer for 2D inputs. @@ -219,7 +222,7 @@ class LocallyConnected2D(Layer): Arguments: filters: Integer, the dimensionality of the output space - (i.e. the number output of filters in the convolution). + (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of 2 integers, specifying the width and height of the 2D convolution window. Can be a single integer to specify the same value for diff --git a/tensorflow/python/keras/_impl/keras/layers/local_test.py b/tensorflow/python/keras/_impl/keras/layers/local_test.py index a815a0fadc8215c00f3db4749e323f96e44b66f3..93741d24b9a74cf9e8a83069f7c4235b1f489818 100644 --- a/tensorflow/python/keras/_impl/keras/layers/local_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/local_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test @@ -27,6 +28,7 @@ from tensorflow.python.platform import test class LocallyConnectedLayersTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_locallyconnected_1d(self): num_samples = 2 num_steps = 8 @@ -39,16 +41,15 @@ class LocallyConnectedLayersTest(test.TestCase): if padding == 'same' and strides != 1: continue - with self.test_session(): - testing_utils.layer_test( - keras.layers.LocallyConnected1D, - kwargs={ - 'filters': filters, - 'kernel_size': filter_length, - 'padding': padding, - 'strides': strides - }, - input_shape=(num_samples, num_steps, input_dim)) + testing_utils.layer_test( + keras.layers.LocallyConnected1D, + kwargs={ + 'filters': filters, + 'kernel_size': filter_length, + 'padding': padding, + 'strides': strides + }, + input_shape=(num_samples, num_steps, input_dim)) def test_locallyconnected_1d_regularization(self): num_samples = 2 @@ -86,6 +87,7 @@ class LocallyConnectedLayersTest(test.TestCase): self.assertEqual(layer.kernel.constraint, k_constraint) self.assertEqual(layer.bias.constraint, b_constraint) + @tf_test_util.run_in_graph_and_eager_modes() def test_locallyconnected_2d(self): num_samples = 8 filters = 3 @@ -98,20 +100,18 @@ class LocallyConnectedLayersTest(test.TestCase): if padding == 'same' and strides != (1, 1): continue - with self.test_session(): - testing_utils.layer_test( - keras.layers.LocallyConnected2D, - kwargs={ - 'filters': filters, - 'kernel_size': 3, - 'padding': padding, - 'kernel_regularizer': 'l2', - 'bias_regularizer': 'l2', - 'activity_regularizer': 'l2', - 'strides': strides, - 'data_format': 'channels_last' - }, - input_shape=(num_samples, num_row, num_col, stack_size)) + testing_utils.layer_test( + keras.layers.LocallyConnected2D, + kwargs={ + 'filters': filters, + 'kernel_size': 3, + 'padding': padding, + 'kernel_regularizer': 'l2', + 'bias_regularizer': 'l2', + 'strides': strides, + 'data_format': 'channels_last' + }, + input_shape=(num_samples, num_row, num_col, stack_size)) def test_locallyconnected_2d_channels_first(self): num_samples = 8 diff --git a/tensorflow/python/keras/_impl/keras/layers/lstm_test.py b/tensorflow/python/keras/_impl/keras/layers/lstm_test.py index 8d359bf17cdb80c98aeeed6d69e301962609ce59..11a5e0aeaacfa7520361ae41ac3d40607e8a9050 100644 --- a/tensorflow/python/keras/_impl/keras/layers/lstm_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/lstm_test.py @@ -20,64 +20,81 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test +from tensorflow.python.training.rmsprop import RMSPropOptimizer class LSTMLayerTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_return_sequences_LSTM(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - testing_utils.layer_test( - keras.layers.LSTM, - kwargs={'units': units, - 'return_sequences': True}, - input_shape=(num_samples, timesteps, embedding_dim)) + testing_utils.layer_test( + keras.layers.LSTM, + kwargs={'units': units, + 'return_sequences': True}, + input_shape=(num_samples, timesteps, embedding_dim)) + + def test_static_shape_inference_LSTM(self): + # Github issue: 15165 + timesteps = 3 + embedding_dim = 4 + units = 2 + + model = keras.models.Sequential() + inputs = keras.layers.Dense(embedding_dim, + input_shape=(timesteps, embedding_dim)) + model.add(inputs) + layer = keras.layers.LSTM(units, return_sequences=True) + model.add(layer) + outputs = model.layers[-1].output + self.assertEquals(outputs.get_shape().as_list(), [None, timesteps, units]) + @tf_test_util.run_in_graph_and_eager_modes() def test_dynamic_behavior_LSTM(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - layer = keras.layers.LSTM(units, input_shape=(None, embedding_dim)) - model = keras.models.Sequential() - model.add(layer) - model.compile('sgd', 'mse') - x = np.random.random((num_samples, timesteps, embedding_dim)) - y = np.random.random((num_samples, units)) - model.train_on_batch(x, y) - + layer = keras.layers.LSTM(units, input_shape=(None, embedding_dim)) + model = keras.models.Sequential() + model.add(layer) + model.compile(RMSPropOptimizer(0.001), 'mse') + x = np.random.random((num_samples, timesteps, embedding_dim)) + y = np.random.random((num_samples, units)) + model.train_on_batch(x, y) + + @tf_test_util.run_in_graph_and_eager_modes() def test_dropout_LSTM(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - testing_utils.layer_test( - keras.layers.LSTM, - kwargs={'units': units, - 'dropout': 0.1, - 'recurrent_dropout': 0.1}, - input_shape=(num_samples, timesteps, embedding_dim)) - + testing_utils.layer_test( + keras.layers.LSTM, + kwargs={'units': units, + 'dropout': 0.1, + 'recurrent_dropout': 0.1}, + input_shape=(num_samples, timesteps, embedding_dim)) + + @tf_test_util.run_in_graph_and_eager_modes() def test_implementation_mode_LSTM(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - for mode in [0, 1, 2]: - testing_utils.layer_test( - keras.layers.LSTM, - kwargs={'units': units, - 'implementation': mode}, - input_shape=(num_samples, timesteps, embedding_dim)) + for mode in [0, 1, 2]: + testing_utils.layer_test( + keras.layers.LSTM, + kwargs={'units': units, + 'implementation': mode}, + input_shape=(num_samples, timesteps, embedding_dim)) def test_statefulness_LSTM(self): num_samples = 2 diff --git a/tensorflow/python/keras/_impl/keras/layers/merge.py b/tensorflow/python/keras/_impl/keras/layers/merge.py index 38b0b302972a31fb29b516680470c587d75a0781..c660cbd449b11a139f64cfa8b3a35310a597491c 100644 --- a/tensorflow/python/keras/_impl/keras/layers/merge.py +++ b/tensorflow/python/keras/_impl/keras/layers/merge.py @@ -21,8 +21,9 @@ from __future__ import division from __future__ import print_function from tensorflow.python.keras._impl.keras import backend as K -from tensorflow.python.keras._impl.keras.engine.topology import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import Layer +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion +from tensorflow.python.util.tf_export import tf_export class _Merge(Layer): @@ -210,6 +211,7 @@ class _Merge(Layer): return K.all(K.concatenate(masks, axis=0), axis=0, keepdims=False) +@tf_export('keras.layers.Add') class Add(_Merge): """Layer that adds a list of inputs. @@ -279,6 +281,7 @@ class Subtract(_Merge): return inputs[0] - inputs[1] +@tf_export('keras.layers.Multiply') class Multiply(_Merge): """Layer that multiplies (element-wise) a list of inputs. @@ -294,6 +297,7 @@ class Multiply(_Merge): return output +@tf_export('keras.layers.Average') class Average(_Merge): """Layer that averages a list of inputs. @@ -309,6 +313,7 @@ class Average(_Merge): return output / len(inputs) +@tf_export('keras.layers.Maximum') class Maximum(_Merge): """Layer that computes the maximum (element-wise) a list of inputs. @@ -339,6 +344,7 @@ class Minimum(_Merge): return output +@tf_export('keras.layers.Concatenate') class Concatenate(_Merge): """Layer that concatenates a list of inputs. @@ -429,6 +435,7 @@ class Concatenate(_Merge): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.Dot') class Dot(_Merge): """Layer that computes a dot product between samples in two tensors. @@ -543,6 +550,7 @@ class Dot(_Merge): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.add') def add(inputs, **kwargs): """Functional interface to the `Add` layer. @@ -599,6 +607,7 @@ def subtract(inputs, **kwargs): return Subtract(**kwargs)(inputs) +@tf_export('keras.layers.multiply') def multiply(inputs, **kwargs): """Functional interface to the `Multiply` layer. @@ -612,6 +621,7 @@ def multiply(inputs, **kwargs): return Multiply(**kwargs)(inputs) +@tf_export('keras.layers.average') def average(inputs, **kwargs): """Functional interface to the `Average` layer. @@ -625,6 +635,7 @@ def average(inputs, **kwargs): return Average(**kwargs)(inputs) +@tf_export('keras.layers.maximum') def maximum(inputs, **kwargs): """Functional interface to the `Maximum` layer. @@ -651,6 +662,7 @@ def minimum(inputs, **kwargs): return Minimum(**kwargs)(inputs) +@tf_export('keras.layers.concatenate') def concatenate(inputs, axis=-1, **kwargs): """Functional interface to the `Concatenate` layer. @@ -665,6 +677,7 @@ def concatenate(inputs, axis=-1, **kwargs): return Concatenate(axis=axis, **kwargs)(inputs) +@tf_export('keras.layers.dot') def dot(inputs, axes, normalize=False, **kwargs): """Functional interface to the `Dot` layer. diff --git a/tensorflow/python/keras/_impl/keras/layers/merge_test.py b/tensorflow/python/keras/_impl/keras/layers/merge_test.py index bb03dda1fc645222c1ced97cfce8d459586dd89d..b2fe06f93e33ed63d6a2aa29522ecb552f582440 100644 --- a/tensorflow/python/keras/_impl/keras/layers/merge_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/merge_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.ops import array_ops from tensorflow.python.platform import test @@ -27,24 +28,25 @@ from tensorflow.python.platform import test class MergeLayersTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_add(self): - with self.test_session(): - i1 = keras.layers.Input(shape=(4, 5)) - i2 = keras.layers.Input(shape=(4, 5)) - i3 = keras.layers.Input(shape=(4, 5)) + i1 = keras.layers.Input(shape=(4, 5)) + i2 = keras.layers.Input(shape=(4, 5)) + i3 = keras.layers.Input(shape=(4, 5)) - o = keras.layers.add([i1, i2, i3]) - self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) - model = keras.models.Model([i1, i2, i3], o) + o = keras.layers.add([i1, i2, i3]) + self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) + model = keras.models.Model([i1, i2, i3], o) - x1 = np.random.random((2, 4, 5)) - x2 = np.random.random((2, 4, 5)) - x3 = np.random.random((2, 4, 5)) - out = model.predict([x1, x2, x3]) - self.assertEqual(out.shape, (2, 4, 5)) - self.assertAllClose(out, x1 + x2 + x3, atol=1e-4) + x1 = np.random.random((2, 4, 5)) + x2 = np.random.random((2, 4, 5)) + x3 = np.random.random((2, 4, 5)) + out = model.predict([x1, x2, x3]) + self.assertEqual(out.shape, (2, 4, 5)) + self.assertAllClose(out, x1 + x2 + x3, atol=1e-4) - # test masking + def test_merge_add_masking(self): + with self.test_session(): i1 = keras.layers.Input(shape=(4, 5)) i2 = keras.layers.Input(shape=(4, 5)) m1 = keras.layers.Masking()(i1) @@ -54,11 +56,13 @@ class MergeLayersTest(test.TestCase): mask = layer.output_mask self.assertListEqual(mask.get_shape().as_list(), [None, 4]) - # test missing shape + def test_merge_add_dynamic_shape(self): + with self.test_session(): i1 = array_ops.placeholder(shape=(4, None), dtype='float32') i2 = array_ops.placeholder(shape=(4, 5), dtype='float32') layer = keras.layers.Add() o = layer([i1, i2]) + self.assertListEqual(o.get_shape().as_list(), [4, 5]) def test_merge_elementwise_errors(self): i1 = keras.layers.Input(shape=(4, 5)) @@ -72,79 +76,82 @@ class MergeLayersTest(test.TestCase): with self.assertRaises(ValueError): keras.layers.add([i1]) + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_multiply(self): - with self.test_session(): - i1 = keras.layers.Input(shape=(4, 5)) - i2 = keras.layers.Input(shape=(4, 5)) - i3 = keras.layers.Input(shape=(4, 5)) - o = keras.layers.multiply([i1, i2, i3]) - self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) - model = keras.models.Model([i1, i2, i3], o) - - x1 = np.random.random((2, 4, 5)) - x2 = np.random.random((2, 4, 5)) - x3 = np.random.random((2, 4, 5)) - out = model.predict([x1, x2, x3]) - self.assertEqual(out.shape, (2, 4, 5)) - self.assertAllClose(out, x1 * x2 * x3, atol=1e-4) - + i1 = keras.layers.Input(shape=(4, 5)) + i2 = keras.layers.Input(shape=(4, 5)) + i3 = keras.layers.Input(shape=(4, 5)) + o = keras.layers.multiply([i1, i2, i3]) + self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) + model = keras.models.Model([i1, i2, i3], o) + + x1 = np.random.random((2, 4, 5)) + x2 = np.random.random((2, 4, 5)) + x3 = np.random.random((2, 4, 5)) + out = model.predict([x1, x2, x3]) + self.assertEqual(out.shape, (2, 4, 5)) + self.assertAllClose(out, x1 * x2 * x3, atol=1e-4) + + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_average(self): - with self.test_session(): - i1 = keras.layers.Input(shape=(4, 5)) - i2 = keras.layers.Input(shape=(4, 5)) - o = keras.layers.average([i1, i2]) - self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) - model = keras.models.Model([i1, i2], o) + i1 = keras.layers.Input(shape=(4, 5)) + i2 = keras.layers.Input(shape=(4, 5)) + o = keras.layers.average([i1, i2]) + self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) + model = keras.models.Model([i1, i2], o) - x1 = np.random.random((2, 4, 5)) - x2 = np.random.random((2, 4, 5)) - out = model.predict([x1, x2]) - self.assertEqual(out.shape, (2, 4, 5)) - self.assertAllClose(out, 0.5 * (x1 + x2), atol=1e-4) + x1 = np.random.random((2, 4, 5)) + x2 = np.random.random((2, 4, 5)) + out = model.predict([x1, x2]) + self.assertEqual(out.shape, (2, 4, 5)) + self.assertAllClose(out, 0.5 * (x1 + x2), atol=1e-4) + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_maximum(self): - with self.test_session(): - i1 = keras.layers.Input(shape=(4, 5)) - i2 = keras.layers.Input(shape=(4, 5)) - o = keras.layers.maximum([i1, i2]) - self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) - model = keras.models.Model([i1, i2], o) + i1 = keras.layers.Input(shape=(4, 5)) + i2 = keras.layers.Input(shape=(4, 5)) + o = keras.layers.maximum([i1, i2]) + self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) + model = keras.models.Model([i1, i2], o) - x1 = np.random.random((2, 4, 5)) - x2 = np.random.random((2, 4, 5)) - out = model.predict([x1, x2]) - self.assertEqual(out.shape, (2, 4, 5)) - self.assertAllClose(out, np.maximum(x1, x2), atol=1e-4) + x1 = np.random.random((2, 4, 5)) + x2 = np.random.random((2, 4, 5)) + out = model.predict([x1, x2]) + self.assertEqual(out.shape, (2, 4, 5)) + self.assertAllClose(out, np.maximum(x1, x2), atol=1e-4) + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_minimum(self): - with self.test_session(): - i1 = keras.layers.Input(shape=(4, 5)) - i2 = keras.layers.Input(shape=(4, 5)) - o = keras.layers.minimum([i1, i2]) - self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) - model = keras.models.Model([i1, i2], o) + i1 = keras.layers.Input(shape=(4, 5)) + i2 = keras.layers.Input(shape=(4, 5)) + o = keras.layers.minimum([i1, i2]) + self.assertListEqual(o.get_shape().as_list(), [None, 4, 5]) + model = keras.models.Model([i1, i2], o) - x1 = np.random.random((2, 4, 5)) - x2 = np.random.random((2, 4, 5)) - out = model.predict([x1, x2]) - self.assertEqual(out.shape, (2, 4, 5)) - self.assertAllClose(out, np.minimum(x1, x2), atol=1e-4) + x1 = np.random.random((2, 4, 5)) + x2 = np.random.random((2, 4, 5)) + out = model.predict([x1, x2]) + self.assertEqual(out.shape, (2, 4, 5)) + self.assertAllClose(out, np.minimum(x1, x2), atol=1e-4) + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_concatenate(self): + i1 = keras.layers.Input(shape=(4, 5)) + i2 = keras.layers.Input(shape=(4, 5)) + o = keras.layers.concatenate([i1, i2], axis=1) + self.assertListEqual(o.get_shape().as_list(), [None, 8, 5]) + model = keras.models.Model([i1, i2], o) + + x1 = np.random.random((2, 4, 5)) + x2 = np.random.random((2, 4, 5)) + out = model.predict([x1, x2]) + self.assertEqual(out.shape, (2, 8, 5)) + self.assertAllClose(out, np.concatenate([x1, x2], axis=1), atol=1e-4) + + def test_merge_concatenate_masking(self): with self.test_session(): i1 = keras.layers.Input(shape=(4, 5)) i2 = keras.layers.Input(shape=(4, 5)) - o = keras.layers.concatenate([i1, i2], axis=1) - self.assertListEqual(o.get_shape().as_list(), [None, 8, 5]) - model = keras.models.Model([i1, i2], o) - - x1 = np.random.random((2, 4, 5)) - x2 = np.random.random((2, 4, 5)) - out = model.predict([x1, x2]) - self.assertEqual(out.shape, (2, 8, 5)) - self.assertAllClose(out, np.concatenate([x1, x2], axis=1), atol=1e-4) - - # test masking m1 = keras.layers.Masking()(i1) layer = keras.layers.Concatenate() o = layer([m1, i2]) @@ -162,35 +169,35 @@ class MergeLayersTest(test.TestCase): with self.assertRaisesRegexp(ValueError, 'called on a list'): keras.layers.concatenate([i1], axis=-1) + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_dot(self): - with self.test_session(): - i1 = keras.layers.Input(shape=(4,)) - i2 = keras.layers.Input(shape=(4,)) - o = keras.layers.dot([i1, i2], axes=1) - self.assertListEqual(o.get_shape().as_list(), [None, 1]) - model = keras.models.Model([i1, i2], o) - _ = keras.layers.Dot(axes=1).get_config() - - x1 = np.random.random((2, 4)) - x2 = np.random.random((2, 4)) - out = model.predict([x1, x2]) - self.assertEqual(out.shape, (2, 1)) - expected = np.zeros((2, 1)) - expected[0, 0] = np.dot(x1[0], x2[0]) - expected[1, 0] = np.dot(x1[1], x2[1]) - self.assertAllClose(out, expected, atol=1e-4) - - # Test with negative tuple of axes. - o = keras.layers.dot([i1, i2], axes=(-1, -1)) - self.assertListEqual(o.get_shape().as_list(), [None, 1]) - model = keras.models.Model([i1, i2], o) - out = model.predict([x1, x2]) - self.assertEqual(out.shape, (2, 1)) - self.assertAllClose(out, expected, atol=1e-4) - - # test compute_output_shape - layer = keras.layers.Dot(axes=-1) - self.assertEqual(layer.compute_output_shape([(4, 5), (4, 5)]), (4, 1)) + i1 = keras.layers.Input(shape=(4,)) + i2 = keras.layers.Input(shape=(4,)) + o = keras.layers.dot([i1, i2], axes=1) + self.assertListEqual(o.get_shape().as_list(), [None, 1]) + model = keras.models.Model([i1, i2], o) + _ = keras.layers.Dot(axes=1).get_config() + + x1 = np.random.random((2, 4)) + x2 = np.random.random((2, 4)) + out = model.predict([x1, x2]) + self.assertEqual(out.shape, (2, 1)) + expected = np.zeros((2, 1)) + expected[0, 0] = np.dot(x1[0], x2[0]) + expected[1, 0] = np.dot(x1[1], x2[1]) + self.assertAllClose(out, expected, atol=1e-4) + + # Test with negative tuple of axes. + o = keras.layers.dot([i1, i2], axes=(-1, -1)) + self.assertListEqual(o.get_shape().as_list(), [None, 1]) + model = keras.models.Model([i1, i2], o) + out = model.predict([x1, x2]) + self.assertEqual(out.shape, (2, 1)) + self.assertAllClose(out, expected, atol=1e-4) + + # test compute_output_shape + layer = keras.layers.Dot(axes=-1) + self.assertEqual(layer.compute_output_shape([(4, 5), (4, 5)]), (4, 1)) def test_dot_errors(self): i1 = keras.layers.Input(shape=(4, 5)) @@ -208,6 +215,7 @@ class MergeLayersTest(test.TestCase): dot = keras.layers.Dot(1) dot.compute_output_shape(1) + @tf_test_util.run_in_graph_and_eager_modes() def test_merge_subtract(self): i1 = keras.layers.Input(shape=(4, 5)) i2 = keras.layers.Input(shape=(4, 5)) diff --git a/tensorflow/python/keras/_impl/keras/layers/noise.py b/tensorflow/python/keras/_impl/keras/layers/noise.py index 04fffcc384cc4d868937d92bbd5a8f6505ca1770..e309d160e5a9be97ff5f5356dad9dfaf85430233 100644 --- a/tensorflow/python/keras/_impl/keras/layers/noise.py +++ b/tensorflow/python/keras/_impl/keras/layers/noise.py @@ -22,9 +22,11 @@ import numpy as np from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.engine import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.GaussianNoise') class GaussianNoise(Layer): """Apply additive zero-centered Gaussian noise. @@ -70,6 +72,7 @@ class GaussianNoise(Layer): return input_shape +@tf_export('keras.layers.GaussianDropout') class GaussianDropout(Layer): """Apply multiplicative 1-centered Gaussian noise. @@ -116,6 +119,7 @@ class GaussianDropout(Layer): return input_shape +@tf_export('keras.layers.AlphaDropout') class AlphaDropout(Layer): """Applies Alpha Dropout to the input. diff --git a/tensorflow/python/keras/_impl/keras/layers/noise_test.py b/tensorflow/python/keras/_impl/keras/layers/noise_test.py index f9b4d9cd090ffec1a5acd9118ea6a65798bd72a6..af4f031ec95bb56b72c1f1018e0e529d8ff55564 100644 --- a/tensorflow/python/keras/_impl/keras/layers/noise_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/noise_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test @@ -39,12 +40,12 @@ class NoiseLayersTest(test.TestCase): kwargs={'rate': 0.5}, input_shape=(3, 2, 3)) + @tf_test_util.run_in_graph_and_eager_modes() def test_AlphaDropout(self): - with self.test_session(): - testing_utils.layer_test( - keras.layers.AlphaDropout, - kwargs={'rate': 0.2}, - input_shape=(3, 2, 3)) + testing_utils.layer_test( + keras.layers.AlphaDropout, + kwargs={'rate': 0.2}, + input_shape=(3, 2, 3)) if __name__ == '__main__': diff --git a/tensorflow/python/keras/_impl/keras/layers/normalization.py b/tensorflow/python/keras/_impl/keras/layers/normalization.py index 965ef70e6e6cb488aa4832462da4a2cb43e964a6..3b44b20bf822429351002c0f81fe8f9596d595d3 100644 --- a/tensorflow/python/keras/_impl/keras/layers/normalization.py +++ b/tensorflow/python/keras/_impl/keras/layers/normalization.py @@ -18,14 +18,17 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.eager import context from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras import constraints from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers from tensorflow.python.keras._impl.keras.engine import Layer from tensorflow.python.layers import normalization as tf_normalization_layers +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.BatchNormalization') class BatchNormalization(tf_normalization_layers.BatchNormalization, Layer): """Batch normalization layer (Ioffe and Szegedy, 2014). @@ -108,7 +111,7 @@ class BatchNormalization(tf_normalization_layers.BatchNormalization, Layer): if training is None: training = K.learning_phase() output = super(BatchNormalization, self).call(inputs, training=training) - if training is K.learning_phase(): + if not context.executing_eagerly() and training is K.learning_phase(): output._uses_learning_phase = True # pylint: disable=protected-access return output diff --git a/tensorflow/python/keras/_impl/keras/layers/normalization_test.py b/tensorflow/python/keras/_impl/keras/layers/normalization_test.py index 39a90e597089b30d110f26f074eba5d6895e52df..2b3628c3f1023612297465bdf3286246261992a2 100644 --- a/tensorflow/python/keras/_impl/keras/layers/normalization_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/normalization_test.py @@ -132,13 +132,19 @@ class NormalizationLayersTest(test.TestCase): model.compile('sgd', 'mse') model.train_on_batch(x, x) - assert len(model.updates) == 2 + self.assertEqual(len(bn.updates), 4) + self.assertEqual(len(model.updates), 2) + self.assertEqual(len(model.get_updates_for(x1)), 0) + self.assertEqual(len(model.get_updates_for(x2)), 2) # Test model-level reuse x3 = keras.layers.Input(shape=(10,)) y3 = model(x3) - new_model = keras.models.Model(x3, y3) - assert len(model.updates) == 2 + new_model = keras.models.Model(x3, y3, name='new_model') + + self.assertEqual(len(new_model.updates), 2) + self.assertEqual(len(model.updates), 4) + self.assertEqual(len(new_model.get_updates_for(x3)), 2) new_model.compile('sgd', 'mse') new_model.train_on_batch(x, x) diff --git a/tensorflow/python/keras/_impl/keras/layers/pooling.py b/tensorflow/python/keras/_impl/keras/layers/pooling.py index b133e2dfaf1bcacd055f6a597bd557f696469ffc..15d53379769d8142f5b2755a07479f60751346d2 100644 --- a/tensorflow/python/keras/_impl/keras/layers/pooling.py +++ b/tensorflow/python/keras/_impl/keras/layers/pooling.py @@ -24,8 +24,10 @@ from tensorflow.python.keras._impl.keras.engine import InputSpec from tensorflow.python.keras._impl.keras.engine import Layer from tensorflow.python.keras._impl.keras.utils import conv_utils from tensorflow.python.layers import pooling as tf_pooling_layers +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.MaxPool1D', 'keras.layers.MaxPooling1D') class MaxPooling1D(tf_pooling_layers.MaxPooling1D, Layer): """Max pooling operation for temporal data. @@ -58,6 +60,7 @@ class MaxPooling1D(tf_pooling_layers.MaxPooling1D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.AveragePooling1D', 'keras.layers.AvgPool1D') class AveragePooling1D(tf_pooling_layers.AveragePooling1D, Layer): """Average pooling for temporal data. @@ -91,6 +94,7 @@ class AveragePooling1D(tf_pooling_layers.AveragePooling1D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.MaxPool2D', 'keras.layers.MaxPooling2D') class MaxPooling2D(tf_pooling_layers.MaxPooling2D, Layer): """Max pooling operation for spatial data. @@ -156,6 +160,7 @@ class MaxPooling2D(tf_pooling_layers.MaxPooling2D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.AveragePooling2D', 'keras.layers.AvgPool2D') class AveragePooling2D(tf_pooling_layers.AveragePooling2D, Layer): """Average pooling operation for spatial data. @@ -221,6 +226,7 @@ class AveragePooling2D(tf_pooling_layers.AveragePooling2D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.MaxPool3D', 'keras.layers.MaxPooling3D') class MaxPooling3D(tf_pooling_layers.MaxPooling3D, Layer): """Max pooling operation for 3D data (spatial or spatio-temporal). @@ -282,6 +288,7 @@ class MaxPooling3D(tf_pooling_layers.MaxPooling3D, Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.AveragePooling3D', 'keras.layers.AvgPool3D') class AveragePooling3D(tf_pooling_layers.AveragePooling3D, Layer): """Average pooling operation for 3D data (spatial or spatio-temporal). @@ -359,6 +366,8 @@ class _GlobalPooling1D(Layer): raise NotImplementedError +@tf_export('keras.layers.GlobalAveragePooling1D', + 'keras.layers.GlobalAvgPool1D') class GlobalAveragePooling1D(_GlobalPooling1D): """Global average pooling operation for temporal data. @@ -374,6 +383,7 @@ class GlobalAveragePooling1D(_GlobalPooling1D): return K.mean(inputs, axis=1) +@tf_export('keras.layers.GlobalMaxPool1D', 'keras.layers.GlobalMaxPooling1D') class GlobalMaxPooling1D(_GlobalPooling1D): """Global max pooling operation for temporal data. @@ -414,6 +424,8 @@ class _GlobalPooling2D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.GlobalAveragePooling2D', + 'keras.layers.GlobalAvgPool2D') class GlobalAveragePooling2D(_GlobalPooling2D): """Global average pooling operation for spatial data. @@ -449,6 +461,7 @@ class GlobalAveragePooling2D(_GlobalPooling2D): return K.mean(inputs, axis=[2, 3]) +@tf_export('keras.layers.GlobalMaxPool2D', 'keras.layers.GlobalMaxPooling2D') class GlobalMaxPooling2D(_GlobalPooling2D): """Global max pooling operation for spatial data. @@ -509,6 +522,8 @@ class _GlobalPooling3D(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.GlobalAveragePooling3D', + 'keras.layers.GlobalAvgPool3D') class GlobalAveragePooling3D(_GlobalPooling3D): """Global Average pooling operation for 3D data. @@ -544,6 +559,7 @@ class GlobalAveragePooling3D(_GlobalPooling3D): return K.mean(inputs, axis=[2, 3, 4]) +@tf_export('keras.layers.GlobalMaxPool3D', 'keras.layers.GlobalMaxPooling3D') class GlobalMaxPooling3D(_GlobalPooling3D): """Global Max pooling operation for 3D data. diff --git a/tensorflow/python/keras/_impl/keras/layers/pooling_test.py b/tensorflow/python/keras/_impl/keras/layers/pooling_test.py index ec0a5ae560f49ee39ecffb64f4ac65d3e800024c..bb003c1dddf80e2a745c1268a3a7d045f4e8b036 100644 --- a/tensorflow/python/keras/_impl/keras/layers/pooling_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/pooling_test.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.eager import context +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test @@ -25,81 +27,85 @@ from tensorflow.python.platform import test class GlobalPoolingTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_globalpooling_1d(self): - with self.test_session(use_gpu=True): - testing_utils.layer_test(keras.layers.pooling.GlobalMaxPooling1D, - input_shape=(3, 4, 5)) - testing_utils.layer_test( - keras.layers.pooling.GlobalAveragePooling1D, input_shape=(3, 4, 5)) + testing_utils.layer_test(keras.layers.pooling.GlobalMaxPooling1D, + input_shape=(3, 4, 5)) + testing_utils.layer_test( + keras.layers.pooling.GlobalAveragePooling1D, input_shape=(3, 4, 5)) + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_globalpooling_2d(self): - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.pooling.GlobalMaxPooling2D, - kwargs={'data_format': 'channels_first'}, - input_shape=(3, 4, 5, 6)) - testing_utils.layer_test( - keras.layers.pooling.GlobalMaxPooling2D, - kwargs={'data_format': 'channels_last'}, - input_shape=(3, 5, 6, 4)) - testing_utils.layer_test( - keras.layers.pooling.GlobalAveragePooling2D, - kwargs={'data_format': 'channels_first'}, - input_shape=(3, 4, 5, 6)) - testing_utils.layer_test( - keras.layers.pooling.GlobalAveragePooling2D, - kwargs={'data_format': 'channels_last'}, - input_shape=(3, 5, 6, 4)) - + testing_utils.layer_test( + keras.layers.pooling.GlobalMaxPooling2D, + kwargs={'data_format': 'channels_first'}, + input_shape=(3, 4, 5, 6)) + testing_utils.layer_test( + keras.layers.pooling.GlobalMaxPooling2D, + kwargs={'data_format': 'channels_last'}, + input_shape=(3, 5, 6, 4)) + testing_utils.layer_test( + keras.layers.pooling.GlobalAveragePooling2D, + kwargs={'data_format': 'channels_first'}, + input_shape=(3, 4, 5, 6)) + testing_utils.layer_test( + keras.layers.pooling.GlobalAveragePooling2D, + kwargs={'data_format': 'channels_last'}, + input_shape=(3, 5, 6, 4)) + + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_globalpooling_3d(self): - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.pooling.GlobalMaxPooling3D, - kwargs={'data_format': 'channels_first'}, - input_shape=(3, 4, 3, 4, 3)) - testing_utils.layer_test( - keras.layers.pooling.GlobalMaxPooling3D, - kwargs={'data_format': 'channels_last'}, - input_shape=(3, 4, 3, 4, 3)) - testing_utils.layer_test( - keras.layers.pooling.GlobalAveragePooling3D, - kwargs={'data_format': 'channels_first'}, - input_shape=(3, 4, 3, 4, 3)) - testing_utils.layer_test( - keras.layers.pooling.GlobalAveragePooling3D, - kwargs={'data_format': 'channels_last'}, - input_shape=(3, 4, 3, 4, 3)) + testing_utils.layer_test( + keras.layers.pooling.GlobalMaxPooling3D, + kwargs={'data_format': 'channels_first'}, + input_shape=(3, 4, 3, 4, 3)) + testing_utils.layer_test( + keras.layers.pooling.GlobalMaxPooling3D, + kwargs={'data_format': 'channels_last'}, + input_shape=(3, 4, 3, 4, 3)) + testing_utils.layer_test( + keras.layers.pooling.GlobalAveragePooling3D, + kwargs={'data_format': 'channels_first'}, + input_shape=(3, 4, 3, 4, 3)) + testing_utils.layer_test( + keras.layers.pooling.GlobalAveragePooling3D, + kwargs={'data_format': 'channels_last'}, + input_shape=(3, 4, 3, 4, 3)) class Pooling2DTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_maxpooling_2d(self): pool_size = (3, 3) - with self.test_session(use_gpu=True): - for strides in [(1, 1), (2, 2)]: - testing_utils.layer_test( - keras.layers.MaxPooling2D, - kwargs={ - 'strides': strides, - 'padding': 'valid', - 'pool_size': pool_size - }, - input_shape=(3, 5, 6, 4)) - - def test_averagepooling_2d(self): - with self.test_session(use_gpu=True): + for strides in [(1, 1), (2, 2)]: testing_utils.layer_test( - keras.layers.AveragePooling2D, - kwargs={'strides': (2, 2), - 'padding': 'same', - 'pool_size': (2, 2)}, - input_shape=(3, 5, 6, 4)) - testing_utils.layer_test( - keras.layers.AveragePooling2D, - kwargs={'strides': (2, 2), - 'padding': 'valid', - 'pool_size': (3, 3)}, + keras.layers.MaxPooling2D, + kwargs={ + 'strides': strides, + 'padding': 'valid', + 'pool_size': pool_size + }, input_shape=(3, 5, 6, 4)) + + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + def test_averagepooling_2d(self): + testing_utils.layer_test( + keras.layers.AveragePooling2D, + kwargs={'strides': (2, 2), + 'padding': 'same', + 'pool_size': (2, 2)}, + input_shape=(3, 5, 6, 4)) + testing_utils.layer_test( + keras.layers.AveragePooling2D, + kwargs={'strides': (2, 2), + 'padding': 'valid', + 'pool_size': (3, 3)}, + input_shape=(3, 5, 6, 4)) + + # This part of the test can only run on GPU but doesn't appear + # to be properly assigned to a GPU when running in eager mode. + if not context.executing_eagerly(): # Only runs on GPU with CUDA, channels_first is not supported on CPU. # TODO(b/62340061): Support channels_first on CPU. if test.is_gpu_available(cuda_only=True): @@ -116,66 +122,66 @@ class Pooling2DTest(test.TestCase): class Pooling3DTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_maxpooling_3d(self): pool_size = (3, 3, 3) - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.MaxPooling3D, - kwargs={'strides': 2, - 'padding': 'valid', - 'pool_size': pool_size}, - input_shape=(3, 11, 12, 10, 4)) - testing_utils.layer_test( - keras.layers.MaxPooling3D, - kwargs={ - 'strides': 3, - 'padding': 'valid', - 'data_format': 'channels_first', - 'pool_size': pool_size - }, - input_shape=(3, 4, 11, 12, 10)) - + testing_utils.layer_test( + keras.layers.MaxPooling3D, + kwargs={'strides': 2, + 'padding': 'valid', + 'pool_size': pool_size}, + input_shape=(3, 11, 12, 10, 4)) + testing_utils.layer_test( + keras.layers.MaxPooling3D, + kwargs={ + 'strides': 3, + 'padding': 'valid', + 'data_format': 'channels_first', + 'pool_size': pool_size + }, + input_shape=(3, 4, 11, 12, 10)) + + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_averagepooling_3d(self): pool_size = (3, 3, 3) - with self.test_session(use_gpu=True): - testing_utils.layer_test( - keras.layers.AveragePooling3D, - kwargs={'strides': 2, - 'padding': 'valid', - 'pool_size': pool_size}, - input_shape=(3, 11, 12, 10, 4)) - testing_utils.layer_test( - keras.layers.AveragePooling3D, - kwargs={ - 'strides': 3, - 'padding': 'valid', - 'data_format': 'channels_first', - 'pool_size': pool_size - }, - input_shape=(3, 4, 11, 12, 10)) + testing_utils.layer_test( + keras.layers.AveragePooling3D, + kwargs={'strides': 2, + 'padding': 'valid', + 'pool_size': pool_size}, + input_shape=(3, 11, 12, 10, 4)) + testing_utils.layer_test( + keras.layers.AveragePooling3D, + kwargs={ + 'strides': 3, + 'padding': 'valid', + 'data_format': 'channels_first', + 'pool_size': pool_size + }, + input_shape=(3, 4, 11, 12, 10)) class Pooling1DTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_maxpooling_1d(self): - with self.test_session(use_gpu=True): - for padding in ['valid', 'same']: - for stride in [1, 2]: - testing_utils.layer_test( - keras.layers.MaxPooling1D, - kwargs={'strides': stride, - 'padding': padding}, - input_shape=(3, 5, 4)) + for padding in ['valid', 'same']: + for stride in [1, 2]: + testing_utils.layer_test( + keras.layers.MaxPooling1D, + kwargs={'strides': stride, + 'padding': padding}, + input_shape=(3, 5, 4)) + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) def test_averagepooling_1d(self): - with self.test_session(use_gpu=True): - for padding in ['valid', 'same']: - for stride in [1, 2]: - testing_utils.layer_test( - keras.layers.AveragePooling1D, - kwargs={'strides': stride, - 'padding': padding}, - input_shape=(3, 5, 4)) + for padding in ['valid', 'same']: + for stride in [1, 2]: + testing_utils.layer_test( + keras.layers.AveragePooling1D, + kwargs={'strides': stride, + 'padding': padding}, + input_shape=(3, 5, 4)) if __name__ == '__main__': diff --git a/tensorflow/python/keras/_impl/keras/layers/recurrent.py b/tensorflow/python/keras/_impl/keras/layers/recurrent.py index 1b0f6cb6cf20fc55d3edefc9684d423cf25b7e0e..791f9b311300ed05591083d551c040eb25ac8e22 100644 --- a/tensorflow/python/keras/_impl/keras/layers/recurrent.py +++ b/tensorflow/python/keras/_impl/keras/layers/recurrent.py @@ -19,8 +19,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numbers import numpy as np +from tensorflow.python.eager import context from tensorflow.python.framework import tensor_shape from tensorflow.python.keras._impl.keras import activations from tensorflow.python.keras._impl.keras import backend as K @@ -29,11 +31,13 @@ from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers from tensorflow.python.keras._impl.keras.engine import InputSpec from tensorflow.python.keras._impl.keras.engine import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion from tensorflow.python.keras._impl.keras.utils.generic_utils import has_arg from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.StackedRNNCells') class StackedRNNCells(Layer): """Wrapper allowing a stack of RNN cells to behave as a single cell. @@ -85,7 +89,7 @@ class StackedRNNCells(Layer): state_size.append(cell.state_size) return tuple(state_size) - def call(self, inputs, states, **kwargs): + def call(self, inputs, states, constants=None, **kwargs): # Recover per-cell states. nested_states = [] for cell in self.cells[::-1]: @@ -100,7 +104,12 @@ class StackedRNNCells(Layer): # Call the cells in order and store the returned states. new_nested_states = [] for cell, states in zip(self.cells, nested_states): - inputs, states = cell.call(inputs, states, **kwargs) + if has_arg(cell.call, 'constants'): + inputs, states = cell.call(inputs, states, constants=constants, + **kwargs) + else: + inputs, states = cell.call(inputs, states, **kwargs) + new_nested_states.append(states) # Format the new states as a flat list @@ -112,9 +121,15 @@ class StackedRNNCells(Layer): @shape_type_conversion def build(self, input_shape): + if isinstance(input_shape, list): + constants_shape = input_shape[1:] + input_shape = input_shape[0] for cell in self.cells: if isinstance(cell, Layer): - cell.build(input_shape) + if has_arg(cell.call, 'constants'): + cell.build([input_shape] + constants_shape) + else: + cell.build(input_shape) if hasattr(cell.state_size, '__len__'): output_dim = cell.state_size[0] else: @@ -200,19 +215,19 @@ class StackedRNNCells(Layer): losses = [] for cell in self.cells: if isinstance(cell, Layer): - cell_losses = cell.losses - losses += cell_losses - return losses + losses += cell.losses + return losses + self._losses - def get_losses_for(self, inputs=None): - losses = [] + @property + def updates(self): + updates = [] for cell in self.cells: if isinstance(cell, Layer): - cell_losses = cell.get_losses_for(inputs) - losses += cell_losses - return losses + updates += cell.updates + return updates + self._updates +@tf_export('keras.layers.RNN') class RNN(Layer): """Base class for recurrent layers. @@ -400,7 +415,7 @@ class RNN(Layer): @property def states(self): if self._states is None: - if isinstance(self.cell.state_size, int): + if isinstance(self.cell.state_size, numbers.Integral): num_states = 1 else: num_states = len(self.cell.state_size) @@ -525,12 +540,14 @@ class RNN(Layer): self._num_constants = len(constants) additional_specs += self.constants_spec # at this point additional_inputs cannot be empty - is_keras_tensor = hasattr(additional_inputs[0], '_keras_history') + is_keras_tensor = K.is_keras_tensor(additional_inputs[0]) for tensor in additional_inputs: - if hasattr(tensor, '_keras_history') != is_keras_tensor: + if K.is_keras_tensor(tensor) != is_keras_tensor: raise ValueError('The initial state or constants of an RNN' ' layer cannot be specified with a mix of' - ' Keras tensors and non-Keras tensors') + ' Keras tensors and non-Keras tensors' + ' (a "Keras tensor" is a tensor that was' + ' returned by a Keras layer, or by `Input`)') if is_keras_tensor: # Compute the full input spec, including state and constants @@ -614,7 +631,7 @@ class RNN(Layer): if self.stateful: updates = [] for i in range(len(states)): - updates.append((self.states[i], states[i])) + updates.append(K.update(self.states[i], states[i])) self.add_update(updates, inputs) if self.return_sequences: @@ -774,24 +791,28 @@ class RNN(Layer): @property def losses(self): + losses = [] if isinstance(self.cell, Layer): - return self.cell.losses - return [] + losses += self.cell.losses + return losses + self._losses - def get_losses_for(self, inputs=None): + @property + def updates(self): + updates = [] if isinstance(self.cell, Layer): - cell_losses = self.cell.get_losses_for(inputs) - return cell_losses + super(RNN, self).get_losses_for(inputs) - return super(RNN, self).get_losses_for(inputs) + updates += self.cell.updates + return updates + self._updates +@tf_export('keras.layers.SimpleRNNCell') class SimpleRNNCell(Layer): """Cell class for SimpleRNN. Arguments: units: Positive integer, dimensionality of the output space. activation: Activation function to use. - If you pass None, no activation is applied + Default: hyperbolic tangent (`tanh`). + If you pass `None`, no activation is applied (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. kernel_initializer: Initializer for the `kernel` weights matrix, @@ -915,7 +936,9 @@ class SimpleRNNCell(Layer): # Properly set learning phase on output tensor. if 0 < self.dropout + self.recurrent_dropout: - if training is None: + if training is None and not context.executing_eagerly(): + # This would be harmless to set in eager mode, but eager tensors + # disallow setting arbitrary attributes. output._uses_learning_phase = True return output, [output] @@ -954,12 +977,14 @@ class SimpleRNNCell(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.SimpleRNN') class SimpleRNN(RNN): """Fully-connected RNN where the output is to be fed back to input. Arguments: units: Positive integer, dimensionality of the output space. activation: Activation function to use. + Default: hyperbolic tangent (`tanh`). If you pass None, no activation is applied (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. @@ -1163,16 +1188,21 @@ class SimpleRNN(RNN): return cls(**config) +@tf_export('keras.layers.GRUCell') class GRUCell(Layer): """Cell class for the GRU layer. Arguments: units: Positive integer, dimensionality of the output space. activation: Activation function to use. + Default: hyperbolic tangent (`tanh`). If you pass None, no activation is applied (ie. "linear" activation: `a(x) = x`). recurrent_activation: Activation function to use for the recurrent step. + Default: hard sigmoid (`hard_sigmoid`). + If you pass `None`, no activation is applied + (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. kernel_initializer: Initializer for the `kernel` weights matrix, used for the linear transformation of the inputs. @@ -1272,23 +1302,6 @@ class GRUCell(Layer): constraint=self.bias_constraint) else: self.bias = None - - self.kernel_z = self.kernel[:, :self.units] - self.recurrent_kernel_z = self.recurrent_kernel[:, :self.units] - self.kernel_r = self.kernel[:, self.units:self.units * 2] - self.recurrent_kernel_r = self.recurrent_kernel[:, self.units: - self.units * 2] - self.kernel_h = self.kernel[:, self.units * 2:] - self.recurrent_kernel_h = self.recurrent_kernel[:, self.units * 2:] - - if self.use_bias: - self.bias_z = self.bias[:self.units] - self.bias_r = self.bias[self.units:self.units * 2] - self.bias_h = self.bias[self.units * 2:] - else: - self.bias_z = None - self.bias_r = None - self.bias_h = None self.built = True def call(self, inputs, states, training=None): @@ -1323,13 +1336,13 @@ class GRUCell(Layer): inputs_z = inputs inputs_r = inputs inputs_h = inputs - x_z = K.dot(inputs_z, self.kernel_z) - x_r = K.dot(inputs_r, self.kernel_r) - x_h = K.dot(inputs_h, self.kernel_h) + x_z = K.dot(inputs_z, self.kernel[:, :self.units]) + x_r = K.dot(inputs_r, self.kernel[:, self.units:self.units * 2]) + x_h = K.dot(inputs_h, self.kernel[:, self.units * 2:]) if self.use_bias: - x_z = K.bias_add(x_z, self.bias_z) - x_r = K.bias_add(x_r, self.bias_r) - x_h = K.bias_add(x_h, self.bias_h) + x_z = K.bias_add(x_z, self.bias[:self.units]) + x_r = K.bias_add(x_r, self.bias[self.units:self.units * 2]) + x_h = K.bias_add(x_h, self.bias[self.units * 2:]) if 0. < self.recurrent_dropout < 1.: h_tm1_z = h_tm1 * rec_dp_mask[0] @@ -1340,11 +1353,14 @@ class GRUCell(Layer): h_tm1_r = h_tm1 h_tm1_h = h_tm1 z = self.recurrent_activation( - x_z + K.dot(h_tm1_z, self.recurrent_kernel_z)) + x_z + K.dot(h_tm1_z, self.recurrent_kernel[:, :self.units])) r = self.recurrent_activation( - x_r + K.dot(h_tm1_r, self.recurrent_kernel_r)) + x_r + K.dot(h_tm1_r, self.recurrent_kernel[:, self.units: + self.units * 2])) - hh = self.activation(x_h + K.dot(r * h_tm1_h, self.recurrent_kernel_h)) + hh = self.activation(x_h + K.dot(r * h_tm1_h, + self.recurrent_kernel[:, + self.units * 2:])) else: if 0. < self.dropout < 1.: inputs *= dp_mask[0] @@ -1368,49 +1384,40 @@ class GRUCell(Layer): hh = self.activation(x_h + recurrent_h) h = z * h_tm1 + (1 - z) * hh if 0 < self.dropout + self.recurrent_dropout: - if training is None: + if training is None and not context.executing_eagerly(): + # This would be harmless to set in eager mode, but eager tensors + # disallow setting arbitrary attributes. h._uses_learning_phase = True return h, [h] def get_config(self): config = { - 'units': - self.units, - 'activation': - activations.serialize(self.activation), + 'units': self.units, + 'activation': activations.serialize(self.activation), 'recurrent_activation': activations.serialize(self.recurrent_activation), - 'use_bias': - self.use_bias, - 'kernel_initializer': - initializers.serialize(self.kernel_initializer), + 'use_bias': self.use_bias, + 'kernel_initializer': initializers.serialize(self.kernel_initializer), 'recurrent_initializer': initializers.serialize(self.recurrent_initializer), - 'bias_initializer': - initializers.serialize(self.bias_initializer), - 'kernel_regularizer': - regularizers.serialize(self.kernel_regularizer), + 'bias_initializer': initializers.serialize(self.bias_initializer), + 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), 'recurrent_regularizer': regularizers.serialize(self.recurrent_regularizer), - 'bias_regularizer': - regularizers.serialize(self.bias_regularizer), - 'kernel_constraint': - constraints.serialize(self.kernel_constraint), + 'bias_regularizer': regularizers.serialize(self.bias_regularizer), + 'kernel_constraint': constraints.serialize(self.kernel_constraint), 'recurrent_constraint': constraints.serialize(self.recurrent_constraint), - 'bias_constraint': - constraints.serialize(self.bias_constraint), - 'dropout': - self.dropout, - 'recurrent_dropout': - self.recurrent_dropout, - 'implementation': - self.implementation + 'bias_constraint': constraints.serialize(self.bias_constraint), + 'dropout': self.dropout, + 'recurrent_dropout': self.recurrent_dropout, + 'implementation': self.implementation } base_config = super(GRUCell, self).get_config() return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.GRU') class GRU(RNN): """Gated Recurrent Unit - Cho et al. @@ -1419,10 +1426,14 @@ class GRU(RNN): Arguments: units: Positive integer, dimensionality of the output space. activation: Activation function to use. - If you pass None, no activation is applied + Default: hyperbolic tangent (`tanh`). + If you pass `None`, no activation is applied (ie. "linear" activation: `a(x) = x`). recurrent_activation: Activation function to use for the recurrent step. + Default: hard sigmoid (`hard_sigmoid`). + If you pass `None`, no activation is applied + (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. kernel_initializer: Initializer for the `kernel` weights matrix, used for the linear transformation of the inputs. @@ -1646,16 +1657,21 @@ class GRU(RNN): return cls(**config) +@tf_export('keras.layers.LSTMCell') class LSTMCell(Layer): """Cell class for the LSTM layer. Arguments: units: Positive integer, dimensionality of the output space. activation: Activation function to use. - If you pass None, no activation is applied + Default: hyperbolic tangent (`tanh`). + If you pass `None`, no activation is applied (ie. "linear" activation: `a(x) = x`). recurrent_activation: Activation function to use for the recurrent step. + Default: hard sigmoid (`hard_sigmoid`). + If you pass `None`, no activation is applied + (ie. "linear" activation: `a(x) = x`).x use_bias: Boolean, whether the layer uses a bias vector. kernel_initializer: Initializer for the `kernel` weights matrix, used for the linear transformation of the inputs. @@ -1772,29 +1788,6 @@ class LSTMCell(Layer): constraint=self.bias_constraint) else: self.bias = None - - self.kernel_i = self.kernel[:, :self.units] - self.kernel_f = self.kernel[:, self.units:self.units * 2] - self.kernel_c = self.kernel[:, self.units * 2:self.units * 3] - self.kernel_o = self.kernel[:, self.units * 3:] - - self.recurrent_kernel_i = self.recurrent_kernel[:, :self.units] - self.recurrent_kernel_f = self.recurrent_kernel[:, self.units: - self.units * 2] - self.recurrent_kernel_c = self.recurrent_kernel[:, self.units * 2: - self.units * 3] - self.recurrent_kernel_o = self.recurrent_kernel[:, self.units * 3:] - - if self.use_bias: - self.bias_i = self.bias[:self.units] - self.bias_f = self.bias[self.units:self.units * 2] - self.bias_c = self.bias[self.units * 2:self.units * 3] - self.bias_o = self.bias[self.units * 3:] - else: - self.bias_i = None - self.bias_f = None - self.bias_c = None - self.bias_o = None self.built = True def call(self, inputs, states, training=None): @@ -1832,15 +1825,15 @@ class LSTMCell(Layer): inputs_f = inputs inputs_c = inputs inputs_o = inputs - x_i = K.dot(inputs_i, self.kernel_i) - x_f = K.dot(inputs_f, self.kernel_f) - x_c = K.dot(inputs_c, self.kernel_c) - x_o = K.dot(inputs_o, self.kernel_o) + x_i = K.dot(inputs_i, self.kernel[:, :self.units]) + x_f = K.dot(inputs_f, self.kernel[:, self.units:self.units * 2]) + x_c = K.dot(inputs_c, self.kernel[:, self.units * 2:self.units * 3]) + x_o = K.dot(inputs_o, self.kernel[:, self.units * 3:]) if self.use_bias: - x_i = K.bias_add(x_i, self.bias_i) - x_f = K.bias_add(x_f, self.bias_f) - x_c = K.bias_add(x_c, self.bias_c) - x_o = K.bias_add(x_o, self.bias_o) + x_i = K.bias_add(x_i, self.bias[:self.units]) + x_f = K.bias_add(x_f, self.bias[self.units:self.units * 2]) + x_c = K.bias_add(x_c, self.bias[self.units * 2:self.units * 3]) + x_o = K.bias_add(x_o, self.bias[self.units * 3:]) if 0 < self.recurrent_dropout < 1.: h_tm1_i = h_tm1 * rec_dp_mask[0] @@ -1853,13 +1846,15 @@ class LSTMCell(Layer): h_tm1_c = h_tm1 h_tm1_o = h_tm1 i = self.recurrent_activation( - x_i + K.dot(h_tm1_i, self.recurrent_kernel_i)) + x_i + K.dot(h_tm1_i, self.recurrent_kernel[:, :self.units])) f = self.recurrent_activation( - x_f + K.dot(h_tm1_f, self.recurrent_kernel_f)) + x_f + K.dot(h_tm1_f, + self.recurrent_kernel[:, self.units: self.units * 2])) c = f * c_tm1 + i * self.activation( - x_c + K.dot(h_tm1_c, self.recurrent_kernel_c)) + x_c + K.dot(h_tm1_c, + self.recurrent_kernel[:, self.units * 2: self.units * 3])) o = self.recurrent_activation( - x_o + K.dot(h_tm1_o, self.recurrent_kernel_o)) + x_o + K.dot(h_tm1_o, self.recurrent_kernel[:, self.units * 3:])) else: if 0. < self.dropout < 1.: inputs *= dp_mask[0] @@ -1882,7 +1877,9 @@ class LSTMCell(Layer): h = o * self.activation(c) if 0 < self.dropout + self.recurrent_dropout: - if training is None: + if training is None and not context.executing_eagerly(): + # This would be harmless to set in eager mode, but eager tensors + # disallow setting arbitrary attributes. h._uses_learning_phase = True return h, [h, c] @@ -1927,16 +1924,21 @@ class LSTMCell(Layer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.layers.LSTM') class LSTM(RNN): """Long-Short Term Memory layer - Hochreiter 1997. Arguments: units: Positive integer, dimensionality of the output space. activation: Activation function to use. - If you pass None, no activation is applied + Default: hyperbolic tangent (`tanh`). + If you pass `None`, no activation is applied (ie. "linear" activation: `a(x) = x`). recurrent_activation: Activation function to use for the recurrent step. + Default: hard sigmoid (`hard_sigmoid`). + If you pass `None`, no activation is applied + (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. kernel_initializer: Initializer for the `kernel` weights matrix, used for the linear transformation of the inputs.. @@ -2454,7 +2456,7 @@ class Recurrent(Layer): if self.stateful: updates = [] for i in range(len(states)): - updates.append((self.states[i], states[i])) + updates.append(K.update(self.states[i], states[i])) self.add_update(updates, inputs) # Properly set learning phase diff --git a/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py b/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py index a1407a24ea895976cbf95f0ea5c7ba98335af329..de022153f6f07240a0dff70e5faeed5b6d4a8c5f 100644 --- a/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py @@ -253,7 +253,7 @@ class RNNTest(test.TestCase): self.assertAllClose(y_np, y_np_2, atol=1e-4) with self.test_session(): - # test flat list inputs + # test flat list inputs. with keras.utils.CustomObjectScope(custom_objects): layer = keras.layers.RNN.from_config(config.copy()) y = layer([x, c]) @@ -262,6 +262,35 @@ class RNNTest(test.TestCase): y_np_3 = model.predict([x_np, c_np]) self.assertAllClose(y_np, y_np_3, atol=1e-4) + with self.test_session(): + # Test stacking. + cells = [keras.layers.recurrent.GRUCell(8), + RNNCellWithConstants(12), + RNNCellWithConstants(32)] + layer = keras.layers.recurrent.RNN(cells) + y = layer(x, constants=c) + model = keras.models.Model([x, c], y) + model.compile(optimizer='rmsprop', loss='mse') + model.train_on_batch( + [np.zeros((6, 5, 5)), np.zeros((6, 3))], + np.zeros((6, 32)) + ) + + with self.test_session(): + # Test stacked RNN serialization + x_np = np.random.random((6, 5, 5)) + c_np = np.random.random((6, 3)) + y_np = model.predict([x_np, c_np]) + weights = model.get_weights() + config = layer.get_config() + with keras.utils.CustomObjectScope(custom_objects): + layer = keras.layers.recurrent.RNN.from_config(config.copy()) + y = layer(x, constants=c) + model = keras.models.Model([x, c], y) + model.set_weights(weights) + y_np_2 = model.predict([x_np, c_np]) + self.assertAllClose(y_np, y_np_2, atol=1e-4) + def test_rnn_cell_with_constants_layer_passing_initial_state(self): class RNNCellWithConstants(keras.layers.Layer): @@ -353,13 +382,10 @@ class RNNTest(test.TestCase): self.assertAllClose(y_np, y_np_3, atol=1e-4) def test_stacked_rnn_attributes(self): - cells = [keras.layers.LSTMCell(3), - keras.layers.LSTMCell(3, kernel_regularizer='l2')] + cells = [keras.layers.LSTMCell(1), + keras.layers.LSTMCell(1)] layer = keras.layers.RNN(cells) - layer.build((None, None, 5)) - - # Test regularization losses - self.assertEqual(len(layer.losses), 1) + layer.build((None, None, 1)) # Test weights self.assertEqual(len(layer.trainable_weights), 6) @@ -367,11 +393,32 @@ class RNNTest(test.TestCase): self.assertEqual(len(layer.trainable_weights), 3) self.assertEqual(len(layer.non_trainable_weights), 3) - # Test `get_losses_for` - x = keras.Input((None, 5)) - y = keras.backend.sum(x) - cells[0].add_loss(y, inputs=x) - self.assertEqual(layer.get_losses_for(x), [y]) + # Test `get_losses_for` and `losses` + x = keras.Input((None, 1)) + loss_1 = keras.backend.sum(x) + loss_2 = keras.backend.sum(cells[0].kernel) + cells[0].add_loss(loss_1, inputs=x) + cells[0].add_loss(loss_2) + self.assertEqual(len(layer.losses), 2) + self.assertEqual(layer.get_losses_for(None), [loss_2]) + self.assertEqual(layer.get_losses_for(x), [loss_1]) + + # Test `get_updates_for` and `updates` + cells = [keras.layers.LSTMCell(1), + keras.layers.LSTMCell(1)] + layer = keras.layers.RNN(cells) + layer.build((None, None, 1)) + + x = keras.Input((None, 1)) + update_1 = keras.backend.update_add( + cells[0].kernel, x[0, 0, 0] * cells[0].kernel) + update_2 = keras.backend.update_add( + cells[0].kernel, keras.backend.ones_like(cells[0].kernel)) + cells[0].add_update(update_1, inputs=x) + cells[0].add_update(update_2) + self.assertEqual(len(layer.updates), 2) + self.assertEqual(layer.get_updates_for(None), [update_2]) + self.assertEqual(layer.get_updates_for(x), [update_1]) def test_rnn_dynamic_trainability(self): layer_class = keras.layers.SimpleRNN diff --git a/tensorflow/python/keras/_impl/keras/layers/simplernn_test.py b/tensorflow/python/keras/_impl/keras/layers/simplernn_test.py index 7edebdacd07d74fe6b5a982d12645fb5556bdf75..8c7189cd4718450a85c015e08ab3a58cc5d86531 100644 --- a/tensorflow/python/keras/_impl/keras/layers/simplernn_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/simplernn_test.py @@ -20,64 +20,66 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.platform import test +from tensorflow.python.training.rmsprop import RMSPropOptimizer class SimpleRNNLayerTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_return_sequences_SimpleRNN(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - testing_utils.layer_test( - keras.layers.SimpleRNN, - kwargs={'units': units, - 'return_sequences': True}, - input_shape=(num_samples, timesteps, embedding_dim)) + testing_utils.layer_test( + keras.layers.SimpleRNN, + kwargs={'units': units, + 'return_sequences': True}, + input_shape=(num_samples, timesteps, embedding_dim)) + @tf_test_util.run_in_graph_and_eager_modes() def test_dynamic_behavior_SimpleRNN(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - layer = keras.layers.SimpleRNN(units, input_shape=(None, embedding_dim)) - model = keras.models.Sequential() - model.add(layer) - model.compile('sgd', 'mse') - x = np.random.random((num_samples, timesteps, embedding_dim)) - y = np.random.random((num_samples, units)) - model.train_on_batch(x, y) - + layer = keras.layers.SimpleRNN(units, input_shape=(None, embedding_dim)) + model = keras.models.Sequential() + model.add(layer) + model.compile(RMSPropOptimizer(0.01), 'mse') + x = np.random.random((num_samples, timesteps, embedding_dim)) + y = np.random.random((num_samples, units)) + model.train_on_batch(x, y) + + @tf_test_util.run_in_graph_and_eager_modes() def test_dropout_SimpleRNN(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - testing_utils.layer_test( - keras.layers.SimpleRNN, - kwargs={'units': units, - 'dropout': 0.1, - 'recurrent_dropout': 0.1}, - input_shape=(num_samples, timesteps, embedding_dim)) - + testing_utils.layer_test( + keras.layers.SimpleRNN, + kwargs={'units': units, + 'dropout': 0.1, + 'recurrent_dropout': 0.1}, + input_shape=(num_samples, timesteps, embedding_dim)) + + @tf_test_util.run_in_graph_and_eager_modes() def test_implementation_mode_SimpleRNN(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 - with self.test_session(): - for mode in [0, 1, 2]: - testing_utils.layer_test( - keras.layers.SimpleRNN, - kwargs={'units': units, - 'implementation': mode}, - input_shape=(num_samples, timesteps, embedding_dim)) + for mode in [0, 1, 2]: + testing_utils.layer_test( + keras.layers.SimpleRNN, + kwargs={'units': units, + 'implementation': mode}, + input_shape=(num_samples, timesteps, embedding_dim)) def test_statefulness_SimpleRNN(self): num_samples = 2 diff --git a/tensorflow/python/keras/_impl/keras/layers/wrappers.py b/tensorflow/python/keras/_impl/keras/layers/wrappers.py index 3667956f8049b8f5e9cc9ab7a9e0037c260a8ce1..76ddd9299dd669da35d89a6fe8fc521ce4c26337 100644 --- a/tensorflow/python/keras/_impl/keras/layers/wrappers.py +++ b/tensorflow/python/keras/_impl/keras/layers/wrappers.py @@ -25,11 +25,13 @@ from tensorflow.python.framework import tensor_shape from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.engine import InputSpec from tensorflow.python.keras._impl.keras.engine import Layer -from tensorflow.python.keras._impl.keras.engine.topology import shape_type_conversion +from tensorflow.python.keras._impl.keras.engine.base_layer import shape_type_conversion from tensorflow.python.keras._impl.keras.utils.generic_utils import has_arg from tensorflow.python.layers import utils as tf_layers_util +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.layers.Wrapper') class Wrapper(Layer): """Abstract wrapper base class. @@ -59,6 +61,14 @@ class Wrapper(Layer): else: return None + @property + def trainable(self): + return self.layer.trainable + + @trainable.setter + def trainable(self, value): + self.layer.trainable = value + @property def trainable_weights(self): return self.layer.trainable_weights @@ -69,34 +79,11 @@ class Wrapper(Layer): @property def updates(self): - if hasattr(self.layer, 'updates'): - return self.layer.updates - return [] - - def get_updates_for(self, inputs=None): - # If the wrapper modifies the inputs, use the modified inputs to - # get the updates from the inner layer. - inner_inputs = inputs - if inputs is not None: - uid = tf_layers_util.object_list_uid(inputs) - if uid in self._input_map: - inner_inputs = self._input_map[uid] - - updates = self.layer.get_updates_for(inner_inputs) - updates += super(Wrapper, self).get_updates_for(inputs) - return updates + return self.layer.updates + self._updates @property def losses(self): - if hasattr(self.layer, 'losses'): - return self.layer.losses - return [] - - def get_losses_for(self, inputs=None): - if inputs is None: - losses = self.layer.get_losses_for(None) - return losses + super(Wrapper, self).get_losses_for(None) - return super(Wrapper, self).get_losses_for(inputs) + return self.layer.losses + self._losses def get_weights(self): return self.layer.get_weights() @@ -122,6 +109,7 @@ class Wrapper(Layer): return cls(layer, **config) +@tf_export('keras.layers.TimeDistributed') class TimeDistributed(Wrapper): """This wrapper allows to apply a layer to every temporal slice of an input. @@ -246,6 +234,7 @@ class TimeDistributed(Wrapper): return y +@tf_export('keras.layers.Bidirectional') class Bidirectional(Wrapper): """Bidirectional wrapper for RNNs. @@ -274,7 +263,6 @@ class Bidirectional(Wrapper): """ def __init__(self, layer, merge_mode='concat', weights=None, **kwargs): - super(Bidirectional, self).__init__(layer, **kwargs) if merge_mode not in ['sum', 'mul', 'ave', 'concat', None]: raise ValueError('Invalid merge mode. ' 'Merge mode should be one of ' @@ -294,6 +282,19 @@ class Bidirectional(Wrapper): self.return_sequences = layer.return_sequences self.return_state = layer.return_state self.supports_masking = True + self._trainable = True + super(Bidirectional, self).__init__(layer, **kwargs) + self.input_spec = layer.input_spec + + @property + def trainable(self): + return self._trainable + + @trainable.setter + def trainable(self, value): + self._trainable = value + self.forward_layer.trainable = value + self.backward_layer.trainable = value def get_weights(self): return self.forward_layer.get_weights() + self.backward_layer.get_weights() @@ -324,6 +325,61 @@ class Bidirectional(Wrapper): return [output_shape] + state_shape + copy.copy(state_shape) return output_shape + def __call__(self, inputs, initial_state=None, **kwargs): + if isinstance(inputs, list): + if len(inputs) > 1: + initial_state = inputs[1:] + inputs = inputs[0] + + if initial_state is None: + return super(Bidirectional, self).__call__(inputs, **kwargs) + + # Standardize `initial_state` into list + if isinstance(initial_state, tuple): + initial_state = list(initial_state) + elif not isinstance(initial_state, list): + initial_state = [initial_state] + + # Check if `initial_state` can be splitted into half + num_states = len(initial_state) + if num_states % 2 > 0: + raise ValueError( + 'When passing `initial_state` to a Bidirectional RNN, the state ' + 'should be a list containing the states of the underlying RNNs. ' + 'Found: ' + str(initial_state)) + + # Applies the same workaround as in `RNN.__call__`, without handling + # constants + kwargs['initial_state'] = initial_state + additional_inputs = initial_state + additional_specs = [InputSpec(shape=K.int_shape(state)) + for state in initial_state] + self.forward_layer.state_spec = additional_specs[:num_states // 2] + self.backward_layer.state_spec = additional_specs[num_states // 2:] + + is_keras_tensor = K.is_keras_tensor(additional_inputs[0]) + for tensor in additional_inputs: + if K.is_keras_tensor(tensor) != is_keras_tensor: + raise ValueError('The initial state of a Bidirectional' + ' layer cannot be specified with a mix of' + ' Keras tensors and non-Keras tensors' + ' (a "Keras tensor" is a tensor that was' + ' returned by a Keras layer, or by `Input`)') + + if is_keras_tensor: + # Compute the full input spec, including state + full_input = [inputs] + additional_inputs + full_input_spec = self.input_spec + additional_specs + + # Perform the call with temporarily replaced input_spec + original_input_spec = self.input_spec + self.input_spec = full_input_spec + output = super(Bidirectional, self).__call__(full_input, **kwargs) + self.input_spec = original_input_spec + return output + else: + return super(Bidirectional, self).__call__(inputs, **kwargs) + def call(self, inputs, training=None, mask=None, initial_state=None): kwargs = {} if has_arg(self.layer.call, 'training'): @@ -332,11 +388,6 @@ class Bidirectional(Wrapper): kwargs['mask'] = mask if initial_state is not None and has_arg(self.layer.call, 'initial_state'): - if not isinstance(initial_state, list): - raise ValueError( - 'When passing `initial_state` to a Bidirectional RNN, the state ' - 'should be a list containing the states of the underlying RNNs. ' - 'Found: ' + str(initial_state)) forward_state = initial_state[:len(initial_state) // 2] backward_state = initial_state[len(initial_state) // 2:] y = self.forward_layer.call(inputs, initial_state=forward_state, **kwargs) diff --git a/tensorflow/python/keras/_impl/keras/layers/wrappers_test.py b/tensorflow/python/keras/_impl/keras/layers/wrappers_test.py index f48c8919a148403874758b618aaa9a662e511240..8fcf66e90ff1289a06a996768ae5de2f1548a27c 100644 --- a/tensorflow/python/keras/_impl/keras/layers/wrappers_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/wrappers_test.py @@ -20,44 +20,43 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.keras._impl import keras from tensorflow.python.platform import test +from tensorflow.python.training.rmsprop import RMSPropOptimizer class TimeDistributedTest(test.TestCase): + @tf_test_util.run_in_graph_and_eager_modes() def test_timedistributed_dense(self): - # first, test with Dense layer - with self.test_session(): - model = keras.models.Sequential() - model.add( - keras.layers.TimeDistributed( - keras.layers.Dense(2), input_shape=(3, 4))) - model.compile(optimizer='rmsprop', loss='mse') - model.fit( - np.random.random((10, 3, 4)), - np.random.random((10, 3, 2)), - epochs=1, - batch_size=10) - - # test config - model.get_config() + model = keras.models.Sequential() + model.add( + keras.layers.TimeDistributed( + keras.layers.Dense(2), input_shape=(3, 4))) + model.compile(optimizer=RMSPropOptimizer(0.01), loss='mse') + model.fit( + np.random.random((10, 3, 4)), + np.random.random((10, 3, 2)), + epochs=1, + batch_size=10) + + # test config + model.get_config() def test_timedistributed_static_batch_size(self): - with self.test_session(): - model = keras.models.Sequential() - model.add( - keras.layers.TimeDistributed( - keras.layers.Dense(2), input_shape=(3, 4), batch_size=10)) - model.compile(optimizer='rmsprop', loss='mse') - model.fit( - np.random.random((10, 3, 4)), - np.random.random((10, 3, 2)), - epochs=1, - batch_size=10) + model = keras.models.Sequential() + model.add( + keras.layers.TimeDistributed( + keras.layers.Dense(2), input_shape=(3, 4), batch_size=10)) + model.compile(optimizer=RMSPropOptimizer(0.01), loss='mse') + model.fit( + np.random.random((10, 3, 4)), + np.random.random((10, 3, 2)), + epochs=1, + batch_size=10) def test_timedistributed_conv2d(self): - # test with Conv2D with self.test_session(): model = keras.models.Sequential() model.add( @@ -73,7 +72,6 @@ class TimeDistributedTest(test.TestCase): model.summary() def test_timedistributed_stacked(self): - # test stacked layers with self.test_session(): model = keras.models.Sequential() model.add( @@ -133,6 +131,20 @@ class TimeDistributedTest(test.TestCase): # Verify input_map has one mapping from inputs to reshaped inputs. self.assertEqual(len(td._input_map.keys()), 1) + def test_TimeDistributed_trainable(self): + # test layers that need learning_phase to be set + x = keras.layers.Input(shape=(3, 2)) + layer = keras.layers.TimeDistributed(keras.layers.BatchNormalization()) + _ = layer(x) + assert len(layer.updates) == 2 + assert len(layer.trainable_weights) == 2 + layer.trainable = False + assert not layer.updates + assert not layer.trainable_weights + layer.trainable = True + assert len(layer.updates) == 2 + assert len(layer.trainable_weights) == 2 + class BidirectionalTest(test.TestCase): @@ -153,7 +165,7 @@ class BidirectionalTest(test.TestCase): model.add( keras.layers.Bidirectional( rnn(output_dim), merge_mode=mode, input_shape=(timesteps, dim))) - model.compile(loss='mse', optimizer='sgd') + model.compile(optimizer=RMSPropOptimizer(0.01), loss='mse') model.fit(x, y, epochs=1, batch_size=1) # test compute output shape @@ -338,23 +350,38 @@ class BidirectionalTest(test.TestCase): units = 3 with self.test_session(): - inputs = keras.Input((timesteps, dim)) + input1 = keras.layers.Input((timesteps, dim)) layer = keras.layers.Bidirectional( rnn(units, return_state=True, return_sequences=True)) - outputs = layer(inputs) - output, state = outputs[0], outputs[1:] + state = layer(input1)[1:] # test passing invalid initial_state: passing a tensor + input2 = keras.layers.Input((timesteps, dim)) with self.assertRaises(ValueError): output = keras.layers.Bidirectional( - rnn(units))(output, initial_state=state[0]) + rnn(units))(input2, initial_state=state[0]) # test valid usage: passing a list - output = keras.layers.Bidirectional( - rnn(units))(output, initial_state=state) - model = keras.Model(inputs, output) - inputs = np.random.rand(samples, timesteps, dim) - outputs = model.predict(inputs) + output = keras.layers.Bidirectional(rnn(units))(input2, + initial_state=state) + model = keras.models.Model([input1, input2], output) + assert len(model.layers) == 4 + assert isinstance(model.layers[-1].input, list) + inputs = [np.random.rand(samples, timesteps, dim), + np.random.rand(samples, timesteps, dim)] + model.predict(inputs) + + def test_Bidirectional_trainable(self): + # test layers that need learning_phase to be set + with self.test_session(): + x = keras.layers.Input(shape=(3, 2)) + layer = keras.layers.Bidirectional(keras.layers.SimpleRNN(3)) + _ = layer(x) + assert len(layer.trainable_weights) == 6 + layer.trainable = False + assert not layer.trainable_weights + layer.trainable = True + assert len(layer.trainable_weights) == 6 def _to_list(ls): diff --git a/tensorflow/python/keras/_impl/keras/losses.py b/tensorflow/python/keras/_impl/keras/losses.py index fe0ef54360608db5a7ae95b1ec89e32640db0931..1576ed7b999f65992f46b357c8ebeda8935c68d0 100644 --- a/tensorflow/python/keras/_impl/keras/losses.py +++ b/tensorflow/python/keras/_impl/keras/losses.py @@ -24,41 +24,54 @@ import six from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.generic_utils import deserialize_keras_object from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_keras_object +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.metrics.mean_squared_error', + 'keras.losses.mean_squared_error') def mean_squared_error(y_true, y_pred): return K.mean(K.square(y_pred - y_true), axis=-1) +@tf_export('keras.metrics.mean_absolute_error', + 'keras.losses.mean_absolute_error') def mean_absolute_error(y_true, y_pred): return K.mean(K.abs(y_pred - y_true), axis=-1) +@tf_export('keras.metrics.mean_absolute_percentage_error', + 'keras.losses.mean_absolute_percentage_error') def mean_absolute_percentage_error(y_true, y_pred): diff = K.abs((y_true - y_pred) / K.clip(K.abs(y_true), K.epsilon(), None)) return 100. * K.mean(diff, axis=-1) +@tf_export('keras.metrics.mean_squared_logarithmic_error', + 'keras.losses.mean_squared_logarithmic_error') def mean_squared_logarithmic_error(y_true, y_pred): first_log = K.log(K.clip(y_pred, K.epsilon(), None) + 1.) second_log = K.log(K.clip(y_true, K.epsilon(), None) + 1.) return K.mean(K.square(first_log - second_log), axis=-1) +@tf_export('keras.metrics.squared_hinge', 'keras.losses.squared_hinge') def squared_hinge(y_true, y_pred): return K.mean(K.square(K.maximum(1. - y_true * y_pred, 0.)), axis=-1) +@tf_export('keras.metrics.hinge', 'keras.losses.hinge') def hinge(y_true, y_pred): return K.mean(K.maximum(1. - y_true * y_pred, 0.), axis=-1) +@tf_export('keras.losses.categorical_hinge') def categorical_hinge(y_true, y_pred): pos = K.sum(y_true * y_pred, axis=-1) neg = K.max((1. - y_true) * y_pred, axis=-1) return K.maximum(0., neg - pos + 1.) +@tf_export('keras.losses.logcosh') def logcosh(y_true, y_pred): """Logarithm of the hyperbolic cosine of the prediction error. @@ -81,28 +94,38 @@ def logcosh(y_true, y_pred): return K.mean(_logcosh(y_pred - y_true), axis=-1) +@tf_export('keras.metrics.categorical_crossentropy', + 'keras.losses.categorical_crossentropy') def categorical_crossentropy(y_true, y_pred): return K.categorical_crossentropy(y_true, y_pred) +@tf_export('keras.metrics.sparse_categorical_crossentropy', + 'keras.losses.sparse_categorical_crossentropy') def sparse_categorical_crossentropy(y_true, y_pred): return K.sparse_categorical_crossentropy(y_true, y_pred) +@tf_export('keras.metrics.binary_crossentropy', + 'keras.losses.binary_crossentropy') def binary_crossentropy(y_true, y_pred): return K.mean(K.binary_crossentropy(y_true, y_pred), axis=-1) +@tf_export('keras.metrics.kullback_leibler_divergence', + 'keras.losses.kullback_leibler_divergence') def kullback_leibler_divergence(y_true, y_pred): y_true = K.clip(y_true, K.epsilon(), 1) y_pred = K.clip(y_pred, K.epsilon(), 1) return K.sum(y_true * K.log(y_true / y_pred), axis=-1) +@tf_export('keras.metrics.poisson', 'keras.losses.poisson') def poisson(y_true, y_pred): return K.mean(y_pred - y_true * K.log(y_pred + K.epsilon()), axis=-1) +@tf_export('keras.metrics.cosine_proximity', 'keras.losses.cosine_proximity') def cosine_proximity(y_true, y_pred): y_true = K.l2_normalize(y_true, axis=-1) y_pred = K.l2_normalize(y_pred, axis=-1) @@ -119,10 +142,12 @@ kld = KLD = kullback_leibler_divergence cosine = cosine_proximity +@tf_export('keras.losses.serialize') def serialize(loss): return serialize_keras_object(loss) +@tf_export('keras.losses.deserialize') def deserialize(name, custom_objects=None): return deserialize_keras_object( name, @@ -131,6 +156,7 @@ def deserialize(name, custom_objects=None): printable_module_name='loss function') +@tf_export('keras.losses.get') def get(identifier): if identifier is None: return None diff --git a/tensorflow/python/keras/_impl/keras/metrics.py b/tensorflow/python/keras/_impl/keras/metrics.py index 3c18e6826076b553136f5457f59e8bfe081bdb40..82778a3dc4fbdc13bb6682d01e28ff68882b6dd9 100644 --- a/tensorflow/python/keras/_impl/keras/metrics.py +++ b/tensorflow/python/keras/_impl/keras/metrics.py @@ -36,12 +36,16 @@ from tensorflow.python.keras._impl.keras.losses import poisson from tensorflow.python.keras._impl.keras.losses import sparse_categorical_crossentropy from tensorflow.python.keras._impl.keras.losses import squared_hinge from tensorflow.python.keras._impl.keras.utils.generic_utils import deserialize_keras_object +from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_keras_object +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.metrics.binary_accuracy') def binary_accuracy(y_true, y_pred): return K.mean(K.equal(y_true, K.round(y_pred)), axis=-1) +@tf_export('keras.metrics.categorical_accuracy') def categorical_accuracy(y_true, y_pred): return K.cast( K.equal(K.argmax(y_true, axis=-1), K.argmax(y_pred, axis=-1)), K.floatx()) @@ -54,10 +58,12 @@ def sparse_categorical_accuracy(y_true, y_pred): K.floatx())), K.floatx()) +@tf_export('keras.metrics.top_k_categorical_accuracy') def top_k_categorical_accuracy(y_true, y_pred, k=5): return K.mean(K.in_top_k(y_pred, K.argmax(y_true, axis=-1), k), axis=-1) +@tf_export('keras.metrics.sparse_top_k_categorical_accuracy') def sparse_top_k_categorical_accuracy(y_true, y_pred, k=5): return K.mean( K.in_top_k(y_pred, K.cast(K.max(y_true, axis=-1), 'int32'), k), axis=-1) @@ -72,24 +78,29 @@ msle = MSLE = mean_squared_logarithmic_error cosine = cosine_proximity +@tf_export('keras.metrics.serialize') def serialize(metric): - return metric.__name__ + return serialize_keras_object(metric) -def deserialize(name, custom_objects=None): +@tf_export('keras.metrics.deserialize') +def deserialize(config, custom_objects=None): return deserialize_keras_object( - name, + config, module_objects=globals(), custom_objects=custom_objects, printable_module_name='metric function') +@tf_export('keras.metrics.get') def get(identifier): - if isinstance(identifier, six.string_types): - identifier = str(identifier) - return deserialize(identifier) + if isinstance(identifier, dict): + config = {'class_name': str(identifier), 'config': {}} + return deserialize(config) + elif isinstance(identifier, six.string_types): + return deserialize(str(identifier)) elif callable(identifier): return identifier else: raise ValueError('Could not interpret ' - 'metric function identifier:', identifier) + 'metric function identifier: %s' % identifier) diff --git a/tensorflow/python/keras/_impl/keras/metrics_test.py b/tensorflow/python/keras/_impl/keras/metrics_test.py index f4792f3543cc5ca8e5e7ad03d9906bbfadd1fb04..44289ea02abf5ae5f8befbe515552aea3d4b231e 100644 --- a/tensorflow/python/keras/_impl/keras/metrics_test.py +++ b/tensorflow/python/keras/_impl/keras/metrics_test.py @@ -72,6 +72,77 @@ class KerasMetricsTest(test.TestCase): keras.metrics.top_k_categorical_accuracy(y_true, y_pred, k=1)) self.assertEqual(result, 0.) + def test_stateful_metrics(self): + np.random.seed(1334) + + class BinaryTruePositives(keras.layers.Layer): + """Stateful Metric to count the total true positives over all batches. + + Assumes predictions and targets of shape `(samples, 1)`. + + Arguments: + threshold: Float, lower limit on prediction value that counts as a + positive class prediction. + name: String, name for the metric. + """ + + def __init__(self, name='true_positives', **kwargs): + super(BinaryTruePositives, self).__init__(name=name, **kwargs) + self.true_positives = keras.backend.variable(value=0, dtype='int32') + + def reset_states(self): + keras.backend.set_value(self.true_positives, 0) + + def __call__(self, y_true, y_pred): + """Computes the number of true positives in a batch. + + Args: + y_true: Tensor, batch_wise labels + y_pred: Tensor, batch_wise predictions + + Returns: + The total number of true positives seen this epoch at the + completion of the batch. + """ + y_true = keras.backend.cast(y_true, 'int32') + y_pred = keras.backend.cast(keras.backend.round(y_pred), 'int32') + correct_preds = keras.backend.cast( + keras.backend.equal(y_pred, y_true), 'int32') + true_pos = keras.backend.cast( + keras.backend.sum(correct_preds * y_true), 'int32') + current_true_pos = self.true_positives * 1 + self.add_update(keras.backend.update_add(self.true_positives, + true_pos), + inputs=[y_true, y_pred]) + return current_true_pos + true_pos + + metric_fn = BinaryTruePositives() + config = keras.metrics.serialize(metric_fn) + metric_fn = keras.metrics.deserialize( + config, custom_objects={'BinaryTruePositives': BinaryTruePositives}) + + # Test on simple model + inputs = keras.Input(shape=(2,)) + outputs = keras.layers.Dense(1, activation='sigmoid')(inputs) + model = keras.Model(inputs, outputs) + model.compile(optimizer='sgd', + loss='binary_crossentropy', + metrics=['acc', metric_fn]) + + # Test fit, evaluate + samples = 1000 + x = np.random.random((samples, 2)) + y = np.random.randint(2, size=(samples, 1)) + model.fit(x, y, epochs=1, batch_size=10) + outs = model.evaluate(x, y, batch_size=10) + preds = model.predict(x) + + def ref_true_pos(y_true, y_pred): + return np.sum(np.logical_and(y_pred > 0.5, y_true == 1)) + + # Test correctness (e.g. updates should have been run) + self.assertAllClose(outs[2], ref_true_pos(y, preds), atol=1e-5) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py new file mode 100644 index 0000000000000000000000000000000000000000..58b144365be6cd8ea5b2ea82e275eacdee6b6c84 --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py @@ -0,0 +1,578 @@ +# 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 Model subclassing.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import tempfile + +import numpy as np + +from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import test_util +from tensorflow.python.keras._impl import keras +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test +from tensorflow.python.training.rmsprop import RMSPropOptimizer + +try: + import h5py # pylint:disable=g-import-not-at-top +except ImportError: + h5py = None + + +class SimpleTestModel(keras.Model): + + def __init__(self, use_bn=False, use_dp=False, num_classes=10): + super(SimpleTestModel, self).__init__(name='test_model') + self.use_bn = use_bn + self.use_dp = use_dp + self.num_classes = num_classes + + self.dense1 = keras.layers.Dense(32, activation='relu') + self.dense2 = keras.layers.Dense(num_classes, activation='softmax') + if self.use_dp: + self.dp = keras.layers.Dropout(0.5) + if self.use_bn: + self.bn = keras.layers.BatchNormalization(axis=-1) + + def call(self, inputs): + x = self.dense1(inputs) + if self.use_dp: + x = self.dp(x) + if self.use_bn: + x = self.bn(x) + return self.dense2(x) + + +class MultiIOTestModel(keras.Model): + + def __init__(self, use_bn=False, use_dp=False, num_classes=(2, 3)): + super(MultiIOTestModel, self).__init__(name='test_model') + self.use_bn = use_bn + self.use_dp = use_dp + self.num_classes = num_classes + + self.dense1 = keras.layers.Dense(32, activation='relu') + self.dense2 = keras.layers.Dense(num_classes[0], activation='softmax') + self.dense3 = keras.layers.Dense(num_classes[1], activation='softmax') + if use_dp: + self.dp = keras.layers.Dropout(0.5) + if use_bn: + self.bn = keras.layers.BatchNormalization() + + def call(self, inputs): + x1, x2 = inputs + x1 = self.dense1(x1) + x2 = self.dense1(x2) + if self.use_dp: + x1 = self.dp(x1) + if self.use_bn: + x2 = self.bn(x2) + return [self.dense2(x1), self.dense3(x2)] + + +class NestedTestModel1(keras.Model): + """A model subclass nested inside a model subclass. + """ + + def __init__(self, num_classes=2): + super(NestedTestModel1, self).__init__(name='nested_model_1') + self.num_classes = num_classes + self.dense1 = keras.layers.Dense(32, activation='relu') + self.dense2 = keras.layers.Dense(num_classes, activation='relu') + self.bn = keras.layers.BatchNormalization() + self.test_net = SimpleTestModel(num_classes=4, + use_bn=True, + use_dp=True) + + def call(self, inputs): + x = self.dense1(inputs) + x = self.bn(x) + x = self.test_net(x) # pylint: disable=not-callable + return self.dense2(x) + + +def get_functional_graph_model(input_dim, num_classes): + # A simple functional-API model (a.k.a. graph network) + inputs = keras.Input(shape=(input_dim,)) + x = keras.layers.Dense(32, activation='relu')(inputs) + x = keras.layers.BatchNormalization()(x) + outputs = keras.layers.Dense(num_classes)(x) + return keras.Model(inputs, outputs) + + +class NestedTestModel2(keras.Model): + """A model subclass with a functional-API graph network inside. + """ + + def __init__(self, num_classes=2): + super(NestedTestModel2, self).__init__(name='nested_model_2') + self.num_classes = num_classes + self.dense1 = keras.layers.Dense(32, activation='relu') + self.dense2 = keras.layers.Dense(num_classes, activation='relu') + self.bn = self.bn = keras.layers.BatchNormalization() + self.test_net = get_functional_graph_model(32, 4) + + def call(self, inputs): + x = self.dense1(inputs) + x = self.bn(x) + x = self.test_net(x) + return self.dense2(x) + + +def get_nested_model_3(input_dim, num_classes): + # A functional-API model with a subclassed model inside. + # NOTE: this requires the inner subclass to implement `compute_output_shape`. + + inputs = keras.Input(shape=(input_dim,)) + x = keras.layers.Dense(32, activation='relu')(inputs) + x = keras.layers.BatchNormalization()(x) + + class Inner(keras.Model): + + def __init__(self): + super(Inner, self).__init__() + self.dense1 = keras.layers.Dense(32, activation='relu') + self.dense2 = keras.layers.Dense(5, activation='relu') + self.bn = keras.layers.BatchNormalization() + + def call(self, inputs): + x = self.dense1(inputs) + x = self.dense2(x) + return self.bn(x) + + def compute_output_shape(self, input_shape): + return tensor_shape.TensorShape((input_shape[0], 5)) + + test_model = Inner() + x = test_model(x) # pylint: disable=not-callable + outputs = keras.layers.Dense(num_classes)(x) + return keras.Model(inputs, outputs, name='nested_model_3') + + +class ModelSubclassingTest(test.TestCase): + + @test_util.run_in_graph_and_eager_modes() + def test_single_io_workflow_with_np_arrays(self): + num_classes = 2 + num_samples = 100 + input_dim = 50 + + model = SimpleTestModel(num_classes=num_classes, + use_dp=True, + use_bn=True) + model.compile(loss='mse', + optimizer=RMSPropOptimizer(learning_rate=0.001), + metrics=['acc']) + + x = np.ones((num_samples, input_dim)) + y = np.zeros((num_samples, num_classes)) + + model.fit(x, y, epochs=2, batch_size=32, verbose=0) + _ = model.evaluate(x, y, verbose=0) + + @test_util.run_in_graph_and_eager_modes() + def test_multi_io_workflow_with_np_arrays(self): + num_classes = (2, 3) + num_samples = 1000 + input_dim = 50 + + model = MultiIOTestModel(num_classes=num_classes, + use_dp=True, + use_bn=True) + model.compile(loss='mse', + optimizer=RMSPropOptimizer(learning_rate=0.001), + metrics=['acc']) + + x1 = np.ones((num_samples, input_dim)) + x2 = np.ones((num_samples, input_dim)) + y1 = np.zeros((num_samples, num_classes[0])) + y2 = np.zeros((num_samples, num_classes[1])) + + model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0) + _ = model.evaluate([x1, x2], [y1, y2], verbose=0) + + def test_single_io_workflow_with_tensors(self): + + num_classes = 2 + num_samples = 10 + input_dim = 50 + + with self.test_session(): + model = SimpleTestModel(num_classes=num_classes, + use_dp=True, + use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + + x = array_ops.ones((num_samples, input_dim)) + y = array_ops.zeros((num_samples, num_classes)) + + model.fit(x, y, epochs=2, steps_per_epoch=10, verbose=0) + _ = model.evaluate(steps=10, verbose=0) + + def test_multi_io_workflow_with_tensors(self): + + num_classes = (2, 3) + num_samples = 10 + input_dim = 50 + + with self.test_session(): + model = MultiIOTestModel(num_classes=num_classes, + use_dp=True, + use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + + x1 = array_ops.ones((num_samples, input_dim)) + x2 = array_ops.ones((num_samples, input_dim)) + y1 = array_ops.zeros((num_samples, num_classes[0])) + y2 = array_ops.zeros((num_samples, num_classes[1])) + + model.fit([x1, x2], [y1, y2], epochs=2, steps_per_epoch=10, verbose=0) + _ = model.evaluate(steps=10, verbose=0) + + def test_multi_io_workflow_with_numpy_arrays_and_custom_placeholders(self): + + num_classes = (2, 3) + num_samples = 1000 + input_dim = 50 + + with self.test_session(): + model = MultiIOTestModel(num_classes=num_classes, + use_dp=True, + use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + + x1 = np.ones((num_samples, input_dim)) + x2 = np.ones((num_samples, input_dim)) + y1 = np.zeros((num_samples, num_classes[0])) + y2 = np.zeros((num_samples, num_classes[1])) + + x2_placeholder = array_ops.placeholder( + dtype='float32', shape=(None, input_dim)) + model._set_inputs([x1, x2_placeholder]) + + model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0) + _ = model.evaluate([x1, x2], [y1, y2], verbose=0) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def test_attributes(self): + # layers, weights, trainable_weights, non_trainable_weights, inputs, outputs + + num_classes = (2, 3) + num_samples = 100 + input_dim = 50 + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + + x1 = np.ones((num_samples, input_dim)) + x2 = np.ones((num_samples, input_dim)) + y1 = np.zeros((num_samples, num_classes[0])) + y2 = np.zeros((num_samples, num_classes[1])) + + self.assertEqual(model.name, 'test_model') + self.assertEqual(model.built, False) + self.assertEqual(len(model.weights), 0) + + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + model.train_on_batch([x1, x2], [y1, y2]) + + self.assertEqual(model.built, True) + self.assertEqual(len(model.layers), 4) + self.assertEqual(len(model.weights), 10) + self.assertEqual(len(model.trainable_weights), 8) + self.assertEqual(len(model.non_trainable_weights), 2) + self.assertEqual(len(model.inputs), 2) + self.assertEqual(len(model.outputs), 2) + + @test_util.run_in_graph_and_eager_modes() + def test_updates(self): + # test that updates get run during training + num_samples = 100 + input_dim = 50 + + class BNNet(keras.Model): + + def __init__(self): + super(BNNet, self).__init__() + self.bn = keras.layers.BatchNormalization(beta_initializer='ones', + gamma_initializer='ones') + + def call(self, inputs): + return self.bn(inputs) + + x = np.ones((num_samples, input_dim)) + y = np.ones((num_samples, input_dim)) + + model = BNNet() + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + y_ref = model.predict(x) + + model.train_on_batch(x, y) + y_new = model.predict(x) + self.assertGreater(np.sum(np.abs(y_ref - y_new)), 0.1) + + @test_util.run_in_graph_and_eager_modes() + def test_training_and_inference_behavior(self): + # test that dropout is applied in training and not inference + + num_samples = 100 + input_dim = 50 + + class DPNet(keras.Model): + + def __init__(self): + super(DPNet, self).__init__() + self.dp = keras.layers.Dropout(0.5) + self.dense = keras.layers.Dense(1, + use_bias=False, + kernel_initializer='ones') + + def call(self, inputs): + x = self.dp(inputs) + return self.dense(x) + + model = DPNet() + x = np.ones((num_samples, input_dim)) + y = model.predict(x) + self.assertEqual(np.sum(y), np.sum(x)) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + loss = model.train_on_batch(x, y) + self.assertGreater(loss, 0.1) + + @test_util.run_in_graph_and_eager_modes() + def test_training_methods(self): + # test fit, train_on_batch + # on different input types: list, dict + + num_classes = (2, 3) + num_samples = 100 + input_dim = 50 + + x1 = np.ones((num_samples, input_dim)) + x2 = np.ones((num_samples, input_dim)) + y1 = np.zeros((num_samples, num_classes[0])) + y2 = np.zeros((num_samples, num_classes[1])) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0) + model.fit({'input_1': x1, 'input_2': x2}, + {'output_1': y1, 'output_2': y2}, + epochs=2, batch_size=32) + model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0, + validation_data=([x1, x2], [y1, y2])) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + model.train_on_batch([x1, x2], [y1, y2]) + model.train_on_batch({'input_1': x1, 'input_2': x2}, + {'output_1': y1, 'output_2': y2}) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def test_inference_methods(self): + # test predict, evaluate, test_on_batch, predict_on_batch + # on different input types: list, dict + num_classes = (2, 3) + num_samples = 100 + input_dim = 50 + + x1 = np.ones((num_samples, input_dim)) + x2 = np.ones((num_samples, input_dim)) + y1 = np.zeros((num_samples, num_classes[0])) + y2 = np.zeros((num_samples, num_classes[1])) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + model.evaluate([x1, x2], [y1, y2]) + model.test_on_batch([x1, x2], [y1, y2]) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + model.predict([x1, x2]) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + model.predict_on_batch([x1, x2]) + + @test_util.run_in_graph_and_eager_modes() + def test_trainable_mutation(self): + # test that you can change `trainable` on a model or layer, and that + # it freezes the model state during training + # TODO(fchollet): add test after we unify BN behavior in eager and symbolic. + pass + + @test_util.run_in_graph_and_eager_modes() + def test_saving(self): + if h5py is None: + return # Skip test if models cannot be saved. + + num_classes = (2, 3) + num_samples = 100 + input_dim = 50 + + x1 = np.ones((num_samples, input_dim)) + x2 = np.ones((num_samples, input_dim)) + y1 = np.zeros((num_samples, num_classes[0])) + y2 = np.zeros((num_samples, num_classes[1])) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0) + y_ref_1, y_ref_2 = model.predict([x1, x2]) + + fd, fname = tempfile.mkstemp('.h5') + model.save_weights(fname) + + model = MultiIOTestModel(num_classes=num_classes, use_bn=True) + # need to build the model before loading weights + # (otherwise no weights to load) + model._set_inputs([x1, x2]) + model.load_weights(fname) + + y1, y2 = model.predict([x1, x2]) + self.assertAllClose(y_ref_1, y1, atol=1e-5) + self.assertAllClose(y_ref_2, y2, atol=1e-5) + os.close(fd) + os.remove(fname) + + @test_util.run_in_graph_and_eager_modes() + def test_summary(self): + + class ToString(object): + + def __init__(self): + self.contents = '' + + def __call__(self, msg): + self.contents += msg + '\n' + + # Single-io + model = SimpleTestModel(num_classes=4, use_bn=True, use_dp=True) + model._set_inputs(np.ones((3, 4))) # need to build model first + print_fn = ToString() + model.summary(print_fn=print_fn) + self.assertTrue('Trainable params: 356' in print_fn.contents) + + # Multi-io + model = MultiIOTestModel(num_classes=(5, 6), use_bn=True, use_dp=True) + model._set_inputs([np.ones((3, 4)), + np.ones((3, 4))]) # need to build model first + print_fn = ToString() + model.summary(print_fn=print_fn) + self.assertTrue('Trainable params: 587' in print_fn.contents) + + @test_util.run_in_graph_and_eager_modes() + def test_subclass_nested_in_subclass(self): + num_classes = 2 + num_samples = 100 + input_dim = 50 + + model = NestedTestModel1(num_classes=num_classes) + model.compile(loss='mse', + optimizer=RMSPropOptimizer(learning_rate=0.001), + metrics=['acc']) + + x = np.ones((num_samples, input_dim)) + y = np.zeros((num_samples, num_classes)) + + model.fit(x, y, epochs=2, batch_size=32, verbose=0) + _ = model.evaluate(x, y, verbose=0) + + self.assertEqual(len(model.weights), 8 + len(model.test_net.weights)) + self.assertEqual(len(model.non_trainable_weights), + 2 + len(model.test_net.non_trainable_weights)) + self.assertEqual(len(model.trainable_weights), + 6 + len(model.test_net.trainable_weights)) + + @test_util.run_in_graph_and_eager_modes() + def test_graph_nested_in_subclass(self): + num_classes = 2 + num_samples = 100 + input_dim = 50 + + model = NestedTestModel2(num_classes=num_classes) + model.compile(loss='mse', + optimizer=RMSPropOptimizer(learning_rate=0.001), + metrics=['acc']) + + x = np.ones((num_samples, input_dim)) + y = np.zeros((num_samples, num_classes)) + + model.fit(x, y, epochs=2, batch_size=32, verbose=0) + _ = model.evaluate(x, y, verbose=0) + + self.assertEqual(len(model.weights), 8 + len(model.test_net.weights)) + self.assertEqual(len(model.non_trainable_weights), + 2 + len(model.test_net.non_trainable_weights)) + self.assertEqual(len(model.trainable_weights), + 6 + len(model.test_net.trainable_weights)) + + @test_util.run_in_graph_and_eager_modes() + def test_subclass_nested_in_graph(self): + num_classes = 2 + num_samples = 100 + input_dim = 50 + + model = get_nested_model_3(input_dim=input_dim, num_classes=num_classes) + model.compile(loss='mse', + optimizer=RMSPropOptimizer(learning_rate=0.001), + metrics=['acc']) + + x = np.ones((num_samples, input_dim)) + y = np.zeros((num_samples, num_classes)) + + model.fit(x, y, epochs=2, batch_size=32, verbose=0) + _ = model.evaluate(x, y, verbose=0) + + self.assertEqual(len(model.weights), 16) + self.assertEqual( + len(model.non_trainable_weights), 4) + self.assertEqual(len(model.trainable_weights), 12) + + @test_util.run_in_graph_and_eager_modes() + def test_support_for_manual_training_arg(self): + # In most cases, the `training` argument is left unspecified, in which + # case it defaults to value corresponding to the Model method being used + # (fit -> True, predict -> False, etc). + # If the user writes their model `call` method to take + # an explicit `training` argument, we must check that the correct value + # is being passed to the model for each method call. + + class DPNet(keras.Model): + + def __init__(self): + super(DPNet, self).__init__() + self.dp = keras.layers.Dropout(0.5) + self.dense = keras.layers.Dense(1, + use_bias=False, + kernel_initializer='ones') + + def call(self, inputs, training=False): + x = self.dp(inputs, training=training) + return self.dense(x) + + model = DPNet() + x = np.ones((10, 10)) + y = model.predict(x) + self.assertEqual(np.sum(y), np.sum(x)) + model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001)) + loss = model.train_on_batch(x, y) + self.assertGreater(loss, 0.1) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/keras/_impl/keras/models.py b/tensorflow/python/keras/_impl/keras/models.py index 9cd547200d89184bb7c2aad28b59d2793f093205..9602e7ba39b290f33c7ca9d0d1b5b35838667531 100644 --- a/tensorflow/python/keras/_impl/keras/models.py +++ b/tensorflow/python/keras/_impl/keras/models.py @@ -13,1326 +13,30 @@ # limitations under the License. # ============================================================================== # pylint: disable=protected-access -"""Home of the Sequential model, and the `save_model`/`load_model` functions. +"""Code for model cloning, plus model-related API entries. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import copy -import json -import os - -import numpy as np - -from tensorflow.python.framework import ops from tensorflow.python.keras._impl.keras import backend as K -from tensorflow.python.keras._impl.keras import layers as layer_module -from tensorflow.python.keras._impl.keras import optimizers -from tensorflow.python.keras._impl.keras.engine import topology -from tensorflow.python.keras._impl.keras.engine.topology import Input -from tensorflow.python.keras._impl.keras.engine.topology import InputLayer -from tensorflow.python.keras._impl.keras.engine.topology import Layer -from tensorflow.python.keras._impl.keras.engine.topology import TFBaseLayer -from tensorflow.python.keras._impl.keras.engine.training import Model +from tensorflow.python.keras._impl.keras.engine import saving +from tensorflow.python.keras._impl.keras.engine import sequential +from tensorflow.python.keras._impl.keras.engine import training +from tensorflow.python.keras._impl.keras.engine.input_layer import Input +from tensorflow.python.keras._impl.keras.engine.input_layer import InputLayer +from tensorflow.python.keras._impl.keras.utils import generic_utils from tensorflow.python.keras._impl.keras.utils.generic_utils import has_arg -from tensorflow.python.keras._impl.keras.utils.io_utils import ask_to_proceed_with_overwrite -from tensorflow.python.platform import tf_logging as logging - - -# pylint: disable=g-import-not-at-top -try: - import h5py -except ImportError: - h5py = None - -try: - import yaml -except ImportError: - yaml = None -# pylint: enable=g-import-not-at-top - - -def save_model(model, filepath, overwrite=True, include_optimizer=True): - """Save a model to a HDF5 file. - - The saved model contains: - - the model's configuration (topology) - - the model's weights - - the model's optimizer's state (if any) - - Thus the saved model can be reinstantiated in - the exact same state, without any of the code - used for model definition or training. - - Arguments: - model: Keras model instance to be saved. - filepath: String, path where to save the model. - overwrite: Whether we should overwrite any existing - model at the target location, or instead - ask the user with a manual prompt. - include_optimizer: If True, save optimizer's state together. - - Raises: - ImportError: if h5py is not available. - """ - - if h5py is None: - raise ImportError('`save_model` requires h5py.') - - def get_json_type(obj): - """Serialize any object to a JSON-serializable structure. - - Arguments: - obj: the object to serialize - - Returns: - JSON-serializable structure representing `obj`. - - Raises: - TypeError: if `obj` cannot be serialized. - """ - # if obj is a serializable Keras class instance - # e.g. optimizer, layer - if hasattr(obj, 'get_config'): - return {'class_name': obj.__class__.__name__, 'config': obj.get_config()} - - # if obj is any numpy type - if type(obj).__module__ == np.__name__: - if isinstance(obj, np.ndarray): - return {'type': type(obj), 'value': obj.tolist()} - else: - return obj.item() - - # misc functions (e.g. loss function) - if callable(obj): - return obj.__name__ - - # if obj is a python 'type' - if type(obj).__name__ == type.__name__: - return obj.__name__ - - raise TypeError('Not JSON Serializable:', obj) - - from tensorflow.python.keras._impl.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top - - # If file exists and should not be overwritten. - if not overwrite and os.path.isfile(filepath): - proceed = ask_to_proceed_with_overwrite(filepath) - if not proceed: - return - - with h5py.File(filepath, mode='w') as f: - f.attrs['keras_version'] = str(keras_version).encode('utf8') - f.attrs['backend'] = K.backend().encode('utf8') - f.attrs['model_config'] = json.dumps( - { - 'class_name': model.__class__.__name__, - 'config': model.get_config() - }, - default=get_json_type).encode('utf8') - - model_weights_group = f.create_group('model_weights') - model_layers = model.layers - topology.save_weights_to_hdf5_group(model_weights_group, model_layers) - - if include_optimizer and hasattr(model, 'optimizer'): - if isinstance(model.optimizer, optimizers.TFOptimizer): - logging.warning( - 'TensorFlow optimizers do not ' - 'make it possible to access ' - 'optimizer attributes or optimizer state ' - 'after instantiation. ' - 'As a result, we cannot save the optimizer ' - 'as part of the model save file.' - 'You will have to compile your model again after loading it. ' - 'Prefer using a Keras optimizer instead ' - '(see keras.io/optimizers).') - else: - f.attrs['training_config'] = json.dumps( - { - 'optimizer_config': { - 'class_name': model.optimizer.__class__.__name__, - 'config': model.optimizer.get_config() - }, - 'loss': model.loss, - 'metrics': model.metrics, - 'sample_weight_mode': model.sample_weight_mode, - 'loss_weights': model.loss_weights, - }, - default=get_json_type).encode('utf8') - - # Save optimizer weights. - symbolic_weights = getattr(model.optimizer, 'weights') - if symbolic_weights: - optimizer_weights_group = f.create_group('optimizer_weights') - weight_values = K.batch_get_value(symbolic_weights) - weight_names = [] - for w, val in zip(symbolic_weights, weight_values): - name = str(w.name) - weight_names.append(name.encode('utf8')) - optimizer_weights_group.attrs['weight_names'] = weight_names - for name, val in zip(weight_names, weight_values): - param_dset = optimizer_weights_group.create_dataset( - name, val.shape, dtype=val.dtype) - if not val.shape: - # scalar - param_dset[()] = val - else: - param_dset[:] = val - f.flush() - - -def load_model(filepath, custom_objects=None, compile=True): # pylint: disable=redefined-builtin - """Loads a model saved via `save_model`. - - Arguments: - filepath: String, path to the saved model. - custom_objects: Optional dictionary mapping names - (strings) to custom classes or functions to be - considered during deserialization. - compile: Boolean, whether to compile the model - after loading. - - Returns: - A Keras model instance. If an optimizer was found - as part of the saved model, the model is already - compiled. Otherwise, the model is uncompiled and - a warning will be displayed. When `compile` is set - to False, the compilation is omitted without any - warning. - - Raises: - ImportError: if h5py is not available. - ValueError: In case of an invalid savefile. - """ - if h5py is None: - raise ImportError('`load_model` requires h5py.') - - if not custom_objects: - custom_objects = {} - - def convert_custom_objects(obj): - """Handles custom object lookup. - - Arguments: - obj: object, dict, or list. - - Returns: - The same structure, where occurrences - of a custom object name have been replaced - with the custom object. - """ - if isinstance(obj, list): - deserialized = [] - for value in obj: - deserialized.append(convert_custom_objects(value)) - return deserialized - if isinstance(obj, dict): - deserialized = {} - for key, value in obj.items(): - deserialized[key] = convert_custom_objects(value) - return deserialized - if obj in custom_objects: - return custom_objects[obj] - return obj - - with h5py.File(filepath, mode='r') as f: - # instantiate model - model_config = f.attrs.get('model_config') - if model_config is None: - raise ValueError('No model found in config file.') - model_config = json.loads(model_config.decode('utf-8')) - model = model_from_config(model_config, custom_objects=custom_objects) - - # set weights - topology.load_weights_from_hdf5_group(f['model_weights'], model.layers) - - # Early return if compilation is not required. - if not compile: - return model - - # instantiate optimizer - training_config = f.attrs.get('training_config') - if training_config is None: - logging.warning('No training configuration found in save file: ' - 'the model was *not* compiled. Compile it manually.') - return model - training_config = json.loads(training_config.decode('utf-8')) - optimizer_config = training_config['optimizer_config'] - optimizer = optimizers.deserialize( - optimizer_config, custom_objects=custom_objects) - - # Recover loss functions and metrics. - loss = convert_custom_objects(training_config['loss']) - metrics = convert_custom_objects(training_config['metrics']) - sample_weight_mode = training_config['sample_weight_mode'] - loss_weights = training_config['loss_weights'] - - # Compile model. - model.compile( - optimizer=optimizer, - loss=loss, - metrics=metrics, - loss_weights=loss_weights, - sample_weight_mode=sample_weight_mode) - - # Set optimizer weights. - if 'optimizer_weights' in f: - # Build train function (to get weight updates). - if isinstance(model, Sequential): - model.model._make_train_function() - else: - model._make_train_function() - optimizer_weights_group = f['optimizer_weights'] - optimizer_weight_names = [ - n.decode('utf8') - for n in optimizer_weights_group.attrs['weight_names'] - ] - optimizer_weight_values = [ - optimizer_weights_group[n] for n in optimizer_weight_names - ] - try: - model.optimizer.set_weights(optimizer_weight_values) - except ValueError: - logging.warning('Error in loading the saved optimizer ' - 'state. As a result, your model is ' - 'starting with a freshly initialized ' - 'optimizer.') - return model - - -def model_from_config(config, custom_objects=None): - """Instantiates a Keras model from its config. - - Arguments: - config: Configuration dictionary. - custom_objects: Optional dictionary mapping names - (strings) to custom classes or functions to be - considered during deserialization. - - Returns: - A Keras model instance (uncompiled). - - Raises: - TypeError: if `config` is not a dictionary. - """ - if isinstance(config, list): - raise TypeError('`model_from_config` expects a dictionary, not a list. ' - 'Maybe you meant to use ' - '`Sequential.from_config(config)`?') - return layer_module.deserialize(config, custom_objects=custom_objects) - - -def model_from_yaml(yaml_string, custom_objects=None): - """Parses a yaml model configuration file and returns a model instance. - - Arguments: - yaml_string: YAML string encoding a model configuration. - custom_objects: Optional dictionary mapping names - (strings) to custom classes or functions to be - considered during deserialization. - - Returns: - A Keras model instance (uncompiled). - - Raises: - ImportError: if yaml module is not found. - """ - if yaml is None: - raise ImportError('Requires yaml module installed.') - config = yaml.load(yaml_string) - return layer_module.deserialize(config, custom_objects=custom_objects) - - -def model_from_json(json_string, custom_objects=None): - """Parses a JSON model configuration file and returns a model instance. - - Arguments: - json_string: JSON string encoding a model configuration. - custom_objects: Optional dictionary mapping names - (strings) to custom classes or functions to be - considered during deserialization. - - Returns: - A Keras model instance (uncompiled). - """ - config = json.loads(json_string) - return layer_module.deserialize(config, custom_objects=custom_objects) - - -class Sequential(Model): - """Linear stack of layers. - - Arguments: - layers: list of layers to add to the model. - - # Note - The first layer passed to a Sequential model - should have a defined input shape. What that - means is that it should have received an `input_shape` - or `batch_input_shape` argument, - or for some type of layers (recurrent, Dense...) - an `input_dim` argument. - - Example: - - ```python - model = Sequential() - # first layer must have a defined input shape - model.add(Dense(32, input_dim=500)) - # afterwards, Keras does automatic shape inference - model.add(Dense(32)) - - # also possible (equivalent to the above): - model = Sequential() - model.add(Dense(32, input_shape=(500,))) - model.add(Dense(32)) - - # also possible (equivalent to the above): - model = Sequential() - # here the batch dimension is None, - # which means any batch size will be accepted by the model. - model.add(Dense(32, batch_input_shape=(None, 500))) - model.add(Dense(32)) - ``` - """ - - def __init__(self, layers=None, name=None): - self.layers = [] # Stack of layers. - self.model = None # Internal Model instance. - self.inputs = [] # List of input tensors - self.outputs = [] # List of length 1: the output tensor (unique). - self._trainable = True - self._initial_weights = None - self._input_layers = [] - - # Model attributes. - self._inbound_nodes = [] - self._outbound_nodes = [] - self.built = False - - # Set model name. - if not name: - prefix = 'sequential_' - name = prefix + str(K.get_uid(prefix)) - self._name = name - - # Used by Layer base class. - self._dtype = None - self._activity_regularizer = None - self._per_input_losses = {} - self._per_input_updates = {} - - # The following properties are not actually used by Keras; - # they exist for compatibility with TF's variable scoping mechanism. - self._updates = [] - self._losses = [] - self._scope = None - self._reuse = None - self._base_name = name - self._graph = ops.get_default_graph() - - # Add to the model any layers passed to the constructor. - if layers: - for layer in layers: - self.add(layer) - - def add(self, layer): - """Adds a layer instance on top of the layer stack. - - Arguments: - layer: layer instance. - - Raises: - TypeError: If `layer` is not a layer instance. - ValueError: In case the `layer` argument does not - know its input shape. - ValueError: In case the `layer` argument has - multiple output tensors, or is already connected - somewhere else (forbidden in `Sequential` models). - """ - if not isinstance(layer, (Layer, TFBaseLayer)): - raise TypeError('The added layer must be ' - 'an instance of class Layer. ' - 'Found: ' + str(layer)) - if not self.outputs: - # First layer in model: check that it is an input layer. - if not isinstance(layer, InputLayer): - # Create an input layer. - # First, we need to infer its expected input shape and dtype. - if isinstance(layer, (Model, Sequential)): - # We were passed a model as first layer. - # This requires a specific way to figure out the - # input shape and dtype. - if not layer.layers: - raise ValueError('Cannot add an empty model ' - 'to a `Sequential` model.') - # In case of nested models: recover the first layer - # of the deepest model to infer input shape and dtype. - first_layer = layer.layers[0] - while isinstance(first_layer, (Model, Sequential)): - first_layer = first_layer.layers[0] - batch_shape = first_layer._batch_input_shape - dtype = first_layer.dtype - else: - # We were passed a regular layer, and it should - # know about its input shape. Otherwise, that's an error. - if not hasattr(layer, '_batch_input_shape'): - raise ValueError('The first layer in a ' - 'Sequential model must ' - 'get an `input_shape` argument.') - batch_shape = layer._batch_input_shape - dtype = layer.dtype - # Instantiate the input layer. - x = Input( - batch_shape=batch_shape, dtype=dtype, name=layer.name + '_input') - # This will build the current layer - # and create the node connecting the current layer - # to the input layer we just created. - layer(x) - - if len(layer._inbound_nodes[-1].output_tensors) != 1: - raise ValueError('All layers in a Sequential model ' - 'should have a single output tensor. ' - 'For multi-output layers, ' - 'use the functional API.') - - self.outputs = [layer._inbound_nodes[-1].output_tensors[0]] - self.inputs = topology.get_source_inputs(self.outputs[0]) - - # We create an input node, which we will keep updated - # as we add more layers - topology.Node( - outbound_layer=self, - inbound_layers=[], - node_indices=[], - tensor_indices=[], - input_tensors=self.inputs, - output_tensors=self.outputs) - else: - output_tensor = layer(self.outputs[0]) - if isinstance(output_tensor, list): - raise TypeError('All layers in a Sequential model ' - 'should have a single output tensor. ' - 'For multi-output layers, ' - 'use the functional API.') - self.outputs = [output_tensor] - # update self._inbound_nodes - self._inbound_nodes[0].output_tensors = self.outputs - self._inbound_nodes[0].output_shapes = [K.int_shape(self.outputs[0])] - - self.layers.append(layer) - self.built = False - - def pop(self): - """Removes the last layer in the model. - - Raises: - TypeError: if there are no layers in the model. - """ - if not self.layers: - raise TypeError('There are no layers in the model.') - - self.layers.pop() - if not self.layers: - self.outputs = [] - self._inbound_nodes = [] - self._outbound_nodes = [] - else: - self.layers[-1]._outbound_nodes = [] - self.outputs = [self.layers[-1].output] - # update self._inbound_nodes - self._inbound_nodes[0].output_tensors = self.outputs - self._inbound_nodes[0].output_shapes = [K.int_shape(self.outputs[0])] - self.built = False - - def get_layer(self, name=None, index=None): - """Retrieve a layer that is part of the model. - - Returns a layer based on either its name (unique) - or its index in the graph. Indices are based on - order of horizontal graph traversal (bottom-up). - - Arguments: - name: string, name of layer. - index: integer, index of layer. - - Returns: - A layer instance. - """ - if not self.built: - self.build() - return self.model.get_layer(name, index) - - def call(self, inputs, mask=None): - if not self.built: - self.build() - return self.model.call(inputs, mask) - - def build(self, input_shape=None): - if not self.inputs or not self.outputs: - raise TypeError('Sequential model cannot be built: model is empty.' - ' Add some layers first.') - # actually create the model - self.model = Model(self.inputs, self.outputs[0], name=self.name + '_model') - self.model.trainable = self.trainable - - # mirror model attributes - self.supports_masking = self.model.supports_masking - self._output_mask_cache = self.model._output_mask_cache - self._output_tensor_cache = self.model._output_tensor_cache - self._output_shape_cache = self.model._output_shape_cache - self._input_layers = self.model._input_layers - self._output_layers = self.model._output_layers - self._input_coordinates = self.model._input_coordinates - self._output_coordinates = self.model._output_coordinates - self._nodes_by_depth = self.model._nodes_by_depth - self._network_nodes = self.model._network_nodes - self.output_names = self.model.output_names - self.input_names = self.model.input_names - self._feed_input_names = self.model._feed_input_names - self._feed_inputs = self.model._feed_inputs - - # Make sure child model callbacks - # will call the parent Sequential model. - self.model.callback_model = self - - self.built = True - - @property - def uses_learning_phase(self): - if not self.built: - self.build() - return self.model.uses_learning_phase - - def _gather_list_attr(self, attr): - all_attrs = [] - for layer in self.layers: - all_attrs += getattr(layer, attr, []) - return all_attrs - - @property - def trainable(self): - return self._trainable - - @trainable.setter - def trainable(self, value): - if self.model: - self.model.trainable = value - self._trainable = value - - @property - def trainable_weights(self): - if not self.trainable: - return [] - return self._gather_list_attr('trainable_weights') - - @property - def non_trainable_weights(self): - weights = self._gather_list_attr('non_trainable_weights') - if not self.trainable: - trainable_weights = self._gather_list_attr('trainable_weights') - return trainable_weights + weights - return weights - - @property - def updates(self): - if not self.built: - self.build() - return self.model.updates - - @property - def state_updates(self): - if not self.built: - self.build() - return self.model.state_updates - - def get_updates_for(self, inputs): - if not self.built: - self.build() - return self.model.get_updates_for(inputs) - - @property - def losses(self): - if not self.built: - self.build() - return self.model.losses - - def get_losses_for(self, inputs): - if not self.built: - self.build() - return self.model.get_losses_for(inputs) - - @property - def regularizers(self): - if not self.built: - self.build() - return self.model.regularizers - - def get_weights(self): - """Retrieves the weights of the model. - - Returns: - A flat list of Numpy arrays - (one array per model weight). - """ - if not self.built: - self.build() - return self.model.get_weights() - - def set_weights(self, weights): - """Sets the weights of the model. - - Arguments: - weights: Should be a list - of Numpy arrays with shapes and types matching - the output of `model.get_weights()`. - """ - if not self.built: - self.build() - self.model.set_weights(weights) - - def load_weights(self, filepath, by_name=False): - if h5py is None: - raise ImportError('`load_weights` requires h5py.') - f = h5py.File(filepath, mode='r') - if 'layer_names' not in f.attrs and 'model_weights' in f: - f = f['model_weights'] - layers = self.layers - if by_name: - topology.load_weights_from_hdf5_group_by_name(f, layers) - else: - topology.load_weights_from_hdf5_group(f, layers) - if hasattr(f, 'close'): - f.close() - - def save_weights(self, filepath, overwrite=True): - if h5py is None: - raise ImportError('`save_weights` requires h5py.') - # If file exists and should not be overwritten: - if not overwrite and os.path.isfile(filepath): - proceed = ask_to_proceed_with_overwrite(filepath) - if not proceed: - return - layers = self.layers - f = h5py.File(filepath, 'w') - topology.save_weights_to_hdf5_group(f, layers) - f.flush() - f.close() - - def compile(self, - optimizer, - loss, - metrics=None, - sample_weight_mode=None, - weighted_metrics=None, - target_tensors=None, - **kwargs): - """Configures the model for training. - - Arguments: - optimizer: String (name of optimizer) or optimizer object. - See [optimizers](/optimizers). - loss: String (name of objective function) or objective function. - See [losses](/losses). - If the model has multiple outputs, you can use a different loss - on each output by passing a dictionary or a list of losses. - The loss value that will be minimized by the model - will then be the sum of all individual losses. - metrics: List of metrics to be evaluated by the model - during training and testing. - Typically you will use `metrics=['accuracy']`. - To specify different metrics for different outputs of a - multi-output model, you could also pass a dictionary, - such as `metrics={'output_a': 'accuracy'}`. - sample_weight_mode: If you need to do timestep-wise - sample weighting (2D weights), set this to `"temporal"`. - `None` defaults to sample-wise weights (1D). - If the model has multiple outputs, you can use a different - `sample_weight_mode` on each output by passing a - dictionary or a list of modes. - weighted_metrics: list of metrics to be evaluated and weighted - by `sample_weight` or `class_weight` during training and testing. - target_tensors: By default, Keras will create a placeholder for the - model's target, which will be fed with the target data during - training. If instead you would like to use your own - target tensor (in turn, Keras will not expect external - Numpy data for these targets at training time), you - can specify them via the `target_tensors` argument. - It should be a single tensor - (for a single-output `Sequential` model). - **kwargs: These arguments are passed into `tf.Session.run`. - - Example: - ```python - model = Sequential() - model.add(Dense(32, input_shape=(500,))) - model.add(Dense(10, activation='softmax')) - model.compile(optimizer='rmsprop', - loss='categorical_crossentropy', - metrics=['accuracy']) - ``` - """ - # create the underlying model - self.build() - # call compile method of Model class - self.model.compile( - optimizer, - loss, - metrics=metrics, - sample_weight_mode=sample_weight_mode, - weighted_metrics=weighted_metrics, - target_tensors=target_tensors, - **kwargs) - self.optimizer = self.model.optimizer - self.loss = self.model.loss - self.metrics = self.model.metrics - self.loss_weights = self.model.loss_weights - self.sample_weight_mode = self.model.sample_weight_mode - self.weighted_metrics = self.model.weighted_metrics - self.targets = self.model.targets - self.metrics_tensors = self.model.metrics_tensors - self.metrics_names = self.model.metrics_names - self.sample_weights = self.model.sample_weights - self.total_loss = self.model.total_loss - - def fit(self, - x=None, - y=None, - batch_size=None, - epochs=1, - verbose=1, - callbacks=None, - validation_split=0., - validation_data=None, - shuffle=True, - class_weight=None, - sample_weight=None, - initial_epoch=0, - steps_per_epoch=None, - validation_steps=None, - **kwargs): - """Trains the model for a fixed number of epochs. - - Arguments: - x: Numpy array of training data. - If the input layer in the model is named, you can also pass a - dictionary mapping the input name to a Numpy array. - `x` can be `None` (default) if feeding from - TensorFlow data tensors. - y: Numpy array of target (label) data. - If the output layer in the model is named, you can also pass a - dictionary mapping the output name to a Numpy array. - `y` can be `None` (default) if feeding from - TensorFlow data tensors. - batch_size: Integer or `None`. - Number of samples per gradient update. - If unspecified, it will default to 32. - epochs: Integer. Number of epochs to train the model. - An epoch is an iteration over the entire `x` and `y` - data provided. - Note that in conjunction with `initial_epoch`, - `epochs` is to be understood as "final epoch". - The model is not trained for a number of iterations - given by `epochs`, but merely until the epoch - of index `epochs` is reached. - verbose: 0, 1, or 2. Verbosity mode. - 0 = silent, 1 = progress bar, 2 = one line per epoch. - callbacks: List of `keras.callbacks.Callback` instances. - List of callbacks to apply during training. - See [callbacks](/callbacks). - validation_split: Float between 0 and 1: - Fraction of the training data to be used as validation data. - The model will set apart this fraction of the training data, - will not train on it, and will evaluate - the loss and any model metrics - on this data at the end of each epoch. - The validation data is selected from the last samples - in the `x` and `y` data provided, before shuffling. - validation_data: tuple `(x_val, y_val)` or tuple - `(x_val, y_val, val_sample_weights)` on which to evaluate - the loss and any model metrics at the end of each epoch. - The model will not be trained on this data. - This will override `validation_split`. - shuffle: Boolean (whether to shuffle the training data - before each epoch) or str (for 'batch'). - 'batch' is a special option for dealing with the - limitations of HDF5 data; it shuffles in batch-sized chunks. - Has no effect when `steps_per_epoch` is not `None`. - class_weight: Optional dictionary mapping class indices (integers) - to a weight (float) value, used for weighting the loss function - (during training only). - This can be useful to tell the model to - "pay more attention" to samples from - an under-represented class. - sample_weight: Optional Numpy array of weights for - the training samples, used for weighting the loss function - (during training only). You can either pass a flat (1D) - Numpy array with the same length as the input samples - (1:1 mapping between weights and samples), - or in the case of temporal data, - you can pass a 2D array with shape - `(samples, sequence_length)`, - to apply a different weight to every timestep of every sample. - In this case you should make sure to specify - `sample_weight_mode="temporal"` in `compile()`. - initial_epoch: Epoch at which to start training - (useful for resuming a previous training run). - steps_per_epoch: Total number of steps (batches of samples) - before declaring one epoch finished and starting the - next epoch. When training with input tensors such as - TensorFlow data tensors, the default `None` is equal to - the number of unique samples in your dataset divided by - the batch size, or 1 if that cannot be determined. - validation_steps: Only relevant if `steps_per_epoch` - is specified. Total number of steps (batches of samples) - to validate before stopping. - **kwargs: Used for backwards compatibility support. - - Returns: - A `History` object. Its `History.history` attribute is - a record of training loss values and metrics values - at successive epochs, as well as validation loss values - and validation metrics values (if applicable). - - Raises: - RuntimeError: If the model was never compiled. - ValueError: In case of mismatch between the provided input data - and what the model expects. - """ - if not self.built: - raise RuntimeError('The model needs to be compiled before being used.') - return self.model.fit( - x, - y, - batch_size=batch_size, - epochs=epochs, - verbose=verbose, - callbacks=callbacks, - validation_split=validation_split, - validation_data=validation_data, - shuffle=shuffle, - class_weight=class_weight, - sample_weight=sample_weight, - initial_epoch=initial_epoch, - steps_per_epoch=steps_per_epoch, - validation_steps=validation_steps) - - def evaluate(self, x, y, batch_size=32, verbose=1, sample_weight=None): - """Computes the loss on some input data, batch by batch. - - Arguments: - x: input data, as a Numpy array or list of Numpy arrays - (if the model has multiple inputs). - y: labels, as a Numpy array. - batch_size: integer. Number of samples per gradient update. - verbose: verbosity mode, 0 or 1. - sample_weight: sample weights, as a Numpy array. - - Returns: - Scalar test loss (if the model has no metrics) - or list of scalars (if the model computes other metrics). - The attribute `model.metrics_names` will give you - the display labels for the scalar outputs. - - Raises: - RuntimeError: if the model was never compiled. - """ - if not self.built: - raise RuntimeError('The model needs to be compiled before being used.') - return self.model.evaluate( - x, - y, - batch_size=batch_size, - verbose=verbose, - sample_weight=sample_weight) - - def predict(self, x, batch_size=32, verbose=0): - """Generates output predictions for the input samples. - - The input samples are processed batch by batch. - - Arguments: - x: the input data, as a Numpy array. - batch_size: integer. - verbose: verbosity mode, 0 or 1. - - Returns: - A Numpy array of predictions. - """ - if not self.built: - self.build() - return self.model.predict(x, batch_size=batch_size, verbose=verbose) - - def predict_on_batch(self, x): - """Returns predictions for a single batch of samples. - - Arguments: - x: input data, as a Numpy array or list of Numpy arrays - (if the model has multiple inputs). - - Returns: - A Numpy array of predictions. - """ - if not self.built: - self.build() - return self.model.predict_on_batch(x) - - def train_on_batch(self, x, y, class_weight=None, sample_weight=None): - """Single gradient update over one batch of samples. - - Arguments: - x: input data, as a Numpy array or list of Numpy arrays - (if the model has multiple inputs). - y: labels, as a Numpy array. - class_weight: dictionary mapping classes to a weight value, - used for scaling the loss function (during training only). - sample_weight: sample weights, as a Numpy array. - - Returns: - Scalar training loss (if the model has no metrics) - or list of scalars (if the model computes other metrics). - The attribute `model.metrics_names` will give you - the display labels for the scalar outputs. - - Raises: - RuntimeError: if the model was never compiled. - """ - if not self.built: - raise RuntimeError('The model needs to be compiled before being used.') - return self.model.train_on_batch( - x, y, sample_weight=sample_weight, class_weight=class_weight) - - def test_on_batch(self, x, y, sample_weight=None): - """Evaluates the model over a single batch of samples. - - Arguments: - x: input data, as a Numpy array or list of Numpy arrays - (if the model has multiple inputs). - y: labels, as a Numpy array. - sample_weight: sample weights, as a Numpy array. - - Returns: - Scalar test loss (if the model has no metrics) - or list of scalars (if the model computes other metrics). - The attribute `model.metrics_names` will give you - the display labels for the scalar outputs. - - Raises: - RuntimeError: if the model was never compiled. - """ - if not self.built: - raise RuntimeError('The model needs to be compiled before being used.') - return self.model.test_on_batch(x, y, sample_weight=sample_weight) - - def predict_proba(self, x, batch_size=32, verbose=0): - """Generates class probability predictions for the input samples. - - The input samples are processed batch by batch. - - Arguments: - x: input data, as a Numpy array or list of Numpy arrays - (if the model has multiple inputs). - batch_size: integer. - verbose: verbosity mode, 0 or 1. - - Returns: - A Numpy array of probability predictions. - """ - preds = self.predict(x, batch_size, verbose) - if preds.min() < 0. or preds.max() > 1.: - logging.warning('Network returning invalid probability values. ' - 'The last layer might not normalize predictions ' - 'into probabilities ' - '(like softmax or sigmoid would).') - return preds - - def predict_classes(self, x, batch_size=32, verbose=0): - """Generate class predictions for the input samples. - - The input samples are processed batch by batch. - - Arguments: - x: input data, as a Numpy array or list of Numpy arrays - (if the model has multiple inputs). - batch_size: integer. - verbose: verbosity mode, 0 or 1. - - Returns: - A numpy array of class predictions. - """ - proba = self.predict(x, batch_size=batch_size, verbose=verbose) - if proba.shape[-1] > 1: - return proba.argmax(axis=-1) - else: - return (proba > 0.5).astype('int32') - - def fit_generator(self, - generator, - steps_per_epoch=None, - epochs=1, - verbose=1, - callbacks=None, - validation_data=None, - validation_steps=None, - class_weight=None, - max_queue_size=10, - workers=1, - use_multiprocessing=False, - shuffle=True, - initial_epoch=0, - **kwargs): - """Fits the model on data generated batch-by-batch by a Python generator. - - The generator is run in parallel to the model, for efficiency. - For instance, this allows you to do real-time data augmentation - on images on CPU in parallel to training your model on GPU. - - Arguments: - generator: A generator. - The output of the generator must be either - - a tuple (inputs, targets) - - a tuple (inputs, targets, sample_weights). - All arrays should contain the same number of samples. - The generator is expected to loop over its data - indefinitely. An epoch finishes when `steps_per_epoch` - batches have been seen by the model. - steps_per_epoch: Total number of steps (batches of samples) - to yield from `generator` before declaring one epoch - finished and starting the next epoch. It should typically - be equal to the number of samples of your dataset - divided by the batch size. - Optional for `Sequence`: if unspecified, will use - the `len(generator)` as a number of steps. - epochs: Integer, total number of iterations on the data. - Note that in conjunction with initial_epoch, the parameter - epochs is to be understood as "final epoch". The model is - not trained for n steps given by epochs, but until the - epoch epochs is reached. - verbose: Verbosity mode, 0, 1, or 2. - callbacks: List of callbacks to be called during training. - validation_data: This can be either - - A generator for the validation data - - A tuple (inputs, targets) - - A tuple (inputs, targets, sample_weights). - validation_steps: Only relevant if `validation_data` - is a generator. - Number of steps to yield from validation generator - at the end of every epoch. It should typically - be equal to the number of samples of your - validation dataset divided by the batch size. - Optional for `Sequence`: if unspecified, will use - the `len(validation_data)` as a number of steps. - class_weight: Dictionary mapping class indices to a weight - for the class. - max_queue_size: Maximum size for the generator queue - workers: Maximum number of processes to spin up - use_multiprocessing: If True, use process based threading. - Note that because - this implementation relies on multiprocessing, - you should not pass - non picklable arguments to the generator - as they can't be passed - easily to children processes. - shuffle: Whether to shuffle the order of the batches at - the beginning of each epoch. Only used with instances - of `Sequence` (keras.utils.Sequence). - initial_epoch: Epoch at which to start training - (useful for resuming a previous training run) - **kwargs: support for legacy arguments. - - Returns: - A `History` object. - - Raises: - RuntimeError: if the model was never compiled. - ValueError: In case the generator yields - data in an invalid format. - - Example: - - ```python - def generate_arrays_from_file(path): - while 1: - f = open(path) - for line in f: - # create Numpy arrays of input data - # and labels, from each line in the file - x, y = process_line(line) - yield (x, y) - f.close() - - model.fit_generator(generate_arrays_from_file('/my_file.txt'), - steps_per_epoch=1000, epochs=10) - ``` - """ - # Legacy support - if 'max_q_size' in kwargs: - max_queue_size = kwargs.pop('max_q_size') - logging.warning('The argument `max_q_size` has been renamed ' - '`max_queue_size`. Update your method calls accordingly.') - if 'pickle_safe' in kwargs: - use_multiprocessing = kwargs.pop('pickle_safe') - logging.warning('The argument `pickle_safe` has been renamed ' - '`use_multiprocessing`. ' - 'Update your method calls accordingly.') - if kwargs: - raise ValueError('Unrecognized keyword arguments: ' + str(kwargs)) - - if not self.built: - raise RuntimeError('The model needs to be compiled before being used.') - return self.model.fit_generator( - generator, - steps_per_epoch, - epochs, - verbose=verbose, - callbacks=callbacks, - validation_data=validation_data, - validation_steps=validation_steps, - class_weight=class_weight, - max_queue_size=max_queue_size, - workers=workers, - use_multiprocessing=use_multiprocessing, - shuffle=shuffle, - initial_epoch=initial_epoch) - - def evaluate_generator(self, - generator, - steps=None, - max_queue_size=10, - workers=1, - use_multiprocessing=False, - **kwargs): - """Evaluates the model on a data generator. - - The generator should return the same kind of data - as accepted by `test_on_batch`. - - Arguments: - generator: Generator yielding tuples (inputs, targets) - or (inputs, targets, sample_weights) - steps: Total number of steps (batches of samples) - to yield from `generator` before stopping. - Optional for `Sequence`: if unspecified, will use - the `len(generator)` as a number of steps. - max_queue_size: maximum size for the generator queue - workers: maximum number of processes to spin up - use_multiprocessing: if True, use process based threading. - Note that because this implementation - relies on multiprocessing, you should not pass - non picklable arguments to the generator - as they can't be passed easily to children processes. - **kwargs: support for legacy arguments. - - Returns: - Scalar test loss (if the model has no metrics) - or list of scalars (if the model computes other metrics). - The attribute `model.metrics_names` will give you - the display labels for the scalar outputs. - - Raises: - RuntimeError: if the model was never compiled. - ValueError: In case the generator yields - data in an invalid format. - """ - # Legacy support - if 'max_q_size' in kwargs: - max_queue_size = kwargs.pop('max_q_size') - logging.warning('The argument `max_q_size` has been renamed ' - '`max_queue_size`. Update your method calls accordingly.') - if 'pickle_safe' in kwargs: - use_multiprocessing = kwargs.pop('pickle_safe') - logging.warning('The argument `pickle_safe` has been renamed ' - '`use_multiprocessing`. ' - 'Update your method calls accordingly.') - if kwargs: - raise ValueError('Unrecognized keyword arguments: ' + str(kwargs)) - - if not self.built: - raise RuntimeError('The model needs to be compiled before being used.') - return self.model.evaluate_generator( - generator, - steps, - max_queue_size=max_queue_size, - workers=workers, - use_multiprocessing=use_multiprocessing) - - def predict_generator(self, - generator, - steps=None, - max_queue_size=10, - workers=1, - use_multiprocessing=False, - verbose=0, - **kwargs): - """Generates predictions for the input samples from a data generator. - - The generator should return the same kind of data as accepted by - `predict_on_batch`. - - Arguments: - generator: generator yielding batches of input samples. - steps: Total number of steps (batches of samples) - to yield from `generator` before stopping. - Optional for `Sequence`: if unspecified, will use - the `len(generator)` as a number of steps. - max_queue_size: maximum size for the generator queue - workers: maximum number of processes to spin up - use_multiprocessing: if True, use process based threading. - Note that because this implementation - relies on multiprocessing, you should not pass - non picklable arguments to the generator - as they can't be passed easily to children processes. - verbose: verbosity mode, 0 or 1. - **kwargs: support for legacy arguments. - - Returns: - A Numpy array of predictions. - - Raises: - ValueError: In case the generator yields - data in an invalid format. - """ - # Legacy support - if 'max_q_size' in kwargs: - max_queue_size = kwargs.pop('max_q_size') - logging.warning('The argument `max_q_size` has been renamed ' - '`max_queue_size`. Update your method calls accordingly.') - if 'pickle_safe' in kwargs: - use_multiprocessing = kwargs.pop('pickle_safe') - logging.warning('The argument `pickle_safe` has been renamed ' - '`use_multiprocessing`. ' - 'Update your method calls accordingly.') - if kwargs: - raise ValueError('Unrecognized keyword arguments: ' + str(kwargs)) - - if not self.built: - self.build() - return self.model.predict_generator( - generator, - steps, - max_queue_size=max_queue_size, - workers=workers, - use_multiprocessing=use_multiprocessing, - verbose=verbose) - def get_config(self): - config = [] - for layer in self.layers: - config.append({ - 'class_name': layer.__class__.__name__, - 'config': layer.get_config() - }) - return copy.deepcopy(config) - @classmethod - def from_config(cls, config, custom_objects=None): - model = cls() - for conf in config: - layer = layer_module.deserialize(conf, custom_objects=custom_objects) - model.add(layer) - return model +# API entries importable from `keras.models`: +Model = training.Model # pylint: disable=invalid-name +Sequential = sequential.Sequential # pylint: disable=invalid-name +save_model = saving.save_model +load_model = saving.load_model +model_from_config = saving.model_from_config +model_from_yaml = saving.model_from_yaml +model_from_json = saving.model_from_json def _clone_functional_model(model, input_tensors=None): @@ -1386,7 +90,7 @@ def _clone_functional_model(model, input_tensors=None): else: # Make sure that all input tensors come from a Keras layer. # If tensor comes from an input layer: cache the input layer. - input_tensors = topology._to_list(input_tensors) + input_tensors = generic_utils.to_list(input_tensors) input_tensors_ = [] for i, x in enumerate(input_tensors): if not K.is_keras_tensor(x): @@ -1423,7 +127,7 @@ def _clone_functional_model(model, input_tensors=None): # Reuse previously cloned layer. layer = layer_map[layer] # Don't call InputLayer multiple times. - if isinstance(layer, topology.InputLayer): + if isinstance(layer, InputLayer): continue # Gather inputs to call the new layer. @@ -1448,8 +152,9 @@ def _clone_functional_model(model, input_tensors=None): if has_arg(layer.call, 'mask'): if 'mask' not in kwargs: kwargs['mask'] = computed_mask - output_tensors = topology._to_list(layer(computed_tensor, **kwargs)) - output_masks = topology._to_list( + output_tensors = generic_utils.to_list(layer(computed_tensor, + **kwargs)) + output_masks = generic_utils.to_list( layer.compute_mask(computed_tensor, computed_mask)) computed_tensors = [computed_tensor] computed_masks = [computed_mask] @@ -1459,8 +164,9 @@ def _clone_functional_model(model, input_tensors=None): if has_arg(layer.call, 'mask'): if 'mask' not in kwargs: kwargs['mask'] = computed_masks - output_tensors = topology._to_list(layer(computed_tensors, **kwargs)) - output_masks = topology._to_list( + output_tensors = generic_utils.to_list(layer(computed_tensors, + **kwargs)) + output_masks = generic_utils.to_list( layer.compute_mask(computed_tensors, computed_masks)) # Update tensor_map. for x, y, mask in zip(reference_output_tensors, output_tensors, @@ -1510,14 +216,14 @@ def _clone_sequential_model(model, input_tensors=None): if input_tensors is None: return Sequential(layers=layers, name=model.name) else: - if len(topology._to_list(input_tensors)) != 1: + if len(generic_utils.to_list(input_tensors)) != 1: raise ValueError('To clone a `Sequential` model, we expect ' ' at most one tensor ' 'as part of `input_tensors`.') - x = topology._to_list(input_tensors)[0] + x = generic_utils.to_list(input_tensors)[0] if K.is_keras_tensor(x): origin_layer = x._keras_history[0] - if isinstance(origin_layer, topology.InputLayer): + if isinstance(origin_layer, InputLayer): return Sequential(layers=[origin_layer] + layers, name=model.name) else: raise ValueError('Cannot clone a `Sequential` model on top ' diff --git a/tensorflow/python/keras/_impl/keras/models_test.py b/tensorflow/python/keras/_impl/keras/models_test.py index 04017e4b28b27e52f88a7746fc44510c29edffce..5978ddd987c63b9d87a31be6837172f08512ef73 100644 --- a/tensorflow/python/keras/_impl/keras/models_test.py +++ b/tensorflow/python/keras/_impl/keras/models_test.py @@ -12,362 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for training routines.""" +"""Tests for `models.py` (model cloning, mainly).""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os -import shutil -import tempfile - import numpy as np from tensorflow.python.keras._impl import keras from tensorflow.python.platform import test -from tensorflow.python.training import training as training_module - -try: - import h5py # pylint:disable=g-import-not-at-top -except ImportError: - h5py = None - - -class TestModelSaving(test.TestCase): - - def test_sequential_model_saving(self): - if h5py is None: - return # Skip test if models cannot be saved. - - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.RepeatVector(3)) - model.add(keras.layers.TimeDistributed(keras.layers.Dense(3))) - model.compile(loss=keras.losses.MSE, - optimizer=keras.optimizers.RMSprop(lr=0.0001), - metrics=[keras.metrics.categorical_accuracy], - sample_weight_mode='temporal') - x = np.random.random((1, 3)) - y = np.random.random((1, 3, 3)) - model.train_on_batch(x, y) - - out = model.predict(x) - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - - new_model = keras.models.load_model(fname) - os.close(fd) - os.remove(fname) - - out2 = new_model.predict(x) - self.assertAllClose(out, out2, atol=1e-05) - - # test that new updates are the same with both models - x = np.random.random((1, 3)) - y = np.random.random((1, 3, 3)) - model.train_on_batch(x, y) - new_model.train_on_batch(x, y) - out = model.predict(x) - out2 = new_model.predict(x) - self.assertAllClose(out, out2, atol=1e-05) - - def test_sequential_model_saving_2(self): - if h5py is None: - return # Skip test if models cannot be saved. - - with self.test_session(): - # test with custom optimizer, loss - - class CustomOp(keras.optimizers.RMSprop): - pass - - def custom_loss(y_true, y_pred): - return keras.losses.mse(y_true, y_pred) - - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.Dense(3)) - model.compile(loss=custom_loss, optimizer=CustomOp(), metrics=['acc']) - - x = np.random.random((1, 3)) - y = np.random.random((1, 3)) - model.train_on_batch(x, y) - - out = model.predict(x) - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - - model = keras.models.load_model( - fname, - custom_objects={'CustomOp': CustomOp, - 'custom_loss': custom_loss}) - os.close(fd) - os.remove(fname) - - out2 = model.predict(x) - self.assertAllClose(out, out2, atol=1e-05) - - def test_functional_model_saving(self): - if h5py is None: - return # Skip test if models cannot be saved. - - with self.test_session(): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - output = keras.layers.Dense(3)(x) - - model = keras.models.Model(inputs, output) - model.compile(loss=keras.losses.MSE, - optimizer=keras.optimizers.RMSprop(lr=0.0001), - metrics=[keras.metrics.categorical_accuracy]) - x = np.random.random((1, 3)) - y = np.random.random((1, 3)) - model.train_on_batch(x, y) - - out = model.predict(x) - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - - model = keras.models.load_model(fname) - os.close(fd) - os.remove(fname) - - out2 = model.predict(x) - self.assertAllClose(out, out2, atol=1e-05) - - def test_saving_without_compilation(self): - if h5py is None: - return # Skip test if models cannot be saved. - - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.Dense(3)) - model.compile(loss='mse', optimizer='sgd', metrics=['acc']) - - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - model = keras.models.load_model(fname) - os.close(fd) - os.remove(fname) - - def test_saving_with_tf_optimizer(self): - if h5py is None: - return # Skip test if models cannot be saved. - - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.Dense(3)) - model.compile(loss='mse', - optimizer=training_module.AdadeltaOptimizer(0.1), - metrics=['acc']) - - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - model = keras.models.load_model(fname) - os.close(fd) - os.remove(fname) - - def test_saving_right_after_compilation(self): - if h5py is None: - return # Skip test if models cannot be saved. - - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.Dense(3)) - model.compile(loss='mse', optimizer='sgd', metrics=['acc']) - model.model._make_train_function() - - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - model = keras.models.load_model(fname) - os.close(fd) - os.remove(fname) - - def test_saving_lambda_numpy_array_arguments(self): - if h5py is None: - return # Skip test if models cannot be saved. - - mean = np.random.random((4, 2, 3)) - std = np.abs(np.random.random((4, 2, 3))) + 1e-5 - inputs = keras.layers.Input(shape=(4, 2, 3)) - output = keras.layers.Lambda(lambda image, mu, std: (image - mu) / std, - arguments={'mu': mean, 'std': std})(inputs) - model = keras.models.Model(inputs, output) - model.compile(loss='mse', optimizer='sgd', metrics=['acc']) - - fd, fname = tempfile.mkstemp('.h5') - keras.models.save_model(model, fname) - - model = keras.models.load_model(fname) - os.close(fd) - os.remove(fname) - - self.assertAllClose(mean, model.layers[1].arguments['mu']) - self.assertAllClose(std, model.layers[1].arguments['std']) - - -class TestSequential(test.TestCase): - """Most Sequential model API tests are covered in `training_test.py`. - """ - - def test_basic_methods(self): - model = keras.models.Sequential() - model.add(keras.layers.Dense(1, input_dim=2)) - model.add(keras.layers.Dropout(0.3, name='dp')) - model.add(keras.layers.Dense(2, kernel_regularizer='l2', - kernel_constraint='max_norm')) - model.build() - self.assertEqual(model.state_updates, model.model.state_updates) - self.assertEqual(model.get_layer(name='dp').name, 'dp') - - def test_sequential_pop(self): - num_hidden = 5 - input_dim = 3 - batch_size = 5 - num_classes = 2 - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) - model.add(keras.layers.Dense(num_classes)) - model.compile(loss='mse', optimizer='sgd') - x = np.random.random((batch_size, input_dim)) - y = np.random.random((batch_size, num_classes)) - model.fit(x, y, epochs=1) - model.pop() - self.assertEqual(len(model.layers), 1) - self.assertEqual(model.output_shape, (None, num_hidden)) - model.compile(loss='mse', optimizer='sgd') - y = np.random.random((batch_size, num_hidden)) - model.fit(x, y, epochs=1) - - # Test popping single-layer model - model = keras.models.Sequential() - model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) - model.pop() - self.assertEqual(len(model.layers), 0) - self.assertEqual(len(model.outputs), 0) - - # Invalid use case - model = keras.models.Sequential() - with self.assertRaises(TypeError): - model.pop() - - def test_sequential_weight_loading(self): - if h5py is None: - return - - temp_dir = self.get_temp_dir() - self.addCleanup(shutil.rmtree, temp_dir) - h5_path = os.path.join(temp_dir, 'test.h5') - - num_hidden = 5 - input_dim = 3 - batch_size = 5 - num_classes = 2 - - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) - model.add(keras.layers.Dense(num_classes)) - - x = np.random.random((batch_size, input_dim)) - ref_y = model.predict(x) - - model.save_weights(h5_path) - - model = keras.models.Sequential() - model.add(keras.layers.Dense(num_hidden, input_dim=input_dim)) - model.add(keras.layers.Dense(num_classes)) - model.load_weights(h5_path) - y = model.predict(x) - - self.assertAllClose(y, ref_y) - - def test_invalid_use_cases(self): - with self.test_session(): - # Added objects must be layer instances - with self.assertRaises(TypeError): - model = keras.models.Sequential() - model.add(None) - - # Added layers must have an inputs shape - with self.assertRaises(ValueError): - model = keras.models.Sequential() - model.add(keras.layers.Dense(1)) - - # Added layers cannot have multiple outputs - class MyLayer(keras.layers.Layer): - - def call(self, inputs): - return [3 * inputs, 2 * inputs] - - def compute_output_shape(self, input_shape): - return [input_shape, input_shape] - - with self.assertRaises(ValueError): - model = keras.models.Sequential() - model.add(MyLayer(input_shape=(3,))) - with self.assertRaises(TypeError): - model = keras.models.Sequential() - model.add(keras.layers.Dense(1, input_dim=1)) - model.add(MyLayer()) - - # Building empty model - model = keras.models.Sequential() - with self.assertRaises(TypeError): - model.build() - - def test_nested_sequential_trainability(self): - input_dim = 20 - num_units = 10 - num_classes = 2 - - inner_model = keras.models.Sequential() - inner_model.add(keras.layers.Dense(num_units, input_shape=(input_dim,))) - - model = keras.models.Sequential() - model.add(inner_model) - model.add(keras.layers.Dense(num_classes)) - - self.assertEqual(len(model.trainable_weights), 4) - inner_model.trainable = False - self.assertEqual(len(model.trainable_weights), 2) - inner_model.trainable = True - self.assertEqual(len(model.trainable_weights), 4) - - def test_sequential_update_disabling(self): - val_a = np.random.random((10, 4)) - val_out = np.random.random((10, 4)) - - with self.test_session(): - model = keras.models.Sequential() - model.add(keras.layers.BatchNormalization(input_shape=(4,))) - - model.trainable = False - assert not model.updates - - model.compile('sgd', 'mse') - assert not model.updates - assert not model.model.updates - - x1 = model.predict(val_a) - model.train_on_batch(val_a, val_out) - x2 = model.predict(val_a) - self.assertAllClose(x1, x2, atol=1e-7) - - model.trainable = True - model.compile('sgd', 'mse') - assert model.updates - assert model.model.updates - - model.train_on_batch(val_a, val_out) - x2 = model.predict(val_a) - assert np.abs(np.sum(x1 - x2)) > 1e-5 class TestModelCloning(test.TestCase): diff --git a/tensorflow/python/keras/_impl/keras/optimizers.py b/tensorflow/python/keras/_impl/keras/optimizers.py index e47987aadc48e1f722558e32929ab81ad82bea0f..b715d722b98b9db3bdf0985da0954356a2facdfe 100644 --- a/tensorflow/python/keras/_impl/keras/optimizers.py +++ b/tensorflow/python/keras/_impl/keras/optimizers.py @@ -32,6 +32,7 @@ from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_ke from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer as tf_optimizer_module +from tensorflow.python.util.tf_export import tf_export def clip_norm(g, c, n): @@ -65,6 +66,7 @@ def clip_norm(g, c, n): return g +@tf_export('keras.optimizers.Optimizer') class Optimizer(object): """Abstract optimizer base class. @@ -93,7 +95,26 @@ class Optimizer(object): raise NotImplementedError def get_gradients(self, loss, params): + """Returns gradients of `loss` with respect to `params`. + + Arguments: + loss: Loss tensor. + params: List of variables. + + Returns: + List of gradient tensors. + + Raises: + ValueError: In case any gradient cannot be computed (e.g. if gradient + function not implemented). + """ grads = K.gradients(loss, params) + if None in grads: + raise ValueError('An operation has `None` for gradient. ' + 'Please make sure that all of your ops have a ' + 'gradient defined (i.e. are differentiable). ' + 'Common ops without gradient: ' + 'K.argmax, K.round, K.eval.') if hasattr(self, 'clipnorm') and self.clipnorm > 0: norm = K.sqrt(sum([K.sum(K.square(g)) for g in grads])) grads = [clip_norm(g, self.clipnorm, norm) for g in grads] @@ -118,6 +139,11 @@ class Optimizer(object): ValueError: in case of incompatible weight shapes. """ params = self.weights + if len(params) != len(weights): + raise ValueError( + 'Length of the specified weight list (' + str(len(weights)) + + ') does not match the number of weights ' + 'of the optimizer (' + str(len(params)) + ')') weight_value_tuples = [] param_values = K.batch_get_value(params) for pv, p, w in zip(param_values, params, weights): @@ -149,6 +175,7 @@ class Optimizer(object): return cls(**config) +@tf_export('keras.optimizers.SGD') class SGD(Optimizer): """Stochastic gradient descent optimizer. @@ -179,8 +206,9 @@ class SGD(Optimizer): lr = self.lr if self.initial_decay > 0: - lr *= (1. / - (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) + lr = lr * (1. / # pylint: disable=g-no-augmented-assignment + (1. + self.decay * K.cast(self.iterations, + K.dtype(self.decay)))) # momentum shapes = [K.int_shape(p) for p in params] moments = [K.zeros(shape) for shape in shapes] @@ -212,6 +240,7 @@ class SGD(Optimizer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.optimizers.RMSprop') class RMSprop(Optimizer): """RMSProp optimizer. @@ -250,8 +279,9 @@ class RMSprop(Optimizer): lr = self.lr if self.initial_decay > 0: - lr *= (1. / - (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) + lr = lr * (1. / # pylint: disable=g-no-augmented-assignment + (1. + self.decay * K.cast(self.iterations, + K.dtype(self.decay)))) for p, g, a in zip(params, grads, accumulators): # update accumulator @@ -277,6 +307,7 @@ class RMSprop(Optimizer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.optimizers.Adagrad') class Adagrad(Optimizer): """Adagrad optimizer. @@ -310,8 +341,9 @@ class Adagrad(Optimizer): lr = self.lr if self.initial_decay > 0: - lr *= (1. / - (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) + lr = lr * (1. / # pylint: disable=g-no-augmented-assignment + (1. + self.decay * K.cast(self.iterations, + K.dtype(self.decay)))) for p, g, a in zip(params, grads, accumulators): new_a = a + K.square(g) # update accumulator @@ -335,6 +367,7 @@ class Adagrad(Optimizer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.optimizers.Adadelta') class Adadelta(Optimizer): """Adadelta optimizer. @@ -372,8 +405,9 @@ class Adadelta(Optimizer): lr = self.lr if self.initial_decay > 0: - lr *= (1. / - (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) + lr = lr * (1. / # pylint: disable=g-no-augmented-assignment + (1. + self.decay * K.cast(self.iterations, + K.dtype(self.decay)))) for p, g, a, d_a in zip(params, grads, accumulators, delta_accumulators): # update accumulator @@ -406,6 +440,7 @@ class Adadelta(Optimizer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.optimizers.Adam') class Adam(Optimizer): """Adam optimizer. @@ -450,8 +485,9 @@ class Adam(Optimizer): lr = self.lr if self.initial_decay > 0: - lr *= (1. / - (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) + lr = lr * (1. / # pylint: disable=g-no-augmented-assignment + (1. + self.decay * K.cast(self.iterations, + K.dtype(self.decay)))) t = K.cast(self.iterations, K.floatx()) + 1 lr_t = lr * ( @@ -499,6 +535,7 @@ class Adam(Optimizer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.optimizers.Adamax') class Adamax(Optimizer): """Adamax optimizer from Adam paper's Section 7. @@ -538,8 +575,9 @@ class Adamax(Optimizer): lr = self.lr if self.initial_decay > 0: - lr *= (1. / - (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) + lr = lr * (1. / # pylint: disable=g-no-augmented-assignment + (1. + self.decay * K.cast(self.iterations, + K.dtype(self.decay)))) t = K.cast(self.iterations, K.floatx()) + 1 lr_t = lr / (1. - K.pow(self.beta_1, t)) @@ -580,6 +618,7 @@ class Adamax(Optimizer): return dict(list(base_config.items()) + list(config.items())) +@tf_export('keras.optimizers.Nadam') class Nadam(Optimizer): """Nesterov Adam optimizer. @@ -682,9 +721,17 @@ class TFOptimizer(Optimizer): with K.name_scope(self.__class__.__name__): self.iterations = K.variable(0, dtype='int64', name='iterations') + def apply_gradients(self, grads): + self.optimizer.apply_gradients(grads) + + def get_grads(self, loss, params): + return self.optimizer.compute_gradients(loss, params) + def get_updates(self, loss, params): - grads = self.optimizer.compute_gradients(loss, params) self.updates = [K.update_add(self.iterations, 1)] + if not params: + return self.updates + grads = self.optimizer.compute_gradients(loss, params) opt_update = self.optimizer.apply_gradients( grads, global_step=self.iterations) self.updates.append(opt_update) @@ -712,10 +759,12 @@ adamax = Adamax nadam = Nadam +@tf_export('keras.optimizers.serialize') def serialize(optimizer): return serialize_keras_object(optimizer) +@tf_export('keras.optimizers.deserialize') def deserialize(config, custom_objects=None): """Inverse of the `serialize` function. @@ -749,6 +798,7 @@ def deserialize(config, custom_objects=None): printable_module_name='optimizer') +@tf_export('keras.optimizers.get') def get(identifier): """Retrieves a Keras Optimizer instance. diff --git a/tensorflow/python/keras/_impl/keras/preprocessing/image.py b/tensorflow/python/keras/_impl/keras/preprocessing/image.py index db1fdd4e6b0fd39536f4f69ab396c7e5552710ea..d12f10863921ee7d635930f34e8bc701c89864e8 100644 --- a/tensorflow/python/keras/_impl/keras/preprocessing/image.py +++ b/tensorflow/python/keras/_impl/keras/preprocessing/image.py @@ -32,6 +32,7 @@ import numpy as np from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.data_utils import Sequence from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export try: from scipy import linalg @@ -62,6 +63,7 @@ if pil_image is not None: _PIL_INTERPOLATION_METHODS['lanczos'] = pil_image.LANCZOS +@tf_export('keras.preprocessing.image.random_rotation') def random_rotation(x, rg, row_axis=1, @@ -96,6 +98,7 @@ def random_rotation(x, return x +@tf_export('keras.preprocessing.image.random_shift') def random_shift(x, wrg, hrg, @@ -132,6 +135,7 @@ def random_shift(x, return x +@tf_export('keras.preprocessing.image.random_shear') def random_shear(x, intensity, row_axis=1, @@ -166,6 +170,7 @@ def random_shear(x, return x +@tf_export('keras.preprocessing.image.random_zoom') def random_zoom(x, zoom_range, row_axis=1, @@ -209,6 +214,7 @@ def random_zoom(x, return x +@tf_export('keras.preprocessing.image.random_channel_shift') def random_channel_shift(x, intensity, channel_axis=0): x = np.rollaxis(x, channel_axis, 0) min_x, max_x = np.min(x), np.max(x) @@ -230,6 +236,7 @@ def transform_matrix_offset_center(matrix, x, y): return transform_matrix +@tf_export('keras.preprocessing.image.apply_transform') def apply_transform(x, transform_matrix, channel_axis=0, @@ -267,6 +274,7 @@ def apply_transform(x, return x +@tf_export('keras.preprocessing.image.flip_axis') def flip_axis(x, axis): x = np.asarray(x).swapaxes(axis, 0) x = x[::-1, ...] @@ -274,6 +282,7 @@ def flip_axis(x, axis): return x +@tf_export('keras.preprocessing.image.array_to_img') def array_to_img(x, data_format=None, scale=True): """Converts a 3D Numpy array to a PIL Image instance. @@ -324,6 +333,7 @@ def array_to_img(x, data_format=None, scale=True): raise ValueError('Unsupported channel number: ', x.shape[2]) +@tf_export('keras.preprocessing.image.img_to_array') def img_to_array(img, data_format=None): """Converts a PIL Image instance to a Numpy array. @@ -358,6 +368,7 @@ def img_to_array(img, data_format=None): return x +@tf_export('keras.preprocessing.image.load_img') def load_img(path, grayscale=False, target_size=None, interpolation='nearest'): """Loads an image into PIL format. @@ -411,6 +422,7 @@ def list_pictures(directory, ext='jpg|jpeg|bmp|png|ppm'): ] +@tf_export('keras.preprocessing.image.ImageDataGenerator') class ImageDataGenerator(object): """Generate minibatches of image data with real-time data augmentation. @@ -824,6 +836,7 @@ class ImageDataGenerator(object): vt.T / np.sqrt(s_expand**2 + self.zca_epsilon)).dot(vt) +@tf_export('keras.preprocessing.image.Iterator') class Iterator(Sequence): """Base class for image data iterators. @@ -913,6 +926,7 @@ class Iterator(Sequence): raise NotImplementedError +@tf_export('keras.preprocessing.image.NumpyArrayIterator') class NumpyArrayIterator(Iterator): """Iterator yielding data from a Numpy array. @@ -1094,6 +1108,7 @@ def _list_valid_filenames_in_directory(directory, white_list_formats, return classes, filenames +@tf_export('keras.preprocessing.image.DirectoryIterator') class DirectoryIterator(Iterator): """Iterator capable of reading images from a directory on disk. diff --git a/tensorflow/python/keras/_impl/keras/preprocessing/sequence.py b/tensorflow/python/keras/_impl/keras/preprocessing/sequence.py index 4d59250af03cd5e264fb8a36ba70311840bd68b0..a423d96d3d8578df347b7ee36fb53dfd335e0d65 100644 --- a/tensorflow/python/keras/_impl/keras/preprocessing/sequence.py +++ b/tensorflow/python/keras/_impl/keras/preprocessing/sequence.py @@ -22,8 +22,10 @@ import random import numpy as np from six.moves import range # pylint: disable=redefined-builtin +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.preprocessing.sequence.pad_sequences') def pad_sequences(sequences, maxlen=None, dtype='int32', @@ -104,6 +106,7 @@ def pad_sequences(sequences, return x +@tf_export('keras.preprocessing.sequence.make_sampling_table') def make_sampling_table(size, sampling_factor=1e-5): """Generates a word rank-based probabilistic sampling table. @@ -137,6 +140,7 @@ def make_sampling_table(size, sampling_factor=1e-5): return np.minimum(1., f / np.sqrt(f)) +@tf_export('keras.preprocessing.sequence.skipgrams') def skipgrams(sequence, vocabulary_size, window_size=4, diff --git a/tensorflow/python/keras/_impl/keras/preprocessing/text.py b/tensorflow/python/keras/_impl/keras/preprocessing/text.py index 8f7f25dc0a3e6fd234abb5fc97b3441ddcf16a4e..1e3828ccf1e3bf9c443691e1c1da5697bedb4653 100644 --- a/tensorflow/python/keras/_impl/keras/preprocessing/text.py +++ b/tensorflow/python/keras/_impl/keras/preprocessing/text.py @@ -28,6 +28,7 @@ from six.moves import range # pylint: disable=redefined-builtin from six.moves import zip # pylint: disable=redefined-builtin from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export if sys.version_info < (3,): @@ -36,6 +37,7 @@ else: maketrans = str.maketrans +@tf_export('keras.preprocessing.text.text_to_word_sequence') def text_to_word_sequence(text, filters='!"#$%&()*+,-./:;<=>?@[\\]^_`{|}~\t\n', lower=True, @@ -64,6 +66,7 @@ def text_to_word_sequence(text, return [i for i in seq if i] +@tf_export('keras.preprocessing.text.one_hot') def one_hot(text, n, filters='!"#$%&()*+,-./:;<=>?@[\\]^_`{|}~\t\n', @@ -129,6 +132,7 @@ def hashing_trick(text, return [(hash_function(w) % (n - 1) + 1) for w in seq] +@tf_export('keras.preprocessing.text.Tokenizer') class Tokenizer(object): """Text tokenization utility class. diff --git a/tensorflow/python/keras/_impl/keras/regularizers.py b/tensorflow/python/keras/_impl/keras/regularizers.py index c53ee8a1aeccbf862324d1e91235ce1a00adb457..2c30844647acdb78d1ca31d052ec7e5ecc6dcc2a 100644 --- a/tensorflow/python/keras/_impl/keras/regularizers.py +++ b/tensorflow/python/keras/_impl/keras/regularizers.py @@ -23,8 +23,10 @@ import six from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.generic_utils import deserialize_keras_object from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_keras_object +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.regularizers.Regularizer') class Regularizer(object): """Regularizer base class. """ @@ -37,6 +39,7 @@ class Regularizer(object): return cls(**config) +@tf_export('keras.regularizers.L1L2') class L1L2(Regularizer): """Regularizer for L1 and L2 regularization. @@ -64,22 +67,27 @@ class L1L2(Regularizer): # Aliases. +@tf_export('keras.regularizers.l1') def l1(l=0.01): return L1L2(l1=l) +@tf_export('keras.regularizers.l2') def l2(l=0.01): return L1L2(l2=l) +@tf_export('keras.regularizers.l1_l2') def l1_l2(l1=0.01, l2=0.01): # pylint: disable=redefined-outer-name return L1L2(l1=l1, l2=l2) +@tf_export('keras.regularizers.serialize') def serialize(regularizer): return serialize_keras_object(regularizer) +@tf_export('keras.regularizers.deserialize') def deserialize(config, custom_objects=None): return deserialize_keras_object( config, @@ -88,6 +96,7 @@ def deserialize(config, custom_objects=None): printable_module_name='regularizer') +@tf_export('keras.regularizers.get') def get(identifier): if identifier is None: return None diff --git a/tensorflow/python/keras/_impl/keras/testing_utils.py b/tensorflow/python/keras/_impl/keras/testing_utils.py index b889e311b37d48732641205a90ca83af34ea4489..60799ee1e038b4466351248bb5de7c8fc0de02a2 100644 --- a/tensorflow/python/keras/_impl/keras/testing_utils.py +++ b/tensorflow/python/keras/_impl/keras/testing_utils.py @@ -22,6 +22,7 @@ import numpy as np from tensorflow.python.framework import tensor_shape from tensorflow.python.keras._impl import keras +from tensorflow.python.training.rmsprop import RMSPropOptimizer from tensorflow.python.util import tf_inspect @@ -105,8 +106,14 @@ def layer_test(layer_cls, kwargs=None, input_shape=None, input_dtype=None, # test in functional API x = keras.layers.Input(shape=input_shape[1:], dtype=input_dtype) y = layer(x) - assert keras.backend.dtype(y) == expected_output_dtype - + if keras.backend.dtype(y) != expected_output_dtype: + raise AssertionError('When testing layer %s, for input %s, found output ' + 'dtype=%s but expected to find %s.\nFull kwargs: %s' % + (layer_cls.__name__, + x, + keras.backend.dtype(y), + expected_output_dtype, + kwargs)) # check shape inference model = keras.models.Model(x, y) expected_output_shape = tuple( @@ -117,7 +124,15 @@ def layer_test(layer_cls, kwargs=None, input_shape=None, input_dtype=None, for expected_dim, actual_dim in zip(expected_output_shape, actual_output_shape): if expected_dim is not None: - assert expected_dim == actual_dim + if expected_dim != actual_dim: + raise AssertionError( + 'When testing layer %s, for input %s, found output_shape=' + '%s but expected to find %s.\nFull kwargs: %s' % + (layer_cls.__name__, + x, + actual_output_shape, + expected_output_shape, + kwargs)) if expected_output is not None: np.testing.assert_allclose(actual_output, expected_output, rtol=1e-3) @@ -131,7 +146,7 @@ def layer_test(layer_cls, kwargs=None, input_shape=None, input_dtype=None, np.testing.assert_allclose(output, actual_output, rtol=1e-3) # test training mode (e.g. useful for dropout tests) - model.compile('rmsprop', 'mse') + model.compile(RMSPropOptimizer(0.01), 'mse') model.train_on_batch(input_data, actual_output) # test as first layer in Sequential API @@ -146,7 +161,15 @@ def layer_test(layer_cls, kwargs=None, input_shape=None, input_dtype=None, for expected_dim, actual_dim in zip(expected_output_shape, actual_output_shape): if expected_dim is not None: - assert expected_dim == actual_dim + if expected_dim != actual_dim: + raise AssertionError( + 'When testing layer %s, for input %s, found output_shape=' + '%s but expected to find %s.\nFull kwargs: %s' % + (layer_cls.__name__, + x, + actual_output_shape, + expected_output_shape, + kwargs)) if expected_output is not None: np.testing.assert_allclose(actual_output, expected_output, rtol=1e-3) @@ -159,9 +182,5 @@ def layer_test(layer_cls, kwargs=None, input_shape=None, input_dtype=None, output = recovered_model.predict(input_data) np.testing.assert_allclose(output, actual_output, rtol=1e-3) - # test training mode (e.g. useful for dropout tests) - model.compile('rmsprop', 'mse') - model.train_on_batch(input_data, actual_output) - # for further checks in the caller function return actual_output diff --git a/tensorflow/python/keras/_impl/keras/utils/__init__.py b/tensorflow/python/keras/_impl/keras/utils/__init__.py index 370ae0dd0f0d00059f1b0cc79459abe75c8ca494..0c9f19a0c8dcf3bf929e102b31679a03b27728f7 100644 --- a/tensorflow/python/keras/_impl/keras/utils/__init__.py +++ b/tensorflow/python/keras/_impl/keras/utils/__init__.py @@ -31,8 +31,8 @@ from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_ke from tensorflow.python.keras._impl.keras.utils.io_utils import HDF5Matrix from tensorflow.python.keras._impl.keras.utils.layer_utils import convert_all_kernels_in_model from tensorflow.python.keras._impl.keras.utils.layer_utils import print_summary +from tensorflow.python.keras._impl.keras.utils.multi_gpu_utils import multi_gpu_model from tensorflow.python.keras._impl.keras.utils.np_utils import normalize from tensorflow.python.keras._impl.keras.utils.np_utils import to_categorical -from tensorflow.python.keras._impl.keras.utils.training_utils import multi_gpu_model from tensorflow.python.keras._impl.keras.utils.vis_utils import plot_model diff --git a/tensorflow/python/keras/_impl/keras/utils/data_utils.py b/tensorflow/python/keras/_impl/keras/utils/data_utils.py index fcee9fbcc32d9356f5776b6f53cae6c125313d62..4c49544c6a63c4e5a0b79d31b074ad352c512bfa 100644 --- a/tensorflow/python/keras/_impl/keras/utils/data_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/data_utils.py @@ -40,6 +40,7 @@ from six.moves.urllib.error import URLError from six.moves.urllib.request import urlopen from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar +from tensorflow.python.util.tf_export import tf_export try: @@ -138,6 +139,7 @@ def _extract_archive(file_path, path='.', archive_format='auto'): return False +@tf_export('keras.utils.get_file') def get_file(fname, origin, untar=False, @@ -318,6 +320,7 @@ def validate_file(fpath, file_hash, algorithm='auto', chunk_size=65535): return False +@tf_export('keras.utils.Sequence') class Sequence(object): """Base object for fitting to a sequence of data, such as a dataset. @@ -390,6 +393,16 @@ class Sequence(object): """ pass + def __iter__(self): + """Creates an infinite generator that iterate over the Sequence. + + Yields: + Sequence items. + """ + while True: + for item in (self[i] for i in range(len(self))): + yield item + # Global variables to be shared across processes _SHARED_SEQUENCES = {} @@ -397,6 +410,11 @@ _SHARED_SEQUENCES = {} _SEQUENCE_COUNTER = None +def init_pool(seqs): + global _SHARED_SEQUENCES + _SHARED_SEQUENCES = seqs + + def get_index(uid, i): """Get the value from the Sequence `uid` at index `i`. @@ -414,6 +432,7 @@ def get_index(uid, i): return _SHARED_SEQUENCES[uid][i] +@tf_export('keras.utils.SequenceEnqueuer') class SequenceEnqueuer(object): """Base class to enqueue inputs. @@ -528,9 +547,11 @@ class OrderedEnqueuer(SequenceEnqueuer): (when full, workers could block on `put()`) """ if self.use_multiprocessing: - self.executor_fn = lambda: multiprocessing.Pool(workers) + self.executor_fn = lambda seqs: multiprocessing.Pool( # pylint: disable=g-long-lambda + workers, initializer=init_pool, initargs=(seqs,)) else: - self.executor_fn = lambda: ThreadPool(workers) + # We do not need the init since it's threads. + self.executor_fn = lambda _: ThreadPool(workers) self.workers = workers self.queue = queue.Queue(max_queue_size) self.stop_signal = threading.Event() @@ -553,7 +574,7 @@ class OrderedEnqueuer(SequenceEnqueuer): if self.shuffle: random.shuffle(sequence) - with closing(self.executor_fn()) as executor: + with closing(self.executor_fn(_SHARED_SEQUENCES)) as executor: for i in sequence: if self.stop_signal.is_set(): return @@ -613,6 +634,7 @@ class OrderedEnqueuer(SequenceEnqueuer): _SHARED_SEQUENCES[self.uid] = None +@tf_export('keras.utils.GeneratorEnqueuer') class GeneratorEnqueuer(SequenceEnqueuer): """Builds a queue out of a data generator. diff --git a/tensorflow/python/keras/_impl/keras/utils/generic_utils.py b/tensorflow/python/keras/_impl/keras/utils/generic_utils.py index adbe6c3288a3eabb858e78267577ddff6d798972..3bbe87f92d8f7eac27033344550ca65397eab986 100644 --- a/tensorflow/python/keras/_impl/keras/utils/generic_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/generic_utils.py @@ -291,55 +291,73 @@ class Progbar(object): Arguments: target: Total number of steps expected, None if unknown. + width: Progress bar width on screen. + verbose: Verbosity mode, 0 (silent), 1 (verbose), 2 (semi-verbose) + stateful_metrics: Iterable of string names of metrics that + should *not* be averaged over time. Metrics in this list + will be displayed as-is. All others will be averaged + by the progbar before display. interval: Minimum visual progress update interval (in seconds). """ - def __init__(self, target, width=30, verbose=1, interval=0.05): - self.width = width - if target is None: - target = -1 + def __init__(self, target, width=30, verbose=1, interval=0.05, + stateful_metrics=None): self.target = target - self.sum_values = {} - self.unique_values = [] - self.start = time.time() - self.last_update = 0 - self.interval = interval - self.total_width = 0 - self.seen_so_far = 0 + self.width = width self.verbose = verbose + self.interval = interval + if stateful_metrics: + self.stateful_metrics = set(stateful_metrics) + else: + self.stateful_metrics = set() + self._dynamic_display = ((hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()) or - 'ipykernel' in sys.modules) - - def update(self, current, values=None, force=False): + 'ipykernel' in sys.modules or + 'posix' in sys.modules) + self._total_width = 0 + self._seen_so_far = 0 + # We use a dict + list to avoid garbage collection + # issues found in OrderedDict + self._values = {} + self._values_order = [] + self._start = time.time() + self._last_update = 0 + + def update(self, current, values=None): """Updates the progress bar. Arguments: current: Index of current step. - values: List of tuples (name, value_for_last_step). - The progress bar will display averages for these values. - force: Whether to force visual progress update. + values: List of tuples: + `(name, value_for_last_step)`. + If `name` is in `stateful_metrics`, + `value_for_last_step` will be displayed as-is. + Else, an average of the metric over time will be displayed. """ values = values or [] for k, v in values: - if k not in self.sum_values: - self.sum_values[k] = [ - v * (current - self.seen_so_far), current - self.seen_so_far - ] - self.unique_values.append(k) + if k not in self._values_order: + self._values_order.append(k) + if k not in self.stateful_metrics: + if k not in self._values: + self._values[k] = [v * (current - self._seen_so_far), + current - self._seen_so_far] + else: + self._values[k][0] += v * (current - self._seen_so_far) + self._values[k][1] += (current - self._seen_so_far) else: - self.sum_values[k][0] += v * (current - self.seen_so_far) - self.sum_values[k][1] += (current - self.seen_so_far) - self.seen_so_far = current + self._values[k] = v + self._seen_so_far = current now = time.time() - info = ' - %.0fs' % (now - self.start) + info = ' - %.0fs' % (now - self._start) if self.verbose == 1: - if (not force and (now - self.last_update) < self.interval and - current < self.target): + if (now - self._last_update < self.interval and + self.target is not None and current < self.target): return - prev_total_width = self.total_width + prev_total_width = self._total_width if self._dynamic_display: sys.stdout.write('\b' * prev_total_width) sys.stdout.write('\r') @@ -360,22 +378,21 @@ class Progbar(object): bar += '=' bar += ('.' * (self.width - prog_width)) bar += ']' - sys.stdout.write(bar) - self.total_width = len(bar) else: bar = '%7d/Unknown' % current - self.total_width = len(bar) + self._total_width = len(bar) sys.stdout.write(bar) if current: - time_per_unit = (now - self.start) / current + time_per_unit = (now - self._start) / current else: time_per_unit = 0 if self.target is not None and current < self.target: eta = time_per_unit * (self.target - current) if eta > 3600: - eta_format = '%d:%02d:%02d' % (eta // 3600, (eta % 3600) // 60, + eta_format = '%d:%02d:%02d' % (eta // 3600, + (eta % 3600) // 60, eta % 60) elif eta > 60: eta_format = '%d:%02d' % (eta // 60, eta % 60) @@ -391,35 +408,32 @@ class Progbar(object): else: info += ' %.0fus/step' % (time_per_unit * 1e6) - for k in self.unique_values: + for k in self._values_order: info += ' - %s:' % k - if isinstance(self.sum_values[k], list): - avg = np.mean(self.sum_values[k][0] / max(1, self.sum_values[k][1])) + if isinstance(self._values[k], list): + avg = np.mean(self._values[k][0] / max(1, self._values[k][1])) if abs(avg) > 1e-3: info += ' %.4f' % avg else: info += ' %.4e' % avg else: - info += ' %s' % self.sum_values[k] + info += ' %s' % self._values[k] + + self._total_width += len(info) + if prev_total_width > self._total_width: + info += (' ' * (prev_total_width - self._total_width)) - self.total_width += len(info) - if prev_total_width > self.total_width: - info += (' ' * (prev_total_width - self.total_width)) if self.target is not None and current >= self.target: info += '\n' sys.stdout.write(info) sys.stdout.flush() - if current >= self.target: - sys.stdout.write('\n') - elif self.verbose == 2: if self.target is None or current >= self.target: - for k in self.unique_values: + for k in self._values_order: info += ' - %s:' % k - avg = np.mean( - self.sum_values[k][0] / max(1, self.sum_values[k][1])) + avg = np.mean(self._values[k][0] / max(1, self._values[k][1])) if avg > 1e-3: info += ' %.4f' % avg else: @@ -429,7 +443,86 @@ class Progbar(object): sys.stdout.write(info) sys.stdout.flush() - self.last_update = now + self._last_update = now def add(self, n, values=None): - self.update(self.seen_so_far + n, values) + self.update(self._seen_so_far + n, values) + + +def make_batches(size, batch_size): + """Returns a list of batch indices (tuples of indices). + + Arguments: + size: Integer, total size of the data to slice into batches. + batch_size: Integer, batch size. + + Returns: + A list of tuples of array indices. + """ + num_batches = int(np.ceil(size / float(batch_size))) + return [(i * batch_size, min(size, (i + 1) * batch_size)) + for i in range(0, num_batches)] + + +def slice_arrays(arrays, start=None, stop=None): + """Slice an array or list of arrays. + + This takes an array-like, or a list of + array-likes, and outputs: + - arrays[start:stop] if `arrays` is an array-like + - [x[start:stop] for x in arrays] if `arrays` is a list + + Can also work on list/array of indices: `slice_arrays(x, indices)` + + Arguments: + arrays: Single array or list of arrays. + start: can be an integer index (start index) + or a list/array of indices + stop: integer (stop index); should be None if + `start` was a list. + + Returns: + A slice of the array(s). + + Raises: + ValueError: If the value of start is a list and stop is not None. + """ + if arrays is None: + return [None] + if isinstance(start, list) and stop is not None: + raise ValueError('The stop argument has to be None if the value of start ' + 'is a list.') + elif isinstance(arrays, list): + if hasattr(start, '__len__'): + # hdf5 datasets only support list objects as indices + if hasattr(start, 'shape'): + start = start.tolist() + return [None if x is None else x[start] for x in arrays] + else: + return [None if x is None else x[start:stop] for x in arrays] + else: + if hasattr(start, '__len__'): + if hasattr(start, 'shape'): + start = start.tolist() + return arrays[start] + elif hasattr(start, '__getitem__'): + return arrays[start:stop] + else: + return [None] + + +def to_list(x): + """Normalizes a list/tensor into a list. + + If a tensor is passed, we return + a list of size 1 containing the tensor. + + Arguments: + x: target object to be normalized. + + Returns: + A list. + """ + if isinstance(x, list): + return x + return [x] diff --git a/tensorflow/python/keras/_impl/keras/utils/io_utils.py b/tensorflow/python/keras/_impl/keras/utils/io_utils.py index b36c769843d13a910efa6cf8c0d5309e3333f69b..bbf1d2a3d9c3948271780ec3fad3316b4e6d53c3 100644 --- a/tensorflow/python/keras/_impl/keras/utils/io_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/io_utils.py @@ -22,6 +22,7 @@ from collections import defaultdict import sys import numpy as np +from tensorflow.python.util.tf_export import tf_export try: @@ -30,6 +31,7 @@ except ImportError: h5py = None +@tf_export('keras.utils.HDF5Matrix') class HDF5Matrix(object): """Representation of HDF5 dataset to be used instead of a Numpy array. diff --git a/tensorflow/python/keras/_impl/keras/utils/layer_utils.py b/tensorflow/python/keras/_impl/keras/utils/layer_utils.py index a2d32424b51ea3160d530a5162ccc269e5815e57..4c8009dfd80e1aec457fa03687f2840c7fe4607b 100644 --- a/tensorflow/python/keras/_impl/keras/utils/layer_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/layer_utils.py @@ -23,6 +23,7 @@ import numpy as np from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras.utils.conv_utils import convert_kernel +from tensorflow.python.util.tf_export import tf_export def count_params(weights): @@ -58,6 +59,10 @@ def print_summary(model, line_length=None, positions=None, print_fn=None): if model.__class__.__name__ == 'Sequential': sequential_like = True + elif not model._is_graph_network: + # We treat subclassed models as a simple sequence of layers, for logging + # purposes. + sequential_like = True else: sequential_like = True nodes_by_depth = model._nodes_by_depth.values() @@ -117,17 +122,24 @@ def print_summary(model, line_length=None, positions=None, print_fn=None): print_fn('=' * line_length) def print_layer_summary(layer): + """Prints a summary for a single layer. + + Arguments: + layer: target layer. + """ try: output_shape = layer.output_shape except AttributeError: output_shape = 'multiple' + except RuntimeError: # output_shape unknown in Eager mode. + output_shape = '?' name = layer.name cls_name = layer.__class__.__name__ fields = [name + ' (' + cls_name + ')', output_shape, layer.count_params()] print_row(fields, positions) def print_layer_summary_with_connections(layer): - """Prints a summary for a single layer. + """Prints a summary for a single layer (including topological connections). Arguments: layer: target layer. @@ -190,6 +202,7 @@ def print_summary(model, line_length=None, positions=None, print_fn=None): print_fn('_' * line_length) +@tf_export('keras.utils.convert_all_kernels_in_model') def convert_all_kernels_in_model(model): """Converts all convolution kernels in a model from Theano to TensorFlow. diff --git a/tensorflow/python/keras/_impl/keras/utils/training_utils.py b/tensorflow/python/keras/_impl/keras/utils/multi_gpu_utils.py similarity index 98% rename from tensorflow/python/keras/_impl/keras/utils/training_utils.py rename to tensorflow/python/keras/_impl/keras/utils/multi_gpu_utils.py index ce7402e9d279278eaaf5aab58a3973eec6de8e99..231ace2a0b4a4f25cebf06a5216cf3d30aadc49b 100644 --- a/tensorflow/python/keras/_impl/keras/utils/training_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/multi_gpu_utils.py @@ -125,7 +125,7 @@ def multi_gpu_model(model, gpus): if gpus <= 1: raise ValueError('For multi-gpu usage to be effective, ' 'call `multi_gpu_model` with `gpus >= 2`. ' - 'Received: `gpus=%d`' % gpus) + 'Received: `gpus=%s`' % gpus) num_gpus = gpus target_gpu_ids = range(num_gpus) @@ -136,7 +136,7 @@ def multi_gpu_model(model, gpus): ] for device in target_devices: if device not in available_devices: - raise ValueError('To call `multi_gpu_model` with `gpus=%d`, ' + raise ValueError('To call `multi_gpu_model` with `gpus=%s`, ' 'we expect the following devices to be available: %s. ' 'However this machine only has: %s. ' 'Try reducing `gpus`.' % (gpus, target_devices, diff --git a/tensorflow/python/keras/_impl/keras/utils/training_utils_test.py b/tensorflow/python/keras/_impl/keras/utils/multi_gpu_utils_test.py similarity index 64% rename from tensorflow/python/keras/_impl/keras/utils/training_utils_test.py rename to tensorflow/python/keras/_impl/keras/utils/multi_gpu_utils_test.py index 12354c49ca72cddc0f395bcfcfabab18c1189227..0a38d6b5228fe791ce14adc7e37e0b7a6926fadf 100644 --- a/tensorflow/python/keras/_impl/keras/utils/training_utils_test.py +++ b/tensorflow/python/keras/_impl/keras/utils/multi_gpu_utils_test.py @@ -19,21 +19,34 @@ from __future__ import print_function import numpy as np - +from tensorflow.python import data from tensorflow.python.keras._impl import keras from tensorflow.python.platform import test +def check_if_compatible_devices(gpus=2): + available_devices = [ + keras.utils.multi_gpu_utils._normalize_device_name(name) + for name in keras.utils.multi_gpu_utils._get_available_devices() + ] + if '/gpu:%d' % (gpus - 1) not in available_devices: + return False + return True + + class TestMultiGPUModel(test.TestCase): - def multi_gpu_test_simple_model(self): + def test_multi_gpu_test_simple_model(self): gpus = 2 num_samples = 1000 input_dim = 10 output_dim = 1 hidden_dim = 10 epochs = 2 - target_gpu_id = [0, 2, 4] + target_gpu_id = [0, 1] + + if not check_if_compatible_devices(gpus=gpus): + return with self.test_session(): model = keras.models.Sequential() @@ -47,12 +60,11 @@ class TestMultiGPUModel(test.TestCase): parallel_model = keras.utils.multi_gpu_model(model, gpus=gpus) parallel_model.compile(loss='mse', optimizer='rmsprop') parallel_model.fit(x, y, epochs=epochs) - parallel_model = keras.utils.multi_gpu_model(model, gpus=target_gpu_id) parallel_model.compile(loss='mse', optimizer='rmsprop') parallel_model.fit(x, y, epochs=epochs) - def multi_gpu_test_multi_io_model(self): + def test_multi_gpu_test_multi_io_model(self): gpus = 2 num_samples = 1000 input_dim_a = 10 @@ -61,7 +73,10 @@ class TestMultiGPUModel(test.TestCase): output_dim_b = 2 hidden_dim = 10 epochs = 2 - target_gpu_id = [0, 2, 4] + target_gpu_id = [0, 1] + + if not check_if_compatible_devices(gpus=gpus): + return with self.test_session(): input_a = keras.Input((input_dim_a,)) @@ -86,7 +101,10 @@ class TestMultiGPUModel(test.TestCase): parallel_model.compile(loss='mse', optimizer='rmsprop') parallel_model.fit([a_x, b_x], [a_y, b_y], epochs=epochs) - def multi_gpu_test_invalid_devices(self): + def test_multi_gpu_test_invalid_devices(self): + if not check_if_compatible_devices(gpus=2): + return + with self.test_session(): input_shape = (1000, 10) model = keras.models.Sequential() @@ -115,3 +133,53 @@ class TestMultiGPUModel(test.TestCase): with self.assertRaises(ValueError): parallel_model = keras.utils.multi_gpu_model(model, gpus=[0]) parallel_model.fit(x, y, epochs=2) + + def test_nested_model_with_tensor_input(self): + gpus = 2 + input_dim = 10 + shape = (input_dim,) + num_samples = 16 + num_classes = 10 + + if not check_if_compatible_devices(gpus=gpus): + return + + with self.test_session(): + input_shape = (num_samples,) + shape + x_train = np.random.randint(0, 255, input_shape) + y_train = np.random.randint(0, num_classes, (input_shape[0],)) + keras.backend.set_learning_phase(True) + + y_train = keras.utils.to_categorical(y_train, num_classes) + + x_train = x_train.astype('float32') + y_train = y_train.astype('float32') + + dataset = data.Dataset.from_tensor_slices((x_train, y_train)) + dataset = dataset.repeat() + dataset = dataset.batch(4) + iterator = dataset.make_one_shot_iterator() + + inputs, targets = iterator.get_next() + + input_tensor = keras.layers.Input(tensor=inputs) + + model = keras.models.Sequential() + model.add(keras.layers.Dense(3, + input_shape=(input_dim,))) + model.add(keras.layers.Dense(num_classes)) + + output = model(input_tensor) + outer_model = keras.Model(input_tensor, output) + parallel_model = keras.utils.multi_gpu_model(outer_model, gpus=gpus) + + parallel_model.compile( + loss='categorical_crossentropy', + optimizer=keras.optimizers.RMSprop(lr=0.0001, decay=1e-6), + metrics=['accuracy'], + target_tensors=[targets]) + parallel_model.fit(epochs=1, steps_per_epoch=3) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/python/keras/_impl/keras/utils/np_utils.py b/tensorflow/python/keras/_impl/keras/utils/np_utils.py index 231833e7760f1824df673ce84bc5d77df91721b0..a611be08aaed824ebb278b4b28ef52ea1872563b 100644 --- a/tensorflow/python/keras/_impl/keras/utils/np_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/np_utils.py @@ -18,8 +18,10 @@ from __future__ import division from __future__ import print_function import numpy as np +from tensorflow.python.util.tf_export import tf_export +@tf_export('keras.utils.to_categorical') def to_categorical(y, num_classes=None): """Converts a class vector (integers) to binary class matrix. @@ -48,6 +50,7 @@ def to_categorical(y, num_classes=None): return categorical +@tf_export('keras.utils.normalize') def normalize(x, axis=-1, order=2): """Normalizes a Numpy array. diff --git a/tensorflow/python/keras/_impl/keras/utils/vis_utils.py b/tensorflow/python/keras/_impl/keras/utils/vis_utils.py index 0c5f2c19c79d5007882dcbc5d14a4cc8dd57ce3d..4761cece82c727e4962d0374f8efb80dfaeac3c6 100644 --- a/tensorflow/python/keras/_impl/keras/utils/vis_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/vis_utils.py @@ -20,6 +20,7 @@ from __future__ import division from __future__ import print_function import os +from tensorflow.python.util.tf_export import tf_export try: @@ -119,7 +120,7 @@ def model_to_dot(model, show_shapes=False, show_layer_names=True, rankdir='TB'): layer_id = str(id(layer)) for i, node in enumerate(layer._inbound_nodes): node_key = layer.name + '_ib-' + str(i) - if node_key in model._container_nodes: + if node_key in model._network_nodes: # pylint: disable=protected-access for inbound_layer in node.inbound_layers: inbound_layer_id = str(id(inbound_layer)) layer_id = str(id(layer)) @@ -127,6 +128,7 @@ def model_to_dot(model, show_shapes=False, show_layer_names=True, rankdir='TB'): return dot +@tf_export('keras.utils.plot_model') def plot_model(model, to_file='model.png', show_shapes=False, diff --git a/tensorflow/python/keras/_impl/keras/wrappers/scikit_learn.py b/tensorflow/python/keras/_impl/keras/wrappers/scikit_learn.py index 223ceac3deae643fa07594d10e551bea581eb641..2884dc84cc5d99511947e6f0f97b0bf8a505221f 100644 --- a/tensorflow/python/keras/_impl/keras/wrappers/scikit_learn.py +++ b/tensorflow/python/keras/_impl/keras/wrappers/scikit_learn.py @@ -26,6 +26,7 @@ import numpy as np from tensorflow.python.keras._impl.keras.models import Sequential from tensorflow.python.keras._impl.keras.utils.generic_utils import has_arg from tensorflow.python.keras._impl.keras.utils.np_utils import to_categorical +from tensorflow.python.util.tf_export import tf_export class BaseWrapper(object): @@ -187,6 +188,7 @@ class BaseWrapper(object): return res +@tf_export('keras.wrappers.scikit_learn.KerasClassifier') class KerasClassifier(BaseWrapper): """Implementation of the scikit-learn classifier API for Keras. """ @@ -309,6 +311,7 @@ class KerasClassifier(BaseWrapper): 'the `model.compile()` method.') +@tf_export('keras.wrappers.scikit_learn.KerasRegressor') class KerasRegressor(BaseWrapper): """Implementation of the scikit-learn regressor API for Keras. """ diff --git a/tensorflow/python/keras/datasets/fashion_mnist/__init__.py b/tensorflow/python/keras/datasets/fashion_mnist/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..7f5ddecc4707334d52ebf4966f2ec6141cce0d46 100644 --- a/tensorflow/python/keras/datasets/fashion_mnist/__init__.py +++ b/tensorflow/python/keras/datasets/fashion_mnist/__init__.py @@ -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. +# ============================================================================== +"""Fashion-MNIST dataset.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.keras._impl.keras.datasets.fashion_mnist import load_data + +del absolute_import +del division +del print_function diff --git a/tensorflow/python/keras/utils/__init__.py b/tensorflow/python/keras/utils/__init__.py index 91cc8607274a80a14dd27a64274da7f8f0aafab1..2f74cf031d0520c8d874b7269c52e3b9e1b9931b 100644 --- a/tensorflow/python/keras/utils/__init__.py +++ b/tensorflow/python/keras/utils/__init__.py @@ -30,9 +30,9 @@ from tensorflow.python.keras._impl.keras.utils.generic_utils import Progbar from tensorflow.python.keras._impl.keras.utils.generic_utils import serialize_keras_object from tensorflow.python.keras._impl.keras.utils.io_utils import HDF5Matrix from tensorflow.python.keras._impl.keras.utils.layer_utils import convert_all_kernels_in_model +from tensorflow.python.keras._impl.keras.utils.multi_gpu_utils import multi_gpu_model from tensorflow.python.keras._impl.keras.utils.np_utils import normalize from tensorflow.python.keras._impl.keras.utils.np_utils import to_categorical -from tensorflow.python.keras._impl.keras.utils.training_utils import multi_gpu_model from tensorflow.python.keras._impl.keras.utils.vis_utils import plot_model del absolute_import diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD index 3a6058054be4fad05daf20703384d57624d95deb..f27ca5c2051be81eb94488ba388c11bcfebafc4b 100644 --- a/tensorflow/python/kernel_tests/BUILD +++ b/tensorflow/python/kernel_tests/BUILD @@ -295,7 +295,6 @@ tf_py_test( "//tensorflow/python:nn_grad", ], data = ["//tensorflow/core:image_testdata"], - tags = ["no_windows"], ) tf_py_test( @@ -712,6 +711,18 @@ cuda_py_test( ], ) +tf_py_test( + name = "regex_replace_op_test", + size = "small", + srcs = ["regex_replace_op_test.py"], + additional_deps = [ + "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:dtypes", + "//tensorflow/python:string_ops", + ], +) + tf_py_test( name = "save_restore_ops_test", size = "small", @@ -1075,6 +1086,7 @@ cuda_py_test( tags = [ "no_windows", "noasan", + "notap", ], ) @@ -1126,7 +1138,6 @@ tf_py_test( "//tensorflow/python:variables", ], data = ["//tensorflow/core:lmdb_testdata"], - tags = ["no_windows"], ) cuda_py_test( @@ -1294,7 +1305,7 @@ cuda_py_test( cuda_py_test( name = "control_flow_ops_py_test", - # TOOD(b/70473603): change this back to "small" once the C API is + # TODO(b/70473603): change this back to "small" once the C API is # permanently enabled size = "medium", srcs = ["control_flow_ops_py_test.py"], @@ -1559,12 +1570,15 @@ cuda_py_test( "//third_party/py/numpy", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", + "//tensorflow/python:layers", "//tensorflow/python:framework", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:init_ops", + "//tensorflow/python:linalg_ops", "//tensorflow/python:math_ops", "//tensorflow/python:nn_ops", "//tensorflow/python:partitioned_variables", + "//tensorflow/python:random_ops", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", ], @@ -2313,7 +2327,6 @@ cuda_py_test( "//tensorflow/python:variables", ], shard_count = 4, - tags = ["no_windows"], ) cuda_py_test( @@ -2444,7 +2457,6 @@ cuda_py_test( "//tensorflow/python/eager:context", ], shard_count = 10, - tags = ["no_windows"], ) cuda_py_test( @@ -2892,6 +2904,40 @@ tf_py_test( ], ) +tf_py_test( + name = "accumulate_n_test", + size = "small", + srcs = ["accumulate_n_test.py"], + additional_deps = [ + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:gradients", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:variables", + ], +) + +tf_py_test( + name = "accumulate_n_eager_test", + size = "small", + srcs = ["accumulate_n_eager_test.py"], + additional_deps = [ + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:gradients", + "//tensorflow/python:math_ops", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python/eager:backprop", + "//tensorflow/python/eager:context", + "//tensorflow/python/eager:tape", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py b/tensorflow/python/kernel_tests/accumulate_n_eager_test.py similarity index 67% rename from tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py rename to tensorflow/python/kernel_tests/accumulate_n_eager_test.py index 8f44698da851b48abf831e957c80fa1643a58bda..dc11b7deceb9040584aca1f629f4d003aef39428 100644 --- a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py +++ b/tensorflow/python/kernel_tests/accumulate_n_eager_test.py @@ -12,53 +12,41 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for new version of accumulate_n op that will eventually go into -`ops.math_ops`. - -These test cases spefically exercise the `eager` APIs. They need to be in a -separate file from the remaining tests because eager mode is currently something -you can turn on but can't turn off for the lifetime of the current process.""" +"""Tests for new version of accumulate_n op.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function 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 - class AccumulateNV2EagerTest(test_util.TensorFlowTestCase): - """Tests of the new, differentiable version of accumulate_n""" + """Tests of the new, differentiable version of accumulate_n.""" def testMinimalEagerMode(self): forty = constant_op.constant(40) two = constant_op.constant(2) - answer = av2.accumulate_n_v2([forty, two]) + answer = math_ops.accumulate_n([forty, two]) self.assertEqual(42, answer.numpy()) - def testFloat(self): np.random.seed(12345) x = [np.random.random((1, 2, 3, 4, 5)) - 0.5 for _ in range(5)] tf_x = ops.convert_n_to_tensor(x) with self.test_session(use_gpu=True): - self.assertAllClose(sum(x), av2.accumulate_n_v2(tf_x).numpy()) - self.assertAllClose(x[0] * 5, av2.accumulate_n_v2([tf_x[0]] * 5).numpy()) + self.assertAllClose(sum(x), math_ops.accumulate_n(tf_x).numpy()) + self.assertAllClose(x[0] * 5, + math_ops.accumulate_n([tf_x[0]] * 5).numpy()) def testGrad(self): np.random.seed(42) @@ -70,16 +58,14 @@ class AccumulateNV2EagerTest(test_util.TensorFlowTestCase): ] def fn(first, second, third): - return av2.accumulate_n_v2([first, second, third]) + return math_ops.accumulate_n([first, second, third]) grad_fn = backprop.gradients_function(fn) grad = grad_fn(input_vars[0], input_vars[1], input_vars[2]) - self.assertAllEqual(np.repeat(1.0, num_inputs), # d/dx (x + y + ...) = 1 + self.assertAllEqual(np.repeat(1.0, num_inputs), # d/dx (x + y + ...) = 1 [elem.numpy() for elem in grad]) - if __name__ == "__main__": ops.enable_eager_execution() test.main() - diff --git a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py b/tensorflow/python/kernel_tests/accumulate_n_test.py similarity index 75% rename from tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py rename to tensorflow/python/kernel_tests/accumulate_n_test.py index 6f65fe771eb77c10d0914faa90886b587adae68c..b793906fac2cd12a5c0c663dd169000ad6067759 100644 --- a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py +++ b/tensorflow/python/kernel_tests/accumulate_n_test.py @@ -12,43 +12,49 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for new version of accumulate_n op that will eventually go into -`ops.math_ops`.""" +"""Tests for new version of accumulate_n op.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function 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 +from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradients +from tensorflow.python.ops import math_ops 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""" + """Tests of the new, differentiable version of accumulate_n.""" def testFloat(self): np.random.seed(12345) x = [np.random.random((1, 2, 3, 4, 5)) - 0.5 for _ in range(5)] tf_x = ops.convert_n_to_tensor(x) with self.test_session(use_gpu=True): - self.assertAllClose(sum(x), av2.accumulate_n_v2(tf_x).eval()) - self.assertAllClose(x[0] * 5, av2.accumulate_n_v2([tf_x[0]] * 5).eval()) + self.assertAllClose(sum(x), math_ops.accumulate_n(tf_x).eval()) + self.assertAllClose(x[0] * 5, + math_ops.accumulate_n([tf_x[0]] * 5).eval()) def testInt(self): np.random.seed(54321) x = [np.random.randint(-128, 128, (5, 4, 3, 2, 1)) for _ in range(6)] tf_x = ops.convert_n_to_tensor(x) with self.test_session(use_gpu=True): - self.assertAllEqual(sum(x), av2.accumulate_n_v2(tf_x).eval()) - self.assertAllEqual(x[0] * 6, av2.accumulate_n_v2([tf_x[0]] * 6).eval()) + self.assertAllEqual(sum(x), math_ops.accumulate_n(tf_x).eval()) + self.assertAllEqual(x[0] * 6, + math_ops.accumulate_n([tf_x[0]] * 6).eval()) + + def testUnknownShape(self): + with self.test_session(use_gpu=True): + x0 = array_ops.placeholder(dtype=dtypes_lib.int32, shape=[None]) + acc = math_ops.accumulate_n([x0, x0], shape=[None]) + self.assertAllEqual([2, 4], acc.eval(feed_dict={x0: [1, 2]})) def testGrad(self): np.random.seed(42) @@ -56,9 +62,9 @@ class AccumulateNV2Test(test_util.TensorFlowTestCase): with self.test_session(use_gpu=True) as sess: input_vars = [ variables.Variable(10.0 * np.random.random()) - for i in range(0, num_inputs) + for _ in range(0, num_inputs) ] - accum_n = av2.accumulate_n_v2(input_vars) + accum_n = math_ops.accumulate_n(input_vars) sess.run(variables.global_variables_initializer()) accum_n_grad = gradients.gradients(accum_n, input_vars) self.assertAllEqual( @@ -78,7 +84,7 @@ class AccumulateNV2Test(test_util.TensorFlowTestCase): ops.convert_to_tensor(x, dtype=dtypes_lib.float32) for x in random_arrays ] - tf_val = av2.accumulate_n_v2(random_tensors) + tf_val = math_ops.accumulate_n(random_tensors) np_val = random_arrays[0] for random_array in random_arrays[1:]: np_val += random_array @@ -87,7 +93,7 @@ class AccumulateNV2Test(test_util.TensorFlowTestCase): def testZeroArgs(self): with self.test_session(): with self.assertRaises(ValueError): - tf_val = av2.accumulate_n_v2([]) + tf_val = math_ops.accumulate_n([]) tf_val.eval() def testWrongShape(self): @@ -95,28 +101,28 @@ 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=[] + math_ops.accumulate_n([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]) + math_ops.accumulate_n([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) + math_ops.accumulate_n([a, b], tensor_dtype=np.int32) def testWrongTypeOneInput(self): # Scenario that used to trigger a bug, even when testWrongType() worked with self.test_session(): with self.assertRaises(TypeError): a = variables.Variable(0.2, dtype=np.float32) - tf_val = av2.accumulate_n_v2([a], tensor_dtype=np.int32) + math_ops.accumulate_n([a], tensor_dtype=np.int32) if __name__ == "__main__": diff --git a/tensorflow/python/kernel_tests/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops_test.py index 68b7c3a98a160017c2d5d84d3d9a8f92a6ab66b0..d0ba8020c1eaee74ded5ad67cae39b51d44097bd 100644 --- a/tensorflow/python/kernel_tests/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops_test.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function import time +import unittest import numpy as np @@ -414,7 +415,7 @@ class MeshgridTest(test_util.TensorFlowTestCase): def _compareDiffType(self, n, np_dtype, use_gpu): inputs = [] for index in ("ij", "xy"): - for i in range(n): + for _ in range(n): x = np.linspace(-10, 10, 5).astype(np_dtype) if np_dtype in (np.complex64, np.complex128): x += 1j @@ -422,8 +423,8 @@ class MeshgridTest(test_util.TensorFlowTestCase): numpy_out = np.meshgrid(*inputs, indexing=index) with self.test_session(use_gpu=use_gpu): tf_out = array_ops.meshgrid(*inputs, indexing=index) - for X, _X in zip(numpy_out, tf_out): - self.assertAllEqual(X, _X.eval()) + for x_np, x_tf in zip(numpy_out, tf_out): + self.assertAllEqual(x_np, x_tf.eval()) def testCompare(self): for t in (np.float16, np.float32, np.float64, np.int32, np.int64, @@ -497,7 +498,7 @@ class StridedSliceTest(test_util.TensorFlowTestCase): def test_basic_slice(self): for tensor_type in STRIDED_SLICE_TYPES: - with self.test_session(use_gpu=True): + with self.test_session(use_gpu=not tensor_type.is_integer): checker = StridedSliceChecker( self, StridedSliceChecker.REF_TENSOR, tensor_type=tensor_type) _ = checker[:, :, :] @@ -883,7 +884,8 @@ class StridedSliceAssignChecker(object): if self.tensor_type.is_complex: value -= 1j * value - with self.test.test_session(use_gpu=True) as sess: + with self.test.test_session( + use_gpu=not self.tensor_type.is_integer) as sess: if self._use_resource: var = resource_variable_ops.ResourceVariable(self.x) else: @@ -952,6 +954,30 @@ class SliceAssignTest(test_util.TensorFlowTestCase): v = variables.Variable([1, 2]) sess.run(v[:].assign([1, 2])) + def testTypeError(self): + init_val = constant_op.constant([1, 2], dtype=dtypes.int32) + too_small_val = constant_op.constant([3, 4], dtype=dtypes.int8) + too_large_val = constant_op.constant([3, 4], dtype=dtypes.int64) + v = variables.Variable(init_val) + with self.assertRaises(TypeError): + v[:].assign(too_small_val) + with self.assertRaises(TypeError): + v[:].assign(too_large_val) + + def testTypeErrorResource(self): + init_val = constant_op.constant([1, 2], dtype=dtypes.int32) + too_small_val = constant_op.constant([3, 4], dtype=dtypes.int8) + too_large_val = constant_op.constant([3, 4], dtype=dtypes.int64) + v = resource_variable_ops.ResourceVariable(init_val) + with self.test_session() as sess: + sess.run(v.initializer) + with self.assertRaisesRegexp( + errors.InvalidArgumentError, + "l-value dtype int32 does not match r-value dtype int64"): + sess.run(v[:].assign(too_large_val)) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(v[:].assign(too_small_val)) + class ShapeSizeRankTest(test_util.TensorFlowTestCase): @@ -989,7 +1015,7 @@ class SequenceMaskTest(test_util.TensorFlowTestCase): with self.assertRaisesRegexp(ValueError, "maxlen must be scalar"): array_ops.sequence_mask([10, 20], [10, 20]) - def testOneDimensional(self): + def testOneDimensionalWithMaxlen(self): with self.test_session(): res = array_ops.sequence_mask(constant_op.constant([1, 3, 2]), 5) self.assertAllEqual(res.get_shape(), [3, 5]) @@ -998,9 +1024,12 @@ class SequenceMaskTest(test_util.TensorFlowTestCase): [[True, False, False, False, False], [True, True, True, False, False], [True, True, False, False, False]]) + @test_util.enable_c_shapes + def testOneDimensionalDtypeWithoutMaxlen(self): + with self.test_session(): # test dtype and default maxlen: - res = array_ops.sequence_mask( - constant_op.constant([0, 1, 4]), dtype=dtypes.float32) + res = array_ops.sequence_mask(constant_op.constant([0, 1, 4]), + dtype=dtypes.float32) if ops._USE_C_API: self.assertAllEqual(res.get_shape().as_list(), [3, 4]) else: @@ -1009,6 +1038,22 @@ class SequenceMaskTest(test_util.TensorFlowTestCase): res.eval(), [[0.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]) + @test_util.enable_c_shapes + def testOneDimensionalWithoutMaxlen(self): + with self.test_session(): + res = array_ops.sequence_mask( + constant_op.constant([0, 1, 4])) + if ops._USE_C_API: + self.assertAllEqual(res.get_shape().as_list(), [3, 4]) + else: + self.assertAllEqual(res.get_shape().as_list(), [3, None]) + self.assertAllEqual( + res.eval(), + [[False, False, False, False], + [True, False, False, False], + [True, True, True, True]]) + + @test_util.enable_c_shapes def testTwoDimensional(self): with self.test_session(): res = array_ops.sequence_mask(constant_op.constant([[1, 3, 2]]), 5) @@ -1029,6 +1074,11 @@ class SequenceMaskTest(test_util.TensorFlowTestCase): [[[0.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]], [[1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]]]) + def testUnknownShape(self): + lengths = array_ops.placeholder(dtype=dtypes.int32) + res = array_ops.sequence_mask(lengths) + self.assertEqual(res.shape, None) + def testDtypes(self): def check_dtypes(lengths_dtype, maxlen_dtype): @@ -1114,8 +1164,11 @@ class InvertPermutationTest(test_util.TensorFlowTestCase): self.assertAllEqual(y.get_shape(), [5]) self.assertAllEqual(y.eval(), [2, 4, 3, 0, 1]) + class UnravelIndexTest(test_util.TensorFlowTestCase): + # TODO(b/73086570): Reenable test. + @unittest.skip("Test does not pass internally.") def testUnravelIndex(self): with self.test_session(): for dtype in [dtypes.int32, dtypes.int64]: @@ -1173,7 +1226,7 @@ class SnapshotOpTest(test_util.TensorFlowTestCase): for dtype in [dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64]: with self.test_session(use_gpu=True): x = constant_op.constant([0, 1, 2, 3], dtype=dtype) - y = gen_array_ops._snapshot(x) + y = gen_array_ops.snapshot(x) self.assertAllEqual(y.eval(), [0, 1, 2, 3]) diff --git a/tensorflow/python/kernel_tests/atrous_convolution_test.py b/tensorflow/python/kernel_tests/atrous_convolution_test.py index 2d1b3d9b7e836591646a2d0e59742bf6139446d1..0ef08581c9f931b991ef0c1218dc503345e248c2 100644 --- a/tensorflow/python/kernel_tests/atrous_convolution_test.py +++ b/tensorflow/python/kernel_tests/atrous_convolution_test.py @@ -83,14 +83,14 @@ class AtrousConvolutionTest(test.TestCase): checks = [] def add_check(check, *args, **kwargs): - if context.in_eager_mode(): + if context.executing_eagerly(): args_val, kwargs_val = self.evaluate([args, kwargs]) check(*args_val, **kwargs_val) else: checks.append((check, args, kwargs)) yield add_check - if context.in_graph_mode(): + if not context.executing_eagerly(): all_values = self.evaluate([[args, kwargs] for _, args, kwargs in checks]) for (check, _, _), (args, kwargs) in zip(checks, all_values): check(*args, **kwargs) diff --git a/tensorflow/python/kernel_tests/basic_gpu_test.py b/tensorflow/python/kernel_tests/basic_gpu_test.py index 405651e8ae97fbc5eefd4aba0a95a99ff8fd8c26..987a6ffcd4b18eb5857ff9e82206de7f6ebe8a27 100644 --- a/tensorflow/python/kernel_tests/basic_gpu_test.py +++ b/tensorflow/python/kernel_tests/basic_gpu_test.py @@ -33,7 +33,7 @@ from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variables -from tensorflow.python.ops.gen_array_ops import _broadcast_gradient_args +from tensorflow.python.ops.gen_array_ops import broadcast_gradient_args from tensorflow.python.platform import test @@ -157,7 +157,7 @@ class BroadcastSimpleTest(test.TestCase): def _GetGradientArgs(self, xs, ys): with self.test_session(use_gpu=True) as sess: - return sess.run(_broadcast_gradient_args(xs, ys)) + return sess.run(broadcast_gradient_args(xs, ys)) def testBroadcast(self): r0, r1 = self._GetGradientArgs([2, 3, 5], [1]) diff --git a/tensorflow/python/kernel_tests/batchtospace_op_test.py b/tensorflow/python/kernel_tests/batchtospace_op_test.py index 0c802476a0e788aff3de84ab736fa8f1de5daab4..6143cd3baa6317fc512d80f94b494710037d4082 100644 --- a/tensorflow/python/kernel_tests/batchtospace_op_test.py +++ b/tensorflow/python/kernel_tests/batchtospace_op_test.py @@ -44,7 +44,7 @@ class CppOpImpl(object): @staticmethod def batch_to_space(*args, **kwargs): - return gen_array_ops._batch_to_space(*args, **kwargs) + return gen_array_ops.batch_to_space(*args, **kwargs) class BatchToSpaceDepthToSpace(test.TestCase, PythonOpImpl): diff --git a/tensorflow/python/kernel_tests/bcast_ops_test.py b/tensorflow/python/kernel_tests/bcast_ops_test.py index 9e512346053a4c3af089170f47313606c4a307c2..3305e55c05bd03d31c46fd333db09dbab9a5d09c 100644 --- a/tensorflow/python/kernel_tests/bcast_ops_test.py +++ b/tensorflow/python/kernel_tests/bcast_ops_test.py @@ -20,8 +20,8 @@ from __future__ import print_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes -from tensorflow.python.ops.gen_array_ops import _broadcast_args -from tensorflow.python.ops.gen_array_ops import _broadcast_gradient_args +from tensorflow.python.ops.gen_array_ops import broadcast_args +from tensorflow.python.ops.gen_array_ops import broadcast_gradient_args from tensorflow.python.platform import test @@ -29,11 +29,11 @@ class BcastOpsTest(test.TestCase): def _GetBroadcastShape(self, xs, ys): with self.test_session() as sess: - return sess.run(_broadcast_args(xs, ys)) + return sess.run(broadcast_args(xs, ys)) def _GetGradientArgs(self, xs, ys): with self.test_session() as sess: - return sess.run(_broadcast_gradient_args(xs, ys)) + return sess.run(broadcast_gradient_args(xs, ys)) def testBasic(self): r = self._GetBroadcastShape([2, 3, 5], [1]) diff --git a/tensorflow/python/kernel_tests/check_ops_test.py b/tensorflow/python/kernel_tests/check_ops_test.py index 2e94603a3f3d4ca9074320cfb4e9bf06b6640e82..26d3df9e635ea49c6f25e3bb1df25ea72df27304 100644 --- a/tensorflow/python/kernel_tests/check_ops_test.py +++ b/tensorflow/python/kernel_tests/check_ops_test.py @@ -102,17 +102,15 @@ class AssertEqualTest(test.TestCase): with self.assertRaisesRegexp(errors.InvalidArgumentError, "fail"): check_ops.assert_equal(static_big, static_small, message="fail") - # Dynamic check - if context.in_graph_mode(): - with self.test_session(): - small = array_ops.placeholder(dtypes.int32, name="small") - big = array_ops.placeholder(dtypes.int32, name="big") - with ops.control_dependencies( - [check_ops.assert_equal( - big, small, message="fail")]): - out = array_ops.identity(small) - with self.assertRaisesOpError("fail.*big.*small"): - out.eval(feed_dict={small: [1, 2], big: [3, 4]}) + def test_raises_when_greater_dynamic(self): + with self.test_session(): + small = array_ops.placeholder(dtypes.int32, name="small") + big = array_ops.placeholder(dtypes.int32, name="big") + with ops.control_dependencies( + [check_ops.assert_equal(big, small, message="fail")]): + out = array_ops.identity(small) + with self.assertRaisesOpError("fail.*big.*small"): + out.eval(feed_dict={small: [1, 2], big: [3, 4]}) def test_error_message_eager(self): expected_error_msg_full = r"""big does not equal small @@ -182,15 +180,14 @@ First 2 elements of y: with self.assertRaisesRegexp(errors.InvalidArgumentError, "fail"): check_ops.assert_equal(static_big, static_small, message="fail") - # Dynamic check - if context.in_graph_mode(): - with self.test_session(): - small = array_ops.placeholder(dtypes.int32, name="small") - big = array_ops.placeholder(dtypes.int32, name="big") - with ops.control_dependencies([check_ops.assert_equal(small, big)]): - out = array_ops.identity(small) - with self.assertRaisesOpError("small.*big"): - out.eval(feed_dict={small: [3, 1], big: [4, 2]}) + def test_raises_when_less_dynamic(self): + with self.test_session(): + small = array_ops.placeholder(dtypes.int32, name="small") + big = array_ops.placeholder(dtypes.int32, name="big") + with ops.control_dependencies([check_ops.assert_equal(small, big)]): + out = array_ops.identity(small) + with self.assertRaisesOpError("small.*big"): + out.eval(feed_dict={small: [3, 1], big: [4, 2]}) @test_util.run_in_graph_and_eager_modes() def test_doesnt_raise_when_equal_and_broadcastable_shapes(self): diff --git a/tensorflow/python/kernel_tests/checkpoint_ops_test.py b/tensorflow/python/kernel_tests/checkpoint_ops_test.py index a786d0a47e569f71812086fb93c21dc12660a2a5..7f147ba53a71539962f424158731e359724f664f 100644 --- a/tensorflow/python/kernel_tests/checkpoint_ops_test.py +++ b/tensorflow/python/kernel_tests/checkpoint_ops_test.py @@ -50,7 +50,7 @@ class GenerateVocabRemappingTest(test.TestCase): def test_generate_remapping_with_no_vocab_changes(self): """Tests where vocab does not change at all.""" - remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping( + remapping, num_present = gen_checkpoint_ops.generate_vocab_remapping( new_vocab_file=self.old_vocab_file, old_vocab_file=self.old_vocab_file, num_new_vocab=3, @@ -63,7 +63,7 @@ class GenerateVocabRemappingTest(test.TestCase): def test_generate_remapping_with_shifted_vocab(self): """Tests where vocab is the same, but shifted / ordered differently.""" - remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping( + remapping, num_present = gen_checkpoint_ops.generate_vocab_remapping( new_vocab_file=self.new_vocab_file, old_vocab_file=self.old_vocab_file, num_new_vocab=3, @@ -76,7 +76,7 @@ class GenerateVocabRemappingTest(test.TestCase): def test_generate_remapping_with_offset(self): """Tests offset and num_new_vocab logic.""" - remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping( + remapping, num_present = gen_checkpoint_ops.generate_vocab_remapping( new_vocab_file=self.new_vocab_file, old_vocab_file=self.old_vocab_file, num_new_vocab=1, @@ -89,7 +89,7 @@ class GenerateVocabRemappingTest(test.TestCase): def test_generate_remapping_with_old_vocab_size(self): """Tests where old_vocab_size is specified.""" - remapping, num_present = gen_checkpoint_ops._generate_vocab_remapping( + remapping, num_present = gen_checkpoint_ops.generate_vocab_remapping( new_vocab_file=self.new_vocab_file, old_vocab_file=self.old_vocab_file, num_new_vocab=3, @@ -132,7 +132,7 @@ class LoadAndRemapMatrixTest(test.TestCase): # No column remapping, new weight matrix has second row, then first row. row_remapping = [1, 0] - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=row_remapping, @@ -147,7 +147,7 @@ class LoadAndRemapMatrixTest(test.TestCase): # No row remapping, new weight matrix has third col, then first col. row_remapping = list(range(self.old_num_rows)) col_remapping = [2, 0] - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=row_remapping, @@ -162,7 +162,7 @@ class LoadAndRemapMatrixTest(test.TestCase): # Both row and column remappings. row_remapping = [1, 0, 4] col_remapping = [1, 15] - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=row_remapping, @@ -177,7 +177,7 @@ class LoadAndRemapMatrixTest(test.TestCase): def test_load_and_remap_with_init(self): """Tests the op's load and remap where there are missing entries.""" init_val = 42 - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=[2, -1, 0], @@ -196,7 +196,7 @@ class LoadAndRemapMatrixTest(test.TestCase): """Tests when all the rows are missing and need to be initialized.""" num_rows = 7 initializing_values = [42] * num_rows * self.old_num_cols - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=[-1] * num_rows, @@ -214,7 +214,7 @@ class LoadAndRemapMatrixTest(test.TestCase): num_rows = 7 num_cols = 4 initializing_values = [42] * num_rows * num_cols - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=[-1] * num_rows, @@ -235,7 +235,7 @@ class LoadAndRemapMatrixTest(test.TestCase): invalid_remapping = [1, 0, 0, 0, 1, 2] # Invalid row remapping. - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=invalid_remapping, @@ -247,7 +247,7 @@ class LoadAndRemapMatrixTest(test.TestCase): remapped_matrix.eval() # Invalid column remapping. - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=list(range(self.old_num_rows)), @@ -260,7 +260,7 @@ class LoadAndRemapMatrixTest(test.TestCase): def test_load_and_remap_incorrect_initializing_values(self): """Tests that errors are raised with incorrect number of init values.""" - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=[2, -1, 0], @@ -275,7 +275,7 @@ class LoadAndRemapMatrixTest(test.TestCase): with self.test_session(), self.assertRaises(errors.InvalidArgumentError): remapped_matrix.eval() - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=[self.bundle_file], old_tensor_name=self.old_tensor_name, row_remapping=[2, -1, 0], @@ -314,7 +314,7 @@ class LoadAndRemapMatrixWithMaxRowsTest(test.TestCase): num_rows, num_cols = np_value.shape # Tests loading the entire tensor (except reversed). - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=ckpt_path, old_tensor_name=old_tensor_name, # Simply reverses the rows of the matrix. @@ -332,7 +332,7 @@ class LoadAndRemapMatrixWithMaxRowsTest(test.TestCase): self.assertGreater(num_rows, 2) prefix_rows = 2 suffix_rows = 3 - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=ckpt_path, old_tensor_name=old_tensor_name, # Reverses the rows of the matrix, then prepends and appends @@ -353,7 +353,7 @@ class LoadAndRemapMatrixWithMaxRowsTest(test.TestCase): # Tests when everything is taken from initializing_values. new_rows = 7 initializing_values = [42] * new_rows * num_cols - remapped_matrix = gen_checkpoint_ops._load_and_remap_matrix( + remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=ckpt_path, old_tensor_name=old_tensor_name, # Nothing is loaded from the old tensor. diff --git a/tensorflow/python/kernel_tests/concat_op_test.py b/tensorflow/python/kernel_tests/concat_op_test.py index a5fd3bc3345f41d9d3f07278dc7979c1103b597f..c22934ce47543ab11b6a5b9acde2e2ec3aec9da7 100644 --- a/tensorflow/python/kernel_tests/concat_op_test.py +++ b/tensorflow/python/kernel_tests/concat_op_test.py @@ -495,9 +495,9 @@ class ConcatOpTest(test.TestCase): p = [] shape = np.array([7, 13]) if test.is_gpu_available(): - num_tensors = 10000 + num_tensors = 5000 else: - num_tensors = 1000 + num_tensors = 500 for i in np.arange(num_tensors): input_shape = shape placeholder = array_ops.placeholder(dtypes.float32, shape=input_shape) @@ -526,7 +526,7 @@ class ConcatOpTest(test.TestCase): with self.test_session(use_gpu=True): t1 = [] t2 = [] - output = gen_array_ops._concat_v2([t1, t2], 0).eval() + output = gen_array_ops.concat_v2([t1, t2], 0).eval() self.assertFalse(output) # Checks that output is empty def testConcatInvalidAxis(self): @@ -534,20 +534,20 @@ class ConcatOpTest(test.TestCase): with self.test_session(use_gpu=True): t1 = [1] t2 = [2] - gen_array_ops._concat_v2([t1, t2], 1).eval() + gen_array_ops.concat_v2([t1, t2], 1).eval() def testConcatNegativeAxis(self): with self.test_session(use_gpu=True): t1 = [[1, 2, 3], [4, 5, 6]] t2 = [[7, 8, 9], [10, 11, 12]] - c = gen_array_ops._concat_v2([t1, t2], -2) + c = gen_array_ops.concat_v2([t1, t2], -2) self.assertEqual([4, 3], c.get_shape().as_list()) output = c.eval() self.assertAllEqual([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], output) - c = gen_array_ops._concat_v2([t1, t2], -1) + c = gen_array_ops.concat_v2([t1, t2], -1) self.assertEqual([2, 6], c.get_shape().as_list()) output = c.eval() self.assertAllEqual([[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]], output) @@ -606,6 +606,17 @@ class ConcatOpTest(test.TestCase): inp_tensors_placeholders, -2, output_shape=[2, 3], gather_indexes=[2, 0], feed_dict=feed_dict) + def testConcatAxisType(self): + for dtype in [dtypes.int32, dtypes.int64]: + with self.test_session(use_gpu=True): + t1 = [[1, 2, 3], [4, 5, 6]] + t2 = [[7, 8, 9], [10, 11, 12]] + + c = gen_array_ops.concat_v2([t1, t2], + constant_op.constant(1, dtype=dtype)) + self.assertEqual([2, 6], c.get_shape().as_list()) + output = c.eval() + self.assertAllEqual([[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]], output) class ConcatOffsetTest(test.TestCase): @@ -615,7 +626,7 @@ class ConcatOffsetTest(test.TestCase): s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5], dtypes.int32) s2 = constant_op.constant([2, 20, 5], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1, s2]) + off = gen_array_ops.concat_offset(cdim, [s0, s1, s2]) ans = sess.run(off) self.assertAllEqual(ans, [[0, 0, 0], [0, 3, 0], [0, 10, 0]]) @@ -624,7 +635,7 @@ class ConcatOffsetTest(test.TestCase): cdim = constant_op.constant(1, dtypes.int32) s0 = constant_op.constant([[2, 3, 5]], dtypes.int32) s1 = constant_op.constant([[2, 7, 5]], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1]) + off = gen_array_ops.concat_offset(cdim, [s0, s1]) with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, r"should be a vector"): sess.run(off) @@ -634,7 +645,7 @@ class ConcatOffsetTest(test.TestCase): cdim = constant_op.constant(4, dtypes.int32) s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1]) + off = gen_array_ops.concat_offset(cdim, [s0, s1]) with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, r"Concat dim is out of range: 4 vs. 3"): sess.run(off) @@ -644,7 +655,7 @@ class ConcatOffsetTest(test.TestCase): cdim = constant_op.constant(1, dtypes.int32) s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5, 10], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1]) + off = gen_array_ops.concat_offset(cdim, [s0, s1]) with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, r"should contain 3 elem"): sess.run(off) @@ -654,7 +665,7 @@ class ConcatOffsetTest(test.TestCase): cdim = constant_op.constant(1, dtypes.int32) s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 10], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1]) + off = gen_array_ops.concat_offset(cdim, [s0, s1]) with self.assertRaisesRegexp( errors_impl.InvalidArgumentError, r"All dimensions except 1 must match. Input 1 has shape \[2 7 10\] " @@ -667,7 +678,7 @@ class ConcatOffsetTest(test.TestCase): s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5], dtypes.int32) s2 = constant_op.constant([2, 20, 5], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1, s2]) + off = gen_array_ops.concat_offset(cdim, [s0, s1, s2]) ans = sess.run(off) self.assertAllEqual(ans, [[0, 0, 0], [0, 3, 0], [0, 10, 0]]) @@ -675,7 +686,7 @@ class ConcatOffsetTest(test.TestCase): s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([1, 3, 5], dtypes.int32) s2 = constant_op.constant([3, 3, 5], dtypes.int32) - off = gen_array_ops._concat_offset(cdim, [s0, s1, s2]) + off = gen_array_ops.concat_offset(cdim, [s0, s1, s2]) ans = sess.run(off) self.assertAllEqual(ans, [[0, 0, 0], [2, 0, 0], [3, 0, 0]]) diff --git a/tensorflow/python/kernel_tests/constant_op_test.py b/tensorflow/python/kernel_tests/constant_op_test.py index 576bb68ba49cf5a5c7618131ad8a567672cb08d8..18796f709566f022258806ce46cc706e8fe34354 100644 --- a/tensorflow/python/kernel_tests/constant_op_test.py +++ b/tensorflow/python/kernel_tests/constant_op_test.py @@ -30,6 +30,7 @@ from tensorflow.python.framework import errors_impl from tensorflow.python.framework import importer 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 gradient_checker from tensorflow.python.ops import logging_ops @@ -180,6 +181,11 @@ class ConstantTest(test.TestCase): shape=[2, 3, 5]) self.assertEqual(c.get_shape(), [2, 3, 5]) + @test_util.assert_no_new_pyobjects_executing_eagerly + def testEagerMemory(self): + """Tests PyObject refs are managed correctly when executing eagerly.""" + constant_op.constant([[1.]]) + def testImplicitShapeNumPy(self): with ops.Graph().as_default(): c = constant_op.constant( @@ -465,9 +471,8 @@ class ZerosLikeTest(test.TestCase): def testZerosLikeGPU(self): for dtype in [ dtypes_lib.half, dtypes_lib.float32, dtypes_lib.float64, - dtypes_lib.int32, dtypes_lib.int64, - dtypes_lib.complex64, dtypes_lib.complex128, - dtypes_lib.bool + dtypes_lib.int32, dtypes_lib.int64, dtypes_lib.complex64, + dtypes_lib.complex128, dtypes_lib.bool ]: self._compareZeros(dtype, fully_defined_shape=False, use_gpu=True) self._compareZeros(dtype, fully_defined_shape=True, use_gpu=True) @@ -876,7 +881,7 @@ versions { class PlaceholderWithDefaultTest(test.TestCase): def testFullShape(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): p = array_ops.placeholder_with_default([[2, 2], [2, 2]], shape=[2, 2]) a = array_ops.identity(p) self.assertAllEqual([[2, 2], [2, 2]], a.eval()) @@ -887,7 +892,7 @@ class PlaceholderWithDefaultTest(test.TestCase): a.eval(feed_dict={p: [[6, 6, 6], [6, 6, 6]]}) def testPartialShape(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): p = array_ops.placeholder_with_default([1, 2, 3], shape=[None]) a = array_ops.identity(p) self.assertAllEqual([1, 2, 3], a.eval()) @@ -897,7 +902,7 @@ class PlaceholderWithDefaultTest(test.TestCase): a.eval(feed_dict={p: [[2, 2], [2, 2]]}) def testNoShape(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): p = array_ops.placeholder_with_default([17], shape=None) a = array_ops.identity(p) self.assertAllEqual([17], a.eval()) @@ -906,11 +911,12 @@ class PlaceholderWithDefaultTest(test.TestCase): [[3, 3], [3, 3]], a.eval(feed_dict={p: [[3, 3], [3, 3]]})) def testGradient(self): - with self.test_session(): + with self.test_session(force_gpu=test_util.is_gpu_available()): x = array_ops.placeholder(dtypes_lib.float32, [5, 7]) y = array_ops.placeholder_with_default(x, None) err = gradient_checker.compute_gradient_error(x, [5, 7], y, [5, 7]) self.assertLess(err, 1e-3) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py index 5d648bb235c8a8a0ae435e0c249bcb64ba787b08..75f8644f694c4cebb7dbdac4599244dda427bc05 100644 --- a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py +++ b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py @@ -44,6 +44,7 @@ from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import data_flow_ops from tensorflow.python.ops import functional_ops from tensorflow.python.ops import gen_array_ops +from tensorflow.python.ops import gen_control_flow_ops from tensorflow.python.ops import gen_data_flow_ops from tensorflow.python.ops import gen_logging_ops from tensorflow.python.ops import gen_state_ops @@ -143,7 +144,7 @@ class ControlFlowTest(test.TestCase): enter_v = control_flow_ops._Enter(v, "foo_1", is_constant=True) nine = constant_op.constant(9) - enter_nine = control_flow_ops.enter(nine, "foo_1") + enter_nine = gen_control_flow_ops.enter(nine, "foo_1") op = state_ops.assign(enter_v, enter_nine) v2 = control_flow_ops.with_dependencies([op], enter_v) v3 = control_flow_ops.exit(v2) @@ -163,9 +164,9 @@ class ControlFlowTest(test.TestCase): def testEnterMulExit(self): with self.test_session(): data = constant_op.constant([1, 2, 3, 4, 5, 6], name="data") - enter_data = control_flow_ops.enter(data, "foo_1", False) + enter_data = gen_control_flow_ops.enter(data, "foo_1", False) five = constant_op.constant(5) - enter_five = control_flow_ops.enter(five, "foo_1", False) + enter_five = gen_control_flow_ops.enter(five, "foo_1", False) mul_op = math_ops.multiply(enter_data, enter_five) exit_op = control_flow_ops.exit(mul_op) @@ -177,11 +178,12 @@ class ControlFlowTest(test.TestCase): v = variables.Variable([0.0, 0.0], dtype=dtypes.float32) # If is_constant=True, the shape information should be propagated. - enter_v_constant = control_flow_ops.enter(v, "frame1", is_constant=True) + enter_v_constant = gen_control_flow_ops.enter( + v, "frame1", is_constant=True) self.assertEqual(enter_v_constant.shape, [2]) # Otherwise, the shape should be unknown. - enter_v_non_constant = control_flow_ops.enter( + enter_v_non_constant = gen_control_flow_ops.enter( v, "frame2", is_constant=False) self.assertEqual(enter_v_non_constant.shape, None) @@ -255,8 +257,8 @@ class ControlFlowTest(test.TestCase): false = ops.convert_to_tensor(False) n = constant_op.constant(10) - enter_false = control_flow_ops.enter(false, "foo_1", False) - enter_n = control_flow_ops.enter(n, "foo_1", False) + enter_false = gen_control_flow_ops.enter(false, "foo_1", False) + enter_n = gen_control_flow_ops.enter(n, "foo_1", False) merge_n = control_flow_ops.merge([enter_n, enter_n], name="merge_n")[0] switch_n = control_flow_ops.switch(merge_n, enter_false) @@ -273,9 +275,9 @@ class ControlFlowTest(test.TestCase): one = constant_op.constant(1) n = constant_op.constant(10) - enter_i = control_flow_ops.enter(zero, "foo", False) - enter_one = control_flow_ops.enter(one, "foo", True) - enter_n = control_flow_ops.enter(n, "foo", True) + enter_i = gen_control_flow_ops.enter(zero, "foo", False) + enter_one = gen_control_flow_ops.enter(one, "foo", True) + enter_n = gen_control_flow_ops.enter(n, "foo", True) with ops.device(test.gpu_device_name()): merge_i = control_flow_ops.merge([enter_i, enter_i])[0] @@ -299,9 +301,9 @@ class ControlFlowTest(test.TestCase): one = constant_op.constant(1) n = constant_op.constant(10) - enter_i = control_flow_ops.enter(zero, "foo", False) - enter_one = control_flow_ops.enter(one, "foo", True) - enter_n = control_flow_ops.enter(n, "foo", True) + enter_i = gen_control_flow_ops.enter(zero, "foo", False) + enter_one = gen_control_flow_ops.enter(one, "foo", True) + enter_n = gen_control_flow_ops.enter(n, "foo", True) merge_i = control_flow_ops.merge([enter_i, enter_i])[0] @@ -322,8 +324,8 @@ class ControlFlowTest(test.TestCase): def testDifferentFrame(self): with self.test_session(): data = array_ops.placeholder(dtypes.float32, shape=[]) - enter_1 = control_flow_ops.enter(data, "foo_1", False) - enter_2 = control_flow_ops.enter(data, "foo_2", False) + enter_1 = gen_control_flow_ops.enter(data, "foo_1", False) + enter_2 = gen_control_flow_ops.enter(data, "foo_2", False) res = math_ops.add(enter_1, enter_2) with self.assertRaisesOpError("has inputs from different frames"): res.eval(feed_dict={data: 1.0}) @@ -550,7 +552,7 @@ class ControlFlowTest(test.TestCase): def testCondRef(self): with self.test_session(): - x = gen_state_ops._variable( + x = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="x", @@ -578,7 +580,7 @@ class ControlFlowTest(test.TestCase): def testUninitializedRefIdentity(self): with self.test_session() as sess: - v = gen_state_ops._variable( + v = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="v", @@ -589,10 +591,10 @@ class ControlFlowTest(test.TestCase): # Both v_f and v_t are uninitialized references. However, an actual use # of the reference in the 'true' branch in the 'tf.identity' op will # not 'fire' when v is uninitialized, so this is a valid construction. - # This test tests that _ref_identity allows uninitialized ref as input + # This test tests that ref_identity allows uninitialized ref as input # so that this construction is allowed. - v_f_op = gen_array_ops._ref_identity(v_f) - v_t_op = gen_array_ops._ref_identity(v_t) + v_f_op = gen_array_ops.ref_identity(v_f) + v_t_op = gen_array_ops.ref_identity(v_t) with ops.control_dependencies([v_f_op]): assign_v = state_ops.assign(v, [1.0]) with ops.control_dependencies([v_t_op]): @@ -631,7 +633,8 @@ class ControlFlowTest(test.TestCase): sess.run(r) def testCondGrad_1(self): - with self.test_session(): + graph = ops.Graph() + with graph.as_default(): x = constant_op.constant(10.0, name="x") pred = math_ops.less(1, 2) fn1 = lambda: array_ops.identity(x) @@ -639,8 +642,14 @@ class ControlFlowTest(test.TestCase): r = control_flow_ops.cond(pred, fn1, fn2) grad = gradients_impl.gradients(r, [x])[0] - result = grad.eval() - self.assertAllEqual(1.0, result) + with self.test_session(): + self.assertAllEqual(1.0, grad.eval()) + # The gradients computation creates a tensor with zeros by broadcasting a + # zeros constant to the required shape. Verify that the zero constant + # feeding into the fill is dominated by a Switch. + zero = graph.get_operation_by_name("gradients/zeros/Const") + self.assertEqual(len(zero.control_inputs), 1) + self.assertEqual(zero.control_inputs[0].type, "Switch") def testCondGrad_2(self): with self.test_session(): @@ -702,6 +711,36 @@ class ControlFlowTest(test.TestCase): r = control_flow_ops.while_loop(c, b, [n], parallel_iterations=20) self.assertEqual(10000, r.eval()) + def testWhileExternalControlDependencies(self): + with self.test_session(): + v = variables.Variable(0.0) + v.initializer.run() + increment = v.assign_add(1.0) + + def body_fn(i): + with ops.control_dependencies([increment]): + return i + i + + result = control_flow_ops.while_loop(cond=lambda i: i < 1, + body=body_fn, loop_vars=[1]) + result.eval() + self.assertAllEqual(v.eval(), 1.0) + + def testWhileExternalControlDependenciesNoInput(self): + with self.test_session(): + v = variables.Variable(0.0) + v.initializer.run() + increment = v.assign_add(1.0) + + def body_fn(unused_i): + with ops.control_dependencies([increment]): + return constant_op.constant(5, name="five") + + result = control_flow_ops.while_loop(cond=lambda i: i < 5, + body=body_fn, loop_vars=[0]) + result.eval() + self.assertAllEqual(v.eval(), 1.0) + def testWhileWithRefs_1(self): with self.test_session() as sess: x = variables.Variable(0)._ref() # pylint: disable=protected-access @@ -712,7 +751,7 @@ class ControlFlowTest(test.TestCase): def b(i, x): self.assertEqual(x.dtype, dtypes.int32_ref) - return (i + 1, gen_array_ops._ref_identity(x)) + return (i + 1, gen_array_ops.ref_identity(x)) r = control_flow_ops.while_loop(c, b, [i, x], parallel_iterations=5) @@ -1588,7 +1627,7 @@ class ControlFlowTest(test.TestCase): def testWhileStack_1(self): with self.test_session(): - s = gen_data_flow_ops._stack_v2(-1, dtypes.int32, stack_name="foo") + s = gen_data_flow_ops.stack_v2(-1, dtypes.int32, stack_name="foo") i = constant_op.constant(0) def c(i): @@ -1597,7 +1636,7 @@ class ControlFlowTest(test.TestCase): def b(i): ni = math_ops.add(i, 1) ni = control_flow_ops.with_dependencies( - [gen_data_flow_ops._stack_push_v2(s, i)], ni) + [gen_data_flow_ops.stack_push_v2(s, i)], ni) return ni r = control_flow_ops.while_loop(c, b, [i], parallel_iterations=1) @@ -1609,7 +1648,7 @@ class ControlFlowTest(test.TestCase): def b1(i, x): ni = math_ops.subtract(i, 1) - nx = x + gen_data_flow_ops._stack_pop_v2(s, dtypes.int32) + nx = x + gen_data_flow_ops.stack_pop_v2(s, dtypes.int32) return [ni, nx] _, rx = control_flow_ops.while_loop( @@ -1808,6 +1847,23 @@ class ControlFlowTest(test.TestCase): [tensor_shape.unknown_shape()]) self.assertAllClose(9.0, r.eval(feed_dict={x: 1.0})) + def testCondGradInNestedWhiles(self): + def outer_body(i, x): + _, x = control_flow_ops.while_loop( + lambda j, x: j < 3, inner_body, [0, 0.0]) + return i + 1, x + + def inner_body(j, x): + y = control_flow_ops.cond(math_ops.less(x, 1), lambda: 2 * x, lambda: x) + return j + 1, gradients_impl.gradients(y, x)[0] + + i, x = control_flow_ops.while_loop(lambda i, x: i < 3, outer_body, [0, 0.0]) + + with self.test_session() as sess: + i_val, x_val = sess.run([i, x]) + self.assertEqual(i_val, 3) + self.assertAllClose(x_val, 1.0) + def testWhile_NestedInput(self): with self.test_session() as sess: named = collections.namedtuple("named", ("a", "b")) @@ -2156,12 +2212,9 @@ class ControlFlowTest(test.TestCase): self.assertEqual(x.dtype, dtypes.int32_ref) - # pylint: disable=protected-access def body(i, x): self.assertEqual(x.dtype, dtypes.int32_ref) - return [i + 1, gen_array_ops._ref_identity(x)] - - # pylint: enable=protected-access + return [i + 1, gen_array_ops.ref_identity(x)] r = control_flow_ops.while_loop(c, body, [i, x], parallel_iterations=5) diff --git a/tensorflow/python/kernel_tests/conv2d_transpose_test.py b/tensorflow/python/kernel_tests/conv2d_transpose_test.py index 7d0bc54b6993daff0298f9d76e9e67dfcbfa5711..b692d3da609fd97a55b8f5fce3334b8e9d97c827 100644 --- a/tensorflow/python/kernel_tests/conv2d_transpose_test.py +++ b/tensorflow/python/kernel_tests/conv2d_transpose_test.py @@ -21,7 +21,6 @@ from __future__ import print_function import numpy as np from six.moves import xrange # pylint: disable=redefined-builtin -from tensorflow.python.client import device_lib from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops @@ -175,7 +174,7 @@ class Conv2DTransposeTest(test.TestCase): self.assertLess(err, err_tolerance) def testConv2DTransposeSingleStrideNCHW(self): - # `NCHW` data fomat is only supported for CUDA device. + # `NCHW` data format is only supported for CUDA device. if test.is_gpu_available(cuda_only=True): with self.test_session(use_gpu=True): strides = [1, 1, 1, 1] @@ -210,7 +209,7 @@ class Conv2DTransposeTest(test.TestCase): self.assertAllClose(target, value[n, k, h, w]) def testConv2DTransposeSameNCHW(self): - # `NCHW` data fomat is only supported for CUDA device. + # `NCHW` data format is only supported for CUDA device. if test.is_gpu_available(cuda_only=True): with self.test_session(use_gpu=True): strides = [1, 1, 2, 2] @@ -246,7 +245,7 @@ class Conv2DTransposeTest(test.TestCase): self.assertAllClose(target, value[n, k, h, w]) def testConv2DTransposeValidNCHW(self): - # `NCHW` data fomat is only supported for CUDA device. + # `NCHW` data format is only supported for CUDA device. if test.is_gpu_available(cuda_only=True): with self.test_session(use_gpu=True): strides = [1, 1, 2, 2] diff --git a/tensorflow/python/kernel_tests/conv_ops_test.py b/tensorflow/python/kernel_tests/conv_ops_test.py index c5446326ba1763b962b32866235251d773e069f2..25525cc1285d8dbc305fdcf838143b6e1c420316 100644 --- a/tensorflow/python/kernel_tests/conv_ops_test.py +++ b/tensorflow/python/kernel_tests/conv_ops_test.py @@ -24,7 +24,7 @@ import time import numpy as np -from six.moves import xrange +from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.contrib import layers from tensorflow.python.client import session as session_lib from tensorflow.python.framework import constant_op @@ -302,25 +302,20 @@ class Conv2DTest(test.TestCase): padding, dilations): expected_results = [] computed_results = [] - default_dilations = (dilations[0] == 1 and dilations[1] == 1) for data_format, use_gpu in GetTestConfigs(): - # If any dilation rate is larger than 1, only do test on the GPU - # because we currently do not have a CPU implementation for arbitrary - # dilation rates. - if default_dilations or use_gpu: - expected, computed = self._ComputeReferenceDilatedConv( - tensor_in_sizes, filter_in_sizes, strides, dilations, padding, - data_format, use_gpu) - expected_results.append(expected) - computed_results.append(computed) - tolerance = 1e-2 if use_gpu else 1e-5 - expected_values = self.evaluate(expected_results) - computed_values = self.evaluate(computed_results) - for e_value, c_value in zip(expected_values, computed_values): - print("expected = ", e_value) - print("actual = ", c_value) - self.assertAllClose( - e_value.flatten(), c_value.flatten(), atol=tolerance, rtol=1e-4) + expected, computed = self._ComputeReferenceDilatedConv( + tensor_in_sizes, filter_in_sizes, strides, dilations, padding, + data_format, use_gpu) + expected_results.append(expected) + computed_results.append(computed) + tolerance = 1e-2 if use_gpu else 1e-5 + expected_values = self.evaluate(expected_results) + computed_values = self.evaluate(computed_results) + for e_value, c_value in zip(expected_values, computed_values): + print("expected = ", e_value) + print("actual = ", c_value) + self.assertAllClose( + e_value.flatten(), c_value.flatten(), atol=tolerance, rtol=1e-4) def _VerifyValues(self, tensor_in_sizes, filter_in_sizes, strides, padding, expected): @@ -365,13 +360,12 @@ class Conv2DTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testConv2D2x2Filter2x1Dilation(self): - if test.is_gpu_available(cuda_only=True): - self._VerifyDilatedConvValues( - tensor_in_sizes=[1, 4, 4, 1], - filter_in_sizes=[2, 2, 1, 1], - strides=[1, 1], - dilations=[2, 1], - padding="VALID") + self._VerifyDilatedConvValues( + tensor_in_sizes=[1, 4, 4, 1], + filter_in_sizes=[2, 2, 1, 1], + strides=[1, 1], + dilations=[2, 1], + padding="VALID") @test_util.run_in_graph_and_eager_modes() def testConv2DEmpty(self): @@ -385,13 +379,12 @@ class Conv2DTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testConv2DEmptyDilation(self): - if test.is_gpu_available(cuda_only=True): - self._VerifyDilatedConvValues( - tensor_in_sizes=[0, 2, 3, 3], - filter_in_sizes=[1, 1, 3, 3], - strides=[1, 1], - dilations=[2, 1], - padding="VALID") + self._VerifyDilatedConvValues( + tensor_in_sizes=[0, 2, 3, 3], + filter_in_sizes=[1, 1, 3, 3], + strides=[1, 1], + dilations=[2, 1], + padding="VALID") @test_util.run_in_graph_and_eager_modes() def testConv2D2x2Filter(self): @@ -406,13 +399,12 @@ class Conv2DTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testConv2D2x2FilterDilation(self): - if test.is_gpu_available(cuda_only=True): - self._VerifyDilatedConvValues( - tensor_in_sizes=[1, 2, 3, 3], - filter_in_sizes=[2, 2, 3, 3], - strides=[1, 1], - dilations=[1, 2], - padding="VALID") + self._VerifyDilatedConvValues( + tensor_in_sizes=[1, 2, 3, 3], + filter_in_sizes=[2, 2, 3, 3], + strides=[1, 1], + dilations=[1, 2], + padding="VALID") @test_util.run_in_graph_and_eager_modes() def testConv2D1x2Filter(self): @@ -430,13 +422,12 @@ class Conv2DTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testConv2D1x2FilterDilation(self): - if test.is_gpu_available(cuda_only=True): - self._VerifyDilatedConvValues( - tensor_in_sizes=[1, 2, 3, 3], - filter_in_sizes=[1, 2, 3, 3], - strides=[1, 1], - dilations=[2, 1], - padding="VALID") + self._VerifyDilatedConvValues( + tensor_in_sizes=[1, 2, 3, 3], + filter_in_sizes=[1, 2, 3, 3], + strides=[1, 1], + dilations=[2, 1], + padding="VALID") @test_util.run_in_graph_and_eager_modes() def testConv2D2x2FilterStride2(self): @@ -512,15 +503,14 @@ class Conv2DTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testConv2DKernelSizeMatchesInputSizeDilation(self): - if test.is_gpu_available(cuda_only=True): - self._VerifyDilatedConvValues( - tensor_in_sizes=[1, 3, 3, 1], - filter_in_sizes=[2, 2, 1, 2], - strides=[1, 1], - dilations=[2, 2], - padding="VALID") + self._VerifyDilatedConvValues( + tensor_in_sizes=[1, 3, 3, 1], + filter_in_sizes=[2, 2, 1, 2], + strides=[1, 1], + dilations=[2, 2], + padding="VALID") - # TODO this currently fails. + # TODO(yzhwang): this currently fails. # self._VerifyValues(tensor_in_sizes=[1, 8, 8, 1], # filter_in_sizes=[2, 2, 1, 1], # strides=[4, 4], padding="SAME", @@ -980,7 +970,7 @@ class Conv2DTest(test.TestCase): self.assertArrayNear(value_2.flatten(), value.flatten(), err) def testConv2D2x2Depth3ValidBackpropFilterStride1x1Dilation2x1(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropFilterDilation( input_sizes=[1, 3, 6, 1], @@ -994,7 +984,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2D2x2Depth1ValidBackpropFilterDilation1x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropFilterDilation( input_sizes=[1, 2, 3, 1], @@ -1008,7 +998,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2DEmptyBackpropFilterDilation1x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropFilterDilation( input_sizes=[1, 2, 3, 1], @@ -1022,7 +1012,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2D2x2Depth3ValidBackpropFilterDilation2x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropFilterDilation( input_sizes=[1, 3, 4, 3], @@ -1036,7 +1026,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2DKernelSizeMatchesInputSizeBackpropFilterDilation2x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropFilterDilation( input_sizes=[1, 3, 3, 1], @@ -1050,7 +1040,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2D2x2Depth3ValidBackpropInputStride1x1Dilation2x1(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropInputDilation( input_sizes=[1, 3, 6, 1], @@ -1064,7 +1054,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2D2x2Depth1ValidBackpropInputDilation1x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropInputDilation( input_sizes=[1, 2, 3, 1], @@ -1078,7 +1068,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2DEmptyBackpropInputDilation1x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropInputDilation( input_sizes=[0, 2, 3, 1], @@ -1092,7 +1082,7 @@ class Conv2DTest(test.TestCase): err=1e-5) def testConv2D2x2Depth3ValidBackpropInputDilation2x1(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): # The GPU version of this test is not very stable. So adjusting the # error threshold to 1e-4. @@ -1108,7 +1098,7 @@ class Conv2DTest(test.TestCase): err=1e-4) def testConv2DKernelSizeMatchesInputSizeBackpropInputDilation2x2(self): - if test.is_gpu_available(cuda_only=True): + if test.is_gpu_available(cuda_only=True) or test_util.IsMklEnabled(): for (data_format, use_gpu) in GetTestConfigs(): self._RunAndVerifyBackpropInputDilation( input_sizes=[1, 3, 3, 1], @@ -1523,36 +1513,6 @@ class Conv2DTest(test.TestCase): strides=[1, 1, 1, 1], padding="VALID")) - def testCPUConv2DNCHWUnimplemented(self): - with self.test_session(use_gpu=False): - with self.assertRaisesRegexp(errors_impl.UnimplementedError, - "NHWC tensor format for now"): - conv = self._SetupValuesForDevice( - tensor_in_sizes=[1, 4, 4, 1], - filter_in_sizes=[2, 2, 1, 1], - dilations=[1, 1], - strides=[1, 1], - padding="VALID", - data_format="NCHW", - dtype=dtypes.float32, - use_gpu=False) - self.evaluate(conv) - - def testCPUConv2DDilatedUnimplemented(self): - with self.test_session(use_gpu=False): - with self.assertRaisesRegexp(errors_impl.UnimplementedError, - "dilated rate of 1 for now"): - conv = self._SetupValuesForDevice( - tensor_in_sizes=[1, 4, 4, 1], - filter_in_sizes=[2, 2, 1, 1], - dilations=[2, 1], - strides=[1, 1], - padding="VALID", - data_format="NHWC", - dtype=dtypes.float32, - use_gpu=False) - self.evaluate(conv) - class DepthwiseConv2DTest(test.TestCase): @@ -1887,7 +1847,7 @@ def GetInceptionFwdTest(input_size, filter_size, stride, padding, def GetInceptionFwdDilatedConvTest(input_size, filter_size, stride, padding): def Test(self): - if test.is_gpu_available(cuda_only=True) and stride == 1: + if stride == 1: tf_logging.info("Testing InceptionFwd with dilations %s", (input_size, filter_size, stride, padding)) self._VerifyDilatedConvValues( diff --git a/tensorflow/python/kernel_tests/cwise_ops_test.py b/tensorflow/python/kernel_tests/cwise_ops_test.py index 0d9b46c30dbbed20dd940e0427fbf6f6d5415106..8db0bb6f0dc495e7be2cd717787acf87156f42af 100644 --- a/tensorflow/python/kernel_tests/cwise_ops_test.py +++ b/tensorflow/python/kernel_tests/cwise_ops_test.py @@ -495,11 +495,11 @@ class UnaryOpTest(test.TestCase): dtype_tols = [(np.float32, 5e-4), (np.float64, 1e-6), (np.complex64, 5e-4), (np.complex128, 1e-6)] op_range = [ - (gen_math_ops._reciprocal_grad, [-2, 2]), - (gen_math_ops._rsqrt_grad, [0.1, 3]), - (gen_math_ops._sigmoid_grad, [-2, 2]), - (gen_math_ops._sqrt_grad, [0.1, 3]), - (gen_math_ops._tanh_grad, [-2, 2]), + (gen_math_ops.reciprocal_grad, [-2, 2]), + (gen_math_ops.rsqrt_grad, [0.1, 3]), + (gen_math_ops.sigmoid_grad, [-2, 2]), + (gen_math_ops.sqrt_grad, [0.1, 3]), + (gen_math_ops.tanh_grad, [-2, 2]), ] def rand(dtype): diff --git a/tensorflow/python/kernel_tests/decode_bmp_op_test.py b/tensorflow/python/kernel_tests/decode_bmp_op_test.py index c67c26b7be0777587eb6d7c49119ad6cd2e22953..35f8f76991a679e4164da4c63bacbe79fb5cd2c2 100644 --- a/tensorflow/python/kernel_tests/decode_bmp_op_test.py +++ b/tensorflow/python/kernel_tests/decode_bmp_op_test.py @@ -20,7 +20,6 @@ from __future__ import print_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors_impl from tensorflow.python.ops import array_ops from tensorflow.python.ops import image_ops from tensorflow.python.platform import test diff --git a/tensorflow/python/kernel_tests/decode_jpeg_op_test.py b/tensorflow/python/kernel_tests/decode_jpeg_op_test.py index 89fd26c544b5f2e8c15cec1b2d21a8c95fd503db..510daf79dc4252c3e2943e2ba23c1012370bf456 100644 --- a/tensorflow/python/kernel_tests/decode_jpeg_op_test.py +++ b/tensorflow/python/kernel_tests/decode_jpeg_op_test.py @@ -21,7 +21,7 @@ from __future__ import print_function import os import time -from six.moves import xrange +from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.python.client import session from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops diff --git a/tensorflow/python/kernel_tests/decode_raw_op_test.py b/tensorflow/python/kernel_tests/decode_raw_op_test.py index 0c7025f54e672bb09e601715a58864673a670d12..122a9ed46967fc9c02c59ea3047216cb73a72293 100644 --- a/tensorflow/python/kernel_tests/decode_raw_op_test.py +++ b/tensorflow/python/kernel_tests/decode_raw_op_test.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import numpy as np -import sys from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops diff --git a/tensorflow/python/kernel_tests/depthtospace_op_test.py b/tensorflow/python/kernel_tests/depthtospace_op_test.py index 7df2366954f3a6f3f37aef447479ba67c263025f..f0beabb4e20e4ec0a2fc7a487bf2541d19568927 100644 --- a/tensorflow/python/kernel_tests/depthtospace_op_test.py +++ b/tensorflow/python/kernel_tests/depthtospace_op_test.py @@ -35,8 +35,8 @@ from tensorflow.python.platform import tf_logging class DepthToSpaceTest(test.TestCase): - def _testOne(self, inputs, block_size, outputs): - input_nhwc = math_ops.to_float(inputs) + def _testOne(self, inputs, block_size, outputs, dtype=dtypes.float32): + input_nhwc = math_ops.cast(inputs, dtype) with self.test_session(use_gpu=False): # test NHWC (default) on CPU x_tf = array_ops.depth_to_space(input_nhwc, block_size) @@ -59,6 +59,12 @@ class DepthToSpaceTest(test.TestCase): x_out = [[[[1], [2]], [[3], [4]]]] self._testOne(x_np, block_size, x_out) + def testBasicFloat16(self): + x_np = [[[[1, 2, 3, 4]]]] + block_size = 2 + x_out = [[[[1], [2]], [[3], [4]]]] + self._testOne(x_np, block_size, x_out, dtype=dtypes.float16) + # Tests for larger input dimensions. To make sure elements are # correctly ordered spatially. def testBlockSize2(self): @@ -90,6 +96,24 @@ class DepthToSpaceTest(test.TestCase): x_out = [batch_output_elt(i) for i in range(batch_size)] self._testOne(x_np, block_size, x_out) + def testBatchSize0(self): + block_size = 2 + batch_size = 0 + input_nhwc = array_ops.ones([batch_size, 2, 3, 12]) + x_out = array_ops.ones([batch_size, 4, 6, 3]) + + with self.test_session(use_gpu=False): + # test NHWC (default) on CPU + x_tf = array_ops.depth_to_space(input_nhwc, block_size) + self.assertAllEqual(x_tf.shape, x_out.shape) + x_tf.eval() + if test.is_gpu_available(): + with self.test_session(use_gpu=True): + # test NHWC (default) on GPU + x_tf = array_ops.depth_to_space(input_nhwc, block_size) + self.assertAllEqual(x_tf.shape, x_out.shape) + x_tf.eval() + # Tests for different width and height. def testNonSquare(self): x_np = [[[[1, 10, 2, 20, 3, 30, 4, 40]], diff --git a/tensorflow/python/kernel_tests/determinant_op_test.py b/tensorflow/python/kernel_tests/determinant_op_test.py index 222038b22ef3c766efd14fd9b1c9044a0b6e9125..a52b2c0dc32c26ecd5ef08aa3f8678f0006cd4fe 100644 --- a/tensorflow/python/kernel_tests/determinant_op_test.py +++ b/tensorflow/python/kernel_tests/determinant_op_test.py @@ -65,7 +65,7 @@ class DeterminantOpTest(test.TestCase): self._compareDeterminantBase(matrix_x, linalg_ops.matrix_determinant(matrix_x)) self._compareLogDeterminantBase( - matrix_x, gen_linalg_ops._log_matrix_determinant(matrix_x)) + matrix_x, gen_linalg_ops.log_matrix_determinant(matrix_x)) def testBasic(self): # 2x2 matrices diff --git a/tensorflow/python/kernel_tests/distributions/bernoulli_test.py b/tensorflow/python/kernel_tests/distributions/bernoulli_test.py index a269d722737866fa5e6ae9feee919be0db71bcf1..09812db8166567403dc966ac9cb4304be0740e50 100644 --- a/tensorflow/python/kernel_tests/distributions/bernoulli_test.py +++ b/tensorflow/python/kernel_tests/distributions/bernoulli_test.py @@ -25,7 +25,6 @@ import numpy as np 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 math_ops from tensorflow.python.ops.distributions import bernoulli from tensorflow.python.ops.distributions import kullback_leibler from tensorflow.python.platform import test @@ -291,12 +290,6 @@ class BernoulliTest(test.TestCase): [np.sqrt(var(0.5)), np.sqrt(var(0.4))]], dtype=np.float32)) - def testBernoulliWithSigmoidProbs(self): - p = np.array([8.3, 4.2]) - dist = bernoulli.BernoulliWithSigmoidProbs(logits=p) - with self.test_session(): - self.assertAllClose(math_ops.sigmoid(p).eval(), dist.probs.eval()) - def testBernoulliBernoulliKL(self): with self.test_session() as sess: batch_size = 6 diff --git a/tensorflow/python/kernel_tests/distributions/beta_test.py b/tensorflow/python/kernel_tests/distributions/beta_test.py index 91a451f033ffbb01d54c3dacce952b406564b7b4..ab5041a6eb477ce231acbd1e6041c354ee17409b 100644 --- a/tensorflow/python/kernel_tests/distributions/beta_test.py +++ b/tensorflow/python/kernel_tests/distributions/beta_test.py @@ -107,8 +107,10 @@ class BetaTest(test.TestCase): 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`"): + with self.assertRaisesOpError("sample must be less than `1`"): dist.prob([.1, .2, 1.2]).eval() + with self.assertRaisesOpError("sample must be less than `1`"): + dist.prob([.1, .2, 1.0]).eval() def testPdfTwoBatches(self): with self.test_session(): diff --git a/tensorflow/python/kernel_tests/dynamic_partition_op_test.py b/tensorflow/python/kernel_tests/dynamic_partition_op_test.py index fedbf9e696923a34968e7a907e4099c520d1447b..5e8937ad2c36afb2b1ddb58ffb238a45e09e4b30 100644 --- a/tensorflow/python/kernel_tests/dynamic_partition_op_test.py +++ b/tensorflow/python/kernel_tests/dynamic_partition_op_test.py @@ -326,6 +326,18 @@ class DynamicPartitionTest(test.TestCase): with self.assertRaises(ValueError): data_flow_ops.dynamic_partition(data, indices, num_partitions=4) + # see https://github.com/tensorflow/tensorflow/issues/17106 + def testCUBBug(self): + x = constant_op.constant(np.random.randn(3072)) + inds = [0]*189 + [1]*184 + [2]*184 + [3]*191 + [4]*192 + [5]*195 + [6]*195 + inds += [7]*195 + [8]*188 + [9]*195 + [10]*188 + [11]*202 + [12]*194 + inds += [13]*194 + [14]*194 + [15]*192 + self.assertEqual(len(inds), x.shape[0]) + partitioned = data_flow_ops.dynamic_partition(x, inds, 16) + with self.test_session() as sess: + res = sess.run(partitioned) + self.assertEqual(res[-1].shape[0], 192) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/extract_image_patches_op_test.py b/tensorflow/python/kernel_tests/extract_image_patches_op_test.py index 5c7624f1f6be4da91ca74d4ef2ed81a21890b35c..6ea9f1badc3b8fac06fe6328f95714b93de97c0e 100644 --- a/tensorflow/python/kernel_tests/extract_image_patches_op_test.py +++ b/tensorflow/python/kernel_tests/extract_image_patches_op_test.py @@ -84,7 +84,7 @@ class ExtractImagePatches(test.TestCase): patches=patches) def testKsize2x2Stride1x1Rate1x1Valid(self): - """Test for 1x1 kernel .""" + """Test for 2x2 kernel with VALID padding.""" # [1, 2, 2, 1] image = [[[[1], [2]], [[3], [4]]]] # [1, 1, 1, 4] @@ -98,7 +98,7 @@ class ExtractImagePatches(test.TestCase): patches=patches) def testKsize2x2Stride1x1Rate1x1Same(self): - """Test for 1x1 kernel .""" + """Test for 2x2 kernel with SAME padding.""" # [1, 2, 2, 1] image = [[[[1], [2]], [[3], [4]]]] # [1, 2, 2, 4] @@ -111,6 +111,20 @@ class ExtractImagePatches(test.TestCase): 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/python/kernel_tests/fifo_queue_test.py b/tensorflow/python/kernel_tests/fifo_queue_test.py index 748135440ec5e8ad387f910e1433f638abf2260a..ce73e7ad3e5f822363c697609dfa163b6f13751a 100644 --- a/tensorflow/python/kernel_tests/fifo_queue_test.py +++ b/tensorflow/python/kernel_tests/fifo_queue_test.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import random -import re import time import numpy as np diff --git a/tensorflow/python/kernel_tests/fractional_avg_pool_op_test.py b/tensorflow/python/kernel_tests/fractional_avg_pool_op_test.py index feec9934e459590bb1dd0bc5c7cf40013d3d8b88..faac7d8365dfaa1b6b32f8fe66a76c3694aa0d5b 100644 --- a/tensorflow/python/kernel_tests/fractional_avg_pool_op_test.py +++ b/tensorflow/python/kernel_tests/fractional_avg_pool_op_test.py @@ -347,7 +347,7 @@ class FractionalAvgPoolGradTest(test.TestCase): Two types of tests for FractionalAvgPoolGrad. 1) Test fractional_avg_pool_grad() directly. - This type of test relies on gen_nn_ops._avg_pool_grad() returns the + This type of test relies on gen_nn_ops.avg_pool_grad() returns the correct result. For example: * input_tensor_shape = (1, 10, 10, 1) * window_size = (1, 2, 2, 1) @@ -404,13 +404,13 @@ class FractionalAvgPoolGradTest(test.TestCase): num_elements *= dim_size output_backprop = (self._PRNG.rand(num_elements) * 1000).reshape(output_data.shape) - input_backprop_tensor = gen_nn_ops._avg_pool_grad( + input_backprop_tensor = gen_nn_ops.avg_pool_grad( input_tensor.get_shape(), output_backprop, window_size, stride_size, padding) input_backprop = input_backprop_tensor.eval() row_seq = list(range(0, num_rows + 1, row_window_size)) col_seq = list(range(0, num_cols + 1, col_window_size)) - fap_input_backprop_tensor = gen_nn_ops._fractional_avg_pool_grad( + fap_input_backprop_tensor = gen_nn_ops.fractional_avg_pool_grad( input_tensor.get_shape(), output_backprop, row_seq, @@ -443,7 +443,7 @@ class FractionalAvgPoolGradTest(test.TestCase): num_elements *= dim_size output_backprop = (self._PRNG.rand(num_elements) * 1000).reshape(output_data.shape) - input_backprop_tensor = gen_nn_ops._avg_pool_grad( + input_backprop_tensor = gen_nn_ops.avg_pool_grad( input_tensor.get_shape(), output_backprop, window_size, stride_size, padding) input_backprop = input_backprop_tensor.eval() @@ -451,7 +451,7 @@ class FractionalAvgPoolGradTest(test.TestCase): col_seq = list(range(0, num_cols, col_window_size - 1)) row_seq[-1] += 1 col_seq[-1] += 1 - fap_input_backprop_tensor = gen_nn_ops._fractional_avg_pool_grad( + fap_input_backprop_tensor = gen_nn_ops.fractional_avg_pool_grad( input_tensor.get_shape(), output_backprop, row_seq, diff --git a/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py b/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py index 5983ae7759dbf3eb2db9867def829ce8dbeb4b73..6477c9ebc4c35fcc5963b27a0f5c50624a73fa09 100644 --- a/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py +++ b/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py @@ -318,7 +318,7 @@ class FractionalMaxPoolGradTest(test.TestCase): Two types of tests for FractionalMaxPoolGrad. 1) Test fractional_max_pool_grad() directly. - This type of test relies on gen_nn_ops._max_pool_grad() returns the correct + This type of test relies on gen_nn_ops.max_pool_grad() returns the correct result. For example: * input_tensor_shape = (1, 10, 10, 1) * window_size = (1, 2, 2, 1) @@ -384,16 +384,13 @@ class FractionalMaxPoolGradTest(test.TestCase): stride_size, padding) output_data = output_tensor.eval() output_backprop = self._PRNG.randint(100, size=output_data.shape) - input_backprop_tensor = gen_nn_ops._max_pool_grad(input_tensor, - output_tensor, - output_backprop, - window_size, - stride_size, - padding) + input_backprop_tensor = gen_nn_ops.max_pool_grad( + input_tensor, output_tensor, output_backprop, window_size, + stride_size, padding) input_backprop = input_backprop_tensor.eval() row_seq = list(range(0, num_rows + 1, row_window_size)) col_seq = list(range(0, num_cols + 1, col_window_size)) - fmp_input_backprop_tensor = gen_nn_ops._fractional_max_pool_grad( + fmp_input_backprop_tensor = gen_nn_ops.fractional_max_pool_grad( input_tensor, output_tensor, output_backprop, @@ -422,18 +419,15 @@ class FractionalMaxPoolGradTest(test.TestCase): stride_size, padding) output_data = output_tensor.eval() output_backprop = self._PRNG.randint(100, size=output_data.shape) - input_backprop_tensor = gen_nn_ops._max_pool_grad(input_tensor, - output_tensor, - output_backprop, - window_size, - stride_size, - padding) + input_backprop_tensor = gen_nn_ops.max_pool_grad( + input_tensor, output_tensor, output_backprop, window_size, + stride_size, padding) input_backprop = input_backprop_tensor.eval() row_seq = list(range(0, num_rows, row_window_size - 1)) col_seq = list(range(0, num_cols, col_window_size - 1)) row_seq[-1] += 1 col_seq[-1] += 1 - fmp_input_backprop_tensor = gen_nn_ops._fractional_max_pool_grad( + fmp_input_backprop_tensor = gen_nn_ops.fractional_max_pool_grad( input_tensor, output_tensor, output_backprop, @@ -591,7 +585,7 @@ class FractionalMaxPoolGradTest(test.TestCase): output_tensor = constant_op.constant( output_data_not_overlapping, shape=output_size) grad = constant_op.constant(output_backprop, shape=output_size) - r = gen_nn_ops._fractional_max_pool_grad( + r = gen_nn_ops.fractional_max_pool_grad( input_tensor, output_tensor, grad, @@ -606,7 +600,7 @@ class FractionalMaxPoolGradTest(test.TestCase): # Test when overlapping is True output_tensor = constant_op.constant( output_data_overlapping, shape=output_size) - r = gen_nn_ops._fractional_max_pool_grad( + r = gen_nn_ops.fractional_max_pool_grad( input_tensor, output_tensor, grad, row_seq, col_seq, overlapping=True) input_backprop_overlapping = r.eval() self.assertShapeEqual( diff --git a/tensorflow/python/kernel_tests/identity_op_py_test.py b/tensorflow/python/kernel_tests/identity_op_py_test.py index 2cfe420bd49ec44815d1386bd873b234d8710e9d..49fb76d5b41de18ed3ba2187e85cb288e7344c38 100644 --- a/tensorflow/python/kernel_tests/identity_op_py_test.py +++ b/tensorflow/python/kernel_tests/identity_op_py_test.py @@ -65,7 +65,7 @@ class IdentityOpTest(test.TestCase): constant_op.constant( [[1, 2, 3], [6, 5, 4]], dtype=dtypes.int32)) self.assertEquals(shape, tensor.get_shape()) - self.assertEquals(shape, gen_array_ops._ref_identity(tensor).get_shape()) + self.assertEquals(shape, gen_array_ops.ref_identity(tensor).get_shape()) if __name__ == "__main__": diff --git a/tensorflow/python/kernel_tests/init_ops_test.py b/tensorflow/python/kernel_tests/init_ops_test.py index 19a7d2f9d51fff46ee817ad03ef62383f6727791..c1755985ee85c62005c8d3d5fb916859193aa5f3 100644 --- a/tensorflow/python/kernel_tests/init_ops_test.py +++ b/tensorflow/python/kernel_tests/init_ops_test.py @@ -25,10 +25,13 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed +from tensorflow.python.layers import convolutional from tensorflow.python.ops import array_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 partitioned_variables +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 @@ -571,6 +574,82 @@ class OrthogonalInitializerTest(test.TestCase): np.dot(t, t.T), np.eye(t.shape[0]), rtol=tol, atol=tol) +class ConvolutionDeltaOrthogonalInitializerTest(test.TestCase): + + def testInitializerIdentical(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_delta_orthogonal(seed=1, dtype=dtype) + init2 = init_ops.convolutional_delta_orthogonal(seed=1, dtype=dtype) + self.assertTrue(identicaltest(self, init1, init2, (3, 3, 10, 10))) + + def testInitializerDifferent(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_delta_orthogonal(seed=1, dtype=dtype) + init2 = init_ops.convolutional_delta_orthogonal(seed=2, dtype=dtype) + self.assertFalse(identicaltest(self, init1, init2, (3, 3, 10, 10))) + + def testDuplicatedInitializer(self): + init = init_ops.convolutional_delta_orthogonal() + self.assertFalse(duplicated_initializer(self, init, 1, (3, 3, 10, 10))) + + def testInvalidDataType(self): + self.assertRaises( + ValueError, init_ops.convolutional_delta_orthogonal, + dtype=dtypes.string) + + def testInvalidShape(self): + init1 = init_ops.convolutional_delta_orthogonal() + with self.test_session(graph=ops.Graph(), use_gpu=True): + self.assertRaises(ValueError, init1, shape=[3, 3, 6, 5]) + + def testGain(self): + shape = (3, 3, 10, 10) + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_delta_orthogonal(seed=1, dtype=dtype) + init2 = init_ops.convolutional_delta_orthogonal(gain=3.14, + seed=1, dtype=dtype) + with self.test_session(graph=ops.Graph(), use_gpu=True): + t1 = init1(shape).eval() + with self.test_session(graph=ops.Graph(), use_gpu=True): + t2 = init2(shape).eval() + return np.allclose(t1, t2 / 3.14, rtol=1e-15, atol=1e-15) + + def testShapesValues(self): + for dtype in [dtypes.float32]: + for kernel_size in [[3], [8], [3, 5], [2, 4], [3, 3, 3], [2, 2, 2]]: + tol = 1e-2 + # Check orthogonality by computing the 2-norms of the inputs and ouputs. + if len(kernel_size) == 1: + shape = [4, 32, 64] + convolution = convolutional.conv1d + elif len(kernel_size) == 2: + convolution = convolutional.conv2d + shape = [4, 32, 32, 64] + else: + shape = [4, 16, 16, 16, 64] + convolution = convolutional.conv3d + inputs = random_ops.random_normal(shape, dtype=dtype) + inputs_2norm = linalg_ops.norm(inputs) + outputs = convolution( + inputs, padding="same", filters=128, + kernel_size=kernel_size, use_bias=False, + kernel_initializer=init_ops.convolutional_delta_orthogonal( + gain=3.14)) + outputs_shape = shape[0:-1] + [128] + outputs_2norm = linalg_ops.norm(outputs) + my_ops = variables.global_variables_initializer() + with self.test_session(use_gpu=True) as sess: + sess.run(my_ops) + # Check the shape of the outputs + t = outputs.eval() + self.assertAllEqual(t.shape, outputs_shape) + # Check isometry of the delta-orthogonal kernel. + self.assertAllClose( + sess.run(inputs_2norm)/np.sqrt(np.prod(shape)), + sess.run(outputs_2norm)/(np.sqrt(np.prod(shape))*np.sqrt(3.14)), + rtol=tol, atol=tol) + + class IdentityInitializerTest(test.TestCase): def testInvalidDataType(self): diff --git a/tensorflow/python/kernel_tests/io_ops_test.py b/tensorflow/python/kernel_tests/io_ops_test.py index f91875c6f0c1a7bfa388ec1b1a58f06b65889c3e..61944f7e3197844d00cbc001459e48b50c9003b4 100644 --- a/tensorflow/python/kernel_tests/io_ops_test.py +++ b/tensorflow/python/kernel_tests/io_ops_test.py @@ -1,4 +1,4 @@ -# -*- coding: utf-8 -*- +# -*- coding: utf-8 -*- # Copyright 2015 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD index 4e18eaa4e8281c799e4669b2d6083c00bc1e2863..fd1b5bab6f5aa072c8821eb053bd8d39391be4d4 100644 --- a/tensorflow/python/kernel_tests/linalg/BUILD +++ b/tensorflow/python/kernel_tests/linalg/BUILD @@ -39,6 +39,7 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], + shard_count = 5, tags = ["noasan"], # times out b/63678675 ) @@ -57,6 +58,7 @@ cuda_py_test( "//tensorflow/python:platform_test", "//tensorflow/python:random_ops", ], + shard_count = 5, ) cuda_py_test( @@ -73,6 +75,7 @@ cuda_py_test( "//tensorflow/python:platform_test", "//tensorflow/python:random_ops", ], + shard_count = 5, ) cuda_py_test( @@ -88,6 +91,7 @@ cuda_py_test( "//tensorflow/python:framework_test_lib", "//tensorflow/python:platform_test", ], + shard_count = 5, ) cuda_py_test( @@ -134,6 +138,7 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], + shard_count = 5, ) filegroup( diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_composition_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_composition_test.py index 4d79365dbefc74fe8412b65ec089fb2af4255aea..f96b9ccdaacae7d8e0552ed3d74ce53808fed963 100644 --- a/tensorflow/python/kernel_tests/linalg/linear_operator_composition_test.py +++ b/tensorflow/python/kernel_tests/linalg/linear_operator_composition_test.py @@ -44,9 +44,9 @@ class SquareLinearOperatorCompositionTest( self._rtol[dtypes.float32] = 1e-4 self._rtol[dtypes.complex64] = 1e-4 - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): sess = ops.get_default_session() - shape = list(shape) + shape = list(build_info.shape) # Either 1 or 2 matrices, depending. num_operators = rng.randint(low=1, high=3) @@ -148,9 +148,9 @@ class NonSquareLinearOperatorCompositionTest( self._rtol[dtypes.float32] = 1e-4 self._rtol[dtypes.complex64] = 1e-4 - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): sess = ops.get_default_session() - shape = list(shape) + shape = list(build_info.shape) # Test only the case of 2 matrices. # The Square test uses either 1 or 2, so we have tested the case of 1 matrix diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_diag_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_diag_test.py index 343d158498833dd92361bc41d433e28296fc4c9a..0a0e31c716ecfa10ed93cff92fa908a240f8495e 100644 --- a/tensorflow/python/kernel_tests/linalg/linear_operator_diag_test.py +++ b/tensorflow/python/kernel_tests/linalg/linear_operator_diag_test.py @@ -34,7 +34,8 @@ class LinearOperatorDiagTest( linear_operator_test_util.SquareLinearOperatorDerivedClassTest): """Most tests done in the base class LinearOperatorDerivedClassTest.""" - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) diag = linear_operator_test_util.random_sign_uniform( shape[:-1], minval=1., maxval=2., dtype=dtype) if use_placeholder: @@ -129,7 +130,7 @@ class LinearOperatorDiagTest( with self.test_session() as sess: x = random_ops.random_normal(shape=(2, 2, 3, 4)) - # This LinearOperatorDiag will be brodacast to (2, 2, 3, 3) during solve + # This LinearOperatorDiag will be broadcast to (2, 2, 3, 3) during solve # and matmul with 'x' as the argument. diag = random_ops.random_uniform(shape=(2, 1, 3)) operator = linalg.LinearOperatorDiag(diag, is_self_adjoint=True) diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_full_matrix_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_full_matrix_test.py index 50d6f524e9ad75715d7a57348638fdfeee667f40..b3da623b5e8d8c99c6777e75e2d49f24dab1c96b 100644 --- a/tensorflow/python/kernel_tests/linalg/linear_operator_full_matrix_test.py +++ b/tensorflow/python/kernel_tests/linalg/linear_operator_full_matrix_test.py @@ -36,11 +36,11 @@ class SquareLinearOperatorFullMatrixTest( linear_operator_test_util.SquareLinearOperatorDerivedClassTest): """Most tests done in the base class LinearOperatorDerivedClassTest.""" - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): - shape = list(shape) + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) - matrix = linear_operator_test_util.random_positive_definite_matrix(shape, - dtype) + matrix = linear_operator_test_util.random_positive_definite_matrix( + shape, dtype) if use_placeholder: matrix_ph = array_ops.placeholder(dtype=dtype) @@ -136,8 +136,8 @@ class SquareLinearOperatorFullMatrixSymmetricPositiveDefiniteTest( def _dtypes_to_test(self): return [dtypes.float32, dtypes.float64] - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): - shape = list(shape) + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) matrix = linear_operator_test_util.random_positive_definite_matrix( shape, dtype, force_well_conditioned=True) @@ -210,7 +210,8 @@ class NonSquareLinearOperatorFullMatrixTest( linear_operator_test_util.NonSquareLinearOperatorDerivedClassTest): """Most tests done in the base class LinearOperatorDerivedClassTest.""" - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) matrix = linear_operator_test_util.random_normal(shape, dtype=dtype) if use_placeholder: matrix_ph = array_ops.placeholder(dtype=dtype) diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_identity_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_identity_test.py index 6d635707683f4500919073a4f43c320a44b65018..59f63f949e96991193412d3574603e58a75cb6e5 100644 --- a/tensorflow/python/kernel_tests/linalg/linear_operator_identity_test.py +++ b/tensorflow/python/kernel_tests/linalg/linear_operator_identity_test.py @@ -43,8 +43,8 @@ class LinearOperatorIdentityTest( # 16bit. return [dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128] - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): - shape = list(shape) + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) assert shape[-1] == shape[-2] batch_shape = shape[:-2] @@ -261,8 +261,8 @@ class LinearOperatorScaledIdentityTest( # 16bit. return [dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128] - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): - shape = list(shape) + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) assert shape[-1] == shape[-2] batch_shape = shape[:-2] diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_low_rank_update_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_low_rank_update_test.py index d3a47da946b12277c4c390a4a320d7c91ed81b32..8095f6419ef0d9543339cf1f4ee9cd4783f852b9 100644 --- a/tensorflow/python/kernel_tests/linalg/linear_operator_low_rank_update_test.py +++ b/tensorflow/python/kernel_tests/linalg/linear_operator_low_rank_update_test.py @@ -55,16 +55,22 @@ class BaseLinearOperatorLowRankUpdatetest(object): return [dtypes.float32, dtypes.float64] @property - def _shapes_to_test(self): + def _operator_build_infos(self): + build_info = linear_operator_test_util.OperatorBuildInfo # Previously we had a (2, 10, 10) shape at the end. We did this to test the # inversion and determinant lemmas on not-tiny matrices, since these are # known to have stability issues. This resulted in test timeouts, so this # shape has been removed, but rest assured, the tests did pass. - return [(0, 0), (1, 1), (1, 3, 3), (3, 4, 4), (2, 1, 4, 4)] - - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + return [ + build_info((0, 0)), + build_info((1, 1)), + build_info((1, 3, 3)), + build_info((3, 4, 4)), + build_info((2, 1, 4, 4))] + + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): # Recall A = L + UDV^H - shape = list(shape) + shape = list(build_info.shape) diag_shape = shape[:-1] k = shape[-2] // 2 + 1 u_perturbation_shape = shape[:-1] + [k] diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_lower_triangular_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_lower_triangular_test.py index db3918f9983c5b7d05fa4ba3bc85b26a485f2f00..a57d2f085e089fb913f09fdd9b07cf13aa7f3c35 100644 --- a/tensorflow/python/kernel_tests/linalg/linear_operator_lower_triangular_test.py +++ b/tensorflow/python/kernel_tests/linalg/linear_operator_lower_triangular_test.py @@ -38,7 +38,8 @@ class LinearOperatorLowerTriangularTest( # matrix_triangular_solve. return [dtypes.float32, dtypes.float64] - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) # Upper triangle will be nonzero, but ignored. # Use a diagonal that ensures this matrix is well conditioned. tril = linear_operator_test_util.random_tril_matrix( diff --git a/tensorflow/python/kernel_tests/list_ops_test.py b/tensorflow/python/kernel_tests/list_ops_test.py index 1577b7bc8021a326eb720bdf059b8d1c568c0cc1..8865e165fd6bb15b53d10f336e526f8e830500b1 100644 --- a/tensorflow/python/kernel_tests/list_ops_test.py +++ b/tensorflow/python/kernel_tests/list_ops_test.py @@ -30,7 +30,9 @@ from tensorflow.python.framework import errors 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 control_flow_ops from tensorflow.python.ops import list_ops +from tensorflow.python.ops import math_ops from tensorflow.python.platform import test from tensorflow.python.training import server_lib @@ -123,6 +125,56 @@ class ListOpsTest(test_util.TensorFlowTestCase): l_cpu, element_dtype=dtypes.float32)[1], 2.0) + def testGraphStack(self): + with context.graph_mode(), self.test_session(): + tl = list_ops.empty_tensor_list( + element_shape=constant_op.constant([1], dtype=dtypes.int32), + element_dtype=dtypes.int32) + tl = list_ops.tensor_list_push_back(tl, [1]) + self.assertAllEqual( + list_ops.tensor_list_stack(tl, element_dtype=dtypes.int32).eval(), + [[1]]) + + def testGraphStackInLoop(self): + with context.graph_mode(), self.test_session(): + t1 = list_ops.empty_tensor_list( + element_shape=constant_op.constant([], dtype=dtypes.int32), + element_dtype=dtypes.int32) + i = constant_op.constant(0, dtype=dtypes.int32) + + def body(i, t1): + t1 = list_ops.tensor_list_push_back(t1, i) + i += 1 + return i, t1 + + i, t1 = control_flow_ops.while_loop(lambda i, t1: math_ops.less(i, 4), + body, [i, t1]) + s1 = list_ops.tensor_list_stack(t1, element_dtype=dtypes.int32).eval() + self.assertAllEqual(s1, [0, 1, 2, 3]) + + def testGraphStackInLoopSwitchDtype(self): + with context.graph_mode(), self.test_session(): + t1 = list_ops.empty_tensor_list( + element_shape=constant_op.constant([], dtype=dtypes.int32), + element_dtype=dtypes.int32) + i = constant_op.constant(0, dtype=dtypes.float32) + m = constant_op.constant([1, 2, 3], dtype=dtypes.float32) + + def body(i, m, t1): + t1 = control_flow_ops.cond( + math_ops.equal(list_ops.tensor_list_length(t1), 0), + lambda: list_ops.empty_tensor_list(m.shape, m.dtype), lambda: t1) + + t1 = list_ops.tensor_list_push_back(t1, m * i) + i += 1.0 + return i, m, t1 + + i, m, t1 = control_flow_ops.while_loop( + lambda i, m, t1: math_ops.less(i, 4), body, [i, m, t1]) + s1 = list_ops.tensor_list_stack(t1, element_dtype=dtypes.float32).eval() + np_s1 = np.vstack([np.arange(1, 4) * i for i in range(4)]) + self.assertAllEqual(s1, np_s1) + def testSerialize(self): # pylint: disable=g-import-not-at-top try: diff --git a/tensorflow/python/kernel_tests/losses_test.py b/tensorflow/python/kernel_tests/losses_test.py index 81af3a0887d09a7736a145a5b3c99c9391691724..1123c20a165ba93bd380fa471a8be91f7005d7bb 100644 --- a/tensorflow/python/kernel_tests/losses_test.py +++ b/tensorflow/python/kernel_tests/losses_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors_impl @@ -32,11 +33,25 @@ from tensorflow.python.ops import random_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.ops.losses import losses +from tensorflow.python.ops.losses import losses_impl from tensorflow.python.ops.losses import util from tensorflow.python.platform import test from tensorflow.python.training import momentum as momentum_lib +safe_div = losses_impl._safe_div # pylint: disable=protected-access + + +class SafeDivTest(test.TestCase): + + def testEager(self): + with context.eager_mode(): + self.assertAllEqual(safe_div(constant_op.constant(1.0), + constant_op.constant(0.0)), 0.0) + self.assertAllEqual(safe_div(constant_op.constant(1.0), + 0.0), 0.0) + + class AbsoluteDifferenceLossTest(test.TestCase): def setUp(self): @@ -953,14 +968,14 @@ class MeanPairwiseSquaredErrorTest(test.TestCase): # Compute the expected loss 'manually'. total = np.zeros((batch_size,)) for b in range(batch_size): - for i in range(dims): - for j in range(dims): + for i in range(dims - 1): + for j in range(i + 1, dims): x = self._predictions[b, i].item() - self._predictions[b, j].item() y = self._labels[b, i].item() - self._labels[b, j].item() diff = (x - y) total[b] += (diff * diff) - self._expected_losses = np.divide(total, 9.0) + self._expected_losses = np.divide(total, 3.0) def testValueErrorThrownWhenWeightIsNone(self): with self.test_session(): @@ -1059,8 +1074,7 @@ class MeanPairwiseSquaredErrorTest(test.TestCase): [[4, 8, 12], [1, 2, 3], [4, 5, 6]], [[8, 1, 3], [7, 8, 9], [10, 11, 12]], ]) - self._test_valid_weights( - labels, predictions, expected_loss=122.22222) + self._test_valid_weights(labels, predictions, expected_loss=137.5) def test3dWeightedScalar(self): labels = np.array([ @@ -1073,8 +1087,7 @@ class MeanPairwiseSquaredErrorTest(test.TestCase): ]) weight = 3.0 self._test_valid_weights( - labels, predictions, expected_loss=weight * 122.22222, - weights=weight) + labels, predictions, expected_loss=weight * 137.5, weights=weight) def _test_invalid_weights( self, labels, predictions, weights=1.0): @@ -1124,7 +1137,9 @@ class MeanPairwiseSquaredErrorTest(test.TestCase): ]) self._test_valid_weights( # TODO(ptucker): This doesn't look right. - labels, predictions, expected_loss=9 * 122.22222, + labels, + predictions, + expected_loss=9 * 137.5, weights=np.ones((2, 3, 3))) def testLossWithAllZeroBatchSpecificWeights(self): @@ -1345,6 +1360,34 @@ class ComputeWeightedLossTest(test.TestCase): self.assertAllClose( np.mean(self._raw_losses), unweighted_loss.eval()) + def testUnweightedFromPlaceholder(self): + for reduction in losses.Reduction.all(): + with ops.Graph().as_default() as g: + self.assertEqual(0, len(util.get_losses())) + raw_losses = array_ops.placeholder(dtype=dtypes.float32) + feed_dict = {raw_losses: self._raw_losses} + unweighted_losses = ( + losses.compute_weighted_loss(raw_losses, reduction=reduction), + losses.compute_weighted_loss( + raw_losses, weights=np.ones((1, 1, 1)), reduction=reduction), + losses.compute_weighted_loss( + raw_losses, weights=np.ones((1, 1, 4)), reduction=reduction), + ) + self.assertEqual(3, len(util.get_losses())) + with self.test_session(g): + for unweighted_loss in unweighted_losses: + if reduction == losses.Reduction.NONE: + self.assertAllClose( + self._raw_losses, unweighted_loss.eval(feed_dict)) + elif reduction == losses.Reduction.SUM: + self.assertAllClose( + np.sum(self._raw_losses), unweighted_loss.eval(feed_dict)) + else: + # reduction one of MEAN, SUM_OVER_NONZERO_WEIGHTS, + # SUM_BY_NONZERO_WEIGHTS or SUM_OVER_BATCH_SIZE. + self.assertAllClose( + np.mean(self._raw_losses), unweighted_loss.eval(feed_dict)) + def testScalarWeight(self): with ops.Graph().as_default(): self.assertEqual(0, len(util.get_losses())) diff --git a/tensorflow/python/kernel_tests/manip_ops_test.py b/tensorflow/python/kernel_tests/manip_ops_test.py index 3044b21aa4271cabda2a174a85b6db16fe13f93b..b8200ac0cb1e4315a56181779c70da1126d8fc15 100644 --- a/tensorflow/python/kernel_tests/manip_ops_test.py +++ b/tensorflow/python/kernel_tests/manip_ops_test.py @@ -17,25 +17,27 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + from tensorflow.python.framework import constant_op from tensorflow.python.framework import errors_impl from tensorflow.python.framework import test_util -from tensorflow.python.ops import manip_ops from tensorflow.python.ops import gradient_checker +from tensorflow.python.ops import manip_ops from tensorflow.python.platform import test as test_lib -import numpy as np - # pylint: disable=g-import-not-at-top try: from distutils.version import StrictVersion as Version # numpy.roll for multiple shifts was introduced in numpy version 1.12.0 - NP_ROLL_CAN_MULTISHIFT = Version(np.version.version) >= Version('1.12.0') + NP_ROLL_CAN_MULTISHIFT = Version(np.version.version) >= Version("1.12.0") except ImportError: NP_ROLL_CAN_MULTISHIFT = False # pylint: enable=g-import-not-at-top + class RollTest(test_util.TensorFlowTestCase): + def _testRoll(self, np_input, shift, axis): expected_roll = np.roll(np_input, shift, axis) with self.test_session(): @@ -62,10 +64,12 @@ class RollTest(test_util.TensorFlowTestCase): for t in [np.int32, np.int64]: self._testAll(np.random.randint(-100, 100, (5)).astype(t), 3, 0) if NP_ROLL_CAN_MULTISHIFT: - self._testAll(np.random.randint(-100, 100, (4, 4, 3)).astype(t), - [1, -2, 3], [0, 1, 2]) - self._testAll(np.random.randint(-100, 100, (4, 2, 1, 3)).astype(t), - [0, 1, -2], [1, 2, 3]) + self._testAll( + np.random.randint(-100, 100, (4, 4, 3)).astype(t), [1, -2, 3], + [0, 1, 2]) + self._testAll( + np.random.randint(-100, 100, (4, 2, 1, 3)).astype(t), [0, 1, -2], + [1, 2, 3]) def testFloatTypes(self): for t in [np.float32, np.float64]: @@ -84,7 +88,6 @@ class RollTest(test_util.TensorFlowTestCase): x = np.random.rand(3, 2, 1, 1).astype(t) self._testAll(x + 1j * x, [2, 1, 1, 0], [0, 3, 1, 2]) - def testRollInputMustVectorHigherRaises(self): tensor = 7 shift = 1 @@ -95,8 +98,7 @@ class RollTest(test_util.TensorFlowTestCase): manip_ops.roll(tensor, shift, axis).eval() def testRollAxisMustBeScalarOrVectorRaises(self): - tensor = [[1, 2], - [3, 4]] + tensor = [[1, 2], [3, 4]] shift = 1 axis = [[0, 1]] with self.test_session(): @@ -105,8 +107,7 @@ class RollTest(test_util.TensorFlowTestCase): manip_ops.roll(tensor, shift, axis).eval() def testRollShiftMustBeScalarOrVectorRaises(self): - tensor = [[1, 2], - [3, 4]] + tensor = [[1, 2], [3, 4]] shift = [[0, 1]] axis = 1 with self.test_session(): @@ -115,8 +116,7 @@ class RollTest(test_util.TensorFlowTestCase): manip_ops.roll(tensor, shift, axis).eval() def testRollShiftAndAxisMustBeSameSizeRaises(self): - tensor = [[1, 2], - [3, 4]] + tensor = [[1, 2], [3, 4]] shift = [1] axis = [0, 1] with self.test_session(): @@ -133,5 +133,6 @@ class RollTest(test_util.TensorFlowTestCase): "is out of range"): manip_ops.roll(tensor, shift, axis).eval() + if __name__ == "__main__": test_lib.main() diff --git a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py index 317b8dc05beac7642c384bf89e6d154be50f6992..68d626de2c5cdd91ee332247c05ddce2a558a35e 100644 --- a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py +++ b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py @@ -21,6 +21,7 @@ import numpy as np from tensorflow.python.client import session 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.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -54,9 +55,13 @@ def _GetMatrixBandPartTest(dtype_, batch_shape_, shape_): band_np = np.tril(band_np, upper) if batch_shape_ is not (): band_np = np.tile(band_np, batch_shape_ + (1, 1)) - with self.test_session(use_gpu=False): - band = array_ops.matrix_band_part(batch_mat, lower, upper) - self.assertAllEqual(band_np, band.eval()) + for index_dtype in [dtypes_lib.int32, dtypes_lib.int64]: + with self.test_session(use_gpu=False): + band = array_ops.matrix_band_part( + batch_mat, + constant_op.constant(lower, index_dtype), + constant_op.constant(upper, index_dtype)) + self.assertAllEqual(band_np, band.eval()) return Test diff --git a/tensorflow/python/kernel_tests/matrix_exponential_op_test.py b/tensorflow/python/kernel_tests/matrix_exponential_op_test.py index 6203a412d7faec4fe9f6179141301579b5900291..a0c66c77d8850d3144678870983730537a253556 100644 --- a/tensorflow/python/kernel_tests/matrix_exponential_op_test.py +++ b/tensorflow/python/kernel_tests/matrix_exponential_op_test.py @@ -48,7 +48,7 @@ class ExponentialOpTest(test.TestCase): def _verifyExponential(self, x, np_type): inp = x.astype(np_type) with self.test_session(use_gpu=True): - tf_ans = gen_linalg_ops._matrix_exponential(inp) + tf_ans = gen_linalg_ops.matrix_exponential(inp) if x.size == 0: np_ans = np.empty(x.shape, dtype=np_type) else: @@ -116,13 +116,13 @@ class ExponentialOpTest(test.TestCase): # When the exponential of a non-square matrix is attempted we should return # an error with self.assertRaises(ValueError): - gen_linalg_ops._matrix_exponential(np.array([[1., 2., 3.], [3., 4., 5.]])) + gen_linalg_ops.matrix_exponential(np.array([[1., 2., 3.], [3., 4., 5.]])) def testWrongDimensions(self): # The input to the exponential should be at least a 2-dimensional tensor. tensor3 = constant_op.constant([1., 2.]) with self.assertRaises(ValueError): - gen_linalg_ops._matrix_exponential(tensor3) + gen_linalg_ops.matrix_exponential(tensor3) def testEmpty(self): self._verifyExponentialReal(np.empty([0, 2, 2])) @@ -143,8 +143,8 @@ class ExponentialOpTest(test.TestCase): with self.test_session(use_gpu=True) as sess: matrix1 = random_ops.random_normal([5, 5], seed=42) matrix2 = random_ops.random_normal([5, 5], seed=42) - expm1 = gen_linalg_ops._matrix_exponential(matrix1) - expm2 = gen_linalg_ops._matrix_exponential(matrix2) + expm1 = gen_linalg_ops.matrix_exponential(matrix1) + expm2 = gen_linalg_ops.matrix_exponential(matrix2) expm = sess.run([expm1, expm2]) self.assertAllEqual(expm[0], expm[1]) @@ -180,7 +180,7 @@ class MatrixExponentialBenchmark(test.Benchmark): session.Session() as sess, \ ops.device("/cpu:0"): matrix = self._GenerateMatrix(shape) - expm = gen_linalg_ops._matrix_exponential(matrix) + expm = gen_linalg_ops.matrix_exponential(matrix) variables.global_variables_initializer().run() self.run_op_benchmark( sess, diff --git a/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py b/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py index 18ed59828c15f5ad21fe054cd6e40991c02bb356..24edc4f59fe6dd84da6732036eb53e2ad367bd06 100644 --- a/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py +++ b/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py @@ -39,8 +39,8 @@ class LogarithmOpTest(test.TestCase): inp = x.astype(np_type) with self.test_session(use_gpu=True): # Verify that expm(logm(A)) == A. - tf_ans = gen_linalg_ops._matrix_exponential( - gen_linalg_ops._matrix_logarithm(inp)) + tf_ans = gen_linalg_ops.matrix_exponential( + gen_linalg_ops.matrix_logarithm(inp)) out = tf_ans.eval() self.assertAllClose(inp, out, rtol=1e-4, atol=1e-3) @@ -85,14 +85,14 @@ class LogarithmOpTest(test.TestCase): # When the logarithm of a non-square matrix is attempted we should return # an error with self.assertRaises(ValueError): - gen_linalg_ops._matrix_logarithm( + gen_linalg_ops.matrix_logarithm( np.array([[1., 2., 3.], [3., 4., 5.]], dtype=np.complex64)) def testWrongDimensions(self): # The input to the logarithm should be at least a 2-dimensional tensor. tensor3 = constant_op.constant([1., 2.], dtype=dtypes.complex64) with self.assertRaises(ValueError): - gen_linalg_ops._matrix_logarithm(tensor3) + gen_linalg_ops.matrix_logarithm(tensor3) def testEmpty(self): self._verifyLogarithmComplex(np.empty([0, 2, 2], dtype=np.complex64)) @@ -115,8 +115,8 @@ class LogarithmOpTest(test.TestCase): random_ops.random_normal([5, 5], seed=42), dtypes.complex64) matrix2 = math_ops.cast( random_ops.random_normal([5, 5], seed=42), dtypes.complex64) - logm1 = gen_linalg_ops._matrix_logarithm(matrix1) - logm2 = gen_linalg_ops._matrix_logarithm(matrix2) + logm1 = gen_linalg_ops.matrix_logarithm(matrix1) + logm2 = gen_linalg_ops.matrix_logarithm(matrix2) logm = sess.run([logm1, logm2]) self.assertAllEqual(logm[0], logm[1]) @@ -152,7 +152,7 @@ class MatrixLogarithmBenchmark(test.Benchmark): session.Session() as sess, \ ops.device("/cpu:0"): matrix = self._GenerateMatrix(shape) - logm = gen_linalg_ops._matrix_logarithm(matrix) + logm = gen_linalg_ops.matrix_logarithm(matrix) variables.global_variables_initializer().run() self.run_op_benchmark( sess, diff --git a/tensorflow/python/kernel_tests/metrics_test.py b/tensorflow/python/kernel_tests/metrics_test.py index e0e752147cdf8690d22fa782aca2561b2935fa8e..ad802f7e1f72f6cbc3dda1ca98e46e6da4e5110a 100644 --- a/tensorflow/python/kernel_tests/metrics_test.py +++ b/tensorflow/python/kernel_tests/metrics_test.py @@ -417,7 +417,7 @@ class MeanTensorTest(test.TestCase): self.assertAllClose([[-0.9 / 4., 3.525]], sess.run(mean), 5) - def testWeighted1d(self): + def testBinaryWeighted1d(self): with self.test_session() as sess: # Create the queue that populates the values. values_queue = data_flow_ops.FIFOQueue( @@ -444,6 +444,33 @@ class MeanTensorTest(test.TestCase): sess.run(update_op) self.assertAllClose([[3.25, 0.5]], sess.run(mean), 5) + def testWeighted1d(self): + with self.test_session() as sess: + # Create the queue that populates the values. + values_queue = data_flow_ops.FIFOQueue( + 4, dtypes=dtypes_lib.float32, shapes=(1, 2)) + _enqueue_vector(sess, values_queue, [0, 1]) + _enqueue_vector(sess, values_queue, [-4.2, 9.1]) + _enqueue_vector(sess, values_queue, [6.5, 0]) + _enqueue_vector(sess, values_queue, [-3.2, 4.0]) + values = values_queue.dequeue() + + # Create the queue that populates the weights. + weights_queue = data_flow_ops.FIFOQueue( + 4, dtypes=dtypes_lib.float32, shapes=(1, 1)) + _enqueue_vector(sess, weights_queue, [[0.0025]]) + _enqueue_vector(sess, weights_queue, [[0.005]]) + _enqueue_vector(sess, weights_queue, [[0.01]]) + _enqueue_vector(sess, weights_queue, [[0.0075]]) + weights = weights_queue.dequeue() + + mean, update_op = metrics.mean_tensor(values, weights) + + sess.run(variables.local_variables_initializer()) + for _ in range(4): + sess.run(update_op) + self.assertAllClose([[0.8, 3.52]], sess.run(mean), 5) + def testWeighted2d_1(self): with self.test_session() as sess: # Create the queue that populates the values. diff --git a/tensorflow/python/kernel_tests/neon_depthwise_conv_op_test.py b/tensorflow/python/kernel_tests/neon_depthwise_conv_op_test.py index 30795eed8a063076a69ec2ec7851788775fe4dc6..d8ce9fffbd2bc0d18033339a02e0ad84f8f4c952 100644 --- a/tensorflow/python/kernel_tests/neon_depthwise_conv_op_test.py +++ b/tensorflow/python/kernel_tests/neon_depthwise_conv_op_test.py @@ -148,7 +148,7 @@ class DepthwiseConv2DTest(test.TestCase): print("depthwise conv_2d: ", tensor_in_sizes, "*", filter_in_sizes, ", stride:", stride, ", padding: ", padding, ", max diff: ", np.amax(np.absolute(native_result - interface_result))) - self.assertArrayNear( + self.assertAllClose( np.ravel(native_result), np.ravel(interface_result), 1e-5) self.assertShapeEqual(native_result, conv_native) self.assertShapeEqual(native_result, conv_interface) @@ -213,7 +213,7 @@ class DepthwiseConv2DTest(test.TestCase): t1, t2, strides=[1, stride, stride, 1], padding=padding) value = sess.run(conv) print("value = ", value) - self.assertArrayNear(expected, np.ravel(value), 1e-5) + self.assertAllClose(expected, np.ravel(value), 1e-5) self.assertShapeEqual(value, conv) def testConv2D2x2Filter(self): diff --git a/tensorflow/python/kernel_tests/pad_op_test.py b/tensorflow/python/kernel_tests/pad_op_test.py index 2c766e364073fc8c92156f19d08753367982e7fc..14632ec29a5e8224de5bf85673904aac4eb482a7 100644 --- a/tensorflow/python/kernel_tests/pad_op_test.py +++ b/tensorflow/python/kernel_tests/pad_op_test.py @@ -215,13 +215,13 @@ class PadOpTest(test.TestCase): def testIntTypes(self): # TODO(touts): Figure out why the padding tests do not work on GPU # for int types and rank > 2. - for t in [np.int32, np.int64]: + for t in [np.int8, np.int32, np.int64]: self._testAll( np.random.randint(-100, 100, (4, 4, 3)).astype(t), [[1, 0], [2, 3], [0, 2]], 0) self._testAll( np.random.randint(-100, 100, (4, 2, 1, 3)).astype(t), - [[0, 0], [0, 0], [0, 0], [0, 0]], -1234) + [[0, 0], [0, 0], [0, 0], [0, 0]], -123) def testFloatTypes(self): for t in [np.float32, np.float64]: @@ -238,6 +238,29 @@ class PadOpTest(test.TestCase): x = np.random.rand(3, 2, 1, 1).astype(t) self._testAll(x + 1j * x, [[0, 0], [0, 0], [0, 0], [0, 0]], 0 + 0j) + def testString(self): + # Numpy does not support padding strings so we compare padding manually. + x = ops.convert_to_tensor([["Hello", "World"], + ["Goodnight", "Moon"]]) + + constant = array_ops.pad(x, [[1, 0], [0, 1]], mode="CONSTANT", + constant_values="PAD") + reflect = array_ops.pad(x, [[1, 0], [0, 1]], mode="REFLECT", + constant_values="PAD") + symmetric = array_ops.pad(x, [[1, 0], [0, 1]], mode="SYMMETRIC", + constant_values="PAD") + with self.test_session(use_gpu=True): + self.assertAllEqual([[b"PAD", b"PAD", b"PAD"], + [b"Hello", b"World", b"PAD"], + [b"Goodnight", b"Moon", b"PAD"]], constant.eval()) + self.assertAllEqual([[b"Goodnight", b"Moon", b"Goodnight"], + [b"Hello", b"World", b"Hello"], + [b"Goodnight", b"Moon", b"Goodnight"]], + reflect.eval()) + self.assertAllEqual([[b"Hello", b"World", b"World"], + [b"Hello", b"World", b"World"], + [b"Goodnight", b"Moon", b"Moon"]], symmetric.eval()) + def testShapeFunctionEdgeCases(self): # Unknown paddings shape. inp = constant_op.constant(0.0, shape=[4, 4, 4, 4]) @@ -313,5 +336,30 @@ class PadOpTest(test.TestCase): self.assertAllEqual(inp, out) self.assertShapeEqual(inp, tf_val) + def testCollapseAdjacentNonPaddedDimensions(self): + # pyformat: disable + for paddings_value in [[[0, 0], [0, 0], [0, 0], [0, 1]], + [[0, 0], [2, 3], [0, 0], [0, 0]], + [[0, 0], [0, 0], [0, 0], [0, 0]]]: + # pyformat: enable + inp = constant_op.constant(1.0, shape=[8, 28, 28, 3]) + paddings = constant_op.constant(paddings_value, dtype=dtypes.int32) + padded = array_ops.pad(inp, paddings) + middle = array_ops.slice(padded, [row[0] for row in paddings_value], + [dim.value for dim in inp.shape.dims]) + left = array_ops.slice(padded, [0, 0, 0, 0], + [row[0] for row in paddings_value]) + right = array_ops.slice( + padded, + [paddings_value[i][0] + inp.shape.dims[i].value for i in range(4)], + [-1, -1, -1, -1]) + with self.test_session(use_gpu=True): + self.assertAllEqual(inp.eval(), middle.eval()) + self.assertAllEqual( + np.zeros([row[0] for row in paddings_value]), left.eval()) + self.assertAllEqual( + np.zeros([row[1] for row in paddings_value]), right.eval()) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/pool_test.py b/tensorflow/python/kernel_tests/pool_test.py index 63848976336f5487cf2a44f7cf62ea316c40d7c8..6ede654aadc7d0d78bc18f13c2d4b3d47fef0402 100644 --- a/tensorflow/python/kernel_tests/pool_test.py +++ b/tensorflow/python/kernel_tests/pool_test.py @@ -96,7 +96,7 @@ def pool_direct_single_axis( def pool_direct( - input, + input, # pylint: disable=redefined-builtin window_shape, pooling_type, padding, # pylint: disable=redefined-builtin diff --git a/tensorflow/python/kernel_tests/pooling_ops_test.py b/tensorflow/python/kernel_tests/pooling_ops_test.py index 4466beeec96509b3761e34d885276e1510c62d10..2f3bea5825f8889c5880c819ebf6b17aaa613f08 100644 --- a/tensorflow/python/kernel_tests/pooling_ops_test.py +++ b/tensorflow/python/kernel_tests/pooling_ops_test.py @@ -31,6 +31,7 @@ from tensorflow.python.ops import gen_nn_ops from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import variables import tensorflow.python.ops.nn_grad # pylint: disable=unused-import from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging @@ -405,7 +406,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 3, 3, 3], ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], @@ -427,7 +428,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 2, 3, 3], ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], @@ -456,7 +457,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 2, 2, 1], ksize=[1, 1, 2, 1], strides=[1, 1, 1, 1], @@ -485,7 +486,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 4, 4, 1], ksize=[1, 2, 2, 1], strides=[1, 1, 2, 1], @@ -494,7 +495,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu, v2=v2) self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 4, 4, 1], ksize=[1, 2, 2, 1], strides=[1, 2, 1, 1], @@ -519,7 +520,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 4, 4, 4], ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], @@ -554,7 +555,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 8, 8, 8], ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1], @@ -565,7 +566,7 @@ class PoolingTest(test.TestCase): def _testMaxPoolEmptyInput(self, use_gpu): self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[0, 8, 8, 8], ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1], @@ -600,7 +601,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 1, 1, 10], ksize=[1, 1, 1, 2], strides=[1, 1, 1, 2], @@ -626,7 +627,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 2, 2, 6], ksize=[1, 1, 1, 3], strides=[1, 1, 1, 3], @@ -648,7 +649,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 7, 7, 1], ksize=[1, 2, 2, 1], strides=[1, 3, 3, 1], @@ -689,7 +690,7 @@ class PoolingTest(test.TestCase): for v2 in [True, False]: self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 3, 3, 1], ksize=[1, 1, 1, 1], strides=[1, 2, 2, 1], @@ -699,7 +700,7 @@ class PoolingTest(test.TestCase): v2=v2) self._VerifyValues( - gen_nn_ops._max_pool_v2, + gen_nn_ops.max_pool_v2, input_sizes=[1, 4, 4, 1], ksize=[1, 1, 1, 1], strides=[1, 2, 2, 1], @@ -731,7 +732,8 @@ class PoolingTest(test.TestCase): [1, 1, 1, 3], "evenly divide") if test.is_gpu_available(): with self.test_session(use_gpu=True): - t = constant_op.constant(1.0, shape=[1, 2, 2, 4]) + t = variables.Variable(np.ones([1, 2, 2, 4])) + variables.global_variables_initializer().run() with self.assertRaisesOpError("for CPU devices"): nn_ops.max_pool( t, ksize=[1, 1, 1, 2], strides=[1, 1, 1, 2], @@ -764,8 +766,8 @@ class PoolingTest(test.TestCase): _, argmax_op = nn_ops.max_pool_with_argmax(t, ksize, strides, padding) argmax = argmax_op.eval() grad_in = constant_op.constant(tensor_output, shape=output_shape) - out_op = gen_nn_ops._max_pool_grad_with_argmax(t, grad_in, argmax, - ksize, strides, padding) + out_op = gen_nn_ops.max_pool_grad_with_argmax(t, grad_in, argmax, ksize, + strides, padding) gpu_val = out_op.eval() self.assertShapeEqual(gpu_val, out_op) with self.test_session(use_gpu=False): @@ -773,8 +775,8 @@ class PoolingTest(test.TestCase): out_op = nn_ops.max_pool(t, ksize, strides, padding) orig_out = out_op.eval() grad_in = constant_op.constant(tensor_output, shape=output_shape) - out_op = gen_nn_ops._max_pool_grad(t, orig_out, grad_in, ksize, strides, - padding) + out_op = gen_nn_ops.max_pool_grad(t, orig_out, grad_in, ksize, strides, + padding) cpu_val = out_op.eval() self.assertShapeEqual(cpu_val, out_op) # The CPU version accumulates its gradient on fp16, so it's less @@ -793,7 +795,7 @@ class PoolingTest(test.TestCase): _, argmax_op = nn_ops.max_pool_with_argmax(t, ksize, strides, padding) argmax = argmax_op.eval() grad_in = constant_op.constant(tensor_input, shape=input_shape) - out_op = gen_nn_ops._max_pool_grad_grad_with_argmax( + out_op = gen_nn_ops.max_pool_grad_grad_with_argmax( t, grad_in, argmax, ksize, strides, padding) gpu_val = out_op.eval() self.assertShapeEqual(gpu_val, out_op) @@ -802,8 +804,8 @@ class PoolingTest(test.TestCase): out_op = nn_ops.max_pool(t, ksize, strides, padding) orig_out = out_op.eval() grad_in = constant_op.constant(tensor_input, shape=input_shape) - out_op = gen_nn_ops._max_pool_grad_grad(t, orig_out, grad_in, ksize, - strides, padding) + out_op = gen_nn_ops.max_pool_grad_grad(t, orig_out, grad_in, ksize, + strides, padding) cpu_val = out_op.eval() self.assertShapeEqual(cpu_val, out_op) # The CPU version accumulates its gradient on fp16, so it's less @@ -842,7 +844,7 @@ class PoolingTest(test.TestCase): t = constant_op.constant(tensor_input, shape=[1, 2, 2, 1]) argmax = constant_op.constant( tensor_argmax, shape=[1, 2, 2, 1], dtype=dtypes.int64) - out_op = gen_nn_ops._max_pool_grad_with_argmax( + out_op = gen_nn_ops.max_pool_grad_with_argmax( orig_in, t, argmax, @@ -865,7 +867,7 @@ class PoolingTest(test.TestCase): t = constant_op.constant(tensor_input, shape=[1, 3, 3, 1]) argmax = constant_op.constant( tensor_argmax, shape=[1, 2, 2, 1], dtype=dtypes.int64) - out_op = gen_nn_ops._max_pool_grad_grad_with_argmax( + out_op = gen_nn_ops.max_pool_grad_grad_with_argmax( orig_in, t, argmax, @@ -1029,7 +1031,7 @@ class PoolingTest(test.TestCase): self.assertLess(err, err_tolerance) def _testMaxPoolGradValidPadding1_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[1, 3, 3, 1], @@ -1043,7 +1045,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradValidPadding2_1_6(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 6, 6, 3], @@ -1057,7 +1059,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradValidPadding2_1_7(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 7, 7, 3], @@ -1071,7 +1073,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradValidPadding1_2(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[1, 3, 3, 1], @@ -1085,7 +1087,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradValidPadding2_2(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 2, 2, 3], @@ -1099,7 +1101,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradSamePadding1_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1113,7 +1115,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradSamePadding1_2(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1127,7 +1129,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradSamePadding2_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1141,7 +1143,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradSamePadding2_2(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1155,7 +1157,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradSamePadding3_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestGradient( pool_func, input_sizes=[1, 7, 7, 1], @@ -1199,7 +1201,7 @@ class PoolingTest(test.TestCase): Returns: A Tensor. """ - pool_func = gen_nn_ops.max_pool_grad_v2 if v2 else gen_nn_ops._max_pool_grad + pool_func = gen_nn_ops.max_pool_grad_v2 if v2 else gen_nn_ops.max_pool_grad return pool_func(orig_input, orig_output, grad, [1, window_rows, window_cols, 1], [1, row_stride, col_stride, 1], padding) @@ -1208,9 +1210,11 @@ class PoolingTest(test.TestCase): expected_input_backprop, input_sizes, output_sizes, window_rows, window_cols, row_stride, col_stride, padding, use_gpu, v2): - pool_func = gen_nn_ops._max_pool_v2 if v2 else nn_ops.max_pool + pool_func = gen_nn_ops.max_pool_v2 if v2 else nn_ops.max_pool with self.test_session(use_gpu=use_gpu): - input_tensor = constant_op.constant(input_data, shape=input_sizes) + input_tensor = variables.Variable( + np.array(input_data, dtype=np.float32).reshape(input_sizes)) + variables.global_variables_initializer().run() output_tensor = pool_func(input_tensor, [1, window_rows, window_cols, 1], [1, row_stride, col_stride, 1], padding) output_backprop_tensor = constant_op.constant( @@ -1504,7 +1508,7 @@ class PoolingTest(test.TestCase): self._testMaxPoolGradDirectWithNans2_2() def _testMaxPoolGradGradValidPadding1_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[1, 3, 3, 1], @@ -1518,7 +1522,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradValidPadding2_1_6(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[2, 6, 6, 3], @@ -1532,7 +1536,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradValidPadding2_1_7(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[2, 7, 7, 3], @@ -1546,7 +1550,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradValidPadding2_2(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[2, 2, 2, 3], @@ -1560,7 +1564,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradSamePadding1_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1574,7 +1578,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradSamePadding2_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1588,7 +1592,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradSamePadding2_2(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[2, 2, 4, 3], @@ -1602,7 +1606,7 @@ class PoolingTest(test.TestCase): use_gpu=use_gpu) def _testMaxPoolGradGradSamePadding3_1(self, data_format, use_gpu): - for pool_func in [gen_nn_ops._max_pool_v2, nn_ops.max_pool]: + for pool_func in [gen_nn_ops.max_pool_v2, nn_ops.max_pool]: self._ConstructAndTestSecondGradient( pool_func, input_sizes=[1, 7, 7, 1], @@ -1644,7 +1648,7 @@ class PoolingTest(test.TestCase): Returns: A Tensor. """ - return gen_nn_ops._max_pool_grad_grad( + return gen_nn_ops.max_pool_grad_grad( orig_input, orig_output, grad, [1, window_rows, window_cols, 1], [1, row_stride, col_stride, 1], padding) diff --git a/tensorflow/python/kernel_tests/py_func_test.py b/tensorflow/python/kernel_tests/py_func_test.py index c7181497d891f6d35a788c90bf59a0ce5a536328..36142801d682d4e99bc1beceee7aee76206db4cf 100644 --- a/tensorflow/python/kernel_tests/py_func_test.py +++ b/tensorflow/python/kernel_tests/py_func_test.py @@ -1,3 +1,4 @@ +# -*- coding: utf-8 -*- # Copyright 2015 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,6 +19,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import re + import numpy as np from six.moves import queue from six.moves import xrange # pylint: disable=redefined-builtin @@ -212,6 +215,16 @@ class PyFuncTest(test.TestCase): value.op.run() self.assertAllEqual(np_array, [1.0, 2.0]) + def testReturnUnicodeString(self): + with self.test_session(): + correct = u"你好 世界" + + def unicode_string(): + return correct + + z, = script_ops.py_func(unicode_string, [], [dtypes.string]) + self.assertEqual(z.eval(), correct.encode("utf8")) + def testBadNumpyReturnType(self): with self.test_session(): @@ -345,12 +358,22 @@ class PyFuncTest(test.TestCase): def _testExceptionHandling(self, py_exp, tf_exp, eager=False): - def raise_exception(): + def inner_exception(): raise py_exp("blah") # pylint: disable=not-callable + def raise_exception(): + inner_exception() + + expected_regexp = r": blah.*" # Error at the top + expected_regexp += r"in raise_exception.*" # Stacktrace outer + expected_regexp += r"in inner_exception.*" # Stacktrace inner + expected_regexp += r": blah" # Stacktrace of raise + def expected_error_check(exception): + return re.search(expected_regexp, str(exception), re.DOTALL) + if eager: - if context.in_eager_mode(): - with self.assertRaisesRegexp(tf_exp, "blah"): + if context.executing_eagerly(): + with self.assertRaisesWithPredicateMatch(tf_exp, expected_error_check): f = script_ops.eager_py_func(raise_exception, [], []) return else: @@ -359,7 +382,7 @@ class PyFuncTest(test.TestCase): f = script_ops.py_func(raise_exception, [], []) with self.test_session(): - with self.assertRaisesRegexp(tf_exp, "blah"): + with self.assertRaisesWithPredicateMatch(tf_exp, expected_error_check): self.evaluate(f) def testExceptionHandling(self): @@ -421,7 +444,7 @@ class PyFuncTest(test.TestCase): output = script_ops.eager_py_func(no_return_value, inp=[], Tout=[]) ret = self.evaluate(output) - if context.in_eager_mode(): + if context.executing_eagerly(): self.assertEquals(len(ret), 0) else: self.assertIsNone(ret) diff --git a/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py b/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py index c4e16ff6280cc7ce121955474fe8ec45acd57f95..b7a79f239cee04b191b78affd002f687b7de851a 100644 --- a/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py +++ b/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import random -import re import time import numpy as np diff --git a/tensorflow/python/kernel_tests/reduction_ops_test.py b/tensorflow/python/kernel_tests/reduction_ops_test.py index 4231a79b2dcef951048ca54e8c8df2f42b44b1a1..589ea54973c10902c461f552d5c54b6fad6ecf67 100644 --- a/tensorflow/python/kernel_tests/reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/reduction_ops_test.py @@ -30,6 +30,7 @@ from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import math_ops +from tensorflow.python.ops import variables from tensorflow.python.platform import test # The maximum input rank to test. @@ -110,10 +111,10 @@ class ReductionUnknownShape(test.TestCase): class BaseReductionTest(test.TestCase): - def _tf_reduce(self, x, reduction_axes, keep_dims): + def _tf_reduce(self, x, reduction_axes, keepdims): raise NotImplementedError() - def _np_reduce(self, x, reduction_axes, keep_dims): + def _np_reduce(self, x, reduction_axes, keepdims): raise NotImplementedError() def _makeIncremental(self, shape, dtype): @@ -128,10 +129,10 @@ class BaseReductionTest(test.TestCase): data -= 2j * data return data - def _compare(self, x, reduction_axes, keep_dims, feed_dict=None): - np_ans = self._np_reduce(x, reduction_axes, keep_dims) + def _compare(self, x, reduction_axes, keepdims, feed_dict=None): + np_ans = self._np_reduce(x, reduction_axes, keepdims) with self.test_session(use_gpu=True) as sess: - tf_ans = self._tf_reduce(x, reduction_axes, keep_dims) + tf_ans = self._tf_reduce(x, reduction_axes, keepdims) out = sess.run(tf_ans, feed_dict) self.assertAllClose(np_ans, out) self.assertShapeEqual(np_ans, tf_ans) @@ -140,8 +141,8 @@ class BaseReductionTest(test.TestCase): if reduction_axes is not None and np.shape(reduction_axes) == (1,): # Test scalar reduction_axes argument self._compareAll(x, reduction_axes[0]) - self._compare(x, reduction_axes, keep_dims=False, feed_dict=feed_dict) - self._compare(x, reduction_axes, keep_dims=True, feed_dict=feed_dict) + self._compare(x, reduction_axes, keepdims=False, feed_dict=feed_dict) + self._compare(x, reduction_axes, keepdims=True, feed_dict=feed_dict) def _compareAllAxes(self, x, feed_dict=None): self._compareAll(x, None) @@ -171,14 +172,14 @@ class BaseReductionTest(test.TestCase): class SumReductionTest(BaseReductionTest): - def _tf_reduce(self, x, reduction_axes, keep_dims): - return math_ops.reduce_sum(x, reduction_axes, keep_dims) + def _tf_reduce(self, x, reduction_axes, keepdims): + return math_ops.reduce_sum(x, reduction_axes, keepdims) - def _np_reduce(self, x, reduction_axes, keep_dims): + def _np_reduce(self, x, reduction_axes, keepdims): if isinstance(reduction_axes, list) or isinstance(reduction_axes, np.ndarray): reduction_axes = tuple(reduction_axes) - return np.sum(x, axis=reduction_axes, keepdims=keep_dims) + return np.sum(x, axis=reduction_axes, keepdims=keepdims) def testAxesType(self): for dtype in [dtypes.int64, dtypes.int32]: @@ -212,7 +213,8 @@ class SumReductionTest(BaseReductionTest): arr = np.ones([68000], dtype=np.float16) with self.test_session(graph=ops.Graph(), use_gpu=True) as sess: - tf_arr = array_ops.constant(arr) + tf_arr = variables.Variable(arr) + variables.global_variables_initializer().run() tf_mean = math_ops.reduce_mean(tf_arr, 0, False) tf_out_mean = sess.run(tf_mean) self.assertAllClose(tf_out_mean, 1.) @@ -298,7 +300,7 @@ class SumReductionTest(BaseReductionTest): c_known_rank = array_ops.placeholder(dtypes.float32) c_known_rank.set_shape(tensor_shape.unknown_shape(ndims=3)) s_known_rank = math_ops.reduce_sum( - c_known_rank, reduction_axes, keep_dims=True) + c_known_rank, reduction_axes, keepdims=True) self.assertEqual(3, s_known_rank.get_shape().ndims) np_input = np.random.randn(3, 3, 3) @@ -308,11 +310,11 @@ class SumReductionTest(BaseReductionTest): unknown_indices = array_ops.placeholder(dtypes.int32) c_unknown_indices = constant_op.constant([[10.0], [20.0]]) s_unknown_indices = math_ops.reduce_sum( - c_unknown_indices, unknown_indices, keep_dims=False) + c_unknown_indices, unknown_indices, keepdims=False) self.assertEqual(tensor_shape.unknown_shape(), s_unknown_indices.get_shape()) s_unknown_indices_keep = math_ops.reduce_sum( - c_unknown_indices, unknown_indices, keep_dims=True) + c_unknown_indices, unknown_indices, keepdims=True) self.assertEqual(2, s_unknown_indices_keep.get_shape().ndims) def testWrongShapeForReductionIndices(self): @@ -372,10 +374,10 @@ class SumReductionTest(BaseReductionTest): class MeanReductionTest(BaseReductionTest): - def _tf_reduce(self, x, reduction_axes, keep_dims): - return math_ops.reduce_mean(x, reduction_axes, keep_dims) + def _tf_reduce(self, x, reduction_axes, keepdims): + return math_ops.reduce_mean(x, reduction_axes, keepdims) - def _np_reduce(self, x, reduction_axes, keep_dims): + def _np_reduce(self, x, reduction_axes, keepdims): if isinstance(reduction_axes, list) or isinstance(reduction_axes, np.ndarray): reduction_axes = tuple(reduction_axes) @@ -389,7 +391,7 @@ class MeanReductionTest(BaseReductionTest): # np.mean automatically converts integer inputs to float, while TensorFlow's # reduce_mean does not. For integer inputs, we emulate TensorFlow's behavior # using np.sum and truncating division. - np_sum = np.sum(x, axis=reduction_axes, keepdims=keep_dims) + np_sum = np.sum(x, axis=reduction_axes, keepdims=keepdims) if np.issubdtype(x.dtype, np.integer): return np_sum // count return np_sum / count @@ -458,14 +460,14 @@ class MeanReductionTest(BaseReductionTest): class ProdReductionTest(BaseReductionTest): - def _tf_reduce(self, x, reduction_axes, keep_dims): - return math_ops.reduce_prod(x, reduction_axes, keep_dims) + def _tf_reduce(self, x, reduction_axes, keepdims): + return math_ops.reduce_prod(x, reduction_axes, keepdims) - def _np_reduce(self, x, reduction_axes, keep_dims): + def _np_reduce(self, x, reduction_axes, keepdims): if isinstance(reduction_axes, list) or isinstance(reduction_axes, np.ndarray): reduction_axes = tuple(reduction_axes) - return np.prod(x, axis=reduction_axes, keepdims=keep_dims) + return np.prod(x, axis=reduction_axes, keepdims=keepdims) def testAxesType(self): for dtype in [dtypes.int64, dtypes.int32]: @@ -549,17 +551,17 @@ class ProdReductionTest(BaseReductionTest): class MinReductionTest(test.TestCase): - def _compare(self, x, reduction_axes, keep_dims, use_gpu=False): + def _compare(self, x, reduction_axes, keepdims, use_gpu=False): np_ans = x if reduction_axes is None: - np_ans = np.amin(np_ans, keepdims=keep_dims) + np_ans = np.amin(np_ans, keepdims=keepdims) else: for ra in reduction_axes[::-1]: - np_ans = np.amin(np_ans, axis=ra, keepdims=keep_dims) + np_ans = np.amin(np_ans, axis=ra, keepdims=keepdims) with self.test_session(use_gpu=use_gpu): if reduction_axes is not None: reduction_axes = np.array(reduction_axes).astype(np.int32) - tf_ans = math_ops.reduce_min(x, reduction_axes, keep_dims) + tf_ans = math_ops.reduce_min(x, reduction_axes, keepdims) out = tf_ans.eval() self.assertAllClose(np_ans, out) self.assertShapeEqual(np_ans, tf_ans) @@ -662,17 +664,17 @@ class MinReductionTest(test.TestCase): class MaxReductionTest(test.TestCase): - def _compare(self, x, reduction_axes, keep_dims, use_gpu=False): + def _compare(self, x, reduction_axes, keepdims, use_gpu=False): np_ans = x if reduction_axes is None: - np_ans = np.amax(np_ans, keepdims=keep_dims) + np_ans = np.amax(np_ans, keepdims=keepdims) else: for ra in reduction_axes[::-1]: - np_ans = np.amax(np_ans, axis=ra, keepdims=keep_dims) + np_ans = np.amax(np_ans, axis=ra, keepdims=keepdims) with self.test_session(use_gpu=use_gpu): if reduction_axes is not None: reduction_axes = np.array(reduction_axes).astype(np.int32) - tf_ans = math_ops.reduce_max(x, reduction_axes, keep_dims) + tf_ans = math_ops.reduce_max(x, reduction_axes, keepdims) out = tf_ans.eval() self.assertAllClose(np_ans, out) self.assertShapeEqual(np_ans, tf_ans) @@ -789,17 +791,17 @@ class MaxReductionTest(test.TestCase): class AllReductionTest(test.TestCase): - def _compare(self, x, reduction_axes, keep_dims, use_gpu=False): + def _compare(self, x, reduction_axes, keepdims, use_gpu=False): np_ans = x if reduction_axes is None: - np_ans = np.all(np_ans, keepdims=keep_dims) + np_ans = np.all(np_ans, keepdims=keepdims) else: for ra in reduction_axes[::-1]: - np_ans = np.all(np_ans, axis=ra, keepdims=keep_dims) + np_ans = np.all(np_ans, axis=ra, keepdims=keepdims) with self.test_session(use_gpu=use_gpu): if reduction_axes is not None: reduction_axes = np.array(reduction_axes).astype(np.int32) - tf_ans = math_ops.reduce_all(x, reduction_axes, keep_dims) + tf_ans = math_ops.reduce_all(x, reduction_axes, keepdims) out = tf_ans.eval() self.assertAllEqual(np_ans, out) self.assertShapeEqual(np_ans, tf_ans) @@ -838,17 +840,17 @@ class AllReductionTest(test.TestCase): class AnyReductionTest(test.TestCase): - def _compare(self, x, reduction_axes, keep_dims, use_gpu=False): + def _compare(self, x, reduction_axes, keepdims, use_gpu=False): np_ans = x if reduction_axes is None: - np_ans = np.any(np_ans, keepdims=keep_dims) + np_ans = np.any(np_ans, keepdims=keepdims) else: for ra in reduction_axes[::-1]: - np_ans = np.any(np_ans, axis=ra, keepdims=keep_dims) + np_ans = np.any(np_ans, axis=ra, keepdims=keepdims) with self.test_session(use_gpu=use_gpu): if reduction_axes is not None: reduction_axes = np.array(reduction_axes).astype(np.int32) - tf_ans = math_ops.reduce_any(x, reduction_axes, keep_dims) + tf_ans = math_ops.reduce_any(x, reduction_axes, keepdims) out = tf_ans.eval() self.assertAllEqual(np_ans, out) self.assertShapeEqual(np_ans, tf_ans) @@ -887,21 +889,17 @@ class AnyReductionTest(test.TestCase): class CountNonzeroReductionTest(test.TestCase): - def _compare(self, - x, - reduction_axes, - keep_dims, - use_gpu=False, + def _compare(self, x, reduction_axes, keepdims, use_gpu=False, feed_dict=None): np_ans = (x != 0).astype(np.int32) if reduction_axes is None: - np_ans = np.sum(np_ans, keepdims=keep_dims) + np_ans = np.sum(np_ans, keepdims=keepdims) else: reduction_axes = np.array(reduction_axes).astype(np.int32) for ra in reduction_axes.ravel()[::-1]: - np_ans = np.sum(np_ans, axis=ra, keepdims=keep_dims) + np_ans = np.sum(np_ans, axis=ra, keepdims=keepdims) with self.test_session(use_gpu=use_gpu) as sess: - tf_ans = math_ops.count_nonzero(x, reduction_axes, keep_dims) + tf_ans = math_ops.count_nonzero(x, reduction_axes, keepdims) out = sess.run(tf_ans, feed_dict) self.assertAllClose(np_ans, out) self.assertShapeEqual(np_ans, tf_ans) diff --git a/tensorflow/python/kernel_tests/reduction_ops_test_big.py b/tensorflow/python/kernel_tests/reduction_ops_test_big.py index 0959adb026e3934713442e6f3487b30a0b252943..d70360775a03caa32eab995371d54786c3c0a0d9 100644 --- a/tensorflow/python/kernel_tests/reduction_ops_test_big.py +++ b/tensorflow/python/kernel_tests/reduction_ops_test_big.py @@ -27,24 +27,24 @@ from tensorflow.python.platform import test class BaseReductionTest(test.TestCase): - def _tf_reduce(self, x, reduction_axes, keep_dims): + def _tf_reduce(self, x, reduction_axes, keepdims): raise NotImplementedError() class BigReductionTest(BaseReductionTest): """Test reductions for sum and boolean all over a wide range of shapes.""" - def _tf_reduce_max(self, x, reduction_axes, keep_dims): - return math_ops.reduce_max(x, reduction_axes, keep_dims) + def _tf_reduce_max(self, x, reduction_axes, keepdims): + return math_ops.reduce_max(x, reduction_axes, keepdims) - def _tf_reduce_all(self, x, reduction_axes, keep_dims): - return math_ops.reduce_all(x, reduction_axes, keep_dims) + def _tf_reduce_all(self, x, reduction_axes, keepdims): + return math_ops.reduce_all(x, reduction_axes, keepdims) - def _tf_reduce_mean(self, x, reduction_axes, keep_dims): - return math_ops.reduce_mean(x, reduction_axes, keep_dims) + def _tf_reduce_mean(self, x, reduction_axes, keepdims): + return math_ops.reduce_mean(x, reduction_axes, keepdims) - def _tf_reduce_sum(self, x, reduction_axes, keep_dims): - return math_ops.reduce_sum(x, reduction_axes, keep_dims) + def _tf_reduce_sum(self, x, reduction_axes, keepdims): + return math_ops.reduce_sum(x, reduction_axes, keepdims) def testFloat32Sum(self): # make sure we test all possible kernel invocations diff --git a/tensorflow/python/kernel_tests/regex_replace_op_test.py b/tensorflow/python/kernel_tests/regex_replace_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6739ac32245668e98d37673fe9e9fe9d55cc0c5f --- /dev/null +++ b/tensorflow/python/kernel_tests/regex_replace_op_test.py @@ -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. +# ============================================================================== +"""Tests for RegexReplace op from string_ops.""" + +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.ops import string_ops +from tensorflow.python.platform import test + + +class RegexReplaceOpTest(test.TestCase): + + def testRemovePrefix(self): + values = ["a:foo", "a:bar", "a:foo", "b:baz", "b:qux", "ca:b"] + with self.test_session(): + input_vector = constant_op.constant(values, dtypes.string) + stripped = string_ops.regex_replace( + input_vector, "^(a:|b:)", "", replace_global=False).eval() + self.assertAllEqual([b"foo", b"bar", b"foo", b"baz", b"qux", b"ca:b"], + stripped) + + def testRegexReplace(self): + values = ["aba\naba", "abcdabcde"] + with self.test_session(): + input_vector = constant_op.constant(values, dtypes.string) + stripped = string_ops.regex_replace(input_vector, "a.*a", "(\\0)").eval() + self.assertAllEqual([b"(aba)\n(aba)", b"(abcda)bcde"], stripped) + + def testEmptyMatch(self): + values = ["abc", "1"] + with self.test_session(): + input_vector = constant_op.constant(values, dtypes.string) + stripped = string_ops.regex_replace(input_vector, "", "x").eval() + self.assertAllEqual([b"xaxbxcx", b"x1x"], stripped) + + def testInvalidPattern(self): + values = ["abc", "1"] + with self.test_session(): + input_vector = constant_op.constant(values, dtypes.string) + invalid_pattern = "A[" + replace = string_ops.regex_replace(input_vector, invalid_pattern, "x") + with self.assertRaisesOpError("Invalid pattern"): + replace.eval() + + def testGlobal(self): + values = ["ababababab", "abcabcabc", ""] + with self.test_session(): + input_vector = constant_op.constant(values, dtypes.string) + stripped = string_ops.regex_replace(input_vector, "ab", "abc", + True).eval() + self.assertAllEqual([b"abcabcabcabcabc", b"abccabccabcc", b""], stripped) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/kernel_tests/relu_op_test.py b/tensorflow/python/kernel_tests/relu_op_test.py index 6b4091ae5d3c6e469a9cd5237b978eae4c75485f..25e947f09e137b37ea129ba6015a060aa01f02e4 100644 --- a/tensorflow/python/kernel_tests/relu_op_test.py +++ b/tensorflow/python/kernel_tests/relu_op_test.py @@ -19,12 +19,14 @@ 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 from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import gradients_impl +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 variables @@ -87,6 +89,35 @@ class ReluTest(test.TestCase): print("relu (float32) gradient err = ", err) self.assertLess(err, 1e-4) + # The gradient for fp16 is inaccurate due to the low-precision. + # Instead of relying on compute_gradient_error, we compare the fp16 analytical + # gradient against their fp32 counterpart. + def testGradientFloat16(self): + with self.test_session(use_gpu=True) as sess: + # Randomly construct a 1D shape from [1, 40) + shape = random_ops.random_uniform( + [1], minval=1, maxval=40, dtype=dtypes.int32) + + # Construct the fp32 graph and its gradient. + x = random_ops.random_uniform(shape, minval=-1, maxval=1, name="x") + y1 = nn_ops.relu(x, name="relu_fp32") + l1 = nn_ops.l2_loss(y1) + dx_f32 = gradients_impl.gradients(l1, x) + + # Construct the fp16 graph and its gradient. + # It starts with the same x, in fp32. But before it reaches Relu, it is + # cast into fp16. So during backprop, the gradient computation is in fp16. + x2 = math_ops.cast(x, dtype=dtypes.float16, name="cast") + y2 = nn_ops.relu(x2, name="relu_fp16") + l2 = nn_ops.l2_loss(y2) + dx_f16 = gradients_impl.gradients(l2, x) + + # Repeat the experiment for 100 times. All tensor shapes and its tensor + # values are randomly generated for each run. + for _ in xrange(100): + dx_f32_v, dx_f16_v = sess.run([dx_f32, dx_f16]) + self.assertAllClose(dx_f32_v, dx_f16_v, atol=3e-4) + def testGradientFloat64(self): with self.test_session(): x = constant_op.constant( diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py index b4b555591d054226210eb6af20036967b240928f..2dc993f8117a41de8f15663ce763ec1d5b7ecdb4 100644 --- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py @@ -36,6 +36,7 @@ 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.util import compat @test_util.with_c_api @@ -63,6 +64,13 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): 0, dtype=dtypes.int32)).run() + def testGPUInt64(self): + if not context.context().num_gpus(): + return + with context.eager_mode(), context.device("gpu:0"): + v = resource_variable_ops.ResourceVariable(1, dtype=dtypes.int64) + self.assertAllEqual(1, v.numpy()) + def testEagerNameNotIdentity(self): with context.eager_mode(): v0 = resource_variable_ops.ResourceVariable(1.0, name="a") @@ -161,14 +169,26 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): @test_util.run_in_graph_and_eager_modes(use_gpu=True) def testScatterAdd(self): + with ops.device("cpu:0"): + handle = resource_variable_ops.var_handle_op( + dtype=dtypes.int32, shape=[1, 1]) + self.evaluate(resource_variable_ops.assign_variable_op( + handle, constant_op.constant([[1]], dtype=dtypes.int32))) + self.evaluate(resource_variable_ops.resource_scatter_add( + handle, [0], constant_op.constant([[2]], dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[3]]) + + def testScatterUpdateString(self): handle = resource_variable_ops.var_handle_op( - dtype=dtypes.int32, shape=[1, 1]) + dtype=dtypes.string, shape=[1, 1]) self.evaluate(resource_variable_ops.assign_variable_op( - handle, constant_op.constant([[1]], dtype=dtypes.int32))) - self.evaluate(resource_variable_ops.resource_scatter_add( - handle, [0], constant_op.constant([[2]], dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[3]]) + handle, constant_op.constant([["a"]], dtype=dtypes.string))) + self.evaluate(resource_variable_ops.resource_scatter_update( + handle, [0], constant_op.constant([["b"]], dtype=dtypes.string))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.string) + self.assertEqual(compat.as_bytes(self.evaluate(read)[0][0]), + compat.as_bytes("b")) # TODO(alive): get this to work in Eager mode. def testGPU(self): @@ -257,6 +277,17 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): self.evaluate(v.assign(2.0)) self.assertEqual(2.0, self.evaluate(v.value())) + # Tests for the 'read_value' argument: + assign_with_read = v.assign(3.0, read_value=True) + self.assertEqual(3.0, self.evaluate(assign_with_read)) + assign_without_read = v.assign(4.0, read_value=False) + if context.executing_eagerly(): + self.assertIsNone(assign_without_read) + else: + self.assertIsInstance(assign_without_read, ops.Operation) + self.evaluate(assign_without_read) + self.assertEqual(4.0, self.evaluate(v.value())) + @test_util.run_in_graph_and_eager_modes() def testLoad(self): v = resource_variable_ops.ResourceVariable(1.0, name="var0") @@ -264,6 +295,32 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): v.load(2.0) self.assertEqual(2.0, self.evaluate(v.value())) + def testVariableDefInitializedInstances(self): + with ops.Graph().as_default(), self.test_session() as sess: + v_def = resource_variable_ops.ResourceVariable( + initial_value=constant_op.constant(3.0)).to_proto() + + with ops.Graph().as_default(), self.test_session() as sess: + # v describes a VariableDef-based variable without an initial value. + v = resource_variable_ops.ResourceVariable(variable_def=v_def) + self.assertEqual(3.0, sess.run(v.initialized_value())) + + # initialized_value should not rerun the initializer_op if the variable + # has already been initialized elsewhere. + sess.run(v.assign(1.0)) + self.assertEqual(1.0, v.initialized_value().eval()) + + v_def.ClearField("initial_value_name") + with ops.Graph().as_default(), self.test_session() as sess: + # Restoring a legacy VariableDef proto that does not have + # initial_value_name set should still work. + v = resource_variable_ops.ResourceVariable(variable_def=v_def) + # We should also be able to re-export the variable to a new meta graph. + self.assertProtoEquals(v_def, v.to_proto()) + # But attempts to use initialized_value will result in errors. + with self.assertRaises(ValueError): + sess.run(v.initialized_value()) + @test_util.run_in_graph_and_eager_modes() def testSparseRead(self): with self.test_session(): @@ -283,6 +340,9 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): w = resource_variable_ops.ResourceVariable.from_proto(v.to_proto()) self.assertEquals(2, math_ops.add(w, 1).eval()) + self.assertEquals(v._handle, w._handle) + self.assertEquals(v._graph_element, w._graph_element) + @test_util.run_in_graph_and_eager_modes() def testAssignAddMethod(self): v = resource_variable_ops.ResourceVariable(1.0, name="var0") @@ -290,6 +350,17 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): self.evaluate(v.assign_add(1.0)) self.assertEqual(2.0, self.evaluate(v.value())) + # Tests for the 'read_value' argument: + assign_with_read = v.assign_add(1.0, read_value=True) + self.assertEqual(3.0, self.evaluate(assign_with_read)) + assign_without_read = v.assign_add(1.0, read_value=False) + if context.executing_eagerly(): + self.assertIsNone(assign_without_read) + else: + self.assertIsInstance(assign_without_read, ops.Operation) + self.evaluate(assign_without_read) + self.assertEqual(4.0, self.evaluate(v.value())) + @test_util.run_in_graph_and_eager_modes() def testAssignSubMethod(self): v = resource_variable_ops.ResourceVariable(3.0, name="var0") @@ -297,6 +368,17 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): self.evaluate(v.assign_sub(1.0)) self.assertEqual(2.0, self.evaluate(v.value())) + # Tests for the 'read_value' argument: + assign_with_read = v.assign_sub(1.0, read_value=True) + self.assertEqual(1.0, self.evaluate(assign_with_read)) + assign_without_read = v.assign_sub(1.0, read_value=False) + if context.executing_eagerly(): + self.assertIsNone(assign_without_read) + else: + self.assertIsInstance(assign_without_read, ops.Operation) + self.evaluate(assign_without_read) + self.assertEqual(0.0, self.evaluate(v.value())) + @test_util.run_in_graph_and_eager_modes() def testDestroyResource(self): v = resource_variable_ops.ResourceVariable(3.0, name="var0") @@ -394,7 +476,7 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): self.assertEqual("(10, 20, 35)", str(v.get_shape())) self.assertEqual("(10, 20, 35)", str(v.value().shape)) self.assertEqual("(3, 20, 35)", str(v.sparse_read([0, 1, 2]).shape)) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( "", str(v.sparse_read(array_ops.placeholder(dtypes.int32)).shape)) @@ -435,7 +517,6 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): self.assertEqual(dtypes.int32, v.dtype) self.assertEqual("foo/var7:0", v.name) self.assertAllEqual([10, 20, 35], v.shape.as_list()) - self.assertEqual(context.get_default_context().device_name, v.device) self.assertTrue(isinstance(v.handle, ops.EagerTensor)) self.assertEqual(constraint, v.constraint) self.assertAllEqual(init.numpy(), v.read_value().numpy()) @@ -505,6 +586,15 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): state_ops.scatter_update(v, [1], [3]) self.assertAllEqual([1.0, 3.0], v.numpy()) + @test_util.run_in_graph_and_eager_modes() + def testScatterUpdateInvalidArgs(self): + v = resource_variable_ops.ResourceVariable([0, 1, 2, 3], name="update") + # The exact error and message differ between graph construction (where the + # error is realized during shape inference at graph construction time) and + # eager execution (where the error is realized during kernel execution). + with self.assertRaisesRegexp(Exception, r"shape.*2.*3"): + state_ops.scatter_update(v, [0, 1], [0, 1, 2]) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/rnn_test.py b/tensorflow/python/kernel_tests/rnn_test.py index a86b65affec9e491fec13577ad6dc3db610df797..9a0409c796ab60da3d47cf7d46ef6fbd5bd82394 100644 --- a/tensorflow/python/kernel_tests/rnn_test.py +++ b/tensorflow/python/kernel_tests/rnn_test.py @@ -23,7 +23,7 @@ import timeit import numpy as np -from six.moves import xrange +from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.contrib import rnn as contrib_rnn from tensorflow.core.protobuf import config_pb2 from tensorflow.python.client import session @@ -111,10 +111,10 @@ class RNNTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testInvalidSequenceLengthShape(self): cell = Plus1RNNCell() - if context.in_graph_mode(): - inputs = [array_ops.placeholder(dtypes.float32, shape=(3, 4))] - else: + if context.executing_eagerly(): inputs = [constant_op.constant(np.ones((3, 4)))] + else: + inputs = [array_ops.placeholder(dtypes.float32, shape=(3, 4))] with self.assertRaisesRegexp(ValueError, "must be a vector"): rnn.dynamic_rnn( cell, @@ -125,38 +125,30 @@ class RNNTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testBatchSizeFromInput(self): cell = Plus1RNNCell() - in_graph_mode = context.in_graph_mode() + in_eager_mode = context.executing_eagerly() # With static batch size - if in_graph_mode: - inputs = array_ops.placeholder(dtypes.float32, shape=(3, 4, 5)) - initial_state = array_ops.placeholder(dtypes.float32, shape=(3, 5)) - else: + if in_eager_mode: inputs = np.zeros((3, 4, 5), dtype=np.float32) initial_state = np.zeros((3, 5), dtype=np.float32) + else: + inputs = array_ops.placeholder(dtypes.float32, shape=(3, 4, 5)) + initial_state = array_ops.placeholder(dtypes.float32, shape=(3, 5)) # - Without initial_state outputs, state = rnn.dynamic_rnn(cell, inputs, dtype=dtypes.float32) - if in_graph_mode: - self.assertEqual(3, outputs.shape[0].value) - self.assertEqual(3, state.shape[0].value) - else: - self.assertEqual(3, outputs.shape[0]) - self.assertEqual(3, state.shape[0]) + self.assertEqual(3, outputs.shape[0]) + self.assertEqual(3, state.shape[0]) # - With initial_state outputs, state = rnn.dynamic_rnn( cell, inputs, initial_state=initial_state) - if in_graph_mode: - self.assertEqual(3, outputs.shape[0].value) - self.assertEqual(3, state.shape[0].value) - else: - self.assertEqual(3, outputs.shape[0]) - self.assertEqual(3, state.shape[0]) + self.assertEqual(3, outputs.shape[0]) + self.assertEqual(3, state.shape[0]) # Without static batch size - # Tensor shapes are fully determined in Eager mode, so only run this - # test in graph mode. - if in_graph_mode: + # Tensor shapes are fully determined with eager execution enabled, + # so only run this test for graph construction. + if not in_eager_mode: inputs = array_ops.placeholder(dtypes.float32, shape=(None, 4, 5)) # - Without initial_state outputs, state = rnn.dynamic_rnn(cell, inputs, dtype=dtypes.float32) @@ -173,56 +165,46 @@ class RNNTest(test.TestCase): @test_util.run_in_graph_and_eager_modes() def testScalarStateIsAccepted(self): cell = ScalarStateRNNCell() - in_graph_mode = context.in_graph_mode() + in_eager_mode = context.executing_eagerly() - if in_graph_mode: - inputs = array_ops.placeholder(dtypes.float32, shape=(1, 4, 1)) - else: + if in_eager_mode: inputs = np.array([[[1], [2], [3], [4]]], dtype=np.float32) + else: + inputs = array_ops.placeholder(dtypes.float32, shape=(1, 4, 1)) with self.test_session() as sess: outputs, state = rnn.dynamic_rnn( cell, inputs, dtype=dtypes.float32, sequence_length=[4]) - if in_graph_mode: + if not in_eager_mode: outputs, state = sess.run( [outputs, state], feed_dict={inputs: [[[1], [2], [3], [4]]]}) - if in_graph_mode: - self.assertAllEqual(outputs, np.array([[[1], [2], [3], [4]]])) - self.assertEqual(state, 4) - else: - self.assertAllEqual(outputs.numpy(), np.array([[[1], [2], [3], [4]]])) - self.assertEqual(state.numpy(), 4) + self.assertAllEqual([[[1], [2], [3], [4]]], outputs) + self.assertAllEqual(4, state) @test_util.run_in_graph_and_eager_modes() def testTensorArrayStateIsAccepted(self): cell = TensorArrayStateRNNCell() - in_graph_mode = context.in_graph_mode() + in_eager_mode = context.executing_eagerly() - if in_graph_mode: - inputs = array_ops.placeholder(dtypes.float32, shape=(1, 4, 1)) - else: + if in_eager_mode: inputs = np.array([[[1], [2], [3], [4]]], dtype=np.float32) + else: + inputs = array_ops.placeholder(dtypes.float32, shape=(1, 4, 1)) with self.test_session() as sess: outputs, state = rnn.dynamic_rnn( cell, inputs, dtype=dtypes.float32, sequence_length=[4]) state = (state[0], state[1].stack()) - if in_graph_mode: + if not in_eager_mode: outputs, state = sess.run( [outputs, state], feed_dict={ inputs: [[[1], [2], [3], [4]]] }) - if in_graph_mode: - self.assertAllEqual(outputs, np.array([[[1], [2], [3], [4]]])) - self.assertEqual(state[0], 4) - self.assertAllEqual(state[1], np.array([[[1]], [[2]], [[3]], [[4]]])) - else: - self.assertAllEqual(outputs.numpy(), np.array([[[1], [2], [3], [4]]])) - self.assertEqual(state[0].numpy(), 4) - self.assertAllEqual(state[1].numpy(), - np.array([[[1]], [[2]], [[3]], [[4]]])) + self.assertAllEqual([[[1], [2], [3], [4]]], outputs) + self.assertAllEqual(4, state[0]) + self.assertAllEqual([[[1]], [[2]], [[3]], [[4]]], state[1]) ######### Benchmarking RNN code diff --git a/tensorflow/python/kernel_tests/save_restore_ops_test.py b/tensorflow/python/kernel_tests/save_restore_ops_test.py index 1bdfa9ebd8e1a4495e67004f59adfb51bf3a6602..cb9aa1e34d6eb82efa94e60e7b56c26b181cef04 100644 --- a/tensorflow/python/kernel_tests/save_restore_ops_test.py +++ b/tensorflow/python/kernel_tests/save_restore_ops_test.py @@ -31,11 +31,10 @@ class ShardedFileOpsTest(test.TestCase): with session.Session( target="", config=config_pb2.ConfigProto(device_count={"CPU": 2})): self.assertEqual( - gen_io_ops._sharded_filename("foo", 4, 100).eval(), + gen_io_ops.sharded_filename("foo", 4, 100).eval(), b"foo-00004-of-00100") self.assertEqual( - gen_io_ops._sharded_filespec("foo", 100).eval(), - b"foo-?????-of-00100") + gen_io_ops.sharded_filespec("foo", 100).eval(), b"foo-?????-of-00100") class ShapeInferenceTest(test.TestCase): @@ -53,7 +52,7 @@ class ShapeInferenceTest(test.TestCase): [dtypes.float32, dtypes.float32]) def testRestoreSlice(self): - op = gen_io_ops._restore_slice("model", "var", "3 4 0,1:-", dtypes.float32) + op = gen_io_ops.restore_slice("model", "var", "3 4 0,1:-", dtypes.float32) self.assertEqual([1, 4], op.get_shape()) diff --git a/tensorflow/python/kernel_tests/scalar_test.py b/tensorflow/python/kernel_tests/scalar_test.py index e65241981eac2d42207c1de7a261f7936e588f2a..0d8fd232946883ac1d95c4c2d9744af69175ab90 100644 --- a/tensorflow/python/kernel_tests/scalar_test.py +++ b/tensorflow/python/kernel_tests/scalar_test.py @@ -92,11 +92,11 @@ class ScalarTest(test.TestCase): self.check(array_ops.reshape, (7, 1), 'sizes input must be 1-D', [7]) def testShardedFilename(self): - self.check(gen_io_ops._sharded_filename, ('foo', 4, [100]), + self.check(gen_io_ops.sharded_filename, ('foo', 4, [100]), 'must be a scalar', b'foo-00004-of-00100') def testShardedFilespec(self): - self.check(gen_io_ops._sharded_filespec, ('foo', [100]), 'must be a scalar', + self.check(gen_io_ops.sharded_filespec, ('foo', [100]), 'must be a scalar', b'foo-?????-of-00100') def testUnsortedSegmentSum(self): diff --git a/tensorflow/python/kernel_tests/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/segment_reduction_ops_test.py index 5a54f448d092093db668570d055801f9f9cd0f9f..3bca5fadc42693f514911c7ffa8f078de8ef9bcd 100644 --- a/tensorflow/python/kernel_tests/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/segment_reduction_ops_test.py @@ -46,7 +46,8 @@ class SegmentReductionHelper(test.TestCase): return constant_op.constant( np_values, shape=input_shape, dtype=dtype), np_values - def _segmentReduce(self, indices, x, op1, op2=None, num_segments=None): + def _segmentReduce(self, indices, x, op1, op2=None, num_segments=None, + initial_value=0): if not x.size: return np.array([]) indices = np.asarray(indices) @@ -64,13 +65,8 @@ class SegmentReductionHelper(test.TestCase): else: output[index] = x_flat[i] # zero initialize values that are still uncalcuated. - # output = [o if o is not None else np.zeros(slice_shape) for o in output] - if not op1 == np.max: - output = [o if o is not None else np.zeros(slice_shape) for o in output] - else: - zeroslice = np.zeros(slice_shape) - zeroslice.fill(dtype.min) - output = [o if o is not None else zeroslice for o in output] + initial_value_slice = np.ones(slice_shape) * initial_value + output = [o if o is not None else initial_value_slice for o in output] if op2 is not None: output = [op2(o) for o in output] output = [o.reshape(slice_shape) for o in output] @@ -82,6 +78,9 @@ class SegmentReductionHelper(test.TestCase): def _mean_reduce_op(self, x): return x[0] / x[1] if isinstance(x, tuple) else x + def _sqrt_n_reduce_op(self, x): + return x[0] / np.sqrt(x[1]) if isinstance(x, tuple) else x + class SegmentReductionOpTest(SegmentReductionHelper): @@ -244,27 +243,61 @@ class SegmentReductionOpTest(SegmentReductionHelper): self.assertAllClose(jacob_t, jacob_n) -class UnsortedSegmentSumTest(SegmentReductionHelper): +class UnsortedSegmentTest(SegmentReductionHelper): + + def __init__(self, methodName='runTest'): + # Each item is np_op1, np_op2, tf_op, initial_value functor + self.ops_list = [(np.add, None, + math_ops.unsorted_segment_sum, lambda t: 0), + (self._mean_cum_op, self._mean_reduce_op, + math_ops.unsorted_segment_mean, lambda t: 0), + (self._mean_cum_op, self._sqrt_n_reduce_op, + math_ops.unsorted_segment_sqrt_n, lambda t: 0), + (np.ndarray.__mul__, None, + math_ops.unsorted_segment_prod, lambda t: 1), + (np.minimum, None, + math_ops.unsorted_segment_min, lambda t: t.max), + (np.maximum, None, + math_ops.unsorted_segment_max, lambda t: t.min)] + + # A subset of ops has been enabled for complex numbers + self.complex_ops_list = [(np.add, None, + math_ops.unsorted_segment_sum, lambda t: 0)] + self.differentiable_dtypes = [dtypes_lib.float16, dtypes_lib.float32, + dtypes_lib.float64] + self.all_dtypes = (self.differentiable_dtypes + + [dtypes_lib.bfloat16, + dtypes_lib.int64, dtypes_lib.int32, + dtypes_lib.complex64, dtypes_lib.complex128]) + super(UnsortedSegmentTest, self).__init__(methodName=methodName) def testValues(self): - dtypes = [ - dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int64, - dtypes_lib.int32, dtypes_lib.complex64, dtypes_lib.complex128 - ] indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) num_segments = 12 for indices in indices_flat, indices_flat.reshape(5, 2): shape = indices.shape + (2,) - for dtype in dtypes: - with self.test_session(use_gpu=True): - tf_x, np_x = self._input(shape, dtype=dtype) - np_ans = self._segmentReduce( - indices, np_x, np.add, op2=None, num_segments=num_segments) - s = math_ops.unsorted_segment_sum( - data=tf_x, segment_ids=indices, num_segments=num_segments) - tf_ans = s.eval() - self.assertAllClose(np_ans, tf_ans) - self.assertShapeEqual(np_ans, s) + for dtype in self.all_dtypes: + ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list + tf_x, np_x = self._input(shape, dtype=dtype) + for use_gpu in [True, False]: + with self.test_session(use_gpu=True): + for np_op1, np_op2, tf_op, init_op in ops_list: + # sqrt_n doesn't support integers + if (np_op2 == self._sqrt_n_reduce_op and dtype.is_integer): + continue + # todo(philjd): enable this test once real_div supports bfloat16 + if (np_op2 in [self._sqrt_n_reduce_op, self._mean_reduce_op] and + dtype == dtypes_lib.bfloat16): + continue + np_ans = self._segmentReduce( + indices, np_x, np_op1, np_op2, num_segments=num_segments, + initial_value=init_op(dtype)) + s = tf_op(tf_x, segment_ids=indices, num_segments=num_segments) + tf_ans = s.eval() + if dtype is dtypes_lib.bfloat16: + tf_ans = tf_ans.astype(np.float32) + self.assertAllClose(np_ans, tf_ans) + self.assertShapeEqual(np_ans, s) def testNumSegmentsTypes(self): dtypes = [dtypes_lib.int32, dtypes_lib.int64] @@ -287,25 +320,51 @@ class UnsortedSegmentSumTest(SegmentReductionHelper): self.assertAllClose(np_ans, tf_ans) self.assertShapeEqual(np_ans, s) - def testGradientSegmentSum(self): + def testGradients(self): num_cols = 2 - indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) + indices_flat = np.array([0, 4, 0, -1, 3, -1, 4, 7, 7, 3]) num_segments = max(indices_flat) + 3 - for dtype in [dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.complex64, - dtypes_lib.complex128]: + for dtype in self.differentiable_dtypes: + ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list for indices in indices_flat, indices_flat.reshape(5, 2): shape = indices.shape + (num_cols,) - with self.test_session(use_gpu=True): - tf_x, np_x = self._input(shape, dtype=dtype) - s = math_ops.unsorted_segment_sum( - data=tf_x, segment_ids=indices, num_segments=num_segments) + # test CPU and GPU as tf.gather behaves differently on each device + for use_gpu in [False, True]: + with self.test_session(use_gpu=use_gpu): + for _, _, tf_op, _ in ops_list: + tf_x, np_x = self._input(shape, dtype=dtype) + s = tf_op(tf_x, indices, num_segments) + jacob_t, jacob_n = gradient_checker.compute_gradient( + tf_x, + shape, + s, [num_segments, num_cols], + x_init_value=np_x, + delta=1) + self.assertAllClose(jacob_t, jacob_n) + + def testProdGrad(self): + # additional test for the prod gradient to ensure correct handling of zeros + values = np.array([0, 0, 1, 0, 2, 2, 3, 3, 3], dtype=np.float32) + indices = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2], dtype=np.int32) + indices_neg = np.array([-1, 0, 0, -1, 1, 1, -1, 2, 2], dtype=np.int32) + values_tf = constant_op.constant(values) + # ground truth partial derivatives + gradients_indices = np.zeros((9, 3), dtype=np.float32) + gradients_indices_neg = np.zeros((9, 3), dtype=np.float32) + # the derivative w.r.t. to the other segments is zero, so here we only + # explicitly set the grad values for the corresponding segment + gradients_indices[range(9), indices] = [0, 0, 0, 4, 0, 0, 9, 9, 9] + gradients_indices_neg[range(9), indices_neg] = [0, 1, 0, 0, 2, 2, 0, 3, 3] + for use_gpu in [False, True]: + with self.test_session(use_gpu=use_gpu): + for ind, grad_gt in [(indices, gradients_indices), + (indices_neg, gradients_indices_neg)]: + s = math_ops.unsorted_segment_prod(values_tf, + constant_op.constant(ind), 3) jacob_t, jacob_n = gradient_checker.compute_gradient( - tf_x, - shape, - s, [num_segments, num_cols], - x_init_value=np_x, - delta=1) - self.assertAllClose(jacob_t, jacob_n) + values_tf, (9,), s, (3,), x_init_value=values, delta=1) + self.assertAllClose(jacob_t, jacob_n) + self.assertAllClose(jacob_t, grad_gt) def testGradientMatchesSegmentSum(self): # Strategy: compute the gradient for UnsortedSegmentSum and SegmentSum @@ -318,8 +377,7 @@ class UnsortedSegmentSumTest(SegmentReductionHelper): num_cols = 2 shape = [n, num_cols] num_segments = max(indices) + 1 - for dtype in [dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.complex64, - dtypes_lib.complex128]: + for dtype in self.differentiable_dtypes: with self.test_session(use_gpu=True): tf_x, np_x = self._input(shape, dtype=dtype) # Results from UnsortedSegmentSum @@ -353,9 +411,8 @@ class UnsortedSegmentSumTest(SegmentReductionHelper): unsorted.eval() def testEmptySecondDimension(self): - dtypes = [ - np.float32, np.float64, np.int64, np.int32, np.complex64, np.complex128 - ] + dtypes = [np.float16, np.float32, np.float64, np.int64, np.int32, + np.complex64, np.complex128] with self.test_session(use_gpu=True): for dtype in dtypes: for itype in (np.int32, np.int64): @@ -364,36 +421,14 @@ class UnsortedSegmentSumTest(SegmentReductionHelper): unsorted = math_ops.unsorted_segment_sum(data, segment_ids, 2) self.assertAllEqual(unsorted.eval(), np.zeros((2, 0), dtype=dtype)) - def testGradientSegmentMax(self): - num_cols = 2 - indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) - num_segments = max(indices_flat) + 3 - for indices in indices_flat, indices_flat.reshape(5, 2): - shape = indices.shape + (num_cols,) - with self.test_session(use_gpu=True): - tf_x, np_x = self._input(shape, dtype=dtypes_lib.float64) - s = math_ops.unsorted_segment_max( - data=tf_x, segment_ids=indices, num_segments=num_segments) - jacob_t, jacob_n = gradient_checker.compute_gradient( - tf_x, - shape, - s, - [num_segments, num_cols], - x_init_value=np_x.astype(np.double), delta=1) - self.assertAllClose(jacob_t, jacob_n) - def testDropNegatives(self): # Note: the test is done by replacing segment_ids with 8 to -1 # for index and replace values generated by numpy with 0. - dtypes = [ - dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int64, - dtypes_lib.int32, dtypes_lib.complex64, dtypes_lib.complex128 - ] indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) num_segments = 12 for indices in indices_flat, indices_flat.reshape(5, 2): shape = indices.shape + (2,) - for dtype in dtypes: + for dtype in self.all_dtypes: with self.test_session(use_gpu=True): tf_x, np_x = self._input(shape, dtype=dtype) np_ans = self._segmentReduce( @@ -507,6 +542,25 @@ class SparseSegmentReductionOpTest(SparseSegmentReductionHelper): tf_ans = s.eval() self.assertAllClose(np_ans, tf_ans) + def testWithEmptySegments(self): + tf_x = constant_op.constant([], shape=[0, 4], dtype=dtypes_lib.float32) + ops_list = [ + math_ops.sparse_segment_sum_with_num_segments, + math_ops.sparse_segment_mean_with_num_segments + ] + segment_indices = [] + tf_indices = [] + num_segments = 5 + with self.test_session(use_gpu=False): + for tf_op in ops_list: + s = tf_op( + data=tf_x, + indices=tf_indices, + segment_ids=segment_indices, + num_segments=num_segments) + tf_ans = s.eval() + self.assertAllClose(np.zeros([5, 4]), tf_ans) + def testSegmentIdsGreaterThanZero(self): tf_x, np_x = self._input([10, 4], dtype=dtypes_lib.float32) ops_list = [(np.add, None, math_ops.sparse_segment_sum), ( diff --git a/tensorflow/python/kernel_tests/slice_op_test.py b/tensorflow/python/kernel_tests/slice_op_test.py index 051a25080b826de05ee3e24a82fbcd1f47995544..5fc9bef21816e3a12f0d274bab1fc82a83546422 100644 --- a/tensorflow/python/kernel_tests/slice_op_test.py +++ b/tensorflow/python/kernel_tests/slice_op_test.py @@ -283,7 +283,7 @@ class SliceTest(test.TestCase): # unintended behavior is prevented. c = constant_op.constant(5.0) with self.assertRaisesWithPredicateMatch( - TypeError, lambda e: "`Tensor` objects are not iterable" in str(e)): + TypeError, lambda e: "Tensor objects are not iterable" in str(e)): for _ in c: pass diff --git a/tensorflow/python/kernel_tests/softmax_op_test.py b/tensorflow/python/kernel_tests/softmax_op_test.py index bb3f6970e4f18ce174062e2fce488af738b93a50..2b8e99e18e6881143ee77c4f1ec5096635e5c1b2 100644 --- a/tensorflow/python/kernel_tests/softmax_op_test.py +++ b/tensorflow/python/kernel_tests/softmax_op_test.py @@ -18,11 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import sys - import numpy as np -from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors_impl from tensorflow.python.framework import test_util @@ -101,10 +98,10 @@ class SoftmaxTest(test.TestCase): def _testOverflow(self, use_gpu=False): if use_gpu: - type = np.float32 + type = np.float32 # pylint: disable=redefined-builtin else: - type = np.float64 - max = np.finfo(type).max + type = np.float64 # pylint: disable=redefined-builtin + max = np.finfo(type).max # pylint: disable=redefined-builtin features = np.array([[1., 1., 1., 1.], [max, 1., 2., 3.]]).astype(type) with self.test_session(use_gpu=use_gpu): tf_log_softmax = nn_ops.log_softmax(features) @@ -168,7 +165,7 @@ class SoftmaxTest(test.TestCase): def testEmptyInput(self): with self.test_session(): - x = constant_op.constant([[]], shape=[0, 3]) + x = array_ops.placeholder(dtypes.float32, shape=[0, 3]) self.assertEqual(0, array_ops.size(x).eval()) # reshape would raise if logits is empty with self.assertRaises(errors_impl.InvalidArgumentError): diff --git a/tensorflow/python/kernel_tests/spacetobatch_op_test.py b/tensorflow/python/kernel_tests/spacetobatch_op_test.py index b943dfa4e5f2a06eddcb3af03764e5e046b715f4..2a9232b6aecb66328f10a62f2251246c4fcec6e6 100644 --- a/tensorflow/python/kernel_tests/spacetobatch_op_test.py +++ b/tensorflow/python/kernel_tests/spacetobatch_op_test.py @@ -86,11 +86,11 @@ class CppOpImpl(object): @staticmethod def space_to_batch(*args, **kwargs): - return gen_array_ops._space_to_batch(*args, **kwargs) + return gen_array_ops.space_to_batch(*args, **kwargs) @staticmethod def batch_to_space(*args, **kwargs): - return gen_array_ops._batch_to_space(*args, **kwargs) + return gen_array_ops.batch_to_space(*args, **kwargs) class SpaceToBatchTest(test.TestCase, PythonOpImpl): diff --git a/tensorflow/python/kernel_tests/spacetodepth_op_test.py b/tensorflow/python/kernel_tests/spacetodepth_op_test.py index 3c98a685e07a1f2d55c3c1035a99ffaa593d35b3..cd90d16aacb4325ed426b0466266d9616b574401 100644 --- a/tensorflow/python/kernel_tests/spacetodepth_op_test.py +++ b/tensorflow/python/kernel_tests/spacetodepth_op_test.py @@ -34,8 +34,8 @@ from tensorflow.python.platform import tf_logging class SpaceToDepthTest(test.TestCase): - def _testOne(self, inputs, block_size, outputs): - input_nhwc = math_ops.to_float(inputs) + def _testOne(self, inputs, block_size, outputs, dtype=dtypes.float32): + input_nhwc = math_ops.cast(inputs, dtype) with self.test_session(use_gpu=False): # test NHWC (default) on CPU x_tf = array_ops.space_to_depth(input_nhwc, block_size) @@ -58,6 +58,12 @@ class SpaceToDepthTest(test.TestCase): x_out = [[[[1, 2, 3, 4]]]] self._testOne(x_np, block_size, x_out) + def testBasicFloat16(self): + x_np = [[[[1], [2]], [[3], [4]]]] + block_size = 2 + x_out = [[[[1, 2, 3, 4]]]] + self._testOne(x_np, block_size, x_out, dtype=dtypes.float16) + # Tests for larger input dimensions. To make sure elements are # correctly ordered spatially. def testLargerInput2x2(self): @@ -126,6 +132,24 @@ class SpaceToDepthTest(test.TestCase): x_out = [batch_output_elt(i) for i in range(batch_size)] self._testOne(x_np, block_size, x_out) + def testBatchSize0(self): + block_size = 2 + batch_size = 0 + input_nhwc = array_ops.ones([batch_size, 4, 6, 3]) + x_out = array_ops.ones([batch_size, 2, 3, 12]) + + with self.test_session(use_gpu=False): + # test NHWC (default) on CPU + x_tf = array_ops.space_to_depth(input_nhwc, block_size) + self.assertAllEqual(x_tf.shape, x_out.shape) + x_tf.eval() + if test.is_gpu_available(): + with self.test_session(use_gpu=True): + # test NHWC (default) on GPU + x_tf = array_ops.space_to_depth(input_nhwc, block_size) + self.assertAllEqual(x_tf.shape, x_out.shape) + x_tf.eval() + # Tests for different width and height. def testNonSquare(self): x_np = [[[[1, 10], [2, 20]], [[3, 30], [4, 40]], [[5, 50], [6, 60]], diff --git a/tensorflow/python/kernel_tests/sparse_xent_op_test.py b/tensorflow/python/kernel_tests/sparse_xent_op_test.py index cd5b711a0ed18aabff543aa4b6ecb1a885618caf..a841fe83a7f585a69ef33c437570359797484a4a 100644 --- a/tensorflow/python/kernel_tests/sparse_xent_op_test.py +++ b/tensorflow/python/kernel_tests/sparse_xent_op_test.py @@ -64,7 +64,7 @@ class SparseXentTest(test.TestCase): def _testXent(self, np_features, np_labels): np_loss, np_backprop = self._npXent(np_features, np_labels) with self.test_session(use_gpu=True) as sess: - loss, backprop = gen_nn_ops._sparse_softmax_cross_entropy_with_logits( + loss, backprop = gen_nn_ops.sparse_softmax_cross_entropy_with_logits( np_features, np_labels) tf_loss, tf_backprop = sess.run([loss, backprop]) self.assertAllCloseAccordingToType(np_loss, tf_loss) @@ -73,7 +73,7 @@ class SparseXentTest(test.TestCase): def testSingleClass(self): for label_dtype in np.int32, np.int64: with self.test_session(use_gpu=True) as sess: - loss, backprop = gen_nn_ops._sparse_softmax_cross_entropy_with_logits( + loss, backprop = gen_nn_ops.sparse_softmax_cross_entropy_with_logits( np.array([[1.], [-1.], [0.]]).astype(np.float32), np.array([0, 0, 0]).astype(label_dtype)) tf_loss, tf_backprop = sess.run([loss, backprop]) @@ -87,8 +87,9 @@ class SparseXentTest(test.TestCase): if test.is_built_with_cuda() and test.is_gpu_available(): with self.test_session(use_gpu=True) as sess: - loss, backprop = (gen_nn_ops._sparse_softmax_cross_entropy_with_logits( - features, labels)) + loss, backprop = ( + gen_nn_ops.sparse_softmax_cross_entropy_with_logits( + features, labels)) tf_loss, tf_backprop = sess.run([loss, backprop]) self.assertAllClose( [[np.nan] * 4, [0.25, 0.25, 0.25, -0.75], @@ -100,8 +101,8 @@ class SparseXentTest(test.TestCase): [np.nan, 1.3862, 3.4420, np.nan], tf_loss, rtol=1e-3, atol=1e-3) with self.test_session(use_gpu=False) as sess: - loss, backprop = (gen_nn_ops._sparse_softmax_cross_entropy_with_logits( - features, labels)) + loss, backprop = ( + gen_nn_ops.sparse_softmax_cross_entropy_with_logits(features, labels)) with self.assertRaisesOpError("Received a label value of"): sess.run([loss, backprop]) diff --git a/tensorflow/python/kernel_tests/split_op_test.py b/tensorflow/python/kernel_tests/split_op_test.py index 6171793b148f8d8f195b9548a13df89d29c5e96e..8cfee3eb933afcea7a58d5632948b87b0c4c10df 100644 --- a/tensorflow/python/kernel_tests/split_op_test.py +++ b/tensorflow/python/kernel_tests/split_op_test.py @@ -336,6 +336,20 @@ class SplitOpTest(test.TestCase): for s in splits: self.assertEqual(None, s.get_shape().ndims) + def testNonexistentDimTensor(self): + x = array_ops.placeholder(dtypes.int32) + values = np.zeros([5, 30]) + splits = array_ops.placeholder(dtypes.int32) + with self.assertRaisesRegexp(ValueError, "Cannot infer"): + y = array_ops.split(values, splits, axis=x) + + splits = array_ops.placeholder(dtypes.int32, [3]) + y = array_ops.split(values, splits, axis=x) + with self.test_session(use_gpu=True) as sess: + with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, + "must have exactly one element"): + sess.run(y, {x: np.array([], dtype=np.int32), splits: [4, 11, 15]}) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/stack_op_test.py b/tensorflow/python/kernel_tests/stack_op_test.py index 347baf81148e9b747a9be4849912d154b220a084..2f27d1839b2218d0cc33d7278116186548ad3420 100644 --- a/tensorflow/python/kernel_tests/stack_op_test.py +++ b/tensorflow/python/kernel_tests/stack_op_test.py @@ -50,7 +50,7 @@ class StackOpTest(test.TestCase): # Convert [data[0], data[1], ...] separately to tensorflow # TODO(irving): Remove list() once we handle maps correctly xs = list(map(constant_op.constant, data)) - # Pack back into a single tensorflow tensor + # Stack back into a single tensorflow tensor c = array_ops.stack(xs) self.assertAllEqual(c.eval(), data) @@ -78,7 +78,7 @@ class StackOpTest(test.TestCase): for shape in (2,), (3,), (2, 3), (3, 2), (4, 3, 2): for dtype in [np.bool, np.float32, np.int32, np.int64]: data = np.random.randn(*shape).astype(dtype) - # Pack back into a single tensorflow tensor directly using np array + # Stack back into a single tensorflow tensor directly using np array c = array_ops.stack(data) # This is implemented via a Const: self.assertEqual(c.op.type, "Const") @@ -223,7 +223,7 @@ class StackOpTest(test.TestCase): array_ops.stack(t, axis=-3) -class AutomaticPackingTest(test.TestCase): +class AutomaticStackingTest(test.TestCase): def testSimple(self): with self.test_session(use_gpu=True): diff --git a/tensorflow/python/kernel_tests/stack_ops_test.py b/tensorflow/python/kernel_tests/stack_ops_test.py index aa409336f5c50178e4d0ca946190119fb0e4188e..afd2eaffab992bca4b3ae7b4f65e0370f325b548 100644 --- a/tensorflow/python/kernel_tests/stack_ops_test.py +++ b/tensorflow/python/kernel_tests/stack_ops_test.py @@ -34,11 +34,11 @@ class StackOpTest(test.TestCase): def _testStackPushPop(self, use_gpu): with self.test_session(use_gpu=use_gpu): - h = gen_data_flow_ops._stack_v2( + h = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push_v2(h, [[4.0, 5.0]]) + c = gen_data_flow_ops.stack_push_v2(h, [[4.0, 5.0]]) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32) self.assertAllClose([[4.0, 5.0]], c1.eval()) def testStackPushPop(self): @@ -49,11 +49,11 @@ class StackOpTest(test.TestCase): with self.test_session(use_gpu=use_gpu): a = np.arange(2000) x = constant_op.constant(a, dtype=dtypes.float32) - h = gen_data_flow_ops._stack_v2( + h = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push_v2(h, x, swap_memory=True) + c = gen_data_flow_ops.stack_push_v2(h, x, swap_memory=True) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32) self.assertAllClose(a, c1.eval()) def testStackPushPopSwap(self): @@ -63,7 +63,7 @@ class StackOpTest(test.TestCase): def _testStackWhileSwap(self, use_gpu): with self.test_session(use_gpu=use_gpu): n = constant_op.constant(0) - h = gen_data_flow_ops._stack_v2( + h = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") def c(x): @@ -72,7 +72,7 @@ class StackOpTest(test.TestCase): def b(x): with ops.control_dependencies([x]): a = constant_op.constant(np.ones(2000), dtype=dtypes.float32) - v = gen_data_flow_ops._stack_push_v2(h, a, swap_memory=True) + v = gen_data_flow_ops.stack_push_v2(h, a, swap_memory=True) with ops.control_dependencies([v]): return math_ops.add(x, 1) @@ -86,7 +86,7 @@ class StackOpTest(test.TestCase): def b1(x, y): nx = math_ops.subtract(x, 1) - ny = y + gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + ny = y + gen_data_flow_ops.stack_pop_v2(h, dtypes.float32) return [nx, ny] _, ry = control_flow_ops.while_loop( @@ -99,16 +99,16 @@ class StackOpTest(test.TestCase): def _testMultiStack(self, use_gpu): with self.test_session(use_gpu=use_gpu): - h1 = gen_data_flow_ops._stack_v2( + h1 = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_push_v2(h1, 4.0) + c1 = gen_data_flow_ops.stack_push_v2(h1, 4.0) with ops.control_dependencies([c1]): - c1 = gen_data_flow_ops._stack_pop_v2(h1, dtypes.float32) - h2 = gen_data_flow_ops._stack_v2( + c1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32) + h2 = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="bar") - c2 = gen_data_flow_ops._stack_push_v2(h2, 5.0) + c2 = gen_data_flow_ops.stack_push_v2(h2, 5.0) with ops.control_dependencies([c2]): - c2 = gen_data_flow_ops._stack_pop_v2(h2, dtypes.float32) + c2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32) r = c1 + c2 self.assertAllClose(9.0, r.eval()) @@ -119,17 +119,17 @@ class StackOpTest(test.TestCase): def _testSameNameStacks(self, use_gpu): """Different stacks with the same name do not interfere.""" with self.test_session(use_gpu=use_gpu) as sess: - h1 = gen_data_flow_ops._stack_v2( + h1 = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - h2 = gen_data_flow_ops._stack_v2( + h2 = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_push_v2(h1, 4.0) + c1 = gen_data_flow_ops.stack_push_v2(h1, 4.0) with ops.control_dependencies([c1]): - c2 = gen_data_flow_ops._stack_push_v2(h2, 5.0) + c2 = gen_data_flow_ops.stack_push_v2(h2, 5.0) with ops.control_dependencies([c2]): - pop1 = gen_data_flow_ops._stack_pop_v2(h1, dtypes.float32) - pop2 = gen_data_flow_ops._stack_pop_v2(h2, dtypes.float32) + pop1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32) + pop2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32) out1, out2 = sess.run([pop1, pop2]) self.assertAllClose(out1, 4.0) @@ -141,9 +141,9 @@ class StackOpTest(test.TestCase): def _testCloseStack(self, use_gpu): with self.test_session(use_gpu=use_gpu) as sess: - h = gen_data_flow_ops._stack_v2( + h = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_close_v2(h) + c1 = gen_data_flow_ops.stack_close_v2(h) sess.run(c1) def testCloseStack(self): @@ -152,11 +152,11 @@ class StackOpTest(test.TestCase): def _testPushCloseStack(self, use_gpu): with self.test_session(use_gpu=use_gpu) as sess: - h = gen_data_flow_ops._stack_v2( + h = gen_data_flow_ops.stack_v2( -1, elem_type=dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push_v2(h, [[4.0, 5.0]]) + c = gen_data_flow_ops.stack_push_v2(h, [[4.0, 5.0]]) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_close_v2(h) + c1 = gen_data_flow_ops.stack_close_v2(h) sess.run(c1) def testPushCloseStack(self): @@ -170,9 +170,9 @@ class StackOpRefTest(test.TestCase): def _testStackPushPop(self, use_gpu): with self.test_session(use_gpu=use_gpu): h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push(h, [[4.0, 5.0]]) + c = gen_data_flow_ops.stack_push(h, [[4.0, 5.0]]) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_pop(h, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop(h, dtypes.float32) self.assertAllClose([[4.0, 5.0]], c1.eval()) def testStackPushPop(self): @@ -184,9 +184,9 @@ class StackOpRefTest(test.TestCase): a = np.arange(2000) x = constant_op.constant(a, dtype=dtypes.float32) h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push(h, x, swap_memory=True) + c = gen_data_flow_ops.stack_push(h, x, swap_memory=True) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_pop(h, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop(h, dtypes.float32) self.assertAllClose(a, c1.eval()) def testStackPushPopSwap(self): @@ -196,13 +196,13 @@ class StackOpRefTest(test.TestCase): def _testMultiStack(self, use_gpu): with self.test_session(use_gpu=use_gpu): h1 = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_push(h1, 4.0) + c1 = gen_data_flow_ops.stack_push(h1, 4.0) with ops.control_dependencies([c1]): - c1 = gen_data_flow_ops._stack_pop(h1, dtypes.float32) + c1 = gen_data_flow_ops.stack_pop(h1, dtypes.float32) h2 = gen_data_flow_ops._stack(dtypes.float32, stack_name="bar") - c2 = gen_data_flow_ops._stack_push(h2, 5.0) + c2 = gen_data_flow_ops.stack_push(h2, 5.0) with ops.control_dependencies([c2]): - c2 = gen_data_flow_ops._stack_pop(h2, dtypes.float32) + c2 = gen_data_flow_ops.stack_pop(h2, dtypes.float32) r = c1 + c2 self.assertAllClose(9.0, r.eval()) @@ -217,7 +217,7 @@ class StackOpRefTest(test.TestCase): def b(x): with ops.control_dependencies([x]): a = constant_op.constant(np.ones(2000), dtype=dtypes.float32) - v = gen_data_flow_ops._stack_push(h, a, swap_memory=True) + v = gen_data_flow_ops.stack_push(h, a, swap_memory=True) with ops.control_dependencies([v]): return math_ops.add(x, 1) @@ -231,7 +231,7 @@ class StackOpRefTest(test.TestCase): def b1(x, y): nx = math_ops.subtract(x, 1) - ny = y + gen_data_flow_ops._stack_pop(h, dtypes.float32) + ny = y + gen_data_flow_ops.stack_pop(h, dtypes.float32) return [nx, ny] _, ry = control_flow_ops.while_loop( @@ -249,9 +249,9 @@ class StackOpRefTest(test.TestCase): def _testSameNameStacks(self, use_gpu): with self.test_session(use_gpu=use_gpu): h1 = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_push(h1, 4.0) + c1 = gen_data_flow_ops.stack_push(h1, 4.0) h2 = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c2 = gen_data_flow_ops._stack_push(h2, 5.0) + c2 = gen_data_flow_ops.stack_push(h2, 5.0) _ = c1 + c2 self.assertNotEqual(h1.eval()[1], h2.eval()[1]) @@ -262,7 +262,7 @@ class StackOpRefTest(test.TestCase): def _testCloseStack(self, use_gpu): with self.test_session(use_gpu=use_gpu) as sess: h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c1 = gen_data_flow_ops._stack_close(h) + c1 = gen_data_flow_ops.stack_close(h) sess.run(c1) def testCloseStack(self): @@ -272,9 +272,9 @@ class StackOpRefTest(test.TestCase): def _testPushCloseStack(self, use_gpu): with self.test_session(use_gpu=use_gpu) as sess: h = gen_data_flow_ops._stack(dtypes.float32, stack_name="foo") - c = gen_data_flow_ops._stack_push(h, [[4.0, 5.0]]) + c = gen_data_flow_ops.stack_push(h, [[4.0, 5.0]]) with ops.control_dependencies([c]): - c1 = gen_data_flow_ops._stack_close(h) + c1 = gen_data_flow_ops.stack_close(h) sess.run(c1) def testPushCloseStack(self): diff --git a/tensorflow/python/kernel_tests/template_test.py b/tensorflow/python/kernel_tests/template_test.py index 8792ab41a07aac2dc8c9fcb956c378054a309a41..1b935d5286729e9e802c56e90e2ae7ab72a6e080 100644 --- a/tensorflow/python/kernel_tests/template_test.py +++ b/tensorflow/python/kernel_tests/template_test.py @@ -17,10 +17,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools import traceback from tensorflow.python.client import session from tensorflow.python.eager import context +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 @@ -354,6 +356,10 @@ class TemplateTest(test.TestCase): self.assertEqual("s1_1/nested/dummy:0", v5.name) self.assertEqual("s1_1/nested_1/dummy:0", v6.name) + self.assertEqual(2, len(tmpl1._checkpoint_dependencies)) + self.assertEqual("nested", tmpl1._checkpoint_dependencies[0].name) + self.assertEqual("nested_1", tmpl1._checkpoint_dependencies[1].name) + @test_util.run_in_graph_and_eager_modes() def test_nested_templates_with_defun(self): @@ -412,6 +418,17 @@ class TemplateTest(test.TestCase): self.assertEqual("s1_1/nested/dummy:0", v3[0].name) self.assertEqual("s1_1/nested_1/dummy:0", v3[1].name) + def test_graph_function_no_name(self): + with context.eager_mode(): + + def f(_, y): + return y + 1 + + partial = functools.partial(f, 1.0) + tmpl = template.make_template_internal( + "a", partial, create_graph_function_=True) + self.assertAllEqual(tmpl(ops.convert_to_tensor(1.0)), 2.0) + @test_util.run_in_graph_and_eager_modes() def test_immediate_scope_creation(self): # Create templates in scope a then call in scope b. make_template should @@ -545,7 +562,7 @@ class TemplateTest(test.TestCase): outputs_b, _ = linear1(inputs) self.assertEquals("foo", linear1.variable_scope.name) self.assertEquals("foo/w:0", w1.name) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEquals("foo/add:0", outputs_a.name, "First application of template should get " "same name scope as variables.") @@ -560,7 +577,7 @@ class TemplateTest(test.TestCase): "New template gets a freshly uniquified variable scope " "because 'foo' is already taken.") self.assertEquals("foo_1/w:0", w2.name) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEquals("foo_1_1/add:0", outputs_c.name, "First application of template would get " "same name scope as variables, but 'foo_1' is already " @@ -575,7 +592,7 @@ class TemplateTest(test.TestCase): with variable_scope.variable_scope("foo"): # Create two templates with the same name, ensure scopes are made unique. ta = template.make_template("bar", variable_scoped_function, True) - if context.in_eager_mode(): + if context.executing_eagerly(): tb = template.make_template("s", function_with_side_create, trainable=False) else: diff --git a/tensorflow/python/kernel_tests/tensor_array_ops_test.py b/tensorflow/python/kernel_tests/tensor_array_ops_test.py index aad2443eea7ad87faf481973e91ca3df32ccfb44..a834675828b67aed4057d1857c546a586cee69c9 100644 --- a/tensorflow/python/kernel_tests/tensor_array_ops_test.py +++ b/tensorflow/python/kernel_tests/tensor_array_ops_test.py @@ -399,28 +399,14 @@ class TensorArrayTest(test.TestCase): def testTensorArrayWriteWrongIndexOrDataTypeFails(self): with self.test_session(use_gpu=True): ta = _make_ta(3, "foo", dtype=dtypes.float32) - in_graph_mode = context.in_graph_mode() # Test writing the wrong datatype - if in_graph_mode: - with self.assertRaisesOpError( - "TensorArray dtype is float but Op is trying to write " - "dtype string"): - self.evaluate(ta.write(0, "wrong_type_scalar").flow) - else: - with self.assertRaisesOpError( - "TensorArray dtype is float32 but Op is trying to write " - "dtype string"): - self.evaluate(ta.write(0, "wrong_type_scalar").flow) + with self.assertRaisesOpError( + "TensorArray dtype is (float|float32) but Op is trying to write " + "dtype string"): + self.evaluate(ta.write(0, "wrong_type_scalar").flow) - if context.in_graph_mode(): - with self.assertRaisesOpError( - "Tried to write to index -1 but array is not " - "resizeable and size is: 3"): - self.evaluate(ta.write(-1, 3.0).flow) - else: - with self.assertRaisesOpError( - r"Writing to negative indices \(index -1\) is not allowed."): - self.evaluate(ta.write(-1, 3.0).flow) + with self.assertRaisesOpError("index -1"): + self.evaluate(ta.write(-1, 3.0).flow) # Test reading from too large an index with self.assertRaisesOpError( @@ -435,23 +421,17 @@ class TensorArrayTest(test.TestCase): w0 = ta.write(0, [[4.0, 5.0]]) - # Test reading wrong datatype, which is only possible in graph mode - if context.in_graph_mode(): - r0_bad = gen_data_flow_ops._tensor_array_read_v3( + # Test reading wrong datatype (only possible when constructing graphs). + if not context.executing_eagerly(): + r0_bad = gen_data_flow_ops.tensor_array_read_v3( handle=w0.handle, index=0, dtype=dtypes.float64, flow_in=w0.flow) with self.assertRaisesOpError( "TensorArray dtype is float but Op requested dtype double."): r0_bad.eval() # Test reading from a negative index, which is not allowed - if context.in_graph_mode(): - with self.assertRaisesOpError( - r"Tried to read from index -1 but array size is: 3"): - self.evaluate(ta.read(-1)) - else: - with self.assertRaisesOpError( - r"Reading from negative indices \(index -1\) is not allowed."): - self.evaluate(ta.read(-1)) + with self.assertRaisesOpError("index -1"): + self.evaluate(ta.read(-1)) # Test reading from too large an index with self.assertRaisesOpError( @@ -467,10 +447,7 @@ class TensorArrayTest(test.TestCase): with self.assertRaisesOpError( "Could not write to TensorArray index 2 because " "it has already been written to."): - if context.in_graph_mode(): - self.evaluate(ta.write(2, 3.0).write(2, 3.0).flow) - else: - self.evaluate(ta.write(2, 3.0).write(2, 3.0)) + self.evaluate(ta.write(2, 3.0).write(2, 3.0).flow) @test_util.run_in_graph_and_eager_modes() def testTensorArrayConcatIncompatibleShapesFails(self): @@ -499,58 +476,40 @@ class TensorArrayTest(test.TestCase): w2 = w1.write(1, [4.0]) w3 = w2.write(2, [[3.0]]) - # The eager-mode implementation just passes up array_op.concat's error - # message. - if context.in_graph_mode(): - with self.assertRaisesOpError( - r"TensorArray has inconsistent shapes. Index 0 has " - r"\(excepting dimension 0\) shape: \[\] but index 2 has " - r"\(excepting dimension 0\) shape: \[1\]"): - self.evaluate(w3.concat()) - else: - with self.assertRaisesOpError( - r".*Ranks of all input tensors should match: shape\[0\] " - r"= \[1\] vs\. shape\[2\] = \[1,1\].*"): - self.evaluate(w3.concat()) + # The exact error messages differ between eager execution and graph + # construction as the former bubbles up the error from array_op.concat. + with self.assertRaisesOpError("shape"): + self.evaluate(w3.concat()) @test_util.run_in_graph_and_eager_modes() def testTensorArraySplitIncompatibleShapesFails(self): with self.test_session(use_gpu=True): - in_graph_mode = context.in_graph_mode() + in_eager_mode = context.executing_eagerly() ta = _make_ta(3, "foo") with self.assertRaisesOpError( r"Expected lengths to be a vector, received shape: \[\]"): - if in_graph_mode: + if in_eager_mode: + self.evaluate(ta.split([1.0, 2.0, 3.0], 1)) + else: lengths = array_ops.placeholder(dtypes.int64) ta.split([1.0, 2.0, 3.0], lengths).flow.eval(feed_dict={lengths: 1}) - else: - self.evaluate(ta.split([1.0, 2.0, 3.0], 1)) with self.assertRaisesOpError( r"Expected sum of lengths to be equal to values.shape\[0\], " r"but sum of lengths is 1 and value's shape is: \[3\]"): - if in_graph_mode: - self.evaluate(ta.split([1.0, 2.0, 3.0], [1]).flow) - else: - self.evaluate(ta.split([1.0, 2.0, 3.0], [1])) + self.evaluate(ta.split([1.0, 2.0, 3.0], [1]).flow) ta = _make_ta(1, "baz") with self.assertRaisesOpError( r"Expected value to be at least a vector, but received shape: \[\]"): - if in_graph_mode: - self.evaluate(ta.split(1.0, [1]).flow) - else: - self.evaluate(ta.split(1.0, [1])) + self.evaluate(ta.split(1.0, [1]).flow) ta = _make_ta(2, "buz") with self.assertRaisesOpError( r"TensorArray's size is not equal to the size of lengths " r"\(2 vs. 1\), and the TensorArray is not marked as " r"dynamically resizeable"): - if in_graph_mode: - self.evaluate(ta.split([1.0], [1]).flow) - else: - self.evaluate(ta.split([1.0], [1])) + self.evaluate(ta.split([1.0], [1]).flow) def _testTensorArrayWriteGradientAddMultipleAdds(self, dtype): with self.test_session(use_gpu=True): @@ -868,14 +827,14 @@ class TensorArrayTest(test.TestCase): vout = func(v0, state0, var) grad_val = -np.arange(3 * 5, dtype=np_dtype).reshape(3, 5) - if context.in_graph_mode(): + if context.executing_eagerly(): + grad_fn = backprop.gradients_function(func) + v0_grad, state0_grad, var_grad = grad_fn(v0, state0, var, dy=grad_val) + else: v0_grad = gradients_impl.gradients([vout], [v0], [grad_val])[0] state0_grad = gradients_impl.gradients([vout], [state0], [grad_val])[0] var_grad = gradients_impl.gradients([vout], [var], [grad_val])[0] variables.global_variables_initializer().run() - else: - grad_fn = backprop.gradients_function(func) - v0_grad, state0_grad, var_grad = grad_fn(v0, state0, var, dy=grad_val) state0_t, var_t, v0_t, vout_t, v0_grad_t, var_grad_t, state0_grad_t = ( self.evaluate( @@ -959,10 +918,10 @@ class TensorArrayTest(test.TestCase): return r x = constant_op.constant(2.0, name="x") - if context.in_graph_mode(): - grad = gradients_impl.gradients(loop(x), [x])[0] - else: + if context.executing_eagerly(): grad = backprop.gradients_function(loop)(x)[0] + else: + grad = gradients_impl.gradients(loop(x), [x])[0] self.assertAllClose(31.0, self.evaluate(grad)) def testSumOfTwoReadVariablesWithoutRepeatGrad(self): @@ -1158,14 +1117,14 @@ class TensorArrayTest(test.TestCase): infer_shape=True) w0 = ta1.split(value, [1, 2]) r0 = w0.read(0) - if context.in_graph_mode(): + if context.executing_eagerly(): + self.assertEqual((1, 2), r0.get_shape()) + self.assertEqual((2, 2), w0.read(1).get_shape()) + else: self.assertEqual(r0.get_shape().ndims, None) self.assertEqual( tensor_shape.TensorShape( ta1.handle.op.get_attr("element_shape")).ndims, None) - else: - self.assertEqual((1, 2), r0.get_shape()) - self.assertEqual((2, 2), w0.read(1).get_shape()) def testWriteUnknownShape(self): with self.test_session(use_gpu=True): @@ -1297,13 +1256,13 @@ class TensorArrayTest(test.TestCase): g = func(values) grad_ys = [[[2.0, 3.0], [4.0, 5.0]]] # Test combined gradients + aggregation of read(0) - if context.in_graph_mode(): - grad = gradients_impl.gradients(ys=[g], xs=[values], grad_ys=grad_ys) - g_vals, grad_vals = session.run([[g], grad]) - else: + if context.executing_eagerly(): g_vals = [g] grad_vals = backprop.gradients_function(func)( values, dy=constant_op.constant(grad_ys[0], dtype=dtypes.float32)) + else: + grad = gradients_impl.gradients(ys=[g], xs=[values], grad_ys=grad_ys) + g_vals, grad_vals = session.run([[g], grad]) # Gradients for 8 of the 10 unread components are zero. expected_grad = np.zeros((10, 2)) @@ -1453,13 +1412,13 @@ class TensorArrayTest(test.TestCase): # Tests correct properties on new TensorArrays. self.assertEqual(dtypes.float32, ta0.dtype) self.assertEqual(dtypes.int32, ta1.dtype) - if context.in_graph_mode(): - self.assertEqual(tensor_shape.unknown_shape(), read0.get_shape()) + if context.executing_eagerly(): + self.assertEqual(tensor_shape.scalar(), read0.get_shape()) else: - self.assertEqual(tensor_shape.scalar(), read1.get_shape()) + self.assertEqual(tensor_shape.unknown_shape(), read0.get_shape()) self.assertEqual(tensor_shape.scalar(), read1.get_shape()) - if context.in_graph_mode(): + if not context.executing_eagerly(): variables.global_variables_initializer().run() read0_v, read1_v, size0_v, size1_v = self.evaluate((read0, read1, size0, diff --git a/tensorflow/python/kernel_tests/tensordot_op_test.py b/tensorflow/python/kernel_tests/tensordot_op_test.py index 084f98581ea843d935c4916de206d93fead8ae2d..8ad29afd0a0f2e7fbaaf2bde956326e578466b1d 100644 --- a/tensorflow/python/kernel_tests/tensordot_op_test.py +++ b/tensorflow/python/kernel_tests/tensordot_op_test.py @@ -105,7 +105,7 @@ class TensordotTest(test_lib.TestCase): self.assertAllEqual(tf_ans, np_ans) def test_partial_shape_inference(self): - for axes in ([1],[0]), 1: + for axes in ([1], [0]), 1: a = array_ops.placeholder(dtypes.float32) b = array_ops.placeholder(dtypes.float32) output = math_ops.tensordot(a, b, axes) diff --git a/tensorflow/python/kernel_tests/topk_op_test.py b/tensorflow/python/kernel_tests/topk_op_test.py index efb5b9f3641ceaebf1fd5285486b4a9bb93615cf..6ab931fdb97a8945ab610fda27a036693f0291e5 100644 --- a/tensorflow/python/kernel_tests/topk_op_test.py +++ b/tensorflow/python/kernel_tests/topk_op_test.py @@ -58,7 +58,7 @@ class TopKTest(test.TestCase): # Do some special casing of equality of indices: if indices # are not the same, but values are floating type, ensure that # the values are within epsilon of each other. - if not np.issubdtype(np_expected_values.dtype, np.float): + if not np.issubdtype(np_expected_values.dtype, np.floating): # Values are not floating point type; check indices exactly self.assertAllEqual(np_expected_indices, indices) else: diff --git a/tensorflow/python/kernel_tests/unique_op_test.py b/tensorflow/python/kernel_tests/unique_op_test.py index 6366d2e181c8cfabba8a78b664c25c85debc67ef..bbc040dc13fc151b970f130eeb76fa1639245416 100644 --- a/tensorflow/python/kernel_tests/unique_op_test.py +++ b/tensorflow/python/kernel_tests/unique_op_test.py @@ -66,9 +66,9 @@ class UniqueTest(test.TestCase): for dtype in [np.int32, np.int64]: x = np.array([[1, 0, 0], [1, 0, 0], [2, 0, 0]]) with self.test_session() as sess: - y0, idx0 = gen_array_ops._unique_v2(x, axis=np.array([0], dtype)) + y0, idx0 = gen_array_ops.unique_v2(x, axis=np.array([0], dtype)) tf_y0, tf_idx0 = sess.run([y0, idx0]) - y1, idx1 = gen_array_ops._unique_v2(x, axis=np.array([1], dtype)) + y1, idx1 = gen_array_ops.unique_v2(x, axis=np.array([1], dtype)) tf_y1, tf_idx1 = sess.run([y1, idx1]) self.assertAllEqual(tf_y0, np.array([[1, 0, 0], [2, 0, 0]])) self.assertAllEqual(tf_idx0, np.array([0, 0, 1])) @@ -80,7 +80,7 @@ class UniqueTest(test.TestCase): # by default, the axis will be wrapped to allow `axis=None`. x = np.random.randint(2, high=10, size=7000) with self.test_session() as sess: - y, idx = gen_array_ops._unique_v2(x, axis=np.array([], np.int32)) + y, idx = gen_array_ops.unique_v2(x, axis=np.array([], np.int32)) tf_y, tf_idx = sess.run([y, idx]) self.assertEqual(len(x), len(tf_idx)) @@ -133,6 +133,39 @@ class UniqueWithCountsTest(test.TestCase): v = [1 if x[i] == value.decode('ascii') else 0 for i in range(7000)] self.assertEqual(count, sum(v)) + def testInt32Axis(self): + for dtype in [np.int32, np.int64]: + x = np.array([[1, 0, 0], [1, 0, 0], [2, 0, 0]]) + with self.test_session() as sess: + y0, idx0, count0 = gen_array_ops.unique_with_counts_v2( + x, axis=np.array([0], dtype)) + tf_y0, tf_idx0, tf_count0 = sess.run([y0, idx0, count0]) + y1, idx1, count1 = gen_array_ops.unique_with_counts_v2( + x, axis=np.array([1], dtype)) + tf_y1, tf_idx1, tf_count1 = sess.run([y1, idx1, count1]) + self.assertAllEqual(tf_y0, np.array([[1, 0, 0], [2, 0, 0]])) + self.assertAllEqual(tf_idx0, np.array([0, 0, 1])) + self.assertAllEqual(tf_count0, np.array([2, 1])) + self.assertAllEqual(tf_y1, np.array([[1, 0], [1, 0], [2, 0]])) + self.assertAllEqual(tf_idx1, np.array([0, 1, 1])) + self.assertAllEqual(tf_count1, np.array([1, 2])) + + def testInt32V2(self): + # This test is only temporary, once V2 is used + # by default, the axis will be wrapped to allow `axis=None`. + x = np.random.randint(2, high=10, size=7000) + with self.test_session() as sess: + y, idx, count = gen_array_ops.unique_with_counts_v2( + x, axis=np.array([], np.int32)) + tf_y, tf_idx, tf_count = sess.run([y, idx, count]) + + self.assertEqual(len(x), len(tf_idx)) + self.assertEqual(len(tf_y), len(np.unique(x))) + for i in range(len(x)): + self.assertEqual(x[i], tf_y[tf_idx[i]]) + for value, count in zip(tf_y, tf_count): + self.assertEqual(count, np.sum(x == value)) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/kernel_tests/unstack_op_test.py b/tensorflow/python/kernel_tests/unstack_op_test.py index 84818755766a435c873f30e96dc0080af4f78b84..1ee6e0866a6b1c7a9b641a95403d45213f5dc0b4 100644 --- a/tensorflow/python/kernel_tests/unstack_op_test.py +++ b/tensorflow/python/kernel_tests/unstack_op_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Functional tests for Unpack Op.""" +"""Functional tests for Unstack Op.""" from __future__ import absolute_import from __future__ import division @@ -49,7 +49,7 @@ class UnstackOpTest(test.TestCase): data = np.random.randn(*shape).astype(dtype) # Convert data to a single tensorflow tensor x = constant_op.constant(data) - # Unpack into a list of tensors + # Unstack into a list of tensors cs = array_ops.unstack(x, num=shape[0]) self.assertEqual(type(cs), list) self.assertEqual(len(cs), shape[0]) @@ -66,7 +66,7 @@ class UnstackOpTest(test.TestCase): data = np.random.randn(*shape).astype(dtype) # Convert data to a single tensorflow tensor x = constant_op.constant(data) - # Unpack into a list of tensors + # Unstack into a list of tensors cs = array_ops.unstack(x, num=shape[0]) self.assertEqual(type(cs), list) self.assertEqual(len(cs), shape[0]) diff --git a/tensorflow/python/kernel_tests/variable_ops_test.py b/tensorflow/python/kernel_tests/variable_ops_test.py index 79071029fd42374964d12f513e9c510bdc7400eb..cf369c071813120fef685b7220292d50b966cf11 100644 --- a/tensorflow/python/kernel_tests/variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variable_ops_test.py @@ -165,26 +165,26 @@ class VariableOpTest(test.TestCase): def testTemporaryVariable(self): with self.test_session(use_gpu=True): - var = gen_state_ops._temporary_variable( + var = gen_state_ops.temporary_variable( [1, 2], dtypes.float32, var_name="foo") var = state_ops.assign(var, [[4.0, 5.0]]) var = state_ops.assign_add(var, [[6.0, 7.0]]) - final = gen_state_ops._destroy_temporary_variable(var, var_name="foo") + final = gen_state_ops.destroy_temporary_variable(var, var_name="foo") self.assertAllClose([[10.0, 12.0]], final.eval()) def testDestroyNonexistentTemporaryVariable(self): with self.test_session(use_gpu=True): - var = gen_state_ops._temporary_variable([1, 2], dtypes.float32) - final = gen_state_ops._destroy_temporary_variable(var, var_name="bad") + var = gen_state_ops.temporary_variable([1, 2], dtypes.float32) + final = gen_state_ops.destroy_temporary_variable(var, var_name="bad") with self.assertRaises(errors.NotFoundError): final.eval() def testDuplicateTemporaryVariable(self): with self.test_session(use_gpu=True): - var1 = gen_state_ops._temporary_variable( + var1 = gen_state_ops.temporary_variable( [1, 2], dtypes.float32, var_name="dup") var1 = state_ops.assign(var1, [[1.0, 2.0]]) - var2 = gen_state_ops._temporary_variable( + var2 = gen_state_ops.temporary_variable( [1, 2], dtypes.float32, var_name="dup") var2 = state_ops.assign(var2, [[3.0, 4.0]]) final = var1 + var2 @@ -193,25 +193,25 @@ class VariableOpTest(test.TestCase): def testDestroyTemporaryVariableTwice(self): with self.test_session(use_gpu=True): - var = gen_state_ops._temporary_variable([1, 2], dtypes.float32) - val1 = gen_state_ops._destroy_temporary_variable(var, var_name="dup") - val2 = gen_state_ops._destroy_temporary_variable(var, var_name="dup") + var = gen_state_ops.temporary_variable([1, 2], dtypes.float32) + val1 = gen_state_ops.destroy_temporary_variable(var, var_name="dup") + val2 = gen_state_ops.destroy_temporary_variable(var, var_name="dup") final = val1 + val2 with self.assertRaises(errors.NotFoundError): final.eval() def testTemporaryVariableNoLeak(self): with self.test_session(use_gpu=True): - var = gen_state_ops._temporary_variable( + var = gen_state_ops.temporary_variable( [1, 2], dtypes.float32, var_name="bar") final = array_ops.identity(var) final.eval() def testTwoTemporaryVariablesNoLeaks(self): with self.test_session(use_gpu=True): - var1 = gen_state_ops._temporary_variable( + var1 = gen_state_ops.temporary_variable( [1, 2], dtypes.float32, var_name="var1") - var2 = gen_state_ops._temporary_variable( + var2 = gen_state_ops.temporary_variable( [1, 2], dtypes.float32, var_name="var2") final = var1 + var2 final.eval() diff --git a/tensorflow/python/kernel_tests/variable_scope_test.py b/tensorflow/python/kernel_tests/variable_scope_test.py index 8527f116f9541942e52ba2ab635ca1212ea38583..531d0cdf9002d966909fe29ffcd538c3510c60ab 100644 --- a/tensorflow/python/kernel_tests/variable_scope_test.py +++ b/tensorflow/python/kernel_tests/variable_scope_test.py @@ -166,12 +166,10 @@ class VariableScopeTest(test.TestCase): self.evaluate(variables_lib.variables_initializer([w])) self.assertAllClose(self.evaluate(w.value()), [1, 2, 3]) - if context.in_graph_mode(): - with self.assertRaises(TypeError): - variable_scope.get_variable("x4", initializer={}) - else: - with self.assertRaises(ValueError): - variable_scope.get_variable("x4", initializer={}) + # A quirk to be revisited? + error = ValueError if context.executing_eagerly() else TypeError + with self.assertRaises(error): + variable_scope.get_variable("x4", initializer={}) @test_util.run_in_graph_and_eager_modes() def testInitFromNonInitializer(self): @@ -267,7 +265,7 @@ class VariableScopeTest(test.TestCase): self.assertAllClose(self.evaluate(losses[2]), 0.5) with variable_scope.variable_scope("foo", reuse=True): # reuse=True is for now only supported when eager execution is disabled. - if context.in_graph_mode(): + if not context.executing_eagerly(): v = variable_scope.get_variable("v", []) # "v" is alredy there, reused losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) @@ -374,7 +372,7 @@ class VariableScopeTest(test.TestCase): v = variable_scope.get_variable("v", []) self.evaluate(variables_lib.variables_initializer([v])) self.assertAllClose(self.evaluate(v.value()), 0.3) - if context.in_graph_mode(): + if not context.executing_eagerly(): # Check that we can set reuse. variable_scope.get_variable_scope().reuse_variables() with self.assertRaises(ValueError): # Fail, w does not exist yet. @@ -408,7 +406,7 @@ class VariableScopeTest(test.TestCase): with variable_scope.variable_scope("tower") as tower: with ops.name_scope("scope2") as sc2: self.assertEqual(sc2, "testVarScopeNameScope1/tower/scope2/") - if context.in_graph_mode(): + if not context.executing_eagerly(): with variable_scope.variable_scope( tower): # Re-entering acts like another "tower". with ops.name_scope("scope2") as sc2: @@ -422,7 +420,7 @@ class VariableScopeTest(test.TestCase): with variable_scope.variable_scope("tower"): with ops.name_scope("scope2") as sc2: self.assertEqual(sc2, "testVarScopeNameScope2/tower/scope2/") - if context.in_graph_mode(): + if not context.executing_eagerly(): with variable_scope.variable_scope(tower): with ops.name_scope("scope2") as sc2: self.assertEqual(sc2, "testVarScopeNameScope2/tower_1/scope2/") @@ -903,17 +901,15 @@ class VariableScopeTest(test.TestCase): "w", [], collections=["foo"]) self.assertEqual(local_var.name, "outer/w:0") - # Since variable is local, it should be in the local variable collection - # but not the trainable collection. - if context.in_graph_mode(): + if not context.executing_eagerly(): + # Since variable is local, it should be in the local variable collection + # but not the trainable collection. self.assertIn(local_var, ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES)) self.assertIn(local_var, ops.get_collection("foo")) self.assertNotIn(local_var, ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)) - - # Check that local variable respects `reuse`. - if context.in_graph_mode(): + # Check that local variable respects `reuse`. with variable_scope.variable_scope(outer, "default", reuse=True): self.assertEqual( variable_scope.get_local_variable("w", []).name, "outer/w:0") diff --git a/tensorflow/python/kernel_tests/variables_test.py b/tensorflow/python/kernel_tests/variables_test.py index b16c8c002c98a0351d1fc55fce061695327a18c9..27599868b74be323189b872c2147c6a33f84d170 100644 --- a/tensorflow/python/kernel_tests/variables_test.py +++ b/tensorflow/python/kernel_tests/variables_test.py @@ -687,7 +687,7 @@ class VariableContainerTest(test.TestCase): v1 = variables.Variable([1]) with ops.container("l2"): v2 = variables.Variable([2]) - special_v = gen_state_ops._variable( + special_v = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="VariableInL3", diff --git a/tensorflow/python/kernel_tests/xent_op_test.py b/tensorflow/python/kernel_tests/xent_op_test.py index e152f02d8e983364603053dc5c8d14b5dfaf3605..e3e120a4eb01885ac5ac5e41f82ad3e480a83a77 100644 --- a/tensorflow/python/kernel_tests/xent_op_test.py +++ b/tensorflow/python/kernel_tests/xent_op_test.py @@ -48,7 +48,7 @@ class XentTest(test.TestCase): def _testXent(self, np_features, np_labels, use_gpu=False): np_loss, np_backprop = self._npXent(np_features, np_labels) with self.test_session(use_gpu=use_gpu) as sess: - loss, backprop = gen_nn_ops._softmax_cross_entropy_with_logits( + loss, backprop = gen_nn_ops.softmax_cross_entropy_with_logits( np_features, np_labels) tf_loss, tf_backprop = sess.run([loss, backprop]) self.assertAllCloseAccordingToType(np_loss, tf_loss) @@ -71,7 +71,7 @@ class XentTest(test.TestCase): def _testSingleClass(self, use_gpu=False): for dtype in np.float16, np.float32: with self.test_session(use_gpu=use_gpu) as sess: - loss, backprop = gen_nn_ops._softmax_cross_entropy_with_logits( + loss, backprop = gen_nn_ops.softmax_cross_entropy_with_logits( np.array([[1.], [-1.], [0.]]).astype(dtype), np.array([[-1.], [0.], [1.]]).astype(dtype)) tf_loss, tf_backprop = sess.run([loss, backprop]) @@ -89,7 +89,7 @@ class XentTest(test.TestCase): np_labels = np.array([[[0., 0., 0., 1.]], [[0., .5, .5, 0.]]]).astype(dtype) self.assertRaisesRegexp(ValueError, "must be rank 2", - gen_nn_ops._softmax_cross_entropy_with_logits, + gen_nn_ops.softmax_cross_entropy_with_logits, np_features, np_labels) def testNpXent(self): @@ -131,14 +131,14 @@ class XentTest(test.TestCase): def testShapeMismatch(self): with self.test_session(): with self.assertRaises(ValueError): - gen_nn_ops._softmax_cross_entropy_with_logits( + gen_nn_ops.softmax_cross_entropy_with_logits( [[0., 1.], [2., 3.]], [[0., 1., 0.], [1., 0., 0.]]) def testNotMatrix(self): with self.test_session(): with self.assertRaises(ValueError): - gen_nn_ops._softmax_cross_entropy_with_logits([0., 1., 2., 3.], - [0., 1., 0., 1.]) + gen_nn_ops.softmax_cross_entropy_with_logits([0., 1., 2., 3.], + [0., 1., 0., 1.]) def testHalf(self): self._testAll( diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py index 04b6056acec0de8dea968128dc9a4c248a98379b..e4395bea92961d348ed3841a31cacb91aaa282ec 100644 --- a/tensorflow/python/layers/base.py +++ b/tensorflow/python/layers/base.py @@ -31,14 +31,18 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.layers import utils as layers_util +from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import variable_scope as vs from tensorflow.python.ops import variables as tf_variables from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpointable from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export -class Layer(object): +@tf_export('layers.Layer') +class Layer(checkpointable.CheckpointableBase): """Base layer class. This is the class from which all layers inherit, implementing common @@ -111,7 +115,7 @@ class Layer(object): # Provides information about which inputs are compatible with the layer. self.input_spec = None - if activity_regularizer and context.in_eager_mode(): + if activity_regularizer and context.executing_eagerly(): raise ValueError( ('Activity regularization is not supported when executing eagerly. ' 'Got activity_regularizer=%s') % (activity_regularizer,)) @@ -123,14 +127,12 @@ class Layer(object): # return tensors. When using graph execution, _losses is a list of ops. self._losses = [] self._reuse = kwargs.get('_reuse') - self._graph = ops.get_default_graph() - self._per_input_losses = {} - self._per_input_updates = {} + self._graph = None # Will be set at build time. self._dtype = None if dtype is None else dtypes.as_dtype(dtype).name - call_fn_args = estimator_util.fn_args(self.call) - self._compute_previous_mask = ('mask' in call_fn_args or + self._call_fn_args = estimator_util.fn_args(self.call) + self._compute_previous_mask = ('mask' in self._call_fn_args or hasattr(self, 'compute_mask')) - self._call_has_scope_arg = 'scope' in call_fn_args + self._call_has_scope_arg = 'scope' in self._call_fn_args # These lists will be filled via successive calls # to self._add_inbound_node(). @@ -226,7 +228,7 @@ class Layer(object): @property def updates(self): - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.updates not supported in Eager mode.') if not self.trainable and not self.stateful: return [] @@ -249,39 +251,34 @@ class Layer(object): Arguments: updates: Update op, or list/tuple of update ops. - inputs: Optional input tensor(s) that the update(s) depend on. Must - match the `inputs` argument passed to the `__call__` method at the time - the updates are created. If `None` is passed, the updates are assumed - to be unconditional, and will apply across all dataflows of the layer. + inputs: If anything other than None is passed, it signals the updates + are conditional on some of the layer's inputs, + and thus they should only be run where these inputs are available. + This is the case for BatchNormalization updates, for instance. + If None, the updates will be taken into account unconditionally, + and you are responsible for making sure that any dependency they might + have is available at runtime. + A step counter might fall into this category. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return # Updates already applied when in eager mode. + updates = _to_list(updates) - if not updates: - return + updates = [x if isinstance(x, ops.Operation) + else ops.convert_to_tensor(x) for x in updates] self._updates += updates - if inputs is not None: - inputs = nest.flatten(inputs) - if not inputs: - inputs = None - if inputs is not None: - # We compute an ID that uniquely identifies the list of tensors. - # This ID is order-sensitive. - inputs_hash = layers_util.object_list_uid(inputs) + if inputs is None: + for u in updates: + u._unconditional_update = True # pylint: disable=protected-access else: - inputs_hash = None - if inputs_hash not in self._per_input_updates: - self._per_input_updates[inputs_hash] = [] - self._per_input_updates[inputs_hash] += updates + for u in updates: + u._unconditional_update = False # pylint: disable=protected-access def get_updates_for(self, inputs): """Retrieves updates relevant to a specific set of inputs. Arguments: inputs: Input tensor or list/tuple of input tensors. - Must match the `inputs` argument passed to the `__call__` method - at the time the updates were created. - If you pass `inputs=None`, unconditional updates are returned. Returns: List of update ops of the layer that depend on `inputs`. @@ -289,19 +286,25 @@ class Layer(object): Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): - raise RuntimeError('Layer.get_updates_for not supported in Eager mode.') + if context.executing_eagerly(): + raise RuntimeError('`get_updates_for()` not supported in Eager mode.') + + # Updates disabled if layer is not trainable and not explicitly stateful. if not self.trainable and not self.stateful: return [] - if inputs is not None: - inputs = nest.flatten(inputs) - if not inputs: - inputs = None - if inputs is not None: - inputs_hash = layers_util.object_list_uid(inputs) - else: - inputs_hash = None - return self._per_input_updates.get(inputs_hash, []) + + if inputs is None: + # Requesting unconditional updates. + return [x for x in self.updates if x._unconditional_update] # pylint: disable=protected-access + + # Requesting input-conditional updates. + inputs = nest.flatten(inputs) + reachable = layers_util.get_reachable_from_inputs(inputs, self.updates) + updates = [] + for update in self.updates: + if update in reachable: + updates.append(update) + return updates @property def losses(self): @@ -314,7 +317,7 @@ class Layer(object): Returns: A list of tensors. """ - if context.in_eager_mode(): + if context.executing_eagerly(): # _losses may only contain variable regularization losses when executing # eagerly, and they have been saved as lambdas to be executed when # requested. @@ -341,34 +344,37 @@ class Layer(object): Arguments: losses: Loss tensor, or list/tuple of tensors. - inputs: Optional input tensor(s) that the loss(es) depend on. Must - match the `inputs` argument passed to the `__call__` method at the time - the losses are created. If `None` is passed, the losses are assumed + inputs: If anything other than None is passed, it signals the losses + are conditional on some of the layer's inputs, + and thus they should only be run where these inputs are available. + This is the case for activity regularization losses, for instance. + If `None` is passed, the losses are assumed to be unconditional, and will apply across all dataflows of the layer (e.g. weight regularization losses). Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): + # TODO(fchollet): it should be possible (and highly desirable) to support + # `add_loss` in eager mode. This allows great convenience and flexibility + # in defining custom losses on the fly (e.g. in VAEs). + # Simply appending the loss value to `self._losses` + # is the correct behavior. + # The only caveat is that we need to force the user to only call + # `add_loss` from inside a model or Layer's `call` method + # (otherwise the loss computation cannot be backproped through). raise RuntimeError('Layer.add_loss not supported in Eager mode.') + losses = _to_list(losses) - if not losses: - return self._losses += losses - if inputs is not None: - inputs = nest.flatten(inputs) - if not inputs: - inputs = None - if inputs is not None: - # We compute an ID that uniquely identifies the list of tensors. - # This ID is order-sensitive. - inputs_hash = layers_util.object_list_uid(inputs) + if inputs is None: + for loss in losses: + loss._unconditional_loss = True # pylint: disable=protected-access else: - inputs_hash = None - if inputs_hash not in self._per_input_losses: - self._per_input_losses[inputs_hash] = [] - self._per_input_losses[inputs_hash] += losses + for loss in losses: + loss._unconditional_loss = False # pylint: disable=protected-access + # TODO(fchollet): deprecate collection below. _add_elements_to_collection(losses, ops.GraphKeys.REGULARIZATION_LOSSES) def get_losses_for(self, inputs): @@ -376,10 +382,6 @@ class Layer(object): Arguments: inputs: Input tensor or list/tuple of input tensors. - Must match the `inputs` argument passed to the `__call__` - method at the time the losses were created. - If you pass `inputs=None`, unconditional losses are returned, - such as weight regularization losses. Returns: List of loss tensors of the layer that depend on `inputs`. @@ -387,17 +389,25 @@ class Layer(object): Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.get_losses_for not supported in Eager mode.') - if inputs is not None: - inputs = nest.flatten(inputs) - if not inputs: - inputs = None - if inputs is not None: - inputs_hash = layers_util.object_list_uid(inputs) - else: - inputs_hash = None - return self._per_input_losses.get(inputs_hash, []) + + if inputs is None: + # Requesting unconditional losses. + return [x for x in self.losses if x._unconditional_loss] # pylint: disable=protected-access + + # Requesting input-conditional losses. + inputs = nest.flatten(inputs) + # Retrieve the set of tensors in the TF graph that depend on `inputs`. + # The losses we want to return will be part of this set. + # To avoid unnecessary work, we stop the search in case all of + # `self.losses` have been retrieved. + reachable = layers_util.get_reachable_from_inputs(inputs, self.losses) + losses = [] + for loss in self.losses: + if loss in reachable: + losses.append(loss) + return losses def build(self, _): """Creates the variables of the layer.""" @@ -499,7 +509,7 @@ class Layer(object): # will occur; it should be None if and only if initialization will take # place in the eager context. init_graph = None - if context.in_graph_mode(): + if not context.executing_eagerly(): default_graph = ops.get_default_graph() if default_graph.building_function: with ops.init_scope(): @@ -507,7 +517,7 @@ class Layer(object): # will be lifted; if initialization ops will be lifted into # the eager context, then there is nothing to retrieve, since variable # collections are not supported when eager execution is enabled. - if context.in_graph_mode(): + if not context.executing_eagerly(): init_graph = ops.get_default_graph() existing_variables = set(tf_variables.global_variables()) else: @@ -523,13 +533,17 @@ class Layer(object): with vs.variable_scope( self._scope, reuse=reuse, auxiliary_name_scope=False) as scope: with ops.name_scope(self._name_scope_name(scope)): - variable = vs.get_variable(name, - shape=shape, - initializer=initializer, - dtype=dtypes.as_dtype(dtype), - constraint=constraint, - trainable=trainable and self.trainable, - partitioner=partitioner) + variable = self._add_variable_with_custom_getter( + name=name, + shape=shape, + getter=vs.get_variable, + # Manage errors in Layer rather than Checkpointable. + overwrite=True, + initializer=initializer, + dtype=dtypes.as_dtype(dtype), + constraint=constraint, + trainable=trainable and self.trainable, + partitioner=partitioner) if init_graph is not None: # pylint: disable=protected-access # The variable was created and initialized in a graph. @@ -564,7 +578,7 @@ class Layer(object): if isinstance(variable, tf_variables.PartitionedVariable): raise RuntimeError( 'Partitioned variable regularization is not yet ' - 'supported when executing eagerly. File a feature request' + 'supported when executing eagerly. File a feature request ' 'if this is important to you.') # Save a zero-argument lambda which runs the regularizer on the # variable, to be executed when `Layer.losses` is requested. @@ -610,16 +624,17 @@ class Layer(object): self._set_scope(kwargs.pop('scope', None)) input_list = nest.flatten(inputs) - in_graph_mode = context.in_graph_mode() + build_graph = not context.executing_eagerly() in_deferred_mode = isinstance(input_list[0], _DeferredTensor) # Ensure the Layer, if being reused, is working with inputs from # the same graph as where it was created. - if in_graph_mode: + if build_graph: try: - ops._get_graph_from_inputs(input_list, graph=self.graph) # pylint: disable=protected-access + # Set layer's "graph" at build time + self._graph = ops._get_graph_from_inputs(input_list, graph=self._graph) # pylint: disable=protected-access except ValueError as e: raise ValueError('Input graph and Layer graph are not the same: %s' % e) - if in_graph_mode or in_deferred_mode: + if build_graph or in_deferred_mode: user_kwargs = copy.copy(kwargs) # Handle Keras mask propagation from previous layer to current layer. @@ -627,8 +642,9 @@ class Layer(object): if (not hasattr(self, '_compute_previous_mask') or self._compute_previous_mask): previous_mask = _collect_previous_mask(inputs) - if ('mask' in estimator_util.fn_args(self.call) and - 'mask' not in kwargs and + if not hasattr(self, '_call_fn_args'): + self._call_fn_args = estimator_util.fn_args(self.call) + if ('mask' in self._call_fn_args and 'mask' not in kwargs and not _is_all_none(previous_mask)): # The previous layer generated a mask, and mask was not explicitly pass # to __call__, hence we set previous_mask as the default value. @@ -649,16 +665,18 @@ class Layer(object): else: scope_context_manager = vs.variable_scope( self._scope, reuse=self._reuse, auxiliary_name_scope=False) + input_shapes = None with scope_context_manager as scope: with ops.name_scope(self._name_scope_name(scope)): if not self.built: - if not in_graph_mode: + if not build_graph: # Activity regularization is currently unsupported in Eager mode. if self._activity_regularizer: - raise ValueError('activity_regularizer currently unsupported in ' - 'Eager mode. Found an activity_regularizer in ' - '%s(%s).' % (self.__class__.__name__, self)) - if not in_graph_mode and not in_deferred_mode: + raise ValueError( + 'activity_regularizer currently unsupported with ' + 'eager execution enabled. Found an activity_regularizer in ' + '%s(%s).' % (self.__class__.__name__, self)) + if not build_graph and not in_deferred_mode: # TODO(agarwal): support _keras_history in Eager mode. for x in input_list: if hasattr(x, '_keras_history'): @@ -683,11 +701,13 @@ class Layer(object): # TODO(agarwal): Fix the sub-classes and avoid this complexity. call_has_scope_arg = self._call_has_scope_arg except AttributeError: - call_has_scope_arg = 'scope' in estimator_util.fn_args(self.call) + self._call_fn_args = estimator_util.fn_args(self.call) + self._call_has_scope_arg = 'scope' in self._call_fn_args + call_has_scope_arg = self._call_has_scope_arg if call_has_scope_arg: kwargs['scope'] = scope # Check input assumptions set after layer building, e.g. input shape. - if in_graph_mode or in_deferred_mode: + if build_graph or in_deferred_mode: self._assert_input_compatibility(inputs) if not in_deferred_mode: @@ -698,6 +718,9 @@ class Layer(object): else: # Deferred mode behavior: use `compute_output_shape` to # infer the number of outputs of the layer and their shapes. + if input_shapes is None: + input_shapes = nest.map_structure(lambda x: x.get_shape(), inputs) + output_shapes = self.compute_output_shape(input_shapes) output_shapes = nest.flatten(output_shapes) outputs = [ @@ -708,7 +731,7 @@ class Layer(object): if len(outputs) == 1: outputs = outputs[0] - if in_graph_mode: + if build_graph: # Apply activity regularization. # Note that it should be applied every time the layer creates a new # output, since it is output-specific. @@ -719,12 +742,10 @@ class Layer(object): activity_regularization = self._activity_regularizer(output) self.add_loss(activity_regularization, inputs=inputs) - if not in_deferred_mode: - # TODO(fchollet): consider how masking will work with deferred mode. - # Handle mask computation and propagation to the next layer. + # TODO(fchollet): consider enabling masking for Eager mode. if hasattr(self, 'compute_mask'): output_mask = self.compute_mask(inputs, previous_mask) - if isinstance(outputs, list): + if isinstance(outputs, (list, tuple)): if output_mask is None: output_mask = [None for _ in range(len(outputs))] for x, m in zip(outputs, output_mask): @@ -732,7 +753,7 @@ class Layer(object): else: outputs._keras_mask = output_mask # pylint: disable=protected-access - if in_graph_mode: + if build_graph: # If all input tensors have history metadata, # we update the output tensors # with corresponding history metadata, thus eventually allowing to use @@ -755,7 +776,7 @@ class Layer(object): # Update global default collections. _add_elements_to_collection(self.updates, ops.GraphKeys.UPDATE_OPS) - if in_deferred_mode or in_graph_mode: + if in_deferred_mode or build_graph: if _have_all_keras_metadata(inputs): # Add an inbound node to the layer, so it can keep track of this call. # This updates the layer history of the output tensor(s). @@ -767,7 +788,7 @@ class Layer(object): @property def graph(self): - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.graph not supported in Eager mode.') return self._graph @@ -871,7 +892,7 @@ class Layer(object): mode. ValueError: If the index provided does not match any node. """ - assert context.in_graph_mode() + assert not context.executing_eagerly() if not self._inbound_nodes: raise RuntimeError('The layer has never been called ' 'and thus has no defined ' + attr_name + '.') @@ -901,7 +922,7 @@ class Layer(object): Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( 'Layer.get_input_shape_at not supported in Eager mode.') return self._get_node_attribute_at_index(node_index, 'input_shapes', @@ -923,7 +944,7 @@ class Layer(object): Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( 'Layer.get_output_shape_at not supported in Eager mode.') return self._get_node_attribute_at_index(node_index, 'output_shapes', @@ -944,7 +965,7 @@ class Layer(object): Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.get_input_at not supported in Eager mode.') return self._get_node_attribute_at_index(node_index, 'input_tensors', 'input') @@ -964,7 +985,7 @@ class Layer(object): Raises: RuntimeError: If called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.get_output_at not supported in Eager mode.') return self._get_node_attribute_at_index(node_index, 'output_tensors', 'output') @@ -987,7 +1008,7 @@ class Layer(object): RuntimeError: If called in Eager mode. AttributeError: If no inbound nodes are found. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.input not supported in Eager mode.') if not self._inbound_nodes: raise AttributeError('Layer ' + self.name + @@ -1009,7 +1030,7 @@ class Layer(object): layers. RuntimeError: if called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.output not supported in Eager mode.') if not self._inbound_nodes: raise AttributeError('Layer ' + self.name + ' has no inbound nodes.') @@ -1031,7 +1052,7 @@ class Layer(object): AttributeError: if the layer has no defined input_shape. RuntimeError: if called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.input_shape not supported in Eager mode.') if not self._inbound_nodes: raise AttributeError('The layer has never been called ' @@ -1092,7 +1113,7 @@ class Layer(object): AttributeError: if the layer has no defined output shape. RuntimeError: if called in Eager mode. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Layer.output_shape not supported in Eager mode.') if not self._inbound_nodes: raise AttributeError('The layer has never been called ' @@ -1223,6 +1244,7 @@ class Layer(object): ', found shape=' + str(shape)) +@tf_export('keras.layers.InputSpec', 'layers.InputSpec') class InputSpec(object): """Specifies the ndim, dtype and shape of every input to a layer. @@ -1393,7 +1415,10 @@ class _DeferredTensor(object): def __init__(self, shape, dtype, name=None): self.shape = tensor_shape.TensorShape(shape) - self.dtype = dtypes.as_dtype(dtype) + if dtype is None: + self.dtype = dtypes.as_dtype(np.float32) + else: + self.dtype = dtypes.as_dtype(dtype) self.name = name def get_shape(self): @@ -1446,7 +1471,7 @@ def _to_list(x): def _add_elements_to_collection(elements, collection_list): - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('Using collections from Layers not supported in Eager ' 'mode. Tried to add %s to %s' % (elements, collection_list)) diff --git a/tensorflow/python/layers/base_test.py b/tensorflow/python/layers/base_test.py index 06ba214c0fc60202c773f8f231b17c3b728f5c52..9ed4afeaba931c47d2a1e65f08489773f0b9eb1b 100644 --- a/tensorflow/python/layers/base_test.py +++ b/tensorflow/python/layers/base_test.py @@ -31,6 +31,7 @@ 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 random_ops +from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import test @@ -43,7 +44,7 @@ class BaseLayerTest(test.TestCase): self.assertEqual(layer.variables, []) self.assertEqual(layer.trainable_variables, []) self.assertEqual(layer.non_trainable_variables, []) - if context.in_graph_mode(): + if not context.executing_eagerly(): # updates, losses only supported in GRAPH mode self.assertEqual(layer.updates, []) self.assertEqual(layer.losses, []) @@ -62,7 +63,7 @@ class BaseLayerTest(test.TestCase): self.assertEqual(layer.variables, [variable]) self.assertEqual(layer.trainable_variables, [variable]) self.assertEqual(layer.non_trainable_variables, []) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( layer.variables, ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)) @@ -76,7 +77,7 @@ class BaseLayerTest(test.TestCase): self.assertEqual(layer.variables, [variable, variable_2]) self.assertEqual(layer.trainable_variables, [variable]) self.assertEqual(layer.non_trainable_variables, [variable_2]) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( len(ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)), 1) @@ -160,7 +161,7 @@ class BaseLayerTest(test.TestCase): inputs = random_ops.random_uniform((5,), seed=1) outputs = layer.apply(inputs) self.assertEqual(layer.built, True) - if context.in_graph_mode(): + if not context.executing_eagerly(): # op is only supported in GRAPH mode self.assertEqual(outputs.op.name, 'my_layer/Square') @@ -209,7 +210,7 @@ class BaseLayerTest(test.TestCase): inputs = random_ops.random_uniform((5,), seed=1) outputs = layer.apply(inputs) self.assertEqual(layer.built, True) - if context.in_graph_mode(): + if not context.executing_eagerly(): # op only supported in GRAPH mode. self.assertEqual(outputs.op.name, 'my_layer/Square') @@ -279,7 +280,7 @@ class BaseLayerTest(test.TestCase): def call(self, inputs): return inputs - if context.in_graph_mode(): + if not context.executing_eagerly(): layer = CustomerLayer() with self.assertRaisesRegexp(ValueError, r'requires a defined rank'): layer.apply(array_ops.placeholder('int32')) @@ -306,7 +307,7 @@ class BaseLayerTest(test.TestCase): def call(self, inputs): return inputs - if context.in_graph_mode(): + if not context.executing_eagerly(): layer = CustomerLayer() with self.assertRaisesRegexp(ValueError, r'requires a defined rank'): layer.apply(array_ops.placeholder('int32')) @@ -334,7 +335,7 @@ class BaseLayerTest(test.TestCase): def call(self, inputs): return inputs - if context.in_graph_mode(): + if not context.executing_eagerly(): layer = CustomerLayer() with self.assertRaisesRegexp(ValueError, r'requires a defined rank'): layer.apply(array_ops.placeholder('int32')) @@ -429,7 +430,7 @@ class BaseLayerTest(test.TestCase): layer.apply(constant_op.constant(1)) # Works - if context.in_graph_mode(): + if not context.executing_eagerly(): layer.apply(array_ops.placeholder('int32')) layer.apply(array_ops.placeholder('int32', shape=(2, 3))) @@ -452,13 +453,7 @@ class BaseLayerTest(test.TestCase): return {'l' + key: inputs[key] for key in inputs} layer = DictLayer() - if context.in_graph_mode(): - i1 = array_ops.placeholder('int32') - i2 = array_ops.placeholder('float32') - result = layer.apply({'abel': i1, 'ogits': i2}) - self.assertTrue(isinstance(result, dict)) - self.assertEqual(set(['label', 'logits']), set(result.keys())) - else: + if context.executing_eagerly(): i1 = constant_op.constant(3) i2 = constant_op.constant(4.0) result = layer.apply({'abel': i1, 'ogits': i2}) @@ -466,6 +461,12 @@ class BaseLayerTest(test.TestCase): self.assertEqual(set(['label', 'logits']), set(result.keys())) self.assertEqual(3, result['label'].numpy()) self.assertEqual(4.0, result['logits'].numpy()) + else: + i1 = array_ops.placeholder('int32') + i2 = array_ops.placeholder('float32') + result = layer.apply({'abel': i1, 'ogits': i2}) + self.assertTrue(isinstance(result, dict)) + self.assertEqual(set(['label', 'logits']), set(result.keys())) def testActivityRegularizer(self): regularizer = math_ops.reduce_sum @@ -555,6 +556,103 @@ class BaseLayerTest(test.TestCase): self.assertEqual(len(layer.trainable_variables), 1) self.assertEqual(layer.variables[0].graph, outer_graph) + def testGetUpdateFor(self): + + class MyLayer(base_layers.Layer): + + def build(self, input_shape): + self.a = self.add_variable('a', + (), + dtypes.float32, + trainable=False) + self.b = self.add_variable('b', + (), + dtypes.float32, + trainable=False) + self.add_update(state_ops.assign_add(self.a, 1., name='b_update')) + self.built = True + + def call(self, inputs): + self.add_update(state_ops.assign_add(self.a, inputs, name='a_update'), + inputs=True) + return inputs + 1 + + layer = MyLayer() + inputs = array_ops.placeholder(dtypes.float32, (), 'inputs') + intermediate_inputs = inputs + 1 + outputs = layer.apply(intermediate_inputs) + + self.assertEqual(len(layer.updates), 2) + self.assertEqual(len(layer.get_updates_for(None)), 1) + self.assertEqual(len(layer.get_updates_for([inputs])), 1) + self.assertEqual(len(layer.get_updates_for([intermediate_inputs])), 1) + self.assertEqual(len(layer.get_updates_for([outputs])), 0) + + # Call same layer on new input, creating one more conditional update + inputs = array_ops.placeholder(dtypes.float32, (), 'inputs') + intermediate_inputs = inputs + 1 + outputs = layer.apply(intermediate_inputs) + + self.assertEqual(len(layer.updates), 3) + self.assertEqual(len(layer.get_updates_for(None)), 1) + # Check that we are successfully filtering out irrelevant updates + self.assertEqual(len(layer.get_updates_for([inputs])), 1) + self.assertEqual(len(layer.get_updates_for([intermediate_inputs])), 1) + self.assertEqual(len(layer.get_updates_for([outputs])), 0) + + def testGetLossesFor(self): + + class MyLayer(base_layers.Layer): + + def build(self, input_shape): + self.a = self.add_variable('a', + (), + dtypes.float32, + trainable=False) + self.b = self.add_variable('b', + (), + dtypes.float32, + trainable=False) + self.add_loss(self.a) + self.built = True + + def call(self, inputs): + self.add_loss(inputs, inputs=True) + return inputs + 1 + + layer = MyLayer() + inputs = array_ops.placeholder(dtypes.float32, (), 'inputs') + intermediate_inputs = inputs + 1 + outputs = layer.apply(intermediate_inputs) + + self.assertEqual(len(layer.losses), 2) + self.assertEqual(len(layer.get_losses_for(None)), 1) + self.assertEqual(len(layer.get_losses_for([inputs])), 1) + self.assertEqual(len(layer.get_losses_for([intermediate_inputs])), 1) + self.assertEqual(len(layer.get_losses_for([outputs])), 0) + + # Call same layer on new input, creating one more conditional loss + inputs = array_ops.placeholder(dtypes.float32, (), 'inputs') + intermediate_inputs = inputs + 1 + outputs = layer.apply(intermediate_inputs) + + self.assertEqual(len(layer.losses), 3) + self.assertEqual(len(layer.get_losses_for(None)), 1) + # Check that we are successfully filtering out irrelevant losses + self.assertEqual(len(layer.get_losses_for([inputs])), 1) + self.assertEqual(len(layer.get_losses_for([intermediate_inputs])), 1) + self.assertEqual(len(layer.get_losses_for([outputs])), 0) + + def testLayerGraphSetInFirstApply(self): + with ops.Graph().as_default(): + layer = core_layers.Dense(1) # Graph at construction time is ignored + with ops.Graph().as_default(): + layer.apply(constant_op.constant([[1]])) + # layer is now bound to second Graph + with ops.Graph().as_default(), self.assertRaisesRegexp( + ValueError, 'Input graph and Layer graph are not the same'): + layer.apply(constant_op.constant([[1]])) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/layers/convolutional.py b/tensorflow/python/layers/convolutional.py index e8dba3cea321a415b84e1ec89fd7b021e2b272d0..74e7c63fb364d9c4475af5efe7d5db95cccf8166 100644 --- a/tensorflow/python/layers/convolutional.py +++ b/tensorflow/python/layers/convolutional.py @@ -29,6 +29,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import nn from tensorflow.python.ops import nn_ops +from tensorflow.python.util.tf_export import tf_export class _Conv(base.Layer): @@ -222,6 +223,7 @@ class _Conv(base.Layer): new_space) +@tf_export('layers.Conv1D') class Conv1D(_Conv): """1D convolution layer (e.g. temporal convolution). @@ -311,6 +313,7 @@ class Conv1D(_Conv): name=name, **kwargs) +@tf_export('layers.conv1d') def conv1d(inputs, filters, kernel_size, @@ -411,6 +414,7 @@ def conv1d(inputs, return layer.apply(inputs) +@tf_export('layers.Conv2D') class Conv2D(_Conv): """2D convolution layer (e.g. spatial convolution over images). @@ -507,6 +511,7 @@ class Conv2D(_Conv): name=name, **kwargs) +@tf_export('layers.conv2d') def conv2d(inputs, filters, kernel_size, @@ -614,6 +619,7 @@ def conv2d(inputs, return layer.apply(inputs) +@tf_export('layers.Conv3D') class Conv3D(_Conv): """3D convolution layer (e.g. spatial convolution over volumes). @@ -711,6 +717,7 @@ class Conv3D(_Conv): name=name, **kwargs) +@tf_export('layers.conv3d') def conv3d(inputs, filters, kernel_size, @@ -980,6 +987,7 @@ class _SeparableConv(_Conv): raise NotImplementedError +@tf_export('layers.SeparableConv1D') class SeparableConv1D(_SeparableConv): """Depthwise separable 1D convolution. @@ -1088,10 +1096,10 @@ class SeparableConv1D(_SeparableConv): def call(self, inputs): if self.data_format == 'channels_last': - strides = (1, 1) + self.strides + (1,) + strides = (1,) + self.strides * 2 + (1,) spatial_start_dim = 1 else: - strides = (1, 1, 1) + self.strides + strides = (1, 1) + self.strides * 2 spatial_start_dim = 2 # Explicitly broadcast inputs and kernels to 4D. @@ -1123,6 +1131,7 @@ class SeparableConv1D(_SeparableConv): return outputs +@tf_export('layers.SeparableConv2D') class SeparableConv2D(_SeparableConv): """Depthwise separable 2D convolution. @@ -1260,6 +1269,7 @@ class SeparableConv2D(_SeparableConv): return outputs +@tf_export('layers.separable_conv1d') def separable_conv1d(inputs, filters, kernel_size, @@ -1376,6 +1386,7 @@ def separable_conv1d(inputs, return layer.apply(inputs) +@tf_export('layers.separable_conv2d') def separable_conv2d(inputs, filters, kernel_size, @@ -1497,6 +1508,7 @@ def separable_conv2d(inputs, return layer.apply(inputs) +@tf_export('layers.Conv2DTranspose') class Conv2DTranspose(Conv2D): """Transposed 2D convolution layer (sometimes called 2D Deconvolution). @@ -1652,7 +1664,7 @@ class Conv2DTranspose(Conv2D): padding=self.padding.upper(), data_format=utils.convert_data_format(self.data_format, ndim=4)) - if context.in_graph_mode(): + if not context.executing_eagerly(): # Infer the static output shape: out_shape = inputs.get_shape().as_list() out_shape[c_axis] = self.filters @@ -1695,6 +1707,7 @@ class Conv2DTranspose(Conv2D): return tensor_shape.TensorShape(output_shape) +@tf_export('layers.conv2d_transpose') def conv2d_transpose(inputs, filters, kernel_size, @@ -1790,6 +1803,7 @@ def conv2d_transpose(inputs, return layer.apply(inputs) +@tf_export('layers.Conv3DTranspose') class Conv3DTranspose(Conv3D): """Transposed 3D convolution layer (sometimes called 3D Deconvolution). @@ -1955,7 +1969,7 @@ class Conv3DTranspose(Conv3D): data_format=utils.convert_data_format(self.data_format, ndim=5), padding=self.padding.upper()) - if context.in_graph_mode(): + if not context.executing_eagerly(): # Infer the static output shape: out_shape = inputs.get_shape().as_list() out_shape[c_axis] = self.filters @@ -2018,6 +2032,7 @@ class Conv3DTranspose(Conv3D): return tensor_shape.TensorShape(output_shape) +@tf_export('layers.conv3d_transpose') def conv3d_transpose(inputs, filters, kernel_size, diff --git a/tensorflow/python/layers/core.py b/tensorflow/python/layers/core.py index 7bf62d45b8e97aaa027467b6c9862ca2a4699fc1..e598d9f83ab21f2dd5fabb3dd37fa0bfb5f003a4 100644 --- a/tensorflow/python/layers/core.py +++ b/tensorflow/python/layers/core.py @@ -35,10 +35,14 @@ from tensorflow.python.layers import utils 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 gen_math_ops from tensorflow.python.ops import nn +from tensorflow.python.ops import nn_ops from tensorflow.python.ops import standard_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export('layers.Dense') class Dense(base.Layer): """Densely-connected layer class. @@ -152,11 +156,11 @@ class Dense(base.Layer): outputs = standard_ops.tensordot(inputs, self.kernel, [[len(shape) - 1], [0]]) # Reshape the output back to the original ndim of the input. - if context.in_graph_mode(): + if not context.executing_eagerly(): output_shape = shape[:-1] + [self.units] outputs.set_shape(output_shape) else: - outputs = standard_ops.matmul(inputs, self.kernel) + outputs = gen_math_ops.mat_mul(inputs, self.kernel) if self.use_bias: outputs = nn.bias_add(outputs, self.bias) if self.activation is not None: @@ -173,6 +177,7 @@ class Dense(base.Layer): return input_shape[:-1].concatenate(self.units) +@tf_export('layers.dense') def dense( inputs, units, activation=None, @@ -248,6 +253,7 @@ def dense( return layer.apply(inputs) +@tf_export('layers.Dropout') class Dropout(base.Layer): """Applies Dropout to the input. @@ -287,13 +293,7 @@ class Dropout(base.Layer): # shapes with dynamically sized inputs. if self.noise_shape is None: return self.noise_shape - - symbolic_shape = array_ops.shape(inputs) - noise_shape = [ - symbolic_shape[axis] if shape is None else shape - for axis, shape in enumerate(self.noise_shape) - ] - return noise_shape + return nn_ops._get_noise_shape(inputs, self.noise_shape) def call(self, inputs, training=False): @@ -309,6 +309,7 @@ class Dropout(base.Layer): return input_shape +@tf_export('layers.dropout') def dropout(inputs, rate=0.5, noise_shape=None, @@ -350,6 +351,7 @@ def dropout(inputs, return layer.apply(inputs, training=training) +@tf_export('layers.Flatten') class Flatten(base.Layer): """Flattens an input tensor while preserving the batch axis (axis 0). @@ -372,7 +374,7 @@ class Flatten(base.Layer): def call(self, inputs): outputs = array_ops.reshape(inputs, (array_ops.shape(inputs)[0], -1)) - if context.in_graph_mode(): + if not context.executing_eagerly(): outputs.set_shape(self.compute_output_shape(inputs.get_shape())) return outputs @@ -386,6 +388,7 @@ class Flatten(base.Layer): return tensor_shape.TensorShape(output_shape) +@tf_export('layers.flatten') def flatten(inputs, name=None): """Flattens an input tensor while preserving the batch axis (axis 0). diff --git a/tensorflow/python/layers/core_test.py b/tensorflow/python/layers/core_test.py index 15ce6cba21fcc78126f7db58ab18934db69c15fd..cf45b07637108422f1c612390bb01efdad6d5bcf 100644 --- a/tensorflow/python/layers/core_test.py +++ b/tensorflow/python/layers/core_test.py @@ -67,7 +67,7 @@ class DenseTest(test.TestCase): variables.global_variables_initializer().run() self.assertAllEqual(x.eval(), [[0.0]]) - @test_util.run_in_graph_and_eager_modes() + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) def testCall(self): dense = core_layers.Dense(2, activation=nn_ops.relu, name='my_dense') inputs = random_ops.random_uniform((5, 4), seed=1) @@ -77,12 +77,20 @@ class DenseTest(test.TestCase): self.assertListEqual(dense.trainable_variables, [dense.kernel, dense.bias]) self.assertListEqual(dense.non_trainable_variables, []) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( len(ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)), 2) self.assertEqual(dense.kernel.name, 'my_dense/kernel:0') self.assertEqual(dense.bias.name, 'my_dense/bias:0') + @test_util.assert_no_new_pyobjects_executing_eagerly + def testNoEagerLeak(self): + # Tests that repeatedly constructing and building a Layer does not leak + # Python objects. + inputs = random_ops.random_uniform((5, 4), seed=1) + core_layers.Dense(5)(inputs) + core_layers.Dense(2, activation=nn_ops.relu, name='my_dense')(inputs) + @test_util.run_in_graph_and_eager_modes() def testCallTensorDot(self): dense = core_layers.Dense(2, activation=nn_ops.relu, name='my_dense') @@ -98,7 +106,7 @@ class DenseTest(test.TestCase): self.assertListEqual(dense.variables, [dense.kernel]) self.assertListEqual(dense.trainable_variables, [dense.kernel]) self.assertListEqual(dense.non_trainable_variables, []) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( len(ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)), 1) self.assertEqual(dense.kernel.name, 'my_dense/kernel:0') @@ -113,7 +121,7 @@ class DenseTest(test.TestCase): self.assertListEqual(dense.non_trainable_variables, [dense.kernel, dense.bias]) self.assertListEqual(dense.trainable_variables, []) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( len(ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)), 0) @@ -162,13 +170,13 @@ class DenseTest(test.TestCase): dense = core_layers.Dense(2, activation=nn_ops.relu, name='dense1') inputs = random_ops.random_uniform((5, 3), seed=1) outputs = dense(inputs) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual(outputs.op.name, 'dense1/Relu') dense = core_layers.Dense(2, name='dense2') inputs = random_ops.random_uniform((5, 3), seed=1) outputs = dense(inputs) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual(outputs.op.name, 'dense2/BiasAdd') def testActivityRegularizer(self): @@ -374,7 +382,7 @@ class DropoutTest(test.TestCase): dp = core_layers.Dropout(0.5) inputs = array_ops.ones((5, 3)) dropped = dp.apply(inputs, training=True) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(variables.global_variables_initializer()) np_output = self.evaluate(dropped) self.assertAlmostEqual(0., np_output.min()) diff --git a/tensorflow/python/layers/layers.py b/tensorflow/python/layers/layers.py index 1555846efde812b9e31f48315decaf1f86aa4f70..13a8e8e39caaf9c74d1c7d0ea4d6856f725256fd 100644 --- a/tensorflow/python/layers/layers.py +++ b/tensorflow/python/layers/layers.py @@ -68,7 +68,6 @@ from tensorflow.python.util.all_util import remove_undocumented # Base objects. from tensorflow.python.layers.base import Layer from tensorflow.python.layers.base import InputSpec -from tensorflow.python.layers.network import Input # Core layers. from tensorflow.python.layers.core import Dense diff --git a/tensorflow/python/layers/maxout.py b/tensorflow/python/layers/maxout.py deleted file mode 100644 index 20ce6c9770087f9cfb90d40235955dfe1b7ee98b..0000000000000000000000000000000000000000 --- a/tensorflow/python/layers/maxout.py +++ /dev/null @@ -1,111 +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. -# ============================================================================= - -# pylint: disable=unused-import,g-bad-import-order -"""Contains the maxout layer -""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.python.eager import context -from tensorflow.python.framework import ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import gen_array_ops - -from tensorflow.python.layers import base - - -def maxout(inputs, num_units, axis=-1, name=None): - """Adds a maxout op from https://arxiv.org/abs/1302.4389 - - "Maxout Networks" Ian J. Goodfellow, David Warde-Farley, Mehdi Mirza, Aaron - Courville, - Yoshua Bengio - - Usually the operation is performed in the filter/channel dimension. This can - also be - used after fully-connected layers to reduce number of features. - - Arguments: - inputs: Tensor input - num_units: Specifies how many features will remain after maxout in the `axis` - dimension - (usually channel). This must be multiple of number of `axis`. - axis: The dimension where max pooling will be performed. Default is the - last dimension. - name: Optional scope for name_scope. - - Returns: - A `Tensor` representing the results of the pooling operation. - - Raises: - ValueError: if num_units is not multiple of number of features. - """ - return MaxOut(num_units=num_units, axis=axis, name=name)(inputs) - - -class MaxOut(base.Layer): - """Adds a maxout op from https://arxiv.org/abs/1302.4389 - - "Maxout Networks" Ian J. Goodfellow, David Warde-Farley, Mehdi Mirza, Aaron - Courville, Yoshua - Bengio - - Usually the operation is performed in the filter/channel dimension. This can - also be - used after fully-connected layers to reduce number of features. - - Arguments: - inputs: Tensor input - num_units: Specifies how many features will remain after maxout in the - `axis` dimension - (usually channel). - This must be multiple of number of `axis`. - axis: The dimension where max pooling will be performed. Default is the - last dimension. - name: Optional scope for name_scope. - - Returns: - A `Tensor` representing the results of the pooling operation. - - Raises: - ValueError: if num_units is not multiple of number of features. - """ - - def __init__(self, num_units, axis=-1, name=None, **kwargs): - super(MaxOut, self).__init__(name=name, trainable=False, **kwargs) - self.axis = axis - self.num_units = num_units - - def call(self, inputs): - inputs = ops.convert_to_tensor(inputs) - shape = inputs.get_shape().as_list() - num_channels = shape[self.axis] - if num_channels % self.num_units: - raise ValueError('number of features({}) is not ' - 'a multiple of num_units({})'.format( - num_channels, self.num_units)) - shape[self.axis] = -1 - shape += [num_channels // self.num_units] - - # Dealing with batches with arbitrary sizes - for i in range(len(shape)): - if shape[i] is None: - shape[i] = gen_array_ops.shape(inputs)[i] - outputs = math_ops.reduce_max( - gen_array_ops.reshape(inputs, shape), -1, keep_dims=False) - - return outputs diff --git a/tensorflow/python/layers/maxout_test.py b/tensorflow/python/layers/maxout_test.py deleted file mode 100644 index 26acac57c41da759f288f255c0cd523f9c6b1dbd..0000000000000000000000000000000000000000 --- a/tensorflow/python/layers/maxout_test.py +++ /dev/null @@ -1,61 +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. -# ============================================================================= - -# pylint: disable=unused-import,g-bad-import-order - - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.python.layers import maxout -from tensorflow.python.layers import convolutional as conv_layers -from tensorflow.python.layers import core as core_layers - -from tensorflow.python.ops import random_ops -from tensorflow.python.platform import test -import numpy as np - -""" -Contains the maxout layer tests -""" - - -class MaxOutTest(test.TestCase): - def test_simple(self): - inputs = random_ops.random_uniform((64, 10, 36), seed=1) - graph = maxout.maxout(inputs, num_units=3) - self.assertEqual(graph.get_shape().as_list(), [64, 10, 3]) - - def test_fully_connected(self): - inputs = random_ops.random_uniform((64, 50), seed=1) - graph = core_layers.dense(inputs, 50) - graph = maxout.maxout(graph, num_units=10) - self.assertEqual(graph.get_shape().as_list(), [64, 10]) - - def test_nchw(self): - inputs = random_ops.random_uniform((10, 100, 100, 3), seed=1) - graph = conv_layers.conv2d(inputs, 10, 3, padding="SAME") - graph = maxout.maxout(graph, num_units=1) - self.assertEqual(graph.get_shape().as_list(), [10, 100, 100, 1]) - - def test_invalid_shape(self): - inputs = random_ops.random_uniform((10, 100, 100, 3), seed=1) - graph = conv_layers.conv2d(inputs, 3, 10, strides=(1, 1)) - with self.assertRaisesRegexp(ValueError, 'number of features'): - graph = maxout.maxout(graph, num_units=2) - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/python/layers/network.py b/tensorflow/python/layers/network.py deleted file mode 100644 index 0a5dd57621b7dc06d9bc2d69c04cd8d6936fb7c8..0000000000000000000000000000000000000000 --- a/tensorflow/python/layers/network.py +++ /dev/null @@ -1,964 +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. -# ============================================================================= -"""Contains Network, a composition of layers.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - - -from tensorflow.python.eager import context -from tensorflow.python.estimator import util as estimator_util -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape -from tensorflow.python.layers import base -from tensorflow.python.layers import utils as layers_util -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.util import nest - - -class InputLayer(base.Layer): - """Layer to be used as an entry point into a Network (a graph of layers). - - It can either wrap an existing tensor (pass an `input_tensor` argument) - or create its a placeholder tensor (pass arguments `input_shape` - as well as `dtype`). - - It is generally recommend to use the functional layer API via `Input`, - (which creates an `InputLayer`) without directly using `InputLayer`. - - Arguments: - input_shape: Shape tuple (not including the batch axis), or `TensorShape` - instance (not including the batch axis). - batch_size: Optional input batch size (integer or None). - dtype: Datatype of the input. - input_tensor: Optional tensor to use as layer input - instead of creating a placeholder. - sparse: Boolean, whether the placeholder created - is meant to be sparse. - name: Name of the layer (string). - - Raises: - RuntimeError: If created in Eager mode. - """ - - def __init__(self, - input_shape=None, - batch_size=None, - dtype=dtypes.float32, - input_tensor=None, - sparse=False, - name=None): - super(InputLayer, self).__init__(dtype=dtype, name=name) - self.built = True - self.sparse = sparse - self.batch_size = batch_size - - if isinstance(input_shape, tensor_shape.TensorShape): - input_shape = tuple(input_shape.as_list()) - - if input_tensor is None: - if input_shape is not None: - batch_input_shape = (batch_size,) + tuple(input_shape) - else: - batch_input_shape = None - - if context.in_eager_mode(): - # In eager mode, create a temporary placeholder to call the layer on. - input_tensor = base._DeferredTensor( # pylint: disable=protected-access - shape=batch_input_shape, - dtype=dtype, - name=self.name) - else: - # In graph mode, create a graph placeholder to call the layer on. - if sparse: - input_tensor = array_ops.sparse_placeholder( - shape=batch_input_shape, - dtype=dtype, - name=self.name) - else: - input_tensor = array_ops.placeholder( - shape=batch_input_shape, - dtype=dtype, - name=self.name) - - # For compatibility with Keras API. - self.is_placeholder = True - self._batch_input_shape = batch_input_shape - else: - # For compatibility with Keras API. - self.is_placeholder = False - self._batch_input_shape = tuple(input_tensor.get_shape().as_list()) - - # Create an input node to add to self.outbound_node - # and set output_tensors' _keras_history. - input_tensor._keras_history = (self, 0, 0) # pylint: disable=protected-access - base.Node( - self, - inbound_layers=[], - node_indices=[], - tensor_indices=[], - input_tensors=[input_tensor], - output_tensors=[input_tensor]) - - -def Input( # pylint: disable=invalid-name - shape=None, - batch_size=None, - name=None, - dtype=dtypes.float32, - sparse=False, - tensor=None): - """`Input()` is used to instantiate an input tensor for use with a `Network`. - - For instance, if a, b and c are tensors created via `Input`, - it becomes possible to do: - - `network = Network(inputs=[a, b], outputs=c)` - - Example: - - ```python - # This is a logistic regression - x = tf.layers.Input(shape=(32,)) - y = tf.layers.Dense(16, activation='softmax')(x) - network = tf.layers.Network(x, y) - ``` - - Arguments: - shape: A shape tuple (integer), not including the batch size. - For instance, `shape=(32,)` indicates that the expected input - will be batches of 32-dimensional vectors. - batch_size: Optional input batch size (integer or None). - name: An optional name string for the layer. - Should be unique in a model (do not reuse the same name twice). - It will be autogenerated if it isn't provided. - dtype: The data type expected by the input, as a string - (`float32`, `float64`, `int32`...) - sparse: A boolean specifying whether the placeholder - to be created is sparse. - tensor: Optional existing tensor to wrap into the `Input` layer. - If set, the layer will not create a placeholder tensor. - - Returns: - A tensor: either a new placeholder (with history metadata) or - `tensor` (if passed), with added history metadata. - - Raises: - RuntimeError: If called in Eager mode. - """ - input_layer = InputLayer( - input_shape=shape, - batch_size=batch_size, - name=name, - dtype=dtype, - sparse=sparse, - input_tensor=tensor) - # Return tensor including `_keras_history` metadata. - # Note that in this case train_output and test_output are the same pointer. - outputs = input_layer._inbound_nodes[0].output_tensors # pylint: disable=protected-access - if len(outputs) == 1: - return outputs[0] - else: - return outputs - - -class GraphNetwork(base.Layer): - """A GraphNetwork is a directed acyclic graph of layers. - - It is the topological form of a `tf.keras.models.Model`. A `Model` is simply a - `GraphNetwork` with added training/evaluation routines. - - A `GraphNetwork` instance implements the full `Layer` API. In particular, a - `GraphNetwork` can be called on new inputs. - - Example: - - ```python - # This is a logistic regression - x = tf.layers.Input(shape=(32,)) - y = tf.layers.Dense(16, activation='softmax')(x) - network = tf.layers.GraphNetwork(x, y) - - # It is then possible to call the network on compatible inputs: - z = tf.layers.Input(shape=(32,)) - w = network(z) - - # It is possible to retrieve the same properties as a layer: - weights = network.trainable_weights - ``` - - Arguments: - inputs: Input tensor or list of input tensors. - Must come from `tf.layers.Input`. - output: Output tensor or list of output tensors. Must come from - tf.layers Layers or Keras layers. - name: Optional name of the model (string). - - Attributes: - GraphNetwork has the same attributes as Layer. On top of it, it also has: - - layers: a list of the children layers of the network, - a list of layer instances, ordered from "earlier in the graph" - to "later in the graph". - - Methods: - GraphNetwork has the same methods as Layer. On top of it, it also has: - - get_layer: retrieves a child layer by name or index in the graph. - - Raises: - RuntimeError: If created in Eager mode. - """ - - def __init__(self, inputs, outputs, name=None): # pylint: disable=super-init-not-called - if context.in_eager_mode(): - # TODO(fchollet): check that all inputs and outputs are DeferredTensors. - pass - - self._init_set_name(name) - self._activity_regularizer = None - with vs.variable_scope( - None, default_name=self._base_name) as captured_scope: - self._scope = captured_scope - call_fn_args = estimator_util.fn_args(self.call) - self._compute_previous_mask = ('mask' in call_fn_args or - hasattr(self, 'compute_mask')) - self._call_has_scope_arg = 'scope' in call_fn_args - - # This acts just like the `trainable` attribute of any layer instance. - # It does not affect users of the underlying layers, only users of the - # GraphNetwork instance. - self.trainable = True - # A GraphNetwork does not create weights of its own, thus it is already - # built. - self.built = True - # A GraphNetwork does not create weights of its own, thus has no dtype. - self._dtype = None - # The following are implemented as property functions: - # self.trainable_weights - # self.non_trainable_weights - # self.input_spec - - # Private attributes to implement compatibility with Layer. - self._per_input_losses = {} - self._per_input_updates = {} - self._updates = [] - self._losses = [] - self._scope = None - self._reuse = None - self._graph = ops.get_default_graph() - - # GraphNetwork-specific properties. - if isinstance(inputs, (list, tuple)): - self.inputs = list(inputs) # Tensor or list of tensors. - else: - self.inputs = [inputs] - if isinstance(outputs, (list, tuple)): - self.outputs = list(outputs) - else: - self.outputs = [outputs] - # All layers in order of horizontal graph traversal. - # Entries are unique. Includes input and output layers. - self.layers = [] - - # Check for redundancy in inputs. - if len(set(self.inputs)) != len(self.inputs): - raise ValueError('The list of inputs passed to the model ' - 'is redundant. ' - 'All inputs should only appear once.' - ' Found: ' + str(self.inputs)) - - # # List of initial layers (1 to 1 mapping with self.inputs, - # # hence the same layer might appear twice) - # self._input_layers = [] - # self._input_layers_node_indices = [] - # self._input_layers_tensor_indices = [] - # # list of layers (1 to 1 mapping with self.inputs, - # # hence the same layer might appear twice) - # self._output_layers = [] - # self._output_layers_node_indices = [] - # self._output_layers_tensor_indices = [] - - self._input_layers = [] - self._output_layers = [] - self._input_coordinates = [] - self._output_coordinates = [] - - # This is for performance optimization when calling the GraphNetwork on new - # inputs. Every time the GraphNetwork is called on a set on input tensors, - # we compute the output tensors, output masks and output shapes in one pass, - # then cache them here. When any of these outputs is queried later, we - # retrieve it from there instead of recomputing it. - self._output_mask_cache = {} - self._output_tensor_cache = {} - self._output_shape_cache = {} - - # User-provided arguments validation. - for x in self.inputs: - # Check that x has appropriate `_keras_history` metadata. - if not hasattr(x, '_keras_history'): - cls_name = self.__class__.__name__ - raise ValueError('Input tensors to a ' + cls_name + ' ' + - 'must come from `tf.layers.Input`. ' - 'Received: ' + str(x) + - ' (missing previous layer metadata).') - # Check that x is an input tensor. - # pylint: disable=protected-access - layer, node_index, tensor_index = x._keras_history - if len(layer._inbound_nodes) > 1 or ( - layer._inbound_nodes and layer._inbound_nodes[0].inbound_layers): - cls_name = self.__class__.__name__ - logging.warning(cls_name + ' inputs must come from ' - '`tf.layers.Input` (thus holding past layer metadata), ' - 'they cannot be the output of ' - 'a previous non-Input layer. ' - 'Here, a tensor specified as ' - 'input to "' + self.name + '" was not an Input tensor, ' - 'it was generated by layer ' + layer.name + '.\n' - 'Note that input tensors are ' - 'instantiated via `tensor = tf.layers.Input(shape)`.\n' - 'The tensor that caused the issue was: ' + str(x.name)) - # pylint: enable=protected-access - for x in self.outputs: - if not hasattr(x, '_keras_history'): - cls_name = self.__class__.__name__ - raise ValueError('Output tensors to a ' + cls_name + ' must be ' - 'the output of a TensorFlow `Layer` ' - '(thus holding past layer metadata). Found: ' + str(x)) - - # Build self._output_layers: - for x in self.outputs: - layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access - self._output_layers.append(layer) - self._output_coordinates.append((layer, node_index, tensor_index)) - - # Build self._input_layers: - for x in self.inputs: - layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access - # It's supposed to be an input layer, so only one node - # and one tensor output. - assert node_index == 0 - assert tensor_index == 0 - self._input_layers.append(layer) - self._input_coordinates.append((layer, node_index, tensor_index)) - - # Network_nodes: set of nodes included in the graph - # (not all nodes included in the layers - # are relevant to the current graph). - network_nodes = set() # ids of all nodes relevant to the GraphNetwork - nodes_depths = {} # dict {node: depth value} - layers_depths = {} # dict {layer: depth value} - layer_indices = {} # dict {layer: index in traversal} - nodes_in_decreasing_depth = [] - - def build_map_of_graph(tensor, - finished_nodes, - nodes_in_progress, - layer, - node_index, - tensor_index): - """Builds a map of the graph of layers. - - This recursively updates the map `layer_indices`, - the list `nodes_in_decreasing_depth` and the set `network_nodes`. - - Arguments: - tensor: Some tensor in a graph. - finished_nodes: Set of nodes whose subgraphs have been traversed - completely. Useful to prevent duplicated work. - nodes_in_progress: Set of nodes that are currently active on the - recursion stack. Useful to detect cycles. - layer: Layer from which `tensor` comes from. If not provided, - will be obtained from `tensor._keras_history`. - node_index: Node index from which `tensor` comes from. - tensor_index: Tensor_index from which `tensor` comes from. - - Raises: - ValueError: if a cycle is detected. - """ - node = layer._inbound_nodes[node_index] # pylint: disable=protected-access - - # Prevent cycles. - if node in nodes_in_progress: - raise ValueError('The tensor ' + str(tensor) + ' at layer "' + - layer.name + '" is part of a cycle.') - - # Don't repeat work for shared subgraphs - if node in finished_nodes: - return - - node_key = _make_node_key(layer.name, node_index) - # Update network_nodes. - network_nodes.add(node_key) - - # Store the traversal order for layer sorting. - if layer not in layer_indices: - layer_indices[layer] = len(layer_indices) - - nodes_in_progress.add(node) - - # Propagate to all previous tensors connected to this node. - for i in range(len(node.inbound_layers)): - x = node.input_tensors[i] - layer = node.inbound_layers[i] - node_index = node.node_indices[i] - tensor_index = node.tensor_indices[i] - build_map_of_graph(x, finished_nodes, nodes_in_progress, layer, - node_index, tensor_index) - - finished_nodes.add(node) - nodes_in_progress.remove(node) - nodes_in_decreasing_depth.append(node) - - finished_nodes = set() - nodes_in_progress = set() - for x in self.outputs: - layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access - build_map_of_graph(x, finished_nodes, nodes_in_progress, - layer=layer, - node_index=node_index, - tensor_index=tensor_index) - - for node in reversed(nodes_in_decreasing_depth): - # If the depth is not set, the node has no outbound nodes (depth 0). - depth = nodes_depths.setdefault(node, 0) - - # Update the depth of the corresponding layer - previous_depth = layers_depths.get(node.outbound_layer, 0) - # If we've seen this layer before at a higher depth, - # we should use that depth instead of the node depth. - # This is necessary for shared layers that have inputs at different - # depth levels in the graph. - depth = max(depth, previous_depth) - layers_depths[node.outbound_layer] = depth - nodes_depths[node] = depth - - # Update the depth of inbound nodes. - # The "depth" of a node is the max of the depths - # of all layers it is connected to. - for i in range(len(node.inbound_layers)): - inbound_layer = node.inbound_layers[i] - node_index = node.node_indices[i] - inbound_node = inbound_layer._inbound_nodes[node_index] # pylint: disable=protected-access - previous_depth = nodes_depths.get(inbound_node, 0) - nodes_depths[inbound_node] = max(depth + 1, previous_depth) - - # Build a dict {depth: list of nodes with this depth} - nodes_by_depth = {} - for node, depth in nodes_depths.items(): - if depth not in nodes_by_depth: - nodes_by_depth[depth] = [] - nodes_by_depth[depth].append(node) - - # Build a dict {depth: list of layers with this depth} - layers_by_depth = {} - for layer, depth in layers_depths.items(): - if depth not in layers_by_depth: - layers_by_depth[depth] = [] - layers_by_depth[depth].append(layer) - - # Get sorted list of layer depths. - depth_keys = list(layers_by_depth.keys()) - depth_keys.sort(reverse=True) - - # Set self.layers and self._layers_by_depth. - layers = [] - for depth in depth_keys: - layers_for_depth = layers_by_depth[depth] - # GraphNetwork.layers needs to have a deterministic order: - # here we order them by traversal order. - layers_for_depth.sort(key=lambda x: layer_indices[x]) - layers.extend(layers_for_depth) - self.layers = layers - self._layers_by_depth = layers_by_depth - - # Get sorted list of node depths. - depth_keys = list(nodes_by_depth.keys()) - depth_keys.sort(reverse=True) - - # Check that all tensors required are computable. - # computable_tensors: all tensors in the graph - # that can be computed from the inputs provided. - computable_tensors = [] - for x in self.inputs: - computable_tensors.append(x) - - layers_with_complete_input = [] # To provide a better error msg. - for depth in depth_keys: - for node in nodes_by_depth[depth]: - layer = node.outbound_layer - if layer: - for x in node.input_tensors: - if x not in computable_tensors: - raise ValueError('Graph disconnected: ' - 'cannot obtain value for tensor ' + str(x) + - ' at layer "' + layer.name + '". ' - 'The following previous layers ' - 'were accessed without issue: ' + - str(layers_with_complete_input)) - for x in node.output_tensors: - computable_tensors.append(x) - layers_with_complete_input.append(layer.name) - - # Keep track of the network's nodes. - self._network_nodes = network_nodes - self._nodes_by_depth = nodes_by_depth - - # Ensure name unicity, which will be crucial for serialization - # (since serialized nodes refer to layers by their name). - all_names = [layer.name for layer in self.layers] - for name in all_names: - if all_names.count(name) != 1: - raise ValueError('The name "' + name + '" is used ' + - str(all_names.count(name)) + ' times in the model. ' - 'All layer names should be unique.') - - # Layer parameters. - # The new network starts with a single inbound node - # for its inputs, and no outbound nodes. - self._outbound_nodes = [] # Will be appended to by future calls to __call__ - self._inbound_nodes = [ - ] # Will be appended to below, and by future calls to __call__ - # Create the node linking internal inputs to internal outputs. - base.Node( - outbound_layer=self, - inbound_layers=[], - node_indices=[], - tensor_indices=[], - input_tensors=self.inputs, - output_tensors=self.outputs) - - def get_layer(self, name=None, index=None): - """Retrieves a layer based on either its name (unique) or index. - - Indices are based on order of horizontal graph traversal (bottom-up). - - Arguments: - name: String, name of layer. - index: Integer, index of layer. - - Returns: - A layer instance. - - Raises: - ValueError: In case of invalid layer name or index. - """ - # TODO(fchollet): We could build a dictionary based on layer names - # since they are constant, but we have not done that yet. - if index is not None: - if len(self.layers) <= index: - raise ValueError('Was asked to retrieve layer at index ' + str(index) + - ' but model only has ' + str(len(self.layers)) + - ' layers.') - else: - return self.layers[index] - else: - if not name: - raise ValueError('Provide either a layer name or layer index.') - for layer in self.layers: - if layer.name == name: - return layer - raise ValueError('No such layer: ' + name) - - @property - def stateful(self): - return any([(hasattr(layer, 'stateful') and layer.stateful) - for layer in self.layers]) - - @property - def updates(self): - """Retrieve the network's updates. - - Will only include updates that are either - unconditional, or conditional on inputs to this model - (e.g. will not include updates that depend on tensors - that aren't inputs to this model). - - Returns: - A list of update ops. - """ - if not self.trainable and not self.stateful: - return [] - updates = [] - for layer in self.layers: - if hasattr(layer, 'updates'): - # Collect updates that are dependent on inputs - # that are part of the model. - for node_index, node in enumerate(layer._inbound_nodes): # pylint: disable=protected-access - node_key = _make_node_key(layer.name, node_index) - if node_key in self._network_nodes: - # The model owns this layer node. - inputs = node.input_tensors - updates += layer.get_updates_for(inputs) - # Collect unconditional updates. - updates += layer.get_updates_for(None) - return updates - - @property - def losses(self): - """Retrieve the network's losses. - - Will only include losses that are either - unconditional, or conditional on inputs to this model - (e.g. will not include losses that depend on tensors - that aren't inputs to this model). - - Returns: - A list of loss tensors. - """ - losses = [] - # Retrieve losses for all internal layers. - for layer in self.layers: - if hasattr(layer, 'losses'): - # Collect losses that are dependent on inputs - # that are part of the model. - for node_index, node in enumerate(layer._inbound_nodes): # pylint: disable=protected-access - node_key = _make_node_key(layer.name, node_index) - if node_key in self._network_nodes: - # The model owns this layer node. - inputs = node.input_tensors - losses += layer.get_losses_for(inputs) - # Collect unconditional losses. - losses += layer.get_losses_for(None) - # Add any potential unconditional model-level loss. - losses += self.get_losses_for(None) - return losses - - @property - def trainable_weights(self): - if not self.trainable: - return [] - weights = [] - for layer in self.layers: - weights += layer.trainable_weights - return weights - - @property - def non_trainable_weights(self): - weights = [] - for layer in self.layers: - weights += layer.non_trainable_weights - if not self.trainable: - trainable_weights = [] - for layer in self.layers: - trainable_weights += layer.trainable_weights - return trainable_weights + weights - return weights - - @property - def input_spec(self): - """Gets the network's input specs. - - Returns: - A list of `InputSpec` instances (one per input to the model) - or a single instance if the model has only one input. - """ - specs = [] - for layer in self._input_layers: - if layer.input_spec is None: - specs.append(None) - else: - if not isinstance(layer.input_spec, list): - raise TypeError('Layer ' + layer.name + - ' has an input_spec attribute that ' - 'is not a list. We expect a list. ' - 'Found input_spec = ' + str(layer.input_spec)) - specs += layer.input_spec - if len(specs) == 1: - return specs[0] - return specs - - def call(self, inputs, mask=None): - """Call the model on new inputs. - - In this case `call` just reapplies - all ops in the graph to the new inputs - (e.g. build a new computational graph from the provided inputs). - - Arguments: - inputs: A tensor or list of tensors. - mask: A mask or list of masks. A mask can be - either a tensor or None (no mask). - - Returns: - A tensor if there is a single output, or - a list of tensors if there are more than one outputs. - """ - inputs = nest.flatten(inputs) - if mask is None: - masks = [None for _ in range(len(inputs))] - else: - masks = nest.flatten(mask) - - if context.in_graph_mode(): - # Try to retrieve cached outputs if the layer has already been called - # on these exact inputs. - cache_key = (layers_util.object_list_uid(inputs) - + '_' + layers_util.object_list_uid(masks)) - if cache_key in self._output_tensor_cache: - # Cache hit. - return self._output_tensor_cache[cache_key] - # Actually apply the network graph to the new inputs. - outputs, _ = self._run_internal_graph(inputs, masks) - return outputs - - def compute_output_shape(self, input_shape): - if isinstance(input_shape, list): - input_shapes = [] - for shape in input_shape: - if shape is not None: - input_shapes.append(tuple(tensor_shape.TensorShape(shape).as_list())) - else: - input_shapes.append(None) - else: - if input_shape is not None: - input_shapes = [tuple(tensor_shape.TensorShape(input_shape).as_list())] - else: - input_shapes = [None] - - if len(input_shapes) != len(self._input_layers): - raise ValueError('Invalid input_shape argument ' + str(input_shape) + - ': model has ' + str(len(self._input_layers)) + - ' tensor inputs.') - - cache_key = layers_util.object_list_uid(input_shapes) - if cache_key not in self._output_shape_cache: - # Cache miss. We have to run the network graph manually (recursive calls - # to `compute_output_shape`). - layers_to_output_shapes = {} - for i in range(len(input_shapes)): - layer = self._input_layers[i] - input_shape = input_shapes[i] - # It's an input layer: then `compute_output_shape` is identity, - # and there is only one node and one tensor output. - shape_key = layer.name + '_0_0' - layers_to_output_shapes[shape_key] = input_shape - - depth_keys = list(self._nodes_by_depth.keys()) - depth_keys.sort(reverse=True) - # Iterate over nodes, by depth level. - if len(depth_keys) > 1: - for depth in depth_keys: - nodes = self._nodes_by_depth[depth] - for node in nodes: - # This is always a single layer, never a list. - layer = node.outbound_layer - if layer in self._input_layers: - # We've already covered the input layers - # a few lines above. - continue - # Potentially redundant list, - # same size as node.input_tensors. - input_shapes = [] - for j in range(len(node.inbound_layers)): - inbound_layer = node.inbound_layers[j] - node_index = node.node_indices[j] - tensor_index = node.tensor_indices[j] - shape_key = inbound_layer.name + '_%s_%s' % (node_index, - tensor_index) - input_shape = layers_to_output_shapes[shape_key] - input_shapes.append(input_shape) - - if len(input_shapes) == 1: - output_shape = layer.compute_output_shape(input_shapes[0]) - else: - output_shape = layer.compute_output_shape(input_shapes) - if isinstance(output_shape, list): - output_shapes = [ - tuple(tensor_shape.TensorShape(shape).as_list()) - for shape in output_shape - ] - else: - output_shapes = [ - tuple(tensor_shape.TensorShape(output_shape).as_list()) - ] - - node_index = layer._inbound_nodes.index(node) # pylint: disable=protected-access - for j in range(len(output_shapes)): - shape_key = layer.name + '_%s_%s' % (node_index, j) - layers_to_output_shapes[shape_key] = output_shapes[j] - - # Read final output shapes from layers_to_output_shapes. - output_shapes = [] - for i in range(len(self._output_layers)): - layer, node_index, tensor_index = self._output_coordinates[i] - shape_key = layer.name + '_%s_%s' % (node_index, tensor_index) - output_shapes.append(layers_to_output_shapes[shape_key]) - - # Store in cache. - self._output_shape_cache[cache_key] = output_shapes - else: - # Cache hit. - output_shapes = self._output_shape_cache[cache_key] - - if isinstance(output_shapes, list): - if len(output_shapes) == 1: - return tensor_shape.TensorShape(output_shapes[0]) - else: - return [tensor_shape.TensorShape(shape) for shape in output_shapes] - else: - return tensor_shape.TensorShape(output_shapes) - - def _run_internal_graph(self, inputs, masks=None): - """Computes output tensors for new inputs. - - # Note: - - Expects `inputs` to be a list (potentially with 1 element). - - Can be run on non-Keras tensors. - - Arguments: - inputs: List of tensors - masks: List of masks (tensors or None). - - Returns: - Three lists: output_tensors, output_masks, output_shapes - """ - # Note: masking support is relevant mainly for Keras. - # It cannot be factored out without having the fully reimplement the network - # calling logic on the Keras side. We choose to incorporate it in - # GraphNetwork because 1) it may be useful to fully support in tf.layers in - # the future and 2) Keras is a major user of GraphNetwork. If you don't - # use masking, it does not interfere with regular behavior at all and you - # can ignore it. - if masks is None: - masks = [None for _ in range(len(inputs))] - - # Dictionary mapping reference tensors to tuples - # (computed tensor, compute mask) - # we assume a 1:1 mapping from tensor to mask - # TODO(fchollet): raise exception when a `.compute_mask()` call - # does not return a list the same size as `call` - tensor_map = {} - for x, y, mask in zip(self.inputs, inputs, masks): - tensor_map[str(id(x))] = (y, mask) - - depth_keys = list(self._nodes_by_depth.keys()) - depth_keys.sort(reverse=True) - for depth in depth_keys: - nodes = self._nodes_by_depth[depth] - for node in nodes: - # This is always a single layer, never a list. - layer = node.outbound_layer - - reference_input_tensors = node.input_tensors - reference_output_tensors = node.output_tensors - - # If all previous input tensors are available in tensor_map, - # then call node.inbound_layer on them. - computed_data = [] # List of tuples (input, mask). - for x in reference_input_tensors: - if str(id(x)) in tensor_map: - computed_data.append(tensor_map[str(id(x))]) - - if len(computed_data) == len(reference_input_tensors): - # Call layer (reapplying ops to new inputs). - with ops.name_scope(layer.name): - if node.arguments: - kwargs = node.arguments - else: - kwargs = {} - if len(computed_data) == 1: - computed_tensor, computed_mask = computed_data[0] - # Ensure mask propagation if applicable. - if 'mask' in estimator_util.fn_args(layer.call): - if 'mask' not in kwargs: - kwargs['mask'] = computed_mask - - output_tensors = nest.flatten( - layer.call(computed_tensor, **kwargs)) - if hasattr(layer, 'compute_mask'): - output_masks = nest.flatten( - layer.compute_mask(computed_tensor, computed_mask)) - else: - output_masks = [None for _ in range(len(output_tensors))] - computed_tensors = [computed_tensor] - computed_masks = [computed_mask] - else: - computed_tensors = [x[0] for x in computed_data] - computed_masks = [x[1] for x in computed_data] - if 'mask' in estimator_util.fn_args(layer.call): - if 'mask' not in kwargs: - kwargs['mask'] = computed_masks - output_tensors = nest.flatten( - layer.call(computed_tensors, **kwargs)) - if hasattr(layer, 'compute_mask'): - output_masks = nest.flatten( - layer.compute_mask(computed_tensors, computed_masks)) - else: - output_masks = [None for _ in range(len(output_tensors))] - - # Apply activity regularizer if any: - if layer.activity_regularizer is not None: - regularization_losses = [ - layer.activity_regularizer(x) for x in computed_tensors - ] - layer.add_loss(regularization_losses, computed_tensors) - - if context.in_graph_mode(): - # Update model updates and losses: - # Keep track of updates that depend on the inputs - # (e.g. BN updates). - self.add_update(layer.get_updates_for(computed_tensors), inputs) - # Keep track of unconditional updates (e.g. a counter). - self.add_update(layer.get_updates_for(None), None) - # Keep track of losses that depend on the inputs - # (e.g. activity regularizers). - self.add_loss(layer.get_losses_for(computed_tensors), inputs) - # Keep track of unconditional losses - # (e.g. weight regularizers). - self.add_loss(layer.get_losses_for(None), None) - - # Update tensor_map. - for x, y, mask in zip(reference_output_tensors, output_tensors, - output_masks): - tensor_map[str(id(x))] = (y, mask) - - output_tensors = [] - output_masks = [] - output_shapes = [] - for x in self.outputs: - assert str(id(x)) in tensor_map, 'Could not compute output ' + str(x) - tensor, mask = tensor_map[str(id(x))] - output_shapes.append(layers_util.static_shape(x)) - output_tensors.append(tensor) - output_masks.append(mask) - - if len(output_tensors) == 1: - output_tensors = output_tensors[0] - if output_shapes is not None: - output_shapes = output_shapes[0] - if output_masks is not None: - output_masks = output_masks[0] - - if context.in_graph_mode(): - # Update cache; - # keys are based on ids on input tensors and inputs masks. - cache_key = (layers_util.object_list_uid(inputs) - + '_' + layers_util.object_list_uid(masks)) - self._output_tensor_cache[cache_key] = output_tensors - if output_masks is not None: - self._output_mask_cache[cache_key] = output_masks - if output_shapes is not None: - input_shapes = [layers_util.static_shape(x) for x in inputs] - cache_key = layers_util.object_list_uid(input_shapes) - self._output_shape_cache[cache_key] = output_shapes - - return output_tensors, output_masks - - -def _make_node_key(layer_name, node_index): - return layer_name + '_ib-' + str(node_index) diff --git a/tensorflow/python/layers/network_test.py b/tensorflow/python/layers/network_test.py deleted file mode 100644 index 7a2c7fb3fc782f6e59b8b483ec43c4abddf4d023..0000000000000000000000000000000000000000 --- a/tensorflow/python/layers/network_test.py +++ /dev/null @@ -1,525 +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. -# ============================================================================== -"""Tests for tf.layers.network.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.python.eager import context -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import test_util -from tensorflow.python.layers import base as base_layers -from tensorflow.python.layers import core as core_layers -from tensorflow.python.layers import network as network_layers -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import sparse_ops -from tensorflow.python.platform import test - - -class BaseLayerCompatibilityTest(test.TestCase): - - def test_get_updates_for(self): - a = network_layers.Input(shape=(2,)) - dense_layer = core_layers.Dense(1) - dense_layer.add_update(0, inputs=a) - dense_layer.add_update(1, inputs=None) - - self.assertEqual(dense_layer.get_updates_for(a), [0]) - self.assertEqual(dense_layer.get_updates_for(None), [1]) - - def test_get_losses_for(self): - a = network_layers.Input(shape=(2,)) - dense_layer = core_layers.Dense(1) - dense_layer.add_loss(0, inputs=a) - dense_layer.add_loss(1, inputs=None) - - self.assertEqual(dense_layer.get_losses_for(a), [0]) - self.assertEqual(dense_layer.get_losses_for(None), [1]) - - def testTopologicalAttributes(self): - # test layer attributes / methods related to cross-layer connectivity. - a = network_layers.Input(shape=(32,), name='input_a') - b = network_layers.Input(shape=(32,), name='input_b') - - # test input, output, input_shape, output_shape - test_layer = core_layers.Dense(16, name='test_layer') - a_test = test_layer(a) - self.assertEqual(test_layer.input, a) - self.assertEqual(test_layer.output, a_test) - self.assertEqual(test_layer.input_shape, (None, 32)) - self.assertEqual(test_layer.output_shape, (None, 16)) - - # test `get_*_at` methods - dense = core_layers.Dense(16, name='dense_1') - a_2 = dense(a) - b_2 = dense(b) - - self.assertEqual(dense.get_input_at(0), a) - self.assertEqual(dense.get_input_at(1), b) - self.assertEqual(dense.get_output_at(0), a_2) - self.assertEqual(dense.get_output_at(1), b_2) - self.assertEqual(dense.get_input_shape_at(0), (None, 32)) - self.assertEqual(dense.get_input_shape_at(1), (None, 32)) - self.assertEqual(dense.get_output_shape_at(0), (None, 16)) - self.assertEqual(dense.get_output_shape_at(1), (None, 16)) - - # Test invalid value for attribute retrieval. - with self.assertRaises(ValueError): - dense.get_input_at(2) - with self.assertRaises(AttributeError): - new_dense = core_layers.Dense(16) - _ = new_dense.input - with self.assertRaises(AttributeError): - new_dense = core_layers.Dense(16) - _ = new_dense.output - with self.assertRaises(AttributeError): - new_dense = core_layers.Dense(16) - _ = new_dense.output_shape - with self.assertRaises(AttributeError): - new_dense = core_layers.Dense(16) - _ = new_dense.input_shape - with self.assertRaises(AttributeError): - new_dense = core_layers.Dense(16) - a = network_layers.Input(shape=(3, 32)) - a = network_layers.Input(shape=(5, 32)) - a_2 = dense(a) - b_2 = dense(b) - _ = new_dense.input_shape - with self.assertRaises(AttributeError): - new_dense = core_layers.Dense(16) - a = network_layers.Input(shape=(3, 32)) - a = network_layers.Input(shape=(5, 32)) - a_2 = dense(a) - b_2 = dense(b) - _ = new_dense.output_shape - - def testTopologicalAttributesMultiOutputLayer(self): - - class PowersLayer(base_layers.Layer): - - def call(self, inputs): - return [inputs**2, inputs**3] - - x = network_layers.Input(shape=(32,)) - test_layer = PowersLayer() - p1, p2 = test_layer(x) # pylint: disable=not-callable - - self.assertEqual(test_layer.input, x) - self.assertEqual(test_layer.output, [p1, p2]) - self.assertEqual(test_layer.input_shape, (None, 32)) - self.assertEqual(test_layer.output_shape, [(None, 32), (None, 32)]) - - def testTopologicalAttributesMultiInputLayer(self): - - class AddLayer(base_layers.Layer): - - def call(self, inputs): - assert len(inputs) == 2 - return inputs[0] + inputs[1] - - a = network_layers.Input(shape=(32,)) - b = network_layers.Input(shape=(32,)) - test_layer = AddLayer() - y = test_layer([a, b]) # pylint: disable=not-callable - - self.assertEqual(test_layer.input, [a, b]) - self.assertEqual(test_layer.output, y) - self.assertEqual(test_layer.input_shape, [(None, 32), (None, 32)]) - self.assertEqual(test_layer.output_shape, (None, 32)) - - -class NetworkTest(test.TestCase): - - def testBasicNetwork(self): - # minimum viable network - x = network_layers.Input(shape=(32,)) - dense = core_layers.Dense(2) - y = dense(x) - network = network_layers.GraphNetwork(x, y, name='dense_network') - - # test basic attributes - self.assertEqual(network.name, 'dense_network') - self.assertEqual(len(network.layers), 2) # InputLayer + Dense - self.assertEqual(network.layers[1], dense) - self.assertEqual(network.weights, dense.weights) - self.assertEqual(network.trainable_weights, dense.trainable_weights) - self.assertEqual(network.non_trainable_weights, dense.non_trainable_weights) - - # test callability on Input - x_2 = network_layers.Input(shape=(32,)) - y_2 = network(x_2) - self.assertEqual(y_2.get_shape().as_list(), [None, 2]) - - # test callability on regular tensor - x_2 = array_ops.placeholder(dtype='float32', shape=(None, 32)) - y_2 = network(x_2) - self.assertEqual(y_2.get_shape().as_list(), [None, 2]) - - # test network `trainable` attribute - network.trainable = False - self.assertEqual(network.weights, dense.weights) - self.assertEqual(network.trainable_weights, []) - self.assertEqual(network.non_trainable_weights, - dense.trainable_weights + dense.non_trainable_weights) - - def test_node_construction(self): - # test graph topology construction basics - a = network_layers.Input(shape=(32,), name='input_a') - b = network_layers.Input(shape=(32,), name='input_b') - - self.assertEqual(a.get_shape().as_list(), [None, 32]) - a_layer, a_node_index, a_tensor_index = a._keras_history - b_layer, _, _ = b._keras_history - self.assertEqual(len(a_layer._inbound_nodes), 1) - self.assertEqual(a_tensor_index, 0) - node = a_layer._inbound_nodes[a_node_index] - self.assertEqual(node.outbound_layer, a_layer) - - self.assertEqual(node.inbound_layers, []) - self.assertEqual(node.input_tensors, [a]) - self.assertEqual(node.input_shapes, [(None, 32)]) - self.assertEqual(node.output_tensors, [a]) - self.assertEqual(node.output_shapes, [(None, 32)]) - - dense = core_layers.Dense(16, name='dense_1') - dense(a) - dense(b) - - self.assertEqual(len(dense._inbound_nodes), 2) - self.assertEqual(len(dense._outbound_nodes), 0) - self.assertEqual(dense._inbound_nodes[0].inbound_layers, [a_layer]) - self.assertEqual(dense._inbound_nodes[0].outbound_layer, dense) - self.assertEqual(dense._inbound_nodes[1].inbound_layers, [b_layer]) - self.assertEqual(dense._inbound_nodes[1].outbound_layer, dense) - self.assertEqual(dense._inbound_nodes[0].input_tensors, [a]) - self.assertEqual(dense._inbound_nodes[1].input_tensors, [b]) - - # Test config - config_0 = dense._inbound_nodes[0].get_config() - self.assertEqual(config_0['outbound_layer'], dense.name) - - def testMultiInputNetwork(self): - a = network_layers.Input(shape=(32,), name='input_a') - b = network_layers.Input(shape=(32,), name='input_b') - - class AddLayer(base_layers.Layer): - - def call(self, inputs): - assert len(inputs) == 2 - return inputs[0] + inputs[1] - - c = AddLayer()([a, b]) # pylint: disable=not-callable - network = network_layers.GraphNetwork([a, b], c) - self.assertEqual(len(network.layers), 3) # 2 * InputLayer + AddLayer - - # Test callability. - a2 = network_layers.Input(shape=(32,)) - b2 = network_layers.Input(shape=(32,)) - c2 = network([a2, b2]) - self.assertEqual(c2.get_shape().as_list(), [None, 32]) - - def testMultiOutputNetwork(self): - x = network_layers.Input(shape=(32,)) - y1 = core_layers.Dense(2)(x) - y2 = core_layers.Dense(3)(x) - network = network_layers.GraphNetwork(x, [y1, y2]) - - self.assertEqual(len(network.layers), 3) # InputLayer + 2 * Dense - - # Test callability. - x2 = network_layers.Input(shape=(32,)) - outputs = network(x2) - - self.assertEqual(type(outputs), list) - self.assertEqual(len(outputs), 2) - self.assertEqual(outputs[0].get_shape().as_list(), [None, 2]) - self.assertEqual(outputs[1].get_shape().as_list(), [None, 3]) - - def testMultiInputMultiOutputNetworkSharedLayer(self): - a = network_layers.Input(shape=(32,), name='input_a') - b = network_layers.Input(shape=(32,), name='input_b') - - dense = core_layers.Dense(2) - - y1 = dense(a) - y2 = dense(b) - network = network_layers.GraphNetwork([a, b], [y1, y2]) - self.assertEqual(len(network.layers), 3) # 2 * InputLayer + Dense - - # Test callability. - a2 = network_layers.Input(shape=(32,)) - b2 = network_layers.Input(shape=(32,)) - outputs = network([a2, b2]) - - self.assertEqual(type(outputs), list) - self.assertEqual(len(outputs), 2) - self.assertEqual(outputs[0].get_shape().as_list(), [None, 2]) - self.assertEqual(outputs[1].get_shape().as_list(), [None, 2]) - - def testCrossDataFlows(self): - # Test the ability to have multi-output layers with outputs that get routed - # to separate layers - - class PowersLayer(base_layers.Layer): - - def call(self, inputs): - return [inputs**2, inputs**3] - - x = network_layers.Input(shape=(32,)) - p1, p2 = PowersLayer()(x) # pylint: disable=not-callable - y1 = core_layers.Dense(2)(p1) - y2 = core_layers.Dense(3)(p2) - network = network_layers.GraphNetwork(x, [y1, y2]) - - self.assertEqual(len(network.layers), 4) # InputLayer + 2 * Dense + PLayer - - # Test callability. - x2 = network_layers.Input(shape=(32,)) - outputs = network(x2) - - self.assertEqual(type(outputs), list) - self.assertEqual(len(outputs), 2) - self.assertEqual(outputs[0].get_shape().as_list(), [None, 2]) - self.assertEqual(outputs[1].get_shape().as_list(), [None, 3]) - - def testNetworkAttributes(self): - x = network_layers.Input(shape=(32,)) - z = core_layers.Dense(2, kernel_regularizer=lambda x: 0.01 * (x**2))(x) - dense = core_layers.Dense(2, name='dense') - dense.add_update(1) - y = dense(z) - net = network_layers.GraphNetwork(x, y) - - # losses - self.assertEqual(len(net.losses), 1) - - # updates - self.assertEqual(len(net.updates), 1) - - # get_layer - self.assertEqual(net.get_layer('dense'), dense) - self.assertEqual(net.get_layer(index=2), dense) - with self.assertRaises(ValueError): - net.get_layer('dense_unknown') - with self.assertRaises(ValueError): - net.get_layer() - with self.assertRaises(ValueError): - net.get_layer(index=4) - - # input, output - self.assertEqual(net.input, x) - self.assertEqual(net.output, y) - - # input_shape, output_shape - self.assertEqual(net.input_shape, (None, 32)) - self.assertEqual(net.output_shape, (None, 2)) - - # get_*_at - self.assertEqual(net.get_input_at(0), x) - self.assertEqual(net.get_output_at(0), y) - - # compute_output_shape - self.assertEqual(net.compute_output_shape((3, 32)).as_list(), [3, 2]) - - def testInvalidNetworks(self): - # redundant inputs - x = network_layers.Input(shape=(32,)) - y = core_layers.Dense(2)(x) - with self.assertRaises(ValueError): - network_layers.GraphNetwork([x, x], y) - - # inputs that don't come from Input - x = array_ops.placeholder(dtype='float32', shape=(None, 32)) - y = core_layers.Dense(2)(x) - with self.assertRaises(ValueError): - network_layers.GraphNetwork(x, y) - - # inputs that don't come from Input but have a layer history - x = network_layers.Input(shape=(32,)) - x = core_layers.Dense(32)(x) - y = core_layers.Dense(2)(x) - with self.assertRaises(ValueError): - network_layers.GraphNetwork(x, y) - - # outputs that don't come from layers - x = network_layers.Input(shape=(32,)) - y = core_layers.Dense(2)(x) - y = 2 * y - with self.assertRaises(ValueError): - network_layers.GraphNetwork(x, y) - - # disconnected graphs - x1 = network_layers.Input(shape=(32,)) - x2 = network_layers.Input(shape=(32,)) - y = core_layers.Dense(2)(x1) - with self.assertRaises(ValueError): - network_layers.GraphNetwork(x2, y) - - # redundant layer names - x = network_layers.Input(shape=(32,)) - z = core_layers.Dense(2, name='dense')(x) - y = core_layers.Dense(2, name='dense')(z) - with self.assertRaises(ValueError): - network_layers.GraphNetwork(x, y) - - def testInputTensorWrapping(self): - x = array_ops.placeholder(dtype='float32', shape=(None, 32)) - x = network_layers.Input(tensor=x) - y = core_layers.Dense(2)(x) - network_layers.GraphNetwork(x, y) - - def testExplicitBatchSize(self): - x = network_layers.Input(shape=(32,), batch_size=3) - y = core_layers.Dense(2)(x) - self.assertEqual(y.get_shape().as_list(), [3, 2]) - - def testNetworkRecursion(self): - # test the ability of networks to be used as layers inside networks. - a = network_layers.Input(shape=(32,)) - b = core_layers.Dense(2)(a) - net = network_layers.GraphNetwork(a, b) - - c = network_layers.Input(shape=(32,)) - d = net(c) - - recursive_net = network_layers.GraphNetwork(c, d) - self.assertEqual(len(recursive_net.layers), 2) - self.assertEqual(recursive_net.layers[1], net) - self.assertEqual(len(recursive_net.weights), 2) - - # test callability - x = array_ops.placeholder(dtype='float32', shape=(None, 32)) - y = recursive_net(x) - self.assertEqual(y.get_shape().as_list(), [None, 2]) - - def testSparseInput(self): - - class SparseSoftmax(base_layers.Layer): - - def call(self, inputs): - return sparse_ops.sparse_softmax(inputs) - - x = network_layers.Input(shape=(32,), sparse=True) - y = SparseSoftmax()(x) # pylint: disable=not-callable - network = network_layers.GraphNetwork(x, y) - - self.assertEqual(len(network.layers), 2) - self.assertEqual(network.layers[0].sparse, True) - - @test_util.run_in_graph_and_eager_modes() - def testMaskingSingleInput(self): - - class MaskedLayer(base_layers.Layer): - - def call(self, inputs, mask=None): - if mask is not None: - return inputs * mask - return inputs - - def compute_mask(self, inputs, mask=None): - return array_ops.ones_like(inputs) - - if context.in_graph_mode(): - x = network_layers.Input(shape=(32,)) - y = MaskedLayer()(x) # pylint: disable=not-callable - network = network_layers.GraphNetwork(x, y) - - # test callability on Input - x_2 = network_layers.Input(shape=(32,)) - y_2 = network(x_2) - self.assertEqual(y_2.get_shape().as_list(), [None, 32]) - - # test callability on regular tensor - x_2 = array_ops.placeholder(dtype='float32', shape=(None, 32)) - y_2 = network(x_2) - self.assertEqual(y_2.get_shape().as_list(), [None, 32]) - else: - a = constant_op.constant([2] * 32) - mask = constant_op.constant([0, 1] * 16) - a._keras_mask = mask - b = MaskedLayer().apply(a) - self.assertTrue(hasattr(b, '_keras_mask')) - self.assertAllEqual(self.evaluate(array_ops.ones_like(mask)), - self.evaluate(getattr(b, '_keras_mask'))) - self.assertAllEqual(self.evaluate(a * mask), self.evaluate(b)) - - -class DeferredModeTest(test.TestCase): - - def testDeferredTensorAttributes(self): - x = base_layers._DeferredTensor(shape=(None, 2), dtype='float32', name='x') - self.assertEqual(str(x), - 'DeferredTensor(\'x\', shape=(?, 2), dtype=float32)') - self.assertEqual(repr(x), - '<_DeferredTensor \'x\' shape=(?, 2) dtype=float32>') - - @test_util.run_in_graph_and_eager_modes() - def testSimpleNetworkBuilding(self): - inputs = network_layers.Input(shape=(32,)) - if context.in_eager_mode(): - self.assertIsInstance(inputs, base_layers._DeferredTensor) - self.assertEqual(inputs.dtype.name, 'float32') - self.assertEqual(inputs.shape.as_list(), [None, 32]) - - x = core_layers.Dense(2)(inputs) - if context.in_eager_mode(): - self.assertIsInstance(x, base_layers._DeferredTensor) - self.assertEqual(x.dtype.name, 'float32') - self.assertEqual(x.shape.as_list(), [None, 2]) - - outputs = core_layers.Dense(4)(x) - network = network_layers.GraphNetwork(inputs, outputs) - self.assertIsInstance(network, network_layers.GraphNetwork) - - if context.in_eager_mode(): - # It should be possible to call such a network on EagerTensors. - inputs = constant_op.constant( - np.random.random((10, 32)).astype('float32')) - outputs = network(inputs) - self.assertEqual(outputs.shape.as_list(), [10, 4]) - - @test_util.run_in_graph_and_eager_modes() - def testMultiIONetworkbuilding(self): - input_a = network_layers.Input(shape=(32,)) - input_b = network_layers.Input(shape=(16,)) - a = core_layers.Dense(16)(input_a) - - class AddLayer(base_layers.Layer): - - def call(self, inputs): - return inputs[0] + inputs[1] - - def compute_output_shape(self, input_shape): - return input_shape[0] - - c = AddLayer()([a, input_b]) # pylint: disable=not-callable - c = core_layers.Dense(2)(c) - - network = network_layers.GraphNetwork([input_a, input_b], [a, c]) - if context.in_eager_mode(): - a_val = constant_op.constant( - np.random.random((10, 32)).astype('float32')) - b_val = constant_op.constant( - np.random.random((10, 16)).astype('float32')) - outputs = network([a_val, b_val]) - self.assertEqual(len(outputs), 2) - self.assertEqual(outputs[0].shape.as_list(), [10, 16]) - self.assertEqual(outputs[1].shape.as_list(), [10, 2]) - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/python/layers/normalization.py b/tensorflow/python/layers/normalization.py index 890c12f6e00daabe7e64c00814fcb3ff8f04ae3a..3aafdbb14b56bbba3c4c9928e09a91d51f18db73 100644 --- a/tensorflow/python/layers/normalization.py +++ b/tensorflow/python/layers/normalization.py @@ -39,8 +39,10 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import state_ops from tensorflow.python.training import moving_averages +from tensorflow.python.util.tf_export import tf_export +@tf_export('layers.BatchNormalization') class BatchNormalization(base.Layer): """Batch Normalization layer from http://arxiv.org/abs/1502.03167. @@ -92,8 +94,8 @@ class BatchNormalization(base.Layer): and should be neither too small (which would add noise) nor too large (which would give stale estimates). Note that `momentum` is still applied to get the means and variances for inference. - 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. trainable: Boolean, if `True` also add variables to the graph collection `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable). virtual_batch_size: An `int`. By default, `virtual_batch_size` is `None`, @@ -336,8 +338,9 @@ class BatchNormalization(base.Layer): return var with ops.device(None): - device = ((lambda _: self.moving_mean.device) - if context.in_graph_mode() else self.moving_mean.device) + device = ( + self.moving_mean.device if context.executing_eagerly() else + (lambda _: self.moving_mean.device)) with ops.device(device): self.renorm_mean = _renorm_variable('renorm_mean', param_shape) self.renorm_mean_weight = _renorm_variable('renorm_mean_weight', ()) @@ -345,8 +348,9 @@ class BatchNormalization(base.Layer): # renorm_stddev_weight. This allows us to (1) mix the average # stddev with the minibatch stddev early in training, and (2) compute # the unbiased average stddev by dividing renorm_stddev by the weight. - device = ((lambda _: self.moving_variance.device) - if context.in_graph_mode() else self.moving_variance.device) + device = ( + self.moving_variance.device if context.executing_eagerly() else + (lambda _: self.moving_variance.device)) with ops.device(device): self.renorm_stddev = _renorm_variable('renorm_stddev', param_shape) self.renorm_stddev_weight = _renorm_variable( @@ -418,7 +422,7 @@ class BatchNormalization(base.Layer): one_minus_decay) variance_update = self._assign_moving_average(self.moving_variance, variance, one_minus_decay) - if context.in_graph_mode(): + if not context.executing_eagerly(): # Note that in Eager mode, the updates are already executed when running # assign_moving_averages. So we do not need to put them into # collections. @@ -491,6 +495,7 @@ class BatchNormalization(base.Layer): return (r, d, new_mean, new_variance) def call(self, inputs, training=False): + in_eager_mode = context.executing_eagerly() if self.virtual_batch_size is not None: # Virtual batches (aka ghost batches) can be simulated by reshaping the # Tensor and reusing the existing batch norm implementation @@ -593,6 +598,9 @@ class BatchNormalization(base.Layer): axis=1, keep_dims=True) def _do_update(var, value): + if in_eager_mode and not self.trainable: + return + return moving_averages.assign_moving_average( var, value, self.momentum, zero_debias=False) @@ -604,7 +612,7 @@ class BatchNormalization(base.Layer): training, lambda: _do_update(self.moving_variance, new_variance), lambda: self.moving_variance) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.add_update(mean_update, inputs=inputs) self.add_update(variance_update, inputs=inputs) @@ -629,6 +637,7 @@ class BatchNormalization(base.Layer): return input_shape +@tf_export('layers.batch_normalization') def batch_normalization(inputs, axis=-1, momentum=0.99, @@ -664,9 +673,16 @@ def batch_normalization(inputs, 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 - need to be added as a dependency to the `train_op`. For example: + need to be added as a dependency to the `train_op`. Also, be sure to add + any batch_normalization ops before getting the update_ops collection. + Otherwise, update_ops will be empty, and training/inference will not work + properly. For example: ```python + x_norm = tf.layers.batch_normalization(x, training=training) + + # ... + update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): train_op = optimizer.minimize(loss) @@ -722,8 +738,8 @@ def batch_normalization(inputs, and should be neither too small (which would add noise) nor too large (which would give stale estimates). Note that `momentum` is still applied to get the means and variances for inference. - 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. virtual_batch_size: An `int`. By default, `virtual_batch_size` is `None`, which means batch normalization is performed across the whole batch. When `virtual_batch_size` is not `None`, instead perform "Ghost Batch diff --git a/tensorflow/python/layers/pooling.py b/tensorflow/python/layers/pooling.py index ab06a3a40826e7d41c040066fd41c56c1ed84ad2..50503ce093fbc251b11c4d5cbccb2a2683d92e7a 100644 --- a/tensorflow/python/layers/pooling.py +++ b/tensorflow/python/layers/pooling.py @@ -26,6 +26,7 @@ from tensorflow.python.layers import base from tensorflow.python.layers import utils from tensorflow.python.ops import array_ops from tensorflow.python.ops import nn +from tensorflow.python.util.tf_export import tf_export class _Pooling1D(base.Layer): @@ -96,6 +97,7 @@ class _Pooling1D(base.Layer): return tensor_shape.TensorShape([input_shape[0], length, input_shape[2]]) +@tf_export('layers.AveragePooling1D') class AveragePooling1D(_Pooling1D): """Average Pooling layer for 1D inputs. @@ -127,6 +129,7 @@ class AveragePooling1D(_Pooling1D): **kwargs) +@tf_export('layers.average_pooling1d') def average_pooling1d(inputs, pool_size, strides, padding='valid', data_format='channels_last', name=None): @@ -161,6 +164,7 @@ def average_pooling1d(inputs, pool_size, strides, return layer.apply(inputs) +@tf_export('layers.MaxPooling1D') class MaxPooling1D(_Pooling1D): """Max Pooling layer for 1D inputs. @@ -192,6 +196,7 @@ class MaxPooling1D(_Pooling1D): **kwargs) +@tf_export('layers.max_pooling1d') def max_pooling1d(inputs, pool_size, strides, padding='valid', data_format='channels_last', name=None): @@ -297,6 +302,7 @@ class _Pooling2D(base.Layer): [input_shape[0], rows, cols, input_shape[3]]) +@tf_export('layers.AveragePooling2D') class AveragePooling2D(_Pooling2D): """Average pooling layer for 2D inputs (e.g. images). @@ -328,6 +334,7 @@ class AveragePooling2D(_Pooling2D): padding=padding, data_format=data_format, name=name, **kwargs) +@tf_export('layers.average_pooling2d') def average_pooling2d(inputs, pool_size, strides, padding='valid', data_format='channels_last', @@ -365,6 +372,7 @@ def average_pooling2d(inputs, return layer.apply(inputs) +@tf_export('layers.MaxPooling2D') class MaxPooling2D(_Pooling2D): """Max pooling layer for 2D inputs (e.g. images). @@ -396,6 +404,7 @@ class MaxPooling2D(_Pooling2D): padding=padding, data_format=data_format, name=name, **kwargs) +@tf_export('layers.max_pooling2d') def max_pooling2d(inputs, pool_size, strides, padding='valid', data_format='channels_last', @@ -515,6 +524,7 @@ class _Pooling3D(base.Layer): [input_shape[0], len_dim1, len_dim2, len_dim3, input_shape[4]]) +@tf_export('layers.AveragePooling3D') class AveragePooling3D(_Pooling3D): """Average pooling layer for 3D inputs (e.g. volumes). @@ -548,6 +558,7 @@ class AveragePooling3D(_Pooling3D): padding=padding, data_format=data_format, name=name, **kwargs) +@tf_export('layers.average_pooling3d') def average_pooling3d(inputs, pool_size, strides, padding='valid', data_format='channels_last', @@ -587,6 +598,7 @@ def average_pooling3d(inputs, return layer.apply(inputs) +@tf_export('layers.MaxPooling3D') class MaxPooling3D(_Pooling3D): """Max pooling layer for 3D inputs (e.g. volumes). @@ -620,6 +632,7 @@ class MaxPooling3D(_Pooling3D): padding=padding, data_format=data_format, name=name, **kwargs) +@tf_export('layers.max_pooling3d') def max_pooling3d(inputs, pool_size, strides, padding='valid', data_format='channels_last', diff --git a/tensorflow/python/layers/utils.py b/tensorflow/python/layers/utils.py index 7407d9a7b30268271810a23a7146c92fd28f8d4e..3b156c36a2ff35fb9e05af1406d7b3f6cf883394 100644 --- a/tensorflow/python/layers/utils.py +++ b/tensorflow/python/layers/utils.py @@ -20,9 +20,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.eager import context from tensorflow.python.ops import variables from tensorflow.python.ops import control_flow_ops from tensorflow.python.framework import ops +from tensorflow.python.framework import smart_cond as smart_module from tensorflow.python.framework import tensor_util from tensorflow.python.util import nest @@ -178,67 +180,56 @@ def deconv_output_length(input_length, filter_size, padding, stride): return input_length -def smart_cond(pred, fn1, fn2, name=None): - """Return either `fn1()` or `fn2()` based on the boolean predicate `pred`. +def smart_cond(pred, true_fn=None, false_fn=None, name=None): + """Return either `true_fn()` if predicate `pred` is true else `false_fn()`. - If `pred` is a bool or has a constant value, we return either `fn1()` - or `fn2()`, otherwise we use `tf.cond` to dynamically route to both. + If `pred` is a bool or has a constant value, we return either `true_fn()` + or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both. Arguments: - pred: A scalar determining whether to return the result of `fn1` or `fn2`. - fn1: The callable to be performed if pred is true. - fn2: The callable to be performed if pred is false. + pred: A scalar determining whether to return the result of `true_fn` or + `false_fn`. + true_fn: The callable to be performed if pred is true. + false_fn: The callable to be performed if pred is false. name: Optional name prefix when using `tf.cond`. Returns: - Tensors returned by the call to either `fn1` or `fn2`. + Tensors returned by the call to either `true_fn` or `false_fn`. Raises: - TypeError: If `fn1` or `fn2` is not callable. + TypeError: If `true_fn` or `false_fn` is not callable. """ - if not callable(fn1): - raise TypeError('`fn1` must be callable.') - if not callable(fn2): - raise TypeError('`fn2` must be callable.') - - pred_value = constant_value(pred) - if pred_value is not None: - if pred_value: - return fn1() - else: - return fn2() - else: - return control_flow_ops.cond(pred, true_fn=fn1, false_fn=fn2, name=name) + if isinstance(pred, variables.Variable): + return control_flow_ops.cond( + pred, true_fn=true_fn, false_fn=false_fn, name=name) + return smart_module.smart_cond( + pred, true_fn=true_fn, false_fn=false_fn, name=name) def constant_value(pred): """Return the bool value for `pred`, or None if `pred` had a dynamic value. - Arguments: - pred: A scalar, either a Python bool or a TensorFlow boolean variable - or tensor, or the Python integer 1 or 0. + Arguments: + pred: A scalar, either a Python bool or a TensorFlow boolean variable + or tensor, or the Python integer 1 or 0. - Returns: - True or False if `pred` has a constant boolean value, None otherwise. + Returns: + True or False if `pred` has a constant boolean value, None otherwise. - Raises: - TypeError: If `pred` is not a Variable, Tensor or bool. - """ + Raises: + TypeError: If `pred` is not a Variable, Tensor or bool, or Python + interger 1 or 0. + """ # Allow integer booleans. - if pred == 0: - pred = False - elif pred == 1: - pred = True - - if isinstance(pred, bool): - pred_value = pred - elif isinstance(pred, variables.Variable): - pred_value = None - elif isinstance(pred, ops.Tensor): - pred_value = tensor_util.constant_value(pred) - else: - raise TypeError('`pred` must be a Tensor, a Variable, or a Python bool.') - return pred_value + if isinstance(pred, int): + if pred == 1: + pred = True + elif pred == 0: + pred = False + + if isinstance(pred, variables.Variable): + return None + return smart_module.smart_constant_value(pred) def object_list_uid(object_list): @@ -255,3 +246,45 @@ def static_shape(x): return tuple(x.get_shape().as_list()) except ValueError: return None + + +def get_reachable_from_inputs(inputs, targets=None): + """Returns the set of tensors reachable from `inputs`. + + Stops if all targets have been found (target is optional). + + Only valid in Symbolic mode, not Eager mode. + + Args: + inputs: List of tensors. + targets: List of tensors. + + Returns: + A set of tensors reachable from the inputs (includes the inputs themselves). + """ + reachable = set(inputs) + if targets: + targets = set(targets) + queue = inputs[:] + + while queue: + x = queue.pop() + outputs = [] + try: + consumers = x.consumers() + except AttributeError: + # Case where x is a variable type + consumers = [x.op] + for z in consumers: + consumer_outputs = z.outputs + if consumer_outputs: # May be None + outputs += consumer_outputs + + for y in outputs: + if y not in reachable: + reachable.add(y) + queue.insert(0, y) + + if targets and targets.issubset(reachable): + return reachable + return reachable diff --git a/tensorflow/python/layers/utils_test.py b/tensorflow/python/layers/utils_test.py index a560f6b6d21efc0c1070d5a9296a7a8e914e2eb9..c941aad7bc63dbb891fbe78cd2a47dd6805bf231 100644 --- a/tensorflow/python/layers/utils_test.py +++ b/tensorflow/python/layers/utils_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function from tensorflow.python.layers import utils +from tensorflow.python.ops import array_ops from tensorflow.python.platform import test @@ -87,5 +88,34 @@ class ConvUtilsTest(test.TestCase): self.assertEqual(3, utils.deconv_output_length(4, 2, 'full', 1)) self.assertEqual(6, utils.deconv_output_length(4, 2, 'full', 2)) + +class GraphUtilsTest(test.TestCase): + + def testGetReachableFromInputs(self): + + with self.test_session(): + pl_1 = array_ops.placeholder(shape=None, dtype='float32') + pl_2 = array_ops.placeholder(shape=None, dtype='float32') + pl_3 = array_ops.placeholder(shape=None, dtype='float32') + x_1 = pl_1 + pl_2 + x_2 = pl_2 * 2 + x_3 = pl_3 + 1 + x_4 = x_1 + x_2 + x_5 = x_3 * pl_1 + + self.assertEqual( + utils.get_reachable_from_inputs([pl_1]), + {pl_1, x_1, x_4, x_5}) + self.assertEqual( + utils.get_reachable_from_inputs([pl_1, pl_2]), + {pl_1, pl_2, x_1, x_2, x_4, x_5}) + self.assertEqual( + utils.get_reachable_from_inputs([pl_3]), + {pl_3, x_3, x_5}) + self.assertEqual( + utils.get_reachable_from_inputs([x_3]), + {x_3, x_5}) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc index e0422ef80add42307268be2743e668eb8c8acb68..02eafd42b35231195a6405c8c3cc11871ed55772 100644 --- a/tensorflow/python/lib/core/py_func.cc +++ b/tensorflow/python/lib/core/py_func.cc @@ -79,10 +79,11 @@ Status MakeArgTuple(const PyCall* call, PyObject** tuple) { const Tensor& t = call->ins[i]; if (call->eager) { if (call->gpu) { - arg = EagerTensorFromHandle(new TFE_TensorHandle(t, call->device)); + arg = EagerTensorFromHandle( + new TFE_TensorHandle(t, call->device, call->device)); } else { // TFE_TensorHandle assumes that CPU is identified by `nullptr`. - arg = EagerTensorFromHandle(new TFE_TensorHandle(t, nullptr)); + arg = EagerTensorFromHandle(new TFE_TensorHandle(t, nullptr, nullptr)); } if (arg == nullptr) { return errors::Internal("Unable to procure EagerTensor from Tensor."); @@ -163,9 +164,9 @@ bool IsSingleNone(PyObject* obj) { } // Retrieves a Tensor from `eager_tensor` and stores it in `output_tensor`. -void ExtractTensorFromEagerTensor(const PyObject* eager_tensor, - Tensor* output_tensor) { - *output_tensor = EagerTensor_Handle(eager_tensor)->t; +tensorflow::Status ExtractTensorFromEagerTensor(const PyObject* eager_tensor, + const Tensor** output_tensor) { + return EagerTensor_Handle(eager_tensor)->Tensor(output_tensor); } // Calls the registered py function through the trampoline. @@ -219,7 +220,9 @@ Status DoCallPyFunc(PyCall* call, bool* out_log_on_error) { if (call->eager) { const PyObject* item = PyList_GetItem(result, i); if (EagerTensor_CheckExact(item)) { - ExtractTensorFromEagerTensor(item, &t); + const Tensor* tensor = nullptr; + s = ExtractTensorFromEagerTensor(item, &tensor); + if (s.ok()) t = *tensor; } else { s = errors::FailedPrecondition( "Expected EagerTensor, found PyObject of type: ", @@ -237,10 +240,10 @@ Status DoCallPyFunc(PyCall* call, bool* out_log_on_error) { } else if (EagerTensor_CheckExact(result) || result == Py_None) { // result is an `EagerTensor` or `None`. DCHECK(call->eager); - Tensor t; if (result != Py_None) { - ExtractTensorFromEagerTensor(result, &t); - call->out.push_back(t); + const Tensor* t = nullptr; + s = ExtractTensorFromEagerTensor(result, &t); + if (s.ok()) call->out.push_back(*t); } } else if (PyArray_Check(result)) { // `result` is a NumPy array. diff --git a/tensorflow/python/lib/core/py_seq_tensor.cc b/tensorflow/python/lib/core/py_seq_tensor.cc index 317bdc2e14747583f372808f48a5928273f5570a..8247d354db62532c10c5acc9875cc08289cd31bf 100644 --- a/tensorflow/python/lib/core/py_seq_tensor.cc +++ b/tensorflow/python/lib/core/py_seq_tensor.cc @@ -84,6 +84,7 @@ bool IsPyDimension(PyObject* obj) { } Status InferShapeAndType(PyObject* obj, TensorShape* shape, DataType* dtype) { + std::vector refs_to_clean; while (true) { // We test strings first, in case a string is considered a sequence. if (IsPyString(obj)) { @@ -93,6 +94,7 @@ Status InferShapeAndType(PyObject* obj, TensorShape* shape, DataType* dtype) { if (length > 0) { shape->AddDim(length); obj = PySequence_GetItem(obj, 0); + refs_to_clean.push_back(make_safe(obj)); continue; } else if (length == 0) { shape->AddDim(length); @@ -167,14 +169,15 @@ const char ErrorFoundFloat[] = if (shape.dims() > 1) { \ /* Iterate over outer dim, and recursively convert each element. */ \ const int64 s = shape.dim_size(0); \ - if (TF_PREDICT_FALSE(s != PySequence_Length(obj))) { \ + Safe_PyObjectPtr seq = make_safe(PySequence_Fast(obj, "")); \ + if (TF_PREDICT_FALSE(s != PySequence_Fast_GET_SIZE(seq.get()))) { \ return ErrorRectangular; \ } \ TensorShape rest = shape; \ rest.RemoveDim(0); \ for (int64 i = 0; i < s; ++i) { \ - const char* error = \ - FUNCTION##Helper(PySequence_GetItem(obj, i), rest, buf); \ + const char* error = FUNCTION##Helper( \ + PySequence_Fast_GET_ITEM(seq.get(), i), rest, buf); \ if (TF_PREDICT_FALSE(error != nullptr)) return error; \ } \ } else { \ diff --git a/tensorflow/python/lib/core/py_util.cc b/tensorflow/python/lib/core/py_util.cc index 2635694e23c07dd8e75d4bb0cfb9e83a2042d921..00cbf0c532cf80d3bb27afe168ecde963ba3591d 100644 --- a/tensorflow/python/lib/core/py_util.cc +++ b/tensorflow/python/lib/core/py_util.cc @@ -41,6 +41,55 @@ const char* ClassName(PyObject* py) { } // end namespace +// Returns a PyObject containing a string, or null +void TryAppendTraceback(PyObject* ptype, PyObject* pvalue, PyObject* ptraceback, + string* out) { + // The "traceback" module is assumed to be imported already by script_ops.py. + PyObject* tb_module = PyImport_AddModule("traceback"); + + if (!tb_module) { + return; + } + + PyObject* format_exception = + PyObject_GetAttrString(tb_module, "format_exception"); + + if (!format_exception) { + return; + } + + if (!PyCallable_Check(format_exception)) { + Py_DECREF(format_exception); + return; + } + + PyObject* ret_val = PyObject_CallFunctionObjArgs(format_exception, ptype, + pvalue, ptraceback, nullptr); + Py_DECREF(format_exception); + + if (!ret_val) { + return; + } + + if (!PyList_Check(ret_val)) { + Py_DECREF(ret_val); + return; + } + + Py_ssize_t n = PyList_GET_SIZE(ret_val); + for (Py_ssize_t i = 0; i < n; ++i) { + PyObject* v = PyList_GET_ITEM(ret_val, i); +#if PY_MAJOR_VERSION < 3 + strings::StrAppend(out, PyString_AS_STRING(v), "\n"); +#else + strings::StrAppend(out, PyUnicode_AsUTF8(v), "\n"); +#endif + } + + // Iterate through ret_val. + Py_DECREF(ret_val); +} + string PyExceptionFetch() { CHECK(PyErr_Occurred()) << "Must only call PyExceptionFetch after an exception."; @@ -52,14 +101,20 @@ string PyExceptionFetch() { string err = ClassName(ptype); if (pvalue) { PyObject* str = PyObject_Str(pvalue); + if (str) { #if PY_MAJOR_VERSION < 3 - strings::StrAppend(&err, ": ", PyString_AS_STRING(str)); + strings::StrAppend(&err, ": ", PyString_AS_STRING(str), "\n"); #else - strings::StrAppend(&err, ": ", PyUnicode_AsUTF8(str)); + strings::StrAppend(&err, ": ", PyUnicode_AsUTF8(str), "\n"); #endif Py_DECREF(str); + } else { + strings::StrAppend(&err, "(unknown error message)\n"); } + + TryAppendTraceback(ptype, pvalue, ptraceback, &err); + Py_DECREF(pvalue); } Py_DECREF(ptype); diff --git a/tensorflow/python/lib/io/file_io.i b/tensorflow/python/lib/io/file_io.i index c0c4e035fc3d6a50334acb9228c13c702ef426c0..891a7b0fd0dc177f5ee439707c9e2c99148e177c 100644 --- a/tensorflow/python/lib/io/file_io.i +++ b/tensorflow/python/lib/io/file_io.i @@ -110,21 +110,15 @@ void RecursivelyCreateDir(const string& dirname, TF_Status* out_status) { } } -void CopyFile(const string& oldpath, const string& newpath, bool overwrite, +void CopyFile(const string& src, const string& target, bool overwrite, TF_Status* out_status) { - // If overwrite is false and the newpath file exists then it's an error. - if (!overwrite && tensorflow::Env::Default()->FileExists(newpath).ok()) { + // If overwrite is false and the target file exists then its an error. + if (!overwrite && tensorflow::Env::Default()->FileExists(target).ok()) { TF_SetStatus(out_status, TF_ALREADY_EXISTS, "file already exists"); return; } - string file_content; - tensorflow::Status status = ReadFileToString(tensorflow::Env::Default(), - oldpath, &file_content); - if (!status.ok()) { - Set_TF_Status_from_Status(out_status, status); - return; - } - status = WriteStringToFile(tensorflow::Env::Default(), newpath, file_content); + tensorflow::Status status = + tensorflow::Env::Default()->CopyFile(src, target); if (!status.ok()) { Set_TF_Status_from_Status(out_status, status); } diff --git a/tensorflow/python/lib/io/file_io.py b/tensorflow/python/lib/io/file_io.py index 4e3071d8513a28b02b70b290c4987bec92b3c32e..59f5075f177ef5335115cb4f24182d28a9b547c8 100644 --- a/tensorflow/python/lib/io/file_io.py +++ b/tensorflow/python/lib/io/file_io.py @@ -31,6 +31,7 @@ from tensorflow.python.framework import c_api_util from tensorflow.python.framework import errors from tensorflow.python.util import compat from tensorflow.python.util import deprecation +from tensorflow.python.util.tf_export import tf_export class FileIO(object): @@ -235,6 +236,7 @@ class FileIO(object): self._writable_file = None +@tf_export("gfile.Exists") def file_exists(filename): """Determines whether a path exists or not. @@ -256,6 +258,7 @@ def file_exists(filename): return True +@tf_export("gfile.Remove") def delete_file(filename): """Deletes the file located at 'filename'. @@ -306,6 +309,7 @@ def write_string_to_file(filename, file_content): f.write(file_content) +@tf_export("gfile.Glob") def get_matching_files(filename): """Returns a list of files that match the given pattern(s). @@ -336,6 +340,7 @@ def get_matching_files(filename): ] +@tf_export("gfile.MkDir") def create_dir(dirname): """Creates a directory with the name 'dirname'. @@ -353,6 +358,7 @@ def create_dir(dirname): pywrap_tensorflow.CreateDir(compat.as_bytes(dirname), status) +@tf_export("gfile.MakeDirs") def recursive_create_dir(dirname): """Creates a directory and all parent/intermediate directories. @@ -368,6 +374,7 @@ def recursive_create_dir(dirname): pywrap_tensorflow.RecursivelyCreateDir(compat.as_bytes(dirname), status) +@tf_export("gfile.Copy") def copy(oldpath, newpath, overwrite=False): """Copies data from oldpath to newpath. @@ -385,6 +392,7 @@ def copy(oldpath, newpath, overwrite=False): compat.as_bytes(oldpath), compat.as_bytes(newpath), overwrite, status) +@tf_export("gfile.Rename") def rename(oldname, newname, overwrite=False): """Rename or move a file / directory. @@ -426,6 +434,7 @@ def atomic_write_string_to_file(filename, contents, overwrite=True): raise +@tf_export("gfile.DeleteRecursively") def delete_recursively(dirname): """Deletes everything under dirname recursively. @@ -439,6 +448,7 @@ def delete_recursively(dirname): pywrap_tensorflow.DeleteRecursively(compat.as_bytes(dirname), status) +@tf_export("gfile.IsDirectory") def is_directory(dirname): """Returns whether the path is a directory or not. @@ -452,6 +462,7 @@ def is_directory(dirname): return pywrap_tensorflow.IsDirectory(compat.as_bytes(dirname), status) +@tf_export("gfile.ListDirectory") def list_directory(dirname): """Returns a list of entries contained within a directory. @@ -479,6 +490,7 @@ def list_directory(dirname): ] +@tf_export("gfile.Walk") def walk(top, in_order=True): """Recursive directory tree generator for directories. @@ -522,6 +534,7 @@ def walk(top, in_order=True): yield here +@tf_export("gfile.Stat") def stat(filename): """Returns file statistics for a given path. diff --git a/tensorflow/python/lib/io/file_io_test.py b/tensorflow/python/lib/io/file_io_test.py index a751607aaa1f47ca7c08674eca2b27ee0cafa3d2..223858edfa84eaa1c7879a9774dcc836de4f4672 100644 --- a/tensorflow/python/lib/io/file_io_test.py +++ b/tensorflow/python/lib/io/file_io_test.py @@ -485,6 +485,11 @@ class FileIoTest(test.TestCase): f.flush() self.assertEqual(content, f.read(len(content) + 1)) + def testUTF8StringPathExists(self): + file_path = os.path.join(self._base_dir, "UTF8测试_file_exist") + file_io.write_string_to_file(file_path, "testing") + v = file_io.file_exists(file_path) + self.assertEqual(v, True) if __name__ == "__main__": test.main() diff --git a/tensorflow/python/lib/io/tf_record.py b/tensorflow/python/lib/io/tf_record.py index df190100689bd864de78f5a2cf52b1ade081a789..6fcf9c91d831e3a89552b522040e8e8647114a2f 100644 --- a/tensorflow/python/lib/io/tf_record.py +++ b/tensorflow/python/lib/io/tf_record.py @@ -22,8 +22,10 @@ from __future__ import print_function from tensorflow.python import pywrap_tensorflow from tensorflow.python.framework import errors from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export +@tf_export("python_io.TFRecordCompressionType") class TFRecordCompressionType(object): """The type of compression for the record.""" NONE = 0 @@ -33,6 +35,7 @@ class TFRecordCompressionType(object): # NOTE(vrv): This will eventually be converted into a proto. to match # the interface used by the C++ RecordWriter. +@tf_export("python_io.TFRecordOptions") class TFRecordOptions(object): """Options used for manipulating TFRecord files.""" compression_type_map = { @@ -51,6 +54,7 @@ class TFRecordOptions(object): return cls.compression_type_map[options.compression_type] +@tf_export("python_io.tf_record_iterator") def tf_record_iterator(path, options=None): """An iterator that read the records from a TFRecords file. @@ -71,16 +75,19 @@ def tf_record_iterator(path, options=None): if reader is None: raise IOError("Could not open %s." % path) - while True: - try: - with errors.raise_exception_on_not_ok_status() as status: - reader.GetNext(status) - except errors.OutOfRangeError: - break - yield reader.record() - reader.Close() - - + try: + while True: + try: + with errors.raise_exception_on_not_ok_status() as status: + reader.GetNext(status) + except errors.OutOfRangeError: + break + yield reader.record() + finally: + reader.Close() + + +@tf_export("python_io.TFRecordWriter") class TFRecordWriter(object): """A class to write records to a TFRecords file. diff --git a/tensorflow/python/ops/accumulate_n_benchmark.py b/tensorflow/python/ops/accumulate_n_benchmark.py index c58d36f39705ecf0f24214ce4ba4574e70a93e77..a709066cae4da2811b3e98d2e93bf44ec12dcee6 100644 --- a/tensorflow/python/ops/accumulate_n_benchmark.py +++ b/tensorflow/python/ops/accumulate_n_benchmark.py @@ -39,7 +39,7 @@ from tensorflow.python.platform import test class AccumulateNBenchmark(test.Benchmark): def _AccumulateNTemplate(self, inputs, init, shape, validate_shape): - var = gen_state_ops._temporary_variable( + var = gen_state_ops.temporary_variable( shape=shape, dtype=inputs[0].dtype.base_dtype) ref = state_ops.assign(var, init, validate_shape=validate_shape) update_ops = [ @@ -47,8 +47,7 @@ class AccumulateNBenchmark(test.Benchmark): ref, tensor, use_locking=True).op for tensor in inputs ] with ops.control_dependencies(update_ops): - return gen_state_ops._destroy_temporary_variable( - ref, var_name=var.op.name) + return gen_state_ops.destroy_temporary_variable(ref, var_name=var.op.name) def _AccumulateNInitializedWithFirst(self, inputs): return self._AccumulateNTemplate( @@ -60,7 +59,7 @@ class AccumulateNBenchmark(test.Benchmark): def _AccumulateNInitializedWithMerge(self, inputs): return self._AccumulateNTemplate( inputs, - init=array_ops.zeros_like(gen_control_flow_ops._merge(inputs)[0]), + init=array_ops.zeros_like(gen_control_flow_ops.merge(inputs)[0]), shape=tensor_shape.vector(0), validate_shape=False) diff --git a/tensorflow/python/ops/array_grad.py b/tensorflow/python/ops/array_grad.py index c9292184e6a9910db6b41022ab01312ce44e0a96..3c6a5c9e562ff9765c2ef47555871c94cd6feb1e 100644 --- a/tensorflow/python/ops/array_grad.py +++ b/tensorflow/python/ops/array_grad.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import gen_array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import sparse_ops @@ -79,7 +80,7 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index): def _ExtractInputShapes(inputs): """Extract the shapes of a set of input tensors.""" - if not context.in_graph_mode(): + if context.executing_eagerly(): return array_ops.shape_n(inputs) sizes = [] fully_known = True @@ -105,7 +106,7 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index): out_grads = [] if isinstance(grad, ops.Tensor): - if context.in_eager_mode(): + if context.executing_eagerly(): # Using mod here for convenience since concat_dim is already verified # in concat implementation to be within the allowed [-rank, rank) range. non_neg_concat_dim = ( @@ -115,6 +116,19 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index): non_neg_concat_dim) out_grads = array_ops.split(grad, sizes, non_neg_concat_dim) else: + if constant_op.is_constant(concat_dim): + # If concat_dim is a constant defined in a different context, + # then we duplicate it in the current context to avoid passing it + # through an Enter node. + # This is a small optimization in general, but it is required when + # compiling with XLA, as XLA needs the concat input to be folded into a + # constant. + grad_context = control_flow_util.GetOutputContext(grad.op) + dim_context = control_flow_util.GetOutputContext(concat_dim.op) + if dim_context != grad_context: + value = tensor_util.constant_value(concat_dim) + concat_dim = constant_op.constant(value=value, dtype=concat_dim.dtype) + # Using mod here for convenience since concat_dim is already verified # in concat implementation to be within the allowed [-rank, rank) range. non_neg_concat_dim = concat_dim % array_ops.rank(input_values[0]) @@ -125,7 +139,6 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index): # on CPUs and a Maxwell TitanX. A speedup was seen in a large majority of # cases when switching implementations at N=16, but it is possible that # there will be a small number of performance regressions. - # pylint: disable=protected-access if len(sizes) > 16: # extract the size of each input along the concat dimension sizes = array_ops.squeeze( @@ -134,10 +147,9 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index): [1, -1])) out_grads = array_ops.split(grad, sizes, non_neg_concat_dim) else: - offset = gen_array_ops._concat_offset(non_neg_concat_dim, sizes) + offset = gen_array_ops.concat_offset(non_neg_concat_dim, sizes) for (begin, size) in zip(offset, sizes): out_grads.append(array_ops.slice(grad, begin, size)) - # pylint: enable=protected-access elif isinstance(grad, ops.IndexedSlices): # Using mod here for convenience since concat_dim is already verified # in concat implementation to be within the allowed [-rank, rank) range. @@ -416,7 +428,7 @@ def _GatherV2Grad(op, grad): # For axis 0 gathers, build an appropriately shaped IndexedSlices. if axis_static == 0: - if context.in_eager_mode(): + if context.executing_eagerly(): params_tail_shape = params_shape.cpu()[1:] else: params_tail_shape = params_shape[1:] @@ -566,7 +578,7 @@ def _TileGrad(op, grad): axes = math_ops.range(0, array_ops.size(split_shape), 2) input_grad = math_ops.reduce_sum(array_ops.reshape(grad, split_shape), axes) # Fix shape inference - if context.in_graph_mode(): + if not context.executing_eagerly(): input_grad.set_shape(op.inputs[0].get_shape()) return [input_grad, None] @@ -613,9 +625,7 @@ def _ReverseSequenceGrad(op, grad): @ops.RegisterGradient("Reverse") def _ReverseGrad(op, grad): reverse_dims = op.inputs[1] - # pylint: disable=protected-access - return gen_array_ops._reverse(grad, reverse_dims), None - # pylint: enable=protected-access + return gen_array_ops.reverse(grad, reverse_dims), None @ops.RegisterGradient("ReverseV2") @@ -686,17 +696,13 @@ ops.NotDifferentiable("OneHot") @ops.RegisterGradient("MirrorPad") def _MirrorPadGrad(op, grad): mode = op.get_attr("mode") - # pylint: disable=protected-access - return [gen_array_ops._mirror_pad_grad(grad, op.inputs[1], mode=mode), None] - # pylint: enable=protected-access + return [gen_array_ops.mirror_pad_grad(grad, op.inputs[1], mode=mode), None] @ops.RegisterGradient("MirrorPadGrad") def _MirrorPadGradGrad(op, grad): mode = op.get_attr("mode") - # pylint: disable=protected-access - return [gen_array_ops._mirror_pad(grad, op.inputs[1], mode=mode), None] - # pylint: enable=protected-access + return [gen_array_ops.mirror_pad(grad, op.inputs[1], mode=mode), None] @ops.RegisterGradient("QuantizeAndDequantize") diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 9541b097a94466861a83cb48ed3111563490cfba..ec7c14f7d8697e61d2acb25a82c0ac9b2fcf28f4 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +# Tests for this file live in python/kernel_tests/array_ops_test.py """Support for manipulating tensors. See the @{$python/array_ops} guide. @@ -127,15 +128,18 @@ def identity(input, name=None): # pylint: disable=redefined-builtin Returns: A `Tensor`. Has the same type as `input`. """ - if context.in_graph_mode(): - return gen_array_ops.identity(input, name=name) - else: + if context.executing_eagerly(): input = ops.convert_to_tensor(input) in_device = input.device # TODO(ashankar): Does 'identity' need to invoke execution callbacks? - if context.context().device_name != in_device: + context_device = context.context().device_name + if not context_device: + context_device = "/job:localhost/replica:0/task:0/device:CPU:0" + if context_device != in_device: return input._copy() # pylint: disable=protected-access return input + else: + return gen_array_ops.identity(input, name=name) # pylint: disable=redefined-builtin,protected-access @@ -194,7 +198,7 @@ def expand_dims(input, axis=None, name=None, dim=None): if axis is not None: raise ValueError("can't specify both 'dim' and 'axis'") axis = dim - return gen_array_ops._expand_dims(input, axis, name) + return gen_array_ops.expand_dims(input, axis, name) # pylint: enable=redefined-builtin,protected-access @@ -207,28 +211,25 @@ def expand_dims(input, axis=None, name=None, dim=None): "This op will be removed after the deprecation date. " "Please switch to tf.setdiff1d().") def listdiff(x, y, out_idx=None, name=None): - return gen_array_ops._list_diff(x, y, out_idx, name) + return gen_array_ops.list_diff(x, y, out_idx, name) -listdiff.__doc__ = gen_array_ops._list_diff.__doc__ + "\n" + listdiff.__doc__ +listdiff.__doc__ = gen_array_ops.list_diff.__doc__ + "\n" + listdiff.__doc__ # pylint: enable=protected-access -# pylint: disable=undefined-variable,protected-access +# pylint: disable=undefined-variable @tf_export("setdiff1d") def setdiff1d(x, y, index_dtype=dtypes.int32, name=None): - return gen_array_ops._list_diff(x, y, index_dtype, name) - + return gen_array_ops.list_diff(x, y, index_dtype, name) -setdiff1d.__doc__ = gen_array_ops._list_diff.__doc__ -# pylint: enable=protected-access +setdiff1d.__doc__ = gen_array_ops.list_diff.__doc__ @tf_export("broadcast_dynamic_shape") def broadcast_dynamic_shape(shape_x, shape_y): - # pylint: disable=protected-access """Returns the broadcasted dynamic shape between `shape_x` and `shape_y`. Args: @@ -238,8 +239,7 @@ def broadcast_dynamic_shape(shape_x, shape_y): Returns: A rank 1 integer `Tensor` representing the broadcasted shape. """ - return gen_array_ops._broadcast_args(shape_x, shape_y) - # pylint: enable=protected-access + return gen_array_ops.broadcast_args(shape_x, shape_y) @tf_export("broadcast_static_shape") @@ -305,7 +305,7 @@ def shape_internal(input, name=None, optimize=True, out_type=dtypes.int32): sparse_tensor.SparseTensorValue)): return gen_math_ops.cast(input.dense_shape, out_type) else: - if context.in_graph_mode(): + if not context.executing_eagerly(): input_tensor = ops.convert_to_tensor(input) input_shape = input_tensor.get_shape() if optimize and input_shape.is_fully_defined(): @@ -330,7 +330,7 @@ def shape_n(input, out_type=dtypes.int32, name=None): """ output = gen_array_ops.shape_n(input, out_type=out_type, name=name) - if context.in_graph_mode(): + if not context.executing_eagerly(): for i, input_tensor in enumerate(input): input_tensor = ops.convert_to_tensor(input_tensor) input_shape = input_tensor.get_shape() @@ -385,16 +385,22 @@ def size_internal(input, name=None, optimize=True, out_type=dtypes.int32): Returns: A `Tensor` of type `out_type`. Defaults to `tf.int32`. """ + if context.executing_eagerly() and not isinstance( + input, (sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)): + return np.prod(ops.convert_to_tensor(input)._shape_tuple()) # pylint: disable=protected-access with ops.name_scope(name, "Size", [input]) as name: if isinstance(input, (sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)): - return gen_math_ops._prod( + return gen_math_ops.prod( gen_math_ops.cast(input.dense_shape, out_type), 0, name=name) else: input_tensor = ops.convert_to_tensor(input) input_shape = input_tensor.get_shape() - if optimize and input_shape.is_fully_defined(): - return constant(input_shape.num_elements(), out_type, name=name) + if optimize: + if input_shape.is_fully_defined(): + return constant(input_shape.num_elements(), out_type, name=name) + if input_shape.dims and any(dim == 0 for dim in input_shape.dims): + return constant(0, out_type, name=name) return gen_array_ops.size(input, name=name, out_type=out_type) @@ -604,7 +610,7 @@ def slice(input_, begin, size, name=None): Note that @{tf.Tensor.__getitem__} is typically a more pythonic way to perform slices, as it allows you to write `foo[3:7, :-2]` instead of - `tf.slice([3, 0], [4, foo.get_shape()[1]-2])`. + `tf.slice(foo, [3, 0], [4, foo.get_shape()[1]-2])`. `begin` is zero-based; `size` is one-based. If `size[i]` is -1, all remaining elements in dimension i are included in the @@ -776,7 +782,7 @@ def strided_slice(input_, new_axis_mask=new_axis_mask, shrink_axis_mask=shrink_axis_mask) - if context.in_graph_mode(): + if not context.executing_eagerly(): # TODO(apassos) In eager mode assignment will be done by overriding # __setitem__ instead. op.assign = assign @@ -787,8 +793,8 @@ def _SliceHelperVar(var, slice_spec): """Creates a slice helper object given a variable. This allows creating a sub-tensor from part of the current contents - of a variable. See ${tf.Tensor$`Tensor.__getitem__`} - for detailed examples of slicing. + of a variable. See @{tf.Tensor.__getitem__} for detailed examples + of slicing. This function in addition also allows assignment to a sliced range. This is similar to `__setitem__` functionality in Python. However, @@ -878,7 +884,7 @@ def parallel_stack(values, name="parallel_stack"): output_shape = tensor_shape.TensorShape([len(values)]) output_shape = output_shape.concatenate(value_shape) # expand_dims converts concat to stack. - return gen_array_ops._parallel_concat( + return gen_array_ops.parallel_concat( [expand_dims(value, 0) for value in values], shape=output_shape) @@ -936,7 +942,7 @@ def stack(values, axis=0, name="stack"): raise ValueError("axis = %d not in [%d, %d)" % (axis, -expanded_num_dims, expanded_num_dims)) - return gen_array_ops._pack(values, axis=axis, name=name) + return gen_array_ops.pack(values, axis=axis, name=name) # pylint: disable=invalid-name @@ -980,7 +986,7 @@ def _autopacking_helper(list_or_tuple, dtype, name): # convertible-to-tensor types, such as numpy arrays. elems_as_tensors.append( constant_op.constant(elem, dtype=dtype, name=str(i))) - return gen_array_ops._pack(elems_as_tensors, name=scope) + return gen_array_ops.pack(elems_as_tensors, name=scope) else: return converted_elems @@ -1075,7 +1081,7 @@ def unstack(value, num=None, axis=0, name="unstack"): num = value_shape[axis].value if num is None: raise ValueError("Cannot infer num from shape %s" % value_shape) - return gen_array_ops._unpack(value, num=num, axis=axis, name=name) + return gen_array_ops.unpack(value, num=num, axis=axis, name=name) @tf_export("concat") @@ -1172,7 +1178,7 @@ def concat(values, axis, name="concat"): dtype=dtypes.int32).get_shape().assert_is_compatible_with( tensor_shape.scalar()) return identity(values[0], name=scope) - return gen_array_ops._concat_v2(values=values, axis=axis, name=name) + return gen_array_ops.concat_v2(values=values, axis=axis, name=name) @tf_export("boolean_mask") @@ -1240,8 +1246,7 @@ def boolean_mask(tensor, mask, name="boolean_mask", axis=None): axis = 0 if axis is None else axis shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask) - leading_size = gen_math_ops._prod( - shape(tensor)[axis:axis + ndims_mask], [0]) + leading_size = gen_math_ops.prod(shape(tensor)[axis:axis + ndims_mask], [0]) tensor = reshape(tensor, concat([ shape(tensor)[:axis], [leading_size], @@ -1305,10 +1310,22 @@ def unique(x, out_idx=dtypes.int32, name=None): # period (3 weeks) pass. # TODO(yongtang): The documentation should also # be updated when switch to v2. - return gen_array_ops._unique(x, out_idx, name) + return gen_array_ops.unique(x, out_idx, name) -unique.__doc__ = gen_array_ops._unique.__doc__ +unique.__doc__ = gen_array_ops.unique.__doc__ + + +@tf_export("unique_with_counts") +def unique_with_counts(x, out_idx=dtypes.int32, name=None): + # TODO(yongtang): switch to v2 once API deprecation + # period (3 weeks) pass. + # TODO(yongtang): The documentation should also + # be updated when switch to v2. + return gen_array_ops.unique_with_counts(x, out_idx, name) + + +unique_with_counts.__doc__ = gen_array_ops.unique_with_counts.__doc__ @tf_export("split") @@ -1362,20 +1379,18 @@ def split(value, num_or_size_splits, axis=0, num=None, name="split"): """ size_splits = ops.convert_to_tensor(num_or_size_splits) if size_splits._rank() == 0 and size_splits.dtype.is_integer: - return gen_array_ops._split( + return gen_array_ops.split( axis=axis, num_split=num_or_size_splits, value=value, name=name) if num is None: - num = size_splits._shape_tuple()[0] + size_splits_shape = size_splits._shape_tuple() + if size_splits_shape: + num = size_splits_shape[0] if num is None: raise ValueError("Cannot infer num from shape %s" % num_or_size_splits) - return gen_array_ops._split_v( - value=value, - size_splits=size_splits, - axis=axis, - num_split=num, - name=name) + return gen_array_ops.split_v( + value=value, size_splits=size_splits, axis=axis, num_split=num, name=name) @tf_export("transpose") @@ -1389,6 +1404,14 @@ def transpose(a, perm=None, name="transpose", conjugate=False): `a.dtype` is either `complex64` or `complex128` then the values of `a` are conjugated and transposed. + @compatibility(numpy) + In `numpy` transposes are memory-efficient constant time operations as they + simply return a new view of the same data with adjusted `strides`. + + TensorFlow does not support strides, so `transpose` returns a new tensor with + the items permuted. + @end_compatibility + For example: ```python @@ -1437,7 +1460,7 @@ def transpose(a, perm=None, name="transpose", conjugate=False): """ with ops.name_scope(name, "transpose", [a]) as name: transpose_fn = ( - gen_array_ops._conjugate_transpose + gen_array_ops.conjugate_transpose if (conjugate and a.dtype.is_complex) else gen_array_ops.transpose) if perm is None: rank = gen_array_ops.rank(a) @@ -1445,7 +1468,7 @@ def transpose(a, perm=None, name="transpose", conjugate=False): ret = transpose_fn(a, perm, name=name) # NOTE(mrry): Setting the shape explicitly because # reverse is not handled by the shape function. - if context.in_graph_mode(): + if not context.executing_eagerly(): input_shape = ret.op.inputs[0].get_shape().dims if input_shape is not None: ret.set_shape(input_shape[::-1]) @@ -1489,6 +1512,14 @@ def matrix_transpose(a, name="matrix_transpose", conjugate=False): tf.matmul(matrix, tf.matrix_transpose(b)) ``` + @compatibility(numpy) + In `numpy` transposes are memory-efficient constant time operations as they + simply return a new view of the same data with adjusted `strides`. + + TensorFlow does not support strides, `matrix_transposes` return a new tensor + with the items permuted. + @end_compatibility + Args: a: A `Tensor` with `rank >= 2`. name: A name for the operation (optional). @@ -1602,12 +1633,12 @@ def zeros_like(tensor, dtype=None, name=None, optimize=True): with ops.name_scope(name, "zeros_like", [tensor]) as name: tensor = ops.convert_to_tensor(tensor, name="tensor") - if context.in_eager_mode(): + if context.executing_eagerly(): if dtype is not None and dtype != tensor.dtype: return zeros( shape_internal(tensor, optimize=optimize), dtype=dtype, name=name) with ops.device(tensor.device): - return gen_array_ops._zeros_like(tensor, name=name) + return gen_array_ops.zeros_like(tensor, name=name) # For now, variant types must be created via zeros_like; as we need to # pass the input variant object to the proper zeros callback. @@ -1622,7 +1653,7 @@ def zeros_like(tensor, dtype=None, name=None, optimize=True): return zeros( shape_internal(tensor, optimize=optimize), dtype=dtype, name=name) else: - return gen_array_ops._zeros_like(tensor, name=name) + return gen_array_ops.zeros_like(tensor, name=name) @tf_export("ones_like") @@ -1658,7 +1689,7 @@ def ones_like(tensor, dtype=None, name=None, optimize=True): if dtype is None: dtype = tensor.dtype ret = ones(ones_shape, dtype=dtype, name=name) - if context.in_graph_mode(): + if not context.executing_eagerly(): ret.set_shape(tensor.get_shape()) return ret @@ -1739,11 +1770,11 @@ def placeholder(dtype, shape=None, name=None): Raises: RuntimeError: if eager execution is enabled """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("tf.placeholder() is not compatible with " "eager execution.") - return gen_array_ops._placeholder(dtype=dtype, shape=shape, name=name) + return gen_array_ops.placeholder(dtype=dtype, shape=shape, name=name) # pylint: disable=redefined-outer-name @@ -1802,7 +1833,7 @@ def sparse_placeholder(dtype, shape=None, name=None): Raises: RuntimeError: if eager execution is enabled """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("tf.placeholder() is not compatible with " "eager execution.") @@ -1887,21 +1918,21 @@ def pad(tensor, paddings, mode="CONSTANT", name=None, constant_values=0): # pyl # TODO(rjryan): Once the forward compatibility period (3 weeks) have passed # remove the "Pad" fallback here. if constant_values != 0: - result = gen_array_ops._pad_v2( + result = gen_array_ops.pad_v2( tensor, paddings, constant_values, name=name) else: - result = gen_array_ops._pad(tensor, paddings, name=name) + result = gen_array_ops.pad(tensor, paddings, name=name) elif mode == "REFLECT": - result = gen_array_ops._mirror_pad( + result = gen_array_ops.mirror_pad( tensor, paddings, mode="REFLECT", name=name) elif mode == "SYMMETRIC": - result = gen_array_ops._mirror_pad( + result = gen_array_ops.mirror_pad( tensor, paddings, mode="SYMMETRIC", name=name) else: raise ValueError("Unknown padding mode: %s" % mode) # Restore shape information where possible. - if context.in_graph_mode(): + if not context.executing_eagerly(): paddings_constant = tensor_util.constant_value( result.op.inputs[1], partial=True) input_shape = result.op.inputs[0].shape @@ -2125,7 +2156,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): sparse_tensor.SparseTensorValue)): raise TypeError("Truth must be a SparseTensor.") - return gen_array_ops._edit_distance( + return gen_array_ops.edit_distance( hypothesis.indices, hypothesis.values, hypothesis.dense_shape, @@ -2262,7 +2293,7 @@ def space_to_batch(input, paddings, block_size, name=None): # pylint: disable=r return result -space_to_batch.__doc__ = gen_array_ops._space_to_batch.__doc__ +space_to_batch.__doc__ = gen_array_ops.space_to_batch.__doc__ @tf_export("space_to_depth") @@ -2292,7 +2323,7 @@ def batch_to_space(input, crops, block_size, name=None): # pylint: disable=rede return result -batch_to_space.__doc__ = gen_array_ops._batch_to_space.__doc__ +batch_to_space.__doc__ = gen_array_ops.batch_to_space.__doc__ @tf_export("one_hot") @@ -2436,8 +2467,8 @@ def one_hot(indices, raise TypeError("dtype {0} of on_value does not match " "dtype {1} of off_value".format(on_dtype, off_dtype)) - return gen_array_ops._one_hot(indices, depth, on_value, off_value, axis, - name) + return gen_array_ops.one_hot(indices, depth, on_value, off_value, axis, + name) def _all_dimensions(x): @@ -2451,8 +2482,8 @@ def _all_dimensions(x): r = x.dense_shape.get_shape()[0].value # sparse.dense_shape is 1-D. return constant_op.constant(np.arange(r), dtype=dtypes.int32) - # Otherwise, we rely on Range and Rank to do the right thing at run-time. - return range(0, rank(x)) + # Otherwise, we rely on `range` and `rank` to do the right thing at runtime. + return gen_math_ops._range(0, rank(x), 1) @tf_export("sequence_mask") @@ -2497,7 +2528,7 @@ def sequence_mask(lengths, maxlen=None, dtype=dtypes.bool, name=None): maxlen = gen_math_ops._max(lengths, _all_dimensions(lengths)) else: maxlen = ops.convert_to_tensor(maxlen) - if maxlen.get_shape().ndims != 0: + if maxlen.get_shape().ndims is not None and maxlen.get_shape().ndims != 0: raise ValueError("maxlen must be scalar for sequence_mask") # The basic idea is to compare a range row vector of size maxlen: @@ -2565,7 +2596,7 @@ def squeeze(input, axis=None, name=None, squeeze_dims=None): axis = squeeze_dims if np.isscalar(axis): axis = [axis] - return gen_array_ops._squeeze(input, axis, name) + return gen_array_ops.squeeze(input, axis, name) @tf_export("where") @@ -2616,7 +2647,7 @@ def where(condition, x=None, y=None, name=None): condition, preferred_dtype=dtypes.bool, name="condition") return gen_array_ops.where(condition=condition, name=name) elif x is not None and y is not None: - return gen_math_ops._select(condition=condition, x=x, y=y, name=name) + return gen_math_ops.select(condition=condition, x=x, y=y, name=name) else: raise ValueError("x and y must both be non-None or both be None.") diff --git a/tensorflow/python/ops/batch_norm_benchmark.py b/tensorflow/python/ops/batch_norm_benchmark.py index c2ee2b383231333239c6e2d4e874a0ad1cdf493e..5d68b47aeaef3a90973387ecd5b265eef1e96a5f 100644 --- a/tensorflow/python/ops/batch_norm_benchmark.py +++ b/tensorflow/python/ops/batch_norm_benchmark.py @@ -41,9 +41,8 @@ def batch_norm_op(tensor, mean, variance, beta, gamma, scale): # _batch_norm_with_global_normalization is deprecated in v9 ops.get_default_graph().graph_def_versions.producer = 8 # pylint: disable=protected-access - return gen_nn_ops._batch_norm_with_global_normalization(tensor, mean, - variance, beta, gamma, - 0.001, scale) + return gen_nn_ops._batch_norm_with_global_normalization( + tensor, mean, variance, beta, gamma, 0.001, scale) # pylint: enable=protected-access diff --git a/tensorflow/python/ops/bitwise_ops_test.py b/tensorflow/python/ops/bitwise_ops_test.py index f9b025b787e4f49e1dcde6c589f66c59d779fcef..c4cfc0da197edcfd143cfee79fd3c3f9b7a2858b 100644 --- a/tensorflow/python/ops/bitwise_ops_test.py +++ b/tensorflow/python/ops/bitwise_ops_test.py @@ -71,7 +71,7 @@ class BitwiseOpTest(test_util.TensorFlowTestCase): def testInvertOp(self): dtype_list = [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, - dtypes.uint8, dtypes.uint16] + dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64] inputs = [0, 5, 3, 14] with self.test_session(use_gpu=True) as sess: for dtype in dtype_list: diff --git a/tensorflow/python/ops/candidate_sampling_ops.py b/tensorflow/python/ops/candidate_sampling_ops.py index 20445c78a290a4fe67cad668dd714dd2c61c5f3d..9ea1ea9c92c9b016a3f9126c89ee4dc1e73c9f27 100644 --- a/tensorflow/python/ops/candidate_sampling_ops.py +++ b/tensorflow/python/ops/candidate_sampling_ops.py @@ -20,9 +20,9 @@ from __future__ import division from __future__ import print_function from tensorflow.python.framework import random_seed -from tensorflow.python.ops import array_ops +from tensorflow.python.ops import array_ops # pylint: disable=unused-import from tensorflow.python.ops import gen_candidate_sampling_ops -from tensorflow.python.ops import math_ops +from tensorflow.python.ops import math_ops # pylint: disable=unused-import from tensorflow.python.util.tf_export import tf_export @@ -77,7 +77,7 @@ def uniform_candidate_sampler(true_classes, num_true, num_sampled, unique, of each of `sampled_candidates`. """ seed1, seed2 = random_seed.get_seed(seed) - return gen_candidate_sampling_ops._uniform_candidate_sampler( + return gen_candidate_sampling_ops.uniform_candidate_sampler( true_classes, num_true, num_sampled, unique, range_max, seed=seed1, seed2=seed2, name=name) @@ -136,7 +136,7 @@ def log_uniform_candidate_sampler(true_classes, num_true, num_sampled, unique, of each of `sampled_candidates`. """ seed1, seed2 = random_seed.get_seed(seed) - return gen_candidate_sampling_ops._log_uniform_candidate_sampler( + return gen_candidate_sampling_ops.log_uniform_candidate_sampler( true_classes, num_true, num_sampled, unique, range_max, seed=seed1, seed2=seed2, name=name) @@ -193,7 +193,7 @@ def learned_unigram_candidate_sampler(true_classes, num_true, num_sampled, """ seed1, seed2 = random_seed.get_seed(seed) - return gen_candidate_sampling_ops._learned_unigram_candidate_sampler( + return gen_candidate_sampling_ops.learned_unigram_candidate_sampler( true_classes, num_true, num_sampled, unique, range_max, seed=seed1, seed2=seed2, name=name) @@ -283,7 +283,7 @@ def fixed_unigram_candidate_sampler(true_classes, """ seed1, seed2 = random_seed.get_seed(seed) - return gen_candidate_sampling_ops._fixed_unigram_candidate_sampler( + return gen_candidate_sampling_ops.fixed_unigram_candidate_sampler( true_classes, num_true, num_sampled, unique, range_max, vocab_file=vocab_file, distortion=distortion, num_reserved_ids=num_reserved_ids, num_shards=num_shards, shard=shard, @@ -321,7 +321,7 @@ def all_candidate_sampler(true_classes, num_true, num_sampled, unique, of each of `sampled_candidates`. All returned values are 1.0. """ seed1, seed2 = random_seed.get_seed(seed) - return gen_candidate_sampling_ops._all_candidate_sampler( + return gen_candidate_sampling_ops.all_candidate_sampler( true_classes, num_true, num_sampled, unique, seed=seed1, seed2=seed2, name=name) @@ -370,6 +370,6 @@ def compute_accidental_hits(true_classes, sampled_candidates, num_true, """ seed1, seed2 = random_seed.get_seed(seed) - return gen_candidate_sampling_ops._compute_accidental_hits( + return gen_candidate_sampling_ops.compute_accidental_hits( true_classes, sampled_candidates, num_true, seed=seed1, seed2=seed2, name=name) diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py index 0fd6e29a49c8e4e31e244bfbbfca525d72e4d811..d6d75e4ef9874e2ed4706d48b7bdeed62baef414 100644 --- a/tensorflow/python/ops/check_ops.py +++ b/tensorflow/python/ops/check_ops.py @@ -169,7 +169,7 @@ def assert_negative(x, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_negative', [x, data]): x = ops.convert_to_tensor(x, name='x') if data is None: - if context.in_eager_mode(): + if context.executing_eagerly(): name = _shape_and_dtype_str(x) else: name = x.name @@ -210,7 +210,7 @@ def assert_positive(x, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_positive', [x, data]): x = ops.convert_to_tensor(x, name='x') if data is None: - if context.in_eager_mode(): + if context.executing_eagerly(): name = _shape_and_dtype_str(x) else: name = x.name @@ -251,7 +251,7 @@ def assert_non_negative(x, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_non_negative', [x, data]): x = ops.convert_to_tensor(x, name='x') if data is None: - if context.in_eager_mode(): + if context.executing_eagerly(): name = _shape_and_dtype_str(x) else: name = x.name @@ -293,7 +293,7 @@ def assert_non_positive(x, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_non_positive', [x, data]): x = ops.convert_to_tensor(x, name='x') if data is None: - if context.in_eager_mode(): + if context.executing_eagerly(): name = _shape_and_dtype_str(x) else: name = x.name @@ -334,16 +334,16 @@ def assert_equal(x, y, data=None, summarize=None, message=None, name=None): @compatibility{eager} returns None Raises: - InvalidArgumentError if the check can be performed immediately and - `x == y` is False. The check can be performed immediately during - eager execution or if `x` and `y` are statically known. + InvalidArgumentError: if the check can be performed immediately and + `x == y` is False. The check can be performed immediately during eager + execution or if `x` and `y` are statically known. """ message = message or '' with ops.name_scope(name, 'assert_equal', [x, y, data]): x = ops.convert_to_tensor(x, name='x') y = ops.convert_to_tensor(y, name='y') - if context.in_eager_mode(): + if context.executing_eagerly(): eq = math_ops.equal(x, y) condition = math_ops.reduce_all(eq) if not condition: @@ -435,7 +435,7 @@ def assert_none_equal( with ops.name_scope(name, 'assert_none_equal', [x, y, data]): x = ops.convert_to_tensor(x, name='x') y = ops.convert_to_tensor(y, name='y') - if context.in_eager_mode(): + if context.executing_eagerly(): x_name = _shape_and_dtype_str(x) y_name = _shape_and_dtype_str(y) else: @@ -512,7 +512,7 @@ def assert_near( rtol = ops.convert_to_tensor(rtol, name='rtol', dtype=x.dtype) atol = ops.convert_to_tensor(atol, name='atol', dtype=x.dtype) - if context.in_eager_mode(): + if context.executing_eagerly(): x_name = _shape_and_dtype_str(x) y_name = _shape_and_dtype_str(y) else: @@ -562,7 +562,7 @@ def assert_less(x, y, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_less', [x, y, data]): x = ops.convert_to_tensor(x, name='x') y = ops.convert_to_tensor(y, name='y') - if context.in_eager_mode(): + if context.executing_eagerly(): x_name = _shape_and_dtype_str(x) y_name = _shape_and_dtype_str(y) else: @@ -610,7 +610,7 @@ def assert_less_equal(x, y, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_less_equal', [x, y, data]): x = ops.convert_to_tensor(x, name='x') y = ops.convert_to_tensor(y, name='y') - if context.in_eager_mode(): + if context.executing_eagerly(): x_name = _shape_and_dtype_str(x) y_name = _shape_and_dtype_str(y) else: @@ -658,7 +658,7 @@ def assert_greater(x, y, data=None, summarize=None, message=None, name=None): with ops.name_scope(name, 'assert_greater', [x, y, data]): x = ops.convert_to_tensor(x, name='x') y = ops.convert_to_tensor(y, name='y') - if context.in_eager_mode(): + if context.executing_eagerly(): x_name = _shape_and_dtype_str(x) y_name = _shape_and_dtype_str(y) else: @@ -708,7 +708,7 @@ def assert_greater_equal(x, y, data=None, summarize=None, message=None, with ops.name_scope(name, 'assert_greater_equal', [x, y, data]): x = ops.convert_to_tensor(x, name='x') y = ops.convert_to_tensor(y, name='y') - if context.in_eager_mode(): + if context.executing_eagerly(): x_name = _shape_and_dtype_str(x) y_name = _shape_and_dtype_str(y) else: @@ -808,7 +808,7 @@ def assert_rank(x, rank, data=None, summarize=None, message=None, name=None): static_condition = lambda actual_rank, given_rank: actual_rank == given_rank dynamic_condition = math_ops.equal - if context.in_eager_mode(): + if context.executing_eagerly(): name = '' else: name = x.name @@ -873,7 +873,7 @@ def assert_rank_at_least( static_condition = lambda actual_rank, given_rank: actual_rank >= given_rank dynamic_condition = math_ops.greater_equal - if context.in_eager_mode(): + if context.executing_eagerly(): name = '' else: name = x.name @@ -1001,7 +1001,7 @@ def assert_rank_in( ranks = tuple([ops.convert_to_tensor(rank, name='rank') for rank in ranks]) message = message or '' - if context.in_eager_mode(): + if context.executing_eagerly(): name = '' else: name = x.name @@ -1054,7 +1054,7 @@ def assert_integer(x, message=None, name=None): with ops.name_scope(name, 'assert_integer', [x]): x = ops.convert_to_tensor(x, name='x') if not x.dtype.is_integer: - if context.in_eager_mode(): + if context.executing_eagerly(): name = 'tensor' else: name = x.name @@ -1087,12 +1087,11 @@ def assert_type(tensor, tf_type, message=None, name=None): with ops.name_scope(name, 'assert_type', [tensor]): tensor = ops.convert_to_tensor(tensor, name='tensor') if tensor.dtype != tf_type: - if context.in_graph_mode(): - raise TypeError( - '%s %s must be of type %s' % (message, tensor.name, tf_type)) + if context.executing_eagerly(): + raise TypeError('%s tensor must be of type %s' % (message, tf_type)) else: - raise TypeError( - '%s tensor must be of type %s' % (message, tf_type)) + raise TypeError('%s %s must be of type %s' % (message, tensor.name, + tf_type)) return control_flow_ops.no_op('statically_determined_correct_type') @@ -1240,7 +1239,7 @@ def assert_scalar(tensor, name=None): tensor = ops.convert_to_tensor(tensor, name=name_scope) shape = tensor.get_shape() if shape.ndims != 0: - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError('Expected scalar shape, saw shape: %s.' % (shape,)) else: diff --git a/tensorflow/python/ops/clip_ops.py b/tensorflow/python/ops/clip_ops.py index dd8c33247c2436413ee8c9a3ceeca4d8a493bb4e..49f8c665313562cb20dbe4494103ded16646c741 100644 --- a/tensorflow/python/ops/clip_ops.py +++ b/tensorflow/python/ops/clip_ops.py @@ -110,7 +110,7 @@ def clip_by_norm(t, clip_norm, axes=None, name=None): t = ops.convert_to_tensor(t, name="t") # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm - l2norm = math_ops.sqrt(math_ops.reduce_sum(t * t, axes, keep_dims=True)) + l2norm = math_ops.sqrt(math_ops.reduce_sum(t * t, axes, keepdims=True)) intermediate = t * clip_norm # Assert that the shape is compatible with the initial shape, # to prevent unintentional broadcasting. diff --git a/tensorflow/python/ops/confusion_matrix.py b/tensorflow/python/ops/confusion_matrix.py index 50690cd891f73df1e345817b834ce6c361bff9e8..b9a93c3bedfff1f398e3b42cedf02a2f0a3ddd5c 100644 --- a/tensorflow/python/ops/confusion_matrix.py +++ b/tensorflow/python/ops/confusion_matrix.py @@ -99,19 +99,16 @@ def confusion_matrix(labels, predictions, num_classes=None, dtype=dtypes.int32, name=None, weights=None): """Computes the confusion matrix from predictions and labels. - Calculate the Confusion Matrix for a pair of prediction and - label 1-D int arrays. - The matrix columns represent the prediction labels and the rows represent the real labels. The confusion matrix is always a 2-D array of shape `[n, n]`, where `n` is the number of valid labels for a given classification task. Both prediction and labels must be 1-D arrays of the same shape in order for this function to work. - If `num_classes` is None, then `num_classes` will be set to the one plus - the maximum value in either predictions or labels. - Class labels are expected to start at 0. E.g., if `num_classes` was - three, then the possible labels would be `[0, 1, 2]`. + If `num_classes` is `None`, then `num_classes` will be set to one plus the + maximum value in either predictions or labels. Class labels are expected to + start at 0. For example, if `num_classes` is 3, then the possible labels + would be `[0, 1, 2]`. If `weights` is not `None`, then each prediction contributes its corresponding weight to the total value of the confusion matrix cell. @@ -119,7 +116,7 @@ def confusion_matrix(labels, predictions, num_classes=None, dtype=dtypes.int32, For example: ```python - tf.contrib.metrics.confusion_matrix([1, 2, 4], [2, 2, 4]) ==> + tf.confusion_matrix([1, 2, 4], [2, 2, 4]) ==> [[0 0 0 0 0] [0 0 1 0 0] [0 0 1 0 0] @@ -141,8 +138,9 @@ def confusion_matrix(labels, predictions, num_classes=None, dtype=dtypes.int32, weights: An optional `Tensor` whose shape matches `predictions`. Returns: - A k X k matrix representing the confusion matrix, where k is the number of - possible labels in the classification task. + A `Tensor` of type `dtype` with shape `[n, n]` representing the confusion + matrix, where `n` is the number of possible labels in the classification + task. Raises: ValueError: If both predictions and labels are not 1-D vectors and have @@ -188,7 +186,7 @@ def confusion_matrix(labels, predictions, num_classes=None, dtype=dtypes.int32, weights = math_ops.cast(weights, dtype) shape = array_ops.stack([num_classes, num_classes]) - indices = array_ops.transpose(array_ops.stack([labels, predictions])) + indices = array_ops.stack([labels, predictions], axis=1) values = (array_ops.ones_like(predictions, dtype) if weights is None else weights) cm_sparse = sparse_tensor.SparseTensor( diff --git a/tensorflow/python/ops/control_flow_grad.py b/tensorflow/python/ops/control_flow_grad.py index 97b57177b29986a006df992f4c0c2b79e11467aa..21354b5ae8ff1724bbb2539aff370b3df6da2598 100644 --- a/tensorflow/python/ops/control_flow_grad.py +++ b/tensorflow/python/ops/control_flow_grad.py @@ -28,7 +28,6 @@ from tensorflow.python.ops import math_ops # go/tf-wildcard-import # pylint: disable=wildcard-import,undefined-variable from tensorflow.python.ops.control_flow_ops import * -from tensorflow.python.ops.gen_control_flow_ops import * # pylint: enable=wildcard-import diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py index 49191c647d59691a59aa5d7dd9cc9dac285b9fea..1278768d8bdc9f039f19cf032f8ee09442ea34a9 100644 --- a/tensorflow/python/ops/control_flow_ops.py +++ b/tensorflow/python/ops/control_flow_ops.py @@ -44,18 +44,20 @@ See the @{$python/control_flow_ops} guide. @@add_check_numerics_ops @@Assert @@Print +@@timestamp """ # pylint: disable=g-bad-name from __future__ import absolute_import from __future__ import division from __future__ import print_function +import abc import collections import functools import six -from six.moves import xrange # pylint: disable=redefined-builtin +from tensorflow.core.framework import attr_value_pb2 from tensorflow.core.protobuf import control_flow_pb2 from tensorflow.python.eager import context from tensorflow.python.framework import constant_op @@ -78,6 +80,7 @@ from tensorflow.python.ops import tensor_array_ops from tensorflow.python.ops.gen_control_flow_ops import * # pylint: enable=wildcard-import from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util import compat from tensorflow.python.util import deprecation from tensorflow.python.util import nest from tensorflow.python.util import tf_should_use @@ -149,7 +152,7 @@ def Assert(condition, data, summarize=None, name=None): @compatibility{eager} `tf.errors.InvalidArgumentError` if `condition` is not true """ - if context.in_eager_mode(): + if context.executing_eagerly(): if not condition: xs = ops.convert_n_to_tensor(data) data_str = [_summarize_eager(x, summarize) for x in xs] @@ -175,7 +178,7 @@ def Assert(condition, data, summarize=None, name=None): condition, data, summarize, name="Assert") guarded_assert = cond(condition, no_op, true_assert, name="AssertGuard") - if context.in_eager_mode(): + if context.executing_eagerly(): return return guarded_assert.op @@ -193,7 +196,7 @@ def _Identity(data, name=None): data = ops.internal_convert_to_tensor_or_indexed_slices(data, as_ref=True) if isinstance(data, ops.Tensor): if data.dtype._is_ref_dtype: # pylint: disable=protected-access - return gen_array_ops._ref_identity(data, name=name) + return gen_array_ops.ref_identity(data, name=name) else: return array_ops.identity(data, name=name) else: @@ -261,10 +264,10 @@ def _Enter(data, data = ops.internal_convert_to_tensor_or_indexed_slices(data, as_ref=True) if isinstance(data, ops.Tensor): if data.dtype._is_ref_dtype and use_ref: # pylint: disable=protected-access - result = ref_enter( + result = gen_control_flow_ops.ref_enter( data, frame_name, is_constant, parallel_iterations, name=name) else: - result = enter( + result = gen_control_flow_ops.enter( data, frame_name, is_constant, parallel_iterations, name=name) if use_input_shape: result.set_shape(data.get_shape()) @@ -279,7 +282,7 @@ def _Enter(data, parallel_iterations=parallel_iterations, use_input_shape=use_input_shape, name=name) - indices = enter( + indices = gen_control_flow_ops.enter( data.indices, frame_name, is_constant, @@ -290,7 +293,7 @@ def _Enter(data, if isinstance(data, ops.IndexedSlices): dense_shape = data.dense_shape if dense_shape is not None: - dense_shape = enter( + dense_shape = gen_control_flow_ops.enter( dense_shape, frame_name, is_constant, @@ -300,7 +303,7 @@ def _Enter(data, dense_shape.set_shape(data.dense_shape.get_shape()) return ops.IndexedSlices(values, indices, dense_shape) else: - dense_shape = enter( + dense_shape = gen_control_flow_ops.enter( data.dense_shape, frame_name, is_constant, @@ -311,7 +314,7 @@ def _Enter(data, return sparse_tensor.SparseTensor(indices, values, dense_shape) -def exit(data, name=None): +def exit(data, name=None): # pylint: disable=redefined-builtin """Exits the current frame to its parent frame. Exit makes its input `data` available to the parent frame. @@ -326,7 +329,7 @@ def exit(data, name=None): data = ops.internal_convert_to_tensor_or_indexed_slices(data, as_ref=True) if isinstance(data, ops.Tensor): if data.dtype._is_ref_dtype: # pylint: disable=protected-access - return gen_control_flow_ops._ref_exit(data, name) + return gen_control_flow_ops.ref_exit(data, name) else: return gen_control_flow_ops._exit(data, name) else: @@ -368,17 +371,17 @@ def switch(data, pred, dtype=None, name=None): data, dtype=dtype, name="data", as_ref=True) pred = ops.convert_to_tensor(pred, name="pred") if isinstance(data, ops.Tensor): - return gen_control_flow_ops._switch(data, pred, name=name) + return gen_control_flow_ops.switch(data, pred, name=name) else: if not isinstance(data, (ops.IndexedSlices, sparse_tensor.SparseTensor)): raise TypeError("Type %s not supported" % type(data)) val, ind = data.values, data.indices - val_f, val_t = gen_control_flow_ops._switch(val, pred, name=name) - ind_f, ind_t = gen_control_flow_ops._switch(ind, pred, name="indices") + val_f, val_t = gen_control_flow_ops.switch(val, pred, name=name) + ind_f, ind_t = gen_control_flow_ops.switch(ind, pred, name="indices") if isinstance(data, ops.IndexedSlices): dense_shape = data.dense_shape if dense_shape is not None: - dense_shape_f, dense_shape_t = gen_control_flow_ops._switch( + dense_shape_f, dense_shape_t = gen_control_flow_ops.switch( dense_shape, pred, name="dense_shape") else: dense_shape_f, dense_shape_t = None, None @@ -386,7 +389,7 @@ def switch(data, pred, dtype=None, name=None): ops.IndexedSlices(val_t, ind_t, dense_shape_t)) else: dense_shape = data.dense_shape - dense_shape_f, dense_shape_t = gen_control_flow_ops._switch( + dense_shape_f, dense_shape_t = gen_control_flow_ops.switch( data.dense_shape, pred, name="dense_shape") return (sparse_tensor.SparseTensor(ind_f, val_f, dense_shape_f), sparse_tensor.SparseTensor(ind_t, val_t, dense_shape_t)) @@ -470,15 +473,15 @@ def merge(inputs, name=None): ] if all([isinstance(v, ops.Tensor) for v in inputs]): if all([v.dtype._is_ref_dtype for v in inputs]): # pylint: disable=protected-access - return gen_control_flow_ops._ref_merge(inputs, name) + return gen_control_flow_ops.ref_merge(inputs, name) else: - return gen_control_flow_ops._merge(inputs, name) + return gen_control_flow_ops.merge(inputs, name) elif all([isinstance(v, sparse_tensor.SparseTensor) for v in inputs]): # Only handle the case when all inputs are SparseTensor. values, _ = merge([inp.values for inp in inputs], name=name) - indices, chosen_index = gen_control_flow_ops._merge( + indices, chosen_index = gen_control_flow_ops.merge( [inp.indices for inp in inputs], name="indices") - dense_shape, _ = gen_control_flow_ops._merge( + dense_shape, _ = gen_control_flow_ops.merge( [inp.dense_shape for inp in inputs], name="dense_shape") return (sparse_tensor.SparseTensor(indices, values, dense_shape), chosen_index) @@ -486,13 +489,13 @@ def merge(inputs, name=None): # For now convert all the inputs as IndexedSlices. inputs = math_ops._as_indexed_slices_list(inputs, optimize=False) values, _ = merge([inp.values for inp in inputs], name=name) - indices, chosen_index = gen_control_flow_ops._merge( + indices, chosen_index = gen_control_flow_ops.merge( [inp.indices for inp in inputs], name="indices") if any(inp.dense_shape is not None for inp in inputs): if any(inp.dense_shape is None for inp in inputs): raise ValueError("Either all merged IndexedSlices must have a " "dense_shape, or none must have a dense_shape.") - dense_shape, _ = gen_control_flow_ops._merge( + dense_shape, _ = gen_control_flow_ops.merge( [inp.dense_shape for inp in inputs], name="dense_shape") else: dense_shape = None @@ -1012,10 +1015,8 @@ class GradLoopState(object): else: max_size = GetMaxSizeFromNestedMaximumIterations( value, self.forward_context) - # pylint: disable=protected-access - acc = gen_data_flow_ops._stack_v2( + acc = gen_data_flow_ops.stack_v2( max_size=max_size, elem_type=value.dtype.base_dtype, name="f_acc") - # pylint: enable=protected-access if curr_ctxt: curr_ctxt.Exit() @@ -1028,10 +1029,8 @@ class GradLoopState(object): if value_ctxt == self.forward_context: # value is not nested in the forward context. self.forward_context.Enter() - # pylint: disable=protected-access - push = gen_data_flow_ops._stack_push_v2( + push = gen_data_flow_ops.stack_push_v2( enter_acc, value, swap_memory=swap_enabled) - # pylint: enable=protected-access self.forward_context.Exit() # Protect stack push and order it before forward_index. self.forward_index.op._add_control_input(push.op) @@ -1043,18 +1042,14 @@ class GradLoopState(object): # The special case for creating a zero tensor for a dead # branch of a switch. See ControlFlowState.ZerosLike(). value_ctxt.outer_context.Enter() - # pylint: disable=protected-access - push = gen_data_flow_ops._stack_push_v2( + push = gen_data_flow_ops.stack_push_v2( enter_acc, value, swap_memory=swap_enabled) - # pylint: enable=protected-access value_ctxt.outer_context.Exit() push.op._set_control_flow_context(value_ctxt) else: value_ctxt.Enter() - # pylint: disable=protected-access - push = gen_data_flow_ops._stack_push_v2( + push = gen_data_flow_ops.stack_push_v2( enter_acc, value, swap_memory=swap_enabled) - # pylint: enable=protected-access value_ctxt.Exit() # Protect stack push and order it before forward_sync. self.forward_sync._add_control_input(push.op) @@ -1101,10 +1096,8 @@ class GradLoopState(object): pred = cond_ctxt.pred branch = (1 - cond_ctxt.branch) if dead_branch else cond_ctxt.branch history_value = _SwitchRefOrTensor(history_value, pred)[branch] - # pylint: disable=protected-access - pop = gen_data_flow_ops._stack_pop_v2(history_value, - value.dtype.base_dtype) - # pylint: enable=protected-access + pop = gen_data_flow_ops.stack_pop_v2(history_value, + value.dtype.base_dtype) pop.set_shape(value.get_shape()) self.grad_context.Exit() parallel_iterations = self.grad_context.parallel_iterations @@ -1474,7 +1467,10 @@ def ZerosLikeOutsideLoop(op, index): branch = op_ctxt.branch switch_val = switch(op.inputs[0], pred)[1 - branch] zeros_shape = array_ops.shape_internal(switch_val, optimize=False) - return array_ops.zeros(zeros_shape, dtype=val.dtype) + # Ensure ops created within array_ops.zeros are dominated by switch in + # cond context. + with ops.control_dependencies([switch_val]): + return array_ops.zeros(zeros_shape, dtype=val.dtype) else: return array_ops.zeros_like(val, optimize=False) @@ -1498,14 +1494,19 @@ class ControlFlowContext(object): """ def __init__(self, values_def=None, import_scope=None): + self._nested_contexts = [] self._outer_context = ops.get_default_graph()._get_control_flow_context() + if self._outer_context: + self._outer_context._nested_contexts.append(self) # pylint: disable=protected-access self._context_stack = [] if values_def: self._init_values_from_proto(values_def, import_scope=import_scope) else: - # Values that have been already seen in this context. + # The names of tensors that have been already seen in this context. self._values = set() - # Values referenced by but external to this context. + # The keys are the names of tensors referenced by but external to this + # context. Each value is the Tensor that should be used by this context to + # access the key value (e.g. a switch output guarding a cond input value). self._external_values = {} def _init_values_from_proto(self, values_def, import_scope=None): @@ -1551,7 +1552,17 @@ class ControlFlowContext(object): def back_prop(self): raise NotImplementedError("Abstract method") - def _to_proto(self, export_scope=None): + @abc.abstractmethod + def to_control_flow_context_def(self, context_def, export_scope=None): + """Serializes this into `context_def`. + + Args: + context_def: a `ControlFlowContextDef` protocol buffer. + export_scope: Optional `string`. Name scope to remove. + """ + raise NotImplementedError("Abstract method") + + def _to_values_def(self, export_scope=None): """Converts the values to a `ValuesDef` protocol buffer. Args: @@ -1568,11 +1579,6 @@ class ControlFlowContext(object): values_def.external_values[k] = ops.strip_name_scope(v.name, export_scope) return values_def - @staticmethod - def _from_proto(values_def, import_scope=None): - """Returns a `ControlFlowContext` created from `values_def`.""" - return ControlFlowContext(values_def=values_def, import_scope=import_scope) - def AddName(self, name): self._values.add(name) @@ -1622,10 +1628,13 @@ class ControlFlowContext(object): ctxt = util.GetOutputContext(x) if ctxt is not None and ctxt.GetWhileContext() == while_ctxt: internal_control_inputs.append(x) + external_control_inputs = [] if len(internal_control_inputs) != len(op.control_inputs): + external_control_inputs = list(set(op.control_inputs) + - set(internal_control_inputs)) op._remove_all_control_inputs() op._add_control_inputs(internal_control_inputs) - return internal_control_inputs + return internal_control_inputs, external_control_inputs # pylint: enable=protected-access @@ -1684,9 +1693,12 @@ class CondContext(ControlFlowContext): self._pivot = pivot # The predicate tensor in this branch self._branch = branch # 0 or 1 representing this branch - # Values considered to have been already seen in this context. + # Values considered to have been already seen in this context. They are + # not included in this context. self._values.add(pred.name) + self._external_values[pred.name] = pred self._values.add(pivot.name) + self._external_values[pivot.name] = pivot def _init_from_proto(self, context_def, import_scope=None): """Creates a new `CondContext` from protocol buffer. @@ -1704,8 +1716,8 @@ class CondContext(ControlFlowContext): self._pivot = g.as_graph_element( ops.prepend_name_scope(context_def.pivot_name, import_scope)) self._branch = context_def.branch - super(CondContext, self).__init__( - values_def=context_def.values_def, import_scope=import_scope) + super(CondContext, self).__init__(values_def=context_def.values_def, + import_scope=import_scope) @property def pred(self): @@ -1751,8 +1763,11 @@ class CondContext(ControlFlowContext): context_def.pivot_name = ops.strip_name_scope(self._pivot.name, export_scope) context_def.branch = self._branch - context_def.values_def.MergeFrom( - super(CondContext, self)._to_proto(export_scope)) + context_def.values_def.MergeFrom(super(CondContext, self)._to_values_def( + export_scope)) + for nested in self._nested_contexts: + nested_def = context_def.nested_contexts.add() + nested.to_control_flow_context_def(nested_def) return context_def else: @@ -1761,7 +1776,17 @@ class CondContext(ControlFlowContext): @staticmethod def from_proto(context_def, import_scope=None): """Returns a `CondContext` object created from `context_def`.""" - return CondContext(context_def=context_def, import_scope=import_scope) + ret = CondContext(context_def=context_def, + import_scope=import_scope) + + ret.Enter() + for nested_def in context_def.nested_contexts: + from_control_flow_context_def(nested_def, import_scope=import_scope) + ret.Exit() + return ret + + def to_control_flow_context_def(self, context_def, export_scope=None): + context_def.cond_ctxt.CopyFrom(self.to_proto(export_scope=export_scope)) def AddValue(self, val): """Add `val` to the current context and its outer context recursively.""" @@ -1776,6 +1801,7 @@ class CondContext(ControlFlowContext): if self._outer_context: result = self._outer_context.AddValue(val) self._values.add(result.name) + self._external_values[result.name] = result with ops.control_dependencies(None): result = _SwitchRefOrTensor(result, self._pred)[self._branch] if self._outer_context: @@ -1801,8 +1827,6 @@ class CondContext(ControlFlowContext): # pylint: disable=protected-access op._add_control_input(self._pivot.op) # pylint: enable=protected-access - for x in op.outputs: - self._values.add(x.name) else: for index in range(len(op.inputs)): x = op.inputs[index] @@ -1813,13 +1837,20 @@ class CondContext(ControlFlowContext): # pylint: enable=protected-access # Remove any external control dependency on this op. self._RemoveExternalControlEdges(op) - for x in op.outputs: - self._values.add(x.name) # pylint: disable=protected-access if op.graph._is_function(op.type) or op.type == "SymbolicGradient": op._add_control_input(self._pivot.op) # pylint: enable=protected-access + # Mark op's outputs as seen by this context and any outer contexts. + output_names = [x.name for x in op.outputs] + ctxt = self + while ctxt is not None: + # pylint: disable=protected-access + ctxt._values.update(output_names) + ctxt = ctxt._outer_context + # pylint: enable=protected-access + if self._outer_context or not util.IsLoopExit(op): op.graph.prevent_fetching(op) @@ -1835,6 +1866,7 @@ class CondContext(ControlFlowContext): if self._outer_context: real_val = self._outer_context.AddValue(val) self._values.add(real_val.name) + self._external_values[real_val.name] = real_val real_val = _SwitchRefOrTensor(real_val, self._pred)[self._branch] self._external_values[val.name] = real_val else: @@ -1996,7 +2028,7 @@ def cond(pred, raise TypeError("false_fn must be callable.") with ops.name_scope(name, "cond", [pred]): - if context.in_eager_mode(): + if context.executing_eagerly(): if pred: return _UnpackIfSingleton(true_fn()) return _UnpackIfSingleton(false_fn()) @@ -2067,9 +2099,12 @@ def cond(pred, merges = [merge(pair)[0] for pair in zip(res_f_flat, res_t_flat)] merges = _convert_flows_to_tensorarrays(nest.flatten(orig_res_t), merges) - # Add to collections - ops.add_to_collection(ops.GraphKeys.COND_CONTEXT, context_t) - ops.add_to_collection(ops.GraphKeys.COND_CONTEXT, context_f) + # Only add non-nested conds to the collection. Any nested control flow will + # be encapsulated in the root context. + assert context_t.outer_context == context_f.outer_context + if context_t.outer_context is None: + ops.add_to_collection(ops.GraphKeys.COND_CONTEXT, context_t) + ops.add_to_collection(ops.GraphKeys.COND_CONTEXT, context_f) merges = nest.pack_sequence_as(structure=orig_res_t, flat_sequence=merges) @@ -2206,6 +2241,17 @@ class WhileContext(ControlFlowContext): super(WhileContext, self).__init__( values_def=context_def.values_def, import_scope=import_scope) + # import_scope causes self.name to be different from the original serialized + # context's name. Rewrite "frame_name" attrs with the new name. + if import_scope: + for tensor_name in self._values: + op = g.as_graph_element(tensor_name).op + if util.IsLoopEnter(op): + # pylint: disable=protected-access + op._set_attr("frame_name", + attr_value_pb2.AttrValue(s=compat.as_bytes(self.name))) + # pylint: enable=protected-access + @property def maximum_iterations(self): """The maximum number of iterations that will be executed.""" @@ -2277,12 +2323,19 @@ class WhileContext(ControlFlowContext): ops.strip_name_scope(l.name, export_scope) for l in self._loop_enters ]) context_def.values_def.MergeFrom( - super(WhileContext, self)._to_proto(export_scope=export_scope)) + super(WhileContext, self)._to_values_def( + export_scope=export_scope)) + for nested in self._nested_contexts: + nested_def = context_def.nested_contexts.add() + nested.to_control_flow_context_def(nested_def) return context_def else: return None + def to_control_flow_context_def(self, context_def, export_scope=None): + context_def.while_ctxt.CopyFrom(self.to_proto(export_scope=export_scope)) + @staticmethod def from_proto(context_def, import_scope=None): """Returns a `WhileContext` object created from `context_def`. @@ -2294,7 +2347,13 @@ class WhileContext(ControlFlowContext): Returns: A `WhileContext` Python object. """ - return WhileContext(context_def=context_def, import_scope=import_scope) + ret = WhileContext(context_def=context_def, + import_scope=import_scope) + ret.Enter() + for nested_def in context_def.nested_contexts: + from_control_flow_context_def(nested_def, import_scope=import_scope) + ret.Exit() + return ret def GetWhileContext(self): return self @@ -2375,14 +2434,12 @@ class WhileContext(ControlFlowContext): def _AddOpInternal(self, op): """Add `op` to the current context. - In the case that op has only external data inputs, we remove all of its - external control inputs so all its inputs are in the same while loop - context. This is valid because op now has an Enter input that has all - the right control dependency. + We move any external control dependencies of the op to the loop pivot, to + ensure they get executed. """ if not op.inputs: # Remove any external control dependency on this op - control_inputs = self._RemoveExternalControlEdges(op) + control_inputs, external_inputs = self._RemoveExternalControlEdges(op) # Add a control edge from the control pivot to this op. if not control_inputs: # pylint: disable=protected-access @@ -2395,14 +2452,23 @@ class WhileContext(ControlFlowContext): x = op.inputs[index] real_x = self.AddValue(x) if real_x != x: - op._update_input(index, real_x) + op._update_input(index, real_x) # pylint: disable=protected-access # Remove any external control dependency on this op. - self._RemoveExternalControlEdges(op) + _, external_inputs = self._RemoveExternalControlEdges(op) # Add a control dependency to prevent loop invariants from # enabling ops that should not be executed. self._MaybeAddControlDependency(op) for x in op.outputs: self._values.add(x.name) + if external_inputs: + # Use an identity to pull control inputs as data inputs. Note that we + # ignore ops which don't have outputs. TODO(apassos): fix that + with ops.control_dependencies(None): + self.Enter() + external_inputs = [array_ops.identity(x.outputs[0]).op + for x in external_inputs if x.outputs] + self.Exit() + op._add_control_inputs(external_inputs) # pylint: disable=protected-access if self._outer_context or not util.IsLoopExit(op): op.graph.prevent_fetching(op) for x in op.outputs: @@ -2870,8 +2936,11 @@ class WhileContext(ControlFlowContext): loop_vars = ops.convert_n_to_tensor_or_indexed_slices(loop_vars) try: self.Enter() - original_body_result, exit_vars = self._BuildLoop( - pred, body, original_loop_vars, loop_vars, shape_invariants) + # _BuildLoop calls _update_input in several places. _lock ensures a + # Session.run call cannot occur between creating and mutating new ops. + with ops.get_default_graph()._lock: # pylint: disable=protected-access + original_body_result, exit_vars = self._BuildLoop( + pred, body, original_loop_vars, loop_vars, shape_invariants) finally: self.Exit() @@ -3041,6 +3110,43 @@ def while_loop(cond, shape_invariants=[i0.get_shape(), tf.TensorShape([None, 2])]) ``` + Example which demonstrates non-strict semantics: In the following + example, the final value of the counter `i` does not depend on `x`. So + the `while_loop` can increment the counter parallel to updates of `x`. + However, because the loop counter at one loop iteration depends + on the value at the previous iteration, the loop counter itself cannot + be incremented in parallel. Hence if we just want the final value of the + counter (which we print on the line `print(sess.run(i))`), then + `x` will never be incremented, but the counter will be updated on a + single thread. Conversely, if we want the value of the output (which we + print on the line `print(sess.run(out).shape)`), then the counter may be + incremented on its own thread, while `x` can be incremented in + parallel on a separate thread. In the extreme case, it is conceivable + that the thread incrementing the counter runs until completion before + `x` is incremented even a single time. The only thing that can never + happen is that the thread updating `x` can never get ahead of the + counter thread because the thread incrementing `x` depends on the value + of the counter. + ```python + import tensorflow as tf + + n = 10000 + x = tf.constant(list(range(n))) + c = lambda i, x: i < n + b = lambda i, x: (tf.Print(i + 1, [i]), tf.Print(x + 1, [i], "x:")) + i, out = tf.while_loop(c, b, (0, x)) + with tf.Session() as sess: + print(sess.run(i)) # prints [0] ... [9999] + + # The following line may increment the counter and x in parallel. + # The counter thread may get ahead of the other thread, but not the + # other way around. So you may see things like + # [9996] x:[9987] + # meaning that the counter thread is on iteration 9996, + # while the other thread is on iteration 9987 + print(sess.run(out).shape) + ``` + """ with ops.name_scope(name, "while", loop_vars): if not loop_vars: @@ -3074,7 +3180,7 @@ def while_loop(cond, math_ops.logical_and(i < maximum_iterations, orig_cond(*lv))) body = lambda i, lv: (i + 1, orig_body(*lv)) - if context.in_eager_mode(): + if context.executing_eagerly(): while cond(*loop_vars): loop_vars = body(*loop_vars) if maximum_iterations is not None: @@ -3092,7 +3198,10 @@ def while_loop(cond, parallel_iterations=parallel_iterations, back_prop=back_prop, swap_memory=swap_memory) - ops.add_to_collection(ops.GraphKeys.WHILE_CONTEXT, loop_context) + # Only add non-nested loops to the collection. Any nested control flow will + # be encapsulated in the root context. + if loop_context.outer_context is None: + ops.add_to_collection(ops.GraphKeys.WHILE_CONTEXT, loop_context) result = loop_context.BuildLoop(cond, body, loop_vars, shape_invariants) if maximum_iterations is not None: return result[1] @@ -3165,7 +3274,7 @@ def with_dependencies(dependencies, output_tensor, name=None): Raises: TypeError: if `output_tensor` is not a `Tensor` or `IndexedSlices`. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return output_tensor with ops.name_scope(name, "control_dependency", list(dependencies) + [output_tensor]) as name: @@ -3210,7 +3319,7 @@ def group(*inputs, **kwargs): Raises: ValueError: If an unknown keyword argument is provided. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return None name = kwargs.pop("name", None) if kwargs: @@ -3258,7 +3367,7 @@ def group(*inputs, **kwargs): @tf_export("tuple") -def tuple(tensors, name=None, control_inputs=None): +def tuple(tensors, name=None, control_inputs=None): # pylint: disable=redefined-builtin """Group tensors together. This creates a tuple of tensors with the same values as the `tensors` @@ -3290,10 +3399,15 @@ def tuple(tensors, name=None, control_inputs=None): objects. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return tensors with ops.name_scope(name, "tuple", tensors) as name: - gating_ops = [t.op for t in tensors if t is not None] + tensors = [t if (isinstance(t, ops.Operation) + or tensor_util.is_tensor(t) + or t is None) + else ops.convert_to_tensor(t) for t in tensors] + gating_ops = [t if isinstance(t, ops.Operation) else t.op for t in tensors + if t is not None] if control_inputs: for c in control_inputs: if isinstance(c, ops.Tensor): @@ -3309,8 +3423,11 @@ def tuple(tensors, name=None, control_inputs=None): gate = group(*gating_ops) tpl = [] for t in tensors: - if t is not None: + if tensor_util.is_tensor(t): tpl.append(with_dependencies([gate], t)) + elif isinstance(t, ops.Operation): + with ops.control_dependencies([gate]): + tpl.append(group(t)) else: tpl.append(None) return tpl @@ -3370,15 +3487,17 @@ def _case_create_default_action(predicates, actions): return default_action, other_predicates, other_actions -def _case_verify_and_canonicalize_args(pred_fn_pairs, exclusive, name): +def _case_verify_and_canonicalize_args(pred_fn_pairs, exclusive, name, + allow_python_preds): """Verifies input arguments for the case function. Args: - pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor and a - callable which returns a list of tensors. + pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor, + and a callable which returns a list of tensors. exclusive: True iff at most one predicate is allowed to evaluate to `True`. name: A name for the case operation. - + allow_python_preds: if true, pred_fn_pairs may contain Python bools in + addition to boolean Tensors Raises: TypeError: If `pred_fn_pairs` is not a list/dictionary. TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples. @@ -3403,14 +3522,69 @@ def _case_verify_and_canonicalize_args(pred_fn_pairs, exclusive, name): if not isinstance(pred_fn_pair, _basetuple) or len(pred_fn_pair) != 2: raise TypeError("Each entry in pred_fn_pairs must be a 2-tuple") pred, fn = pred_fn_pair - if pred.dtype != dtypes.bool: - raise TypeError("pred must be of type bool: %s", pred.name) + + if isinstance(pred, ops.Tensor): + if pred.dtype != dtypes.bool: + raise TypeError("pred must be Tensor of type bool: %s" % pred.name) + elif not allow_python_preds: + raise TypeError("pred must be a Tensor, got: %s" % pred) + elif not isinstance(pred, bool): + raise TypeError("pred must be a Tensor or bool, got: %s" % pred) + if not callable(fn): raise TypeError("fn for pred %s must be callable." % pred.name) + predicates, actions = zip(*pred_fn_pairs) return predicates, actions +def _case_helper(cond_fn, pred_fn_pairs, default, + exclusive, name, allow_python_preds=False, **cond_kwargs): + """Implementation of case that allows for different cond functions. + + Args: + cond_fn: method that has signature and semantics of `cond` above. + pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor, and a + callable which returns a list of tensors. + default: Optional callable that returns a list of tensors. + exclusive: True iff at most one predicate is allowed to evaluate to `True`. + name: A name for this operation (optional). + allow_python_preds: if true, pred_fn_pairs may contain Python bools in + addition to boolean Tensors + **cond_kwargs: keyword arguments that will be passed to `cond_fn`. + + Returns: + The tensors returned by the first pair whose predicate evaluated to True, or + those returned by `default` if none does. + + Raises: + TypeError: If `pred_fn_pairs` is not a list/dictionary. + TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples. + TypeError: If `fns[i]` is not callable for any i, or `default` is not + callable. + """ + predicates, actions = _case_verify_and_canonicalize_args( + pred_fn_pairs, exclusive, name, allow_python_preds) + with ops.name_scope(name, "case", [predicates]): + if default is None: + default, predicates, actions = _case_create_default_action( + predicates, actions) + fn = default + # To eval conditions in direct order we create nested conditions in reverse: + # cond_fn(c[0], true_fn=.., false_fn=cond_fn(c[1], ...)) + for predicate, action in reversed(list(zip(predicates, actions))): + fn = functools.partial( + cond_fn, predicate, true_fn=action, false_fn=fn, **cond_kwargs) + if exclusive: + with ops.control_dependencies([ + _assert_at_most_n_true( + predicates, n=1, msg="Input error: exclusive=True") + ]): + return fn() + else: + return fn() + + @tf_export("case") def case(pred_fn_pairs, default=None, @@ -3501,26 +3675,8 @@ def case(pred_fn_pairs, TypeError: If `fns[i]` is not callable for any i, or `default` is not callable. """ - predicates, actions = _case_verify_and_canonicalize_args( - pred_fn_pairs, exclusive, name) - with ops.name_scope(name, "case", [predicates]): - if default is None: - default, predicates, actions = _case_create_default_action( - predicates, actions) - fn = default - # To eval conditions in direct order we create nested conditions in reverse: - # cond(c[0], true_fn=.., false_fn=cond(c[1], ...)) - for predicate, action in reversed(list(zip(predicates, actions))): - fn = functools.partial( - cond, predicate, true_fn=action, false_fn=fn, strict=strict) - if exclusive: - with ops.control_dependencies([ - _assert_at_most_n_true( - predicates, n=1, msg="Input error: exclusive=True") - ]): - return fn() - else: - return fn() + return _case_helper(cond, pred_fn_pairs, default, exclusive, name, + allow_python_preds=False, strict=strict) class XLAControlFlowContext(ControlFlowContext): @@ -3540,6 +3696,26 @@ class XLAControlFlowContext(ControlFlowContext): return x +def from_control_flow_context_def(context_def, import_scope=None): + """Deserializes `context_def` into the appropriate ControlFlowContext. + + Args: + context_def: ControlFlowContextDef proto + import_scope: Optional `string`. Name scope to add. + + Returns: + A ControlFlowContext subclass + """ + if context_def.HasField("cond_ctxt"): + return CondContext.from_proto(context_def.cond_ctxt, + import_scope=import_scope) + if context_def.HasField("while_ctxt"): + return WhileContext.from_proto(context_def.while_ctxt, + import_scope=import_scope) + raise NotImplementedError("Unknown ControlFlowContextDef field: %s" + % context_def.WhichOneof("ctxt")) + + ops.register_proto_function( ops.GraphKeys.COND_CONTEXT, proto_type=control_flow_pb2.CondContextDef, diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py index cc5a42bf3ddd4b37d037f8d28a2fe6af79f79ba1..f22f3059d139d1bb7c7db57a2939184f1089f397 100644 --- a/tensorflow/python/ops/control_flow_ops_test.py +++ b/tensorflow/python/ops/control_flow_ops_test.py @@ -189,7 +189,7 @@ class SwitchTestCase(test_util.TensorFlowTestCase): zero = constant_op.constant(0) one = constant_op.constant(1) less_op = math_ops.less(zero, one) - switch_false, switch_true = control_flow_ops.switch(data, less_op) + _, switch_true = control_flow_ops.switch(data, less_op) self.assertAllEqual([1, 2, 3], switch_true.values.eval()) self.assertAllEqual([0, 1], switch_true.indices.eval()) @@ -199,16 +199,17 @@ class SwitchTestCase(test_util.TensorFlowTestCase): "embedding_matrix", [5, 5], initializer=init_ops.random_normal_initializer()) - def Cond(it, _): + def cond(it, _): return it < 5 - def Body(it, cost): + def body(it, cost): embedding = embedding_ops.embedding_lookup(embedding_matrix + 0.0, [0]) cost += math_ops.reduce_sum(embedding) return it + 1, cost _, cost = control_flow_ops.while_loop( - Cond, Body, [constant_op.constant(0), constant_op.constant(0.0)]) + cond, body, [constant_op.constant(0), + constant_op.constant(0.0)]) optimizer = momentum.MomentumOptimizer(0.1, 0.9) train_op = optimizer.minimize(cost) with self.test_session() as sess: @@ -223,16 +224,17 @@ class SwitchTestCase(test_util.TensorFlowTestCase): initializer=[[2.0], [3.0]], use_resource=True) - def Cond(it, _): + def cond(it, _): return it < 5 - def Body(it, cost): + def body(it, cost): embedding = embedding_ops.embedding_lookup(embedding_matrix, [0]) cost += math_ops.reduce_sum(embedding) return it + 1, cost _, cost = control_flow_ops.while_loop( - Cond, Body, [constant_op.constant(0), constant_op.constant(0.0)]) + cond, body, [constant_op.constant(0), + constant_op.constant(0.0)]) with self.test_session() as sess: sess.run(variables.global_variables_initializer()) self.assertAllEqual(10.0, cost.eval()) @@ -244,10 +246,10 @@ class SwitchTestCase(test_util.TensorFlowTestCase): initializer=init_ops.random_normal_initializer(), use_resource=use_resource) - def Cond(it, _): + def cond(it, _): return it < 5 - def Body(it, cost): + def body(it, cost): embedding = embedding_ops.embedding_lookup(embedding_matrix, [0]) cost = control_flow_ops.cond( math_ops.equal(it, 3), lambda: math_ops.square(cost), @@ -255,7 +257,8 @@ class SwitchTestCase(test_util.TensorFlowTestCase): return it + 1, cost _, cost = control_flow_ops.while_loop( - Cond, Body, [constant_op.constant(0), constant_op.constant(0.0)]) + cond, body, [constant_op.constant(0), + constant_op.constant(0.0)]) dynamic_grads = gradients_impl.gradients(cost, [embedding_matrix])[0] dynamic_grads = math_ops.segment_sum(dynamic_grads.values, @@ -289,15 +292,15 @@ class SwitchTestCase(test_util.TensorFlowTestCase): dtype=dtype, size=num_steps) initial_i = constant_op.constant(0, dtype=dtypes.int32) - def Cond(i, _): + def cond(i, _): return i < num_steps # pylint: disable=cell-var-from-loop - def Body(i, outputs): + def body(i, outputs): x = array_ops.gather(inputs, i) # pylint: disable=cell-var-from-loop outputs = outputs.write(i, x) return i + 1, outputs - _, outputs = control_flow_ops.while_loop(Cond, Body, + _, outputs = control_flow_ops.while_loop(cond, body, [initial_i, initial_outputs]) outputs = math_ops.reduce_sum(outputs.stack()) @@ -316,15 +319,15 @@ class SwitchTestCase(test_util.TensorFlowTestCase): dtype=dtype, dynamic_size=True, size=1) initial_i = constant_op.constant(0, dtype=dtypes.int32) - def Cond(i, _): + def cond(i, _): return i < array_ops.size(inputs) # pylint: disable=cell-var-from-loop - def Body(i, outputs): + def body(i, outputs): x = array_ops.gather(inputs, i) # pylint: disable=cell-var-from-loop outputs = outputs.write(i, x) return i + 1, outputs - _, outputs = control_flow_ops.while_loop(Cond, Body, + _, outputs = control_flow_ops.while_loop(cond, body, [initial_i, initial_outputs]) outputs = math_ops.reduce_sum(outputs.stack()) @@ -460,11 +463,12 @@ class ContextTest(test_util.TensorFlowTestCase): control_flow_ops.while_loop( c, b, [i], maximum_iterations=maximum_iterations) for op in sess.graph.get_operations(): - context = op._get_control_flow_context() - if context: - self.assertProtoEquals(context.to_proto(), - control_flow_ops.WhileContext.from_proto( - context.to_proto()).to_proto()) + control_flow_context = op._get_control_flow_context() + if control_flow_context: + self.assertProtoEquals( + control_flow_context.to_proto(), + control_flow_ops.WhileContext.from_proto( + control_flow_context.to_proto()).to_proto()) def testWhileContext(self): self._testWhileContextHelper() @@ -483,8 +487,8 @@ class ContextTest(test_util.TensorFlowTestCase): c._values = ["a", "b"] c._external_values = {"a": b1} - c_with_scope = control_flow_ops.ControlFlowContext._from_proto( - c._to_proto(), import_scope="test_scope") + c_with_scope = control_flow_ops.ControlFlowContext( + values_def=c._to_values_def(), import_scope="test_scope") # _values and _external_values should be have scope prepended. self.assertEquals( @@ -494,12 +498,13 @@ class ContextTest(test_util.TensorFlowTestCase): # Calling _to_proto() with export_scope should remove "test_scope". self.assertProtoEquals( - c._to_proto(), - c_with_scope._to_proto(export_scope="test_scope")) + c._to_values_def(), + c_with_scope._to_values_def(export_scope="test_scope")) + +def _get_nested_shape(nested): -def _GetNestedShape(nested): - def _GetShape(tensor): + def _get_shape(tensor): if isinstance(tensor, tensor_array_ops.TensorArray): return tensor_array_ops.TensorArray elif isinstance(tensor, ops.IndexedSlices): @@ -507,10 +512,10 @@ def _GetNestedShape(nested): else: return tensor.get_shape() - return nest.map_structure(_GetShape, nested) + return nest.map_structure(_get_shape, nested) -def _CreateTensorArray(size, shape): +def _create_tensor_array(size, shape): ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=size, clear_after_read=False) for i in range(size): @@ -518,13 +523,15 @@ def _CreateTensorArray(size, shape): return ta -def _RawNestedShape(nested_shape): - def _RawShape(shape): +def _raw_nested_shape(nested_shape): + + def _raw_shape(shape): if isinstance(shape, tensor_shape.TensorShape) and shape.ndims is not None: return [x.value for x in shape] else: return None - return nest.map_structure(_RawShape, nested_shape) + + return nest.map_structure(_raw_shape, nested_shape) # TODO(yori): Add tests for indexed slices. @@ -543,13 +550,15 @@ class DataTypesTest(test_util.TensorFlowTestCase): condition = array_ops.placeholder(dtypes.bool) output_cond = control_flow_ops.cond(condition, fn_true, fn_false, strict=strict) - self.assertEqual(_RawNestedShape(_GetNestedShape(output_cond)), - _RawNestedShape(expected_shape)) + self.assertEqual( + _raw_nested_shape(_get_nested_shape(output_cond)), + _raw_nested_shape(expected_shape)) output_case = control_flow_ops.case([(condition, fn_true)], fn_false, strict=strict) - self.assertEqual(_RawNestedShape(_GetNestedShape(output_case)), - _RawNestedShape(expected_shape)) + self.assertEqual( + _raw_nested_shape(_get_nested_shape(output_case)), + _raw_nested_shape(expected_shape)) def _testReturnValues(self, fn_true, fn_false, expected_value_true, expected_value_false, strict=False, @@ -626,45 +635,55 @@ class DataTypesTest(test_util.TensorFlowTestCase): control_flow_ops.cond(constant_op.constant(True), fn_tensor, fn_none) def test_tensors(self): - def _BuildTrueBranch(dtype): - def _Build(): + + def _build_true_branch(dtype): + + def _build(): return (array_ops.zeros([2, 2], dtype=dtype), array_ops.ones([3, 3], dtype=dtype)) - return _Build - def _BuildFalseBranch(dtype): - def _Build(): + return _build + + def _build_false_branch(dtype): + + def _build(): return (array_ops.ones([2, 2], dtype=dtype), array_ops.zeros([3, 3], dtype=dtype)) - return _Build + + return _build for dtype in (dtypes.float16, dtypes.int8, dtypes.int32, dtypes.uint8): shape = (tensor_shape.TensorShape([2, 2]), tensor_shape.TensorShape([3, 3])) - fn_true = _BuildTrueBranch(dtype) - fn_false = _BuildFalseBranch(dtype) + fn_true = _build_true_branch(dtype) + fn_false = _build_false_branch(dtype) self._testShape(fn_true, fn_false, shape) self._testReturnValues(fn_true, fn_false, (np.zeros([2, 2]), np.ones([3, 3])), (np.ones([2, 2]), np.zeros([3, 3]))) def test_tensors_unknown_shape(self): - def _BuildTrueBranch(dtype): + + def _build_true_branch(dtype): tensor = array_ops.placeholder(dtype=dtype, shape=None) - def _Build(): + + def _build(): return tensor - return _Build, tensor - def _BuildFalseBranch(dtype): + return _build, tensor + + def _build_false_branch(dtype): tensor = array_ops.placeholder(dtype=dtype, shape=None) - def _Build(): + + def _build(): return tensor - return _Build, tensor + + return _build, tensor for dtype in (dtypes.float16, dtypes.int8, dtypes.int32, dtypes.uint8): shape = tensor_shape.TensorShape(None) - fn_true, true_tensor = _BuildTrueBranch(dtype) - fn_false, false_tensor = _BuildFalseBranch(dtype) + fn_true, true_tensor = _build_true_branch(dtype) + fn_false, false_tensor = _build_false_branch(dtype) self._testShape(fn_true, fn_false, shape) self._testReturnValues(fn_true, fn_false, np.zeros([2, 2]), np.ones([2, 2]), @@ -674,11 +693,11 @@ class DataTypesTest(test_util.TensorFlowTestCase): def test_sparse_tensors(self): shape = tensor_shape.TensorShape([None, None]) - def FnTrue(): + def true_fn(): return [sparse_tensor.SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])] - def FnFalse(): + def false_fn(): return [sparse_tensor.SparseTensor(indices=[[0, 0], [2, 1]], values=[3, 4], dense_shape=[3, 4])] @@ -686,26 +705,29 @@ class DataTypesTest(test_util.TensorFlowTestCase): values=[1, 2], dense_shape=[3, 4]) value2 = sparse_tensor.SparseTensorValue(indices=[[0, 0], [2, 1]], values=[3, 4], dense_shape=[3, 4]) - self._testShape(FnTrue, FnFalse, shape) - self._testReturnValues(FnTrue, FnFalse, value1, value2) - self._testShape(FnTrue, FnFalse, [shape], strict=True) - self._testReturnValues(FnTrue, FnFalse, [value1], [value2], strict=True) + self._testShape(true_fn, false_fn, shape) + self._testReturnValues(true_fn, false_fn, value1, value2) + self._testShape(true_fn, false_fn, [shape], strict=True) + self._testReturnValues(true_fn, false_fn, [value1], [value2], strict=True) def test_tensors_with_partially_specified_shapes(self): - def _BuildBranch(dtype, shape): + + def _build_branch(dtype, shape): a = array_ops.placeholder(dtype=dtype, shape=shape[0]) b = array_ops.placeholder(dtype=dtype, shape=shape[1]) c = array_ops.placeholder(dtype=dtype, shape=shape[2]) - def _Build(): + + def _build(): return a, b, c - return _Build, (a, b, c) + + return _build, (a, b, c) for dtype in (dtypes.float16, dtypes.int8, dtypes.int32, dtypes.uint8): shape = (tensor_shape.TensorShape([None, 2]), tensor_shape.TensorShape([None]), tensor_shape.TensorShape([3, None])) - fn_true, true_tensors = _BuildBranch(dtype, shape) - fn_false, false_tensors = _BuildBranch(dtype, shape) + fn_true, true_tensors = _build_branch(dtype, shape) + fn_false, false_tensors = _build_branch(dtype, shape) self._testShape(fn_true, fn_false, shape) self._testReturnValues(fn_true, fn_false, (np.zeros([2, 2]), np.zeros(5), np.ones([3, 3])), @@ -719,8 +741,8 @@ class DataTypesTest(test_util.TensorFlowTestCase): def test_tensor_arrays(self): element_shape = tensor_shape.TensorShape([2]) - ta1 = _CreateTensorArray(4, element_shape) - ta2 = _CreateTensorArray(4, element_shape) + ta1 = _create_tensor_array(4, element_shape) + ta2 = _create_tensor_array(4, element_shape) shape = tensor_array_ops.TensorArray fn_true = lambda: ta1 fn_false = lambda: ta2 @@ -728,7 +750,7 @@ class DataTypesTest(test_util.TensorFlowTestCase): def test_tensor_array_reads(self): shape = tensor_shape.TensorShape([2]) - ta = _CreateTensorArray(4, shape) + ta = _create_tensor_array(4, shape) fn_true = lambda: ta.read(0) fn_false = lambda: ta.read(1) self._testShape(fn_true, fn_false, shape) @@ -827,23 +849,26 @@ class DataTypesTest(test_util.TensorFlowTestCase): tensor_shape.TensorShape([5, 5]), tensor_shape.TensorShape([])] - def FnTrue(): + def true_fn(): return [constant_op.constant(1), TestTuple(constant_op.constant(2), [3, 4]), array_ops.zeros([5, 5]), 6] - def FnFalse(): + def false_fn(): return [constant_op.constant(11), TestTuple(constant_op.constant(12), [13, 14]), array_ops.ones([5, 5]), 16] - self._testShape(FnTrue, FnFalse, shape) - self._testReturnValues(FnTrue, FnFalse, - [1, TestTuple(2, [3, 4]), np.zeros([5, 5]), 6], - [11, TestTuple(12, [13, 14]), np.ones([5, 5]), 16]) + self._testShape(true_fn, false_fn, shape) + self._testReturnValues( + true_fn, false_fn, + [1, TestTuple(2, [3, 4]), np.zeros([5, 5]), 6], + [11, TestTuple(12, [13, 14]), + np.ones([5, 5]), 16]) def test_cond_inside_while_loop(self): - def Body(i, matrix): + + def body(i, matrix): result_tuple, unused_matrix = control_flow_ops.cond( constant_op.constant(True), lambda: (TestTuple(matrix * 2, matrix * 4), matrix), @@ -852,8 +877,9 @@ class DataTypesTest(test_util.TensorFlowTestCase): iteration, matrix = control_flow_ops.while_loop( lambda i, matrix: i < 10, - Body, - loop_vars=[constant_op.constant(0), array_ops.ones([2, 2])]) + body, + loop_vars=[constant_op.constant(0), + array_ops.ones([2, 2])]) self.assertEqual(iteration.get_shape(), tensor_shape.TensorShape([])) self.assertEqual(matrix.get_shape(), tensor_shape.TensorShape([2, 2])) diff --git a/tensorflow/python/ops/ctc_ops.py b/tensorflow/python/ops/ctc_ops.py index 83da6739db673644f59fda3044769b18b2138fbc..4b57e2de790af13499bc73cfcfa98e999eab1603 100644 --- a/tensorflow/python/ops/ctc_ops.py +++ b/tensorflow/python/ops/ctc_ops.py @@ -148,7 +148,7 @@ def ctc_loss(labels, inputs, sequence_length, if not time_major: inputs = array_ops.transpose(inputs, [1, 0, 2]) # (B,T,N) => (T,B,N) - loss, _ = gen_ctc_ops._ctc_loss( + loss, _ = gen_ctc_ops.ctc_loss( inputs, labels.indices, labels.values, @@ -224,7 +224,7 @@ def ctc_greedy_decoder(inputs, sequence_length, merge_repeated=True): sequence found, the negative of the sum of the greatest logit at each timeframe. """ - outputs = gen_ctc_ops._ctc_greedy_decoder( + outputs = gen_ctc_ops.ctc_greedy_decoder( inputs, sequence_length, merge_repeated=merge_repeated) (decoded_ix, decoded_val, decoded_shape, log_probabilities) = outputs return ([sparse_tensor.SparseTensor(decoded_ix, decoded_val, decoded_shape)], @@ -272,7 +272,7 @@ def ctc_beam_search_decoder(inputs, sequence_length, beam_width=100, """ decoded_ixs, decoded_vals, decoded_shapes, log_probabilities = ( - gen_ctc_ops._ctc_beam_search_decoder( + gen_ctc_ops.ctc_beam_search_decoder( inputs, sequence_length, beam_width=beam_width, top_paths=top_paths, merge_repeated=merge_repeated)) diff --git a/tensorflow/python/ops/custom_gradient.py b/tensorflow/python/ops/custom_gradient.py new file mode 100644 index 0000000000000000000000000000000000000000..9eacac1b3704c43cbeb5ecd0cbe827cac3a7cc8b --- /dev/null +++ b/tensorflow/python/ops/custom_gradient.py @@ -0,0 +1,134 @@ +# 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. +# ============================================================================== +"""Decorator to overrides the gradient for a function.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.eager import context +from tensorflow.python.eager import tape +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_array_ops +from tensorflow.python.util import nest +from tensorflow.python.util import tf_decorator +from tensorflow.python.util.tf_export import tf_export + + +@tf_export("custom_gradient") +def custom_gradient(f): + """Decorator to define a function with a custom gradient. + + This decorator allows fine grained control over the gradients of a sequence + for operations. This may be useful for multiple reasons, including providing + a more efficient or numerically stable gradient for a sequence of operations. + + For example, consider the following function that commonly occurs in the + computation of cross entropy and log likelihoods: + + ```python + def log1pexp(x): + return tf.log(1 + tf.exp(x)) + ``` + + Due to numerical instability, the gradient this function evaluated at x=100 is + NaN. For example: + + ```python + x = tf.constant(100.) + y = log1pexp(x) + dy = tf.gradients(y, x) # Will be NaN when evaluated. + ``` + + The gradient expression can be analytically simplified to provide numerical + stability: + + ```python + @tf.custom_gradient + def log1pexp(x): + e = tf.exp(x) + def grad(dy): + return dy * (1 - 1 / (1 + e)) + return tf.log(1 + e), grad + ``` + + With this definition, the gradient at x=100 will be correctly evaluated as + 1.0. + + See also @{tf.RegisterGradient} which registers a gradient function for a + primitive TensorFlow operation. `tf.custom_gradient` on the other hand allows + for fine grained control over the gradient computation of a sequence of + operations. + + Args: + f: function `f(x)` that returns a tuple `(y, grad_fn)` where: + - `x` is a `Tensor` or sequence of `Tensor` inputs to the function. + - `y` is a `Tensor` or sequence of `Tensor` outputs of applying + TensorFlow + operations in `f` to `x`. + - `grad_fn` is a function with the signature `g(grad_ys)` which returns + a list of `Tensor`s - the derivatives of `Tensor`s in `y` with respect + to the `Tensor`s in `x. `grad_ys` is a `Tensor` or sequence of + `Tensor`s the same size as `y` holding the initial value gradients for + each `Tensor` in `y`. + + Returns: + A function `h(x)` which returns the same value as `f(x)[0]` and whose + gradient (as calculated by @{tf.gradients}) is determined by `f(x)[1]`. + """ + + def decorated(*args, **kwargs): + """Decorated function with custom gradient.""" + if not context.executing_eagerly(): + if kwargs: + raise ValueError( + "The custom_gradient decorator currently suports keywords " + "arguments only when eager execution is enabled.") + name = "CustomGradient-%s" % ops.uid() + args = [ops.convert_to_tensor(x) for x in args] + result, grad_fn = f(*args) + flat_result = nest.flatten(result) + all_tensors = flat_result + args + + @ops.RegisterGradient(name) + def internal_grad_fn(unused_op, *result_grads): # pylint: disable=unused-variable + gradients = nest.flatten(grad_fn(*result_grads[:len(flat_result)])) + # Need to return one value per input to the IdentityN, so pad the + # gradients of the inputs of the custom_gradient function with the + # gradients of the outputs as well. + return ([None] * len(flat_result)) + gradients + + with ops.get_default_graph().gradient_override_map({"IdentityN": name}): + all_tensors = array_ops.identity_n(all_tensors) + return nest.pack_sequence_as( + structure=result, flat_sequence=all_tensors[:len(flat_result)]) + + input_tensors = [ops.convert_to_tensor(x) for x in args] + + result, grad_fn = f(*args, **kwargs) + flat_result = nest.flatten(result) + # TODO(apassos) consider removing the identity below. + flat_result = [gen_array_ops.identity(x) for x in flat_result] + + def actual_grad_fn(*outputs): + return nest.flatten(grad_fn(*outputs)) + + tape.record_operation(f.__name__, flat_result, input_tensors, + actual_grad_fn) + flat_result = list(flat_result) + return nest.pack_sequence_as(result, flat_result) + + return tf_decorator.make_decorator(f, decorated) diff --git a/tensorflow/python/ops/data_flow_ops.py b/tensorflow/python/ops/data_flow_ops.py index 95e45bff066d4b2653e5de67684a6277006345f2..d2cc87555f6321432261b32f08431c23ce707eff 100644 --- a/tensorflow/python/ops/data_flow_ops.py +++ b/tensorflow/python/ops/data_flow_ops.py @@ -159,7 +159,7 @@ class QueueBase(object): ValueError: If one of the arguments is invalid. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "Queues are not supported when eager execution is enabled. " "Instead, please use tf.data to get data into your model.") @@ -177,10 +177,10 @@ class QueueBase(object): else: self._names = None self._queue_ref = queue_ref - if context.in_graph_mode(): - self._name = self._queue_ref.op.name.split("/")[-1] - else: + if context.executing_eagerly(): self._name = context.context().scope_name + else: + self._name = self._queue_ref.op.name.split("/")[-1] @staticmethod def from_list(index, queues): @@ -231,9 +231,9 @@ class QueueBase(object): @property def name(self): """The name of the underlying queue.""" - if context.in_graph_mode(): - return self._queue_ref.op.name - return self._name + if context.executing_eagerly(): + return self._name + return self._queue_ref.op.name @property def dtypes(self): @@ -342,10 +342,10 @@ class QueueBase(object): val.get_shape().assert_is_compatible_with(shape) if self._queue_ref.dtype == _dtypes.resource: - return gen_data_flow_ops._queue_enqueue_v2( + return gen_data_flow_ops.queue_enqueue_v2( self._queue_ref, vals, name=scope) else: - return gen_data_flow_ops._queue_enqueue( + return gen_data_flow_ops.queue_enqueue( self._queue_ref, vals, name=scope) def enqueue_many(self, vals, name=None): @@ -387,7 +387,7 @@ class QueueBase(object): val.get_shape().with_rank_at_least(1)[0]) val.get_shape()[1:].assert_is_compatible_with(shape) - return gen_data_flow_ops._queue_enqueue_many_v2( + return gen_data_flow_ops.queue_enqueue_many_v2( self._queue_ref, vals, name=scope) def _dequeue_return_value(self, tensors): @@ -436,15 +436,15 @@ class QueueBase(object): if name is None: name = "%s_Dequeue" % self._name if self._queue_ref.dtype == _dtypes.resource: - ret = gen_data_flow_ops._queue_dequeue_v2( + ret = gen_data_flow_ops.queue_dequeue_v2( self._queue_ref, self._dtypes, name=name) else: - ret = gen_data_flow_ops._queue_dequeue( + ret = gen_data_flow_ops.queue_dequeue( self._queue_ref, self._dtypes, name=name) # NOTE(mrry): Not using a shape function because we need access to # the `QueueBase` object. - if context.in_graph_mode(): + if not context.executing_eagerly(): op = ret[0].op for output, shape in zip(op.values(), self._shapes): output.set_shape(shape) @@ -474,17 +474,17 @@ class QueueBase(object): name: A name for the operation (optional). Returns: - The tuple of concatenated tensors that was dequeued. + The list of concatenated tensors that was dequeued. """ if name is None: name = "%s_DequeueMany" % self._name - ret = gen_data_flow_ops._queue_dequeue_many_v2( + ret = gen_data_flow_ops.queue_dequeue_many_v2( self._queue_ref, n=n, component_types=self._dtypes, name=name) # NOTE(mrry): Not using a shape function because we need access to # the Queue object. - if context.in_graph_mode(): + if not context.executing_eagerly(): op = ret[0].op batch_dim = tensor_shape.Dimension( tensor_util.constant_value(op.inputs[1])) @@ -523,12 +523,12 @@ class QueueBase(object): if name is None: name = "%s_DequeueUpTo" % self._name - ret = gen_data_flow_ops._queue_dequeue_up_to_v2( + ret = gen_data_flow_ops.queue_dequeue_up_to_v2( self._queue_ref, n=n, component_types=self._dtypes, name=name) # NOTE(mrry): Not using a shape function because we need access to # the Queue object. - if context.in_graph_mode(): + if not context.executing_eagerly(): op = ret[0].op for output, shape in zip(op.values(), self._shapes): output.set_shape(tensor_shape.TensorShape([None]).concatenate(shape)) @@ -560,12 +560,12 @@ class QueueBase(object): if name is None: name = "%s_Close" % self._name if self._queue_ref.dtype == _dtypes.resource: - return gen_data_flow_ops._queue_close_v2( + return gen_data_flow_ops.queue_close_v2( self._queue_ref, cancel_pending_enqueues=cancel_pending_enqueues, name=name) else: - return gen_data_flow_ops._queue_close( + return gen_data_flow_ops.queue_close( self._queue_ref, cancel_pending_enqueues=cancel_pending_enqueues, name=name) @@ -601,9 +601,9 @@ class QueueBase(object): if name is None: name = "%s_Size" % self._name if self._queue_ref.dtype == _dtypes.resource: - return gen_data_flow_ops._queue_size_v2(self._queue_ref, name=name) + return gen_data_flow_ops.queue_size_v2(self._queue_ref, name=name) else: - return gen_data_flow_ops._queue_size(self._queue_ref, name=name) + return gen_data_flow_ops.queue_size(self._queue_ref, name=name) @tf_export("RandomShuffleQueue") @@ -683,7 +683,7 @@ class RandomShuffleQueue(QueueBase): # the id of the last op created.) string = (str(seed1) + shared_name).encode("utf-8") seed2 = int(hashlib.md5(string).hexdigest()[:8], 16) & 0x7FFFFFFF - queue_ref = gen_data_flow_ops._random_shuffle_queue_v2( + queue_ref = gen_data_flow_ops.random_shuffle_queue_v2( component_types=dtypes, shapes=shapes, capacity=capacity, @@ -748,7 +748,7 @@ class FIFOQueue(QueueBase): dtypes = _as_type_list(dtypes) shapes = _as_shape_list(shapes, dtypes) names = _as_name_list(names, dtypes) - queue_ref = gen_data_flow_ops._fifo_queue_v2( + queue_ref = gen_data_flow_ops.fifo_queue_v2( component_types=dtypes, shapes=shapes, capacity=capacity, @@ -827,7 +827,7 @@ class PaddingFIFOQueue(QueueBase): "but received %d dtypes and %d shapes." % (len(dtypes), len(shapes))) - queue_ref = gen_data_flow_ops._padding_fifo_queue_v2( + queue_ref = gen_data_flow_ops.padding_fifo_queue_v2( component_types=dtypes, shapes=shapes, capacity=capacity, @@ -895,7 +895,7 @@ class PriorityQueue(QueueBase): types = _as_type_list(types) shapes = _as_shape_list(shapes, types) - queue_ref = gen_data_flow_ops._priority_queue_v2( + queue_ref = gen_data_flow_ops.priority_queue_v2( component_types=types, shapes=shapes, capacity=capacity, @@ -985,15 +985,15 @@ class Barrier(object): else: self._shapes = [tensor_shape.unknown_shape() for _ in self._types] - self._barrier_ref = gen_data_flow_ops._barrier( + self._barrier_ref = gen_data_flow_ops.barrier( component_types=self._types, shapes=self._shapes, shared_name=shared_name, name=name) - if context.in_graph_mode(): - self._name = self._barrier_ref.op.name.split("/")[-1] - else: + if context.executing_eagerly(): self._name = context.context().scope_name + else: + self._name = self._barrier_ref.op.name.split("/")[-1] @property def barrier_ref(self): @@ -1003,9 +1003,9 @@ class Barrier(object): @property def name(self): """The name of the underlying barrier.""" - if context.in_graph_mode(): - return self._barrier_ref.op.name - return self._name + if context.executing_eagerly(): + return self._name + return self._barrier_ref.op.name def insert_many(self, component_index, keys, values, name=None): """For each key, assigns the respective value to the specified component. @@ -1026,7 +1026,7 @@ class Barrier(object): """ if name is None: name = "%s_BarrierInsertMany" % self._name - return gen_data_flow_ops._barrier_insert_many( + return gen_data_flow_ops.barrier_insert_many( self._barrier_ref, keys, values, component_index, name=name) def take_many(self, @@ -1073,7 +1073,7 @@ class Barrier(object): """ if name is None: name = "%s_BarrierTakeMany" % self._name - ret = gen_data_flow_ops._barrier_take_many( + ret = gen_data_flow_ops.barrier_take_many( self._barrier_ref, num_elements, self._types, @@ -1083,7 +1083,7 @@ class Barrier(object): # NOTE(mrry): Not using a shape function because we need access to # the Barrier object. - if context.in_graph_mode(): + if not context.executing_eagerly(): op = ret[0].op if allow_small_batch: batch_dim = None @@ -1122,7 +1122,7 @@ class Barrier(object): """ if name is None: name = "%s_BarrierClose" % self._name - return gen_data_flow_ops._barrier_close( + return gen_data_flow_ops.barrier_close( self._barrier_ref, cancel_pending_enqueues=cancel_pending_enqueues, name=name) @@ -1139,7 +1139,7 @@ class Barrier(object): """ if name is None: name = "%s_BarrierReadySize" % self._name - return gen_data_flow_ops._barrier_ready_size(self._barrier_ref, name=name) + return gen_data_flow_ops.barrier_ready_size(self._barrier_ref, name=name) def incomplete_size(self, name=None): """Compute the number of incomplete elements in the given barrier. @@ -1153,7 +1153,7 @@ class Barrier(object): """ if name is None: name = "%s_BarrierIncompleteSize" % self._name - return gen_data_flow_ops._barrier_incomplete_size( + return gen_data_flow_ops.barrier_incomplete_size( self._barrier_ref, name=name) @@ -1183,10 +1183,10 @@ class ConditionalAccumulatorBase(object): else: self._shape = tensor_shape.unknown_shape() self._accumulator_ref = accumulator_ref - if context.in_graph_mode(): - self._name = self._accumulator_ref.op.name.split("/")[-1] - else: + if context.executing_eagerly(): self._name = context.context().scope_name + else: + self._name = self._accumulator_ref.op.name.split("/")[-1] @property def accumulator_ref(self): diff --git a/tensorflow/python/ops/distributions/bernoulli.py b/tensorflow/python/ops/distributions/bernoulli.py index 1f300b7147be505a316c38ae57cadeae2bd7ea10..68aaf3815e7e2b21c9550562aa49195569c8ea43 100644 --- a/tensorflow/python/ops/distributions/bernoulli.py +++ b/tensorflow/python/ops/distributions/bernoulli.py @@ -22,7 +22,6 @@ 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 control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn from tensorflow.python.ops import random_ops @@ -137,21 +136,12 @@ class Bernoulli(distribution.Distribution): return (array_ops.ones_like(event) * logits, array_ops.ones_like(logits) * event) - # First check static shape. - if (event.get_shape().is_fully_defined() and - logits.get_shape().is_fully_defined()): - if event.get_shape() != logits.get_shape(): - logits, event = _broadcast(logits, event) - else: - logits, event = control_flow_ops.cond( - distribution_util.same_dynamic_shape(logits, event), - lambda: (logits, event), - lambda: _broadcast(logits, event)) + if not (event.get_shape().is_fully_defined() and + logits.get_shape().is_fully_defined() and + event.get_shape() == logits.get_shape()): + logits, event = _broadcast(logits, event) return -nn.sigmoid_cross_entropy_with_logits(labels=event, logits=logits) - def _prob(self, event): - return math_ops.exp(self._log_prob(event)) - def _entropy(self): return (-self.logits * (math_ops.sigmoid(self.logits) - 1) + nn.softplus(-self.logits)) @@ -167,26 +157,6 @@ class Bernoulli(distribution.Distribution): return math_ops.cast(self.probs > 0.5, self.dtype) -class BernoulliWithSigmoidProbs(Bernoulli): - """Bernoulli with `probs = nn.sigmoid(logits)`.""" - - def __init__(self, - logits=None, - dtype=dtypes.int32, - validate_args=False, - allow_nan_stats=True, - name="BernoulliWithSigmoidProbs"): - parameters = locals() - with ops.name_scope(name): - super(BernoulliWithSigmoidProbs, self).__init__( - probs=nn.sigmoid(logits, name="sigmoid_probs"), - dtype=dtype, - validate_args=validate_args, - allow_nan_stats=allow_nan_stats, - name=name) - self._parameters = parameters - - @kullback_leibler.RegisterKL(Bernoulli, Bernoulli) def _kl_bernoulli_bernoulli(a, b, name=None): """Calculate the batched KL divergence KL(a || b) with a and b Bernoulli. diff --git a/tensorflow/python/ops/distributions/beta.py b/tensorflow/python/ops/distributions/beta.py index 6d6b40b04557a4483f60d8c06c35f937d38a24b9..469bcadb8ea3a0ec2a85d3a72c0ca5ba08796856 100644 --- a/tensorflow/python/ops/distributions/beta.py +++ b/tensorflow/python/ops/distributions/beta.py @@ -304,12 +304,11 @@ class Beta(distribution.Distribution): if not self.validate_args: return x return control_flow_ops.with_dependencies([ - check_ops.assert_positive( - x, - message="sample must be positive"), + check_ops.assert_positive(x, message="sample must be positive"), check_ops.assert_less( - x, array_ops.ones([], self.dtype), - message="sample must be no larger than `1`."), + x, + array_ops.ones([], self.dtype), + message="sample must be less than `1`."), ], x) diff --git a/tensorflow/python/ops/distributions/bijector_impl.py b/tensorflow/python/ops/distributions/bijector_impl.py index 44d64070ce48c0c115ea7edb1237124bc6698e90..ed435557fde7a2e8a0a4f7eef4e240daef0565e7 100644 --- a/tensorflow/python/ops/distributions/bijector_impl.py +++ b/tensorflow/python/ops/distributions/bijector_impl.py @@ -114,7 +114,7 @@ class _Mapping(collections.namedtuple( @six.add_metaclass(abc.ABCMeta) @tf_export("distributions.bijectors.Bijector") class Bijector(object): - """Interface for transformations of a `Distribution` sample. + r"""Interface for transformations of a `Distribution` sample. Bijectors can be used to represent any differentiable and injective (one to one) function defined on an open subset of `R^n`. Some non-injective @@ -122,27 +122,24 @@ class Bijector(object): #### Mathematical Details - A `Bijector` implements a - [diffeomorphism](https://en.wikipedia.org/wiki/Diffeomorphism), i.e., a - bijective, differentiable function. A `Bijector` is used by - `TransformedDistribution` but can be generally used for transforming a - `Distribution` generated `Tensor`. A `Bijector` is characterized by three - operations: - - 1. Forward Evaluation + A `Bijector` implements a [smooth covering map]( + https://en.wikipedia.org/wiki/Local_diffeomorphism), i.e., a local + diffeomorphism such that every point in the target has a neighborhood evenly + covered by a map ([see also]( + https://en.wikipedia.org/wiki/Covering_space#Covering_of_a_manifold)). + A `Bijector` is used by `TransformedDistribution` but can be generally used + for transforming a `Distribution` generated `Tensor`. A `Bijector` is + characterized by three operations: + 1. Forward\ Useful for turning one random outcome into another random outcome from a different distribution. - - 2. Inverse Evaluation - + 2. Inverse\ Useful for "reversing" a transformation to compute one probability in terms of another. - - 3. (log o det o Jacobian o inverse)(x) - + 3. `(log o det o Jacobian o inverse)(x)`\ "The log of the determinant of the matrix of all first-order partial - derivatives of the inverse function." + derivatives of the inverse function."\ Useful for inverting a transformation to compute one probability in terms of another. Geometrically, the det(Jacobian) is the volume of the transformation and is used to scale the probability. diff --git a/tensorflow/python/ops/distributions/gamma.py b/tensorflow/python/ops/distributions/gamma.py index 8fb218be3ac7e17e18d85b8e1c100ccd58aa1034..adb1f4f9a879e44cf8cb4cafd22b92554f487712 100644 --- a/tensorflow/python/ops/distributions/gamma.py +++ b/tensorflow/python/ops/distributions/gamma.py @@ -193,12 +193,6 @@ class Gamma(distribution.Distribution): def _log_prob(self, x): return self._log_unnormalized_prob(x) - self._log_normalization() - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - - def _log_cdf(self, x): - return math_ops.log(self._cdf(x)) - def _cdf(self, x): x = self._maybe_assert_valid_sample(x) # Note that igamma returns the regularized incomplete gamma function, diff --git a/tensorflow/python/ops/distributions/multinomial.py b/tensorflow/python/ops/distributions/multinomial.py index 26b5c5aef98fc11b07a8c8357e7ec37819587da9..4ae67a009b0a4052f6e23e2e42262bb7c42f1c14 100644 --- a/tensorflow/python/ops/distributions/multinomial.py +++ b/tensorflow/python/ops/distributions/multinomial.py @@ -238,7 +238,7 @@ class Multinomial(distribution.Distribution): n_draws = math_ops.cast(self.total_count, dtype=dtypes.int32) k = self.event_shape_tensor()[0] - # boardcast the total_count and logits to same shape + # broadcast the total_count and logits to same shape n_draws = array_ops.ones_like( self.logits[..., 0], dtype=n_draws.dtype) * n_draws logits = array_ops.ones_like( diff --git a/tensorflow/python/ops/distributions/normal.py b/tensorflow/python/ops/distributions/normal.py index e7f120ea2da525e20a1ae42e6418cf2ac83686af..32e8a49c81bc4b23d8897639998dd33942b41a80 100644 --- a/tensorflow/python/ops/distributions/normal.py +++ b/tensorflow/python/ops/distributions/normal.py @@ -188,9 +188,6 @@ class Normal(distribution.Distribution): def _log_prob(self, x): return self._log_unnormalized_prob(x) - self._log_normalization() - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - def _log_cdf(self, x): return special_math.log_ndtr(self._z(x)) diff --git a/tensorflow/python/ops/distributions/special_math.py b/tensorflow/python/ops/distributions/special_math.py index bed4cbb2c1a43b6952861f4fab82957229e23c9c..1d605c5dfcca9b709a9178ccbe56619f6a92f869 100644 --- a/tensorflow/python/ops/distributions/special_math.py +++ b/tensorflow/python/ops/distributions/special_math.py @@ -213,7 +213,7 @@ def _ndtri(p): # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z), # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different - # arrays based on wether p < exp(-32). + # arrays based on whether p < exp(-32). z = math_ops.sqrt(-2. * math_ops.log(sanitized_mcp)) first_term = z - math_ops.log(z) / z second_term_small_p = (_create_polynomial(1. / z, p2) diff --git a/tensorflow/python/ops/distributions/student_t.py b/tensorflow/python/ops/distributions/student_t.py index 778fefb8c2991153b7e7a1f20df61680153dab2a..9d9e65b4e8d6d2e40bf9c263339f899439c842c3 100644 --- a/tensorflow/python/ops/distributions/student_t.py +++ b/tensorflow/python/ops/distributions/student_t.py @@ -248,9 +248,6 @@ class StudentT(distribution.Distribution): math_ops.lgamma(0.5 * self.df) - math_ops.lgamma(0.5 * (self.df + 1.))) - def _prob(self, x): - return math_ops.exp(self._log_prob(x)) - def _cdf(self, x): # Take Abs(scale) to make subsequent where work correctly. y = (x - self.loc) / math_ops.abs(self.scale) diff --git a/tensorflow/python/ops/distributions/uniform.py b/tensorflow/python/ops/distributions/uniform.py index 3580af18f241d777c81340f1c565074914838029..ec623b55eb0067e16599c18c9c504635da863907 100644 --- a/tensorflow/python/ops/distributions/uniform.py +++ b/tensorflow/python/ops/distributions/uniform.py @@ -45,11 +45,12 @@ class Uniform(distribution.Distribution): Z = b - a ``` - where: - * `low = a`, - * `high = b`, - * `Z` is the normalizing constant, and, - * `I[predicate]` is the [indicator function]( + where + + - `low = a`, + - `high = b`, + - `Z` is the normalizing constant, and + - `I[predicate]` is the [indicator function]( https://en.wikipedia.org/wiki/Indicator_function) for `predicate`. The parameters `low` and `high` must be shaped in a way that supports @@ -164,9 +165,6 @@ class Uniform(distribution.Distribution): seed=seed) return self.low + self.range() * samples - def _log_prob(self, x): - return math_ops.log(self._prob(x)) - def _prob(self, x): broadcasted_x = x * array_ops.ones(self.batch_shape_tensor()) return array_ops.where( @@ -178,9 +176,6 @@ class Uniform(distribution.Distribution): array_ops.zeros_like(broadcasted_x), array_ops.ones_like(broadcasted_x) / self.range())) - def _log_cdf(self, x): - return math_ops.log(self.cdf(x)) - def _cdf(self, x): broadcast_shape = array_ops.broadcast_dynamic_shape( array_ops.shape(x), self.batch_shape_tensor()) diff --git a/tensorflow/python/ops/distributions/util.py b/tensorflow/python/ops/distributions/util.py index 5bc25128a8d6f77895fc4decc98a8978ae8400f3..0fe6aa30f945dc7682a53fa6495823288cf111b7 100644 --- a/tensorflow/python/ops/distributions/util.py +++ b/tensorflow/python/ops/distributions/util.py @@ -1041,14 +1041,14 @@ def reduce_weighted_logsumexp( with ops.name_scope(name, "reduce_weighted_logsumexp", [logx, w]): logx = ops.convert_to_tensor(logx, name="logx") if w is None: - lswe = math_ops.reduce_logsumexp(logx, axis=axis, keep_dims=keep_dims) + lswe = math_ops.reduce_logsumexp(logx, axis=axis, keepdims=keep_dims) if return_sign: sgn = array_ops.ones_like(lswe) return lswe, sgn return lswe w = ops.convert_to_tensor(w, dtype=logx.dtype, name="w") log_absw_x = logx + math_ops.log(math_ops.abs(w)) - max_log_absw_x = math_ops.reduce_max(log_absw_x, axis=axis, keep_dims=True) + max_log_absw_x = math_ops.reduce_max(log_absw_x, axis=axis, keepdims=True) # If the largest element is `-inf` or `inf` then we don't bother subtracting # off the max. We do this because otherwise we'd get `inf - inf = NaN`. That # this is ok follows from the fact that we're actually free to subtract any @@ -1060,9 +1060,7 @@ def reduce_weighted_logsumexp( wx_over_max_absw_x = ( math_ops.sign(w) * math_ops.exp(log_absw_x - max_log_absw_x)) sum_wx_over_max_absw_x = math_ops.reduce_sum( - wx_over_max_absw_x, - axis=axis, - keep_dims=keep_dims) + wx_over_max_absw_x, axis=axis, keepdims=keep_dims) if not keep_dims: max_log_absw_x = array_ops.squeeze(max_log_absw_x, axis) sgn = math_ops.sign(sum_wx_over_max_absw_x) @@ -1180,8 +1178,7 @@ def process_quadrature_grid_and_probs( grid = ops.convert_to_tensor(grid, name="grid", dtype=dtype) probs = ops.convert_to_tensor(probs, name="unnormalized_probs", dtype=dtype) - probs /= linalg_ops.norm(probs, ord=1, axis=-1, keep_dims=True, - name="probs") + probs /= linalg_ops.norm(probs, ord=1, axis=-1, keepdims=True, name="probs") def _static_event_size(x): """Returns the static size of a specific dimension or `None`.""" diff --git a/tensorflow/python/ops/embedding_ops.py b/tensorflow/python/ops/embedding_ops.py index 3826585f59c31133b12c365816729e090c9ab561..20e4a28b9c0f5427e331a69cba52503492a8420a 100644 --- a/tensorflow/python/ops/embedding_ops.py +++ b/tensorflow/python/ops/embedding_ops.py @@ -396,8 +396,8 @@ def embedding_lookup_sparse(params, with `combiner`="mean", then the output will be a 3x20 matrix where output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5) - output[1, :] = params[0, :] * 1.0 - output[2, :] = params[1, :] * 3.0 + output[1, :] = (params[0, :] * 1.0) / 1.0 + output[2, :] = (params[1, :] * 3.0) / 3.0 Raises: TypeError: If sp_ids is not a SparseTensor, or if sp_weights is neither diff --git a/tensorflow/python/ops/functional_ops.py b/tensorflow/python/ops/functional_ops.py index ac03d30fcd2e65f032937d9259bc8fff18626619..a840b1eddfc6922dc310490e8166efd73480c437 100644 --- a/tensorflow/python/ops/functional_ops.py +++ b/tensorflow/python/ops/functional_ops.py @@ -41,7 +41,7 @@ from tensorflow.python.ops import variable_scope as vs from tensorflow.python.ops.gen_functional_ops import * # pylint: enable=wildcard-import # pylint: disable=unused-import -from tensorflow.python.ops.gen_functional_ops import _symbolic_gradient +from tensorflow.python.ops.gen_functional_ops import symbolic_gradient # pylint: enable=unused-import from tensorflow.python.util import nest from tensorflow.python.util.tf_export import tf_export @@ -90,7 +90,7 @@ def foldl(fn, elems, initializer=None, parallel_iterations=10, back_prop=True, if not callable(fn): raise TypeError("fn must be callable.") - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() with ops.name_scope(name, "foldl", [elems]): # TODO(akshayka): Remove the in_graph_mode check once caching devices are # supported in Eager @@ -178,7 +178,7 @@ def foldr(fn, elems, initializer=None, parallel_iterations=10, back_prop=True, if not callable(fn): raise TypeError("fn must be callable.") - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() with ops.name_scope(name, "foldr", [elems]): # TODO(akshayka): Remove the in_graph_mode check once caching devices are # supported in Eager @@ -343,7 +343,7 @@ def map_fn(fn, elems, dtype=None, parallel_iterations=10, back_prop=True, elems_flat = input_flatten(elems) - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() with ops.name_scope(name, "map", elems_flat): # TODO(akshayka): Remove the in_graph_mode check once caching devices are # supported in Eager @@ -364,8 +364,8 @@ def map_fn(fn, elems, dtype=None, parallel_iterations=10, back_prop=True, dtype = dtype or input_pack([elem.dtype for elem in elems_flat]) dtype_flat = output_flatten(dtype) - # Convert elems to tensor array. - n = array_ops.shape(elems_flat[0])[0] + # Convert elems to tensor array. n may be known statically. + n = elems_flat[0].shape[0].value or array_ops.shape(elems_flat[0])[0] # TensorArrays are always flat elems_ta = [ @@ -536,7 +536,7 @@ def scan(fn, elems, initializer=None, parallel_iterations=10, back_prop=True, elems_flat = input_flatten(elems) - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() with ops.name_scope(name, "scan", elems_flat): # TODO(akshayka): Remove the in_graph_mode check once caching devices are # supported in Eager @@ -555,7 +555,8 @@ def scan(fn, elems, initializer=None, parallel_iterations=10, back_prop=True, elems_flat = [ ops.convert_to_tensor(elem, name="elem") for elem in elems_flat] - n = array_ops.shape(elems_flat[0])[0] + # Convert elems to tensor array. n may be known statically. + n = elems_flat[0].shape[0].value or array_ops.shape(elems_flat[0])[0] # TensorArrays are always flat elems_ta = [ @@ -615,7 +616,8 @@ def scan(fn, elems, initializer=None, parallel_iterations=10, back_prop=True, _, _, r_a = control_flow_ops.while_loop( lambda i, _1, _2: i < n, compute, (i, a_flat, accs_ta), parallel_iterations=parallel_iterations, - back_prop=back_prop, swap_memory=swap_memory) + back_prop=back_prop, swap_memory=swap_memory, + maximum_iterations=n) results_flat = [r.stack() for r in r_a] diff --git a/tensorflow/python/ops/gradients.py b/tensorflow/python/ops/gradients.py index 921fd50aa9fdd1a1e493708f4bc8c66996e26e2c..63d9a232224cc70c9ce6f04787a8dcc921b3e2f3 100644 --- a/tensorflow/python/ops/gradients.py +++ b/tensorflow/python/ops/gradients.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function # pylint: disable=unused-import +from tensorflow.python.ops.custom_gradient import custom_gradient from tensorflow.python.ops.gradients_impl import AggregationMethod from tensorflow.python.ops.gradients_impl import gradients from tensorflow.python.ops.gradients_impl import hessians @@ -28,6 +29,7 @@ from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ # TODO(drpng): find a good place to reference this. "AggregationMethod", + "custom_gradient", "gradients", # tf.gradients.gradients. "hessians", # tf.gradients.hessians ] diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py index 230b6c5946ab7e0e3529994f7199026865cecde2..44473ec69c8ac6cf565f635621eebff7bc403225 100644 --- a/tensorflow/python/ops/gradients_impl.py +++ b/tensorflow/python/ops/gradients_impl.py @@ -35,7 +35,7 @@ from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_grad # pylint: disable=unused-import from tensorflow.python.ops import array_ops -from tensorflow.python.ops import check_ops +from tensorflow.python.ops import check_ops # pylint: disable=unused-import from tensorflow.python.ops import control_flow_grad # pylint: disable=unused-import from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import control_flow_util @@ -86,17 +86,19 @@ def _IndexedSlicesToTensor(value, dtype=None, name=None, as_ref=False): % str(value)) # TODO(mrry): Consider adding static shape information to # IndexedSlices, to avoid using numpy here. - dense_shape_value = tensor_util.constant_value(value.dense_shape) - if dense_shape_value is not None: - num_elements = np.prod(dense_shape_value) - if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: + if not context.executing_eagerly(): + dense_shape_value = tensor_util.constant_value(value.dense_shape) + if dense_shape_value is not None: + num_elements = np.prod(dense_shape_value) + if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: + warnings.warn( + "Converting sparse IndexedSlices to a dense Tensor with %d " + "elements. This may consume a large amount of memory." % + num_elements) + else: warnings.warn( - "Converting sparse IndexedSlices to a dense Tensor with %d elements. " - "This may consume a large amount of memory." % num_elements) - else: - warnings.warn( - "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " - "This may consume a large amount of memory.") + "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " + "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name) @@ -354,7 +356,7 @@ def _SymGrad(op, out_grads): for k in op.node_def.attr: f.attr[k].CopyFrom(op.node_def.attr[k]) # pylint: disable=protected-access - in_grads = functional_ops._symbolic_gradient(input=f_in, Tout=f_types, f=f) + in_grads = functional_ops.symbolic_gradient(input=f_in, Tout=f_types, f=f) # pylint: enable=protected-access return in_grads @@ -478,9 +480,21 @@ def gradients(ys, RuntimeError: if called in Eager mode. """ - if context.in_eager_mode(): - raise RuntimeError("tf.gradients not supported in EAGER mode. Use " - "functions in tf.contrib.eager.backprop instead.") + # Creating the gradient graph for control flow mutates Operations. _lock + # ensures a Session.run call cannot occur between creating and mutating new + # ops. + with ops.get_default_graph()._lock: # pylint: disable=protected-access + return _GradientsHelper(ys, xs, grad_ys, name, colocate_gradients_with_ops, + gate_gradients, aggregation_method, stop_gradients) + + +def _GradientsHelper(ys, xs, grad_ys, name, colocate_gradients_with_ops, + gate_gradients, aggregation_method, stop_gradients): + """Implementation of gradients().""" + if context.executing_eagerly(): + raise RuntimeError("tf.gradients not supported when eager execution " + "is enabled. Use tf.contrib.eager.GradientTape " + "instead.") ys = _AsList(ys) xs = _AsList(xs) stop_gradients = [] if stop_gradients is None else _AsList(stop_gradients) @@ -494,7 +508,7 @@ def gradients(ys, list(ys) + list(xs) + list(stop_gradients) + list(grad_ys)) as grad_scope: ys = ops.convert_n_to_tensor_or_indexed_slices(ys, name="y") xs = [ - x.handle if isinstance(x, resource_variable_ops.ResourceVariable) else x + x.handle if resource_variable_ops.is_resource_variable(x) else x for x in xs ] xs = ops.internal_convert_n_to_tensor_or_indexed_slices( diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py index d39b934819177e3c15af95a0777ba96869c5e9cf..c94f1396b28e2124c6e5123cf711ac86abf174ab 100644 --- a/tensorflow/python/ops/gradients_test.py +++ b/tensorflow/python/ops/gradients_test.py @@ -35,6 +35,7 @@ from tensorflow.python.ops import array_grad # pylint: disable=unused-import from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_grad # pylint: disable=unused-import from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import custom_gradient from tensorflow.python.ops import data_flow_grad # pylint: disable=unused-import from tensorflow.python.ops import data_flow_ops # pylint: disable=unused-import from tensorflow.python.ops import functional_ops # pylint: disable=unused-import @@ -661,6 +662,7 @@ class HessianTest(test_util.TensorFlowTestCase): self.assertAllEqual((m, n, m, n), hess_actual.shape) self.assertAllClose(hess_value, hess_actual.reshape((m * n, m * n))) + @test_util.with_c_api class IndexedSlicesToTensorTest(test_util.TensorFlowTestCase): @@ -741,6 +743,59 @@ class IndexedSlicesToTensorTest(test_util.TensorFlowTestCase): "of unknown shape. This may consume a large amount of memory." in str(w[0].message)) + def testCustomGradientTrivial(self): + + @custom_gradient.custom_gradient + def MyIdentity(x): + + def Grad(dy): + return [3 * dy] + + return x, Grad + + with ops.Graph().as_default(): + x = constant(3.) + y = MyIdentity(MyIdentity(x)) + dy = gradients.gradients(y, x)[0] + with session.Session(): + self.assertEqual(9., dy.eval()) + + def testCustomGradient(self): + + @custom_gradient.custom_gradient + def MyMultiply(x1, x2): + result = x1 * x2 + + def Grad(dy): + # Switched the ordering here. + return [dy * x1, dy * x2] + + return result, Grad + + with ops.Graph().as_default(): + x1 = constant(3.) + x2 = constant(5.) + y = MyMultiply(x1, x2) + dy = gradients.gradients(y, [x1, x2]) + with session.Session() as sess: + self.assertAllEqual([3., 5.], sess.run(dy)) + + def testCustomGradientErrors(self): + + @custom_gradient.custom_gradient + def F(x): + + def Grad(_): + raise RuntimeError("x") + + return x, Grad + + with ops.Graph().as_default(): + x = constant(1.0) + y = F(x) + with self.assertRaises(RuntimeError): + gradients.gradients(y, x) + @test_util.with_c_api class OnlyRealGradientsTest(test_util.TensorFlowTestCase): diff --git a/tensorflow/python/ops/hidden_ops.txt b/tensorflow/python/ops/hidden_ops.txt deleted file mode 100644 index f6ef6f3f3da4389a16a84fa0b3570d3cd1262472..0000000000000000000000000000000000000000 --- a/tensorflow/python/ops/hidden_ops.txt +++ /dev/null @@ -1,392 +0,0 @@ -# array_ops -BatchToSpace -BroadcastArgs -BroadcastGradientArgs -ConcatOffset -Concat -ConcatV2 -ConjugateTranspose -Const -DebugGradientIdentity -DebugGradientRefIdentity -EditDistance -ExpandDims -ListDiff -MirrorPad -MirrorPadGrad -OneHot -Pack -Pad -PadV2 -ParallelConcat -Placeholder -RefIdentity -Reverse -Snapshot -SpaceToBatch -Split -SplitV -Squeeze -Slice -TileGrad # Exported through array_grad instead of array_ops. -ZerosLike # TODO(josh11b): Use this instead of the Python version. -Unique -UniqueV2 -Unpack - -# candidate_sampling_ops -AllCandidateSampler -ComputeAccidentalHits -FixedUnigramCandidateSampler -LearnedUnigramCandidateSampler -LogUniformCandidateSampler -ThreadUnsafeUnigramCandidateSampler -UniformCandidateSampler - -# checkpoint_ops -GenerateVocabRemapping -LoadAndRemapMatrix - - -# control_flow_ops -Switch -Merge -RefMerge -Exit -RefExit - -# ctc_ops -CTCLoss -CTCGreedyDecoder -CTCBeamSearchDecoder - -# data_flow_ops -Barrier -BarrierClose -BarrierIncompleteSize -BarrierInsertMany -BarrierReadySize -BarrierTakeMany -DeleteSessionTensor -FakeQueue -FIFOQueue -FIFOQueueV2 -GetSessionHandle -GetSessionHandleV2 -GetSessionTensor -HashTable -HashTableV2 -InitializeTable -InitializeTableV2 -InitializeTableFromTextFile -InitializeTableFromTextFileV2 -LookupTableExport -LookupTableExportV2 -LookupTableFind -LookupTableFindV2 -LookupTableImport -LookupTableImportV2 -LookupTableInsert -LookupTableInsertV2 -LookupTableSize -LookupTableSizeV2 -MutableDenseHashTable -MutableDenseHashTableV2 -MutableHashTable -MutableHashTableV2 -MutableHashTableOfTensors -MutableHashTableOfTensorsV2 -Mutex -MutexAcquire -MutexRelease -PaddingFIFOQueue -PaddingFIFOQueueV2 -PriorityQueue -PriorityQueueV2 -QueueClose -QueueCloseV2 -QueueDequeue -QueueDequeueV2 -QueueDequeueMany -QueueDequeueManyV2 -QueueDequeueUpTo -QueueDequeueUpToV2 -QueueEnqueue -QueueEnqueueV2 -QueueEnqueueMany -QueueEnqueueManyV2 -QueueSize -QueueSizeV2 -RandomShuffleQueue -RandomShuffleQueueV2 -Stack -StackClose -StackPop -StackPush -StackV2 -StackCloseV2 -StackPopV2 -StackPushV2 -TensorArray -TensorArrayClose -TensorArrayCloseV2 -TensorArrayConcat -TensorArrayConcatV2 -TensorArrayGather -TensorArrayGatherV2 -TensorArrayGrad -TensorArrayGradV2 -TensorArrayPack -TensorArrayPackV2 -TensorArrayRead -TensorArrayReadV2 -TensorArrayScatter -TensorArrayScatterV2 -TensorArraySize -TensorArraySizeV2 -TensorArraySplit -TensorArraySplitV2 -TensorArrayUnpack -TensorArrayUnpackV2 -TensorArrayV2 -TensorArrayWrite -TensorArrayWriteV2 -TensorArrayV3 -TensorArrayCloseV3 -TensorArrayConcatV3 -TensorArrayGatherV3 -TensorArrayGradV3 -TensorArrayReadV3 -TensorArrayPackV3 -TensorArrayScatterV3 -TensorArraySizeV3 -TensorArraySplitV3 -TensorArrayUnpackV3 -TensorArrayWriteV3 - -# functional_ops -SymbolicGradient - -# image_ops -AdjustContrastv2 -NonMaxSuppression -NonMaxSuppressionV2 -RandomCrop -ResizeBilinearGrad -ResizeBicubicGrad -ResizeNearestNeighborGrad -SampleDistortedBoundingBox -SampleDistortedBoundingBoxV2 -ScaleImageGrad - -# io_ops -FixedLengthRecordReader -IdentityReader -ReaderNumRecordsProduced -ReaderNumWorkUnitsCompleted -ReaderRead -ReaderReadUpTo -ReaderReset -ReaderRestoreState -ReaderSerializeState -ReaderWorkQueueLength -FixedLengthRecordReaderV2 -IdentityReaderV2 -ReaderNumRecordsProducedV2 -ReaderNumWorkUnitsCompletedV2 -ReaderReadV2 -ReaderReadUpToV2 -ReaderResetV2 -ReaderRestoreStateV2 -ReaderSerializeStateV2 -ReaderWorkQueueLengthV2 -Restore -RestoreSlice -Save -SaveSlices -ShardedFilename -ShardedFilespec -TextLineReader -TFRecordReader -WholeFileReader -TextLineReaderV2 -TFRecordReaderV2 -WholeFileReaderV2 -LMDBReader -DecodeCSV - -# linalg_ops -BatchCholesky -BatchCholeskyGrad -BatchMatrixDeterminant -BatchMatrixInverse -BatchMatrixSolve -BatchMatrixSolveLs -BatchMatrixTriangularSolve -BatchSelfAdjointEig -BatchSelfAdjointEigV2 -BatchSvd -LogMatrixDeterminant -MatrixExponential -MatrixLogarithm -MatrixSolveLs -SelfAdjointEig -SelfAdjointEigV2 -Svd - -# logging_ops -Assert -AudioSummary -AudioSummaryV2 -HistogramSummary -ImageSummary -MergeSummary -Print -ScalarSummary -TensorSummary -TensorSummaryV2 - -# math_ops -Abs -AccumulateNV2 -AddN -AddV2 -All -Any -BatchMatMul -BatchFFT -BatchFFT2D -BatchFFT3D -BatchIFFT -BatchIFFT2D -BatchIFFT3D -Bucketize -Complex -ComplexAbs -Conj -FloorDiv -FloorMod -HistogramFixedWidth -Max -Mean -Min -Mul -Neg -Pow -Prod -Range -RealDiv -Select -SparseMatMul -Sub -Sum -MatMul -Sigmoid -Tanh -SigmoidGrad -TanhGrad -InvGrad -ReciprocalGrad -SqrtGrad -RsqrtGrad -TruncateDiv -TruncateMod - -# nn_ops -AvgPoolGrad # "*Grad" accessible through nn_grad instead of nn_ops. -AvgPool3DGrad -BatchNormWithGlobalNormalization -BatchNormWithGlobalNormalizationGrad -FusedBatchNorm -FusedBatchNormV2 -SoftmaxCrossEntropyWithLogits -SparseSoftmaxCrossEntropyWithLogits -LRNGrad -MaxPoolGrad -MaxPoolGradWithArgmax -MaxPoolGradGrad -MaxPoolGradGradWithArgmax -MaxPool3DGrad -MaxPool3DGradGrad -ReluGrad -Relu6Grad -EluGrad -SeluGrad -SoftplusGrad -SoftsignGrad -TopK -TopKV2 -BiasAdd -BiasAddV1 -Relu6 -AvgPool -MaxPool -MaxPoolV2 -Softmax -LogSoftmax -FractionalAvgPoolGrad -FractionalMaxPoolGrad -InTopK -InTopKV2 - -# parsing_ops -ParseExample -ParseSingleSequenceExample - -# random_ops -RandomGamma -RandomPoisson -RandomUniform -RandomUniformInt -RandomShuffle -RandomStandardNormal -ParameterizedTruncatedNormal -TruncatedNormal - -# script_ops -PyFunc -PyFuncStateless -EagerPyFunc - -# sdca_ops - -# state_ops -Variable -VariableV2 -TemporaryVariable -DestroyTemporaryVariable - -# sparse_ops -AddSparseToTensorsMap -AddManySparseToTensorsMap -TakeManySparseFromTensorsMap -DeserializeManySparse -DeserializeSparse -SerializeManySparse -SerializeSparse -SparseAdd -SparseAddGrad -SparseConcat -SparseCross -SparseFillEmptyRows -SparseFillEmptyRowsGrad -SparseSplit -SparseSelectLastK -SparseReorder -SparseReshape -SparseToDense -SparseTensorDenseAdd -SparseTensorDenseMatMul - -# string_ops -StringSplit - -# user_ops -Fact - -# training_ops -# (None) - -# word2vec deprecated ops -NegTrain -Skipgram diff --git a/tensorflow/python/ops/histogram_ops.py b/tensorflow/python/ops/histogram_ops.py index f079e56b10ed484225d8f09c6eaf7cf85a02d12a..4a1ef54fb50013881aa832f83674ac66ecccd9bc 100644 --- a/tensorflow/python/ops/histogram_ops.py +++ b/tensorflow/python/ops/histogram_ops.py @@ -32,8 +32,10 @@ from tensorflow.python.ops import clip_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import math_ops from tensorflow.python.util.tf_export import tf_export +from tensorflow.python.util.tf_export import tf_export +@tf_export('histogram_fixed_width_bins') def histogram_fixed_width_bins(values, value_range, nbins=100, @@ -139,5 +141,7 @@ def histogram_fixed_width(values, """ with ops.name_scope(name, 'histogram_fixed_width', [values, value_range, nbins]) as name: - return gen_math_ops._histogram_fixed_width( # pylint: disable=protected-access + # pylint: disable=protected-access + return gen_math_ops._histogram_fixed_width( values, value_range, nbins, dtype=dtype, name=name) + # pylint: enable=protected-access diff --git a/tensorflow/python/ops/image_grad.py b/tensorflow/python/ops/image_grad.py index d17f1a87d9759d5e83393f40e9e027dee8c15979..9f43e3f1466d900ae6d39f3b9ef48043421cb777 100644 --- a/tensorflow/python/ops/image_grad.py +++ b/tensorflow/python/ops/image_grad.py @@ -41,12 +41,10 @@ def _ResizeNearestNeighborGrad(op, grad): else: image_shape = array_ops.shape(image)[1:3] - # pylint: disable=protected-access - grads = gen_image_ops._resize_nearest_neighbor_grad( + grads = gen_image_ops.resize_nearest_neighbor_grad( grad, image_shape, align_corners=op.get_attr("align_corners")) - # pylint: enable=protected-access return [grads, None] @@ -61,15 +59,8 @@ def _ResizeBilinearGrad(op, grad): Returns: The gradients w.r.t. the input. """ - allowed_types = [dtypes.float32, dtypes.float64] - grad0 = None - if op.inputs[0].dtype in allowed_types: - # pylint: disable=protected-access - grad0 = gen_image_ops._resize_bilinear_grad( - grad, - op.inputs[0], - align_corners=op.get_attr("align_corners")) - # pylint: enable=protected-access + grad0 = gen_image_ops.resize_bilinear_grad( + grad, op.inputs[0], align_corners=op.get_attr("align_corners")) return [grad0, None] @@ -87,10 +78,8 @@ def _ResizeBicubicGrad(op, grad): allowed_types = [dtypes.float32, dtypes.float64] grad0 = None if op.inputs[0].dtype in allowed_types: - # pylint: disable=protected-access - grad0 = gen_image_ops._resize_bicubic_grad( + grad0 = gen_image_ops.resize_bicubic_grad( grad, op.inputs[0], align_corners=op.get_attr("align_corners")) - # pylint: enable=protected-access return [grad0, None] diff --git a/tensorflow/python/ops/image_grad_test.py b/tensorflow/python/ops/image_grad_test.py index 05e8fa1d72851caee522bba470bb40f430152464..75d00c8ed17c26c2c1acb4d92961a2206d959ebb 100644 --- a/tensorflow/python/ops/image_grad_test.py +++ b/tensorflow/python/ops/image_grad_test.py @@ -142,18 +142,6 @@ class ResizeBilinearOpTest(test.TestCase): input_tensor, in_shape, resize_out, out_shape, x_init_value=x) self.assertLess(err, 1e-3) - def testGradOnUnsupportedType(self): - in_shape = [1, 4, 6, 1] - out_shape = [1, 2, 3, 1] - - x = np.arange(0, 24).reshape(in_shape).astype(np.uint8) - - with self.test_session(): - input_tensor = constant_op.constant(x, shape=in_shape) - resize_out = image_ops.resize_bilinear(input_tensor, out_shape[1:3]) - grad = gradients_impl.gradients(input_tensor, [resize_out]) - self.assertEqual([None], grad) - def testCompareGpuVsCpu(self): in_shape = [2, 4, 6, 3] out_shape = [2, 8, 16, 3] @@ -172,6 +160,26 @@ class ResizeBilinearOpTest(test.TestCase): self.assertAllClose(grad[False], grad[True], rtol=1e-4, atol=1e-4) + def testTypes(self): + in_shape = [1, 4, 6, 1] + out_shape = [1, 2, 3, 1] + x = np.arange(0, 24).reshape(in_shape) + + with self.test_session() as sess: + for dtype in [np.float16, np.float32, np.float64]: + input_tensor = constant_op.constant(x.astype(dtype), shape=in_shape) + resize_out = image_ops.resize_bilinear(input_tensor, out_shape[1:3]) + grad = sess.run(gradients_impl.gradients(resize_out, input_tensor))[0] + self.assertAllEqual(in_shape, grad.shape) + # Not using gradient_checker.compute_gradient as I didn't work out + # the changes required to compensate for the lower precision of + # float16 when computing the numeric jacobian. + # Instead, we just test the theoretical jacobian. + self.assertAllEqual([[[[1.], [0.], [1.], [0.], [1.], [0.]], [[0.], [ + 0. + ], [0.], [0.], [0.], [0.]], [[1.], [0.], [1.], [0.], [1.], [0.]], + [[0.], [0.], [0.], [0.], [0.], [0.]]]], grad) + class ResizeBicubicOpTest(test.TestCase): diff --git a/tensorflow/python/ops/image_ops.py b/tensorflow/python/ops/image_ops.py index de12c5f63f4357e0982dd2e16999caf2de0b30f8..68be9ccdd642823e7a9c2294f209accd16f45be5 100644 --- a/tensorflow/python/ops/image_ops.py +++ b/tensorflow/python/ops/image_ops.py @@ -26,6 +26,7 @@ See the @{$python/image} guide. @@extract_jpeg_shape @@decode_png @@encode_png +@@is_jpeg @@decode_image @@resize_images @@resize_area @@ -68,6 +69,11 @@ See the @{$python/image} guide. @@non_max_suppression @@sample_distorted_bounding_box @@total_variation +@@psnr +@@ssim +@@ssim_multiscale +@@image_gradients +@@sobel_edges """ from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py index cab1025df11d26064c3d2939598f1c58ab104736..3369fe3c9b37ca05311c5548dbfa3228ba04ee80 100644 --- a/tensorflow/python/ops/image_ops_impl.py +++ b/tensorflow/python/ops/image_ops_impl.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.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -29,6 +31,8 @@ from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_image_ops from tensorflow.python.ops import gen_nn_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 import string_ops from tensorflow.python.ops import variables @@ -167,6 +171,28 @@ def _Assert3DImage(image): _Check3DImage(image, require_static=False), image) +def _AssertAtLeast3DImage(image): + """Assert that we are working with a properly shaped image. + + Performs the check statically if possible (i.e. if the shape + is statically known). Otherwise adds a control dependency + to an assert op that checks the dynamic shape. + + Args: + image: >= 3-D Tensor of size [*, height, width, depth] + + Raises: + ValueError: if image.shape is not a [>= 3] vector. + + Returns: + If the shape of `image` could be verified statically, `image` is + returned unchanged, otherwise there will be a control dependency + added that asserts the correct dynamic shape. + """ + return control_flow_ops.with_dependencies( + _CheckAtLeast3DImage(image, require_static=False), image) + + def _CheckAtLeast3DImage(image, require_static=True): """Assert that we are working with properly shaped image. @@ -292,108 +318,185 @@ def random_flip_left_right(image, seed=None): def flip_left_right(image): """Flip an image horizontally (left to right). - Outputs the contents of `image` flipped along the second dimension, which is - `width`. + Outputs the contents of `image` flipped along the width dimension. See also `reverse()`. Args: - image: A 3-D tensor of shape `[height, width, channels].` + image: 4-D Tensor of shape `[batch, height, width, channels]` or + 3-D Tensor of shape `[height, width, channels]`. Returns: - A 3-D tensor of the same type and shape as `image`. + A tensor of the same type and shape as `image`. Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'flip_left_right', [image]) as scope: + with ops.name_scope(None, 'flip_left_right', [image]): image = ops.convert_to_tensor(image, name='image') - image = _Assert3DImage(image) - return fix_image_flip_shape(image, array_ops.reverse( - image, [1], name=scope)) + image = _AssertAtLeast3DImage(image) + shape = image.get_shape() + if shape.ndims == 3 or shape.ndims is None: + return fix_image_flip_shape(image, array_ops.reverse(image, [1])) + elif shape.ndims == 4: + return array_ops.reverse(image, [2]) + else: + raise ValueError('\'image\' must have either 3 or 4 dimensions.') @tf_export('image.flip_up_down') def flip_up_down(image): """Flip an image vertically (upside down). - Outputs the contents of `image` flipped along the first dimension, which is - `height`. + Outputs the contents of `image` flipped along the height dimension. See also `reverse()`. Args: - image: A 3-D tensor of shape `[height, width, channels].` + image: 4-D Tensor of shape `[batch, height, width, channels]` or + 3-D Tensor of shape `[height, width, channels]`. Returns: - A 3-D tensor of the same type and shape as `image`. + A tensor of the same type and shape as `image`. Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'flip_up_down', [image]) as scope: + with ops.name_scope(None, 'flip_up_down', [image]): image = ops.convert_to_tensor(image, name='image') - image = _Assert3DImage(image) - return fix_image_flip_shape(image, array_ops.reverse( - image, [0], name=scope)) + image = _AssertAtLeast3DImage(image) + shape = image.get_shape() + if shape.ndims == 3 or shape.ndims is None: + return fix_image_flip_shape(image, array_ops.reverse(image, [0])) + elif shape.ndims == 4: + return array_ops.reverse(image, [1]) + else: + raise ValueError('\'image\' must have either 3 or 4 dimensions.') @tf_export('image.rot90') def rot90(image, k=1, name=None): - """Rotate an image counter-clockwise by 90 degrees. + """Rotate image(s) counter-clockwise by 90 degrees. Args: - image: A 3-D tensor of shape `[height, width, channels]`. + image: 4-D Tensor of shape `[batch, height, width, channels]` or + 3-D Tensor of shape `[height, width, channels]`. k: A scalar integer. The number of times the image is rotated by 90 degrees. name: A name for this operation (optional). Returns: - A rotated 3-D tensor of the same type and shape as `image`. + A rotated tensor of the same type and shape as `image`. + + Raises: + ValueError: if the shape of `image` not supported. """ with ops.name_scope(name, 'rot90', [image, k]) as scope: image = ops.convert_to_tensor(image, name='image') - image = _Assert3DImage(image) + image = _AssertAtLeast3DImage(image) k = ops.convert_to_tensor(k, dtype=dtypes.int32, name='k') k.get_shape().assert_has_rank(0) k = math_ops.mod(k, 4) - def _rot90(): - return array_ops.transpose(array_ops.reverse_v2(image, [1]), [1, 0, 2]) + shape = image.get_shape() + if shape.ndims == 3 or shape.ndims is None: + return _rot90_3D(image, k, scope) + elif shape.ndims == 4: + return _rot90_4D(image, k, scope) + else: + raise ValueError('\'image\' must have either 3 or 4 dimensions.') - def _rot180(): - return array_ops.reverse_v2(image, [0, 1]) - def _rot270(): - return array_ops.reverse_v2(array_ops.transpose(image, [1, 0, 2]), [1]) +def _rot90_3D(image, k, name_scope): + """Rotate image counter-clockwise by 90 degrees `k` times. - cases = [(math_ops.equal(k, 1), _rot90), (math_ops.equal(k, 2), _rot180), - (math_ops.equal(k, 3), _rot270)] + Args: + image: 3-D Tensor of shape `[height, width, channels]`. + k: A scalar integer. The number of times the image is rotated by 90 degrees. + name_scope: A valid TensorFlow name scope. + + Returns: + A 3-D tensor of the same type and shape as `image`. + + """ + + def _rot90(): + return array_ops.transpose(array_ops.reverse_v2(image, [1]), [1, 0, 2]) + + def _rot180(): + return array_ops.reverse_v2(image, [0, 1]) + + def _rot270(): + return array_ops.reverse_v2(array_ops.transpose(image, [1, 0, 2]), [1]) + + cases = [(math_ops.equal(k, 1), _rot90), (math_ops.equal(k, 2), _rot180), + (math_ops.equal(k, 3), _rot270)] + + result = control_flow_ops.case( + cases, default=lambda: image, exclusive=True, name=name_scope) + result.set_shape([None, None, image.get_shape()[2]]) + return result + + +def _rot90_4D(images, k, name_scope): + """Rotate batch of images counter-clockwise by 90 degrees `k` times. + + Args: + images: 4-D Tensor of shape `[height, width, channels]`. + k: A scalar integer. The number of times the images are rotated by 90 + degrees. + name_scope: A valid TensorFlow name scope. + + Returns: + A 4-D tensor of the same type and shape as `images`. + + """ + + def _rot90(): + return array_ops.transpose(array_ops.reverse_v2(images, [2]), [0, 2, 1, 3]) - ret = control_flow_ops.case( - cases, default=lambda: image, exclusive=True, name=scope) - ret.set_shape([None, None, image.get_shape()[2]]) - return ret + def _rot180(): + return array_ops.reverse_v2(images, [1, 2]) + def _rot270(): + return array_ops.reverse_v2(array_ops.transpose(images, [0, 2, 1, 3]), [2]) + cases = [(math_ops.equal(k, 1), _rot90), (math_ops.equal(k, 2), _rot180), + (math_ops.equal(k, 3), _rot270)] + + result = control_flow_ops.case( + cases, default=lambda: images, exclusive=True, name=name_scope) + shape = result.get_shape() + result.set_shape([shape[0], None, None, shape[3]]) + return result @tf_export('image.transpose_image') def transpose_image(image): - """Transpose an image by swapping the first and second dimension. + """Transpose image(s) by swapping the height and width dimension. See also `transpose()`. Args: - image: 3-D tensor of shape `[height, width, channels]` + image: 4-D Tensor of shape `[batch, height, width, channels]` or + 3-D Tensor of shape `[height, width, channels]`. Returns: - A 3-D tensor of shape `[width, height, channels]` + If `image` was 4-D, a 4-D float Tensor of shape + `[batch, width, height, channels]` + If `image` was 3-D, a 3-D float Tensor of shape + `[width, height, channels]` Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'transpose_image', [image]) as scope: + with ops.name_scope(None, 'transpose_image', [image]): image = ops.convert_to_tensor(image, name='image') - image = _Assert3DImage(image) - return array_ops.transpose(image, [1, 0, 2], name=scope) + image = _AssertAtLeast3DImage(image) + shape = image.get_shape() + if shape.ndims == 3 or shape.ndims is None: + return array_ops.transpose(image, [1, 0, 2], name='transpose_image') + elif shape.ndims == 4: + return array_ops.transpose(image, [0, 2, 1, 3], name='transpose_image') + else: + raise ValueError('\'image\' must have either 3 or 4 dimensions.') @tf_export('image.central_crop') @@ -1014,10 +1117,8 @@ def adjust_contrast(images, contrast_factor): orig_dtype = images.dtype flt_images = convert_image_dtype(images, dtypes.float32) - # pylint: disable=protected-access - adjusted = gen_image_ops._adjust_contrastv2( + adjusted = gen_image_ops.adjust_contrastv2( flt_images, contrast_factor=contrast_factor, name=name) - # pylint: enable=protected-access return convert_image_dtype(adjusted, orig_dtype, saturate=True) @@ -1026,9 +1127,9 @@ def adjust_contrast(images, contrast_factor): def adjust_gamma(image, gamma=1, gain=1): """Performs Gamma Correction on the input image. - Also known as Power Law Transform. This function transforms the - input image pixelwise according to the equation Out = In**gamma - after scaling each pixel to the range 0 to 1. + Also known as Power Law Transform. This function transforms the + input image pixelwise according to the equation `Out = In**gamma` + after scaling each pixel to the range 0 to 1. Args: image : A Tensor. @@ -1339,6 +1440,26 @@ def adjust_saturation(image, saturation_factor, name=None): orig_dtype) +@tf_export('image.is_jpeg') +def is_jpeg(contents, name=None): + r"""Convenience function to check if the 'contents' encodes a JPEG image. + + Args: + contents: 0-D `string`. The encoded image bytes. + name: A name for the operation (optional) + + Returns: + A scalar boolean tensor indicating if 'contents' may be a JPEG image. + is_jpeg is susceptible to false positives. + """ + # Normal JPEGs start with \xff\xd8\xff\xe0 + # JPEG with EXIF stats with \xff\xd8\xff\xe1 + # Use \xff\xd8\xff to cover both. + with ops.name_scope(name, 'is_jpeg'): + substr = string_ops.substr(contents, 0, 3) + return math_ops.equal(substr, b'\xff\xd8\xff', name=name) + + @tf_export('image.decode_image') def decode_image(contents, channels=None, name=None): """Convenience function for `decode_bmp`, `decode_gif`, `decode_jpeg`, @@ -1427,8 +1548,8 @@ def decode_image(contents, channels=None, name=None): # Decode normal JPEG images (start with \xff\xd8\xff\xe0) # as well as JPEG images with EXIF data (start with \xff\xd8\xff\xe1). - is_jpeg = math_ops.equal(substr, b'\xff\xd8\xff', name='is_jpeg') - return control_flow_ops.cond(is_jpeg, _jpeg, check_png, name='cond_jpeg') + return control_flow_ops.cond( + is_jpeg(contents), _jpeg, check_png, name='cond_jpeg') @tf_export('image.total_variation') @@ -1611,7 +1732,7 @@ def sample_distorted_bounding_box(image_size, Provide as input to `tf.image.draw_bounding_boxes`. """ with ops.name_scope(name, 'sample_distorted_bounding_box'): - return gen_image_ops._sample_distorted_bounding_box_v2( # pylint: disable=protected-access + return gen_image_ops.sample_distorted_bounding_box_v2( image_size, bounding_boxes, seed=seed, @@ -1665,17 +1786,16 @@ def non_max_suppression(boxes, """ with ops.name_scope(name, 'non_max_suppression'): iou_threshold = ops.convert_to_tensor(iou_threshold, name='iou_threshold') - # pylint: disable=protected-access - return gen_image_ops._non_max_suppression_v2(boxes, scores, max_output_size, - iou_threshold) - # pylint: enable=protected-access + return gen_image_ops.non_max_suppression_v2(boxes, scores, max_output_size, + iou_threshold) -_rgb_to_yiq_kernel = [[0.299, 0.59590059, 0.2115], - [0.587, -0.27455667, -0.52273617], +_rgb_to_yiq_kernel = [[0.299, 0.59590059, + 0.2115], [0.587, -0.27455667, -0.52273617], [0.114, -0.32134392, 0.31119955]] +@tf_export('image.rgb_to_yiq') def rgb_to_yiq(images): """Converts one or more images from RGB to YIQ. @@ -1691,16 +1811,17 @@ def rgb_to_yiq(images): images: tensor with the same shape as `images`. """ images = ops.convert_to_tensor(images, name='images') - kernel = ops.convert_to_tensor(_rgb_to_yiq_kernel, dtype=images.dtype, name='kernel') + kernel = ops.convert_to_tensor( + _rgb_to_yiq_kernel, dtype=images.dtype, name='kernel') ndims = images.get_shape().ndims - return math_ops.tensordot(images, kernel, axes=[[ndims-1], [0]]) + return math_ops.tensordot(images, kernel, axes=[[ndims - 1], [0]]) -_yiq_to_rgb_kernel = [[1, 1, 1], - [0.95598634, -0.27201283, -1.10674021], +_yiq_to_rgb_kernel = [[1, 1, 1], [0.95598634, -0.27201283, -1.10674021], [0.6208248, -0.64720424, 1.70423049]] +@tf_export('image.yiq_to_rgb') def yiq_to_rgb(images): """Converts one or more images from YIQ to RGB. @@ -1717,16 +1838,18 @@ def yiq_to_rgb(images): images: tensor with the same shape as `images`. """ images = ops.convert_to_tensor(images, name='images') - kernel = ops.convert_to_tensor(_yiq_to_rgb_kernel, dtype=images.dtype, name='kernel') + kernel = ops.convert_to_tensor( + _yiq_to_rgb_kernel, dtype=images.dtype, name='kernel') ndims = images.get_shape().ndims - return math_ops.tensordot(images, kernel, axes=[[ndims-1], [0]]) + return math_ops.tensordot(images, kernel, axes=[[ndims - 1], [0]]) -_rgb_to_yuv_kernel = [[0.299, -0.14714119, 0.61497538], - [0.587, -0.28886916, -0.51496512], +_rgb_to_yuv_kernel = [[0.299, -0.14714119, + 0.61497538], [0.587, -0.28886916, -0.51496512], [0.114, 0.43601035, -0.10001026]] +@tf_export('image.rgb_to_yuv') def rgb_to_yuv(images): """Converts one or more images from RGB to YUV. @@ -1742,16 +1865,17 @@ def rgb_to_yuv(images): images: tensor with the same shape as `images`. """ images = ops.convert_to_tensor(images, name='images') - kernel = ops.convert_to_tensor(_rgb_to_yuv_kernel, dtype=images.dtype, name='kernel') + kernel = ops.convert_to_tensor( + _rgb_to_yuv_kernel, dtype=images.dtype, name='kernel') ndims = images.get_shape().ndims - return math_ops.tensordot(images, kernel, axes=[[ndims-1], [0]]) + return math_ops.tensordot(images, kernel, axes=[[ndims - 1], [0]]) -_yuv_to_rgb_kernel = [[1, 1, 1], - [0, -0.394642334, 2.03206185], +_yuv_to_rgb_kernel = [[1, 1, 1], [0, -0.394642334, 2.03206185], [1.13988303, -0.58062185, 0]] +@tf_export('image.yuv_to_rgb') def yuv_to_rgb(images): """Converts one or more images from YUV to RGB. @@ -1768,7 +1892,493 @@ def yuv_to_rgb(images): images: tensor with the same shape as `images`. """ images = ops.convert_to_tensor(images, name='images') - kernel = ops.convert_to_tensor(_yuv_to_rgb_kernel, dtype=images.dtype, name='kernel') + kernel = ops.convert_to_tensor( + _yuv_to_rgb_kernel, dtype=images.dtype, name='kernel') ndims = images.get_shape().ndims - return math_ops.tensordot(images, kernel, axes=[[ndims-1], [0]]) + return math_ops.tensordot(images, kernel, axes=[[ndims - 1], [0]]) + +def _verify_compatible_image_shapes(img1, img2): + """Checks if two image tensors are compatible for applying SSIM or PSNR. + + This function checks if two sets of images have ranks at least 3, and if the + last three dimensions match. + + Args: + img1: Tensor containing the first image batch. + img2: Tensor containing the second image batch. + + Returns: + A tuple containing: the first tensor shape, the second tensor shape, and a + list of control_flow_ops.Assert() ops implementing the checks. + + Raises: + ValueError: When static shape check fails. + """ + shape1 = img1.get_shape().with_rank_at_least(3) + shape2 = img2.get_shape().with_rank_at_least(3) + shape1[-3:].assert_is_compatible_with(shape2[-3:]) + + if shape1.ndims is not None and shape2.ndims is not None: + for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])): + if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)): + raise ValueError( + 'Two images are not compatible: %s and %s' % (shape1, shape2)) + + # Now assign shape tensors. + shape1, shape2 = array_ops.shape_n([img1, img2]) + + # TODO(sjhwang): Check if shape1[:-3] and shape2[:-3] are broadcastable. + checks = [] + checks.append(control_flow_ops.Assert( + math_ops.greater_equal(array_ops.size(shape1), 3), + [shape1, shape2], summarize=10)) + checks.append(control_flow_ops.Assert( + math_ops.reduce_all(math_ops.equal(shape1[-3:], shape2[-3:])), + [shape1, shape2], summarize=10)) + return shape1, shape2, checks + + +@tf_export('image.psnr') +def psnr(a, b, max_val, name=None): + """Returns the Peak Signal-to-Noise Ratio between a and b. + + This is intended to be used on signals (or images). Produces a PSNR value for + each image in batch. + + The last three dimensions of input are expected to be [height, width, depth]. + + Example: + + ```python + # Read images from file. + im1 = tf.decode_png('path/to/im1.png') + im2 = tf.decode_png('path/to/im2.png') + # Compute PSNR over tf.uint8 Tensors. + psnr1 = tf.image.psnr(im1, im2, max_val=255) + + # Compute PSNR over tf.float32 Tensors. + im1 = tf.image.convert_image_dtype(im1, tf.float32) + im2 = tf.image.convert_image_dtype(im2, tf.float32) + psnr2 = tf.image.psnr(im1, im2, max_val=1.0) + # psnr1 and psnr2 both have type tf.float32 and are almost equal. + ``` + + Arguments: + a: First set of images. + b: Second set of images. + max_val: The dynamic range of the images (i.e., the difference between the + maximum the and minimum allowed values). + name: Namespace to embed the computation in. + + Returns: + The scalar PSNR between a and b. The returned tensor has type `tf.float32` + and shape [batch_size, 1]. + """ + with ops.name_scope(name, 'PSNR', [a, b]): + # Need to convert the images to float32. Scale max_val accordingly so that + # PSNR is computed correctly. + max_val = math_ops.cast(max_val, a.dtype) + max_val = convert_image_dtype(max_val, dtypes.float32) + a = convert_image_dtype(a, dtypes.float32) + b = convert_image_dtype(b, dtypes.float32) + mse = math_ops.reduce_mean(math_ops.squared_difference(a, b), [-3, -2, -1]) + psnr_val = math_ops.subtract( + 20 * math_ops.log(max_val) / math_ops.log(10.0), + np.float32(10 / np.log(10)) * math_ops.log(mse), + name='psnr') + + _, _, checks = _verify_compatible_image_shapes(a, b) + with ops.control_dependencies(checks): + return array_ops.identity(psnr_val) + +_SSIM_K1 = 0.01 +_SSIM_K2 = 0.03 + + +def _ssim_helper(x, y, reducer, max_val, compensation=1.0): + r"""Helper function for computing SSIM. + + SSIM estimates covariances with weighted sums. The default parameters + use a biased estimate of the covariance: + Suppose `reducer` is a weighted sum, then the mean estimators are + \mu_x = \sum_i w_i x_i, + \mu_y = \sum_i w_i y_i, + where w_i's are the weighted-sum weights, and covariance estimator is + cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y) + with assumption \sum_i w_i = 1. This covariance estimator is biased, since + E[cov_{xy}] = (1 - \sum_i w_i ^ 2) Cov(X, Y). + For SSIM measure with unbiased covariance estimators, pass as `compensation` + argument (1 - \sum_i w_i ^ 2). + + Arguments: + x: First set of images. + y: Second set of images. + reducer: Function that computes 'local' averages from set of images. + For non-covolutional version, this is usually tf.reduce_mean(x, [1, 2]), + and for convolutional version, this is usually tf.nn.avg_pool or + tf.nn.conv2d with weighted-sum kernel. + max_val: The dynamic range (i.e., the difference between the maximum + possible allowed value and the minimum allowed value). + compensation: Compensation factor. See above. + + Returns: + A pair containing the luminance measure, and the contrast-structure measure. + """ + c1 = (_SSIM_K1 * max_val) ** 2 + c2 = (_SSIM_K2 * max_val) ** 2 + + # SSIM luminance measure is + # (2 * mu_x * mu_y + c1) / (mu_x ** 2 + mu_y ** 2 + c1). + mean0 = reducer(x) + mean1 = reducer(y) + num0 = mean0 * mean1 * 2.0 + den0 = math_ops.square(mean0) + math_ops.square(mean1) + luminance = (num0 + c1) / (den0 + c1) + + # SSIM contrast-structure measure is + # (2 * cov_{xy} + c2) / (cov_{xx} + cov_{yy} + c2). + # Note that `reducer` is a weighted sum with weight w_k, \sum_i w_i = 1, then + # cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y) + # = \sum_i w_i x_i y_i - (\sum_i w_i x_i) (\sum_j w_j y_j). + num1 = reducer(x * y) * 2.0 + den1 = reducer(math_ops.square(x) + math_ops.square(y)) + c2 *= compensation + cs = (num1 - num0 + c2) / (den1 - den0 + c2) + + # SSIM score is the product of the luminance and contrast-structure measures. + return luminance, cs + + +def _fspecial_gauss(size, sigma): + """Function to mimic the 'fspecial' gaussian MATLAB function.""" + size = ops.convert_to_tensor(size, dtypes.int32) + sigma = ops.convert_to_tensor(sigma) + + coords = math_ops.cast(math_ops.range(size), sigma.dtype) + coords -= math_ops.cast(size - 1, sigma.dtype) / 2.0 + + g = math_ops.square(coords) + g *= -0.5 / math_ops.square(sigma) + + g = array_ops.reshape(g, shape=[1, -1]) + array_ops.reshape(g, shape=[-1, 1]) + g = array_ops.reshape(g, shape=[1, -1]) # For tf.nn.softmax(). + g = nn_ops.softmax(g) + return array_ops.reshape(g, shape=[size, size, 1, 1]) + + +def _ssim_per_channel(img1, img2, max_val=1.0): + """Computes SSIM index between img1 and img2 per color channel. + + This function matches the standard SSIM implementation from: + Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image + quality assessment: from error visibility to structural similarity. IEEE + transactions on image processing. + + Details: + - 11x11 Gaussian filter of width 1.5 is used. + - k1 = 0.01, k2 = 0.03 as in the original paper. + + Args: + img1: First image batch. + img2: Second image batch. + max_val: The dynamic range of the images (i.e., the difference between the + maximum the and minimum allowed values). + + Returns: + A pair of tensors containing and channel-wise SSIM and contrast-structure + values. The shape is [..., channels]. + """ + filter_size = constant_op.constant(11, dtype=dtypes.int32) + filter_sigma = constant_op.constant(1.5, dtype=img1.dtype) + + shape1, shape2 = array_ops.shape_n([img1, img2]) + checks = [ + control_flow_ops.Assert(math_ops.reduce_all(math_ops.greater_equal( + shape1[-3:-1], filter_size)), [shape1, filter_size], summarize=8), + control_flow_ops.Assert(math_ops.reduce_all(math_ops.greater_equal( + shape2[-3:-1], filter_size)), [shape2, filter_size], summarize=8)] + + # Enforce the check to run before computation. + with ops.control_dependencies(checks): + img1 = array_ops.identity(img1) + + # TODO(sjhwang): Try to cache kernels and compensation factor. + kernel = _fspecial_gauss(filter_size, filter_sigma) + kernel = array_ops.tile(kernel, multiples=[1, 1, shape1[-1], 1]) + + # The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`, + # but to match MATLAB implementation of MS-SSIM, we use 1.0 instead. + compensation = 1.0 + + # TODO(sjhwang): Try FFT. + # TODO(sjhwang): Gaussian kernel is separable in space. Consider applying + # 1-by-n and n-by-1 Gaussain filters instead of an n-by-n filter. + def reducer(x): + shape = array_ops.shape(x) + x = array_ops.reshape(x, shape=array_ops.concat([[-1], shape[-3:]], 0)) + y = nn.depthwise_conv2d(x, kernel, strides=[1, 1, 1, 1], padding='VALID') + return array_ops.reshape(y, array_ops.concat([shape[:-3], + array_ops.shape(y)[1:]], 0)) + + luminance, cs = _ssim_helper(img1, img2, reducer, max_val, compensation) + + # Average over the second and the third from the last: height, width. + axes = constant_op.constant([-3, -2], dtype=dtypes.int32) + ssim_val = math_ops.reduce_mean(luminance * cs, axes) + cs = math_ops.reduce_mean(cs, axes) + return ssim_val, cs + + +@tf_export('image.ssim') +def ssim(img1, img2, max_val): + """Computes SSIM index between img1 and img2. + + This function is based on the standard SSIM implementation from: + Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image + quality assessment: from error visibility to structural similarity. IEEE + transactions on image processing. + + Note: The true SSIM is only defined on grayscale. This function does not + perform any colorspace transform. (If input is already YUV, then it will + compute YUV SSIM average.) + + Details: + - 11x11 Gaussian filter of width 1.5 is used. + - k1 = 0.01, k2 = 0.03 as in the original paper. + + The image sizes must be at least 11x11 because of the filter size. + + Example: + + ```python + # Read images from file. + im1 = tf.decode_png('path/to/im1.png') + im2 = tf.decode_png('path/to/im2.png') + # Compute SSIM over tf.uint8 Tensors. + ssim1 = tf.image.ssim(im1, im2, max_val=255) + + # Compute SSIM over tf.float32 Tensors. + im1 = tf.image.convert_image_dtype(im1, tf.float32) + im2 = tf.image.convert_image_dtype(im2, tf.float32) + ssim2 = tf.image.ssim(im1, im2, max_val=1.0) + # ssim1 and ssim2 both have type tf.float32 and are almost equal. + ``` + + Args: + img1: First image batch. + img2: Second image batch. + max_val: The dynamic range of the images (i.e., the difference between the + maximum the and minimum allowed values). + + Returns: + A tensor containing an SSIM value for each image in batch. Returned SSIM + values are in range (-1, 1], when pixel values are non-negative. Returns + a tensor with shape: broadcast(img1.shape[:-3], img2.shape[:-3]). + """ + _, _, checks = _verify_compatible_image_shapes(img1, img2) + with ops.control_dependencies(checks): + img1 = array_ops.identity(img1) + + # Need to convert the images to float32. Scale max_val accordingly so that + # SSIM is computed correctly. + max_val = math_ops.cast(max_val, img1.dtype) + max_val = convert_image_dtype(max_val, dtypes.float32) + img1 = convert_image_dtype(img1, dtypes.float32) + img2 = convert_image_dtype(img2, dtypes.float32) + ssim_per_channel, _ = _ssim_per_channel(img1, img2, max_val) + # Compute average over color channels. + return math_ops.reduce_mean(ssim_per_channel, [-1]) + + +# Default values obtained by Wang et al. +_MSSSIM_WEIGHTS = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) + + +@tf_export('image.ssim_multiscale') +def ssim_multiscale(img1, img2, max_val, power_factors=_MSSSIM_WEIGHTS): + """Computes the MS-SSIM between img1 and img2. + + This function assumes that `img1` and `img2` are image batches, i.e. the last + three dimensions are [height, width, channels]. + + Note: The true SSIM is only defined on grayscale. This function does not + perform any colorspace transform. (If input is already YUV, then it will + compute YUV SSIM average.) + + Original paper: Wang, Zhou, Eero P. Simoncelli, and Alan C. Bovik. "Multiscale + structural similarity for image quality assessment." Signals, Systems and + Computers, 2004. + + Arguments: + img1: First image batch. + img2: Second image batch. Must have the same rank as img1. + max_val: The dynamic range of the images (i.e., the difference between the + maximum the and minimum allowed values). + power_factors: Iterable of weights for each of the scales. The number of + scales used is the length of the list. Index 0 is the unscaled + resolution's weight and each increasing scale corresponds to the image + being downsampled by 2. Defaults to (0.0448, 0.2856, 0.3001, 0.2363, + 0.1333), which are the values obtained in the original paper. + + Returns: + A tensor containing an MS-SSIM value for each image in batch. The values + are in range [0, 1]. Returns a tensor with shape: + broadcast(img1.shape[:-3], img2.shape[:-3]). + """ + # Shape checking. + shape1 = img1.get_shape().with_rank_at_least(3) + shape2 = img2.get_shape().with_rank_at_least(3) + shape1[-3:].merge_with(shape2[-3:]) + + with ops.name_scope(None, 'MS-SSIM', [img1, img2]): + shape1, shape2, checks = _verify_compatible_image_shapes(img1, img2) + with ops.control_dependencies(checks): + img1 = array_ops.identity(img1) + + # Need to convert the images to float32. Scale max_val accordingly so that + # SSIM is computed correctly. + max_val = math_ops.cast(max_val, img1.dtype) + max_val = convert_image_dtype(max_val, dtypes.float32) + img1 = convert_image_dtype(img1, dtypes.float32) + img2 = convert_image_dtype(img2, dtypes.float32) + + imgs = [img1, img2] + shapes = [shape1, shape2] + + # img1 and img2 are assumed to be a (multi-dimensional) batch of + # 3-dimensional images (height, width, channels). `heads` contain the batch + # dimensions, and `tails` contain the image dimensions. + heads = [s[:-3] for s in shapes] + tails = [s[-3:] for s in shapes] + + divisor = [1, 2, 2, 1] + divisor_tensor = constant_op.constant(divisor[1:], dtype=dtypes.int32) + + def do_pad(images, remainder): + padding = array_ops.expand_dims(remainder, -1) + padding = array_ops.pad(padding, [[1, 0], [1, 0]]) + return [array_ops.pad(x, padding, mode='SYMMETRIC') for x in images] + + mcs = [] + for k in range(len(power_factors)): + with ops.name_scope(None, 'Scale%d' % k, imgs): + if k > 0: + # Avg pool takes rank 4 tensors. Flatten leading dimensions. + flat_imgs = [ + array_ops.reshape(x, array_ops.concat([[-1], t], 0)) + for x, t in zip(imgs, tails) + ] + + remainder = tails[0] % divisor_tensor + need_padding = math_ops.reduce_any(math_ops.not_equal(remainder, 0)) + # pylint: disable=cell-var-from-loop + padded = control_flow_ops.cond(need_padding, + lambda: do_pad(flat_imgs, remainder), + lambda: flat_imgs) + # pylint: enable=cell-var-from-loop + + downscaled = [nn_ops.avg_pool(x, ksize=divisor, strides=divisor, + padding='VALID') + for x in padded] + tails = [x[1:] for x in array_ops.shape_n(downscaled)] + imgs = [ + array_ops.reshape(x, array_ops.concat([h, t], 0)) + for x, h, t in zip(downscaled, heads, tails) + ] + + # Overwrite previous ssim value since we only need the last one. + ssim_per_channel, cs = _ssim_per_channel(*imgs, max_val=max_val) + mcs.append(nn_ops.relu(cs)) + + # Remove the cs score for the last scale. In the MS-SSIM calculation, + # we use the l(p) at the highest scale. l(p) * cs(p) is ssim(p). + mcs.pop() # Remove the cs score for the last scale. + mcs_and_ssim = array_ops.stack(mcs + [nn_ops.relu(ssim_per_channel)], + axis=-1) + # Take weighted geometric mean across the scale axis. + ms_ssim = math_ops.reduce_prod(math_ops.pow(mcs_and_ssim, power_factors), + [-1]) + + return math_ops.reduce_mean(ms_ssim, [-1]) # Avg over color channels. + + +@tf_export('image.image_gradients') +def image_gradients(image): + """Returns image gradients (dy, dx) for each color channel. + + Both output tensors have the same shape as the input: [batch_size, h, w, + d]. The gradient values are organized so that [I(x+1, y) - I(x, y)] is in + location (x, y). That means that dy will always have zeros in the last row, + and dx will always have zeros in the last column. + + Arguments: + image: Tensor with shape [batch_size, h, w, d]. + + Returns: + Pair of tensors (dy, dx) holding the vertical and horizontal image + gradients (1-step finite difference). + + Raises: + ValueError: If `image` is not a 4D tensor. + """ + if image.get_shape().ndims != 4: + raise ValueError('image_gradients expects a 4D tensor ' + '[batch_size, h, w, d], not %s.', image.get_shape()) + image_shape = array_ops.shape(image) + batch_size, height, width, depth = array_ops.unstack(image_shape) + dy = image[:, 1:, :, :] - image[:, :-1, :, :] + dx = image[:, :, 1:, :] - image[:, :, :-1, :] + + # Return tensors with same size as original image by concatenating + # zeros. Place the gradient [I(x+1,y) - I(x,y)] on the base pixel (x, y). + shape = array_ops.stack([batch_size, 1, width, depth]) + dy = array_ops.concat([dy, array_ops.zeros(shape, image.dtype)], 1) + dy = array_ops.reshape(dy, image_shape) + + shape = array_ops.stack([batch_size, height, 1, depth]) + dx = array_ops.concat([dx, array_ops.zeros(shape, image.dtype)], 2) + dx = array_ops.reshape(dx, image_shape) + + return dy, dx + + +@tf_export('image.sobel_edges') +def sobel_edges(image): + """Returns a tensor holding Sobel edge maps. + + Arguments: + image: Image tensor with shape [batch_size, h, w, d] and type float32 or + float64. The image(s) must be 2x2 or larger. + + Returns: + Tensor holding edge maps for each channel. Returns a tensor with shape + [batch_size, h, w, d, 2] where the last two dimensions hold [[dy[0], dx[0]], + [dy[1], dx[1]], ..., [dy[d-1], dx[d-1]]] calculated using the Sobel filter. + """ + # Define vertical and horizontal Sobel filters. + static_image_shape = image.get_shape() + image_shape = array_ops.shape(image) + kernels = [[[-1, -2, -1], [0, 0, 0], [1, 2, 1]], + [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]] + num_kernels = len(kernels) + kernels = np.transpose(np.asarray(kernels), (1, 2, 0)) + kernels = np.expand_dims(kernels, -2) + kernels_tf = constant_op.constant(kernels, dtype=image.dtype) + + kernels_tf = array_ops.tile(kernels_tf, [1, 1, image_shape[-1], 1], + name='sobel_filters') + + # Use depth-wise convolution to calculate edge maps per channel. + pad_sizes = [[0, 0], [1, 1], [1, 1], [0, 0]] + padded = array_ops.pad(image, pad_sizes, mode='REFLECT') + + # Output tensor has shape [batch_size, h, w, d * num_kernels]. + strides = [1, 1, 1, 1] + output = nn.depthwise_conv2d(padded, kernels_tf, strides, 'VALID') + + # Reshape to [batch_size, h, w, d, num_kernels]. + shape = array_ops.concat([image_shape, [num_kernels]], 0) + output = array_ops.reshape(output, shape=shape) + output.set_shape(static_image_shape.concatenate([num_kernels])) + return output diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py index 6a516a99110d2578c905b71020d9242e5636a72b..c437c12c2744792eaee197bf7d2a5f2b75d280bf 100644 --- a/tensorflow/python/ops/image_ops_test.py +++ b/tensorflow/python/ops/image_ops_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import colorsys import functools +import itertools import math import os import time @@ -37,7 +38,9 @@ from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_image_ops +from tensorflow.python.ops import gradients from tensorflow.python.ops import image_ops +from tensorflow.python.ops import image_ops_impl from tensorflow.python.ops import io_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops @@ -934,7 +937,7 @@ class AdjustSaturationTest(test_util.TensorFlowTestCase): class FlipTransposeRotateTest(test_util.TensorFlowTestCase): - def testIdempotentLeftRight(self): + def testInvolutionLeftRight(self): x_np = np.array([[1, 2, 3], [1, 2, 3]], dtype=np.uint8).reshape([2, 3, 1]) with self.test_session(use_gpu=True): x_tf = constant_op.constant(x_np, shape=x_np.shape) @@ -942,6 +945,16 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_tf = y.eval() self.assertAllEqual(y_tf, x_np) + def testInvolutionLeftRightWithBatch(self): + x_np = np.array( + [[[1, 2, 3], [1, 2, 3]], [[1, 2, 3], [1, 2, 3]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + with self.test_session(use_gpu=True): + x_tf = constant_op.constant(x_np, shape=x_np.shape) + y = image_ops.flip_left_right(image_ops.flip_left_right(x_tf)) + y_tf = y.eval() + self.assertAllEqual(y_tf, x_np) + def testLeftRight(self): x_np = np.array([[1, 2, 3], [1, 2, 3]], dtype=np.uint8).reshape([2, 3, 1]) y_np = np.array([[3, 2, 1], [3, 2, 1]], dtype=np.uint8).reshape([2, 3, 1]) @@ -953,9 +966,24 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_tf = y.eval() self.assertAllEqual(y_tf, y_np) + def testLeftRightWithBatch(self): + x_np = np.array( + [[[1, 2, 3], [1, 2, 3]], [[1, 2, 3], [1, 2, 3]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + y_np = np.array( + [[[3, 2, 1], [3, 2, 1]], [[3, 2, 1], [3, 2, 1]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + + with self.test_session(use_gpu=True): + x_tf = constant_op.constant(x_np, shape=x_np.shape) + y = image_ops.flip_left_right(x_tf) + y_tf = y.eval() + self.assertAllEqual(y_tf, y_np) + def testRandomFlipLeftRight(self): x_np = np.array([[1, 2, 3], [1, 2, 3]], dtype=np.uint8).reshape([2, 3, 1]) y_np = np.array([[3, 2, 1], [3, 2, 1]], dtype=np.uint8).reshape([2, 3, 1]) + seed = 42 with self.test_session(use_gpu=True): x_tf = constant_op.constant(x_np, shape=x_np.shape) @@ -964,7 +992,7 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): count_flipped = 0 count_unflipped = 0 - for _ in range(50): + for _ in range(100): y_tf = y.eval() if y_tf[0][0] == 1: self.assertAllEqual(y_tf, x_np) @@ -972,10 +1000,15 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): else: self.assertAllEqual(y_tf, y_np) count_flipped += 1 - self.assertGreaterEqual(count_flipped, 1) - self.assertGreaterEqual(count_unflipped, 1) - def testIdempotentUpDown(self): + # 100 trials + # Mean: 50 + # Std Dev: ~5 + # Six Sigma: 50 - (5 * 6) = 20 + self.assertGreaterEqual(count_flipped, 20) + self.assertGreaterEqual(count_unflipped, 20) + + def testInvolutionUpDown(self): x_np = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8).reshape([2, 3, 1]) with self.test_session(use_gpu=True): @@ -984,6 +1017,17 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_tf = y.eval() self.assertAllEqual(y_tf, x_np) + def testInvolutionUpDownWithBatch(self): + x_np = np.array( + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + + with self.test_session(use_gpu=True): + x_tf = constant_op.constant(x_np, shape=x_np.shape) + y = image_ops.flip_up_down(image_ops.flip_up_down(x_tf)) + y_tf = y.eval() + self.assertAllEqual(y_tf, x_np) + def testUpDown(self): x_np = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8).reshape([2, 3, 1]) y_np = np.array([[4, 5, 6], [1, 2, 3]], dtype=np.uint8).reshape([2, 3, 1]) @@ -995,17 +1039,31 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_tf = y.eval() self.assertAllEqual(y_tf, y_np) + def testUpDownWithBatch(self): + x_np = np.array( + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + y_np = np.array( + [[[4, 5, 6], [1, 2, 3]], [[10, 11, 12], [7, 8, 9]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + + with self.test_session(use_gpu=True): + x_tf = constant_op.constant(x_np, shape=x_np.shape) + y = image_ops.flip_up_down(x_tf) + y_tf = y.eval() + self.assertAllEqual(y_tf, y_np) + def testRandomFlipUpDown(self): x_np = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8).reshape([2, 3, 1]) y_np = np.array([[4, 5, 6], [1, 2, 3]], dtype=np.uint8).reshape([2, 3, 1]) with self.test_session(use_gpu=True): x_tf = constant_op.constant(x_np, shape=x_np.shape) - y = image_ops.random_flip_up_down(x_tf) + y = image_ops.random_flip_up_down(x_tf, seed=42) self.assertTrue(y.op.name.startswith("random_flip_up_down")) count_flipped = 0 count_unflipped = 0 - for _ in range(50): + for _ in range(100): y_tf = y.eval() if y_tf[0][0] == 1: self.assertAllEqual(y_tf, x_np) @@ -1013,10 +1071,15 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): else: self.assertAllEqual(y_tf, y_np) count_flipped += 1 - self.assertGreaterEqual(count_flipped, 1) - self.assertGreaterEqual(count_unflipped, 1) - def testIdempotentTranspose(self): + # 100 trials + # Mean: 50 + # Std Dev: ~5 + # Six Sigma: 50 - (5 * 6) = 20 + self.assertGreaterEqual(count_flipped, 20) + self.assertGreaterEqual(count_unflipped, 20) + + def testInvolutionTranspose(self): x_np = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8).reshape([2, 3, 1]) with self.test_session(use_gpu=True): @@ -1025,6 +1088,17 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_tf = y.eval() self.assertAllEqual(y_tf, x_np) + def testInvolutionTransposeWithBatch(self): + x_np = np.array( + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + + with self.test_session(use_gpu=True): + x_tf = constant_op.constant(x_np, shape=x_np.shape) + y = image_ops.transpose_image(image_ops.transpose_image(x_tf)) + y_tf = y.eval() + self.assertAllEqual(y_tf, x_np) + def testTranspose(self): x_np = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint8).reshape([2, 3, 1]) y_np = np.array([[1, 4], [2, 5], [3, 6]], dtype=np.uint8).reshape([3, 2, 1]) @@ -1036,15 +1110,34 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_tf = y.eval() self.assertAllEqual(y_tf, y_np) + def testTransposeWithBatch(self): + x_np = np.array( + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + dtype=np.uint8).reshape([2, 2, 3, 1]) + + y_np = np.array( + [[[1, 4], [2, 5], [3, 6]], [[7, 10], [8, 11], [9, 12]]], + dtype=np.uint8).reshape([2, 3, 2, 1]) + + with self.test_session(use_gpu=True): + x_tf = constant_op.constant(x_np, shape=x_np.shape) + y = image_ops.transpose_image(x_tf) + y_tf = y.eval() + self.assertAllEqual(y_tf, y_np) + def testPartialShapes(self): p_unknown_rank = array_ops.placeholder(dtypes.uint8) - p_unknown_dims = array_ops.placeholder( + p_unknown_dims_3 = array_ops.placeholder( dtypes.uint8, shape=[None, None, None]) + p_unknown_dims_4 = array_ops.placeholder( + dtypes.uint8, shape=[None, None, None, None]) p_unknown_width = array_ops.placeholder(dtypes.uint8, shape=[64, None, 3]) - + p_unknown_batch = array_ops.placeholder( + dtypes.uint8, shape=[None, 64, 64, 3]) p_wrong_rank = array_ops.placeholder(dtypes.uint8, shape=[None, None]) p_zero_dim = array_ops.placeholder(dtypes.uint8, shape=[64, 0, 3]) + #Ops that support 3D input for op in [ image_ops.flip_left_right, image_ops.flip_up_down, image_ops.random_flip_left_right, image_ops.random_flip_up_down, @@ -1052,16 +1145,35 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): ]: transformed_unknown_rank = op(p_unknown_rank) self.assertEqual(3, transformed_unknown_rank.get_shape().ndims) - transformed_unknown_dims = op(p_unknown_dims) - self.assertEqual(3, transformed_unknown_dims.get_shape().ndims) + transformed_unknown_dims_3 = op(p_unknown_dims_3) + self.assertEqual(3, transformed_unknown_dims_3.get_shape().ndims) transformed_unknown_width = op(p_unknown_width) self.assertEqual(3, transformed_unknown_width.get_shape().ndims) - with self.assertRaisesRegexp(ValueError, "must be three-dimensional"): - op(p_wrong_rank) with self.assertRaisesRegexp(ValueError, "must be > 0"): op(p_zero_dim) + #Ops that support 4D input + for op in [ + image_ops.flip_left_right, image_ops.flip_up_down, + image_ops.transpose_image, image_ops.rot90 + ]: + transformed_unknown_dims_4 = op(p_unknown_dims_4) + self.assertEqual(4, transformed_unknown_dims_4.get_shape().ndims) + transformed_unknown_batch = op(p_unknown_batch) + self.assertEqual(4, transformed_unknown_batch.get_shape().ndims) + with self.assertRaisesRegexp(ValueError, + "must be at least three-dimensional"): + op(p_wrong_rank) + + for op in [ + image_ops.random_flip_left_right, + image_ops.random_flip_up_down, + ]: + with self.assertRaisesRegexp(ValueError, "must be three-dimensional"): + op(p_wrong_rank) + + def testRot90GroupOrder(self): image = np.arange(24, dtype=np.uint8).reshape([2, 4, 3]) with self.test_session(use_gpu=True): @@ -1070,6 +1182,14 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): rotated = image_ops.rot90(rotated) self.assertAllEqual(image, rotated.eval()) + def testRot90GroupOrderWithBatch(self): + image = np.arange(48, dtype=np.uint8).reshape([2, 2, 4, 3]) + with self.test_session(use_gpu=True): + rotated = image + for _ in xrange(4): + rotated = image_ops.rot90(rotated) + self.assertAllEqual(image, rotated.eval()) + def testRot90NumpyEquivalence(self): image = np.arange(24, dtype=np.uint8).reshape([2, 4, 3]) with self.test_session(use_gpu=True): @@ -1079,6 +1199,14 @@ class FlipTransposeRotateTest(test_util.TensorFlowTestCase): y_np = np.rot90(image, k=k) self.assertAllEqual(y_np, y_tf.eval({k_placeholder: k})) + def testRot90NumpyEquivalenceWithBatch(self): + image = np.arange(48, dtype=np.uint8).reshape([2, 2, 4, 3]) + with self.test_session(use_gpu=True): + k_placeholder = array_ops.placeholder(dtypes.int32, shape=[]) + y_tf = image_ops.rot90(image, k_placeholder) + for k in xrange(4): + y_np = np.rot90(image, k=k, axes=(1, 2)) + self.assertAllEqual(y_np, y_tf.eval({k_placeholder: k})) class RandomFlipTest(test_util.TensorFlowTestCase): @@ -1869,8 +1997,8 @@ class SelectDistortedCropBoxTest(test_util.TensorFlowTestCase): image_size = constant_op.constant( [40, 50, 1], shape=[3], dtype=dtypes.int32) bounding_box = constant_op.constant( - [0.0, 0.0, 1.0, 1.0], - shape=[4], + [[[0.0, 0.0, 1.0, 1.0]]], + shape=[1, 1, 4], dtype=dtypes.float32, ) begin, end, bbox_for_drawing = image_ops.sample_distorted_bounding_box( @@ -1884,6 +2012,10 @@ class SelectDistortedCropBoxTest(test_util.TensorFlowTestCase): self.assertAllEqual([3], begin.get_shape().as_list()) self.assertAllEqual([3], end.get_shape().as_list()) self.assertAllEqual([1, 1, 4], bbox_for_drawing.get_shape().as_list()) + # Actual run to make sure shape is correct inside Compute(). + begin = begin.eval() + end = end.eval() + bbox_for_drawing = bbox_for_drawing.eval() begin, end, bbox_for_drawing = image_ops.sample_distorted_bounding_box( image_size=image_size, @@ -1903,9 +2035,10 @@ class SelectDistortedCropBoxTest(test_util.TensorFlowTestCase): image_size = constant_op.constant( [40, 50, 1], shape=[3], dtype=dtypes.int32) bounding_box = constant_op.constant( - [0.0, 0.0, 1.0, 1.0], - shape=[4], - dtype=dtypes.float32,) + [[[0.0, 0.0, 1.0, 1.0]]], + shape=[1, 1, 4], + dtype=dtypes.float32, + ) begin, end, bbox_for_drawing = image_ops.sample_distorted_bounding_box( image_size=image_size, bounding_boxes=bounding_box, @@ -1915,6 +2048,10 @@ class SelectDistortedCropBoxTest(test_util.TensorFlowTestCase): self.assertAllEqual([3], begin.get_shape().as_list()) self.assertAllEqual([3], end.get_shape().as_list()) self.assertAllEqual([1, 1, 4], bbox_for_drawing.get_shape().as_list()) + # Actual run to make sure shape is correct inside Compute(). + begin = begin.eval() + end = end.eval() + bbox_for_drawing = bbox_for_drawing.eval() class ResizeImagesTest(test_util.TensorFlowTestCase): @@ -2821,20 +2958,9 @@ class PngTest(test_util.TensorFlowTestCase): class GifTest(test_util.TensorFlowTestCase): - def testOptimizedGifErrorString(self): - filename = "tensorflow/core/lib/gif/testdata/optimized.gif" - - with self.test_session(use_gpu=True) as sess: - gif = io_ops.read_file(filename) - image = image_ops.decode_gif(gif) - with self.assertRaisesRegexp(errors.InvalidArgumentError, - "can't process optimized gif"): - gif, image = sess.run([gif, image]) - - def testValid(self): + def _testValid(self, filename): # Read some real GIFs prefix = "tensorflow/core/lib/gif/testdata/" - filename = "scan.gif" WIDTH = 20 HEIGHT = 40 STRIDE = 5 @@ -2861,16 +2987,9 @@ class GifTest(test_util.TensorFlowTestCase): self.assertAllClose(frame, gt) - def testInValid(self): - # Read some real GIFs - prefix = "tensorflow/core/lib/gif/testdata/" - filename = "optimized.gif" - - with self.test_session(use_gpu=True) as sess: - gif0 = io_ops.read_file(prefix + filename) - image0 = image_ops.decode_gif(gif0) - with self.assertRaises(errors.InvalidArgumentError): - gif0, image0 = sess.run([gif0, image0]) + def testValid(self): + self._testValid("scan.gif") + self._testValid("optimized.gif") def testShape(self): with self.test_session(use_gpu=True) as sess: @@ -3169,6 +3288,463 @@ class NonMaxSuppressionTest(test_util.TensorFlowTestCase): boxes, scores, max_output_size, iou_threshold).eval() self.assertAllClose(selected_indices, [3, 0, 5]) + def testInvalidShape(self): + # The boxes should be 2D of shape [num_boxes, 4]. + with self.assertRaisesRegexp(ValueError, + "Shape must be rank 2 but is rank 1"): + boxes = constant_op.constant([0.0, 0.0, 1.0, 1.0]) + scores = constant_op.constant([0.9]) + image_ops.non_max_suppression(boxes, scores, 3, 0.5) + + with self.assertRaisesRegexp(ValueError, "Dimension must be 4 but is 3"): + boxes = constant_op.constant([[0.0, 0.0, 1.0]]) + scores = constant_op.constant([0.9]) + image_ops.non_max_suppression(boxes, scores, 3, 0.5) + + # The boxes is of shape [num_boxes, 4], and the scores is + # of shape [num_boxes]. So an error will thrown. + with self.assertRaisesRegexp(ValueError, + "Dimensions must be equal, but are 1 and 2"): + boxes = constant_op.constant([[0.0, 0.0, 1.0, 1.0]]) + scores = constant_op.constant([0.9, 0.75]) + selected_indices = image_ops.non_max_suppression(boxes, scores, 3, 0.5) + + # The scores should be 1D of shape [num_boxes]. + with self.assertRaisesRegexp(ValueError, + "Shape must be rank 1 but is rank 2"): + boxes = constant_op.constant([[0.0, 0.0, 1.0, 1.0]]) + scores = constant_op.constant([[0.9]]) + image_ops.non_max_suppression(boxes, scores, 3, 0.5) + + # The max_output_size should be a scaler (0-D). + with self.assertRaisesRegexp(ValueError, + "Shape must be rank 0 but is rank 1"): + boxes = constant_op.constant([[0.0, 0.0, 1.0, 1.0]]) + scores = constant_op.constant([0.9]) + image_ops.non_max_suppression(boxes, scores, [3], 0.5) + + # The iou_threshold should be a scaler (0-D). + with self.assertRaisesRegexp(ValueError, + "Shape must be rank 0 but is rank 2"): + boxes = constant_op.constant([[0.0, 0.0, 1.0, 1.0]]) + scores = constant_op.constant([0.9]) + image_ops.non_max_suppression(boxes, scores, 3, [[0.5]]) + + +class VerifyCompatibleImageShapesTest(test_util.TensorFlowTestCase): + """Tests utility function used by ssim() and psnr().""" + + def testWrongDims(self): + img = array_ops.placeholder(dtype=dtypes.float32) + img_np = np.array((2, 2)) + + with self.test_session(use_gpu=True) as sess: + _, _, checks = image_ops_impl._verify_compatible_image_shapes(img, img) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(checks, {img: img_np}) + + def testShapeMismatch(self): + img1 = array_ops.placeholder(dtype=dtypes.float32) + img2 = array_ops.placeholder(dtype=dtypes.float32) + + img1_np = np.array([1, 2, 2, 1]) + img2_np = np.array([1, 3, 3, 1]) + + with self.test_session(use_gpu=True) as sess: + _, _, checks = image_ops_impl._verify_compatible_image_shapes(img1, img2) + with self.assertRaises(errors.InvalidArgumentError): + sess.run(checks, {img1: img1_np, img2: img2_np}) + + +class PSNRTest(test_util.TensorFlowTestCase): + """Tests for PSNR.""" + + def _LoadTestImage(self, sess, filename): + content = io_ops.read_file(os.path.join( + "tensorflow/core/lib/psnr/testdata", filename)) + im = image_ops.decode_jpeg(content, dct_method="INTEGER_ACCURATE") + im = image_ops.convert_image_dtype(im, dtypes.float32) + im, = sess.run([im]) + return np.expand_dims(im, axis=0) + + def _LoadTestImages(self): + with self.test_session(use_gpu=True) as sess: + q20 = self._LoadTestImage(sess, "cat_q20.jpg") + q72 = self._LoadTestImage(sess, "cat_q72.jpg") + q95 = self._LoadTestImage(sess, "cat_q95.jpg") + return q20, q72, q95 + + def _PSNR_NumPy(self, orig, target, max_value): + """Numpy implementation of PSNR.""" + mse = ((orig - target) ** 2).mean(axis=(-3, -2, -1)) + return 20 * np.log10(max_value) - 10 * np.log10(mse) + + def _RandomImage(self, shape, max_val): + """Returns an image or image batch with given shape.""" + return np.random.rand(*shape).astype(np.float32) * max_val + + def testPSNRSingleImage(self): + image1 = self._RandomImage((8, 8, 1), 1) + image2 = self._RandomImage((8, 8, 1), 1) + psnr = self._PSNR_NumPy(image1, image2, 1) + + with self.test_session(use_gpu=True): + tf_image1 = constant_op.constant(image1, shape=image1.shape, + dtype=dtypes.float32) + tf_image2 = constant_op.constant(image2, shape=image2.shape, + dtype=dtypes.float32) + tf_psnr = image_ops.psnr(tf_image1, tf_image2, 1.0, "psnr").eval() + self.assertAllClose(psnr, tf_psnr, atol=0.001) + + def testPSNRMultiImage(self): + image1 = self._RandomImage((10, 8, 8, 1), 1) + image2 = self._RandomImage((10, 8, 8, 1), 1) + psnr = self._PSNR_NumPy(image1, image2, 1) + + with self.test_session(use_gpu=True): + tf_image1 = constant_op.constant(image1, shape=image1.shape, + dtype=dtypes.float32) + tf_image2 = constant_op.constant(image2, shape=image2.shape, + dtype=dtypes.float32) + tf_psnr = image_ops.psnr(tf_image1, tf_image2, 1, "psnr").eval() + self.assertAllClose(psnr, tf_psnr, atol=0.001) + + def testGoldenPSNR(self): + q20, q72, q95 = self._LoadTestImages() + + # Verify NumPy implementation first. + # Golden values are generated using GNU Octave's psnr() function. + psnr1 = self._PSNR_NumPy(q20, q72, 1) + self.assertNear(30.321, psnr1, 0.001, msg="q20.dtype=" + str(q20.dtype)) + psnr2 = self._PSNR_NumPy(q20, q95, 1) + self.assertNear(29.994, psnr2, 0.001) + psnr3 = self._PSNR_NumPy(q72, q95, 1) + self.assertNear(35.302, psnr3, 0.001) + + # Test TensorFlow implementation. + with self.test_session(use_gpu=True): + tf_q20 = constant_op.constant(q20, shape=q20.shape, dtype=dtypes.float32) + tf_q72 = constant_op.constant(q72, shape=q72.shape, dtype=dtypes.float32) + tf_q95 = constant_op.constant(q95, shape=q95.shape, dtype=dtypes.float32) + tf_psnr1 = image_ops.psnr(tf_q20, tf_q72, 1, "psnr1").eval() + tf_psnr2 = image_ops.psnr(tf_q20, tf_q95, 1, "psnr2").eval() + tf_psnr3 = image_ops.psnr(tf_q72, tf_q95, 1, "psnr3").eval() + self.assertAllClose(psnr1, tf_psnr1, atol=0.001) + self.assertAllClose(psnr2, tf_psnr2, atol=0.001) + self.assertAllClose(psnr3, tf_psnr3, atol=0.001) + + def testInfinity(self): + q20, _, _ = self._LoadTestImages() + psnr = self._PSNR_NumPy(q20, q20, 1) + with self.test_session(use_gpu=True): + tf_q20 = constant_op.constant(q20, shape=q20.shape, dtype=dtypes.float32) + tf_psnr = image_ops.psnr(tf_q20, tf_q20, 1, "psnr").eval() + self.assertAllClose(psnr, tf_psnr, atol=0.001) + + def testInt(self): + img1 = self._RandomImage((10, 8, 8, 1), 255) + img2 = self._RandomImage((10, 8, 8, 1), 255) + img1 = constant_op.constant(img1, dtypes.uint8) + img2 = constant_op.constant(img2, dtypes.uint8) + psnr_uint8 = image_ops.psnr(img1, img2, 255) + img1 = image_ops.convert_image_dtype(img1, dtypes.float32) + img2 = image_ops.convert_image_dtype(img2, dtypes.float32) + psnr_float32 = image_ops.psnr(img1, img2, 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(psnr_uint8.eval(), psnr_float32.eval(), atol=0.001) + + +class SSIMTest(test_util.TensorFlowTestCase): + """Tests for SSIM.""" + + _filenames = ["checkerboard1.png", + "checkerboard2.png", + "checkerboard3.png",] + + _ssim = np.asarray([[1.000000, 0.230880, 0.231153], + [0.230880, 1.000000, 0.996828], + [0.231153, 0.996828, 1.000000]]) + + def _LoadTestImage(self, sess, filename): + content = io_ops.read_file(os.path.join( + "tensorflow/core/lib/ssim/testdata", filename)) + im = image_ops.decode_png(content) + im = image_ops.convert_image_dtype(im, dtypes.float32) + im, = sess.run([im]) + return np.expand_dims(im, axis=0) + + def _LoadTestImages(self): + with self.test_session(use_gpu=True) as sess: + return [self._LoadTestImage(sess, f) for f in self._filenames] + + def _RandomImage(self, shape, max_val): + """Returns an image or image batch with given shape.""" + return np.random.rand(*shape).astype(np.float32) * max_val + + def testAgainstMatlab(self): + """Tests against values produced by Matlab.""" + img = self._LoadTestImages() + expected = self._ssim[np.triu_indices(3)] + + ph = [array_ops.placeholder(dtype=dtypes.float32) for _ in range(2)] + ssim = image_ops.ssim(*ph, max_val=1.0) + with self.test_session(use_gpu=True): + scores = [ssim.eval(dict(zip(ph, t))) + for t in itertools.combinations_with_replacement(img, 2)] + self.assertAllClose(expected, np.squeeze(scores), atol=1e-4) + + def testBatch(self): + img = self._LoadTestImages() + expected = self._ssim[np.triu_indices(3, k=1)] + + img1, img2 = zip(*itertools.combinations(img, 2)) + img1 = np.concatenate(img1) + img2 = np.concatenate(img2) + + ssim = image_ops.ssim(constant_op.constant(img1), + constant_op.constant(img2), 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(expected, ssim.eval(), atol=1e-4) + + def testBroadcast(self): + img = self._LoadTestImages()[:2] + expected = self._ssim[:2, :2] + + img = constant_op.constant(np.concatenate(img)) + img1 = array_ops.expand_dims(img, axis=0) # batch dims: 1, 2. + img2 = array_ops.expand_dims(img, axis=1) # batch dims: 2, 1. + + ssim = image_ops.ssim(img1, img2, 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(expected, ssim.eval(), atol=1e-4) + + def testNegative(self): + """Tests against negative SSIM index.""" + step = np.expand_dims(np.arange(0, 256, 16, dtype=np.uint8), axis=0) + img1 = np.tile(step, (16, 1)) + img2 = np.fliplr(img1) + + img1 = img1.reshape((1, 16, 16, 1)) + img2 = img2.reshape((1, 16, 16, 1)) + + ssim = image_ops.ssim(constant_op.constant(img1), + constant_op.constant(img2), 255) + with self.test_session(use_gpu=True): + self.assertLess(ssim.eval(), 0) + + def testInt(self): + img1 = self._RandomImage((1, 16, 16, 3), 255) + img2 = self._RandomImage((1, 16, 16, 3), 255) + img1 = constant_op.constant(img1, dtypes.uint8) + img2 = constant_op.constant(img2, dtypes.uint8) + ssim_uint8 = image_ops.ssim(img1, img2, 255) + img1 = image_ops.convert_image_dtype(img1, dtypes.float32) + img2 = image_ops.convert_image_dtype(img2, dtypes.float32) + ssim_float32 = image_ops.ssim(img1, img2, 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(ssim_uint8.eval(), ssim_float32.eval(), atol=0.001) + + +class MultiscaleSSIMTest(test_util.TensorFlowTestCase): + """Tests for MS-SSIM.""" + + _filenames = ["checkerboard1.png", + "checkerboard2.png", + "checkerboard3.png",] + + _msssim = np.asarray([[1.000000, 0.091016, 0.091025], + [0.091016, 1.000000, 0.999567], + [0.091025, 0.999567, 1.000000]]) + + def _LoadTestImage(self, sess, filename): + content = io_ops.read_file(os.path.join( + "tensorflow/core/lib/ssim/testdata", filename)) + im = image_ops.decode_png(content) + im = image_ops.convert_image_dtype(im, dtypes.float32) + im, = sess.run([im]) + return np.expand_dims(im, axis=0) + + def _LoadTestImages(self): + with self.test_session(use_gpu=True) as sess: + return [self._LoadTestImage(sess, f) for f in self._filenames] + + def _RandomImage(self, shape, max_val): + """Returns an image or image batch with given shape.""" + return np.random.rand(*shape).astype(np.float32) * max_val + + def testAgainstMatlab(self): + """Tests against MS-SSIM computed with Matlab implementation. + + For color images, MS-SSIM scores are averaged over color channels. + """ + img = self._LoadTestImages() + expected = self._msssim[np.triu_indices(3)] + + ph = [array_ops.placeholder(dtype=dtypes.float32) for _ in range(2)] + msssim = image_ops.ssim_multiscale(*ph, max_val=1.0) + with self.test_session(use_gpu=True): + scores = [msssim.eval(dict(zip(ph, t))) + for t in itertools.combinations_with_replacement(img, 2)] + + self.assertAllClose(expected, np.squeeze(scores), atol=1e-4) + + def testUnweightedIsDifferentiable(self): + img = self._LoadTestImages() + ph = [array_ops.placeholder(dtype=dtypes.float32) for _ in range(2)] + scalar = constant_op.constant(1.0, dtype=dtypes.float32) + scaled_ph = [x * scalar for x in ph] + msssim = image_ops.ssim_multiscale(*scaled_ph, max_val=1.0, + power_factors=(1, 1, 1, 1, 1)) + grads = gradients.gradients(msssim, scalar) + with self.test_session(use_gpu=True) as sess: + np_grads = sess.run(grads, feed_dict={ph[0]: img[0], ph[1]: img[1]}) + self.assertTrue(np.isfinite(np_grads).all()) + + def testBatch(self): + """Tests MS-SSIM computed in batch.""" + img = self._LoadTestImages() + expected = self._msssim[np.triu_indices(3, k=1)] + + img1, img2 = zip(*itertools.combinations(img, 2)) + img1 = np.concatenate(img1) + img2 = np.concatenate(img2) + + msssim = image_ops.ssim_multiscale(constant_op.constant(img1), + constant_op.constant(img2), 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(expected, msssim.eval(), 1e-4) + + def testBroadcast(self): + """Tests MS-SSIM broadcasting.""" + img = self._LoadTestImages()[:2] + expected = self._msssim[:2, :2] + + img = constant_op.constant(np.concatenate(img)) + img1 = array_ops.expand_dims(img, axis=0) # batch dims: 1, 2. + img2 = array_ops.expand_dims(img, axis=1) # batch dims: 2, 1. + + score_tensor = image_ops.ssim_multiscale(img1, img2, 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(expected, score_tensor.eval(), 1e-4) + + def testRange(self): + """Tests against low MS-SSIM score. + + MS-SSIM is a geometric mean of SSIM and CS scores of various scales. + If any of the value is negative so that the geometric mean is not + well-defined, then treat the MS-SSIM score as zero. + """ + with self.test_session(use_gpu=True) as sess: + img1 = self._LoadTestImage(sess, "checkerboard1.png") + img2 = self._LoadTestImage(sess, "checkerboard3.png") + images = [img1, img2, np.zeros_like(img1), + np.full_like(img1, fill_value=255)] + + images = [ops.convert_to_tensor(x, dtype=dtypes.float32) for x in images] + msssim_ops = [image_ops.ssim_multiscale(x, y, 1.0) + for x, y in itertools.combinations(images, 2)] + msssim = sess.run(msssim_ops) + msssim = np.squeeze(msssim) + + self.assertTrue(np.all(msssim >= 0.0)) + self.assertTrue(np.all(msssim <= 1.0)) + + def testInt(self): + img1 = self._RandomImage((1, 180, 240, 3), 255) + img2 = self._RandomImage((1, 180, 240, 3), 255) + img1 = constant_op.constant(img1, dtypes.uint8) + img2 = constant_op.constant(img2, dtypes.uint8) + ssim_uint8 = image_ops.ssim_multiscale(img1, img2, 255) + img1 = image_ops.convert_image_dtype(img1, dtypes.float32) + img2 = image_ops.convert_image_dtype(img2, dtypes.float32) + ssim_float32 = image_ops.ssim_multiscale(img1, img2, 1.0) + with self.test_session(use_gpu=True): + self.assertAllClose(ssim_uint8.eval(), ssim_float32.eval(), atol=0.001) + + +class ImageGradientsTest(test_util.TensorFlowTestCase): + + def testImageGradients(self): + shape = [1, 2, 4, 1] + img = constant_op.constant([[1, 3, 4, 2], [8, 7, 5, 6]]) + img = array_ops.reshape(img, shape) + + expected_dy = np.reshape([[7, 4, 1, 4], [0, 0, 0, 0]], shape) + expected_dx = np.reshape([[2, 1, -2, 0], [-1, -2, 1, 0]], shape) + + dy, dx = image_ops.image_gradients(img) + with self.test_session(): + actual_dy = dy.eval() + actual_dx = dx.eval() + self.assertAllClose(expected_dy, actual_dy) + self.assertAllClose(expected_dx, actual_dx) + + def testImageGradientsMultiChannelBatch(self): + batch = [[[[1, 2], [2, 5], [3, 3]], + [[8, 4], [5, 1], [9, 8]]], + [[[5, 3], [7, 9], [1, 6]], + [[1, 2], [6, 3], [6, 3]]]] + + expected_dy = [[[[7, 2], [3, -4], [6, 5]], + [[0, 0], [0, 0], [0, 0]]], + [[[-4, -1], [-1, -6], [5, -3]], + [[0, 0], [0, 0], [0, 0]]]] + + expected_dx = [[[[1, 3], [1, -2], [0, 0]], + [[-3, -3], [4, 7], [0, 0]]], + [[[2, 6], [-6, -3], [0, 0]], + [[5, 1], [0, 0], [0, 0]]]] + + batch = constant_op.constant(batch) + assert batch.get_shape().as_list() == [2, 2, 3, 2] + dy, dx = image_ops.image_gradients(batch) + with self.test_session(use_gpu=True): + actual_dy = dy.eval() + actual_dx = dx.eval() + self.assertAllClose(expected_dy, actual_dy) + self.assertAllClose(expected_dx, actual_dx) + + def testImageGradientsBadShape(self): + # [2 x 4] image but missing batch and depth dimensions. + img = constant_op.constant([[1, 3, 4, 2], [8, 7, 5, 6]]) + with self.assertRaises(ValueError): + image_ops.image_gradients(img) + + +class SobelEdgesTest(test_util.TensorFlowTestCase): + + def testSobelEdges1x2x3x1(self): + img = constant_op.constant([[1, 3, 6], [4, 1, 5]], + dtype=dtypes.float32, shape=[1, 2, 3, 1]) + expected = np.reshape([[[0, 0], [0, 12], [0, 0]], + [[0, 0], [0, 12], [0, 0]]], [1, 2, 3, 1, 2]) + sobel = image_ops.sobel_edges(img) + with self.test_session(use_gpu=True): + actual_sobel = sobel.eval() + self.assertAllClose(expected, actual_sobel) + + def testSobelEdges5x3x4x2(self): + batch_size = 5 + plane = np.reshape([[1, 3, 6, 2], [4, 1, 5, 7], [2, 5, 1, 4]], + [1, 3, 4, 1]) + two_channel = np.concatenate([plane, plane], axis=3) + batch = np.concatenate([two_channel] * batch_size, axis=0) + img = constant_op.constant(batch, dtype=dtypes.float32, + shape=[batch_size, 3, 4, 2]) + + expected_plane = np.reshape([[[0, 0], [0, 12], [0, 10], [0, 0]], + [[6, 0], [0, 6], [-6, 10], [-6, 0]], + [[0, 0], [0, 0], [0, 10], [0, 0]]], + [1, 3, 4, 1, 2]) + expected_two_channel = np.concatenate( + [expected_plane, expected_plane], axis=3) + expected_batch = np.concatenate([expected_two_channel] * batch_size, axis=0) + + sobel = image_ops.sobel_edges(img) + with self.test_session(use_gpu=True): + actual_sobel = sobel.eval() + self.assertAllClose(expected_batch, actual_sobel) + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/python/ops/init_ops.py b/tensorflow/python/ops/init_ops.py index c7502d0fda5c38079362d30877a917e3965e6ca0..40ab22951b1aa04a61e09aac155b6449ae358d7b 100644 --- a/tensorflow/python/ops/init_ops.py +++ b/tensorflow/python/ops/init_ops.py @@ -542,6 +542,62 @@ class Orthogonal(Initializer): return {"gain": self.gain, "seed": self.seed, "dtype": self.dtype.name} +class ConvolutionDeltaOrthogonal(Initializer): + """Initializer that generates a delta orthogonal kernel for ConvNets. + + The shape of the tensor must have length 3, 4 or 5. The number of input + filters must not exceed the number of output filters. The center pixels of the + tensor form an orthogonal matrix. Other pixels are set to be zero. + + Args: + gain: multiplicative factor to apply to the orthogonal matrix. Default is 1. + The 2-norm of an input is multiplied by a factor of 'sqrt(gain)' after + applying this convolution. + dtype: The type of the output. + seed: A Python integer. Used to create random seeds. See + @{tf.set_random_seed} + for behavior. + """ + + def __init__(self, gain=1.0, seed=None, dtype=dtypes.float32): + self.gain = gain + self.dtype = _assert_float_dtype(dtypes.as_dtype(dtype)) + self.seed = seed + + def __call__(self, shape, dtype=None, partition_info=None): + if dtype is None: + dtype = self.dtype + # Check the shape + if len(shape) < 3 or len(shape) > 5: + raise ValueError("The tensor to initialize must be at least " + "three-dimensional and at most five-dimensional") + + if shape[-2] > shape[-1]: + raise ValueError("In_filters cannot be greater than out_filters.") + + # Generate a random matrix + a = random_ops.random_normal([shape[-1], shape[-1]], + dtype=dtype, seed=self.seed) + # Compute the qr factorization + q, _ = linalg_ops.qr(a, full_matrices=False) + q = q[:shape[-2], :] + q *= math_ops.sqrt(math_ops.cast(self.gain, dtype=dtype)) + if len(shape) == 3: + weight = array_ops.scatter_nd([[(shape[0]-1)//2]], + array_ops.expand_dims(q, 0), shape) + elif len(shape) == 4: + weight = array_ops.scatter_nd([[(shape[0]-1)//2, (shape[1]-1)//2]], + array_ops.expand_dims(q, 0), shape) + else: + weight = array_ops.scatter_nd([[(shape[0]-1)//2, (shape[1]-1)//2, + (shape[2]-1)//2]], + array_ops.expand_dims(q, 0), shape) + return weight + + def get_config(self): + return {"gain": self.gain, "seed": self.seed, "dtype": self.dtype.name} + + @tf_export("keras.initializers.Identity", "initializers.identity") class Identity(Initializer): """Initializer that generates the identity matrix. @@ -586,7 +642,7 @@ uniform_unit_scaling_initializer = UniformUnitScaling variance_scaling_initializer = VarianceScaling orthogonal_initializer = Orthogonal identity_initializer = Identity - +convolutional_delta_orthogonal = ConvolutionDeltaOrthogonal # pylint: enable=invalid-name diff --git a/tensorflow/python/ops/initializers_ns.py b/tensorflow/python/ops/initializers_ns.py index c21079f2971a4bdd76b4be1a803055c12b243903..e7996efe93eb2f33306a52ded91c273009192789 100644 --- a/tensorflow/python/ops/initializers_ns.py +++ b/tensorflow/python/ops/initializers_ns.py @@ -39,5 +39,8 @@ global_variables = _variables.global_variables_initializer local_variables = _variables.local_variables_initializer # Seal API. +del absolute_import +del division +del print_function del init_ops del _variables diff --git a/tensorflow/python/ops/io_ops.py b/tensorflow/python/ops/io_ops.py index 5e70b3186f382a0c795b1795b2db27bb2058ee41..f6a25610c5a2ee8b76d06e286365cb957ab643cd 100644 --- a/tensorflow/python/ops/io_ops.py +++ b/tensorflow/python/ops/io_ops.py @@ -111,10 +111,10 @@ def _save(filename, tensor_names, tensors, tensor_slices=None, name="save"): An Operation that saves the tensors. """ if tensor_slices is None: - return gen_io_ops._save(filename, tensor_names, tensors, name=name) + return gen_io_ops.save(filename, tensor_names, tensors, name=name) else: - return gen_io_ops._save_slices(filename, tensor_names, tensor_slices, - tensors, name=name) + return gen_io_ops.save_slices(filename, tensor_names, tensor_slices, + tensors, name=name) def _restore_slice(file_pattern, tensor_name, shape_and_slice, tensor_type, @@ -136,7 +136,7 @@ def _restore_slice(file_pattern, tensor_name, shape_and_slice, tensor_type, A tensor of type "tensor_type". """ base_type = dtypes.as_dtype(tensor_type).base_dtype - return gen_io_ops._restore_slice( + return gen_io_ops.restore_slice( file_pattern, tensor_name, shape_and_slice, base_type, preferred_shard, name=name) @@ -173,7 +173,7 @@ class ReaderBase(object): Raises: RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "Readers are not supported when eager execution is enabled. " "Instead, please use tf.data to get data into your model.") @@ -208,12 +208,12 @@ class ReaderBase(object): else: queue_ref = queue.queue_ref if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_read_v2(self._reader_ref, queue_ref, name=name) + return gen_io_ops.reader_read_v2(self._reader_ref, queue_ref, name=name) else: # For compatibility with pre-resource queues, create a ref(string) tensor # which can be looked up as the same queue by a resource manager. - old_queue_op = gen_data_flow_ops._fake_queue(queue_ref) - return gen_io_ops._reader_read(self._reader_ref, old_queue_op, name=name) + old_queue_op = gen_data_flow_ops.fake_queue(queue_ref) + return gen_io_ops.reader_read(self._reader_ref, old_queue_op, name=name) def read_up_to(self, queue, num_records, # pylint: disable=invalid-name name=None): @@ -240,18 +240,18 @@ class ReaderBase(object): else: queue_ref = queue.queue_ref if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_read_up_to_v2(self._reader_ref, - queue_ref, - num_records, - name=name) + return gen_io_ops.reader_read_up_to_v2(self._reader_ref, + queue_ref, + num_records, + name=name) else: # For compatibility with pre-resource queues, create a ref(string) tensor # which can be looked up as the same queue by a resource manager. - old_queue_op = gen_data_flow_ops._fake_queue(queue_ref) - return gen_io_ops._reader_read_up_to(self._reader_ref, - old_queue_op, - num_records, - name=name) + old_queue_op = gen_data_flow_ops.fake_queue(queue_ref) + return gen_io_ops.reader_read_up_to(self._reader_ref, + old_queue_op, + num_records, + name=name) def num_records_produced(self, name=None): """Returns the number of records this reader has produced. @@ -267,11 +267,11 @@ class ReaderBase(object): """ if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_num_records_produced_v2(self._reader_ref, - name=name) + return gen_io_ops.reader_num_records_produced_v2(self._reader_ref, + name=name) else: - return gen_io_ops._reader_num_records_produced(self._reader_ref, - name=name) + return gen_io_ops.reader_num_records_produced(self._reader_ref, + name=name) def num_work_units_completed(self, name=None): """Returns the number of work units this reader has finished processing. @@ -283,11 +283,11 @@ class ReaderBase(object): An int64 Tensor. """ if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_num_work_units_completed_v2(self._reader_ref, - name=name) + return gen_io_ops.reader_num_work_units_completed_v2(self._reader_ref, + name=name) else: - return gen_io_ops._reader_num_work_units_completed(self._reader_ref, - name=name) + return gen_io_ops.reader_num_work_units_completed(self._reader_ref, + name=name) def serialize_state(self, name=None): """Produce a string tensor that encodes the state of a reader. @@ -302,9 +302,9 @@ class ReaderBase(object): A string Tensor. """ if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_serialize_state_v2(self._reader_ref, name=name) + return gen_io_ops.reader_serialize_state_v2(self._reader_ref, name=name) else: - return gen_io_ops._reader_serialize_state(self._reader_ref, name=name) + return gen_io_ops.reader_serialize_state(self._reader_ref, name=name) def restore_state(self, state, name=None): """Restore a reader to a previously saved state. @@ -321,11 +321,10 @@ class ReaderBase(object): The created Operation. """ if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_restore_state_v2( + return gen_io_ops.reader_restore_state_v2( self._reader_ref, state, name=name) else: - return gen_io_ops._reader_restore_state( - self._reader_ref, state, name=name) + return gen_io_ops.reader_restore_state(self._reader_ref, state, name=name) @property def supports_serialize(self): @@ -342,9 +341,9 @@ class ReaderBase(object): The created Operation. """ if self._reader_ref.dtype == dtypes.resource: - return gen_io_ops._reader_reset_v2(self._reader_ref, name=name) + return gen_io_ops.reader_reset_v2(self._reader_ref, name=name) else: - return gen_io_ops._reader_reset(self._reader_ref, name=name) + return gen_io_ops.reader_reset(self._reader_ref, name=name) ops.NotDifferentiable("ReaderRead") @@ -377,7 +376,7 @@ class WholeFileReader(ReaderBase): Args: name: A name for the operation (optional). """ - rr = gen_io_ops._whole_file_reader_v2(name=name) + rr = gen_io_ops.whole_file_reader_v2(name=name) super(WholeFileReader, self).__init__(rr, supports_serialize=True) @@ -406,8 +405,8 @@ class TextLineReader(ReaderBase): to skip from the beginning of every file. name: A name for the operation (optional). """ - rr = gen_io_ops._text_line_reader_v2(skip_header_lines=skip_header_lines, - name=name) + rr = gen_io_ops.text_line_reader_v2(skip_header_lines=skip_header_lines, + name=name) super(TextLineReader, self).__init__(rr) @@ -444,7 +443,7 @@ class FixedLengthRecordReader(ReaderBase): name: A name for the operation (optional). encoding: The type of encoding for the file. Defaults to none. """ - rr = gen_io_ops._fixed_length_record_reader_v2( + rr = gen_io_ops.fixed_length_record_reader_v2( record_bytes=record_bytes, header_bytes=header_bytes, footer_bytes=footer_bytes, @@ -480,7 +479,7 @@ class TFRecordReader(ReaderBase): compression_type = python_io.TFRecordOptions.get_compression_type_string( options) - rr = gen_io_ops._tf_record_reader_v2( + rr = gen_io_ops.tf_record_reader_v2( name=name, compression_type=compression_type) super(TFRecordReader, self).__init__(rr) @@ -506,7 +505,7 @@ class LMDBReader(ReaderBase): name: A name for the operation (optional). options: A LMDBRecordOptions object (optional). """ - rr = gen_io_ops._lmdb_reader(name=name) + rr = gen_io_ops.lmdb_reader(name=name) super(LMDBReader, self).__init__(rr) @@ -534,7 +533,7 @@ class IdentityReader(ReaderBase): Args: name: A name for the operation (optional). """ - rr = gen_io_ops._identity_reader_v2(name=name) + rr = gen_io_ops.identity_reader_v2(name=name) super(IdentityReader, self).__init__(rr, supports_serialize=True) diff --git a/tensorflow/python/ops/linalg/linalg.py b/tensorflow/python/ops/linalg/linalg.py index 5369007a56c89ef8601f8144c2fe18717e2e78fe..14319025ff275944cf34e30128df96254d06072b 100644 --- a/tensorflow/python/ops/linalg/linalg.py +++ b/tensorflow/python/ops/linalg/linalg.py @@ -41,4 +41,5 @@ del gen_linalg_ops del linalg_ops del math_ops del special_math_ops +del tf_export # pylint: enable=undefined-variable diff --git a/tensorflow/python/ops/linalg/linalg_impl.py b/tensorflow/python/ops/linalg/linalg_impl.py index db33a08137e1d2508314c2d28bdbbb001198e6c1..8343c62816c6aeadc77dae701ae9917a86e68954 100644 --- a/tensorflow/python/ops/linalg/linalg_impl.py +++ b/tensorflow/python/ops/linalg/linalg_impl.py @@ -24,24 +24,26 @@ from tensorflow.python.ops import gen_linalg_ops from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import special_math_ops +from tensorflow.python.util.tf_export import tf_export # Linear algebra ops. band_part = array_ops.matrix_band_part cholesky = linalg_ops.cholesky cholesky_solve = linalg_ops.cholesky_solve det = linalg_ops.matrix_determinant -# pylint: disable=protected-access -slogdet = gen_linalg_ops._log_matrix_determinant -# pylint: disable=protected-access +slogdet = gen_linalg_ops.log_matrix_determinant +tf_export('linalg.slogdet')(slogdet) diag = array_ops.matrix_diag diag_part = array_ops.matrix_diag_part eigh = linalg_ops.self_adjoint_eig eigvalsh = linalg_ops.self_adjoint_eigvals einsum = special_math_ops.einsum -expm = gen_linalg_ops._matrix_exponential +expm = gen_linalg_ops.matrix_exponential +tf_export('linalg.expm')(expm) eye = linalg_ops.eye inv = linalg_ops.matrix_inverse -logm = gen_linalg_ops._matrix_logarithm +logm = gen_linalg_ops.matrix_logarithm +tf_export('linalg.logm')(logm) lstsq = linalg_ops.matrix_solve_ls norm = linalg_ops.norm qr = linalg_ops.qr @@ -54,6 +56,7 @@ transpose = array_ops.matrix_transpose triangular_solve = linalg_ops.matrix_triangular_solve +@tf_export('linalg.logdet') def logdet(matrix, name=None): """Computes log of the determinant of a hermitian positive definite matrix. @@ -65,8 +68,8 @@ def logdet(matrix, name=None): ``` Args: - matrix: A `Tensor`. Must be `float32`, `float64`, `complex64`, or - `complex128` with shape `[..., M, M]`. + matrix: A `Tensor`. Must be `float16`, `float32`, `float64`, `complex64`, + or `complex128` with shape `[..., M, M]`. name: A name to give this `Op`. Defaults to `logdet`. Returns: @@ -86,6 +89,7 @@ def logdet(matrix, name=None): reduction_indices=[-1]) +@tf_export('linalg.adjoint') def adjoint(matrix, name=None): """Transposes the last two dimensions of and conjugates tensor `matrix`. @@ -99,8 +103,8 @@ def adjoint(matrix, name=None): # [3 - 3j, 6 - 6j]] Args: - matrix: A `Tensor`. Must be `float32`, `float64`, `complex64`, or - `complex128` with shape `[..., M, M]`. + matrix: A `Tensor`. Must be `float16`, `float32`, `float64`, `complex64`, + or `complex128` with shape `[..., M, M]`. name: A name to give this `Op` (optional). Returns: diff --git a/tensorflow/python/ops/linalg/linear_operator.py b/tensorflow/python/ops/linalg/linear_operator.py index 27e0f17020afa0fd44ec11c49b7a77d4426933dd..957a7959181efe3bbc319e62582053329b763dc3 100644 --- a/tensorflow/python/ops/linalg/linear_operator.py +++ b/tensorflow/python/ops/linalg/linear_operator.py @@ -32,11 +32,13 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops.linalg import linalg_impl as linalg from tensorflow.python.ops.linalg import linear_operator_util from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export __all__ = ["LinearOperator"] # TODO(langmore) Use matrix_solve_ls for singular or non-square matrices. +@tf_export("linalg.LinearOperator") class LinearOperator(object): """Base class defining a [batch of] linear operator[s]. @@ -478,7 +480,6 @@ class LinearOperator(object): cond, self._max_condition_number_to_be_non_singular(), message="Singular matrix up to precision epsilon.") - raise NotImplementedError("assert_non_singular is not implemented.") def _max_condition_number_to_be_non_singular(self): """Return the maximum condition number that we consider nonsingular.""" diff --git a/tensorflow/python/ops/linalg/linear_operator_composition.py b/tensorflow/python/ops/linalg/linear_operator_composition.py index 14411291d4fddeb2242e243d9a611e9c2fcd171a..ecd30e4d7e4dd7cfd4b109ad6e60aacb172700f6 100644 --- a/tensorflow/python/ops/linalg/linear_operator_composition.py +++ b/tensorflow/python/ops/linalg/linear_operator_composition.py @@ -25,10 +25,12 @@ from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops from tensorflow.python.ops.linalg import linear_operator +from tensorflow.python.util.tf_export import tf_export __all__ = ["LinearOperatorComposition"] +@tf_export("linalg.LinearOperatorComposition") class LinearOperatorComposition(linear_operator.LinearOperator): """Composes one or more `LinearOperators`. diff --git a/tensorflow/python/ops/linalg/linear_operator_diag.py b/tensorflow/python/ops/linalg/linear_operator_diag.py index a4724d030f388230cf85cc68bf60b6553b409c17..e180e830263c44fb5ae290d307f1ef80106c31d5 100644 --- a/tensorflow/python/ops/linalg/linear_operator_diag.py +++ b/tensorflow/python/ops/linalg/linear_operator_diag.py @@ -26,10 +26,12 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops.linalg import linalg_impl as linalg from tensorflow.python.ops.linalg import linear_operator from tensorflow.python.ops.linalg import linear_operator_util +from tensorflow.python.util.tf_export import tf_export __all__ = ["LinearOperatorDiag",] +@tf_export("linalg.LinearOperatorDiag") class LinearOperatorDiag(linear_operator.LinearOperator): """`LinearOperator` acting like a [batch] square diagonal matrix. @@ -65,7 +67,7 @@ class LinearOperatorDiag(linear_operator.LinearOperator): operator = LinearOperatorDiag(diag) # Create a shape [2, 1, 4, 2] vector. Note that this shape is compatible - # since the batch dimensions, [2, 1], are brodcast to + # since the batch dimensions, [2, 1], are broadcast to # operator.batch_shape = [2, 3]. y = tf.random_normal(shape=[2, 1, 4, 2]) x = operator.solve(y) @@ -121,8 +123,8 @@ class LinearOperatorDiag(linear_operator.LinearOperator): Args: diag: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`. - The diagonal of the operator. Allowed dtypes: `float32`, `float64`, - `complex64`, `complex128`. + The diagonal of the operator. Allowed dtypes: `float16`, `float32`, + `float64`, `complex64`, `complex128`. is_non_singular: Expect that this operator is non-singular. is_self_adjoint: Expect that this operator is equal to its hermitian transpose. If `diag.dtype` is real, this is auto-set to `True`. @@ -167,7 +169,12 @@ class LinearOperatorDiag(linear_operator.LinearOperator): def _check_diag(self, diag): """Static check of diag.""" allowed_dtypes = [ - dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128] + dtypes.float16, + dtypes.float32, + dtypes.float64, + dtypes.complex64, + dtypes.complex128, + ] dtype = diag.dtype if dtype not in allowed_dtypes: diff --git a/tensorflow/python/ops/linalg/linear_operator_full_matrix.py b/tensorflow/python/ops/linalg/linear_operator_full_matrix.py index dd4c7cb0413013f3f54f6085a7adcb523755a603..f979fb37d6c69a2683af08a1f6722b98da0b6650 100644 --- a/tensorflow/python/ops/linalg/linear_operator_full_matrix.py +++ b/tensorflow/python/ops/linalg/linear_operator_full_matrix.py @@ -23,10 +23,12 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops.linalg import linear_operator +from tensorflow.python.util.tf_export import tf_export __all__ = ["LinearOperatorFullMatrix"] +@tf_export("linalg.LinearOperatorFullMatrix") class LinearOperatorFullMatrix(linear_operator.LinearOperator): """`LinearOperator` that wraps a [batch] matrix. @@ -114,7 +116,8 @@ class LinearOperatorFullMatrix(linear_operator.LinearOperator): Args: matrix: Shape `[B1,...,Bb, M, N]` with `b >= 0`, `M, N >= 0`. - Allowed dtypes: `float32`, `float64`, `complex64`, `complex128`. + Allowed dtypes: `float16`, `float32`, `float64`, `complex64`, + `complex128`. is_non_singular: Expect that this operator is non-singular. is_self_adjoint: Expect that this operator is equal to its hermitian transpose. @@ -147,7 +150,12 @@ class LinearOperatorFullMatrix(linear_operator.LinearOperator): def _check_matrix(self, matrix): """Static check of the `matrix` argument.""" allowed_dtypes = [ - dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128] + dtypes.float16, + dtypes.float32, + dtypes.float64, + dtypes.complex64, + dtypes.complex128, + ] matrix = ops.convert_to_tensor(matrix, name="matrix") diff --git a/tensorflow/python/ops/linalg/linear_operator_identity.py b/tensorflow/python/ops/linalg/linear_operator_identity.py index 740c6c811f2d98f62c200cda7242c6ad00de499d..50f3d407e85e4cca22ad6326931b5a2a736819a8 100644 --- a/tensorflow/python/ops/linalg/linear_operator_identity.py +++ b/tensorflow/python/ops/linalg/linear_operator_identity.py @@ -31,6 +31,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops.linalg import linalg_impl as linalg from tensorflow.python.ops.linalg import linear_operator from tensorflow.python.ops.linalg import linear_operator_util +from tensorflow.python.util.tf_export import tf_export __all__ = [ "LinearOperatorIdentity", @@ -97,6 +98,7 @@ class BaseLinearOperatorIdentity(linear_operator.LinearOperator): return array_ops.ones(shape=d_shape, dtype=self.dtype) +@tf_export("linalg.LinearOperatorIdentity") class LinearOperatorIdentity(BaseLinearOperatorIdentity): """`LinearOperator` acting like a [batch] square identity matrix. @@ -460,6 +462,7 @@ class LinearOperatorIdentity(BaseLinearOperatorIdentity): "%s" % self._batch_shape_static) +@tf_export("linalg.LinearOperatorScaledIdentity") class LinearOperatorScaledIdentity(BaseLinearOperatorIdentity): """`LinearOperator` acting like a scaled [batch] identity matrix `A = c I`. diff --git a/tensorflow/python/ops/linalg/linear_operator_low_rank_update.py b/tensorflow/python/ops/linalg/linear_operator_low_rank_update.py index ad3bb2efa94bfa9751c31ff0c704aad8faa58ba7..be911029095920d424ac90b406e7b85b73884b3b 100644 --- a/tensorflow/python/ops/linalg/linear_operator_low_rank_update.py +++ b/tensorflow/python/ops/linalg/linear_operator_low_rank_update.py @@ -27,12 +27,14 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops.linalg import linear_operator from tensorflow.python.ops.linalg import linear_operator_diag from tensorflow.python.ops.linalg import linear_operator_identity +from tensorflow.python.util.tf_export import tf_export __all__ = [ "LinearOperatorLowRankUpdate", ] +@tf_export("linalg.LinearOperatorLowRankUpdate") class LinearOperatorLowRankUpdate(linear_operator.LinearOperator): """Perturb a `LinearOperator` with a rank `K` update. @@ -150,8 +152,8 @@ class LinearOperatorLowRankUpdate(linear_operator.LinearOperator): `is_X` matrix property hints, which will trigger the appropriate code path. Args: - base_operator: Shape `[B1,...,Bb, M, N]` real `float32` or `float64` - `LinearOperator`. This is `L` above. + base_operator: Shape `[B1,...,Bb, M, N]` real `float16`, `float32` or + `float64` `LinearOperator`. This is `L` above. u: Shape `[B1,...,Bb, M, K]` `Tensor` of same `dtype` as `base_operator`. This is `U` above. diag_update: Optional shape `[B1,...,Bb, K]` `Tensor` with same `dtype` @@ -188,7 +190,11 @@ class LinearOperatorLowRankUpdate(linear_operator.LinearOperator): # because if diag has non-zero imaginary part, it will not be # self-adjoint positive definite. dtype = base_operator.dtype - allowed_dtypes = [dtypes.float32, dtypes.float64] + allowed_dtypes = [ + dtypes.float16, + dtypes.float32, + dtypes.float64, + ] if dtype not in allowed_dtypes: raise TypeError( "Argument matrix must have dtype in %s. Found: %s" diff --git a/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py b/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py index 6ea55f0367bd55379b280f81f22df2c3a0dcfb1e..a5130188b681813e1ccd4818dabdffeeb663e20a 100644 --- a/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py +++ b/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py @@ -26,12 +26,14 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops.linalg import linalg_impl as linalg from tensorflow.python.ops.linalg import linear_operator from tensorflow.python.ops.linalg import linear_operator_util +from tensorflow.python.util.tf_export import tf_export __all__ = [ "LinearOperatorLowerTriangular", ] +@tf_export("linalg.LinearOperatorLowerTriangular") class LinearOperatorLowerTriangular(linear_operator.LinearOperator): """`LinearOperator` acting like a [batch] square lower triangular matrix. @@ -118,7 +120,8 @@ class LinearOperatorLowerTriangular(linear_operator.LinearOperator): Args: tril: Shape `[B1,...,Bb, N, N]` with `b >= 0`, `N >= 0`. The lower triangular part of `tril` defines this operator. The strictly - upper triangle is ignored. Allowed dtypes: `float32`, `float64`. + upper triangle is ignored. Allowed dtypes: `float16`, `float32`, + `float64`. is_non_singular: Expect that this operator is non-singular. This operator is non-singular if and only if its diagonal elements are all non-zero. @@ -164,7 +167,11 @@ class LinearOperatorLowerTriangular(linear_operator.LinearOperator): """Static check of the `tril` argument.""" # TODO(langmore) Add complex types once matrix_triangular_solve works for # them. - allowed_dtypes = [dtypes.float32, dtypes.float64] + allowed_dtypes = [ + dtypes.float16, + dtypes.float32, + dtypes.float64, + ] dtype = tril.dtype if dtype not in allowed_dtypes: raise TypeError( diff --git a/tensorflow/python/ops/linalg/linear_operator_test_util.py b/tensorflow/python/ops/linalg/linear_operator_test_util.py index 2c11f90e6d9de280e6020edfaa4d8ef237126705..ce1a112ad584a14298be6e471578858ef31573d5 100644 --- a/tensorflow/python/ops/linalg/linear_operator_test_util.py +++ b/tensorflow/python/ops/linalg/linear_operator_test_util.py @@ -35,6 +35,18 @@ from tensorflow.python.ops.linalg import linalg_impl as linalg from tensorflow.python.platform import test +class OperatorBuildInfo(object): + """Object encoding expected shape for a test. + + Encodes the expected shape of a matrix for a test. Also + allows additional metadata for the test harness. + """ + + def __init__(self, shape, **kwargs): + self.shape = shape + self.__dict__.update(kwargs) + + @six.add_metaclass(abc.ABCMeta) # pylint: disable=no-init class LinearOperatorDerivedClassTest(test.TestCase): """Tests for derived classes. @@ -84,19 +96,20 @@ class LinearOperatorDerivedClassTest(test.TestCase): return [False, True] @abc.abstractproperty - def _shapes_to_test(self): - """Returns list of tuples, each is one shape that will be tested.""" - raise NotImplementedError("shapes_to_test has not been implemented.") + def _operator_build_infos(self): + """Returns list of OperatorBuildInfo, encapsulating the shape to test.""" + raise NotImplementedError("operator_build_infos has not been implemented.") @abc.abstractmethod - def _operator_and_mat_and_feed_dict(self, shape, dtype, use_placeholder): + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): """Build a batch matrix and an Operator that should have similar behavior. Every operator acts like a (batch) matrix. This method returns both together, and is used by tests. Args: - shape: List-like of Python integers giving full shape of operator. + build_info: `OperatorBuildInfo`, encoding shape information about the + operator. dtype: Numpy dtype. Data type of returned array/operator. use_placeholder: Python bool. If True, initialize the operator with a placeholder of undefined shape and correct dtype. @@ -164,30 +177,30 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_to_dense(self): self._skip_if_tests_to_skip_contains("to_dense") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) op_dense = operator.to_dense() if not use_placeholder: - self.assertAllEqual(shape, op_dense.get_shape()) + self.assertAllEqual(build_info.shape, op_dense.get_shape()) op_dense_v, mat_v = sess.run([op_dense, mat], feed_dict=feed_dict) self.assertAC(op_dense_v, mat_v) def test_det(self): self._skip_if_tests_to_skip_contains("det") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) op_det = operator.determinant() if not use_placeholder: - self.assertAllEqual(shape[:-2], op_det.get_shape()) + self.assertAllEqual(build_info.shape[:-2], op_det.get_shape()) op_det_v, mat_det_v = sess.run( [op_det, linalg_ops.matrix_determinant(mat)], feed_dict=feed_dict) @@ -196,16 +209,17 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_log_abs_det(self): self._skip_if_tests_to_skip_contains("log_abs_det") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) op_log_abs_det = operator.log_abs_determinant() _, mat_log_abs_det = linalg.slogdet(mat) if not use_placeholder: - self.assertAllEqual(shape[:-2], op_log_abs_det.get_shape()) + self.assertAllEqual( + build_info.shape[:-2], op_log_abs_det.get_shape()) op_log_abs_det_v, mat_log_abs_det_v = sess.run( [op_log_abs_det, mat_log_abs_det], feed_dict=feed_dict) self.assertAC(op_log_abs_det_v, mat_log_abs_det_v) @@ -213,14 +227,14 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_matmul(self): self._skip_if_tests_to_skip_contains("matmul") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: for adjoint in self._adjoint_options: for adjoint_arg in self._adjoint_arg_options: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) x = self._make_x(operator, adjoint=adjoint) # If adjoint_arg, compute A X^H^H = A X. if adjoint_arg: @@ -241,14 +255,14 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_solve(self): self._skip_if_tests_to_skip_contains("solve") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: for adjoint in self._adjoint_options: for adjoint_arg in self._adjoint_arg_options: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) rhs = self._make_rhs(operator, adjoint=adjoint) # If adjoint_arg, solve A X = (rhs^H)^H = rhs. if adjoint_arg: @@ -270,12 +284,12 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_trace(self): self._skip_if_tests_to_skip_contains("trace") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) op_trace = operator.trace() mat_trace = math_ops.trace(mat) if not use_placeholder: @@ -287,16 +301,16 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_add_to_tensor(self): self._skip_if_tests_to_skip_contains("add_to_tensor") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) op_plus_2mat = operator.add_to_tensor(2 * mat) if not use_placeholder: - self.assertAllEqual(shape, op_plus_2mat.get_shape()) + self.assertAllEqual(build_info.shape, op_plus_2mat.get_shape()) op_plus_2mat_v, mat_v = sess.run( [op_plus_2mat, mat], feed_dict=feed_dict) @@ -306,12 +320,12 @@ class LinearOperatorDerivedClassTest(test.TestCase): def test_diag_part(self): self._skip_if_tests_to_skip_contains("diag_part") for use_placeholder in self._use_placeholder_options: - for shape in self._shapes_to_test: + for build_info in self._operator_build_infos: for dtype in self._dtypes_to_test: with self.test_session(graph=ops.Graph()) as sess: sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( - shape, dtype, use_placeholder=use_placeholder) + build_info, dtype, use_placeholder=use_placeholder) op_diag_part = operator.diag_part() mat_diag_part = array_ops.matrix_diag_part(mat) @@ -334,9 +348,15 @@ class SquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): """ @property - def _shapes_to_test(self): + def _operator_build_infos(self): + build_info = OperatorBuildInfo # non-batch operators (n, n) and batch operators. - return [(0, 0), (1, 1), (1, 3, 3), (3, 4, 4), (2, 1, 4, 4)] + return [ + build_info((0, 0)), + build_info((1, 1)), + build_info((1, 3, 3)), + build_info((3, 4, 4)), + build_info((2, 1, 4, 4))] def _make_rhs(self, operator, adjoint): # This operator is square, so rhs and x will have same shape. @@ -387,9 +407,15 @@ class NonSquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): return ["solve", "det", "log_abs_det"] @property - def _shapes_to_test(self): + def _operator_build_infos(self): + build_info = OperatorBuildInfo # non-batch operators (n, n) and batch operators. - return [(2, 1), (1, 2), (1, 3, 2), (3, 3, 4), (2, 1, 2, 4)] + return [ + build_info((2, 1)), + build_info((1, 2)), + build_info((1, 3, 2)), + build_info((3, 3, 4)), + build_info((2, 1, 2, 4))] def _make_rhs(self, operator, adjoint): # TODO(langmore) Add once we're testing solve_ls. diff --git a/tensorflow/python/ops/linalg_ops.py b/tensorflow/python/ops/linalg_ops.py index 9803eed6aefe072cbe0841dff2de3f640a440dd5..5b4fb4f7c8d56e49bd4a7d2f5078ee978e7a8765 100644 --- a/tensorflow/python/ops/linalg_ops.py +++ b/tensorflow/python/ops/linalg_ops.py @@ -248,7 +248,7 @@ def matrix_solve_ls(matrix, rhs, l2_regularizer=0.0, fast=True, name=None): and l2_regularizer != 0 due to poor accuracy. """ - # pylint: disable=protected-access,long-lambda + # pylint: disable=long-lambda def _use_composite_impl(fast, tensor_shape): """Determines whether to use the composite or specialized CPU kernel. @@ -323,9 +323,8 @@ def matrix_solve_ls(matrix, rhs, l2_regularizer=0.0, fast=True, name=None): if _use_composite_impl(fast, tensor_shape): return _composite_impl(matrix, rhs, l2_regularizer) else: - return gen_linalg_ops._matrix_solve_ls( + return gen_linalg_ops.matrix_solve_ls( matrix, rhs, l2_regularizer, fast=fast, name=name) - # pylint: enable=protected-access @tf_export('self_adjoint_eig', 'linalg.eigh') @@ -342,12 +341,11 @@ def self_adjoint_eig(tensor, name=None): name: string, optional name of the operation. Returns: - e: Eigenvalues. Shape is `[..., N]`. + e: Eigenvalues. Shape is `[..., N]`. Sorted in non-decreasing order. v: Eigenvectors. Shape is `[..., N, N]`. The columns of the inner most matrices contain eigenvectors of the corresponding matrices in `tensor` """ - # pylint: disable=protected-access - e, v = gen_linalg_ops._self_adjoint_eig_v2(tensor, compute_v=True, name=name) + e, v = gen_linalg_ops.self_adjoint_eig_v2(tensor, compute_v=True, name=name) return e, v @@ -369,8 +367,7 @@ def self_adjoint_eigvals(tensor, name=None): e: Eigenvalues. Shape is `[..., N]`. The vector `e[..., :]` contains the `N` eigenvalues of `tensor[..., :, :]`. """ - # pylint: disable=protected-access - e, _ = gen_linalg_ops._self_adjoint_eig_v2(tensor, compute_v=False, name=name) + e, _ = gen_linalg_ops.self_adjoint_eig_v2(tensor, compute_v=False, name=name) return e @@ -435,10 +432,8 @@ def svd(tensor, full_matrices=False, compute_uv=True, name=None): ```` @end_compatibility """ - # pylint: disable=protected-access - s, u, v = gen_linalg_ops._svd( + s, u, v = gen_linalg_ops.svd( tensor, compute_uv=compute_uv, full_matrices=full_matrices, name=name) - # pylint: enable=protected-access if compute_uv: return math_ops.real(s), u, v else: diff --git a/tensorflow/python/ops/logging_ops.py b/tensorflow/python/ops/logging_ops.py index eadbc1b7c3b6e66aa76c9afd860b2274ac1976ae..222b8ebc9da6b076f012f8febbd50cc3c4c86c08 100644 --- a/tensorflow/python/ops/logging_ops.py +++ b/tensorflow/python/ops/logging_ops.py @@ -109,7 +109,7 @@ def histogram_summary(tag, values, collections=None, name=None): buffer. """ with ops.name_scope(name, "HistogramSummary", [tag, values]) as scope: - val = gen_logging_ops._histogram_summary( + val = gen_logging_ops.histogram_summary( tag=tag, values=values, name=scope) _Collect(val, collections, [ops.GraphKeys.SUMMARIES]) return val @@ -170,7 +170,7 @@ def image_summary(tag, tensor, max_images=3, collections=None, name=None): buffer. """ with ops.name_scope(name, "ImageSummary", [tag, tensor]) as scope: - val = gen_logging_ops._image_summary( + val = gen_logging_ops.image_summary( tag=tag, tensor=tensor, max_images=max_images, name=scope) _Collect(val, collections, [ops.GraphKeys.SUMMARIES]) return val @@ -226,11 +226,12 @@ def audio_summary(tag, with ops.name_scope(name, "AudioSummary", [tag, tensor]) as scope: sample_rate = ops.convert_to_tensor(sample_rate, dtype=dtypes.float32, name="sample_rate") - val = gen_logging_ops._audio_summary_v2(tag=tag, - tensor=tensor, - max_outputs=max_outputs, - sample_rate=sample_rate, - name=scope) + val = gen_logging_ops.audio_summary_v2( + tag=tag, + tensor=tensor, + max_outputs=max_outputs, + sample_rate=sample_rate, + name=scope) _Collect(val, collections, [ops.GraphKeys.SUMMARIES]) return val @@ -263,7 +264,7 @@ def merge_summary(inputs, collections=None, name=None): buffer resulting from the merging. """ with ops.name_scope(name, "MergeSummary", inputs): - val = gen_logging_ops._merge_summary(inputs=inputs, name=name) + val = gen_logging_ops.merge_summary(inputs=inputs, name=name) _Collect(val, collections, []) return val @@ -345,7 +346,7 @@ def scalar_summary(tags, values, collections=None, name=None): buffer. """ with ops.name_scope(name, "ScalarSummary", [tags, values]) as scope: - val = gen_logging_ops._scalar_summary(tags=tags, values=values, name=scope) + val = gen_logging_ops.scalar_summary(tags=tags, values=values, name=scope) _Collect(val, collections, [ops.GraphKeys.SUMMARIES]) return val @@ -356,3 +357,4 @@ ops.NotDifferentiable("AudioSummary") ops.NotDifferentiable("AudioSummaryV2") ops.NotDifferentiable("MergeSummary") ops.NotDifferentiable("ScalarSummary") +ops.NotDifferentiable("Timestamp") diff --git a/tensorflow/python/ops/lookup_ops.py b/tensorflow/python/ops/lookup_ops.py index f539a7bb68da57e31746bc80fb25339a03a4fafe..6f043f60e677eac560004619464905cd616256b2 100644 --- a/tensorflow/python/ops/lookup_ops.py +++ b/tensorflow/python/ops/lookup_ops.py @@ -157,10 +157,10 @@ class InitializableLookupTableBase(LookupInterface): default_value: The value to use if a key is missing in the table. initializer: The table initializer to use. """ - if context.in_graph_mode(): - name = table_ref.op.name.split("/")[-1] - else: + if context.executing_eagerly(): name = context.context().scope_name + else: + name = table_ref.op.name.split("/")[-1] super(InitializableLookupTableBase, self).__init__(initializer.key_dtype, initializer.value_dtype, name) @@ -196,9 +196,7 @@ class InitializableLookupTableBase(LookupInterface): """ with ops.name_scope(name, "%s_Size" % self._name, [self._table_ref]) as scope: - # pylint: disable=protected-access - return gen_lookup_ops._lookup_table_size_v2(self._table_ref, name=scope) - # pylint: enable=protected-access + return gen_lookup_ops.lookup_table_size_v2(self._table_ref, name=scope) def lookup(self, keys, name=None): """Looks up `keys` in a table, outputs the corresponding values. @@ -227,10 +225,8 @@ class InitializableLookupTableBase(LookupInterface): with ops.name_scope(name, "%s_Lookup" % self._name, (self._table_ref, key_tensor, self._default_value)) as scope: - # pylint: disable=protected-access - values = gen_lookup_ops._lookup_table_find_v2( + values = gen_lookup_ops.lookup_table_find_v2( self._table_ref, key_tensor, self._default_value, name=scope) - # pylint: enable=protected-access values.set_shape(key_tensor.get_shape()) if isinstance(keys, sparse_tensor.SparseTensor): @@ -274,13 +270,11 @@ class HashTable(InitializableLookupTableBase): """ with ops.name_scope(name, "hash_table", (initializer, default_value)) as scope: - # pylint: disable=protected-access - table_ref = gen_lookup_ops._hash_table_v2( + table_ref = gen_lookup_ops.hash_table_v2( shared_name=shared_name, key_dtype=initializer.key_dtype, value_dtype=initializer.value_dtype, name=scope) - # pylint: enable=protected-access super(HashTable, self).__init__(table_ref, default_value, initializer) @@ -352,10 +346,8 @@ class KeyValueTensorInitializer(TableInitializerBase): with ops.name_scope( self._name, values=(table.table_ref, self._keys, self._values)) as scope: - # pylint: disable=protected-access - init_op = gen_lookup_ops._initialize_table_v2( + init_op = gen_lookup_ops.initialize_table_v2( table.table_ref, self._keys, self._values, name=scope) - # pylint: enable=protected-access ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op) return init_op @@ -518,8 +510,7 @@ class TextFileInitializer(TableInitializerBase): (table.table_ref,)) as scope: filename = ops.convert_to_tensor( self._filename, dtypes.string, name="asset_filepath") - # pylint: disable=protected-access - init_op = gen_lookup_ops._initialize_table_from_text_file_v2( + init_op = gen_lookup_ops.initialize_table_from_text_file_v2( table.table_ref, filename, self._key_index, @@ -527,11 +518,10 @@ class TextFileInitializer(TableInitializerBase): -1 if self._vocab_size is None else self._vocab_size, self._delimiter, name=scope) - # pylint: enable=protected-access ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op) # If the filename tensor is anything other than a string constant (e.g., if # it is a placeholder) then it does not make sense to track it as an asset. - if context.in_graph_mode() and constant_op.is_constant(filename): + if not context.executing_eagerly() and constant_op.is_constant(filename): ops.add_to_collection(ops.GraphKeys.ASSET_FILEPATHS, filename) return init_op diff --git a/tensorflow/python/ops/losses/losses_impl.py b/tensorflow/python/ops/losses/losses_impl.py index 72508eb4350f57bb06b3829890f92554677c98d5..0840760810c86a6393ea6b4ab0b9410233275f11 100644 --- a/tensorflow/python/ops/losses/losses_impl.py +++ b/tensorflow/python/ops/losses/losses_impl.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import confusion_matrix @@ -28,8 +29,10 @@ from tensorflow.python.ops import nn_ops from tensorflow.python.ops import weights_broadcast_ops from tensorflow.python.ops.losses import util from tensorflow.python.util.deprecation import deprecated_args +from tensorflow.python.util.tf_export import tf_export +@tf_export("losses.Reduction") class Reduction(object): """Types of loss reduction. @@ -132,6 +135,10 @@ def _num_present(losses, weights, per_batch=False): `per_batch` is `True`, the value is returned as a tensor of size `[batch_size]`. Otherwise, a single scalar tensor is returned. """ + if ((isinstance(weights, float) and weights != 0.0) or + (context.executing_eagerly() and weights._rank() == 0 # pylint: disable=protected-access + and not math_ops.equal(weights, 0.0))): + return _num_elements(losses) with ops.name_scope(None, "num_present", (losses, weights)) as scope: weights = math_ops.to_float(weights) present = array_ops.where( @@ -141,17 +148,20 @@ def _num_present(losses, weights, per_batch=False): present = weights_broadcast_ops.broadcast_weights(present, losses) if per_batch: return math_ops.reduce_sum( - present, axis=math_ops.range(1, array_ops.rank(present)), - keep_dims=True, name=scope) + present, + axis=math_ops.range(1, array_ops.rank(present)), + keepdims=True, + name=scope) return math_ops.reduce_sum(present, name=scope) def _num_elements(losses): """Computes the number of elements in `losses` tensor.""" with ops.name_scope(None, "num_elements", values=[losses]) as scope: - return array_ops.size(losses, name=scope, out_type=losses.dtype) + return math_ops.cast(array_ops.size(losses, name=scope), dtype=losses.dtype) +@tf_export("losses.compute_weighted_loss") def compute_weighted_loss( losses, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, reduction=Reduction.SUM_BY_NONZERO_WEIGHTS): @@ -211,6 +221,7 @@ def compute_weighted_loss( return loss +@tf_export("losses.absolute_difference") def absolute_difference( labels, predictions, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, @@ -258,6 +269,7 @@ def absolute_difference( losses, weights, scope, loss_collection, reduction=reduction) +@tf_export("losses.cosine_distance") @deprecated_args(None, "dim is deprecated, use axis instead", "dim") def cosine_distance( labels, predictions, axis=None, weights=1.0, scope=None, @@ -306,11 +318,12 @@ def cosine_distance( predictions.get_shape().assert_is_compatible_with(labels.get_shape()) radial_diffs = math_ops.multiply(predictions, labels) - losses = 1 - math_ops.reduce_sum(radial_diffs, axis=(axis,), keep_dims=True) + losses = 1 - math_ops.reduce_sum(radial_diffs, axis=(axis,), keepdims=True) return compute_weighted_loss( losses, weights, scope, loss_collection, reduction=reduction) +@tf_export("losses.hinge_loss") def hinge_loss(labels, logits, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, reduction=Reduction.SUM_BY_NONZERO_WEIGHTS): @@ -352,6 +365,7 @@ def hinge_loss(labels, logits, weights=1.0, scope=None, losses, weights, scope, loss_collection, reduction=reduction) +@tf_export("losses.huber_loss") def huber_loss(labels, predictions, weights=1.0, delta=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, reduction=Reduction.SUM_BY_NONZERO_WEIGHTS): @@ -370,7 +384,7 @@ def huber_loss(labels, predictions, weights=1.0, delta=1.0, scope=None, `weights` acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If `weights` is a tensor of size - [batch_size], then the total loss for each sample of the batch is rescaled + `[batch_size]`, then the total loss for each sample of the batch is rescaled by the corresponding element in the `weights` vector. If the shape of `weights` matches the shape of `predictions`, then the loss of each measurable element of `predictions` is scaled by the corresponding value of @@ -414,12 +428,17 @@ def huber_loss(labels, predictions, weights=1.0, delta=1.0, scope=None, # expression when abs_error == delta is 0 (for tf.maximum it would be 1). # This is necessary to avoid doubling the gradient, since there is already a # nonzero contribution to the gradient from the quadratic term. - linear = (abs_error - quadratic) - losses = 0.5 * quadratic**2 + delta * linear + linear = math_ops.subtract(abs_error, quadratic) + losses = math_ops.add( + math_ops.multiply( + ops.convert_to_tensor(0.5, dtype=quadratic.dtype), + math_ops.multiply(quadratic, quadratic)), + math_ops.multiply(delta, linear)) return compute_weighted_loss( losses, weights, scope, loss_collection, reduction=reduction) +@tf_export("losses.log_loss") def log_loss(labels, predictions, weights=1.0, epsilon=1e-7, scope=None, loss_collection=ops.GraphKeys.LOSSES, reduction=Reduction.SUM_BY_NONZERO_WEIGHTS): @@ -427,7 +446,7 @@ def log_loss(labels, predictions, weights=1.0, epsilon=1e-7, scope=None, `weights` acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If `weights` is a tensor of size - [batch_size], then the total loss for each sample of the batch is rescaled + `[batch_size]`, then the total loss for each sample of the batch is rescaled by the corresponding element in the `weights` vector. If the shape of `weights` matches the shape of `predictions`, then the loss of each measurable element of `predictions` is scaled by the corresponding value of @@ -471,6 +490,7 @@ def log_loss(labels, predictions, weights=1.0, epsilon=1e-7, scope=None, # TODO(b/37208492): Add reduction arg. +@tf_export("losses.mean_pairwise_squared_error") def mean_pairwise_squared_error( labels, predictions, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES): @@ -493,7 +513,7 @@ def mean_pairwise_squared_error( `weights` acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If `weights` is a tensor of size - [batch_size], then the total loss for each sample of the batch is rescaled + `[batch_size]`, then the total loss for each sample of the batch is rescaled by the corresponding element in the `weights` vector. Args: @@ -533,17 +553,19 @@ def mean_pairwise_squared_error( reduction_indices = math_ops.range(1, array_ops.rank(diffs)) sum_squares_diff_per_batch = math_ops.reduce_sum( - math_ops.square(diffs), reduction_indices=reduction_indices, - keep_dims=True) + math_ops.square(diffs), + reduction_indices=reduction_indices, + keepdims=True) 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) + num_present_per_batch - 1) sum_diff = math_ops.reduce_sum( - diffs, reduction_indices=reduction_indices, keep_dims=True) - term2 = 2.0 * _safe_div(math_ops.square(sum_diff), - math_ops.square(num_present_per_batch)) + diffs, reduction_indices=reduction_indices, keepdims=True) + term2 = 2.0 * _safe_div( + math_ops.square(sum_diff), + math_ops.multiply(num_present_per_batch, num_present_per_batch - 1)) weighted_losses = math_ops.multiply(term1 - term2, weights) loss = math_ops.reduce_sum(weighted_losses) @@ -557,6 +579,7 @@ def mean_pairwise_squared_error( return mean_loss +@tf_export("losses.mean_squared_error") def mean_squared_error( labels, predictions, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, @@ -565,7 +588,7 @@ def mean_squared_error( `weights` acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If `weights` is a tensor of size - [batch_size], then the total loss for each sample of the batch is rescaled + `[batch_size]`, then the total loss for each sample of the batch is rescaled by the corresponding element in the `weights` vector. If the shape of `weights` matches the shape of `predictions`, then the loss of each measurable element of `predictions` is scaled by the corresponding value of @@ -604,6 +627,7 @@ def mean_squared_error( losses, weights, scope, loss_collection, reduction=reduction) +@tf_export("losses.sigmoid_cross_entropy") def sigmoid_cross_entropy( multi_class_labels, logits, weights=1.0, label_smoothing=0, scope=None, loss_collection=ops.GraphKeys.LOSSES, @@ -622,7 +646,7 @@ def sigmoid_cross_entropy( Args: multi_class_labels: `[batch_size, num_classes]` target integer labels in - `(0, 1)`. + `{0, 1}`. logits: Float `[batch_size, num_classes]` logits outputs of the network. weights: Optional `Tensor` whose rank is either 0, or the same rank as `labels`, and must be broadcastable to `labels` (i.e., all dimensions must @@ -662,6 +686,7 @@ def sigmoid_cross_entropy( losses, weights, scope, loss_collection, reduction=reduction) +@tf_export("losses.softmax_cross_entropy") def softmax_cross_entropy( onehot_labels, logits, weights=1.0, label_smoothing=0, scope=None, loss_collection=ops.GraphKeys.LOSSES, @@ -713,9 +738,11 @@ 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") + onehot_labels = array_ops.stop_gradient( + onehot_labels, name="labels_stop_gradient") + losses = nn.softmax_cross_entropy_with_logits_v2( + labels=onehot_labels, logits=logits, name="xentropy") + return compute_weighted_loss( losses, weights, scope, loss_collection, reduction=reduction) @@ -771,6 +798,7 @@ def _remove_squeezable_dimensions( return labels, predictions, weights +@tf_export("losses.sparse_softmax_cross_entropy") def sparse_softmax_cross_entropy( labels, logits, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, @@ -779,7 +807,7 @@ def sparse_softmax_cross_entropy( `weights` acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If `weights` is a - tensor of shape [`batch_size`], then the loss weights apply to each + tensor of shape `[batch_size]`, then the loss weights apply to each corresponding sample. Args: diff --git a/tensorflow/python/ops/losses/util.py b/tensorflow/python/ops/losses/util.py index 3718c481c26afdd9f007ffc22a9e6ec44a1eb10e..b835d963869704f053de6c2f8a75ae1fa72e6a5d 100644 --- a/tensorflow/python/ops/losses/util.py +++ b/tensorflow/python/ops/losses/util.py @@ -30,8 +30,10 @@ from __future__ import print_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("losses.add_loss") def add_loss(loss, loss_collection=ops.GraphKeys.LOSSES): """Adds a externally defined loss to the collection of losses. @@ -43,6 +45,7 @@ def add_loss(loss, loss_collection=ops.GraphKeys.LOSSES): ops.add_to_collection(loss_collection, loss) +@tf_export("losses.get_losses") def get_losses(scope=None, loss_collection=ops.GraphKeys.LOSSES): """Gets the list of losses from the loss_collection. @@ -56,6 +59,7 @@ def get_losses(scope=None, loss_collection=ops.GraphKeys.LOSSES): return ops.get_collection(loss_collection, scope) +@tf_export("losses.get_regularization_losses") def get_regularization_losses(scope=None): """Gets the list of regularization losses. @@ -68,6 +72,7 @@ def get_regularization_losses(scope=None): return ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES, scope) +@tf_export("losses.get_regularization_loss") def get_regularization_loss(scope=None, name="total_regularization_loss"): """Gets the total regularization loss. @@ -85,6 +90,7 @@ def get_regularization_loss(scope=None, name="total_regularization_loss"): return constant_op.constant(0.0) +@tf_export("losses.get_total_loss") def get_total_loss(add_regularization_losses=True, name="total_loss"): """Returns a tensor whose value represents the total loss. diff --git a/tensorflow/python/ops/manip_grad.py b/tensorflow/python/ops/manip_grad.py index 573e8c0a0d4e0b6e5e2ca555e34234d9bea760dc..bb2069359dd6fbe4874e228e6f2f58ea8444744d 100644 --- a/tensorflow/python/ops/manip_grad.py +++ b/tensorflow/python/ops/manip_grad.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Gradients for operators defined in manip_ops.py.""" from __future__ import absolute_import diff --git a/tensorflow/python/ops/manip_ops.py b/tensorflow/python/ops/manip_ops.py index c5f39784f436be361ba0d06c87239ae9b409fb9e..6d335cdc212f368e7667a030791c7b634113a9c6 100644 --- a/tensorflow/python/ops/manip_ops.py +++ b/tensorflow/python/ops/manip_ops.py @@ -23,11 +23,15 @@ from __future__ import print_function from tensorflow.python.ops import gen_manip_ops as _gen_manip_ops from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export + # pylint: disable=protected-access -def roll(input, shift, axis): +@tf_export('manip.roll') +def roll(input, shift, axis): # pylint: disable=redefined-builtin return _gen_manip_ops.roll(input, shift, axis) + roll.__doc__ = _gen_manip_ops.roll.__doc__ # pylint: enable=protected-access diff --git a/tensorflow/python/ops/math_grad.py b/tensorflow/python/ops/math_grad.py index 53308484c427e715f649a09f0dbe3f1448f18f5b..eb33687cb5a41e4b772e69ae24966c8c84571c21 100644 --- a/tensorflow/python/ops/math_grad.py +++ b/tensorflow/python/ops/math_grad.py @@ -35,6 +35,12 @@ def _safe_shape_div(x, y): return x // math_ops.maximum(y, 1) +@ops.RegisterGradient("ArgMax") +def _ArgMaxGrad(op, grad): + del op, grad + return [None, None] + + @ops.RegisterGradient("Sum") def _SumGrad(op, grad): """Gradient for Sum.""" @@ -46,10 +52,18 @@ def _SumGrad(op, grad): if axes is not None: rank = len(input_0_shape) if np.array_equal(axes, np.arange(rank)): # Reduce all dims. - grad = array_ops.reshape(grad, [1] * rank) + if context.executing_eagerly(): + ctx = context.context() + new_shape = ctx.ones_rank_cache().get(rank) + if new_shape is None: + new_shape = constant_op.constant([1] * rank, dtype=dtypes.int32) + ctx.ones_rank_cache().put(rank, new_shape) + else: + new_shape = [1] * rank + grad = array_ops.reshape(grad, new_shape) # If shape is not fully defined (but rank is), we use Shape. if None not in input_0_shape: - input_shape = input_0_shape + input_shape = constant_op.constant(input_0_shape, dtype=dtypes.int32) else: input_shape = array_ops.shape(op.inputs[0]) return [array_ops.tile(grad, input_shape), None] @@ -228,56 +242,142 @@ def _SparseSegmentSqrtNWithNumSegmentsGrad(op, grad): dim0), None, None, None) -def _SegmentMinOrMaxGrad(op, grad, is_sorted): - """Gradient for SegmentMin and (unsorted) SegmentMax. - - They share similar code. - """ - zeros = array_ops.zeros( - array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype) - +def _SegmentMinOrMaxGrad(op, grad): + """ Gradient for SegmentMin and SegmentMax. """ + zeros = array_ops.zeros_like(op.inputs[0], dtype=op.inputs[0].dtype) # Get the number of selected (minimum or maximum) elements in each segment. gathered_outputs = array_ops.gather(op.outputs[0], op.inputs[1]) is_selected = math_ops.equal(op.inputs[0], gathered_outputs) - if is_sorted: - num_selected = math_ops.segment_sum( - math_ops.cast(is_selected, grad.dtype), op.inputs[1]) - else: - num_selected = math_ops.unsorted_segment_sum( - math_ops.cast(is_selected, grad.dtype), op.inputs[1], op.inputs[2]) - + num_selected = math_ops.segment_sum(math_ops.cast(is_selected, grad.dtype), + op.inputs[1]) # Compute the gradient for each segment. The gradient for the ith segment is # divided evenly among the selected elements in that segment. weighted_grads = math_ops.div(grad, num_selected) gathered_grads = array_ops.gather(weighted_grads, op.inputs[1]) - - if is_sorted: - return array_ops.where(is_selected, gathered_grads, zeros), None - else: - return array_ops.where(is_selected, gathered_grads, zeros), None, None + return array_ops.where(is_selected, gathered_grads, zeros), None @ops.RegisterGradient("SegmentMin") def _SegmentMinGrad(op, grad): """Gradient for SegmentMin.""" - return _SegmentMinOrMaxGrad(op, grad, True) + return _SegmentMinOrMaxGrad(op, grad) @ops.RegisterGradient("SegmentMax") def _SegmentMaxGrad(op, grad): """Gradient for SegmentMax.""" - return _SegmentMinOrMaxGrad(op, grad, True) + return _SegmentMinOrMaxGrad(op, grad) + + +def _GatherDropNegatives(params, ids, zero_clipped_indices=None, + is_positive=None): + """ Helper function for unsorted segment ops. Gathers params for + positive segment ids and gathers 0 for inputs with negative segment id. + Also returns the clipped indices and a boolean mask with the same shape + as ids where a positive id is masked as true. With this, the latter two + can be passed as arguments to this function to reuse them. + """ + if zero_clipped_indices is None: + zero_clipped_indices = math_ops.maximum(ids, array_ops.zeros_like(ids)) + gathered = array_ops.gather(params, zero_clipped_indices) + if is_positive is None: + is_positive = math_ops.greater_equal(ids, 0) + # tf.where(condition, x, y) requires condition to have the same shape as x + # and y. + # todo(philjd): remove this if tf.where supports broadcasting (#9284) + for _ in range(gathered.shape.ndims - is_positive.shape.ndims): + is_positive = array_ops.expand_dims(is_positive, -1) + is_positive = (is_positive & + array_ops.ones_like(gathered, dtype=dtypes.bool)) + # replace gathered params of negative indices with 0 + zero_slice = array_ops.zeros_like(gathered) + return (array_ops.where(is_positive, gathered, zero_slice), + zero_clipped_indices, is_positive) + + +def _UnsortedSegmentMinOrMaxGrad(op, grad): + """ Gradient for UnsortedSegmentMin and UnsortedSegmentMax. """ + # Get the number of selected (minimum or maximum) elements in each segment. + gathered_outputs, zero_clipped_indices, is_positive = \ + _GatherDropNegatives(op.outputs[0], op.inputs[1]) + is_selected = math_ops.equal(op.inputs[0], gathered_outputs) + is_selected = math_ops.logical_and(is_selected, is_positive) + num_selected = math_ops.unsorted_segment_sum( + math_ops.cast(is_selected, grad.dtype), op.inputs[1], op.inputs[2]) + # Compute the gradient for each segment. The gradient for the ith segment is + # divided evenly among the selected elements in that segment. + weighted_grads = math_ops.div(grad, num_selected) + gathered_grads, _, _ = _GatherDropNegatives(weighted_grads, None, + zero_clipped_indices, + is_positive) + zeros = array_ops.zeros_like(gathered_grads) + return array_ops.where(is_selected, gathered_grads, zeros), None, None @ops.RegisterGradient("UnsortedSegmentSum") def _UnsortedSegmentSumGrad(op, grad): - """Gradient for SegmentSum.""" - return array_ops.gather(grad, op.inputs[1]), None, None + """Gradient for UnsortedSegmentSum.""" + return _GatherDropNegatives(grad, op.inputs[1])[0], None, None @ops.RegisterGradient("UnsortedSegmentMax") def _UnsortedSegmentMaxGrad(op, grad): - return _SegmentMinOrMaxGrad(op, grad, False) + """ Gradient for UnsortedSegmentMax. """ + return _UnsortedSegmentMinOrMaxGrad(op, grad) + + +@ops.RegisterGradient("UnsortedSegmentMin") +def _UnsortedSegmentMinGrad(op, grad): + """ Gradient for UnsortedSegmentMin. """ + return _UnsortedSegmentMinOrMaxGrad(op, grad) + + +@ops.RegisterGradient("UnsortedSegmentProd") +def _UnsortedSegmentProdGrad(op, grad): + """ Gradient for UnsortedSegmentProd. + The gradient can be expressed for each segment by dividing the segment's + product by each element of the segment input tensor, but this approach can't + deal with zeros in the input. + Unlike reduce_prod we can't use cumsum here as individual segments may have + a different number of elements. Therefore we consider three cases: + 1) A segment input contains no zeros and we can safely divide by the input + tensor. + 2) A segment contains exactly one zero. Then the gradient of each input of + the segment is zero except for the 0-input, there the gradient is + the product of the remaining segment entries. + 3) A segment contains at least two zeros. The gradient is zero for all + segment inputs. + """ + # Note that unsorted_segment_sum will filter out the negative indices, + # so we don't need to do a logical_and with is_positive here + is_zero = math_ops.equal(op.inputs[0], 0) + num_zeros = gen_math_ops.unsorted_segment_sum( + math_ops.cast(is_zero, dtype=dtypes.int32), op.inputs[1], op.inputs[2]) + # handle case 3 and set the gradient to 0 for segments with more than one + # 0 as input + grad = array_ops.where(math_ops.greater(num_zeros, 1), + array_ops.zeros_like(grad), grad) + # replace all zeros with ones and compute the unsorted_segment_prod + non_zero_data = array_ops.where(is_zero, array_ops.ones_like(op.inputs[0]), + op.inputs[0]) + non_zero_prod = gen_math_ops.unsorted_segment_prod( + non_zero_data, op.inputs[1], op.inputs[2]) + # clip the indices for gather to be positive + zero_clipped_indices = math_ops.maximum(op.inputs[1], + array_ops.zeros_like(op.inputs[1])) + gathered_prod = array_ops.gather(op.outputs[0], zero_clipped_indices) + gathered_non_zero_prod = array_ops.gather(non_zero_prod, + zero_clipped_indices) + prod_divided_by_el = gathered_prod / op.inputs[0] # May contain nan/inf. + # Now fetch the individual results for segments containing 0 and those that + # don't. is_zero will also fetch results for entries with negative index + # but the following gather_drop_negatives sets the corresponding entry in + # grad to 0 for these + partial_derivative = array_ops.where(is_zero, gathered_non_zero_prod, + prod_divided_by_el) + gathered_grad = _GatherDropNegatives(grad, op.inputs[1], + zero_clipped_indices)[0] + return gathered_grad * partial_derivative, None, None @ops.RegisterGradient("Abs") @@ -296,16 +396,14 @@ def _NegGrad(_, grad): def _InvGrad(op, grad): """Returns -grad * (1 / x^2).""" y = op.outputs[0] # y = 1 / x - # pylint: disable=protected-access - return gen_math_ops._reciprocal_grad(y, grad) + return gen_math_ops.reciprocal_grad(y, grad) @ops.RegisterGradient("Reciprocal") def _ReciprocalGrad(op, grad): """Returns -grad * (1 / x^2).""" y = op.outputs[0] # y = 1 / x - # pylint: disable=protected-access - return gen_math_ops._reciprocal_grad(y, grad) + return gen_math_ops.reciprocal_grad(y, grad) @ops.RegisterGradient("InvGrad") @@ -315,8 +413,7 @@ def _InvGradGrad(op, grad): with ops.control_dependencies([grad]): ca = math_ops.conj(op.inputs[0]) cg = math_ops.conj(grad) - # pylint: disable=protected-access - return cg * -2.0 * b * ca, gen_math_ops._reciprocal_grad(ca, grad) + return cg * -2.0 * b * ca, gen_math_ops.reciprocal_grad(ca, grad) @ops.RegisterGradient("ReciprocalGrad") @@ -326,8 +423,7 @@ def _ReciprocalGradGrad(op, grad): with ops.control_dependencies([grad]): ca = math_ops.conj(op.inputs[0]) cg = math_ops.conj(grad) - # pylint: disable=protected-access - return cg * -2.0 * b * ca, gen_math_ops._reciprocal_grad(ca, grad) + return cg * -2.0 * b * ca, gen_math_ops.reciprocal_grad(ca, grad) @ops.RegisterGradient("Square") @@ -336,15 +432,14 @@ def _SquareGrad(op, grad): # Added control dependencies to prevent 2*x from being computed too early. with ops.control_dependencies([grad]): x = math_ops.conj(x) - return math_ops.multiply(grad, math_ops.multiply(x, 2.0)) + y = constant_op.constant(2.0, dtype=x.dtype) + return math_ops.multiply(grad, math_ops.multiply(x, y)) @ops.RegisterGradient("Sqrt") def _SqrtGrad(op, grad): y = op.outputs[0] # y = x^(1/2) - # pylint: disable=protected-access - return gen_math_ops._sqrt_grad(y, grad) - # pylint: enable=protected-access + return gen_math_ops.sqrt_grad(y, grad) @ops.RegisterGradient("SqrtGrad") @@ -360,9 +455,7 @@ def _SqrtGradGrad(op, grad): def _RsqrtGrad(op, grad): """Returns -0.5 * grad * conj(y)^3.""" y = op.outputs[0] # y = x^(-1/2) - # pylint: disable=protected-access - return gen_math_ops._rsqrt_grad(y, grad) - # pylint: enable=protected-access + return gen_math_ops.rsqrt_grad(y, grad) @ops.RegisterGradient("RsqrtGrad") @@ -374,8 +467,7 @@ def _RsqrtGradGrad(op, grad): ca = math_ops.conj(a) cg = math_ops.conj(grad) grad_a = -1.5 * cg * b * math_ops.square(ca) - # pylint: disable=protected-access - grad_b = gen_math_ops._rsqrt_grad(ca, grad) + grad_b = gen_math_ops.rsqrt_grad(ca, grad) return grad_a, grad_b @@ -440,8 +532,7 @@ def _TanhGrad(op, grad): y = op.outputs[0] # y = tanh(x) with ops.control_dependencies([grad]): y = math_ops.conj(y) - # pylint: disable=protected-access - return gen_math_ops._tanh_grad(y, grad) + return gen_math_ops.tanh_grad(y, grad) @ops.RegisterGradient("Asinh") @@ -479,8 +570,7 @@ def _TanhGradGrad(op, grad): with ops.control_dependencies([grad]): a = math_ops.conj(op.inputs[0]) b = math_ops.conj(op.inputs[1]) - # pylint: disable=protected-access - return grad * -2.0 * b * a, gen_math_ops._tanh_grad(a, grad) + return grad * -2.0 * b * a, gen_math_ops.tanh_grad(a, grad) @ops.RegisterGradient("Erf") @@ -530,9 +620,7 @@ def _IgammaGrad(op, grad): x = op.inputs[1] sa = array_ops.shape(a) sx = array_ops.shape(x) - # pylint: disable=protected-access - unused_ra, rx = gen_array_ops._broadcast_gradient_args(sa, sx) - # pylint: enable=protected-access + unused_ra, rx = gen_array_ops.broadcast_gradient_args(sa, sx) # Perform operations in log space before summing, because Gamma(a) # and Gamma'(a) can grow large. @@ -559,9 +647,7 @@ def _BetaincGrad(op, grad): # versa; so its sufficient to check against shape(a). sa = array_ops.shape(a) sx = array_ops.shape(x) - # pylint: disable=protected-access - _, rx = gen_array_ops._broadcast_gradient_args(sa, sx) - # pylint: enable=protected-access + _, rx = gen_array_ops.broadcast_gradient_args(sa, sx) # Perform operations in log space before summing, because terms # can grow large. @@ -587,9 +673,7 @@ def _ZetaGrad(op, grad): # Broadcast gradients sx = array_ops.shape(x) sq = array_ops.shape(q) - # pylint: disable=protected-access - unused_rx, rq = gen_array_ops._broadcast_gradient_args(sx, sq) - # pylint: enable=protected-access + unused_rx, rq = gen_array_ops.broadcast_gradient_args(sx, sq) # Evaluate gradient with ops.control_dependencies([grad]): x = math_ops.conj(x) @@ -609,9 +693,7 @@ def _PolygammaGrad(op, grad): # Broadcast gradients sn = array_ops.shape(n) sx = array_ops.shape(x) - # pylint: disable=protected-access - unused_rn, rx = gen_array_ops._broadcast_gradient_args(sn, sx) - # pylint: enable=protected-access + unused_rn, rx = gen_array_ops.broadcast_gradient_args(sn, sx) # Evaluate gradient with ops.control_dependencies([grad]): n = math_ops.conj(n) @@ -628,8 +710,7 @@ def _SigmoidGrad(op, grad): y = op.outputs[0] # y = sigmoid(x) with ops.control_dependencies([grad]): y = math_ops.conj(y) - # pylint: disable=protected-access - return gen_math_ops._sigmoid_grad(y, grad) + return gen_math_ops.sigmoid_grad(y, grad) @ops.RegisterGradient("SigmoidGrad") @@ -638,8 +719,7 @@ def _SigmoidGradGrad(op, grad): a = math_ops.conj(op.inputs[0]) b = math_ops.conj(op.inputs[1]) gb = grad * b - # pylint: disable=protected-access - return gb - 2.0 * gb * a, gen_math_ops._sigmoid_grad(a, grad) + return gb - 2.0 * gb * a, gen_math_ops.sigmoid_grad(a, grad) @ops.RegisterGradient("Sign") @@ -753,9 +833,7 @@ def _AddGrad(op, grad): return grad, grad sx = array_ops.shape(x) sy = array_ops.shape(y) - # pylint: disable=protected-access - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) return (array_ops.reshape(math_ops.reduce_sum(grad, rx), sx), array_ops.reshape(math_ops.reduce_sum(grad, ry), sy)) @@ -770,9 +848,7 @@ def _SubGrad(op, grad): return grad, -grad sx = array_ops.shape(x) sy = array_ops.shape(y) - # pylint: disable=protected-access - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) return (array_ops.reshape(math_ops.reduce_sum(grad, rx), sx), array_ops.reshape(-math_ops.reduce_sum(grad, ry), sy)) @@ -782,20 +858,20 @@ def _MulGrad(op, grad): """The gradient of scalar multiplication.""" x = op.inputs[0] y = op.inputs[1] - # pylint: disable=protected-access if (isinstance(grad, ops.Tensor) and _ShapesFullySpecifiedAndEqual(x, y, grad) and grad.dtype in (dtypes.int32, dtypes.float32)): - return gen_math_ops._mul(grad, y), gen_math_ops._mul(grad, x) + return gen_math_ops.mul(grad, y), gen_math_ops.mul(grad, x) assert x.dtype.base_dtype == y.dtype.base_dtype, (x.dtype, " vs. ", y.dtype) sx = array_ops.shape(x) sy = array_ops.shape(y) - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) x = math_ops.conj(x) y = math_ops.conj(y) - return (array_ops.reshape(math_ops.reduce_sum(grad * y, rx), sx), - array_ops.reshape(math_ops.reduce_sum(x * grad, ry), sy)) + return (array_ops.reshape( + math_ops.reduce_sum(gen_math_ops.mul(grad, y), rx), sx), + array_ops.reshape( + math_ops.reduce_sum(gen_math_ops.mul(x, grad), ry), sy)) @ops.RegisterGradient("Div") @@ -805,9 +881,7 @@ def _DivGrad(op, grad): y = op.inputs[1] sx = array_ops.shape(x) sy = array_ops.shape(y) - # pylint: disable=protected-access - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) x = math_ops.conj(x) y = math_ops.conj(y) return (array_ops.reshape(math_ops.reduce_sum(math_ops.div(grad, y), rx), sx), @@ -830,9 +904,7 @@ def _FloorModGrad(op, grad): sx = array_ops.shape(x) sy = array_ops.shape(y) - # pylint: disable=protected-access - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) floor_xy = math_ops.floor_div(x, y) gx = array_ops.reshape(math_ops.reduce_sum(grad, rx), sx) gy = array_ops.reshape( @@ -852,9 +924,7 @@ def _RealDivGrad(op, grad): y = op.inputs[1] sx = array_ops.shape(x) sy = array_ops.shape(y) - # pylint: disable=protected-access - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) x = math_ops.conj(x) y = math_ops.conj(y) return (array_ops.reshape( @@ -872,7 +942,7 @@ def _PowGrad(op, grad): z = op.outputs[0] sx = array_ops.shape(x) sy = array_ops.shape(y) - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) x = math_ops.conj(x) y = math_ops.conj(y) z = math_ops.conj(z) @@ -900,7 +970,7 @@ def _MaximumMinimumGrad(op, grad, selector_op): gradshape = array_ops.shape(grad) zeros = array_ops.zeros(gradshape, gdtype) xmask = selector_op(x, y) - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) xgrad = array_ops.where(xmask, grad, zeros) ygrad = array_ops.where(xmask, zeros, grad) gx = array_ops.reshape(math_ops.reduce_sum(xgrad, rx), sx) @@ -927,9 +997,7 @@ def _SquaredDifferenceGrad(op, grad): y = op.inputs[1] sx = array_ops.shape(x) sy = array_ops.shape(y) - # pylint: disable=protected-access - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) - # pylint: enable=protected-access + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) with ops.control_dependencies([grad]): # The parens ensure that if grad is IndexedSlices, it'll get multiplied by # Tensor (not a number like 2.0) which causes it to convert to Tensor. @@ -969,17 +1037,17 @@ def _MatMulGrad(op, grad): a = math_ops.conj(op.inputs[0]) b = math_ops.conj(op.inputs[1]) if not t_a and not t_b: - grad_a = math_ops.matmul(grad, b, transpose_b=True) - grad_b = math_ops.matmul(a, grad, transpose_a=True) + grad_a = gen_math_ops.mat_mul(grad, b, transpose_b=True) + grad_b = gen_math_ops.mat_mul(a, grad, transpose_a=True) elif not t_a and t_b: - grad_a = math_ops.matmul(grad, b) - grad_b = math_ops.matmul(grad, a, transpose_a=True) + grad_a = gen_math_ops.mat_mul(grad, b) + grad_b = gen_math_ops.mat_mul(grad, a, transpose_a=True) elif t_a and not t_b: - grad_a = math_ops.matmul(b, grad, transpose_b=True) - grad_b = math_ops.matmul(a, grad) + grad_a = gen_math_ops.mat_mul(b, grad, transpose_b=True) + grad_b = gen_math_ops.mat_mul(a, grad) elif t_a and t_b: - grad_a = math_ops.matmul(b, grad, transpose_a=True, transpose_b=True) - grad_b = math_ops.matmul(grad, a, transpose_a=True, transpose_b=True) + grad_a = gen_math_ops.mat_mul(b, grad, transpose_a=True, transpose_b=True) + grad_b = gen_math_ops.mat_mul(grad, a, transpose_a=True, transpose_b=True) return grad_a, grad_b @@ -993,7 +1061,7 @@ def _SparseMatMulGrad(op, grad): op.inputs[0]: op.get_attr("a_is_sparse"), op.inputs[1]: op.get_attr("b_is_sparse"), # Use heuristic to figure out if grad might be sparse - grad: context.in_graph_mode() and (grad.op.type == "ReluGrad") + grad: not context.executing_eagerly() and (grad.op.type == "ReluGrad") } def _SparseMatMul(t1, t2, out_dtype, transpose_a=False, transpose_b=False): @@ -1093,7 +1161,7 @@ def _ComplexGrad(op, grad): y = op.inputs[1] sx = array_ops.shape(x) sy = array_ops.shape(y) - rx, ry = gen_array_ops._broadcast_gradient_args(sx, sy) + rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy) return (array_ops.reshape(math_ops.reduce_sum(math_ops.real(grad), rx), sx), array_ops.reshape(math_ops.reduce_sum(math_ops.imag(grad), ry), sy)) diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py index 9a8ac93de9dcc12c513b5ddd07cca9d863d19b8a..e18d0e95015556acc66b0b50f0c406c0e2cf41aa 100644 --- a/tensorflow/python/ops/math_ops.py +++ b/tensorflow/python/ops/math_ops.py @@ -89,8 +89,6 @@ See the @{$python/math_ops} guide. @@matrix_inverse @@cholesky @@cholesky_solve -@@matrix_exponential -@@matrix_logarithm @@matrix_solve @@matrix_triangular_solve @@matrix_solve_ls @@ -129,8 +127,14 @@ See the @{$python/math_ops} guide. @@segment_min @@segment_max @@segment_mean +@@to_complex128 +@@to_complex64 @@unsorted_segment_sum @@unsorted_segment_max +@@unsorted_segment_mean +@@unsorted_segment_min +@@unsorted_segment_prod +@@unsorted_segment_sqrt_n @@sparse_segment_sum @@sparse_segment_mean @@sparse_segment_sqrt_n @@ -158,14 +162,11 @@ 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 gen_control_flow_ops from tensorflow.python.ops import gen_data_flow_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import gen_nn_ops from tensorflow.python.ops import gen_sparse_ops from tensorflow.python.ops import gen_spectral_ops -from tensorflow.python.ops import gen_state_ops -from tensorflow.python.ops import state_ops # go/tf-wildcard-import # pylint: disable=wildcard-import from tensorflow.python.ops.gen_math_ops import * @@ -181,6 +182,11 @@ arg_max = deprecation.deprecated(None, "Use `argmax` instead")(arg_max) # pylin arg_min = deprecation.deprecated(None, "Use `argmin` instead")(arg_min) # pylint: disable=used-before-assignment +# This is set by resource_variable_ops.py. It is included in this way since +# there is a circular dependency between math_ops and resource_variable_ops +_resource_variable_type = None + + def _set_doc(doc): def _decorator(func): @@ -237,7 +243,7 @@ def argmin(input, # pylint: disable=anomalous-backslash-in-string,protected-access # pylint: disable=g-docstring-has-escape @tf_export("abs") -def abs(x, name=None): +def abs(x, name=None): # pylint: disable=redefined-builtin r"""Computes the absolute value of a tensor. Given a tensor `x` of complex numbers, this operation returns a tensor of type @@ -263,7 +269,7 @@ def abs(x, name=None): with ops.name_scope(name, "Abs", [x]) as name: if isinstance(x, sparse_tensor.SparseTensor): if x.values.dtype.is_complex: - x_abs = gen_math_ops._complex_abs( + x_abs = gen_math_ops.complex_abs( x.values, Tout=x.values.dtype.real_dtype, name=name) return sparse_tensor.SparseTensor( indices=x.indices, values=x_abs, dense_shape=x.dense_shape) @@ -273,7 +279,7 @@ def abs(x, name=None): else: x = ops.convert_to_tensor(x, name="x") if x.dtype.is_complex: - return gen_math_ops._complex_abs(x, Tout=x.dtype.real_dtype, name=name) + return gen_math_ops.complex_abs(x, Tout=x.dtype.real_dtype, name=name) return gen_math_ops._abs(x, name=name) @@ -282,7 +288,7 @@ def abs(x, name=None): # pylint: disable=redefined-builtin def _bucketize(input, boundaries, name=None): - return gen_math_ops._bucketize(input=input, boundaries=boundaries, name=name) + return gen_math_ops.bucketize(input=input, boundaries=boundaries, name=name) # pylint: enable=redefined-builtin @@ -325,10 +331,10 @@ def divide(x, y, name=None): @tf_export("multiply") def multiply(x, y, name=None): - return gen_math_ops._mul(x, y, name) + return gen_math_ops.mul(x, y, name) -multiply.__doc__ = gen_math_ops._mul.__doc__.replace("Mul", "`tf.multiply`") +multiply.__doc__ = gen_math_ops.mul.__doc__.replace("Multiply", "`tf.multiply`") # TODO(aselle): put deprecation in after another round of global code changes @@ -336,19 +342,19 @@ multiply.__doc__ = gen_math_ops._mul.__doc__.replace("Mul", "`tf.multiply`") "2016-12-30", "`tf.mul(x, y)` is deprecated, please use `tf.multiply(x, y)` or `x * y`") def _mul(x, y, name=None): - return gen_math_ops._mul(x, y, name) + return gen_math_ops.mul(x, y, name) _mul.__doc__ = ( - gen_math_ops._mul.__doc__ + ("" if _mul.__doc__ is None else _mul.__doc__)) + gen_math_ops.mul.__doc__ + ("" if _mul.__doc__ is None else _mul.__doc__)) @tf_export("subtract") def subtract(x, y, name=None): - return gen_math_ops._sub(x, y, name) + return gen_math_ops.sub(x, y, name) -subtract.__doc__ = gen_math_ops._sub.__doc__.replace("`Sub`", "`tf.subtract`") +subtract.__doc__ = gen_math_ops.sub.__doc__.replace("`Sub`", "`tf.subtract`") # TODO(aselle): put deprecation in after another round of global code changes @@ -356,11 +362,11 @@ subtract.__doc__ = gen_math_ops._sub.__doc__.replace("`Sub`", "`tf.subtract`") "2016-12-30", "`tf.sub(x, y)` is deprecated, please use `tf.subtract(x, y)` or `x - y`") def _sub(x, y, name=None): - return gen_math_ops._sub(x, y, name) + return gen_math_ops.sub(x, y, name) _sub.__doc__ = ( - gen_math_ops._sub.__doc__ + ("" if _sub.__doc__ is None else _sub.__doc__)) + gen_math_ops.sub.__doc__ + ("" if _sub.__doc__ is None else _sub.__doc__)) # pylint: disable=g-docstring-has-escape @@ -380,11 +386,11 @@ def negative(x, name=None): """ with ops.name_scope(name, "Neg", [x]) as name: if isinstance(x, sparse_tensor.SparseTensor): - x_neg = gen_math_ops._neg(x.values, name=name) + x_neg = gen_math_ops.neg(x.values, name=name) return sparse_tensor.SparseTensor( indices=x.indices, values=x_neg, dense_shape=x.dense_shape) else: - return gen_math_ops._neg(x, name=name) + return gen_math_ops.neg(x, name=name) # pylint: enable=g-docstring-has-escape @@ -542,7 +548,7 @@ def scalar_mul(scalar, x): @tf_export("pow") -def pow(x, y, name=None): +def pow(x, y, name=None): # pylint: disable=redefined-builtin r"""Computes the power of one value to another. Given a tensor `x` and a tensor `y`, this operation computes \\(x^y\\) for @@ -712,7 +718,7 @@ def angle(input, name=None): @tf_export("round") -def round(x, name=None): +def round(x, name=None): # pylint: disable=redefined-builtin """Rounds the values of a tensor to the nearest integer, element-wise. Rounds half to even. Also known as bankers rounding. If you want to round @@ -898,7 +904,41 @@ def to_bfloat16(x, name="ToBFloat16"): return cast(x, dtypes.bfloat16, name=name) -ops.Tensor._override_operator("__neg__", gen_math_ops._neg) +@tf_export("to_complex64") +def to_complex64(x, name="ToComplex64"): + """Casts a tensor to type `complex64`. + + Args: + x: A `Tensor` or `SparseTensor`. + name: A name for the operation (optional). + + Returns: + A `Tensor` or `SparseTensor` with same shape as `x` with type `complex64`. + + Raises: + TypeError: If `x` cannot be cast to the `complex64`. + """ + return cast(x, dtypes.complex64, name=name) + + +@tf_export("to_complex128") +def to_complex128(x, name="ToComplex128"): + """Casts a tensor to type `complex128`. + + Args: + x: A `Tensor` or `SparseTensor`. + name: A name for the operation (optional). + + Returns: + A `Tensor` or `SparseTensor` with same shape as `x` with type `complex128`. + + Raises: + TypeError: If `x` cannot be cast to the `complex128`. + """ + return cast(x, dtypes.complex128, name=name) + + +ops.Tensor._override_operator("__neg__", gen_math_ops.neg) ops.Tensor._override_operator("__abs__", abs) # __invert__ corresponds to the ~ operator. Here we follow the numpy convention # ~ marks an elementwise bit-wise inverse. This is only implemented for boolean @@ -1027,7 +1067,7 @@ def _truediv_python3(x, y, name=None): if dtype is not None: x = cast(x, dtype) y = cast(y, dtype) - return gen_math_ops._real_div(x, y, name=name) + return gen_math_ops.real_div(x, y, name=name) def _div_python2(x, y, name=None): @@ -1050,9 +1090,9 @@ def _div_python2(x, y, name=None): raise TypeError("x and y must have the same dtype, got %r != %r" % (x_dtype, y_dtype)) if x_dtype.is_floating or x_dtype.is_complex: - return gen_math_ops._real_div(x, y, name=name) + return gen_math_ops.real_div(x, y, name=name) else: - return gen_math_ops._floor_div(x, y, name=name) + return gen_math_ops.floor_div(x, y, name=name) @tf_export("truediv") @@ -1110,7 +1150,7 @@ def div(x, y, name=None): # TODO(aselle): This should be removed -mod = gen_math_ops._floor_mod +mod = gen_math_ops.floor_mod # TODO(aselle): Deprecate this once all internal functionality uses @@ -1143,22 +1183,27 @@ def floordiv(x, y, name=None): TypeError: If the inputs are complex. """ with ops.name_scope(name, "floordiv", [x, y]) as name: - return gen_math_ops._floor_div(x, y, name=name) + return gen_math_ops.floor_div(x, y, name=name) -realdiv = gen_math_ops._real_div -truncatediv = gen_math_ops._truncate_div +realdiv = gen_math_ops.real_div +tf_export("realdiv")(realdiv) +truncatediv = gen_math_ops.truncate_div +tf_export("truncatediv")(truncatediv) # TODO(aselle): Rename this to floordiv when we can. -floor_div = gen_math_ops._floor_div -truncatemod = gen_math_ops._truncate_mod -floormod = gen_math_ops._floor_mod +floor_div = gen_math_ops.floor_div +tf_export("floor_div")(floor_div) +truncatemod = gen_math_ops.truncate_mod +tf_export("truncatemod")(truncatemod) +floormod = gen_math_ops.floor_mod +tf_export("floormod")(floormod) def _mul_dispatch(x, y, name=None): """Dispatches cwise mul for "Dense*Dense" and "Dense*Sparse".""" is_tensor_y = isinstance(y, ops.Tensor) if is_tensor_y: - return gen_math_ops._mul(x, y, name=name) + return gen_math_ops.mul(x, y, name=name) else: assert isinstance(y, sparse_tensor.SparseTensor) # Case: Dense * Sparse. new_vals = gen_sparse_ops.sparse_dense_cwise_mul(y.indices, y.values, @@ -1177,12 +1222,12 @@ _OverrideBinaryOperatorHelper(gen_sparse_ops.sparse_dense_cwise_mul, "mul", sparse_tensor.SparseTensor) _OverrideBinaryOperatorHelper(gen_math_ops.add, "add") -_OverrideBinaryOperatorHelper(gen_math_ops._sub, "sub") +_OverrideBinaryOperatorHelper(gen_math_ops.sub, "sub") _OverrideBinaryOperatorHelper(_mul_dispatch, "mul") _OverrideBinaryOperatorHelper(_div_python2, "div") _OverrideBinaryOperatorHelper(_truediv_python3, "truediv") _OverrideBinaryOperatorHelper(floordiv, "floordiv") -_OverrideBinaryOperatorHelper(gen_math_ops._floor_mod, "mod") +_OverrideBinaryOperatorHelper(gen_math_ops.floor_mod, "mod") _OverrideBinaryOperatorHelper(pow, "pow") @@ -1207,7 +1252,7 @@ ops.Tensor._override_operator("__ge__", gen_math_ops.greater_equal) @tf_export("range") -def range(start, limit=None, delta=1, dtype=None, name="range"): +def range(start, limit=None, delta=1, dtype=None, name="range"): # pylint: disable=redefined-builtin """Creates a sequence of numbers. Creates a sequence of numbers that begins at `start` and extends by @@ -1292,9 +1337,9 @@ def _ReductionDims(x, axis, reduction_indices): return axis else: # Fast path: avoid creating Rank and Range ops if ndims is known. - if isinstance(x, ops.Tensor) and x.get_shape().ndims is not None: + if isinstance(x, ops.Tensor) and x._rank() is not None: # pylint: disable=protected-access return constant_op.constant( - np.arange(x.get_shape().ndims), dtype=dtypes.int32) + np.arange(x._rank()), dtype=dtypes.int32) # pylint: disable=protected-access if (isinstance(x, sparse_tensor.SparseTensor) and x.dense_shape.get_shape().is_fully_defined()): rank = x.dense_shape.get_shape()[0].value # sparse.dense_shape is 1-D. @@ -1504,7 +1549,7 @@ def reduce_mean(input_tensor, if keepdims is None: keepdims = False return _may_reduce_to_scalar(keepdims, axis, reduction_indices, - gen_math_ops._mean( + gen_math_ops.mean( input_tensor, _ReductionDims(input_tensor, axis, reduction_indices), @@ -1554,7 +1599,7 @@ def reduce_prod(input_tensor, if keepdims is None: keepdims = False return _may_reduce_to_scalar(keepdims, axis, reduction_indices, - gen_math_ops._prod( + gen_math_ops.prod( input_tensor, _ReductionDims(input_tensor, axis, reduction_indices), @@ -2007,8 +2052,15 @@ def matmul(a, if transpose_b and adjoint_b: raise ValueError("Only one of transpose_b and adjoint_b can be True.") - a = ops.convert_to_tensor(a, name="a") - b = ops.convert_to_tensor(b, name="b") + if context.executing_eagerly(): + if not isinstance(a, (ops.EagerTensor, _resource_variable_type)): + a = ops.convert_to_tensor(a, name="a") + if not isinstance(b, (ops.EagerTensor, _resource_variable_type)): + b = ops.convert_to_tensor(b, name="b") + else: + a = ops.convert_to_tensor(a, name="a") + b = ops.convert_to_tensor(b, name="b") + # TODO(apassos) remove _shape_tuple here when it is not needed. a_shape = a._shape_tuple() # pylint: disable=protected-access b_shape = b._shape_tuple() # pylint: disable=protected-access @@ -2023,7 +2075,7 @@ def matmul(a, if transpose_b: b = conj(b) adjoint_b = True - return gen_math_ops._batch_mat_mul( + return gen_math_ops.batch_mat_mul( a, b, adj_x=adjoint_a, adj_y=adjoint_b, name=name) # Neither matmul nor sparse_matmul support adjoint, so we conjugate @@ -2060,13 +2112,14 @@ def matmul(a, ret = cast(ret, dtypes.bfloat16) return ret else: - return gen_math_ops._mat_mul( + return gen_math_ops.mat_mul( a, b, transpose_a=transpose_a, transpose_b=transpose_b, name=name) _OverrideBinaryOperatorHelper(matmul, "matmul") -sparse_matmul = gen_math_ops._sparse_mat_mul +sparse_matmul = gen_math_ops.sparse_mat_mul +tf_export("sparse_matmul")(sparse_matmul) @ops.RegisterStatistics("MatMul", "flops") @@ -2171,7 +2224,7 @@ def add_n(inputs, name=None): if name: return array_ops.identity(inputs[0], name=name) return inputs[0] - return gen_math_ops._add_n(inputs, name=name) + return gen_math_ops.add_n(inputs, name=name) @tf_export("accumulate_n") @@ -2181,14 +2234,12 @@ def accumulate_n(inputs, shape=None, tensor_dtype=None, name=None): Optionally, pass `shape` and `tensor_dtype` for shape and type checking, otherwise, these are inferred. - NOTE: This operation is not differentiable and cannot be used if inputs depend - on trainable variables. Please use `tf.add_n` for such cases. + `tf.accumulate_n` performs the same operation as `tf.add_n`, but does not + wait for all of its inputs to be ready before beginning to sum. This can + save memory if inputs are ready at different times, since minimum temporary + storage is proportional to the output size rather than the inputs size. - Aside from differentiability, `tf.accumulate_n` performs the same operation as - `tf.add_n`, but does not wait for all of its inputs to be ready before - beginning to sum. This can save memory if inputs are ready at different times, - since minimum temporary storage is proportional to the output size rather than - the inputs size. + `accumulate_n` is differentiable (but wasn't previous to TensorFlow 1.7). For example: @@ -2198,8 +2249,9 @@ def accumulate_n(inputs, shape=None, tensor_dtype=None, name=None): tf.accumulate_n([a, b, a]) # [[7, 4], [6, 14]] # Explicitly pass shape and type - tf.accumulate_n([a, b, a], shape=[2, 2], tensor_dtype=tf.int32) # [[7, 4], - # [6, 14]] + tf.accumulate_n([a, b, a], shape=[2, 2], tensor_dtype=tf.int32) + # [[7, 4], + # [6, 14]] ``` Args: @@ -2215,20 +2267,17 @@ def accumulate_n(inputs, shape=None, tensor_dtype=None, name=None): ValueError: If `inputs` don't all have same shape and dtype or the shape cannot be inferred. """ - if context.in_eager_mode(): - # TODO(apassos) remove this once the lifetime of eager variables gets - # addressed. - raise ValueError("accumulate_n not supported in eager mode") + def _input_error(): + return ValueError( + "inputs must be a list of at least one Tensor with the " + "same dtype and shape") if not inputs or not isinstance(inputs, (list, tuple)): - raise ValueError("inputs must be a list of at least one Tensor with the " - "same dtype and shape") + raise _input_error() inputs = ops.convert_n_to_tensor_or_indexed_slices(inputs) if not all(isinstance(x, ops.Tensor) for x in inputs): - raise ValueError("inputs must be a list of at least one Tensor with the " - "same dtype and shape") + raise _input_error() if not all(x.dtype == inputs[0].dtype for x in inputs): - raise ValueError("inputs must be a list of at least one Tensor with the " - "same dtype and shape") + raise _input_error() if shape is not None: shape = tensor_shape.as_shape(shape) else: @@ -2236,27 +2285,31 @@ def accumulate_n(inputs, shape=None, tensor_dtype=None, name=None): for input_tensor in inputs: if isinstance(input_tensor, ops.Tensor): shape = shape.merge_with(input_tensor.get_shape()) - if tensor_dtype is None: - tensor_dtype = inputs[0].dtype - if tensor_dtype != inputs[0].dtype: - raise TypeError("tensor_dtype is {}, but input is of type {}".format( - tensor_dtype, inputs[0].dtype)) - if len(inputs) == 1: + + # tensor_dtype is for safety only; operator's output type computed in C++ + if tensor_dtype is not None and tensor_dtype != inputs[0].dtype: + raise TypeError("tensor_dtype is {}, but input is of type {}" + .format(tensor_dtype, inputs[0].dtype)) + + if len(inputs) == 1 and name is None: return inputs[0] - with ops.name_scope(name, "AccumulateN", inputs) as name: - var = gen_state_ops._temporary_variable( - shape=tensor_shape.vector(0), dtype=tensor_dtype) - with ops.colocate_with(var): - zeros = array_ops.zeros_like(gen_control_flow_ops._merge(inputs)[0]) - zeros.set_shape(shape) - ref = state_ops.assign(var, zeros, validate_shape=False) - update_ops = [ - state_ops.assign_add(ref, input_tensor, use_locking=True) - for input_tensor in inputs - ] - with ops.control_dependencies(update_ops): - return gen_state_ops._destroy_temporary_variable( - ref, var_name=var.op.name, name=name) + elif len(inputs) == 1 and name is not None: + return array_ops.identity(inputs[0], name=name) + elif context.executing_eagerly(): + # TemporaryVariable not currently supported in eager mode; fall back + # onto AddN for now. + # TODO(frreiss) remove this once the lifetime of eager variables gets + # addressed + return add_n(inputs, name=name) + else: + return gen_math_ops.accumulate_nv2(inputs, name=name, shape=shape) # pylint: disable=protected-access + + +@ops.RegisterGradient("AccumulateNV2") +def _accumulate_n_grad(op, grad): + """Same as gradient for AddN. Copies the gradient to all inputs.""" + # Not broadcasting. + return [grad] * len(op.inputs) @tf_export("nn.sigmoid", "sigmoid") @@ -2279,7 +2332,7 @@ def sigmoid(x, name=None): """ with ops.name_scope(name, "Sigmoid", [x]) as name: x = ops.convert_to_tensor(x, name="x") - return gen_math_ops._sigmoid(x, name=name) + return gen_math_ops.sigmoid(x, name=name) @tf_export("log_sigmoid") @@ -2298,7 +2351,7 @@ def log_sigmoid(x, name=None): """ with ops.name_scope(name, "LogSigmoid", [x]) as name: x = ops.convert_to_tensor(x, name="x") - return gen_math_ops._neg(gen_nn_ops.softplus(-x), name=name) + return gen_math_ops.neg(gen_nn_ops.softplus(-x), name=name) @tf_export("nn.tanh", "tanh") @@ -2315,11 +2368,11 @@ def tanh(x, name=None): """ with ops.name_scope(name, "Tanh", [x]) as name: if isinstance(x, sparse_tensor.SparseTensor): - x_tanh = gen_math_ops._tanh(x.values, name=name) + x_tanh = gen_math_ops.tanh(x.values, name=name) return sparse_tensor.SparseTensor( indices=x.indices, values=x_tanh, dense_shape=x.dense_shape) else: - return gen_math_ops._tanh(x, name=name) + return gen_math_ops.tanh(x, name=name) @tf_export("bincount") @@ -2508,7 +2561,7 @@ def conj(x, name=None): with ops.name_scope(name, "Conj", [x]) as name: x = ops.convert_to_tensor(x, name="x") if x.dtype.is_complex or x.dtype == dtypes.variant: - return gen_math_ops._conj(x, name=name) + return gen_math_ops.conj(x, name=name) elif x.dtype.is_floating or x.dtype.is_integer: return x else: @@ -2552,12 +2605,93 @@ def reduced_shape(input_shape, axes): ]) # [1, 1] +def _unsorted_segment_N(data, segment_ids, num_segments): + """ Helper function for unsorted_segment_mean/_sqrtN. Computes the number + of segment entries with 0-entries set to 1 to allow division by N. + """ + # bincount doesn't support negative indices so we use unsorted_segment_sum + ones_tensor = array_ops.ones(segment_ids.shape, dtype=data.dtype) + N = gen_math_ops.unsorted_segment_sum(ones_tensor, segment_ids, num_segments) + # add dimensions for all non-reduced axes + ndims_output = data.shape.ndims - segment_ids.shape.ndims + broadcast_shape = [num_segments] + [1] * ndims_output + N = array_ops.reshape(N, broadcast_shape) + return gen_math_ops.maximum(N, 1) + + +@tf_export("unsorted_segment_mean") +def unsorted_segment_mean(data, segment_ids, num_segments, name=None): + r""" Computes the mean along segments of a tensor. + + Read @{$math_ops#segmentation$the section on segmentation} for an explanation + of segments. + + This operator is similar to the unsorted segment sum operator found + [here](../../../api_docs/python/math_ops.md#UnsortedSegmentSum). + Instead of computing the sum over segments, it computes the mean of all + entries belonging to a segment such that: + + \\(output_i = 1/N_i \sum data_j\\) where the sum is over `j` such + that `segment_ids[j] == i` with \\N_i\\ being the number of occurrences + of id \\i\\. + + If there is no entry for a given segment ID `i`, it outputs 0. + + segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s + first dimension. + + output: Has same shape as data, except for dimension 0 which + has size `num_segments`. + """ + with ops.name_scope(name, "UnsortedSegmentMean"): + data = ops.convert_to_tensor(data) + segment_ids = ops.convert_to_tensor(segment_ids) + N = _unsorted_segment_N(data, segment_ids, num_segments) + summed = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments) + return summed / N + + +@tf_export("unsorted_segment_sqrt_n") +def unsorted_segment_sqrt_n(data, segment_ids, num_segments, name=None): + r"""Computes the sum along segments of a tensor divided by the sqrt(N). + + Read @{$math_ops#segmentation$the section on segmentation} for an explanation + of segments. + + This operator is similar to the unsorted segment sum operator found + [here](../../../api_docs/python/math_ops.md#UnsortedSegmentSum). + Additionally to computing the sum over segments, it divides the results by + sqrt(N). + + \\(output_i = 1/sqrt(N_i) \sum data_j\\) where the sum is over `j` such + that `segment_ids[j] == i` with \\N_i\\ being the number of occurrences + of id \\i\\. + + If there is no entry for a given segment ID `i`, it outputs 0. + + Note that this op only supports floating point and complex dtypes, + due to tf.sqrt only supporting these types. + + segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s + first dimension. + + output: Has same shape as data, except for dimension 0 which + has size `num_segments`. + """ + with ops.name_scope(name, "UnsortedSegmentSqrtN"): + data = ops.convert_to_tensor(data) + segment_ids = ops.convert_to_tensor(segment_ids) + N = _unsorted_segment_N(data, segment_ids, num_segments) + summed = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments) + return summed / gen_math_ops.sqrt(N) + + @tf_export("sparse_segment_sum") def sparse_segment_sum(data, indices, segment_ids, name=None, num_segments=None): r"""Computes the sum along sparse segments of a tensor. - Read @{$math_ops#segmentation$the section on segmentation} for an explanation + Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of segments. Like `SegmentSum`, but `segment_ids` can have rank less than `data`'s first @@ -2632,7 +2766,7 @@ def sparse_segment_mean(data, indices, segment_ids, name=None, num_segments=None): r"""Computes the mean along sparse segments of a tensor. - Read @{$math_ops#segmentation$the section on segmentation} for an explanation + Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of segments. Like `SegmentMean`, but `segment_ids` can have rank less than `data`'s first diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py index bd26ff66961c858865c8a61469abac0b783ed645..9f85188b3513563a7444f7a0e908f11af985498b 100644 --- a/tensorflow/python/ops/math_ops_test.py +++ b/tensorflow/python/ops/math_ops_test.py @@ -60,7 +60,7 @@ class ReduceTest(test_util.TensorFlowTestCase): @test_util.run_in_graph_and_eager_modes() def testReduceInvalidAxis(self): - if context.in_eager_mode(): + if context.executing_eagerly(): # The shape check is in run a graph construction time. In eager mode, # it misses the check, magically return result given wrong shape. return @@ -105,7 +105,7 @@ class LogSumExpTest(test_util.TensorFlowTestCase): for dtype in [np.float16, np.float32, np.double]: x_np = np.random.rand(5, 5).astype(dtype) with self.test_session(use_gpu=True): - y_tf_np = math_ops.reduce_logsumexp(x_np, keep_dims=True).eval() + y_tf_np = math_ops.reduce_logsumexp(x_np, keepdims=True).eval() self.assertEqual(y_tf_np.ndim, x_np.ndim) y_np = log(np.sum(exp(x_np), keepdims=True)) self.assertAllClose(y_tf_np, y_np) @@ -249,7 +249,7 @@ class ScalarMulTest(test_util.TensorFlowTestCase): @test_util.run_in_graph_and_eager_modes() def testAcceptsRefs(self): - if context.in_eager_mode(): + if context.executing_eagerly(): var = resource_variable_ops.ResourceVariable(10, name="var") else: var = variables.Variable(10) diff --git a/tensorflow/python/ops/metrics_impl.py b/tensorflow/python/ops/metrics_impl.py index 7776ff08c4f55c43947010f313d8167596b15db7..9ec49545796cfa7a603b31c23bfd0d495639898d 100644 --- a/tensorflow/python/ops/metrics_impl.py +++ b/tensorflow/python/ops/metrics_impl.py @@ -308,7 +308,7 @@ def mean(values, or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean is not supported when eager execution ' 'is enabled.') @@ -394,7 +394,7 @@ def accuracy(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.accuracy is not supported when eager ' 'execution is enabled.') @@ -644,7 +644,7 @@ def auc(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.auc is not supported when eager execution ' 'is enabled.') @@ -758,7 +758,7 @@ def mean_absolute_error(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_absolute_error is not supported ' 'when eager execution is enabled.') @@ -818,7 +818,7 @@ def mean_cosine_distance(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_cosine_distance is not supported when ' 'eager execution is enabled.') @@ -891,7 +891,7 @@ def mean_per_class_accuracy(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_per_class_accuracy is not supported ' 'when eager execution is enabled.') @@ -996,7 +996,7 @@ def mean_iou(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_iou is not supported when ' 'eager execution is enabled.') @@ -1098,7 +1098,7 @@ def mean_relative_error(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_relative_error is not supported when ' 'eager execution is enabled.') @@ -1165,7 +1165,7 @@ def mean_squared_error(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_squared_error is not supported when ' 'eager execution is enabled.') @@ -1223,7 +1223,7 @@ def mean_tensor(values, or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.mean_tensor is not supported when ' 'eager execution is enabled.') @@ -1247,13 +1247,8 @@ def mean_tensor(values, with ops.control_dependencies([values]): update_count_op = state_ops.assign_add(count, num_values) - def compute_mean(total, count, name): - non_zero_count = math_ops.maximum( - count, array_ops.ones_like(count), name=name) - return math_ops.truediv(total, non_zero_count, name=name) - - mean_t = compute_mean(total, count, 'value') - update_op = compute_mean(update_total_op, update_count_op, 'update_op') + mean_t = _safe_div(total, count, 'value') + update_op = _safe_div(update_total_op, update_count_op, 'update_op') if metrics_collections: ops.add_to_collections(metrics_collections, mean_t) @@ -1309,7 +1304,7 @@ def percentage_below(values, or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.percentage_below is not supported when ' 'eager execution is enabled.') @@ -1402,7 +1397,7 @@ def false_negatives(labels, or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.false_negatives is not supported when ' 'eager execution is enabled.') @@ -1458,7 +1453,7 @@ def false_negatives_at_thresholds(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.false_negatives_at_thresholds is not ' 'supported when eager execution is enabled.') @@ -1512,7 +1507,7 @@ def false_positives(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.false_positives is not supported when ' 'eager execution is enabled.') @@ -1568,7 +1563,7 @@ def false_positives_at_thresholds(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.false_positives_at_thresholds is not ' 'supported when eager execution is enabled.') @@ -1622,7 +1617,7 @@ def true_negatives(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.true_negatives is not ' 'supported when eager execution is enabled.') @@ -1678,7 +1673,7 @@ def true_negatives_at_thresholds(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.true_negatives_at_thresholds is not ' 'supported when eager execution is enabled.') @@ -1732,7 +1727,7 @@ def true_positives(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.true_positives is not ' 'supported when eager execution is enabled.') @@ -1788,7 +1783,7 @@ def true_positives_at_thresholds(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.true_positives_at_thresholds is not ' 'supported when eager execution is enabled.') @@ -1856,7 +1851,7 @@ def precision(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.precision is not ' 'supported when eager execution is enabled.') @@ -1952,7 +1947,7 @@ def precision_at_thresholds(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.precision_at_thresholds is not ' 'supported when eager execution is enabled.') @@ -2028,7 +2023,7 @@ def recall(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.recall is not supported is not ' 'supported when eager execution is enabled.') @@ -2405,7 +2400,7 @@ def recall_at_k(labels, are not a list or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.recall_at_k is not ' 'supported when eager execution is enabled.') @@ -2554,7 +2549,7 @@ def recall_at_thresholds(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.recall_at_thresholds is not ' 'supported when eager execution is enabled.') @@ -2631,7 +2626,7 @@ def root_mean_squared_error(labels, tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.root_mean_squared_error is not ' 'supported when eager execution is enabled.') @@ -2712,7 +2707,7 @@ def sensitivity_at_specificity(labels, or `updates_collections` are not a list or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.sensitivity_at_specificity is not ' 'supported when eager execution is enabled.') @@ -3103,7 +3098,7 @@ def average_precision_at_k(labels, ValueError: if k is invalid. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.sparse_average_precision_at_k is not ' 'supported when eager execution is enabled.') @@ -3272,7 +3267,7 @@ def precision_at_top_k(labels, are not a list or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.precision_at_top_k is not ' 'supported when eager execution is enabled.') @@ -3401,7 +3396,7 @@ def precision_at_k(labels, are not a list or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.sparse_precision_at_k is not ' 'supported when eager execution is enabled.') @@ -3478,7 +3473,7 @@ def specificity_at_sensitivity(labels, or `updates_collections` are not a list or tuple. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError('tf.metrics.specificity_at_sensitivity is not ' 'supported when eager execution is enabled.') diff --git a/tensorflow/python/ops/nn_batchnorm_test.py b/tensorflow/python/ops/nn_batchnorm_test.py index eebfb17085a568f48769f6df7dddd3ae2f799efc..3ac2c8eb17ef31b46638ce50e0e9f9705adce189 100644 --- a/tensorflow/python/ops/nn_batchnorm_test.py +++ b/tensorflow/python/ops/nn_batchnorm_test.py @@ -57,7 +57,6 @@ class BatchNormalizationTest(test.TestCase): test_util.set_producer_version(ops.get_default_graph(), 8) return gen_nn_ops._batch_norm_with_global_normalization( x, m, v, beta, gamma, epsilon, scale_after_normalization) - # pylint: enable=protected-access def _tfBatchNormV1BW(self, x, m, v, beta, gamma, epsilon, scale_after_normalization): @@ -223,7 +222,7 @@ class BatchNormalizationTest(test.TestCase): for scale_after_normalization in [True, False]: # _batch_norm_with_global_normalization_grad is deprecated in v9 test_util.set_producer_version(ops.get_default_graph(), 8) - grad = gen_nn_ops._batch_norm_with_global_normalization_grad( + grad = gen_nn_ops.batch_norm_with_global_normalization_grad( x, m, v, gamma, backprop, epsilon, scale_after_normalization) dx, dm, dv, db, dg = grad self.assertEqual(grad.dx, dx) diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index 5e6cafd6aad4a80d3436d9c270ac5e2341c2a5aa..4af5bd26dd80b984b1c898411c2a23827bed1b4b 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -150,7 +150,7 @@ def _Conv3DBackpropFilterGrad(op, grad): @ops.RegisterGradient("AvgPool3D") def _AvgPool3DGrad(op, grad): - return gen_nn_ops._avg_pool3d_grad( + return gen_nn_ops.avg_pool3d_grad( array_ops.shape(op.inputs[0]), grad, ksize=op.get_attr("ksize"), @@ -172,7 +172,7 @@ def _AvgPool3DGradGrad(op, grad): @ops.RegisterGradient("MaxPool3D") def _MaxPool3DGrad(op, grad): - return gen_nn_ops._max_pool3d_grad( + return gen_nn_ops.max_pool3d_grad( op.inputs[0], op.outputs[0], grad, @@ -188,7 +188,7 @@ def _MaxPool3DGradGrad(op, grad): shape=array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype), array_ops.zeros( shape=array_ops.shape(op.inputs[1]), dtype=op.inputs[1].dtype), - gen_nn_ops._max_pool3d_grad_grad( + gen_nn_ops.max_pool3d_grad_grad( op.inputs[0], op.inputs[1], grad, @@ -204,7 +204,7 @@ def _MaxPool3DGradGradGrad(op, grad): shape=array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype), array_ops.zeros( shape=array_ops.shape(op.inputs[1]), dtype=op.inputs[1].dtype), - gen_nn_ops._max_pool3d_grad( + gen_nn_ops.max_pool3d_grad( op.inputs[0], op.inputs[1], grad, @@ -352,13 +352,13 @@ def _BiasAddGradV1(unused_bias_op, received_grad): @ops.RegisterGradient("Relu") def _ReluGrad(op, grad): - return gen_nn_ops._relu_grad(grad, op.outputs[0]) + return gen_nn_ops.relu_grad(grad, op.outputs[0]) @ops.RegisterGradient("EluGrad") def _EluGradGrad(op, grad): elu_x = op.inputs[1] - return (gen_nn_ops._elu_grad(grad, op.outputs[0]), + return (gen_nn_ops.elu_grad(grad, op.outputs[0]), array_ops.where(elu_x < 0, grad * op.inputs[0], array_ops.zeros( shape=array_ops.shape(elu_x), dtype=elu_x.dtype))) @@ -368,63 +368,63 @@ def _EluGradGrad(op, grad): def _SeluGradGrad(op, grad): x = op.inputs[1] scale_alpha = 1.7580993408473768599402175208123 - return (gen_nn_ops._elu_grad(grad, op.outputs[0]), + return (gen_nn_ops.elu_grad(grad, op.outputs[0]), array_ops.where(x < 0., - gen_nn_ops._elu_grad(grad, - op.outputs[0] + scale_alpha), + gen_nn_ops.elu_grad(grad, + op.outputs[0] + scale_alpha), array_ops.zeros( shape=array_ops.shape(x), dtype=x.dtype))) @ops.RegisterGradient("Relu6") def _Relu6Grad(op, grad): - return gen_nn_ops._relu6_grad(grad, op.outputs[0]) # pylint: disable=protected-access + return gen_nn_ops.relu6_grad(grad, op.outputs[0]) @ops.RegisterGradient("Relu6Grad") def _Relu6GradGrad(op, grad): x = op.inputs[1] - return (gen_nn_ops._relu6_grad(grad, x), + return (gen_nn_ops.relu6_grad(grad, x), array_ops.zeros(shape=array_ops.shape(x), dtype=x.dtype)) @ops.RegisterGradient("Elu") def _EluGrad(op, grad): - return gen_nn_ops._elu_grad(grad, op.outputs[0]) + return gen_nn_ops.elu_grad(grad, op.outputs[0]) @ops.RegisterGradient("Selu") def _SeluGrad(op, grad): - return gen_nn_ops._selu_grad(grad, op.outputs[0]) + return gen_nn_ops.selu_grad(grad, op.outputs[0]) @ops.RegisterGradient("Softplus") def _SoftplusGrad(op, grad): - return gen_nn_ops._softplus_grad(grad, op.inputs[0]) + return gen_nn_ops.softplus_grad(grad, op.inputs[0]) @ops.RegisterGradient("SoftplusGrad") def _SoftplusGradGrad(op, grad): # Let: # y = tf.nn.softplus(x) - # dx = gen_nn_ops._softplus_grad(dy, x) = dy / (1 + exp(-x)) + # dx = gen_nn_ops.softplus_grad(dy, x) = dy / (1 + exp(-x)) # This op computes (ddy, d2x) from op.inputs == [dy, x] and grad == ddx. dy, x = op.inputs with ops.control_dependencies([grad]): - ddy = gen_nn_ops._softplus_grad(grad, x) # pylint: disable=protected-access + ddy = gen_nn_ops.softplus_grad(grad, x) d2x = grad * dy / (math_ops.exp(-x) + 2.0 + math_ops.exp(x)) return (ddy, d2x) @ops.RegisterGradient("Softsign") def _SoftsignGrad(op, grad): - return gen_nn_ops._softsign_grad(grad, op.inputs[0]) + return gen_nn_ops.softsign_grad(grad, op.inputs[0]) @ops.RegisterGradient("ReluGrad") def _ReluGradGrad(op, grad): x = op.inputs[1] - return (gen_nn_ops._relu_grad(grad, x), + return (gen_nn_ops.relu_grad(grad, x), array_ops.zeros(shape=array_ops.shape(x), dtype=x.dtype)) @@ -456,7 +456,7 @@ def _SoftmaxCrossEntropyWithLogitsGrad(op, grad_loss, grad_grad): def IsZero(g): # Some introspection to check if the gradient is feeding zeros - if context.in_eager_mode(): + if context.executing_eagerly(): # TODO(apassos) add an efficient way to detect eager zeros here. return False if g.op.type in ("ZerosLike", "Zeros"): @@ -565,14 +565,14 @@ def _LRNGrad(op, grad): alpha = op.get_attr("alpha") beta = op.get_attr("beta") return [ - gen_nn_ops._lrn_grad(grad, op.inputs[0], op.outputs[0], depth_radius, - bias, alpha, beta) + gen_nn_ops.lrn_grad(grad, op.inputs[0], op.outputs[0], depth_radius, bias, + alpha, beta) ] @ops.RegisterGradient("AvgPool") def _AvgPoolGrad(op, grad): - return gen_nn_ops._avg_pool_grad( + return gen_nn_ops.avg_pool_grad( array_ops.shape(op.inputs[0]), grad, op.get_attr("ksize"), @@ -584,7 +584,7 @@ def _AvgPoolGrad(op, grad): @ops.RegisterGradient("AvgPoolGrad") def _AvgPoolGradGrad(op, grad): return (array_ops.stop_gradient(op.inputs[0]), - gen_nn_ops._avg_pool( + gen_nn_ops.avg_pool( grad, op.get_attr("ksize"), op.get_attr("strides"), @@ -594,7 +594,7 @@ def _AvgPoolGradGrad(op, grad): @ops.RegisterGradient("MaxPool") def _MaxPoolGrad(op, grad): - return gen_nn_ops._max_pool_grad( + return gen_nn_ops.max_pool_grad( op.inputs[0], op.outputs[0], grad, @@ -620,7 +620,7 @@ def _MaxPoolGradV2(op, grad): @ops.RegisterGradient("MaxPoolWithArgmax") def _MaxPoolGradWithArgmax(op, grad, unused_argmax_grad): - return gen_nn_ops._max_pool_grad_with_argmax( + return gen_nn_ops.max_pool_grad_with_argmax( op.inputs[0], grad, op.outputs[1], @@ -635,7 +635,7 @@ def _MaxPoolGradGrad(op, grad): shape=array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype), array_ops.zeros( shape=array_ops.shape(op.inputs[1]), dtype=op.inputs[1].dtype), - gen_nn_ops._max_pool_grad_grad( + gen_nn_ops.max_pool_grad_grad( op.inputs[0], op.inputs[1], grad, @@ -669,7 +669,7 @@ def _MaxPoolGradGradGrad(op, grad): shape=array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype), array_ops.zeros( shape=array_ops.shape(op.inputs[1]), dtype=op.inputs[1].dtype), - gen_nn_ops._max_pool_grad( + gen_nn_ops.max_pool_grad( op.inputs[0], op.inputs[1], grad, @@ -696,8 +696,7 @@ def _FractionalMaxPoolGrad(op, grad_0, unused_grad_1, unused_grad_2): Returns: Input backprop for FractionalMaxPool op. """ - # pylint: disable=protected-access - return gen_nn_ops._fractional_max_pool_grad( + return gen_nn_ops.fractional_max_pool_grad( op.inputs[0], op.outputs[0], grad_0, op.outputs[1], op.outputs[2], op.get_attr("overlapping")) @@ -719,10 +718,9 @@ def _FractionalAvgPoolGrad(op, grad_0, unused_grad_1, unused_grad_2): Returns: Input backprop for FractionalAvgPool op. """ - # pylint: disable=protected-access - return gen_nn_ops._fractional_avg_pool_grad(op.inputs[0].get_shape(), grad_0, - op.outputs[1], op.outputs[2], - op.get_attr("overlapping")) + return gen_nn_ops.fractional_avg_pool_grad(op.inputs[0].get_shape(), grad_0, + op.outputs[1], op.outputs[2], + op.get_attr("overlapping")) @ops.RegisterGradient("BatchNormWithGlobalNormalization") @@ -746,7 +744,7 @@ def _BatchNormWithGlobalNormalizationGrad(op, grad): last dimension. dg: Backprop for gamma, which is (grad * ((x - m) * rsqrt(v + epsilon))) """ - dx, dm, dv, db, dg = gen_nn_ops._batch_norm_with_global_normalization_grad( + dx, dm, dv, db, dg = gen_nn_ops.batch_norm_with_global_normalization_grad( op.inputs[0], op.inputs[1], op.inputs[2], op.inputs[4], grad, op.get_attr("variance_epsilon"), op.get_attr("scale_after_normalization")) return dx, dm, dv, db, dg @@ -863,27 +861,27 @@ def _BatchNormGrad(grad_y, grad_y = math_ops.cast(grad_y, dtypes.float32) if is_training: if data_format == b"NHWC": - keep_dims = False + keepdims = False reduce_axis = [0, 1, 2] else: - keep_dims = True + keepdims = True reduce_axis = [0, 2, 3] shape = [1, array_ops.size(scale), 1, 1] scale = array_ops.reshape(scale, shape) - mean_grad_y = math_ops.reduce_mean(grad_y, reduce_axis, keep_dims=keep_dims) - mean_x = math_ops.reduce_mean(x, reduce_axis, keep_dims=keep_dims) + mean_grad_y = math_ops.reduce_mean(grad_y, reduce_axis, keepdims=keepdims) + mean_x = math_ops.reduce_mean(x, reduce_axis, keepdims=keepdims) var_x = math_ops.reduce_mean( math_ops.squared_difference(x, array_ops.stop_gradient(mean_x)), reduce_axis, - keep_dims=keep_dims) + keepdims=keepdims) grad_y_offset = grad_y - mean_grad_y x_offset = x - mean_x mean = math_ops.reduce_mean( - grad_y * x_offset, axis=reduce_axis, keep_dims=keep_dims) + grad_y * x_offset, axis=reduce_axis, keepdims=keepdims) grad_x = scale * math_ops.rsqrt(var_x + epsilon) * ( grad_y_offset - math_ops.reciprocal(var_x + epsilon) * mean * x_offset) grad_scale = math_ops.rsqrt(var_x + epsilon) * math_ops.reduce_sum( - grad_y * x_offset, axis=reduce_axis, keep_dims=keep_dims) + grad_y * x_offset, axis=reduce_axis, keepdims=keepdims) if data_format == b"NCHW": grad_scale = array_ops.squeeze(grad_scale) grad_offset = math_ops.reduce_sum(grad_y, axis=reduce_axis) @@ -1010,7 +1008,7 @@ def _NthElementGrad(op, grad): A list of two tensors, the first being the gradient w.r.t. the input, the second being the gradient w.r.t. the N (None). """ - input = op.inputs[0] + input = op.inputs[0] # pylint: disable=redefined-builtin output = op.outputs[0] # Compute the number of elements which equal to output in each reduction diff --git a/tensorflow/python/ops/nn_grad_test.py b/tensorflow/python/ops/nn_grad_test.py index aa7539ae9f09163bf1a2cc9f7dfcc6fc06737ae8..49d54beb20073162279576e1e1011e10392378e0 100644 --- a/tensorflow/python/ops/nn_grad_test.py +++ b/tensorflow/python/ops/nn_grad_test.py @@ -24,7 +24,7 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import nn_grad +from tensorflow.python.ops import nn_grad # pylint: disable=unused-import from tensorflow.python.ops import nn_ops from tensorflow.python.platform import test diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py index 55fcd176d62009b9c29afb763dc20daf78cdb5d9..47cc4da7f2abd1f5b00e193a76c8391be94ca27d 100644 --- a/tensorflow/python/ops/nn_impl.py +++ b/tensorflow/python/ops/nn_impl.py @@ -27,7 +27,7 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import candidate_sampling_ops from tensorflow.python.ops import embedding_ops -from tensorflow.python.ops import gen_array_ops +from tensorflow.python.ops import gen_array_ops # pylint: disable=unused-import from tensorflow.python.ops import gen_nn_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops @@ -303,12 +303,12 @@ def _swish_grad(features, grad): # @Defun decorator with noinline=True so that sigmoid(features) is re-computed # during backprop, and we can free the sigmoid(features) expression immediately # after use during the forward pass. +@tf_export("nn.swish") @function.Defun( grad_func=_swish_grad, shape_func=_swish_shape, func_name="swish", noinline=True) -@tf_export("nn.swish") def swish(features): # pylint: disable=g-doc-args """Computes the Swish activation function: `x * sigmoid(x)`. @@ -888,12 +888,10 @@ def fused_batch_norm( # TODO(reedwm): In a few weeks, switch to using the V2 version exclusively. We # currently only use the V2 version for float16 inputs, which is not supported # by the V1 version. - # pylint: disable=protected-access if x.dtype == dtypes.float16 or x.dtype == dtypes.bfloat16: - fused_batch_norm_func = gen_nn_ops._fused_batch_norm_v2 + fused_batch_norm_func = gen_nn_ops.fused_batch_norm_v2 else: - fused_batch_norm_func = gen_nn_ops._fused_batch_norm - # pylint: enable=protected-access + fused_batch_norm_func = gen_nn_ops._fused_batch_norm # pylint: disable=protected-access y, batch_mean, batch_var, _, _ = fused_batch_norm_func( x, scale, diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py index a691e281ee7f0ce4ee3253069b095a88df5723c5..a74de39eab34a1a27df90f70adf0f4c68ec29465 100644 --- a/tensorflow/python/ops/nn_ops.py +++ b/tensorflow/python/ops/nn_ops.py @@ -29,6 +29,7 @@ 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 array_ops +from tensorflow.python.ops import check_ops from tensorflow.python.ops import gen_nn_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops @@ -48,8 +49,8 @@ local_response_normalization = gen_nn_ops.lrn def _non_atrous_convolution( - input, - filter, + input, # pylint: disable=redefined-builtin + filter, # pylint: disable=redefined-builtin padding, data_format=None, # pylint: disable=redefined-builtin strides=None, @@ -94,9 +95,9 @@ def _non_atrous_convolution( """ with ops.name_scope(name, "non_atrous_convolution", [input, filter]) as scope: - input = ops.convert_to_tensor(input, name="input") + input = ops.convert_to_tensor(input, name="input") # pylint: disable=redefined-builtin input_shape = input.get_shape() - filter = ops.convert_to_tensor(filter, name="filter") + filter = ops.convert_to_tensor(filter, name="filter") # pylint: disable=redefined-builtin filter_shape = filter.get_shape() op = _NonAtrousConvolution( input_shape, @@ -149,14 +150,12 @@ class _NonAtrousConvolution(object): conv_dims)) if conv_dims == 1: # conv1d uses the 2-d data format names - if data_format is None or data_format == "NWC": - data_format_2d = "NHWC" - elif data_format == "NCW": - data_format_2d = "NCHW" - else: + if data_format is None: + data_format = "NWC" + elif data_format not in {"NCW", "NWC", "NCHW", "NHWC"}: raise ValueError("data_format must be \"NWC\" or \"NCW\".") self.strides = strides[0] - self.data_format = data_format_2d + self.data_format = data_format self.conv_op = self._conv1d elif conv_dims == 2: if data_format is None or data_format == "NHWC": @@ -348,7 +347,7 @@ def with_space_to_batch( ValueError: if `spatial_dims` are invalid. """ - input = ops.convert_to_tensor(input, name="input") + input = ops.convert_to_tensor(input, name="input") # pylint: disable=redefined-builtin input_shape = input.get_shape() def build_op(num_spatial_dims, padding): @@ -645,7 +644,7 @@ def _get_strides_and_dilation_rate(num_spatial_dims, strides, dilation_rate): @tf_export("nn.convolution") def convolution( - input, + input, # pylint: disable=redefined-builtin filter, # pylint: disable=redefined-builtin padding, strides=None, @@ -698,7 +697,7 @@ def convolution( `padded_input` is obtained by zero padding the input using an effective spatial filter shape of `(spatial_filter_shape-1) * dilation_rate + 1` and output striding `strides` as described in the - @{tf.nn.convolution$comment here}. + @{$python/nn#Convolution$comment here}. In the case that `data_format` does start with `"NC"`, the `input` and output (but not the `filter`) are simply transposed as follows: @@ -766,9 +765,9 @@ def convolution( """ # pylint: enable=line-too-long with ops.name_scope(name, "convolution", [input, filter]) as name: - input = ops.convert_to_tensor(input, name="input") + input = ops.convert_to_tensor(input, name="input") # pylint: disable=redefined-builtin input_shape = input.get_shape() - filter = ops.convert_to_tensor(filter, name="filter") + filter = ops.convert_to_tensor(filter, name="filter") # pylint: disable=redefined-builtin filter_shape = filter.get_shape() op = Convolution( input_shape, @@ -962,7 +961,7 @@ def pool( # pylint: enable=line-too-long with ops.name_scope(name, "%s_pool" % (pooling_type.lower()), [input]) as scope: - input = ops.convert_to_tensor(input, name="input") + input = ops.convert_to_tensor(input, name="input") # pylint: disable=redefined-builtin num_spatial_dims = len(window_shape) if num_spatial_dims < 1 or num_spatial_dims > 3: @@ -1042,9 +1041,7 @@ def pool( @tf_export("nn.atrous_conv2d") def atrous_conv2d(value, filters, rate, padding, name=None): - """Atrous convolution (a.k.a. - - convolution with holes or dilated convolution). + """Atrous convolution (a.k.a. convolution with holes or dilated convolution). This function is a simpler wrapper around the more general @{tf.nn.convolution}, and exists only for backwards compatibility. You can @@ -1223,7 +1220,7 @@ def conv2d_transpose( if data_format not in ("NCHW", "NHWC"): raise ValueError("data_format has to be either NCHW or NHWC.") value = ops.convert_to_tensor(value, name="value") - filter = ops.convert_to_tensor(filter, name="filter") + filter = ops.convert_to_tensor(filter, name="filter") # pylint: disable=redefined-builtin axis = 3 if data_format == "NHWC" else 1 if not value.get_shape()[axis].is_compatible_with(filter.get_shape()[3]): raise ValueError("input channels does not match filter's input channels, " @@ -1447,7 +1444,7 @@ def conv3d_transpose( with ops.name_scope(name, "conv3d_transpose", [value, filter, output_shape]) as name: value = ops.convert_to_tensor(value, name="value") - filter = ops.convert_to_tensor(filter, name="filter") + filter = ops.convert_to_tensor(filter, name="filter") # pylint: disable=redefined-builtin axis = 1 if data_format == "NCDHW" else 4 if not value.get_shape()[axis].is_compatible_with(filter.get_shape()[4]): raise ValueError("input channels does not match filter's input channels, " @@ -1481,7 +1478,6 @@ def conv3d_transpose( name=name) -# pylint: disable=protected-access @tf_export("nn.bias_add") def bias_add(value, bias, data_format=None, name=None): """Adds `bias` to `value`. @@ -1504,12 +1500,12 @@ def bias_add(value, bias, data_format=None, name=None): A `Tensor` with the same type as `value`. """ with ops.name_scope(name, "BiasAdd", [value, bias]) as name: - value = ops.convert_to_tensor(value, name="input") - bias = ops.convert_to_tensor(bias, dtype=value.dtype, name="bias") - return gen_nn_ops._bias_add(value, bias, data_format=data_format, name=name) + if not context.executing_eagerly(): + value = ops.convert_to_tensor(value, name="input") + bias = ops.convert_to_tensor(bias, dtype=value.dtype, name="bias") + return gen_nn_ops.bias_add(value, bias, data_format=data_format, name=name) -# pylint: disable=protected-access def bias_add_v1(value, bias, name=None): """Adds `bias` to `value`. @@ -1534,7 +1530,7 @@ def bias_add_v1(value, bias, name=None): with ops.name_scope(name, "BiasAddV1", [value, bias]) as name: value = ops.convert_to_tensor(value, name="input") bias = ops.convert_to_tensor(bias, dtype=value.dtype, name="bias") - return gen_nn_ops._bias_add_v1(value, bias, name=name) + return gen_nn_ops.bias_add_v1(value, bias, name=name) @tf_export("nn.crelu") @@ -1580,7 +1576,7 @@ def relu6(features, name=None): """ with ops.name_scope(name, "Relu6", [features]) as name: features = ops.convert_to_tensor(features, name="features") - return gen_nn_ops._relu6(features, name=name) + return gen_nn_ops.relu6(features, name=name) @tf_export("nn.leaky_relu") @@ -1616,7 +1612,7 @@ def _flatten_outer_dims(logits): output = array_ops.reshape(logits, array_ops.concat([[-1], last_dim_size], 0)) # Set output shape if known. - if context.in_graph_mode(): + if not context.executing_eagerly(): shape = logits.get_shape() if shape is not None and shape.dims is not None: shape = shape.as_list() @@ -1645,7 +1641,7 @@ def _softmax(logits, compute_op, dim=-1, name=None): Args: logits: A non-empty `Tensor`. Must be one of the following types: `half`, `float32`, `float64`. - compute_op: Either gen_nn_ops._softmax or gen_nn_ops._log_softmax + compute_op: Either gen_nn_ops.softmax or gen_nn_ops.log_softmax dim: The dimension softmax would be performed on. The default is -1 which indicates the last dimension. name: A name for the operation (optional). @@ -1739,7 +1735,7 @@ def softmax(logits, axis=None, name=None, dim=None): axis = deprecation.deprecated_argument_lookup("axis", axis, "dim", dim) if axis is None: axis = -1 - return _softmax(logits, gen_nn_ops._softmax, axis, name) + return _softmax(logits, gen_nn_ops.softmax, axis, name) @tf_export("nn.log_softmax") @@ -1769,7 +1765,7 @@ def log_softmax(logits, axis=None, name=None, dim=None): axis = deprecation.deprecated_argument_lookup("axis", axis, "dim", dim) if axis is None: axis = -1 - return _softmax(logits, gen_nn_ops._log_softmax, axis, name) + return _softmax(logits, gen_nn_ops.log_softmax, axis, name) def _ensure_xent_args(name, sentinel, labels, logits): @@ -1871,7 +1867,7 @@ def softmax_cross_entropy_with_logits_v2( # Do the actual op computation. # The second output tensor contains the gradients. We use it in # _CrossEntropyGrad() in nn_grad but not here. - cost, unused_backprop = gen_nn_ops._softmax_cross_entropy_with_logits( + cost, unused_backprop = gen_nn_ops.softmax_cross_entropy_with_logits( precise_logits, labels, name=name) # The output cost shape should be the input minus dim. @@ -1881,7 +1877,8 @@ def softmax_cross_entropy_with_logits_v2( # Make shape inference work since reshape and transpose may erase its static # shape. - if context.in_graph_mode() and shape is not None and shape.dims is not None: + if not context.executing_eagerly( + ) and shape is not None and shape.dims is not None: shape = shape.as_list() del shape[dim] cost.set_shape(shape) @@ -2027,6 +2024,9 @@ def sparse_softmax_cross_entropy_with_logits( # Store label shape for result later. labels_static_shape = labels.get_shape() labels_shape = array_ops.shape(labels) + static_shapes_fully_defined = ( + labels_static_shape.is_fully_defined() and + logits.get_shape()[:-1].is_fully_defined()) if logits.get_shape().ndims is not None and logits.get_shape().ndims == 0: raise ValueError( "Logits cannot be scalars - received shape %s." % logits.get_shape()) @@ -2036,29 +2036,44 @@ def sparse_softmax_cross_entropy_with_logits( raise ValueError("Rank mismatch: Rank of labels (received %s) should " "equal rank of logits minus 1 (received %s)." % (labels_static_shape.ndims, logits.get_shape().ndims)) + if (static_shapes_fully_defined and + labels_static_shape != logits.get_shape()[:-1]): + raise ValueError("Shape mismatch: The shape of labels (received %s) " + "should equal the shape of logits except for the last " + "dimension (received %s)." % (labels_static_shape, + logits.get_shape())) # Check if no reshapes are required. if logits.get_shape().ndims == 2: - cost, _ = gen_nn_ops._sparse_softmax_cross_entropy_with_logits( + cost, _ = gen_nn_ops.sparse_softmax_cross_entropy_with_logits( precise_logits, labels, name=name) if logits.dtype == dtypes.float16: return math_ops.cast(cost, dtypes.float16) else: return cost - # Reshape logits to 2 dim, labels to 1 dim. - num_classes = array_ops.shape(logits)[array_ops.rank(logits) - 1] - precise_logits = array_ops.reshape(precise_logits, [-1, num_classes]) - labels = array_ops.reshape(labels, [-1]) - # The second output tensor contains the gradients. We use it in - # _CrossEntropyGrad() in nn_grad but not here. - cost, _ = gen_nn_ops._sparse_softmax_cross_entropy_with_logits( - precise_logits, labels, name=name) - cost = array_ops.reshape(cost, labels_shape) - cost.set_shape(labels_static_shape) - if logits.dtype == dtypes.float16: - return math_ops.cast(cost, dtypes.float16) - else: - return cost + # Perform a check of the dynamic shapes if the static shapes are not fully + # defined. + shape_checks = [] + if not static_shapes_fully_defined: + shape_checks.append( + check_ops.assert_equal( + array_ops.shape(labels), + array_ops.shape(logits)[:-1])) + with ops.control_dependencies(shape_checks): + # Reshape logits to 2 dim, labels to 1 dim. + num_classes = array_ops.shape(logits)[array_ops.rank(logits) - 1] + precise_logits = array_ops.reshape(precise_logits, [-1, num_classes]) + labels = array_ops.reshape(labels, [-1]) + # The second output tensor contains the gradients. We use it in + # _CrossEntropyGrad() in nn_grad but not here. + cost, _ = gen_nn_ops.sparse_softmax_cross_entropy_with_logits( + precise_logits, labels, name=name) + cost = array_ops.reshape(cost, labels_shape) + cost.set_shape(labels_static_shape) + if logits.dtype == dtypes.float16: + return math_ops.cast(cost, dtypes.float16) + else: + return cost @tf_export("nn.avg_pool") @@ -2086,7 +2101,7 @@ def avg_pool(value, ksize, strides, padding, data_format="NHWC", name=None): """ with ops.name_scope(name, "AvgPool", [value]) as name: value = ops.convert_to_tensor(value, name="input") - return gen_nn_ops._avg_pool( + return gen_nn_ops.avg_pool( value, ksize=ksize, strides=strides, @@ -2116,12 +2131,13 @@ def max_pool(value, ksize, strides, padding, data_format="NHWC", name=None): """ with ops.name_scope(name, "MaxPool", [value]) as name: value = ops.convert_to_tensor(value, name="input") - return gen_nn_ops._max_pool(value, - ksize=ksize, - strides=strides, - padding=padding, - data_format=data_format, - name=name) + return gen_nn_ops.max_pool( + value, + ksize=ksize, + strides=strides, + padding=padding, + data_format=data_format, + name=name) @ops.RegisterStatistics("Conv2D", "flops") @@ -2215,6 +2231,31 @@ def xw_plus_b_v1(x, weights, biases, name=None): # pylint: disable=invalid-name return bias_add_v1(mm, biases, name=name) +def _get_noise_shape(x, noise_shape): + # If noise_shape is none return immediately. + if noise_shape is None: + return array_ops.shape(x) + + try: + # Best effort to figure out the intended shape. + # If not possible, let the op to handle it. + # In eager mode exception will show up. + noise_shape_ = tensor_shape.as_shape(noise_shape) + except (TypeError, ValueError): + return noise_shape + + if x.shape.dims is not None and len(x.shape.dims) == len(noise_shape_.dims): + new_dims = [] + for i, dim in enumerate(x.shape.dims): + if noise_shape_.dims[i].value is None and dim.value is not None: + new_dims.append(dim.value) + else: + new_dims.append(noise_shape_.dims[i].value) + return tensor_shape.TensorShape(new_dims) + + return noise_shape + + @tf_export("nn.dropout") def dropout(x, keep_prob, noise_shape=None, seed=None, name=None): # pylint: disable=invalid-name """Computes dropout. @@ -2265,7 +2306,8 @@ def dropout(x, keep_prob, noise_shape=None, seed=None, name=None): # pylint: di if tensor_util.constant_value(keep_prob) == 1: return x - noise_shape = noise_shape if noise_shape is not None else array_ops.shape(x) + noise_shape = _get_noise_shape(x, noise_shape) + # uniform [keep_prob, 1.0 + keep_prob) random_tensor = keep_prob random_tensor += random_ops.random_uniform( @@ -2273,13 +2315,13 @@ def dropout(x, keep_prob, noise_shape=None, seed=None, name=None): # pylint: di # 0. if [keep_prob, 1.0) and 1. if [1.0, 1.0 + keep_prob) binary_tensor = math_ops.floor(random_tensor) ret = math_ops.div(x, keep_prob) * binary_tensor - if context.in_graph_mode(): + if not context.executing_eagerly(): ret.set_shape(x.get_shape()) return ret @tf_export("nn.top_k") -def top_k(input, k=1, sorted=True, name=None): +def top_k(input, k=1, sorted=True, name=None): # pylint: disable=redefined-builtin """Finds values and indices of the `k` largest entries for the last dimension. If the input is a vector (rank=1), finds the `k` largest entries in the vector @@ -2305,10 +2347,10 @@ def top_k(input, k=1, sorted=True, name=None): values: The `k` largest elements along each last dimensional slice. indices: The indices of `values` within the last dimension of `input`. """ - return gen_nn_ops._top_kv2(input, k=k, sorted=sorted, name=name) + return gen_nn_ops.top_kv2(input, k=k, sorted=sorted, name=name) -def nth_element(input, n, reverse=False, name=None): +def nth_element(input, n, reverse=False, name=None): # pylint: disable=redefined-builtin r"""Finds values of the `n`-th order statistic for the last dmension. If the input is a vector (rank-1), finds the entries which is the nth-smallest @@ -2380,7 +2422,7 @@ def conv1d(value, Args: value: A 3D `Tensor`. Must be of type `float16` or `float32`. - filters: A 3D `Tensor`. Must have the same type as `input`. + filters: A 3D `Tensor`. Must have the same type as `value`. stride: An `integer`. The number of entries by which the filter is moved right at each step. padding: 'SAME' or 'VALID' @@ -2505,7 +2547,7 @@ def conv1d_transpose( spatial_start_dim = 2 strides = [1, 1, 1, stride] value = array_ops.expand_dims(value, spatial_start_dim) - filter = array_ops.expand_dims(filter, 0) + filter = array_ops.expand_dims(filter, 0) # pylint: disable=redefined-builtin result = gen_nn_ops.conv2d_backprop_input( input_sizes=output_shape_, @@ -2624,4 +2666,4 @@ def in_top_k(predictions, targets, k, name=None): A `Tensor` of type `bool`. Computed Precision at `k` as a `bool Tensor`. """ with ops.name_scope(name, "in_top_k"): - return gen_nn_ops._in_top_kv2(predictions, targets, k, name=name) + return gen_nn_ops.in_top_kv2(predictions, targets, k, name=name) diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py index 5a45bdc1e5e1d38a34176ed9443fcd1713f38e1e..af9dae2aa64f0994f403ac81dcba800699d3c960 100644 --- a/tensorflow/python/ops/nn_test.py +++ b/tensorflow/python/ops/nn_test.py @@ -383,6 +383,31 @@ class DropoutTest(test_lib.TestCase): x, keep_prob, noise_shape=array_ops.placeholder(dtypes.int32)) self.assertEqual(x.get_shape(), dropout_x.get_shape()) + def testPartialShapedDropout(self): + x_dim = 40 * 30 + y_dim = 3 + num_iter = 10 + for keep_prob in [0.1, 0.5, 0.8]: + with self.test_session(): + t = constant_op.constant( + 1.0, shape=[x_dim, y_dim], dtype=dtypes.float32) + # Set noise_shape=[None, 1] which means [x_dim, 1]. + dropout = nn_ops.dropout(t, keep_prob, noise_shape=[None, 1]) + self.assertEqual([x_dim, y_dim], dropout.get_shape()) + final_count = 0 + for _ in xrange(0, num_iter): + value = dropout.eval() + final_count += np.count_nonzero(value) + # Verifies that there are only two values: 0 and 1/keep_prob. + sorted_value = np.unique(np.sort(value)) + self.assertEqual(0, sorted_value[0]) + self.assertAllClose(1 / keep_prob, sorted_value[1]) + # Check that we are in the 15% error range + expected_count = x_dim * y_dim * keep_prob * num_iter + rel_error = math.fabs(final_count - expected_count) / expected_count + print(rel_error) + self.assertTrue(rel_error < 0.15) + def testInvalidKeepProb(self): x_dim = 40 y_dim = 30 @@ -1024,6 +1049,22 @@ class DataFormatVectorPermuteTest(test_lib.TestCase): y_val = sess.run(y) self.assertAllEqual(y_val, [7, 9, 3, 4]) + def testNHWCToHWNC(self): + x_val = [7, 4, 9, 3] + x = constant_op.constant(x_val) + y = nn_ops.data_format_vec_permute(x, src_format="NHWC", dst_format="HWNC") + with self.test_session(use_gpu=test_lib.is_gpu_available()) as sess: + y_val = sess.run(y) + self.assertAllEqual(y_val, [4, 9, 7, 3]) + + def testHWNCToNHWC(self): + x_val = [7, 4, 9, 3] + x = constant_op.constant(x_val) + y = nn_ops.data_format_vec_permute(x, src_format="HWNC", dst_format="NHWC") + with self.test_session(use_gpu=test_lib.is_gpu_available()) as sess: + y_val = sess.run(y) + self.assertAllEqual(y_val, [9, 7, 4, 3]) + def testNHWCToNCHW2D(self): x_val = [[7, 4], [9, 3], [4, 5], [5, 1]] x = constant_op.constant(x_val) @@ -1032,6 +1073,22 @@ class DataFormatVectorPermuteTest(test_lib.TestCase): y_val = sess.run(y) self.assertAllEqual(y_val, [[7, 4], [5, 1], [9, 3], [4, 5]]) + def testNHWCToHWNC2D(self): + x_val = [[7, 4], [9, 3], [4, 5], [5, 1]] + x = constant_op.constant(x_val) + y = nn_ops.data_format_vec_permute(x, src_format="NHWC", dst_format="HWNC") + with self.test_session(use_gpu=test_lib.is_gpu_available()) as sess: + y_val = sess.run(y) + self.assertAllEqual(y_val, [[9, 3], [4, 5], [7, 4], [5, 1]]) + + def testHWNCToNHWC2D(self): + x_val = [[7, 4], [9, 3], [4, 5], [5, 1]] + x = constant_op.constant(x_val) + y = nn_ops.data_format_vec_permute(x, src_format="HWNC", dst_format="NHWC") + with self.test_session(use_gpu=test_lib.is_gpu_available()) as sess: + y_val = sess.run(y) + self.assertAllEqual(y_val, [[4, 5], [7, 4], [9, 3], [5, 1]]) + def testNCHWToNHWC2D(self): x_val = [[7, 4], [9, 3], [4, 5], [5, 1]] x = constant_op.constant(x_val) diff --git a/tensorflow/python/ops/numerics.py b/tensorflow/python/ops/numerics.py index b4ce1cbf25346412e2781a520b7e2cdcf720bcd5..d348e47f57b703138aabfc3463e750b795113335 100644 --- a/tensorflow/python/ops/numerics.py +++ b/tensorflow/python/ops/numerics.py @@ -74,7 +74,7 @@ def add_check_numerics_ops(): the checked operations. @enc_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "add_check_numerics_ops() is not compatible with eager execution. " "To check for Inf's and NaN's under eager execution, call " diff --git a/tensorflow/python/ops/parsing_ops.py b/tensorflow/python/ops/parsing_ops.py index b0315ceee268be8ac1813dae5a262a7d9496e154..075b38d743d13329e646c0b268e938b5c5704e47 100644 --- a/tensorflow/python/ops/parsing_ops.py +++ b/tensorflow/python/ops/parsing_ops.py @@ -700,8 +700,7 @@ def _parse_example_raw(serialized, # Finally, convert dense_shapes to TensorShapeProto dense_shapes = [shape.as_proto() for shape in dense_shapes] - # pylint: disable=protected-access - outputs = gen_parsing_ops._parse_example( + outputs = gen_parsing_ops.parse_example( serialized=serialized, names=names, dense_defaults=dense_defaults_vec, @@ -710,7 +709,6 @@ def _parse_example_raw(serialized, dense_keys=dense_keys, dense_shapes=dense_shapes, name=name) - # pylint: enable=protected-access (sparse_indices, sparse_values, sparse_shapes, dense_values) = outputs @@ -1132,8 +1130,7 @@ def _parse_single_sequence_example_raw(serialized, feature_list_dense_shapes = [tensor_shape.as_shape(shape).as_proto() for shape in feature_list_dense_shapes] - # pylint: disable=protected-access - outputs = gen_parsing_ops._parse_single_sequence_example( + outputs = gen_parsing_ops.parse_single_sequence_example( serialized=serialized, debug_name=debug_name, context_dense_defaults=context_dense_defaults_vec, @@ -1149,7 +1146,6 @@ def _parse_single_sequence_example_raw(serialized, feature_list_dense_missing_assumed_empty=( feature_list_dense_missing_assumed_empty), name=name) - # pylint: enable=protected-access (context_sparse_indices, context_sparse_values, context_sparse_shapes, context_dense_values, @@ -1182,7 +1178,6 @@ def _parse_single_sequence_example_raw(serialized, @tf_export("decode_csv") def decode_csv(records, record_defaults, field_delim=",", use_quote_delim=True, name=None, na_value=""): - # pylint: disable=protected-access """Convert CSV records to tensors. Each column maps to one tensor. RFC 4180 format is expected for the CSV records. @@ -1211,11 +1206,13 @@ def decode_csv(records, record_defaults, field_delim=",", Each tensor will have the same shape as records. """ # TODO(martinwicke), remove the wrapper when new Python API generator is done. - return gen_parsing_ops._decode_csv( - records=records, record_defaults=record_defaults, - field_delim=field_delim, use_quote_delim=use_quote_delim, - na_value=na_value, name=name) - # pylint: enable=protected-access + return gen_parsing_ops.decode_csv( + records=records, + record_defaults=record_defaults, + field_delim=field_delim, + use_quote_delim=use_quote_delim, + na_value=na_value, + name=name) # TODO(b/70890287): Combine the implementation of this op and @@ -1391,7 +1388,6 @@ def _parse_single_example_v2_raw(serialized, sparse_keys, sparse_types, # Finally, convert dense_shapes to TensorShapeProto dense_shapes = [shape.as_proto() for shape in dense_shapes] - # pylint: disable=protected-access outputs = gen_parsing_ops.parse_single_example( serialized=serialized, dense_defaults=dense_defaults_vec, @@ -1401,7 +1397,6 @@ def _parse_single_example_v2_raw(serialized, sparse_keys, sparse_types, dense_keys=dense_keys, dense_shapes=dense_shapes, name=name) - # pylint: enable=protected-access (sparse_indices, sparse_values, sparse_shapes, dense_values) = outputs diff --git a/tensorflow/python/ops/random_ops.py b/tensorflow/python/ops/random_ops.py index 2c86358d21b1c280b8d7ade625fd4b7a44c5de26..6a2dd3f1cd55eea1d3b652a31cd2784c411c2ce0 100644 --- a/tensorflow/python/ops/random_ops.py +++ b/tensorflow/python/ops/random_ops.py @@ -43,7 +43,6 @@ def _ShapeTensor(shape): return ops.convert_to_tensor(shape, dtype=dtype, name="shape") -# pylint: disable=protected-access @tf_export("random_normal") def random_normal(shape, mean=0.0, @@ -74,7 +73,7 @@ def random_normal(shape, mean_tensor = ops.convert_to_tensor(mean, dtype=dtype, name="mean") stddev_tensor = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev") seed1, seed2 = random_seed.get_seed(seed) - rnd = gen_random_ops._random_standard_normal( + rnd = gen_random_ops.random_standard_normal( shape_tensor, dtype, seed=seed1, seed2=seed2) mul = rnd * stddev_tensor value = math_ops.add(mul, mean_tensor, name=name) @@ -126,7 +125,7 @@ def parameterized_truncated_normal(shape, minvals_tensor = ops.convert_to_tensor(minvals, dtype=dtype, name="minvals") maxvals_tensor = ops.convert_to_tensor(maxvals, dtype=dtype, name="maxvals") seed1, seed2 = random_seed.get_seed(seed) - rnd = gen_random_ops._parameterized_truncated_normal( + rnd = gen_random_ops.parameterized_truncated_normal( shape_tensor, means_tensor, stddevs_tensor, @@ -171,7 +170,7 @@ def truncated_normal(shape, mean_tensor = ops.convert_to_tensor(mean, dtype=dtype, name="mean") stddev_tensor = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev") seed1, seed2 = random_seed.get_seed(seed) - rnd = gen_random_ops._truncated_normal( + rnd = gen_random_ops.truncated_normal( shape_tensor, dtype, seed=seed1, seed2=seed2) mul = rnd * stddev_tensor value = math_ops.add(mul, mean_tensor, name=name) @@ -210,7 +209,7 @@ def random_uniform(shape, maxval: A 0-D Tensor or Python value of type `dtype`. The upper bound on the range of random values to generate. Defaults to 1 if `dtype` is floating point. - dtype: The type of the output: 'float16`, `float32`, `float64`, `int32`, + dtype: The type of the output: `float16`, `float32`, `float64`, `int32`, or `int64`. seed: A Python integer. Used to create a random seed for the distribution. See @{tf.set_random_seed} @@ -237,11 +236,10 @@ def random_uniform(shape, maxval = ops.convert_to_tensor(maxval, dtype=dtype, name="max") seed1, seed2 = random_seed.get_seed(seed) if dtype.is_integer: - return gen_random_ops._random_uniform_int( + return gen_random_ops.random_uniform_int( shape, minval, maxval, seed=seed1, seed2=seed2, name=name) else: - rnd = gen_random_ops._random_uniform( - shape, dtype, seed=seed1, seed2=seed2) + rnd = gen_random_ops.random_uniform(shape, dtype, seed=seed1, seed2=seed2) return math_ops.add(rnd * (maxval - minval), minval, name=name) @@ -275,7 +273,7 @@ def random_shuffle(value, seed=None, name=None): dimension. """ seed1, seed2 = random_seed.get_seed(seed) - return gen_random_ops._random_shuffle( + return gen_random_ops.random_shuffle( value, seed=seed1, seed2=seed2, name=name) @@ -420,7 +418,7 @@ def random_gamma(shape, seed1, seed2 = random_seed.get_seed(seed) return math_ops.maximum( np.finfo(dtype.as_numpy_dtype).tiny, - gen_random_ops._random_gamma( + gen_random_ops.random_gamma( shape, alpha_broadcast, seed=seed1, seed2=seed2) / beta) ops.NotDifferentiable("RandomGamma") diff --git a/tensorflow/python/ops/resource_variable_ops.py b/tensorflow/python/ops/resource_variable_ops.py index bdf41cd75d6432750b7b23391c28892e2d6b9ffc..affa7ae62995ba944ab9447d6d1df50e54935612 100644 --- a/tensorflow/python/ops/resource_variable_ops.py +++ b/tensorflow/python/ops/resource_variable_ops.py @@ -21,6 +21,7 @@ from __future__ import print_function from tensorflow.core.framework import attr_value_pb2 from tensorflow.core.framework import variable_pb2 +from tensorflow.python import pywrap_tensorflow from tensorflow.python.eager import context from tensorflow.python.eager import tape from tensorflow.python.framework import dtypes @@ -30,11 +31,13 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_array_ops from tensorflow.python.ops import gen_resource_variable_ops from tensorflow.python.ops import gen_state_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import variables # go/tf-wildcard-import # pylint: disable=wildcard-import from tensorflow.python.ops.gen_resource_variable_ops import * # pylint: enable=wildcard-import +from tensorflow.python.training import checkpointable from tensorflow.python.util import compat @@ -43,10 +46,6 @@ def _eager_safe_variable_handle(shape, dtype, shared_name, name, graph_mode): container = ops.get_default_graph()._container # pylint: disable=protected-access if container is None: container = "" - if not graph_mode: - # When in eager mode use a uid for the shared_name, to prevent accidental - # sharing. - shared_name = str(ops.uid()) handle = gen_resource_variable_ops.var_handle_op(shape=shape, dtype=dtype, shared_name=shared_name, name=name, @@ -107,13 +106,16 @@ class EagerResourceDeleter(object): """ def __init__(self, handle, handle_device): + if not isinstance(handle, ops.Tensor): + raise ValueError( + ("Passed handle=%s to EagerResourceDeleter. Was expecting a handle " + "Tensor." % (handle,))) self._handle = handle self._handle_device = handle_device def __del__(self): # Resources follow object-identity when executing eagerly, so it is safe to - # delete the resource we have a handle to. Each Graph has a unique container - # name, which prevents resource sharing. + # delete the resource we have a handle to. try: # This resource was created in eager mode. However, this destructor may be # running in graph mode (especially during unit tests). To clean up @@ -129,10 +131,10 @@ class EagerResourceDeleter(object): # valid, and so on. Printing warnings in these cases is silly # (exceptions raised from __del__ are printed as warnings to stderr). pass # 'NoneType' object is not callable when the handle has been - # partially unloaded. + # partially unloaded. except AttributeError: pass # 'NoneType' object has no attribute 'eager_mode' when context has - # been unloaded. Will catch other module unloads as well. + # been unloaded. Will catch other module unloads as well. def shape_safe_assign_variable_handle(handle, shape, value, name=None): @@ -147,7 +149,7 @@ def shape_safe_assign_variable_handle(handle, shape, value, name=None): class ResourceVariable(variables.Variable): """Variable based on resource handles. - See the ${variables} documentation for more details. + See the @{$python/state_ops$`Variables`} documentation for more details. A `ResourceVariable` allows you to maintain state across subsequent calls to session.run. @@ -261,9 +263,9 @@ class ResourceVariable(variables.Variable): if initial_value is not None: raise ValueError("variable_def and initial_value are mutually " "exclusive.") - if not context.in_graph_mode(): - raise ValueError("Creating ResourceVariable from variable_def" - " only supported in GRAPH mode.") + if context.executing_eagerly(): + raise ValueError("Creating ResourceVariable from variable_def is " + "not supported when eager execution is enabled.") self._init_from_proto(variable_def, import_scope=import_scope) else: self._init_from_args( @@ -344,19 +346,30 @@ class ResourceVariable(variables.Variable): if constraint is not None and not callable(constraint): raise ValueError("The `constraint` argument must be a callable.") + if isinstance(initial_value, checkpointable.CheckpointInitialValue): + self._maybe_initialize_checkpointable() + self._update_uid = initial_value.checkpoint_position.restore_uid + initial_value = initial_value.wrapped_value + self._trainable = trainable if trainable and ops.GraphKeys.TRAINABLE_VARIABLES not in collections: collections = list(collections) + [ops.GraphKeys.TRAINABLE_VARIABLES] self._save_slice_info = None - # Save the graph's container prefix for error checking. Reading the value of - # the ResourceVariable from another Graph in Eager mode is an error. - self._container_prefix = ops.get_default_graph()._container_prefix # pylint: disable=protected-access + # Store the graph key so optimizers know how to only retrieve variables from + # this graph. + self._graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access with ops.init_scope(): - self._in_graph_mode = context.in_graph_mode() + self._in_graph_mode = not context.executing_eagerly() with ops.name_scope(name, "Variable", [] if init_from_fn else [initial_value]) as name: # pylint: disable=protected-access handle_name = ops._name_from_scope_name(name) + if self._in_graph_mode: + shared_name = handle_name + else: + # When in eager mode use a uid for the shared_name, to prevent + # accidental sharing. + shared_name = "%s_%d" % (handle_name, ops.uid()) if init_from_fn: # Use attr_scope and device(None) to simulate the behavior of # colocate_with when the variable we want to colocate with doesn't @@ -372,12 +385,9 @@ class ResourceVariable(variables.Variable): self._handle = _eager_safe_variable_handle( shape=initial_value.get_shape(), dtype=initial_value.dtype.base_dtype, - shared_name=handle_name, + shared_name=shared_name, name=name, graph_mode=self._in_graph_mode) - self._handle_device = ( - self._handle.device if self._in_graph_mode else - context.get_default_context().device_name) self._shape = initial_value.get_shape() else: initial_value = initial_value() @@ -387,12 +397,9 @@ class ResourceVariable(variables.Variable): self._handle = _eager_safe_variable_handle( shape=initial_value.get_shape(), dtype=initial_value.dtype.base_dtype, - shared_name=handle_name, + shared_name=shared_name, name=name, graph_mode=False) - self._handle_device = ( - self._handle.device if self._in_graph_mode else - context.get_default_context().device_name) self._shape = initial_value.get_shape() # pylint: enable=protected-access @@ -413,13 +420,12 @@ class ResourceVariable(variables.Variable): self._handle = _eager_safe_variable_handle( shape=initial_value.get_shape(), dtype=initial_value.dtype.base_dtype, - shared_name=handle_name, + shared_name=shared_name, name=name, graph_mode=self._in_graph_mode) - self._handle_device = (self._handle.device if self._in_graph_mode else - context.get_default_context().device_name) self._shape = initial_value.get_shape() + self._unique_id = shared_name self._initial_value = initial_value if self._in_graph_mode else None self._handle_name = handle_name + ":0" self._dtype = initial_value.dtype.base_dtype @@ -440,7 +446,7 @@ class ResourceVariable(variables.Variable): with ops.name_scope("Read"), ops.colocate_with(self._handle): # Manually assign reads to the handle's device to avoid log # messages. - with ops.device(self._handle_device): + with ops.device(self._handle.device): value = self._read_variable_op() self._graph_element = value if caching_device is not None: @@ -467,7 +473,7 @@ class ResourceVariable(variables.Variable): self._cached_value = self._read_variable_op() else: self._cached_value = None - if context.in_graph_mode(): + if not context.executing_eagerly(): ops.add_to_collections(collections, self) elif ops.GraphKeys.GLOBAL_STEP in collections: ops.add_to_collections(ops.GraphKeys.GLOBAL_STEP, self) @@ -480,12 +486,13 @@ class ResourceVariable(variables.Variable): # cycles being uncollectable, and means that no __del__ will be defined at # all in graph mode. self._handle_deleter = EagerResourceDeleter( - handle=self._handle, handle_device=self._handle_device) + handle=self._handle, handle_device=self._handle.device) + self._cached_shape_as_list = None def _init_from_proto(self, variable_def, import_scope=None): """Initializes from `VariableDef` proto.""" # Note that init_from_proto is currently not supported in Eager mode. - assert context.in_graph_mode() + assert not context.executing_eagerly() self._in_graph_mode = True assert isinstance(variable_def, variable_pb2.VariableDef) if not variable_def.is_resource: @@ -498,11 +505,19 @@ class ResourceVariable(variables.Variable): variable_def.variable_name, import_scope=import_scope)) self._shape = tensor_shape.TensorShape( self._handle.op.get_attr("shape")) - self._handle_device = self._handle.device self._handle_name = self._handle.name + self._unique_id = self._handle_name self._initializer_op = g.as_graph_element( ops.prepend_name_scope( variable_def.initializer_name, import_scope=import_scope)) + # Check whether initial_value_name exists for backwards compatibility. + if (hasattr(variable_def, "initial_value_name") and + variable_def.initial_value_name): + self._initial_value = g.as_graph_element( + ops.prepend_name_scope(variable_def.initial_value_name, + import_scope=import_scope)) + else: + self._initial_value = None if variable_def.snapshot_name: self._cached_value = g.as_graph_element( ops.prepend_name_scope( @@ -517,8 +532,10 @@ class ResourceVariable(variables.Variable): self._save_slice_info = None self._caching_device = None self._dtype = dtypes.as_dtype(self._handle.op.get_attr("dtype")) - self._graph_element = self.value() + self._graph_element = g.get_tensor_by_name( + self._handle.op.name + "/Read/ReadVariableOp:0") self._constraint = None + self._cached_shape_as_list = None def __nonzero__(self): return self.__bool__() @@ -534,7 +551,7 @@ class ResourceVariable(variables.Variable): @property def device(self): """The device this variable is on.""" - return self._handle_device + return self._handle.device @property def graph(self): @@ -551,11 +568,26 @@ class ResourceVariable(variables.Variable): """The shape of this variable.""" return self._shape + def _shape_as_list(self): + if self._cached_shape_as_list: + return self._cached_shape_as_list + if self.shape.ndims is None: + return None + self._cached_shape_as_list = [dim.value for dim in self.shape.dims] + return self._cached_shape_as_list + + def _shape_tuple(self): + shape = self._shape_as_list() + if shape is None: + return None + return tuple(shape) + @property def create(self): """The op responsible for initializing this variable.""" if not self._in_graph_mode: - raise RuntimeError("Calling create in EAGER mode not supported.") + raise RuntimeError("Calling create is not supported when eager execution" + " is enabled.") return self._initializer_op @property @@ -568,7 +600,7 @@ class ResourceVariable(variables.Variable): if self._cached_value is not None: return self._cached_value with ops.colocate_with(None, ignore_existing=True): - with ops.device(self._handle_device): + with ops.device(self._handle.device): return self._read_variable_op() def _as_graph_element(self): @@ -583,7 +615,7 @@ class ResourceVariable(variables.Variable): @property def initial_value(self): """Returns the Tensor used as the initial value for the variable.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("initial_value not supported in EAGER mode.") return self._initial_value @@ -604,15 +636,15 @@ class ResourceVariable(variables.Variable): def eval(self, session=None): """Evaluates and returns the value of this variable.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Trying to eval in EAGER mode") return self._graph_element.eval(session=session) def numpy(self): - if context.in_graph_mode(): - raise NotImplementedError( - "numpy() is only available when eager execution is enabled.") - return self.read_value().numpy() + if context.executing_eagerly(): + return self.read_value().numpy() + raise NotImplementedError( + "numpy() is only available when eager execution is enabled.") def count_up_to(self, limit): """Increments this variable until it reaches `limit`. @@ -662,18 +694,10 @@ class ResourceVariable(variables.Variable): Returns: the read operation. - Raises: - ValueError: if the ResourceVariable was created in another isolation - environment or graph. """ - if (not self._in_graph_mode and - self._container_prefix != ops.get_default_graph()._container_prefix): # pylint: disable=protected-access - raise ValueError( - "Attempted to read a variable from another isolation environment" - " or Graph") with ops.name_scope("Read"): # Ensure we read the variable in the same device as the handle. - with ops.device(self._handle_device): + with ops.device(self._handle.device): value = self._read_variable_op() # Return an identity so it can get placed on whatever device the context # specifies instead of the device where the variable is. @@ -701,12 +725,18 @@ class ResourceVariable(variables.Variable): A `VariableDef` protocol buffer, or `None` if the `Variable` is not in the specified name scope. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("to_proto not supported in EAGER mode.") if export_scope is None or self.handle.name.startswith(export_scope): var_def = variable_pb2.VariableDef() var_def.variable_name = ops.strip_name_scope(self.handle.name, export_scope) + if self._initial_value is not None: + # This is inside an if-statement for backwards compatibility, since + # self._initial_value might be None for variables constructed from old + # protos. + var_def.initial_value_name = ops.strip_name_scope( + self._initial_value.name, export_scope) var_def.initializer_name = ops.strip_name_scope(self.initializer.name, export_scope) if self._cached_value is not None: @@ -722,7 +752,7 @@ class ResourceVariable(variables.Variable): @staticmethod def from_proto(variable_def, import_scope=None): - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("from_proto not supported in EAGER mode.") return ResourceVariable( variable_def=variable_def, import_scope=import_scope) @@ -773,42 +803,89 @@ class ResourceVariable(variables.Variable): __array_priority__ = 100 - def assign_sub(self, delta, use_locking=None, name=None): + def assign_sub(self, delta, use_locking=None, name=None, read_value=True): + """Subtracts a value from this variable. + + Args: + delta: A `Tensor`. The value to subtract from this variable. + use_locking: If `True`, use locking during the operation. + name: The name to use for the operation. + read_value: A `bool`. Whether to read and return the new value of the + variable or not. + + Returns: + If `read_value` is `True`, this method will return the new value of the + variable after the assignment has completed. Otherwise, when in graph mode + it will return the `Operation` that does the assignment, and when in eager + mode it will return `None`. + """ # TODO(apassos): this here and below is not atomic. Consider making it # atomic if there's a way to do so without a performance cost for those who # don't need it. - with ops.control_dependencies([ - gen_resource_variable_ops.assign_sub_variable_op( - self.handle, - ops.convert_to_tensor(delta, dtype=self.dtype), - name=name) - ]): - return self.read_value() - - def assign_add(self, delta, use_locking=None, name=None): - with ops.control_dependencies([ - gen_resource_variable_ops.assign_add_variable_op( - self.handle, - ops.convert_to_tensor(delta, dtype=self.dtype), - name=name) - ]): - return self.read_value() - - def assign(self, value, use_locking=None, name=None): + assign_sub_op = gen_resource_variable_ops.assign_sub_variable_op( + self.handle, ops.convert_to_tensor(delta, dtype=self.dtype), name=name) + if read_value: + return self._lazy_read(assign_sub_op) + return assign_sub_op + + def assign_add(self, delta, use_locking=None, name=None, read_value=True): + """Adds a value to this variable. + + Args: + delta: A `Tensor`. The value to add to this variable. + use_locking: If `True`, use locking during the operation. + name: The name to use for the operation. + read_value: A `bool`. Whether to read and return the new value of the + variable or not. + + Returns: + If `read_value` is `True`, this method will return the new value of the + variable after the assignment has completed. Otherwise, when in graph mode + it will return the `Operation` that does the assignment, and when in eager + mode it will return `None`. + """ + assign_add_op = gen_resource_variable_ops.assign_add_variable_op( + self.handle, ops.convert_to_tensor(delta, dtype=self.dtype), name=name) + if read_value: + return self._lazy_read(assign_add_op) + return assign_add_op + + def _lazy_read(self, op): + if hasattr(self, "_trainable") and self._trainable: + tape.watch_variable(self) + return _UnreadVariable( + self._handle, self.dtype, self._shape, self._in_graph_mode, + self._handle_deleter if not self._in_graph_mode else None, op, + self._unique_id) + + def assign(self, value, use_locking=None, name=None, read_value=True): + """Assigns a new value to this variable. + + Args: + value: A `Tensor`. The new value for this variable. + use_locking: If `True`, use locking during the assignment. + name: The name to use for the assignment. + read_value: A `bool`. Whether to read and return the new value of the + variable or not. + + Returns: + If `read_value` is `True`, this method will return the new value of the + variable after the assignment has completed. Otherwise, when in graph mode + it will return the `Operation` that does the assignment, and when in eager + mode it will return `None`. + """ value_tensor = ops.convert_to_tensor(value, dtype=self.dtype) self._shape.assert_is_compatible_with(value_tensor.shape) - with ops.control_dependencies([ - gen_resource_variable_ops.assign_variable_op( - self.handle, - value_tensor, - name=name) - ]): - return self.read_value() + assign_op = gen_resource_variable_ops.assign_variable_op( + self.handle, value_tensor, name=name) + if read_value: + return self._lazy_read(assign_op) + return assign_op def _strided_slice_assign(self, begin, end, strides, value, name, begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask): - with ops.control_dependencies([ + return self._lazy_read( gen_array_ops.resource_strided_slice_assign( ref=self.handle, begin=begin, @@ -820,9 +897,12 @@ class ResourceVariable(variables.Variable): end_mask=end_mask, ellipsis_mask=ellipsis_mask, new_axis_mask=new_axis_mask, - shrink_axis_mask=shrink_axis_mask) - ]): - return self.value() + shrink_axis_mask=shrink_axis_mask)) + + def __int__(self): + if self.dtype != dtypes.int32 and self.dtype != dtypes.int64: + raise TypeError("Non-integer variable can't be converted to integer.") + return int(self.value().numpy()) def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): del name @@ -876,10 +956,70 @@ class ResourceVariable(variables.Variable): "Tensor object.") +pywrap_tensorflow.TFE_Py_RegisterResourceVariableType(ResourceVariable) +math_ops._resource_variable_type = ResourceVariable # pylint: disable=protected-access + + def _dense_var_to_tensor(var, dtype=None, name=None, as_ref=False): return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access +class _UnreadVariable(ResourceVariable): + """Represents a future for a read of a variable. + + Pretends to be the tensor if anyone looks. + """ + + def __init__(self, handle, dtype, # pylint: disable=super-init-not-called + shape, in_graph_mode, deleter, parent_op, unique_id): + # We do not call super init on purpose. + self._trainable = False + self._save_slice_info = None + self._graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access + self._in_graph_mode = in_graph_mode + self._handle = handle + self._shape = shape + self._initial_value = None + if isinstance(self._handle, ops.EagerTensor): + self._handle_name = "" + else: + self._handle_name = self._handle.name + self._unique_id = unique_id + self._dtype = dtype + self._constraint = None + self._cached_value = None + self._is_initialized_op = None + self._initializer_op = None + self._parent_op = parent_op + if context.executing_eagerly(): + self._graph_element = None + else: + self._graph_element = self.read_value() + self._handle_deleter = deleter + + def value(self): + return self._read_variable_op() + + def read_value(self): + return self._read_variable_op() + + def _read_variable_op(self): + with ops.control_dependencies([self._parent_op]): + return gen_resource_variable_ops.read_variable_op(self._handle, + self._dtype) + + def set_shape(self, shape): + self._shape = shape + self._cached_shape_as_list = None + + @property + def op(self): + """The op for this variable.""" + return self._parent_op + +ops.register_tensor_conversion_function(_UnreadVariable, _dense_var_to_tensor) +ops.register_dense_tensor_like_type(_UnreadVariable) + # Register a conversion function which reads the value of the variable, # allowing instances of the class to be used as tensors. @@ -951,3 +1091,9 @@ ops.register_proto_function( proto_type=variable_pb2.VariableDef, to_proto=_to_proto_fn, from_proto=_from_proto_fn) + + +def is_resource_variable(var): + """"Returns True if `var` is to be considered a ResourceVariable.""" + return isinstance(var, ResourceVariable) or hasattr( + var, "_should_act_as_resource_variable") diff --git a/tensorflow/python/ops/rnn.py b/tensorflow/python/ops/rnn.py index e0052b8869dd2cf331c14e2355d4b40dd217c561..42af7f8b274c555e6375ab8e937a8cc06ffbaa8e 100644 --- a/tensorflow/python/ops/rnn.py +++ b/tensorflow/python/ops/rnn.py @@ -45,7 +45,6 @@ from tensorflow.python.util.tf_export import tf_export # pylint: disable=protected-access _concat = rnn_cell_impl._concat -_like_rnncell = rnn_cell_impl._like_rnncell # pylint: enable=protected-access @@ -83,8 +82,9 @@ def _best_effort_input_batch_size(flat_input): """Get static input batch size if available, with fallback to the dynamic one. Args: - flat_input: An iterable of time major input Tensors of shape [max_time, - batch_size, ...]. All inputs should have compatible batch sizes. + flat_input: An iterable of time major input Tensors of shape + `[max_time, batch_size, ...]`. + All inputs should have compatible batch sizes. Returns: The batch size in Python integer if available, or a scalar Tensor otherwise. @@ -171,11 +171,11 @@ def _rnn_step( return (final_output, final_state) Args: - time: Python int, the current time step - sequence_length: int32 `Tensor` vector of size [batch_size] - min_sequence_length: int32 `Tensor` scalar, min of sequence_length - max_sequence_length: int32 `Tensor` scalar, max of sequence_length - zero_output: `Tensor` vector of shape [output_size] + time: int32 `Tensor` scalar. + sequence_length: int32 `Tensor` vector of size [batch_size]. + min_sequence_length: int32 `Tensor` scalar, min of sequence_length. + max_sequence_length: int32 `Tensor` scalar, max of sequence_length. + zero_output: `Tensor` vector of shape [output_size]. state: Either a single `Tensor` matrix of shape `[batch_size, state_size]`, or a list/tuple of such tensors. call_cell: lambda returning tuple of (new_output, new_state) where @@ -202,6 +202,9 @@ def _rnn_step( flat_state = nest.flatten(state) flat_zero_output = nest.flatten(zero_output) + # Vector describing which batch entries are finished. + copy_cond = time >= sequence_length + def _copy_one_through(output, new_output): # TensorArray and scalar get passed through. if isinstance(output, tensor_array_ops.TensorArray): @@ -209,7 +212,6 @@ def _rnn_step( if output.shape.ndims == 0: return new_output # Otherwise propagate the old or the new value. - copy_cond = (time >= sequence_length) with ops.colocate_with(new_output): return array_ops.where(copy_cond, output, new_output) @@ -400,11 +402,8 @@ def bidirectional_dynamic_rnn(cell_fw, cell_bw, inputs, sequence_length=None, Raises: TypeError: If `cell_fw` or `cell_bw` is not an instance of `RNNCell`. """ - - if not _like_rnncell(cell_fw): - raise TypeError("cell_fw must be an instance of RNNCell") - if not _like_rnncell(cell_bw): - raise TypeError("cell_bw must be an instance of RNNCell") + rnn_cell_impl.assert_like_rnncell("cell_fw", cell_fw) + rnn_cell_impl.assert_like_rnncell("cell_bw", cell_bw) with vs.variable_scope(scope or "bidirectional_rnn"): # Forward direction @@ -565,14 +564,13 @@ def dynamic_rnn(cell, inputs, sequence_length=None, initial_state=None, TypeError: If `cell` is not an instance of RNNCell. ValueError: If inputs is None or an empty list. """ - if not _like_rnncell(cell): - raise TypeError("cell must be an instance of RNNCell") + rnn_cell_impl.assert_like_rnncell("cell", cell) with vs.variable_scope(scope or "rnn") as varscope: # Create a new scope in which the caching device is either # determined by the parent scope, or is set to place the cached # Variable using the same placement as for the rest of the RNN. - if context.in_graph_mode(): + if not context.executing_eagerly(): if varscope.caching_device is None: varscope.set_caching_device(lambda op: op.device) @@ -613,7 +611,7 @@ def dynamic_rnn(cell, inputs, sequence_length=None, initial_state=None, ["Expected shape for Tensor %s is " % x.name, packed_shape, " but saw shape: ", x_shape]) - if context.in_graph_mode() and sequence_length is not None: + if not context.executing_eagerly() and sequence_length is not None: # Perform some shape validation with ops.control_dependencies( [_assert_has_shape(sequence_length, [batch_size])]): @@ -739,7 +737,7 @@ def _dynamic_rnn_loop(cell, element_shape=element_shape, tensor_array_name=base_name + name) - in_graph_mode = context.in_graph_mode() + in_graph_mode = not context.executing_eagerly() if in_graph_mode: output_ta = tuple( _create_ta( @@ -869,7 +867,7 @@ def raw_rnn(cell, loop_fn, ```python time = tf.constant(0, dtype=tf.int32) - (finished, next_input, initial_state, _, loop_state) = loop_fn( + (finished, next_input, initial_state, emit_structure, loop_state) = loop_fn( time=time, cell_output=None, cell_state=None, loop_state=None) emit_ta = TensorArray(dynamic_size=True, dtype=initial_state.dtype) state = initial_state @@ -880,7 +878,7 @@ def raw_rnn(cell, loop_fn, loop_state=loop_state) # Emit zeros and copy forward state for minibatch entries that are finished. state = tf.where(finished, state, next_state) - emit = tf.where(finished, tf.zeros_like(emit), emit) + emit = tf.where(finished, tf.zeros_like(emit_structure), emit) emit_ta = emit_ta.write(time, emit) # If any new minibatch entries are marked as finished, mark these. finished = tf.logical_or(finished, next_finished) @@ -940,10 +938,15 @@ def raw_rnn(cell, loop_fn, and `emit_output`: the output to store for this iteration. Note that `emit_output` should be a `Tensor` or (possibly nested) - tuple of tensors with shapes and structure matching `cell.output_size` - and `cell_output` above. The parameter `cell_state` and output - `next_cell_state` may be either a single or (possibly nested) tuple - of tensors. The parameter `loop_state` and + tuple of tensors which is aggregated in the `emit_ta` inside the + `while_loop`. For the first call to `loop_fn`, the `emit_output` + corresponds to the `emit_structure` which is then used to determine the + size of the `zero_tensor` for the `emit_ta` (defaults to + `cell.output_size`). For the subsequent calls to the `loop_fn`, the + `emit_output` corresponds to the actual output tensor + that is to be aggregated in the `emit_ta`. The parameter `cell_state` + and output `next_cell_state` may be either a single or (possibly nested) + tuple of tensors. The parameter `loop_state` and output `next_loop_state` may be either a single or (possibly nested) tuple of `Tensor` and `TensorArray` objects. This last parameter may be ignored by `loop_fn` and the return value may be `None`. If it @@ -1012,9 +1015,8 @@ def raw_rnn(cell, loop_fn, TypeError: If `cell` is not an instance of RNNCell, or `loop_fn` is not a `callable`. """ + rnn_cell_impl.assert_like_rnncell("cell", cell) - if not _like_rnncell(cell): - raise TypeError("cell must be an instance of RNNCell") if not callable(loop_fn): raise TypeError("loop_fn must be a callable") @@ -1024,7 +1026,7 @@ def raw_rnn(cell, loop_fn, # determined by the parent scope, or is set to place the cached # Variable using the same placement as for the rest of the RNN. with vs.variable_scope(scope or "rnn") as varscope: - if context.in_graph_mode(): + if not context.executing_eagerly(): if varscope.caching_device is None: varscope.set_caching_device(lambda op: op.device) @@ -1226,9 +1228,7 @@ def static_rnn(cell, ValueError: If `inputs` is `None` or an empty list, or if the input depth (column size) cannot be inferred from inputs via shape inference. """ - - if not _like_rnncell(cell): - raise TypeError("cell must be an instance of RNNCell") + rnn_cell_impl.assert_like_rnncell("cell", cell) if not nest.is_sequence(inputs): raise TypeError("inputs must be a sequence") if not inputs: @@ -1239,7 +1239,7 @@ def static_rnn(cell, # determined by the parent scope, or is set to place the cached # Variable using the same placement as for the rest of the RNN. with vs.variable_scope(scope or "rnn") as varscope: - if context.in_graph_mode(): + if not context.executing_eagerly(): if varscope.caching_device is None: varscope.set_caching_device(lambda op: op.device) @@ -1466,11 +1466,8 @@ def static_bidirectional_rnn(cell_fw, TypeError: If `cell_fw` or `cell_bw` is not an instance of `RNNCell`. ValueError: If inputs is None or an empty list. """ - - if not _like_rnncell(cell_fw): - raise TypeError("cell_fw must be an instance of RNNCell") - if not _like_rnncell(cell_bw): - raise TypeError("cell_bw must be an instance of RNNCell") + rnn_cell_impl.assert_like_rnncell("cell_fw", cell_fw) + rnn_cell_impl.assert_like_rnncell("cell_bw", cell_bw) if not nest.is_sequence(inputs): raise TypeError("inputs must be a sequence") if not inputs: diff --git a/tensorflow/python/ops/rnn_cell_impl.py b/tensorflow/python/ops/rnn_cell_impl.py index f1ac3e9bafa09e4647b4a4263e74fad29b643fd5..fe380c44dafdad6dc25d50102bacba610132674d 100644 --- a/tensorflow/python/ops/rnn_cell_impl.py +++ b/tensorflow/python/ops/rnn_cell_impl.py @@ -46,6 +46,7 @@ from tensorflow.python.ops import tensor_array_ops from tensorflow.python.ops import variable_scope as vs from tensorflow.python.ops import variables as tf_variables from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpointable from tensorflow.python.util import nest from tensorflow.python.util.tf_export import tf_export @@ -54,6 +55,8 @@ _BIAS_VARIABLE_NAME = "bias" _WEIGHTS_VARIABLE_NAME = "kernel" +# TODO(jblespiau): Remove this function when we are sure there are no longer +# any usage (even if protected, it is being used). Prefer assert_like_rnncell. def _like_rnncell(cell): """Checks that a given object is an RNNCell by using duck typing.""" conditions = [hasattr(cell, "output_size"), hasattr(cell, "state_size"), @@ -61,6 +64,45 @@ def _like_rnncell(cell): return all(conditions) +# This can be used with self.assertRaisesRegexp for assert_like_rnncell. +ASSERT_LIKE_RNNCELL_ERROR_REGEXP = "is not an RNNCell" + + +def assert_like_rnncell(cell_name, cell): + """Raises a TypeError if cell is not like an RNNCell. + + NOTE: Do not rely on the error message (in particular in tests) which can be + subject to change to increase readability. Use + ASSERT_LIKE_RNNCELL_ERROR_REGEXP. + + Args: + cell_name: A string to give a meaningful error referencing to the name + of the functionargument. + cell: The object which should behave like an RNNCell. + + Raises: + TypeError: A human-friendly exception. + """ + conditions = [ + hasattr(cell, "output_size"), + hasattr(cell, "state_size"), + hasattr(cell, "zero_state"), + callable(cell), + ] + errors = [ + "'output_size' property is missing", + "'state_size' property is missing", + "'zero_state' method is missing", + "is not callable" + ] + + if not all(conditions): + + errors = [error for error, cond in zip(errors, conditions) if not cond] + raise TypeError("The argument {!r} ({}) is not an RNNCell: {}.".format( + cell_name, cell, ", ".join(errors))) + + def _concat(prefix, suffix, static=False): """Concat that enables int, Tensor, or TensorShape values. @@ -127,7 +169,7 @@ def _zero_state_tensors(state_size, batch_size, dtype): """Combine s with batch_size to get a proper tensor shape.""" c = _concat(batch_size, s) size = array_ops.zeros(c, dtype=dtype) - if context.in_graph_mode(): + if not context.executing_eagerly(): c_static = _concat(batch_size, s, static=True) size.set_shape(c_static) return size @@ -191,12 +233,13 @@ class RNNCell(base_layer.Layer): def _rnn_get_variable(self, getter, *args, **kwargs): variable = getter(*args, **kwargs) - if context.in_graph_mode(): - trainable = (variable in tf_variables.trainable_variables() or - (isinstance(variable, tf_variables.PartitionedVariable) and - list(variable)[0] in tf_variables.trainable_variables())) - else: + if context.executing_eagerly(): trainable = variable._trainable # pylint: disable=protected-access + else: + trainable = ( + variable in tf_variables.trainable_variables() or + (isinstance(variable, tf_variables.PartitionedVariable) and + list(variable)[0] in tf_variables.trainable_variables())) if trainable and variable not in self._trainable_weights: self._trainable_weights.append(variable) elif not trainable and variable not in self._non_trainable_weights: @@ -240,7 +283,7 @@ class RNNCell(base_layer.Layer): # Try to use the last cached zero_state. This is done to avoid recreating # zeros, especially when eager execution is enabled. state_size = self.state_size - is_eager = context.in_eager_mode() + is_eager = context.executing_eagerly() if is_eager and hasattr(self, "_last_zero_state"): (last_state_size, last_batch_size, last_dtype, last_output) = getattr(self, "_last_zero_state") @@ -255,7 +298,7 @@ class RNNCell(base_layer.Layer): return output -class _LayerRNNCell(RNNCell): +class LayerRNNCell(RNNCell): """Subclass of RNNCells that act like proper `tf.Layer` objects. For backwards compatibility purposes, most `RNNCell` instances allow their @@ -297,7 +340,7 @@ class _LayerRNNCell(RNNCell): @tf_export("nn.rnn_cell.BasicRNNCell") -class BasicRNNCell(_LayerRNNCell): +class BasicRNNCell(LayerRNNCell): """The most basic RNN cell. Args: @@ -355,7 +398,7 @@ class BasicRNNCell(_LayerRNNCell): @tf_export("nn.rnn_cell.GRUCell") -class GRUCell(_LayerRNNCell): +class GRUCell(LayerRNNCell): """Gated Recurrent Unit cell (cf. http://arxiv.org/abs/1406.1078). Args: @@ -473,7 +516,7 @@ class LSTMStateTuple(_LSTMStateTuple): @tf_export("nn.rnn_cell.BasicLSTMCell") -class BasicLSTMCell(_LayerRNNCell): +class BasicLSTMCell(LayerRNNCell): """Basic LSTM recurrent network cell. The implementation is based on: http://arxiv.org/abs/1409.2329. @@ -598,7 +641,7 @@ class BasicLSTMCell(_LayerRNNCell): @tf_export("nn.rnn_cell.LSTMCell") -class LSTMCell(_LayerRNNCell): +class LSTMCell(LayerRNNCell): """Long short-term memory unit (LSTM) recurrent network cell. The default non-peephole implementation is based on: @@ -912,8 +955,8 @@ class DropoutWrapper(RNNCell): but not `callable`. ValueError: if any of the keep_probs are not between 0 and 1. """ - if not _like_rnncell(cell): - raise TypeError("The parameter cell is not a RNNCell.") + assert_like_rnncell("cell", cell) + if (dropout_state_filter_visitor is not None and not callable(dropout_state_filter_visitor)): raise TypeError("dropout_state_filter_visitor must be callable") @@ -1187,6 +1230,12 @@ class MultiRNNCell(RNNCell): "cells must be a list or tuple, but saw: %s." % cells) self._cells = cells + for cell_number, cell in enumerate(self._cells): + # Add Checkpointable dependencies on these cells so their variables get + # saved with this object when using object-based saving. + if isinstance(cell, checkpointable.CheckpointableBase): + # TODO(allenl): Track down non-Checkpointable callers. + self._track_checkpointable(cell, name="cell-%d" % (cell_number,)) self._state_is_tuple = state_is_tuple if not state_is_tuple: if any(nest.is_sequence(c.state_size) for c in self._cells): diff --git a/tensorflow/python/ops/script_ops.py b/tensorflow/python/ops/script_ops.py index 1b9071ee93c21f8d6bdc9ace11dbf57f3eb3e218..fb59bbba5eb7a668b6b959c8519f2b52f0ce26fb 100644 --- a/tensorflow/python/ops/script_ops.py +++ b/tensorflow/python/ops/script_ops.py @@ -25,6 +25,9 @@ from __future__ import print_function import threading +# Used by py_util.cc to get tracebacks. +import traceback # pylint: disable=unused-import + import numpy as np import six @@ -33,6 +36,7 @@ from tensorflow.python.eager import context from tensorflow.python.framework import function from tensorflow.python.framework import ops from tensorflow.python.ops import gen_script_ops +from tensorflow.python.util import nest from tensorflow.python.util.tf_export import tf_export @@ -97,7 +101,7 @@ class FuncRegistry(object): components of a tensor have different lengths. This is bad: ignoring the padding is wrong for text data, and removing the padding is wrong for binary data. To avoid this bug, we redo the conversion using an object dtype. - Additionally, we convert unicode strings to (byte-)strings for Python3 + Additionally, we convert unicode strings to (byte-)strings for compatibility. Args: @@ -111,7 +115,7 @@ class FuncRegistry(object): if result.dtype.char == "S" and result is not value: return np.asarray(value, order="C", dtype=object) elif result.dtype.char == "U" and result is not value: - value = np.vectorize(lambda x: x.encode())(value) + value = np.vectorize(lambda x: x.encode("utf8"))(value) return np.asarray(value, order="C", dtype=object) elif result.dtype.char == "U": return result.astype(np.bytes_) @@ -176,7 +180,10 @@ class CleanupFunc(object): self._token = token def __del__(self): - _py_funcs.remove(self._token) + if _py_funcs is not None: + # If _py_funcs is None, the program is most likely in shutdown, and the + # _py_funcs object has been destroyed already. + _py_funcs.remove(self._token) def _internal_py_func(func, inp, Tout, stateful=None, eager=False, name=None): @@ -215,18 +222,16 @@ def _internal_py_func(func, inp, Tout, stateful=None, eager=False, name=None): graph._cleanup_py_funcs_used_in_graph.append(cleanup) # pylint: enable=protected-access - # pylint: disable=protected-access if eager: - result = gen_script_ops._eager_py_func( + result = gen_script_ops.eager_py_func( input=inp, token=token, Tout=Tout, name=name) else: if stateful: - result = gen_script_ops._py_func( + result = gen_script_ops.py_func( input=inp, token=token, Tout=Tout, name=name) else: - result = gen_script_ops._py_func_stateless( + result = gen_script_ops.py_func_stateless( input=inp, token=token, Tout=Tout, name=name) - # pylint: enable=protected-access return result if is_list_or_tuple else result[0] @@ -264,7 +269,7 @@ def py_func(func, inp, Tout, stateful=True, name=None): """Wraps a python function and uses it as a TensorFlow op. Given a python function `func`, which takes numpy arrays as its - inputs and returns numpy arrays as its outputs, wrap this function as an + arguments 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: @@ -273,8 +278,8 @@ def py_func(func, inp, Tout, stateful=True, name=None): 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) + input = tf.placeholder(tf.float32) + y = tf.py_func(my_func, [input], tf.float32) ``` **N.B.** The `tf.py_func()` operation has the following known limitations: @@ -290,10 +295,12 @@ def py_func(func, inp, Tout, stateful=True, name=None): 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`. + func: A Python function, which accepts `ndarray` objects as arguments and + returns a list of `ndarray` objects (or a single `ndarray`). This function + must accept as many arguments as there are tensors in `inp`, and these + argument types will match the corresponding `tf.Tensor` objects + in `inp`. The returns `ndarray`s must match the number and types defined + `Tout`. Important Note: Input and output numpy `ndarray`s of `func` are not guaranteed to be copies. In some cases their underlying memory will be shared with the corresponding TensorFlow tensors. @@ -313,6 +320,12 @@ def py_func(func, inp, Tout, stateful=True, name=None): Returns: A list of `Tensor` or a single `Tensor` which `func` computes. """ + if context.executing_eagerly(): + result = func(*[x.numpy() for x in inp]) + result = nest.flatten(result) + + return [x if x is None else ops.convert_to_tensor(x) for x in result] + return _internal_py_func( func=func, inp=inp, Tout=Tout, stateful=stateful, eager=False, name=name) diff --git a/tensorflow/python/ops/session_ops.py b/tensorflow/python/ops/session_ops.py index cedd36c1deed541adcf601ff9447345e2279e8f9..ad38845153c94e9bb31e6e3ee05ebed0a4313efc 100644 --- a/tensorflow/python/ops/session_ops.py +++ b/tensorflow/python/ops/session_ops.py @@ -16,7 +16,6 @@ """Tensor Handle Operations. See the @{$python/session_ops} guide. @@get_session_handle -@@get_session_handle_v2 @@get_session_tensor @@delete_session_tensor """ @@ -182,7 +181,7 @@ def get_session_handle(data, name=None): # Colocate this operation with data. with ops.colocate_with(data): - return gen_data_flow_ops._get_session_handle(data, name=name) # pylint: disable=protected-access + return gen_data_flow_ops.get_session_handle(data, name=name) @tf_export("get_session_tensor") @@ -222,7 +221,7 @@ def get_session_tensor(handle, dtype, name=None): with ops.device(handle_device): holder = array_ops.placeholder(dtypes.string) _register_handle_feeder(holder.graph, holder, dtype) - tensor = gen_data_flow_ops._get_session_tensor(holder, dtype, name=name) + tensor = gen_data_flow_ops.get_session_tensor(holder, dtype, name=name) return (holder, tensor) @@ -246,7 +245,7 @@ def delete_session_tensor(handle, name=None): handle_device = TensorHandle._get_device_name(handle) with ops.device(handle_device): holder = array_ops.placeholder(dtypes.string) - deleter = gen_data_flow_ops._delete_session_tensor(holder, name=name) + deleter = gen_data_flow_ops.delete_session_tensor(holder, name=name) return (holder, deleter) @@ -268,7 +267,7 @@ def _get_handle_reader(graph, handle, dtype): with graph.as_default(), graph.device(handle_device): holder = array_ops.placeholder(dtypes.string) _register_handle_feeder(holder.graph, holder, dtype) - reader = gen_data_flow_ops._get_session_tensor(holder, dtype) + reader = gen_data_flow_ops.get_session_tensor(holder, dtype) result = (holder, reader) graph._handle_readers[graph_key] = result return result @@ -289,7 +288,7 @@ def _get_handle_mover(graph, feeder, handle): # Create mover if we haven't done it. holder, reader = _get_handle_reader(graph, handle, dtype) with graph.as_default(), graph.device(feeder.op.device): - mover = gen_data_flow_ops._get_session_handle(reader) # pylint: disable=protected-access + mover = gen_data_flow_ops.get_session_handle(reader) result = (holder, mover) graph._handle_movers[graph_key] = result return result @@ -303,7 +302,7 @@ def _get_handle_deleter(graph, deleter_key, handle): handle_device = TensorHandle._get_device_name(handle) with graph.as_default(), graph.device(handle_device): holder = array_ops.placeholder(dtypes.string) - deleter = gen_data_flow_ops._delete_session_tensor(holder) + deleter = gen_data_flow_ops.delete_session_tensor(holder) result = (holder, deleter) graph._handle_deleters[deleter_key] = result return result diff --git a/tensorflow/python/ops/sparse_grad.py b/tensorflow/python/ops/sparse_grad.py index 5295e7d21c2b5810422ec36f5aced63c9039feca..97353d6c747cb7e4d3c1fa92ad61af24fb17de91 100644 --- a/tensorflow/python/ops/sparse_grad.py +++ b/tensorflow/python/ops/sparse_grad.py @@ -88,10 +88,8 @@ def _SparseAddGrad(op, *grads): # the non-zero elements of the sum, and we will peek into `sum_indices` in the # gradient op. - # pylint: disable=protected-access - a_val_grad, b_val_grad = gen_sparse_ops._sparse_add_grad(val_grad, a_indices, - b_indices, - sum_indices) + a_val_grad, b_val_grad = gen_sparse_ops.sparse_add_grad( + val_grad, a_indices, b_indices, sum_indices) a_val_grad.set_shape(op.inputs[1].get_shape()) b_val_grad.set_shape(op.inputs[4].get_shape()) # (a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh) @@ -151,7 +149,7 @@ def _SparseTensorDenseMatMulGrad(op, grad): "complex gradients.") # gradient w.r.t. dense - b_grad = gen_sparse_ops._sparse_tensor_dense_mat_mul( # pylint: disable=protected-access + b_grad = gen_sparse_ops.sparse_tensor_dense_mat_mul( a_indices, a_values, a_shape, grad, adjoint_a=not adj_a) if adj_b: b_grad = array_ops.transpose(b_grad) @@ -278,8 +276,7 @@ def _SparseFillEmptyRowsGrad(op, unused_grad_output_indices, output_grad_values, """Gradients for SparseFillEmptyRows.""" reverse_index_map = op.outputs[3] - # pylint: disable=protected-access - d_values, d_default_value = gen_sparse_ops._sparse_fill_empty_rows_grad( + d_values, d_default_value = gen_sparse_ops.sparse_fill_empty_rows_grad( reverse_index_map=reverse_index_map, grad_values=output_grad_values) # d_indices, d_values, d_dense_shape, d_default_value. diff --git a/tensorflow/python/ops/sparse_ops.py b/tensorflow/python/ops/sparse_ops.py index 0fbbf5a805f1439d85ad53f02bdb665c04248606..c580052c32c8b61467b857af3d237be41718c1a1 100644 --- a/tensorflow/python/ops/sparse_ops.py +++ b/tensorflow/python/ops/sparse_ops.py @@ -234,7 +234,7 @@ def sparse_concat(axis, ] output_ind, output_val, output_shape = ( - gen_sparse_ops._sparse_concat(inds, vals, shapes, axis, name=name)) + gen_sparse_ops.sparse_concat(inds, vals, shapes, axis, name=name)) return sparse_tensor.SparseTensor(output_ind, output_val, output_shape) @@ -302,8 +302,8 @@ def sparse_add(a, b, thresh=0): thresh = ops.convert_to_tensor( thresh, dtype=a.values.dtype.real_dtype.base_dtype, name="thresh") output_ind, output_val, output_shape = ( - gen_sparse_ops._sparse_add(a.indices, a.values, a.dense_shape, - b.indices, b.values, b.dense_shape, thresh)) + gen_sparse_ops.sparse_add(a.indices, a.values, a.dense_shape, + b.indices, b.values, b.dense_shape, thresh)) # Attempt to get output_shape statically. a.get_shape().assert_is_compatible_with(b.get_shape()) @@ -317,8 +317,8 @@ def sparse_add(a, b, thresh=0): # swap to make `a` the SparseTensor. if isinstance(b, sparse_classes): a, b = b, a - return gen_sparse_ops._sparse_tensor_dense_add(a.indices, a.values, - a.dense_shape, b) + return gen_sparse_ops.sparse_tensor_dense_add(a.indices, a.values, + a.dense_shape, b) def _sparse_cross(inputs, name=None): @@ -402,7 +402,7 @@ def _sparse_cross_internal(inputs, num_buckets=0, hash_key=None, name=None): - """See gen_sparse_ops._sparse_cross.""" + """See gen_sparse_ops.sparse_cross.""" if not isinstance(inputs, list): raise TypeError("Inputs must be a list") if not all( @@ -432,7 +432,7 @@ def _sparse_cross_internal(inputs, dense_inputs[i] = math_ops.to_int64(dense_inputs[i]) internal_type = dtypes.int64 - indices_out, values_out, shape_out = gen_sparse_ops._sparse_cross( + indices_out, values_out, shape_out = gen_sparse_ops.sparse_cross( indices=indices, values=values, shapes=shapes, @@ -511,7 +511,7 @@ def sparse_reorder(sp_input, name=None): sp_input = _convert_to_sparse_tensor(sp_input) reordered_ind, reordered_val = ( - gen_sparse_ops._sparse_reorder( + gen_sparse_ops.sparse_reorder( sp_input.indices, sp_input.values, sp_input.dense_shape, name=name)) if sp_input.get_shape().is_fully_defined(): @@ -575,7 +575,7 @@ def sparse_reshape(sp_input, shape, name=None): shape = math_ops.cast(shape, dtype=dtypes.int64) with ops.name_scope(name, "SparseReshape", [sp_input]) as name: - reshaped_ind, reshaped_shape = gen_sparse_ops._sparse_reshape( + reshaped_ind, reshaped_shape = gen_sparse_ops.sparse_reshape( sp_input.indices, sp_input.dense_shape, shape, name=name) reshaped_shape_const = tensor_util.constant_value(shape) @@ -671,7 +671,7 @@ def sparse_split(keyword_required=KeywordRequired(), sp_input = _convert_to_sparse_tensor(sp_input) output_inds, output_vals, output_shapes = ( - gen_sparse_ops._sparse_split( + gen_sparse_ops.sparse_split( axis, sp_input.indices, sp_input.values, @@ -782,7 +782,7 @@ def sparse_to_dense(sparse_indices, Dense `Tensor` of shape `output_shape`. Has the same type as `sparse_values`. """ - return gen_sparse_ops._sparse_to_dense( + return gen_sparse_ops.sparse_to_dense( sparse_indices, output_shape, sparse_values, @@ -1412,7 +1412,7 @@ def sparse_fill_empty_rows(sp_input, default_value, name=None): default_value = ops.convert_to_tensor( default_value, dtype=sp_input.values.dtype) (output_indices, output_values, empty_row_indicator, - unused_reverse_index_map) = gen_sparse_ops._sparse_fill_empty_rows( + unused_reverse_index_map) = gen_sparse_ops.sparse_fill_empty_rows( indices=sp_input.indices, values=sp_input.values, dense_shape=sp_input.dense_shape, @@ -1441,7 +1441,7 @@ def serialize_sparse(sp_input, name=None, out_type=dtypes.string): """ sp_input = _convert_to_sparse_tensor(sp_input) - return gen_sparse_ops._serialize_sparse( + return gen_sparse_ops.serialize_sparse( sp_input.indices, sp_input.values, sp_input.dense_shape, @@ -1476,7 +1476,7 @@ def serialize_many_sparse(sp_input, name=None, out_type=dtypes.string): """ sp_input = _convert_to_sparse_tensor(sp_input) - return gen_sparse_ops._serialize_many_sparse( + return gen_sparse_ops.serialize_many_sparse( sp_input.indices, sp_input.values, sp_input.dense_shape, @@ -1541,7 +1541,7 @@ def deserialize_sparse(serialized_sparse, dtype, rank=None, name=None): """ output_indices, output_values, output_shape = ( - gen_sparse_ops._deserialize_sparse(serialized_sparse, dtype, name=name)) + gen_sparse_ops.deserialize_sparse(serialized_sparse, dtype, name=name)) # Feed rank data back in, if available output_indices.set_shape([None, rank]) @@ -1610,7 +1610,7 @@ def deserialize_many_sparse(serialized_sparse, dtype, rank=None, name=None): All of the serialized `SparseTensor`s must have had the same rank and type. """ output_indices, output_values, output_shape = ( - gen_sparse_ops._deserialize_many_sparse( + gen_sparse_ops.deserialize_many_sparse( serialized_sparse, dtype, name=name)) # Feed rank data back in, if available @@ -1828,7 +1828,7 @@ def sparse_tensor_dense_matmul(sp_a, with ops.name_scope(name, "SparseTensorDenseMatMul", [sp_a.indices, sp_a.values, b]) as name: b = ops.convert_to_tensor(b, name="b") - return gen_sparse_ops._sparse_tensor_dense_mat_mul( + return gen_sparse_ops.sparse_tensor_dense_mat_mul( a_indices=sp_a.indices, a_values=sp_a.values, a_shape=sp_a.dense_shape, @@ -2046,7 +2046,7 @@ def _add_sparse_to_tensors_map(sp_input, """ sp_input = _convert_to_sparse_tensor(sp_input) - return gen_sparse_ops._add_sparse_to_tensors_map( + return gen_sparse_ops.add_sparse_to_tensors_map( sp_input.indices, sp_input.values, sp_input.dense_shape, @@ -2086,7 +2086,7 @@ def _add_many_sparse_to_tensors_map(sp_input, """ sp_input = _convert_to_sparse_tensor(sp_input) - return gen_sparse_ops._add_many_sparse_to_tensors_map( + return gen_sparse_ops.add_many_sparse_to_tensors_map( sp_input.indices, sp_input.values, sp_input.dense_shape, @@ -2167,7 +2167,7 @@ def _take_many_sparse_from_tensors_map(sparse_map_op, with ops.colocate_with(sparse_map_op): shared_name = sparse_map_op.get_attr("shared_name") or sparse_map_op.name output_indices, output_values, output_shape = ( - gen_sparse_ops._take_many_sparse_from_tensors_map( + gen_sparse_ops.take_many_sparse_from_tensors_map( sparse_handles, dtype=sparse_map_op.get_attr("T"), container=sparse_map_op.get_attr("container"), diff --git a/tensorflow/python/ops/special_math_ops.py b/tensorflow/python/ops/special_math_ops.py index 6d7eaababcd94d687ff20dddc35c68a98320a19b..5e2146b79f08e6671c429f388b05634b1727b4ed 100644 --- a/tensorflow/python/ops/special_math_ops.py +++ b/tensorflow/python/ops/special_math_ops.py @@ -163,7 +163,7 @@ def einsum(equation, *inputs, **kwargs): if '...' in equation: raise ValueError('Subscripts with ellipses are not yet supported.') - match = re.match('([a-z,]+)(->[a-z]*)?', equation) + match = re.match('^([a-zA-Z,]+)(->[a-zA-Z]*)?$', equation) if not match: raise ValueError('Indices have incorrect format: %s' % equation) @@ -402,7 +402,7 @@ def _exponential_space_einsum(equation, *inputs): if '...' in equation: raise ValueError('Subscripts with ellipses are not yet supported.') - match = re.match('([a-z,]+)(->[a-z]*)?', equation) + match = re.match('^([a-zA-Z,]+)(->[a-zA-Z]*)?$', equation) if not match: raise ValueError('Indices have incorrect format: %s' % equation) diff --git a/tensorflow/python/ops/special_math_ops_test.py b/tensorflow/python/ops/special_math_ops_test.py index 2c212f45483eacfd3fd27eecb8d7b2c846b5fe96..d7c3a7e8dc7c2ad611cf47718dddcf38700ce304 100644 --- a/tensorflow/python/ops/special_math_ops_test.py +++ b/tensorflow/python/ops/special_math_ops_test.py @@ -192,6 +192,9 @@ class EinsumTest(test.TestCase): 'abc,cba', 'dba,ead,cad->bce', 'aef,fbc,dca->bde', + 'iJ,Jk->ik', + 'iJ,Ki->JK', + 'iJk,Jklm->Jk' ] long_cases = [ @@ -208,6 +211,8 @@ class EinsumTest(test.TestCase): 'ijk ijk', 'ij.jk->ik', 'ij...,jk...->ik...', + 'ij,k ->kji', + 'ij,k-> kji', # axis in output that does not exist 'ij,jk->im', diff --git a/tensorflow/python/ops/standard_ops.py b/tensorflow/python/ops/standard_ops.py index 737b923415f4a36ee7886ec0eb2e6dfa78be1de4..60a98aca7f822802363de0efd0d2974e45dc810e 100644 --- a/tensorflow/python/ops/standard_ops.py +++ b/tensorflow/python/ops/standard_ops.py @@ -25,6 +25,7 @@ import sys as _sys # Imports the following modules so that @RegisterGradient get executed. from tensorflow.python.ops import array_grad from tensorflow.python.ops import data_flow_grad +from tensorflow.python.ops import manip_grad from tensorflow.python.ops import math_grad from tensorflow.python.ops import manip_grad from tensorflow.python.ops import sparse_grad @@ -43,11 +44,12 @@ from tensorflow.python.ops.special_math_ops import * # TODO(vrv): Switch to import * once we're okay with exposing the module. from tensorflow.python.ops.confusion_matrix import confusion_matrix from tensorflow.python.ops.control_flow_ops import Assert +from tensorflow.python.ops.control_flow_ops import case +from tensorflow.python.ops.control_flow_ops import cond from tensorflow.python.ops.control_flow_ops import group from tensorflow.python.ops.control_flow_ops import no_op -from tensorflow.python.ops.control_flow_ops import tuple -from tensorflow.python.ops.control_flow_ops import cond -from tensorflow.python.ops.control_flow_ops import case +from tensorflow.python.ops.control_flow_ops import tuple # pylint: disable=redefined-builtin +# pylint: enable=redefined-builtin from tensorflow.python.ops.control_flow_ops import while_loop from tensorflow.python.ops.data_flow_ops import * from tensorflow.python.ops.functional_ops import * @@ -58,6 +60,7 @@ from tensorflow.python.ops.io_ops import * from tensorflow.python.ops.linalg_ops import * from tensorflow.python.ops.logging_ops import Print from tensorflow.python.ops.logging_ops import get_summary_op +from tensorflow.python.ops.logging_ops import timestamp from tensorflow.python.ops.lookup_ops import initialize_all_tables from tensorflow.python.ops.lookup_ops import tables_initializer from tensorflow.python.ops.manip_ops import * @@ -183,7 +186,6 @@ _allowed_symbols_array_ops = [ "quantize_and_dequantize", # to-doc # TODO(drpng): legacy symbols to be removed. - "list_diff", # Use tf.listdiff instead. "batch_matrix_diag", "batch_matrix_band_part", "batch_matrix_diag_part", @@ -216,6 +218,7 @@ _allowed_symbols_gradients = [ # Documented in training.py: # Not importing training.py to avoid complex graph dependencies. "AggregationMethod", + "custom_gradient", "gradients", # tf.gradients = gradients.gradients "hessians", ] @@ -230,7 +233,7 @@ _allowed_symbols_clip_ops = [ "global_norm", ] -_allowed_symbols_image_ops = [ +_allowed_symbols_logging_ops = [ # Documented in training.py. # We are not importing training.py to avoid complex dependencies. "audio_summary", @@ -260,42 +263,43 @@ _allowed_symbols = (_allowed_symbols_array_ops + _allowed_symbols_clip_ops + _allowed_symbols_control_flow_ops + _allowed_symbols_functional_ops + - _allowed_symbols_image_ops + _allowed_symbols_gradients + + _allowed_symbols_logging_ops + _allowed_symbols_math_ops + _allowed_symbols_variable_scope_ops + _allowed_symbols_misc + _allowed_symbols_partitioned_variables) -remove_undocumented(__name__, _allowed_symbols, - [_sys.modules[__name__], - _array_ops, - _check_ops, - _clip_ops, - _confusion_matrix, - _control_flow_ops, - _constant_op, - _data_flow_ops, - _functional_ops, - _gradients, - _histogram_ops, - _init_ops, - _io_ops, - _linalg_ops, - _logging_ops, - _manip_ops, - _math_ops, - _numerics, - _parsing_ops, - _partitioned_variables, - _random_ops, - _script_ops, - _session_ops, - _sparse_ops, - _special_math_ops, - _state_ops, - _string_ops, - _template, - _tensor_array_ops, - _variable_scope, - _variables,]) +remove_undocumented(__name__, _allowed_symbols, [ + _sys.modules[__name__], + _array_ops, + _check_ops, + _clip_ops, + _confusion_matrix, + _control_flow_ops, + _constant_op, + _data_flow_ops, + _functional_ops, + _gradients, + _histogram_ops, + _init_ops, + _io_ops, + _linalg_ops, + _logging_ops, + _manip_ops, + _math_ops, + _numerics, + _parsing_ops, + _partitioned_variables, + _random_ops, + _script_ops, + _session_ops, + _sparse_ops, + _special_math_ops, + _state_ops, + _string_ops, + _template, + _tensor_array_ops, + _variable_scope, + _variables, +]) diff --git a/tensorflow/python/ops/state_ops.py b/tensorflow/python/ops/state_ops.py index 3cc76fdbf34ff6de47d98400cd826d671c9178eb..c3ad5831b4dfa2160a198429c60c7a6ac00f6357 100644 --- a/tensorflow/python/ops/state_ops.py +++ b/tensorflow/python/ops/state_ops.py @@ -99,8 +99,8 @@ def variable_op(shape, dtype, name="Variable", set_shape=True, container="", """Deprecated. Used variable_op_v2 instead.""" if not set_shape: shape = tensor_shape.unknown_shape() - ret = gen_state_ops._variable(shape=shape, dtype=dtype, name=name, - container=container, shared_name=shared_name) + ret = gen_state_ops.variable(shape=shape, dtype=dtype, name=name, + container=container, shared_name=shared_name) # TODO(mrry): Move this to where it is used, so we can get rid of this op # wrapper? if set_shape: @@ -127,11 +127,12 @@ def variable_op_v2(shape, dtype, name="Variable", container="", shared_name=""): Returns: A variable tensor. """ - return gen_state_ops._variable_v2(shape=shape, - dtype=dtype, - name=name, - container=container, - shared_name=shared_name) + return gen_state_ops.variable_v2( + shape=shape, + dtype=dtype, + name=name, + container=container, + shared_name=shared_name) def init_variable(v, init, name="init"): @@ -185,7 +186,7 @@ def is_variable_initialized(ref, name=None): if ref.dtype._is_ref_dtype: return gen_state_ops.is_variable_initialized(ref=ref, name=name) # Handle resource variables. - if context.in_eager_mode() or ref.op.type == "VarHandleOp": + if context.executing_eagerly() or ref.op.type == "VarHandleOp": return gen_resource_variable_ops.var_is_initialized_op(ref.handle, name=name) @@ -278,7 +279,7 @@ def assign(ref, value, validate_shape=None, use_locking=None, name=None): return gen_state_ops.assign( ref, value, use_locking=use_locking, name=name, validate_shape=validate_shape) - return ref.assign(value) + return ref.assign(value, name=name) @tf_export("count_up_to") @@ -353,11 +354,9 @@ def scatter_update(ref, indices, updates, use_locking=True, name=None): if ref.dtype._is_ref_dtype: return gen_state_ops.scatter_update(ref, indices, updates, use_locking=use_locking, name=name) - with ops.control_dependencies( - [gen_resource_variable_ops.resource_scatter_update( - ref.handle, indices, ops.convert_to_tensor(updates, ref.dtype), - name=name)]): - return ref.read_value() + return ref._lazy_read(gen_resource_variable_ops.resource_scatter_update( # pylint: disable=protected-access + ref.handle, indices, ops.convert_to_tensor(updates, ref.dtype), + name=name)) @tf_export("scatter_nd_update") diff --git a/tensorflow/python/ops/string_ops.py b/tensorflow/python/ops/string_ops.py index b8c39d91b41790c6441594b175e8eaa03620e1ec..5bd75b9215fdbccd5882ea39c2b35ccbbe29d5b0 100644 --- a/tensorflow/python/ops/string_ops.py +++ b/tensorflow/python/ops/string_ops.py @@ -17,6 +17,7 @@ See the @{$python/string_ops} guide. +@@regex_replace @@string_to_hash_bucket_fast @@string_to_hash_bucket_strong @@string_to_hash_bucket @@ -93,10 +94,8 @@ def string_split(source, delimiter=" ", skip_empty=True): # pylint: disable=inv delimiter = ops.convert_to_tensor(delimiter, dtype=dtypes.string) source = ops.convert_to_tensor(source, dtype=dtypes.string) - # pylint: disable=protected-access - indices, values, shape = gen_string_ops._string_split( + indices, values, shape = gen_string_ops.string_split( source, delimiter=delimiter, skip_empty=skip_empty) - # pylint: enable=protected-access indices.set_shape([None, 2]) values.set_shape([None]) shape.set_shape([2]) @@ -141,6 +140,7 @@ def reduce_join(inputs, axis=None, reduce_join.__doc__ = deprecation.rewrite_argument_docstring( gen_string_ops.reduce_join.__doc__, "reduction_indices", "axis") +ops.NotDifferentiable("RegexReplace") ops.NotDifferentiable("StringToHashBucket") ops.NotDifferentiable("StringToHashBucketFast") ops.NotDifferentiable("StringToHashBucketStrong") diff --git a/tensorflow/python/ops/summary_ops.py b/tensorflow/python/ops/summary_ops.py index 7f4f4ce5ab4ee2bd309932cb81f05775996371d6..037bc9845a3f734f65b73b0c4b4ca19fb653731d 100644 --- a/tensorflow/python/ops/summary_ops.py +++ b/tensorflow/python/ops/summary_ops.py @@ -13,7 +13,6 @@ # limitations under the License. # ============================================================================== """Summary Operations.""" -# pylint: disable=protected-access from __future__ import absolute_import from __future__ import division from __future__ import print_function @@ -74,7 +73,7 @@ def tensor_summary(name, with summary_op_util.summary_scope( name, family, values=[tensor]) as (tag, scope): - val = gen_logging_ops._tensor_summary_v2( + val = gen_logging_ops.tensor_summary_v2( tensor=tensor, tag=tag, name=scope, diff --git a/tensorflow/python/ops/template.py b/tensorflow/python/ops/template.py index 806fdd3da7aa6de01b7cd4d9d36dbf43f6139db6..0a391d896a388f7b40b574f6e1367e9d22fb1d0a 100644 --- a/tensorflow/python/ops/template.py +++ b/tensorflow/python/ops/template.py @@ -26,6 +26,7 @@ from tensorflow.python.eager import function from tensorflow.python.framework import ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpointable from tensorflow.python.util import tf_contextlib from tensorflow.python.util import tf_decorator from tensorflow.python.util.deprecation import deprecated @@ -203,7 +204,7 @@ def make_template_internal(name_, if kwargs: func_ = tf_decorator.make_decorator(func_, functools.partial( func_, **kwargs)) - if context.in_eager_mode(): + if context.executing_eagerly(): if unique_name_ is not None: raise ValueError( "unique_name_ cannot be used when eager exeuction is enabled.") @@ -230,7 +231,7 @@ def _skip_common_stack_elements(stacktrace, base_case): return stacktrace[-1:] -class Template(object): +class Template(checkpointable.CheckpointableBase): """Wrap a function to aid in variable sharing. Templates are functions that create variables the first time they are called @@ -294,12 +295,115 @@ class Template(object): # which is not the same as whether the scope has been created. self._variables_created = False + @property + def _checkpoint_dependencies(self): + """Sanity checking for object-based saving. + + Does not override Checkpointable dependency tracking, but checks that + variables accessible through Checkpointable dependencies on other `Template` + objects include all of the variable_scope-filtered `Template.variables`. + + Returns: + A list of checkpointable.CheckpointableReference objects. + Raises: + ValueError: If this object is not compatible with object-based saving. + """ + dependencies = super(Template, self)._checkpoint_dependencies + dependency_variables = [] + for _, dependency in dependencies: + if isinstance(dependency, Template): + dependency_variables.extend(dependency.variables) + else: + dependency_variables.append(dependency) + dependency_variables = set(dependency_variables) + not_included_variables = [] + for expected_variable in sorted(self.variables, key=lambda v: v.name): + if expected_variable not in dependency_variables: + not_included_variables.append(expected_variable) + if not_included_variables: + # Trying to save a Template which improperly tracks its variables. + raise ValueError( + ("The Template '%s' references variables which are not included via " + "object-based dependency tracking. Most likely a custom " + "getter/creator was registered which does not call Template's " + "custom variable creator (which is responsible for tracking " + "dependencies).\n\nExpected these variables to be dependencies: %s") + % (self, not_included_variables)) + return dependencies + + def _checkpointable_custom_creator(self, next_creator, name, initial_value, + checkpointable_parent=None, **kwargs): + """A variable creation hook which adds Checkpointable dependencies. + + Set during the `Template`'s first wrapped function execution. Ensures that + (a) `Template` objects depend on `Template`s created inside them which + create variables, and (b) that any variables not in a more deeply nested + `Template` are added as dependencies directly. + + The `checkpointable_parent` argument is passed between `Template` custom + creators but ignored when the variable object itself is created. This + argument indicates (if not `None`) that a more deeply nested `Template` has + already added the variable as a dependency, and that parent `Template`s + should add a dependency on that `Template` rather than on the variable + directly. + + Args: + next_creator: See `variable_scope.variable_creator_scope`; the next + creator in the chain. + name: The (full, scope-influenced) name of the variable. The scope name + for the Template itself is stripped for the purposes of object-based + dependency tracking, but scopes within Templates are respected. + initial_value: See `variable_scope.variable_creator_scope`. Taken + explicitly so the argument can be re-named and used with + `Checkpointable._add_variable_with_custom_getter`. + checkpointable_parent: If not None, a more deeply nested Template object + to add a dependency on (rather than depending on the variable directly). + **kwargs: Passed through to the next creator. + Returns: + The output of `next_creator`: the fetched/created variable object. + """ + def _call_next_creator_renaming_initializer(initializer, **inner_kwargs): + inner_kwargs.pop("name") # Ignored; this is the scope-stripped name which + # we don't want to propagate. + return next_creator( + initial_value=initializer, + name=name, + **inner_kwargs) + if name.startswith(self._variable_scope.name): + scope_stripped_name = name[len(self._variable_scope.name) + 1:] + if not checkpointable_parent: + return self._add_variable_with_custom_getter( + initializer=initial_value, + name=scope_stripped_name, + getter=_call_next_creator_renaming_initializer, + # Disable error checking for Checkpointable. Exceptions are instead + # raised if necessary when the object-based saver tries to + # save/restore the object. + overwrite=True, + checkpointable_parent=self, + **kwargs) + else: + self._track_checkpointable( + checkpointable_parent, + name=checkpointable_parent._variable_scope.name[ # pylint: disable=protected-access + len(self._variable_scope.name) + 1:], + overwrite=True) + return next_creator(name=name, initial_value=initial_value, + checkpointable_parent=self, **kwargs) + def _call_func(self, args, kwargs): try: vars_at_start = len(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)) trainable_at_start = len( ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)) - result = self._func(*args, **kwargs) + if self._variables_created: + result = self._func(*args, **kwargs) + else: + # The first time we run, restore variables if necessary (via + # Checkpointable). + with variable_scope.variable_creator_scope( + self._checkpointable_custom_creator): + result = self._func(*args, **kwargs) if self._variables_created: # Variables were previously created, implying this is not the first @@ -543,7 +647,7 @@ class EagerTemplate(Template): Raises: RuntimeError: if eager execution is not enabled. """ - if not context.in_eager_mode(): + if not context.executing_eagerly(): raise RuntimeError( "{} objects can only be used when eager execution is enabled, use " "tf.Template for graph construction". @@ -557,12 +661,20 @@ class EagerTemplate(Template): # is created in __call__. variable_scope_name = None self._template_store = _EagerTemplateVariableStore(variable_scope_name) + self._variable_scope_context_manager = None def _call_func(self, args, kwargs): try: vars_at_start = self._template_store.variables() trainable_at_start = self._template_store.trainable_variables() - result = self._func(*args, **kwargs) + if self._variables_created: + result = self._func(*args, **kwargs) + else: + # The first time we run, restore variables if necessary (via + # Checkpointable). + with variable_scope.variable_creator_scope( + self._checkpointable_custom_creator): + result = self._func(*args, **kwargs) if self._variables_created: # Variables were previously created, implying this is not the first @@ -611,8 +723,12 @@ class EagerTemplate(Template): # the variable scope is opened in order to ensure that templates nested at # the same level correctly uniquify lower variable scope names. if self._variable_scope: - with variable_scope.variable_scope( - self._variable_scope, reuse=variable_scope.AUTO_REUSE): + # Create a cache for the variable scope context manager the first time + # around so that we don't have to keep recreating it. + if not self._variable_scope_context_manager: + self._variable_scope_context_manager = variable_scope.variable_scope( + self._variable_scope, reuse=variable_scope.AUTO_REUSE) + with self._variable_scope_context_manager: with self._template_store.as_default(): result = self._call_func(args, kwargs) return result diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py index 5cdf03509e3c427deec7e26345059211001e2131..2f6badcb532c0ef9d82b211d0c7b11a67e8e3010 100644 --- a/tensorflow/python/ops/tensor_array_ops.py +++ b/tensorflow/python/ops/tensor_array_ops.py @@ -148,7 +148,7 @@ class _GraphTensorArray(object): # will retroactively set the device value of this op. def create(): """Create the TensorArray op.""" - return gen_data_flow_ops._tensor_array_v3( + return gen_data_flow_ops.tensor_array_v3( dtype=dtype, size=size, element_shape=element_shape, @@ -237,7 +237,7 @@ class _GraphTensorArray(object): flow = self.flow with ops.name_scope(name, "TensorArrayGrad", [self._handle]): with ops.colocate_with(self._handle): - g_handle, unused_flow = gen_data_flow_ops._tensor_array_grad_v3( + g_handle, unused_flow = gen_data_flow_ops.tensor_array_grad_v3( handle=self._handle, source=source, flow_in=flow, name=name) with ops.control_dependencies([g_handle]): flow = array_ops.identity(flow, name="gradient_flow") @@ -252,7 +252,7 @@ class _GraphTensorArray(object): def read(self, index, name=None): """See TensorArray.""" - value = gen_data_flow_ops._tensor_array_read_v3( + value = gen_data_flow_ops.tensor_array_read_v3( handle=self._handle, index=index, flow_in=self._flow, @@ -270,7 +270,7 @@ class _GraphTensorArray(object): if self._infer_shape: self._merge_element_shape(value.shape) with self._maybe_colocate_with(value): - flow_out = gen_data_flow_ops._tensor_array_write_v3( + flow_out = gen_data_flow_ops.tensor_array_write_v3( handle=self._handle, index=index, value=value, @@ -296,7 +296,7 @@ class _GraphTensorArray(object): element_shape = self._element_shape[0] else: element_shape = tensor_shape.TensorShape(None) - value = gen_data_flow_ops._tensor_array_gather_v3( + value = gen_data_flow_ops.tensor_array_gather_v3( handle=self._handle, indices=indices, flow_in=self._flow, @@ -314,7 +314,7 @@ class _GraphTensorArray(object): tensor_shape.TensorShape(self._element_shape[0].dims[1:])) else: element_shape_except0 = tensor_shape.TensorShape(None) - value, _ = gen_data_flow_ops._tensor_array_concat_v3( + value, _ = gen_data_flow_ops.tensor_array_concat_v3( handle=self._handle, flow_in=self._flow, dtype=self._dtype, @@ -338,10 +338,10 @@ class _GraphTensorArray(object): with ops.name_scope(name, "TensorArrayScatter", [self._handle, value, indices]): value = ops.convert_to_tensor(value, name="value") - if self._infer_shape and context.in_graph_mode(): + if self._infer_shape and not context.executing_eagerly(): self._merge_element_shape(value.shape[1:]) with self._maybe_colocate_with(value): - flow_out = gen_data_flow_ops._tensor_array_scatter_v3( + flow_out = gen_data_flow_ops.tensor_array_scatter_v3( handle=self._handle, indices=indices, value=value, @@ -363,14 +363,14 @@ class _GraphTensorArray(object): value = ops.convert_to_tensor(value, name="value") with self._maybe_colocate_with(value): lengths_64 = math_ops.to_int64(lengths) - if self._infer_shape and context.in_graph_mode(): + if self._infer_shape and not context.executing_eagerly(): clengths = tensor_util.constant_value(lengths_64) if value.shape.dims is not None: if clengths is not None and clengths.max() == clengths.min(): self._merge_element_shape( tensor_shape.TensorShape([clengths[0]]).concatenate( value.shape[1:])) - flow_out = gen_data_flow_ops._tensor_array_split_v3( + flow_out = gen_data_flow_ops.tensor_array_split_v3( handle=self._handle, value=value, lengths=lengths_64, @@ -386,13 +386,13 @@ class _GraphTensorArray(object): def size(self, name=None): """See TensorArray.""" - return gen_data_flow_ops._tensor_array_size_v3( + return gen_data_flow_ops.tensor_array_size_v3( handle=self._handle, flow_in=self.flow, name=name) @tf_should_use.should_use_result def close(self, name=None): """See TensorArray.""" - return gen_data_flow_ops._tensor_array_close_v3( + return gen_data_flow_ops.tensor_array_close_v3( handle=self._handle, name=name) # pylint: enable=protected-access @@ -653,7 +653,7 @@ class _EagerTensorArray(object): if len(tensors) > len(self._tensor_array) and not self._dynamic_size: raise ValueError( "Cannot unstack %d tensors into a TensorArray of static size %d" % - (len(tensors), len(self._tensors))) + (len(tensors), len(self._tensor_array))) ta = self._identity_without_array() ta._implementation._tensor_array = tensors # pylint: disable=protected-access return ta @@ -774,10 +774,10 @@ class TensorArray(object): ValueError: if both handle and tensor_array_name are provided. TypeError: if handle is provided but is not a Tensor. """ - if context.in_graph_mode(): - implementation = _GraphTensorArray - else: + if context.executing_eagerly(): implementation = _EagerTensorArray + else: + implementation = _GraphTensorArray self._implementation = implementation( dtype, diff --git a/tensorflow/python/ops/variable_scope.py b/tensorflow/python/ops/variable_scope.py index 81565a63774da49628d100ef071b02f6311f6af2..c1af8ff8d3108d60b9007552ef1aadf98fea32e6 100644 --- a/tensorflow/python/ops/variable_scope.py +++ b/tensorflow/python/ops/variable_scope.py @@ -321,7 +321,7 @@ class _VariableStore(object): raise ValueError( "Passed a custom_getter which is not callable: %s" % custom_getter) - if context.in_eager_mode(): + if context.executing_eagerly(): if not self._store_eager_variables and reuse: raise RuntimeError( "When eager execution is enabled variable reuse is only supported" @@ -518,7 +518,7 @@ class _VariableStore(object): when violating reuse during variable creation, or if an existing sharded variable exists for the given name but with different sharding. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise NotImplementedError("Partitioned variables are not yet supported " "when eager execution is enabled.") @@ -798,7 +798,7 @@ class _VariableStore(object): validate_shape=validate_shape, constraint=constraint, use_resource=use_resource) - if context.in_graph_mode() or self._store_eager_variables: + if not context.executing_eagerly() or self._store_eager_variables: # In eager mode we do not want to keep default references to Variable # objects as this will prevent their memory from being released. self._vars[name] = v @@ -811,12 +811,12 @@ class _VariableStore(object): with ops.name_scope(name + "/Regularizer/"): loss = regularizer(v) if loss is not None: - if context.in_graph_mode(): - v_name = v.name - loss_name = loss.name - else: + if context.executing_eagerly(): v_name = "v_%s" % type(v) loss_name = "loss_%s" % type(loss) + else: + v_name = v.name + loss_name = loss.name logging.vlog(1, "Applied regularizer to %s and added the result %s " "to REGULARIZATION_LOSSES.", v_name, loss_name) ops.add_to_collection(ops.GraphKeys.REGULARIZATION_LOSSES, loss) @@ -920,7 +920,7 @@ class VariableScope(object): self._dtype = dtype self._use_resource = use_resource self._constraint = constraint - if context.in_eager_mode(): + if context.executing_eagerly(): if self._caching_device is not None: raise NotImplementedError("Caching devices is not yet supported " "when eager execution is enabled.") @@ -988,7 +988,7 @@ class VariableScope(object): def set_use_resource(self, use_resource): """Sets whether to use ResourceVariables for this scope.""" - if context.in_eager_mode() and not use_resource: + if context.executing_eagerly() and not use_resource: raise ValueError("When eager execution is enabled, " "use_resource cannot be set to false.") self._use_resource = use_resource @@ -999,14 +999,14 @@ class VariableScope(object): def set_caching_device(self, caching_device): """Set caching_device for this scope.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise NotImplementedError("Caching devices are not yet supported " "when eager execution is enabled.") self._caching_device = caching_device def set_partitioner(self, partitioner): """Set partitioner for this scope.""" - if partitioner and context.in_eager_mode(): + if partitioner and context.executing_eagerly(): raise NotImplementedError("Partitioned variables are not yet supported " "when eager execution is enabled.") self._partitioner = partitioner @@ -1057,14 +1057,14 @@ class VariableScope(object): partitioner = self._partitioner if custom_getter is None: custom_getter = self._custom_getter - if context.in_graph_mode(): + if context.executing_eagerly(): + reuse = False + use_resource = True + else: if reuse is None: reuse = self._reuse if use_resource is None: use_resource = self._use_resource - else: - reuse = False - use_resource = True full_name = self.name + "/" + name if self.name else name # Variable names only depend on variable_scope (full_name here), @@ -1107,7 +1107,7 @@ class VariableScope(object): use_resource=None, constraint=None): """Gets an existing variable with this name or create a new one.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise NotImplementedError("Partitioned variables are not yet supported " "when eager execution is enabled.") if initializer is None: @@ -1274,6 +1274,9 @@ class EagerVariableStore(object): # pylint: enable=protected-access +# The argument list for get_variable must match arguments to get_local_variable. +# So, if you are updating the arguments, also update arguments to +# get_local_variable below. @tf_export("get_variable") def get_variable(name, shape=None, @@ -1385,15 +1388,32 @@ get_variable.__doc__ = get_variable_or_local_docstring % ( "GraphKeys.GLOBAL_VARIABLES") -@functools.wraps(get_variable) +# The argument list for get_local_variable must match arguments to get_variable. +# So, if you are updating the arguments, also update arguments to get_variable. @tf_export("get_local_variable") -def get_local_variable(*args, **kwargs): - kwargs["trainable"] = False - if "collections" in kwargs: - kwargs["collections"] += [ops.GraphKeys.LOCAL_VARIABLES] +def get_local_variable(name, + shape=None, + dtype=None, + initializer=None, + regularizer=None, + trainable=False, # pylint: disable=unused-argument + collections=None, + caching_device=None, + partitioner=None, + validate_shape=True, + use_resource=None, + custom_getter=None, + constraint=None): + if collections: + collections += [ops.GraphKeys.LOCAL_VARIABLES] else: - kwargs["collections"] = [ops.GraphKeys.LOCAL_VARIABLES] - return get_variable(*args, **kwargs) + collections = [ops.GraphKeys.LOCAL_VARIABLES] + return get_variable( + name, shape=shape, dtype=dtype, initializer=initializer, + regularizer=regularizer, trainable=False, collections=collections, + caching_device=caching_device, partitioner=partitioner, + validate_shape=validate_shape, use_resource=use_resource, + custom_getter=custom_getter, constraint=constraint) get_local_variable.__doc__ = get_variable_or_local_docstring % ( "Gets an existing *local* variable or creates a new one.", "Behavior is the same as in `get_variable`, except that variables are\n" @@ -1871,7 +1891,7 @@ class variable_scope(object): raise ValueError("The reuse parameter must be True or False or None.") if self._values is None: self._values = [] - self._in_graph_mode = not context.in_eager_mode() + self._in_graph_mode = not context.executing_eagerly() if self._in_graph_mode: self._graph = ops._get_graph_from_inputs(self._values) # pylint: disable=protected-access self._cached_pure_variable_scope = None @@ -2111,13 +2131,13 @@ def default_variable_creator(next_creator=None, **kwargs): use_resource = kwargs.get("use_resource", None) if use_resource is None: use_resource = get_variable_scope().use_resource - if use_resource or (use_resource is None and context.in_eager_mode()): + if use_resource or (use_resource is None and context.executing_eagerly()): return resource_variable_ops.ResourceVariable( initial_value=initial_value, trainable=trainable, collections=collections, validate_shape=validate_shape, caching_device=caching_device, name=name, dtype=dtype, constraint=constraint) - elif not use_resource and context.in_eager_mode(): + elif not use_resource and context.executing_eagerly(): raise RuntimeError( "VariableScope should use resource variable when eager execution is" " enabled, but use_resource is False." @@ -2145,7 +2165,7 @@ def variable(initial_value=None, constraint=None, use_resource=None): previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs) - for getter in ops.get_default_graph()._get_variable_creator_stack(): # pylint: disable=protected-access + for getter in ops.get_default_graph()._variable_creator_stack: # pylint: disable=protected-access previous_getter = _make_getter(getter, previous_getter) return previous_getter(initial_value=initial_value, trainable=trainable, diff --git a/tensorflow/python/ops/variables.py b/tensorflow/python/ops/variables.py index 19e3298e4019f94132db25ab0dae5ed458bfbeb3..5b9947f4417a0f7621fda885c685021d91701c2b 100644 --- a/tensorflow/python/ops/variables.py +++ b/tensorflow/python/ops/variables.py @@ -29,6 +29,7 @@ from tensorflow.python.ops import gen_array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import state_ops from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpointable from tensorflow.python.util import compat from tensorflow.python.util import tf_should_use from tensorflow.python.util.deprecation import deprecated @@ -36,7 +37,7 @@ from tensorflow.python.util.tf_export import tf_export @tf_export("Variable") -class Variable(object): +class Variable(checkpointable.CheckpointableBase): """See the @{$variables$Variables How To} for a high level overview. A variable maintains state in the graph across calls to `run()`. You add a @@ -209,10 +210,11 @@ class Variable(object): for details on how variables work in eager execution. @end_compatibility """ - if not context.in_graph_mode(): - raise RuntimeError("tf.Variable not supported in Eager mode. " - "Please use tfe.Variable instead") - self._in_graph_mode = context.in_graph_mode() + if context.executing_eagerly(): + raise RuntimeError( + "tf.Variable not supported when eager execution is enabled. " + "Please use tf.contrib.eager.Variable instead") + self._in_graph_mode = True if variable_def: # If variable_def is provided, recreates the variable from its fields. if initial_value: @@ -233,7 +235,7 @@ class Variable(object): constraint=constraint) def __repr__(self): - if context.in_eager_mode(): + if context.executing_eagerly(): return "" % ( self.name, self.get_shape(), self.dtype.name, ops.numpy_text(self.read_value(), is_repr=True)) @@ -306,6 +308,14 @@ class Variable(object): if constraint is not None and not callable(constraint): raise ValueError("The `constraint` argument must be a callable.") + # Store the graph key so optimizers know how to only retrieve variables from + # this graph. + self._graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access + if isinstance(initial_value, checkpointable.CheckpointInitialValue): + self._maybe_initialize_checkpointable() + self._update_uid = initial_value.checkpoint_position.restore_uid + initial_value = initial_value.wrapped_value + if trainable and ops.GraphKeys.TRAINABLE_VARIABLES not in collections: collections = list(collections) + [ops.GraphKeys.TRAINABLE_VARIABLES] with ops.init_scope(): @@ -731,15 +741,15 @@ class Variable(object): Raises: ValueError: Session is not passed and no default session """ - if context.in_graph_mode(): + if context.executing_eagerly(): + self.assign(value) + else: session = session or ops.get_default_session() if session is None: raise ValueError( "Either session argument should be provided or default session " "should be established") session.run(self._initializer_op, {self._initializer_op.inputs[1]: value}) - else: - self.assign(value) # Conversion to tensor. @staticmethod @@ -786,6 +796,10 @@ class Variable(object): setattr(Variable, operator, _run_op) + def _gather_saveables_for_checkpoint(self): + """For implementing `Checkpointable`. This object is saveable on its own.""" + return {checkpointable.VARIABLE_VALUE_KEY: self} + def _try_guard_against_uninitialized_dependencies(self, initial_value): """Attempt to guard against dependencies on uninitialized variables. @@ -1235,9 +1249,9 @@ class PartitionedVariable(object): information does not match `shape`, or `partitions` has invalid values. RuntimeError: If eager execution is enabled """ - if not context.in_graph_mode(): - raise RuntimeError("tf.PartitionedVariable not supported in " - "eager mode. Please use tfe.Variable instead") + if context.executing_eagerly(): + raise RuntimeError( + "tf.PartitionedVariable not supported with eager execution enabled.") if not isinstance(variable_list, (list, tuple)): raise TypeError( "variable_list is not a list or tuple: %s" % variable_list) @@ -1528,7 +1542,7 @@ def variables_initializer(var_list, name="init"): Returns: An Op that run the initializers of all the specified variables. """ - if var_list and context.in_graph_mode(): + if var_list and not context.executing_eagerly(): return control_flow_ops.group(*[v.initializer for v in var_list], name=name) return control_flow_ops.no_op(name=name) @@ -1550,7 +1564,7 @@ def global_variables_initializer(): Returns: An Op that initializes global variables in the graph. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return control_flow_ops.no_op(name="global_variables_initializer") return variables_initializer(global_variables()) @@ -1572,7 +1586,7 @@ def local_variables_initializer(): Returns: An Op that initializes all local variables in the graph. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return control_flow_ops.no_op(name="local_variables_initializer") return variables_initializer(local_variables()) diff --git a/tensorflow/python/platform/app.py b/tensorflow/python/platform/app.py index 9b92d9a18005ca5e6be3820427e3a3ba60a8ec2d..cce64c0ccafc29a9d0d0b51b4c97c5673264657b 100644 --- a/tensorflow/python/platform/app.py +++ b/tensorflow/python/platform/app.py @@ -23,6 +23,7 @@ import sys as _sys from tensorflow.python.platform import flags from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export def _usage(shorthelp): @@ -108,6 +109,7 @@ def _define_help_flags(): _define_help_flags_called = True +@tf_export('app.run') def run(main=None, argv=None): """Runs the program with an optional 'main' function and 'argv' list.""" diff --git a/tensorflow/python/platform/googletest.py b/tensorflow/python/platform/googletest.py index 96219faab719e28a5fa8a9a21c83f81a6f8478e6..8141cf92c568f257a5e9810318182d71f445dfa1 100644 --- a/tensorflow/python/platform/googletest.py +++ b/tensorflow/python/platform/googletest.py @@ -36,6 +36,7 @@ from tensorflow.python.platform import benchmark from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect +from tensorflow.python.util.tf_export import tf_export Benchmark = benchmark.TensorFlowBenchmark # pylint: disable=invalid-name @@ -138,6 +139,7 @@ def StatefulSessionAvailable(): return False +@tf_export('test.StubOutForTesting') class StubOutForTesting(object): """Support class for stubbing methods out for unit testing. diff --git a/tensorflow/python/platform/resource_loader.py b/tensorflow/python/platform/resource_loader.py index 2455acb4c0c469acbb928c4ec44571e50e06de1f..8f7b12e2b2b92d9b2bfe397d0e7cba59e11bc1f6 100644 --- a/tensorflow/python/platform/resource_loader.py +++ b/tensorflow/python/platform/resource_loader.py @@ -29,8 +29,10 @@ import sys as _sys from tensorflow.python.util import tf_inspect as _inspect from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export +@tf_export('resource_loader.load_resource') def load_resource(path): """Load the resource at given path, where path is relative to tensorflow/. @@ -52,6 +54,7 @@ def load_resource(path): # pylint: disable=protected-access +@tf_export('resource_loader.get_data_files_path') def get_data_files_path(): """Get a direct path to the data files colocated with the script. @@ -62,6 +65,7 @@ def get_data_files_path(): return _os.path.dirname(_inspect.getfile(_sys._getframe(1))) +@tf_export('resource_loader.get_root_dir_with_all_resources') def get_root_dir_with_all_resources(): """Get a root directory containing all the data attributes in the build rule. @@ -101,6 +105,7 @@ def get_root_dir_with_all_resources(): return data_files_dir or script_dir +@tf_export('resource_loader.get_path_to_datafile') def get_path_to_datafile(path): """Get the path to the specified file in the data dependencies. @@ -120,6 +125,7 @@ def get_path_to_datafile(path): return _os.path.join(data_files_path, path) +@tf_export('resource_loader.readahead_file_path') def readahead_file_path(path, readahead='128M'): # pylint: disable=unused-argument """Readahead files not implemented; simply returns given path.""" return path diff --git a/tensorflow/python/platform/sysconfig.py b/tensorflow/python/platform/sysconfig.py index 5c50fa023dc3b216838390d9356a39e70e2362d2..fdd2b903fc79c40a26392714328f74756f3fff92 100644 --- a/tensorflow/python/platform/sysconfig.py +++ b/tensorflow/python/platform/sysconfig.py @@ -68,7 +68,6 @@ def get_compile_flags(): """ flags = [] flags.append('-I%s' % get_include()) - flags.append('-I%s/external/nsync/public' % get_include()) flags.append('-D_GLIBCXX_USE_CXX11_ABI=%d' % _CXX11_ABI_FLAG) return flags diff --git a/tensorflow/python/platform/test.py b/tensorflow/python/platform/test.py index 9b7655722ac5a917f2753617f8e99bf2bd2f8d11..1660791febc9da93f3a3a977a17ca876e772a9a5 100644 --- a/tensorflow/python/platform/test.py +++ b/tensorflow/python/platform/test.py @@ -62,6 +62,8 @@ if sys.version_info.major == 2: else: from unittest import mock # pylint: disable=g-import-not-at-top +tf_export('test.mock')(mock) + # Import Benchmark class Benchmark = _googletest.Benchmark # pylint: disable=invalid-name diff --git a/tensorflow/python/platform/tf_logging.py b/tensorflow/python/platform/tf_logging.py index 85ed4f071c7022801f20db75d538e5917b8eea66..22aabfd7121ac9b2eebeae2693f174e044d504ef 100644 --- a/tensorflow/python/platform/tf_logging.py +++ b/tensorflow/python/platform/tf_logging.py @@ -35,6 +35,7 @@ import threading import six from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export # Don't use this directly. Use _get_logger() instead. @@ -90,30 +91,37 @@ def _get_logger(): _logger_lock.release() +@tf_export('logging.log') def log(level, msg, *args, **kwargs): _get_logger().log(level, msg, *args, **kwargs) +@tf_export('logging.debug') def debug(msg, *args, **kwargs): _get_logger().debug(msg, *args, **kwargs) +@tf_export('logging.error') def error(msg, *args, **kwargs): _get_logger().error(msg, *args, **kwargs) +@tf_export('logging.fatal') def fatal(msg, *args, **kwargs): _get_logger().fatal(msg, *args, **kwargs) +@tf_export('logging.info') def info(msg, *args, **kwargs): _get_logger().info(msg, *args, **kwargs) +@tf_export('logging.warn') def warn(msg, *args, **kwargs): _get_logger().warn(msg, *args, **kwargs) +@tf_export('logging.warning') def warning(msg, *args, **kwargs): _get_logger().warning(msg, *args, **kwargs) @@ -136,15 +144,18 @@ _log_prefix = None # later set to google2_log_prefix _log_counter_per_token = {} +@tf_export('logging.TaskLevelStatusMessage') def TaskLevelStatusMessage(msg): error(msg) +@tf_export('logging.flush') def flush(): raise NotImplementedError() # Code below is taken from pyglib/logging +@tf_export('logging.vlog') def vlog(level, msg, *args, **kwargs): _get_logger().log(level, msg, *args, **kwargs) @@ -164,6 +175,7 @@ def _GetNextLogCountPerToken(token): return _log_counter_per_token[token] +@tf_export('logging.log_every_n') def log_every_n(level, msg, n, *args): """Log 'msg % args' at level 'level' once per 'n' times. @@ -180,6 +192,7 @@ def log_every_n(level, msg, n, *args): log_if(level, msg, not (count % n), *args) +@tf_export('logging.log_first_n') def log_first_n(level, msg, n, *args): # pylint: disable=g-bad-name """Log 'msg % args' at level 'level' only first 'n' times. @@ -195,6 +208,7 @@ def log_first_n(level, msg, n, *args): # pylint: disable=g-bad-name log_if(level, msg, count < n, *args) +@tf_export('logging.log_if') def log_if(level, msg, condition, *args): """Log 'msg % args' at level 'level' only if condition is fulfilled.""" if condition: @@ -251,11 +265,13 @@ def google2_log_prefix(level, timestamp=None, file_and_line=None): return s +@tf_export('logging.get_verbosity') def get_verbosity(): """Return how much logging output will be produced.""" return _get_logger().getEffectiveLevel() +@tf_export('logging.set_verbosity') def set_verbosity(v): """Sets the threshold for what messages will be logged.""" _get_logger().setLevel(v) @@ -296,4 +312,10 @@ _allowed_symbols = [ 'warning', ] +tf_export('logging.DEBUG').export_constant(__name__, 'DEBUG') +tf_export('logging.ERROR').export_constant(__name__, 'ERROR') +tf_export('logging.FATAL').export_constant(__name__, 'FATAL') +tf_export('logging.INFO').export_constant(__name__, 'INFO') +tf_export('logging.WARN').export_constant(__name__, 'WARN') + remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/python/profiler/model_analyzer.py b/tensorflow/python/profiler/model_analyzer.py index 8f780545607f7ba2337c83ad2c3740f542b802f6..acf02096fffe8b38e68824878fa698ed69d3895c 100644 --- a/tensorflow/python/profiler/model_analyzer.py +++ b/tensorflow/python/profiler/model_analyzer.py @@ -33,6 +33,7 @@ from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.profiler import option_builder from tensorflow.python.profiler import tfprof_logger +from tensorflow.python.util.tf_export import tf_export _DEFAULT_PROFILE_OPTIONS = 0 _DEFAULT_ADVISE_OPTIONS = 0 @@ -121,6 +122,7 @@ def _build_advisor_options(options): return opts +@tf_export('profiler.Profiler') class Profiler(object): """TensorFlow multi-step profiler. @@ -170,7 +172,7 @@ class Profiler(object): op_log: optional. tensorflow::tfprof::OpLogProto proto. Used to define extra op types. """ - if not graph and context.in_graph_mode(): + if not graph and not context.executing_eagerly(): graph = ops.get_default_graph() self._coverage = 0.0 self._graph = graph @@ -304,6 +306,7 @@ class Profiler(object): print_mdl.WriteProfile(filename) +@tf_export('profiler.profile') def profile(graph=None, run_meta=None, op_log=None, @@ -333,7 +336,7 @@ def profile(graph=None, If cmd is 'op' or 'code', returns MultiGraphNodeProto proto. Side effect: stdout/file/timeline.json depending on options['output'] """ - if not graph and context.in_graph_mode(): + if not graph and not context.executing_eagerly(): graph = ops.get_default_graph() if options == _DEFAULT_PROFILE_OPTIONS: @@ -378,6 +381,7 @@ def profile(graph=None, return tfprof_node +@tf_export('profiler.advise') def advise(graph=None, run_meta=None, options=_DEFAULT_ADVISE_OPTIONS): """Auto profile and advise. diff --git a/tensorflow/python/profiler/option_builder.py b/tensorflow/python/profiler/option_builder.py index 13942ad6a2adc1f1d1cad778ebd280d358f64a59..2ad7adf76933df65ca795dca361397f436adb995 100644 --- a/tensorflow/python/profiler/option_builder.py +++ b/tensorflow/python/profiler/option_builder.py @@ -20,8 +20,10 @@ from __future__ import print_function import copy from tensorflow.python.profiler import tfprof_logger +from tensorflow.python.util.tf_export import tf_export +@tf_export('profiler.ProfileOptionBuilder') class ProfileOptionBuilder(object): # pylint: disable=line-too-long """Option Builder for Profiling API. @@ -298,7 +300,7 @@ class ProfileOptionBuilder(object): # pylint: disable=line-too-long """Only show profiler nodes consuming no less than 'min_float_ops'. - Please see https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profilerg3doc/profile_model_architecture.md + Please see https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/g3doc/profile_model_architecture.md on the caveats of calculating float operations. Args: diff --git a/tensorflow/python/profiler/profiler.py b/tensorflow/python/profiler/profiler.py index 130dcb5134d6f7e6eb43aebea803b366a5ce27d8..fa7f30b236997cecd6d5df98c334aa7f5cc571e4 100644 --- a/tensorflow/python/profiler/profiler.py +++ b/tensorflow/python/profiler/profiler.py @@ -31,6 +31,7 @@ from tensorflow.python.profiler.option_builder import ProfileOptionBuilder from tensorflow.python.profiler.tfprof_logger import write_op_log from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export _allowed_symbols = [ @@ -48,6 +49,12 @@ _allowed_symbols.extend([ 'OpLogProto', ]) +# Export protos +tf_export('profiler.GraphNodeProto')(GraphNodeProto) +tf_export('profiler.MultiGraphNodeProto')(MultiGraphNodeProto) +tf_export('profiler.AdviceProto')(AdviceProto) +tf_export('profiler.OpLogProto')(OpLogProto) + remove_undocumented(__name__, _allowed_symbols, [ Profiler, profile, diff --git a/tensorflow/python/profiler/tfprof_logger.py b/tensorflow/python/profiler/tfprof_logger.py index ffda7ddad759ce68bf718bcfa6e568cfadd59b53..e651de32ea3bce32a965bfbeefc76ff08a79ac38 100644 --- a/tensorflow/python/profiler/tfprof_logger.py +++ b/tensorflow/python/profiler/tfprof_logger.py @@ -30,6 +30,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.platform import gfile from tensorflow.python.profiler.internal import flops_registry # pylint: disable=unused-import +from tensorflow.python.util.tf_export import tf_export TRAINABLE_VARIABLES = '_trainable_variables' REGISTERED_FLOP_STATS = 'flops' @@ -155,7 +156,7 @@ def merge_default_with_oplog(graph, op_log=None, run_meta=None, Returns: tmp_op_log: Merged OpLogProto proto. """ - if not graph and context.in_graph_mode(): + if not graph and not context.executing_eagerly(): graph = ops.get_default_graph() tmp_op_log = tfprof_log_pb2.OpLogProto() @@ -187,6 +188,7 @@ def merge_default_with_oplog(graph, op_log=None, run_meta=None, return tmp_op_log +@tf_export('profiler.write_op_log') def write_op_log(graph, log_dir, op_log=None, run_meta=None, add_trace=True): """Log provided 'op_log', and add additional model information below. @@ -208,7 +210,7 @@ def write_op_log(graph, log_dir, op_log=None, run_meta=None, add_trace=True): add_trace: Whether to add python code trace information. Used to support "code" view. """ - if not graph and context.in_graph_mode(): + if not graph and not context.executing_eagerly(): graph = ops.get_default_graph() op_log = merge_default_with_oplog(graph, op_log, run_meta, add_trace) diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i index 3f25311a8361d11fbc583413708e148648d95906..39fabb9c1bc646a09557293c1f645a8b97f5bbdd 100644 --- a/tensorflow/python/pywrap_tfe.i +++ b/tensorflow/python/pywrap_tfe.i @@ -26,11 +26,18 @@ limitations under the License. %rename("%s") TFE_ContextClearCaches; %rename("%s") TFE_ContextGetDevicePlacementPolicy; %rename("%s") TFE_ContextSetThreadLocalDevicePlacementPolicy; +%rename("%s") TFE_ContextSetAsyncForThread; +%rename("%s") TFE_ContextAsyncWait; +%rename("%s") TFE_ContextAsyncClearError; %rename("%s") TFE_OpNameGetAttrType; %rename("%s") TFE_Py_InitEagerTensor; %rename("%s") TFE_Py_RegisterExceptionClass; +%rename("%s") TFE_Py_RegisterBackwardFunctionGetter; +%rename("%s") TFE_Py_RegisterFallbackExceptionClass; +%rename("%s") TFE_Py_RegisterResourceVariableType; %rename("%s") TFE_Py_Execute; %rename("%s") TFE_Py_FastPathExecute; +%rename("%s") TFE_Py_RecordGradient; %rename("%s") TFE_Py_UID; %rename("%s") TFE_Py_TapeSetNew; %rename("%s") TFE_Py_TapeSetRemove; @@ -47,6 +54,7 @@ limitations under the License. %rename("%s") TFE_NewContextOptions; %rename("%s") TFE_ContextOptionsSetConfig; %rename("%s") TFE_ContextOptionsSetDevicePlacementPolicy; +%rename("%s") TFE_ContextOptionsSetAsync; %rename("%s") TFE_DeleteContextOptions; %rename("%s") TFE_Py_TensorShapeSlice; diff --git a/tensorflow/python/saved_model/BUILD b/tensorflow/python/saved_model/BUILD index e34aa7cc2ca41ecdd7c9ff52ab8f3d552f26fe69..30e0a099d8b2e30cff36b69164ba9f1789dd8916 100644 --- a/tensorflow/python/saved_model/BUILD +++ b/tensorflow/python/saved_model/BUILD @@ -148,6 +148,7 @@ py_test( "//tensorflow/python:math_ops", "//tensorflow/python:saver_test_utils", "//tensorflow/python:state_ops", + "//tensorflow/python:test_ops", "//tensorflow/python:util", "//tensorflow/python:variables", ], diff --git a/tensorflow/python/saved_model/builder_impl.py b/tensorflow/python/saved_model/builder_impl.py index 62ee53b816c2a38327fa116d2924446e6bf24a1e..3447d917e9bf2dace3de784106dadb1fcc3a9647 100644 --- a/tensorflow/python/saved_model/builder_impl.py +++ b/tensorflow/python/saved_model/builder_impl.py @@ -34,8 +34,10 @@ from tensorflow.python.platform import tf_logging from tensorflow.python.saved_model import constants from tensorflow.python.training import saver as tf_saver from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export +@tf_export("saved_model.builder.SavedModelBuilder") class SavedModelBuilder(object): """Builds the `SavedModel` protocol buffer and saves variables and assets. @@ -191,7 +193,8 @@ class SavedModelBuilder(object): def _validate_tensor_info(self, tensor_info): """Validates the `TensorInfo` proto. - Checks if the `name` and `dtype` fields exist and are non-empty. + Checks if the `encoding` (`name` or `coo_sparse`) and `dtype` fields exist + and are non-empty. Args: tensor_info: `TensorInfo` protocol buffer to validate. @@ -204,10 +207,12 @@ class SavedModelBuilder(object): raise AssertionError( "All TensorInfo protos used in the SignatureDefs must have the name " "and dtype fields set.") - if not tensor_info.name: + if tensor_info.WhichOneof("encoding") is None: + # TODO(soergel) validate each of the fields of coo_sparse raise AssertionError( - "All TensorInfo protos used in the SignatureDefs must have the name " - "field set: %s" % tensor_info) + "All TensorInfo protos used in the SignatureDefs must have one of " + "the 'encoding' fields (e.g., name or coo_sparse) set: %s" + % tensor_info) if tensor_info.dtype is types_pb2.DT_INVALID: raise AssertionError( "All TensorInfo protos used in the SignatureDefs must have the dtype " diff --git a/tensorflow/python/saved_model/constants.py b/tensorflow/python/saved_model/constants.py index 7e3e8df47fb0e024eae8add6a788d632709740af..ec49a0539ff52f6cc69bb24483ede657b698ab8d 100644 --- a/tensorflow/python/saved_model/constants.py +++ b/tensorflow/python/saved_model/constants.py @@ -20,33 +20,52 @@ from __future__ import division from __future__ import print_function from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export # Subdirectory name containing the asset files. ASSETS_DIRECTORY = "assets" +tf_export("saved_model.constants.ASSETS_DIRECTORY").export_constant( + __name__, "ASSETS_DIRECTORY") # CollectionDef key containing SavedModel assets. ASSETS_KEY = "saved_model_assets" +tf_export("saved_model.constants.ASSETS_KEY").export_constant( + __name__, "ASSETS_KEY") # CollectionDef key for the legacy init op. LEGACY_INIT_OP_KEY = "legacy_init_op" +tf_export("saved_model.constants.LEGACY_INIT_OP_KEY").export_constant( + __name__, "LEGACY_INIT_OP_KEY") # CollectionDef key for the SavedModel main op. MAIN_OP_KEY = "saved_model_main_op" +tf_export("saved_model.constants.MAIN_OP_KEY").export_constant( + __name__, "MAIN_OP_KEY") # Schema version for SavedModel. SAVED_MODEL_SCHEMA_VERSION = 1 +tf_export("saved_model.constants.SAVED_MODEL_SCHEMA_VERSION").export_constant( + __name__, "SAVED_MODEL_SCHEMA_VERSION") # File name for SavedModel protocol buffer. SAVED_MODEL_FILENAME_PB = "saved_model.pb" +tf_export("saved_model.constants.SAVED_MODEL_FILENAME_PB").export_constant( + __name__, "SAVED_MODEL_FILENAME_PB") # File name for text version of SavedModel protocol buffer. SAVED_MODEL_FILENAME_PBTXT = "saved_model.pbtxt" +tf_export("saved_model.constants.SAVED_MODEL_FILENAME_PBTXT").export_constant( + __name__, "SAVED_MODEL_FILENAME_PBTXT") # Subdirectory name containing the variables/checkpoint files. VARIABLES_DIRECTORY = "variables" +tf_export("saved_model.constants.VARIABLES_DIRECTORY").export_constant( + __name__, "VARIABLES_DIRECTORY") # File name used for variables. VARIABLES_FILENAME = "variables" +tf_export("saved_model.constants.VARIABLES_FILENAME").export_constant( + __name__, "VARIABLES_FILENAME") _allowed_symbols = [ diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py index 6e85df0cbf5623691d38b030036958e5955399ee..bebf1d5e0d3cc6ac0e431230577704365d37a437 100644 --- a/tensorflow/python/saved_model/loader_impl.py +++ b/tensorflow/python/saved_model/loader_impl.py @@ -32,6 +32,7 @@ from tensorflow.python.platform import tf_logging from tensorflow.python.saved_model import constants from tensorflow.python.training import saver as tf_saver from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export def _parse_saved_model(export_dir): @@ -156,6 +157,7 @@ def _get_legacy_init_op_tensor(meta_graph_def_to_load): return legacy_init_op_tensor +@tf_export("saved_model.loader.maybe_saved_model_directory") def maybe_saved_model_directory(export_dir): """Checks whether the provided export directory could contain a SavedModel. @@ -176,6 +178,7 @@ def maybe_saved_model_directory(export_dir): return file_io.file_exists(txt_path) or file_io.file_exists(pb_path) +@tf_export("saved_model.loader.load") def load(sess, tags, export_dir, **saver_kwargs): """Loads the model from a SavedModel as specified by tags. @@ -232,8 +235,9 @@ def load(sess, tags, export_dir, **saver_kwargs): asset_tensors_dictionary = _get_asset_tensors(export_dir, meta_graph_def_to_load) - main_op_tensor = (_get_main_op_tensor(meta_graph_def_to_load) or - (_get_legacy_init_op_tensor(meta_graph_def_to_load))) + main_op_tensor = ( + _get_main_op_tensor(meta_graph_def_to_load) or + (_get_legacy_init_op_tensor(meta_graph_def_to_load))) if main_op_tensor is not None: sess.run(fetches=[main_op_tensor], feed_dict=asset_tensors_dictionary) diff --git a/tensorflow/python/saved_model/main_op_impl.py b/tensorflow/python/saved_model/main_op_impl.py index 355fd57bf1d2166f58a5fdc95d04695ea05b56b3..631ee63729513d24c2ddae71b771f7cf1695358f 100644 --- a/tensorflow/python/saved_model/main_op_impl.py +++ b/tensorflow/python/saved_model/main_op_impl.py @@ -22,8 +22,10 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import lookup_ops from tensorflow.python.ops import variables +from tensorflow.python.util.tf_export import tf_export +@tf_export('saved_model.main_op.main_op') def main_op(): """Returns a main op to init variables and tables. @@ -40,6 +42,7 @@ def main_op(): # TODO(sukritiramesh): Integrate with Saver for complete restore functionality. +@tf_export('saved_model.main_op.main_op_with_restore') def main_op_with_restore(restore_op_name): """Returns a main op to init variables, tables and restore the graph. diff --git a/tensorflow/python/saved_model/saved_model_test.py b/tensorflow/python/saved_model/saved_model_test.py index f92247d52e4150b2347a95e84d4bbf9c6ffc258a..804255375e7c5215597a5dcca02f3b32f2c0a497 100644 --- a/tensorflow/python/saved_model/saved_model_test.py +++ b/tensorflow/python/saved_model/saved_model_test.py @@ -20,7 +20,6 @@ from __future__ import print_function import os -from tensorflow.core.framework import op_def_pb2 from tensorflow.core.framework import types_pb2 from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf import meta_graph_pb2 @@ -28,8 +27,8 @@ from tensorflow.python.client import session from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors -from tensorflow.python.framework import op_def_registry from tensorflow.python.framework import ops +from tensorflow.python.framework import test_ops from tensorflow.python.framework import test_util from tensorflow.python.lib.io import file_io from tensorflow.python.ops import control_flow_ops @@ -95,7 +94,7 @@ class SavedModelTest(test.TestCase): self.assertEqual(expected_asset_file_name, asset.filename) self.assertEqual(expected_asset_tensor_name, asset.tensor_info.name) - def _validate_inputs_tensor_info(self, builder, tensor_info): + def _validate_inputs_tensor_info_fail(self, builder, tensor_info): with self.test_session(graph=ops.Graph()) as sess: self._init_and_validate_variable(sess, "v", 42) @@ -108,7 +107,18 @@ class SavedModelTest(test.TestCase): sess, ["foo"], signature_def_map={"foo_key": foo_signature}) - def _validate_outputs_tensor_info(self, builder, tensor_info): + def _validate_inputs_tensor_info_accept(self, builder, tensor_info): + with self.test_session(graph=ops.Graph()) as sess: + self._init_and_validate_variable(sess, "v", 42) + + foo_signature = signature_def_utils.build_signature_def({ + "foo_inputs": tensor_info + }, dict(), "foo") + builder.add_meta_graph_and_variables( + sess, ["foo"], + signature_def_map={"foo_key": foo_signature}) + + def _validate_outputs_tensor_info_fail(self, builder, tensor_info): with self.test_session(graph=ops.Graph()) as sess: self._init_and_validate_variable(sess, "v", 42) @@ -120,6 +130,16 @@ class SavedModelTest(test.TestCase): sess, ["foo"], signature_def_map={"foo_key": foo_signature}) + def _validate_outputs_tensor_info_accept(self, builder, tensor_info): + with self.test_session(graph=ops.Graph()) as sess: + self._init_and_validate_variable(sess, "v", 42) + + foo_signature = signature_def_utils.build_signature_def( + dict(), {"foo_outputs": tensor_info}, "foo") + builder.add_meta_graph_and_variables( + sess, ["foo"], + signature_def_map={"foo_key": foo_signature}) + def testMaybeSavedModelDir(self): base_path = test.test_src_dir_path("/python/saved_model") self.assertFalse(loader.maybe_saved_model_directory(base_path)) @@ -539,23 +559,50 @@ class SavedModelTest(test.TestCase): self.assertEqual("bar", bar_signature["bar_key"].method_name) self.assertEqual("foo_new", bar_signature["foo_key"].method_name) - def testSignatureDefValidation(self): - export_dir = self._get_export_dir("test_signature_def_validation") + def testSignatureDefValidationFails(self): + export_dir = self._get_export_dir("test_signature_def_validation_fail") builder = saved_model_builder.SavedModelBuilder(export_dir) - tensor_without_name = meta_graph_pb2.TensorInfo() - tensor_without_name.dtype = types_pb2.DT_FLOAT - self._validate_inputs_tensor_info(builder, tensor_without_name) - self._validate_outputs_tensor_info(builder, tensor_without_name) + tensor_without_encoding = meta_graph_pb2.TensorInfo() + tensor_without_encoding.dtype = types_pb2.DT_FLOAT + self._validate_inputs_tensor_info_fail(builder, tensor_without_encoding) + self._validate_outputs_tensor_info_fail(builder, tensor_without_encoding) tensor_without_dtype = meta_graph_pb2.TensorInfo() tensor_without_dtype.name = "x" - self._validate_inputs_tensor_info(builder, tensor_without_dtype) - self._validate_outputs_tensor_info(builder, tensor_without_dtype) + self._validate_inputs_tensor_info_fail(builder, tensor_without_dtype) + self._validate_outputs_tensor_info_fail(builder, tensor_without_dtype) tensor_empty = meta_graph_pb2.TensorInfo() - self._validate_inputs_tensor_info(builder, tensor_empty) - self._validate_outputs_tensor_info(builder, tensor_empty) + self._validate_inputs_tensor_info_fail(builder, tensor_empty) + self._validate_outputs_tensor_info_fail(builder, tensor_empty) + + def testSignatureDefValidationSucceedsWithName(self): + tensor_with_name = meta_graph_pb2.TensorInfo() + tensor_with_name.name = "foo" + tensor_with_name.dtype = types_pb2.DT_FLOAT + + export_dir = self._get_export_dir("test_signature_def_validation_name_1") + builder = saved_model_builder.SavedModelBuilder(export_dir) + self._validate_inputs_tensor_info_accept(builder, tensor_with_name) + + export_dir = self._get_export_dir("test_signature_def_validation_name_2") + builder = saved_model_builder.SavedModelBuilder(export_dir) + self._validate_outputs_tensor_info_accept(builder, tensor_with_name) + + def testSignatureDefValidationSucceedsWithCoo(self): + tensor_with_coo = meta_graph_pb2.TensorInfo() + # TODO(soergel) test validation of each of the fields of coo_sparse + tensor_with_coo.coo_sparse.values_tensor_name = "foo" + tensor_with_coo.dtype = types_pb2.DT_FLOAT + + export_dir = self._get_export_dir("test_signature_def_validation_coo_1") + builder = saved_model_builder.SavedModelBuilder(export_dir) + self._validate_inputs_tensor_info_accept(builder, tensor_with_coo) + + export_dir = self._get_export_dir("test_signature_def_validation_coo_2") + builder = saved_model_builder.SavedModelBuilder(export_dir) + self._validate_outputs_tensor_info_accept(builder, tensor_with_coo) def testAssets(self): export_dir = self._get_export_dir("test_assets") @@ -874,8 +921,6 @@ class SavedModelTest(test.TestCase): 42, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval()) def testStripDefaultAttrs(self): - if ops._USE_C_API: return # TODO(skyewm): get this working - export_dir = self._get_export_dir("test_strip_default_attrs") builder = saved_model_builder.SavedModelBuilder(export_dir) @@ -945,61 +990,75 @@ class SavedModelTest(test.TestCase): self.assertIn("T", node_def.attr) self.assertIn("Tout", node_def.attr) - def testStripDefaultAttrsInconsistentConsumerDefaults(self): - if ops._USE_C_API: return # TODO(skyewm): get this working - + # Tests the behavior of loading SavedModels that having missing attrs or attrs + # with incorrect types. + def testInconsistentConsumerDefaultAttrs(self): export_dir = self._get_export_dir( "test_strip_default_attrs_no_consumer_defaults") builder = saved_model_builder.SavedModelBuilder(export_dir) - # Add a graph with two float32 variables and a Complex Op composing them - # with strip_default_attrs enabled. This must remove the following - # defaults for the "Complex" Op: - # o "T" : float32. (input type) - # o "Tout" : complex64. (output type) + # Add a graph with a single variable and a test op with a defaultless + # float32 attr, "test_attr". with session.Session(graph=ops.Graph()) as sess: - real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real") - imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag") - math_ops.complex(real_num, imag_num, name="complex") + variables.Variable(1.0, dtype=dtypes.float64, name="var") + test_ops.test_attr(T=dtypes.float32, name="test_attr") sess.run(variables.global_variables_initializer()) - builder.add_meta_graph_and_variables( - sess, ["foo"], strip_default_attrs=True) + builder.add_meta_graph_and_variables(sess, ["foo"]) # Save the SavedModel to disk in text format. builder.save(as_text=True) - # Update the Op registry to remove defaults for all attrs("T", "Tout") from - # the "Complex" OpDef. - complex_op_def = op_def_registry.get_registered_ops()["Complex"] - original_complex_op_def = op_def_pb2.OpDef() - original_complex_op_def.CopyFrom(complex_op_def) - for attr_def in complex_op_def.attr: - attr_def.ClearField("default_value") + # Rewrite the SavedModel to remove the T attr from "test_attr". + saved_model_file = os.path.join( + export_dir, constants.SAVED_MODEL_FILENAME_PBTXT) + with open(saved_model_file) as f: + original_saved_model = f.read() + + no_attr_saved_model = original_saved_model.replace(""" + attr { + key: "T" + value { + type: DT_FLOAT + } + }""", "") + with open(saved_model_file, "w") as f: + f.write(no_attr_saved_model) # Loading the SavedModel via the loader must fail because the SavedModel - # does not have any attr values for the "Complex" node and the current - # op registry does not have have any default values for the "Complex" op. + # does not have any attr values for the "TestAttr" node, and there is no + # default specified in the TestAttr OpDef. sess = session.Session(graph=ops.Graph()) - with self.assertRaisesRegexp( - ValueError, - "Expected one attr with name .*T(out)?.* in name: \"complex\".*"): + if ops._USE_C_API: + error_message = "NodeDef missing attr 'T' from Op complex128). - complex_op_def.CopyFrom(original_complex_op_def) - for attr_def in complex_op_def.attr: - if attr_def.name == "Tout": - attr_def.default_value.type = types_pb2.DT_COMPLEX128 - - # Loading the SavedModel via the loader must set "Tout" attr_value for the - # "Complex" node according to the latest defaults (complex128). This is - # expected to fail the model import as there is no OpKernel registered to - # handle attrs "T" (float32) and "Tout" (complex128). + # Rewrite the SavedModel to change the type of the T attr in "test_attr" + bad_type_saved_model = original_saved_model.replace(""" + attr { + key: "T" + value { + type: DT_FLOAT + } + }""", """ + attr { + key: "T" + value { + type: DT_DOUBLE + } + }""") + with open(saved_model_file, "w") as f: + f.write(bad_type_saved_model) + + # Loading the SavedModel via the loader must fail because there is no + # OpKernel registered to handle T = double. sess = session.Session(graph=ops.Graph()) with self.assertRaisesRegexp( errors.InvalidArgumentError, - ".*No OpKernel was registered to support Op \'Complex\' with these " + ".*No OpKernel was registered to support Op \'TestAttr\' with these " "attrs..*"): loader.load(sess, ["foo"], export_dir) diff --git a/tensorflow/python/saved_model/signature_constants.py b/tensorflow/python/saved_model/signature_constants.py index 935a124645bde509a1b5a7751a285a85acbe8cab..6461fe8a7e7bef1a2fc787879da9e3324e2655c8 100644 --- a/tensorflow/python/saved_model/signature_constants.py +++ b/tensorflow/python/saved_model/signature_constants.py @@ -20,51 +20,79 @@ from __future__ import division from __future__ import print_function from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export # Key in the signature def map for `default` serving signatures. The default # signature is used in inference requests where a specific signature was not # specified. DEFAULT_SERVING_SIGNATURE_DEF_KEY = "serving_default" +tf_export("saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY" + ).export_constant(__name__, "DEFAULT_SERVING_SIGNATURE_DEF_KEY") ################################################################################ # Classification API constants. # Classification inputs. CLASSIFY_INPUTS = "inputs" +tf_export("saved_model.signature_constants.CLASSIFY_INPUTS").export_constant( + __name__, "CLASSIFY_INPUTS") # Classification method name used in a SignatureDef. CLASSIFY_METHOD_NAME = "tensorflow/serving/classify" +tf_export( + "saved_model.signature_constants.CLASSIFY_METHOD_NAME").export_constant( + __name__, "CLASSIFY_METHOD_NAME") # Classification classes output. CLASSIFY_OUTPUT_CLASSES = "classes" +tf_export( + "saved_model.signature_constants.CLASSIFY_OUTPUT_CLASSES").export_constant( + __name__, "CLASSIFY_OUTPUT_CLASSES") # Classification scores output. CLASSIFY_OUTPUT_SCORES = "scores" +tf_export( + "saved_model.signature_constants.CLASSIFY_OUTPUT_SCORES").export_constant( + __name__, "CLASSIFY_OUTPUT_SCORES") ################################################################################ # Prediction API constants. # Predict inputs. PREDICT_INPUTS = "inputs" +tf_export("saved_model.signature_constants.PREDICT_INPUTS").export_constant( + __name__, "PREDICT_INPUTS") # Prediction method name used in a SignatureDef. PREDICT_METHOD_NAME = "tensorflow/serving/predict" +tf_export( + "saved_model.signature_constants.PREDICT_METHOD_NAME").export_constant( + __name__, "PREDICT_METHOD_NAME") # Predict outputs. PREDICT_OUTPUTS = "outputs" +tf_export("saved_model.signature_constants.PREDICT_OUTPUTS").export_constant( + __name__, "PREDICT_OUTPUTS") ################################################################################ # Regression API constants. # Regression inputs. REGRESS_INPUTS = "inputs" +tf_export("saved_model.signature_constants.REGRESS_INPUTS").export_constant( + __name__, "REGRESS_INPUTS") # Regression method name used in a SignatureDef. REGRESS_METHOD_NAME = "tensorflow/serving/regress" +tf_export( + "saved_model.signature_constants.REGRESS_METHOD_NAME").export_constant( + __name__, "REGRESS_METHOD_NAME") # Regression outputs. REGRESS_OUTPUTS = "outputs" +tf_export("saved_model.signature_constants.REGRESS_OUTPUTS").export_constant( + __name__, "REGRESS_OUTPUTS") ################################################################################ diff --git a/tensorflow/python/saved_model/signature_def_utils_impl.py b/tensorflow/python/saved_model/signature_def_utils_impl.py index 240ea61aa5f8553852044f84b61d010bfbca69d1..d0331591889110df86bdb2ac69c037bc3b968f91 100644 --- a/tensorflow/python/saved_model/signature_def_utils_impl.py +++ b/tensorflow/python/saved_model/signature_def_utils_impl.py @@ -26,8 +26,10 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model import utils +from tensorflow.python.util.tf_export import tf_export +@tf_export('saved_model.signature_def_utils.build_signature_def') def build_signature_def(inputs=None, outputs=None, method_name=None): """Utility function to build a SignatureDef protocol buffer. @@ -53,6 +55,7 @@ def build_signature_def(inputs=None, outputs=None, method_name=None): return signature_def +@tf_export('saved_model.signature_def_utils.regression_signature_def') def regression_signature_def(examples, predictions): """Creates regression signature from given examples and predictions. @@ -94,6 +97,7 @@ def regression_signature_def(examples, predictions): return signature_def +@tf_export('saved_model.signature_def_utils.classification_signature_def') def classification_signature_def(examples, classes, scores): """Creates classification signature from given examples and predictions. @@ -146,6 +150,7 @@ def classification_signature_def(examples, classes, scores): return signature_def +@tf_export('saved_model.signature_def_utils.predict_signature_def') def predict_signature_def(inputs, outputs): """Creates prediction signature from given inputs and outputs. @@ -180,6 +185,7 @@ def predict_signature_def(inputs, outputs): return signature_def +@tf_export('saved_model.signature_def_utils.is_valid_signature') def is_valid_signature(signature_def): """Determine whether a SignatureDef can be served by TensorFlow Serving.""" if signature_def is None: diff --git a/tensorflow/python/saved_model/simple_save.py b/tensorflow/python/saved_model/simple_save.py index 1e4cc733706c774107184c2fb19037da3d4e733d..042b8fa8e22703d8ffb5e12de3f844d22fb1b1ce 100644 --- a/tensorflow/python/saved_model/simple_save.py +++ b/tensorflow/python/saved_model/simple_save.py @@ -23,8 +23,10 @@ from tensorflow.python.saved_model import builder from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model import signature_def_utils from tensorflow.python.saved_model import tag_constants +from tensorflow.python.util.tf_export import tf_export +@tf_export('saved_model.simple_save') def simple_save(session, export_dir, inputs, outputs, legacy_init_op=None): """Convenience function to build a SavedModel suitable for serving. diff --git a/tensorflow/python/saved_model/tag_constants.py b/tensorflow/python/saved_model/tag_constants.py index e2facafda51919d3f1e0ccbe646db522ed0bc49b..d164e2c23f24469d7536f87cb431afe618ddcc06 100644 --- a/tensorflow/python/saved_model/tag_constants.py +++ b/tensorflow/python/saved_model/tag_constants.py @@ -20,19 +20,26 @@ from __future__ import division from __future__ import print_function from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export # Tag for the `serving` graph. SERVING = "serve" +tf_export("saved_model.tag_constants.SERVING").export_constant( + __name__, "SERVING") # Tag for the `training` graph. TRAINING = "train" +tf_export("saved_model.tag_constants.TRAINING").export_constant( + __name__, "TRAINING") # Tag for the `gpu` graph. GPU = "gpu" +tf_export("saved_model.tag_constants.GPU").export_constant(__name__, "GPU") # Tag for the `tpu` graph. TPU = "tpu" +tf_export("saved_model.tag_constants.TPU").export_constant(__name__, "TPU") _allowed_symbols = [ "SERVING", diff --git a/tensorflow/python/saved_model/utils_impl.py b/tensorflow/python/saved_model/utils_impl.py index 73ca8c9c1c6d8fddc8a9c7dbee56682999281c28..cddce29a08a6c4c79a4c7c5dbfb48a86131530b2 100644 --- a/tensorflow/python/saved_model/utils_impl.py +++ b/tensorflow/python/saved_model/utils_impl.py @@ -22,11 +22,13 @@ from tensorflow.core.protobuf import meta_graph_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor +from tensorflow.python.util.tf_export import tf_export # TensorInfo helpers. +@tf_export("saved_model.utils.build_tensor_info") def build_tensor_info(tensor): """Utility function to build TensorInfo proto. @@ -50,6 +52,7 @@ def build_tensor_info(tensor): return tensor_info +@tf_export("saved_model.utils.get_tensor_from_tensor_info") def get_tensor_from_tensor_info(tensor_info, graph=None, import_scope=None): """Returns the Tensor or SparseTensor described by a TensorInfo proto. diff --git a/tensorflow/python/summary/summary.py b/tensorflow/python/summary/summary.py index 92c1fcadd29c7858da1d31375c209bf1b21f3103..f1b2be0a1a38f5c3f07d1d33e6e4754398687d33 100644 --- a/tensorflow/python/summary/summary.py +++ b/tensorflow/python/summary/summary.py @@ -72,8 +72,10 @@ from tensorflow.python.summary.writer.writer_cache import FileWriterCache from tensorflow.python.util import compat as _compat from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export +@tf_export('summary.scalar') def scalar(name, tensor, collections=None, family=None): """Outputs a `Summary` protocol buffer containing a single scalar value. @@ -96,12 +98,12 @@ def scalar(name, tensor, collections=None, family=None): """ with _summary_op_util.summary_scope( name, family, values=[tensor]) as (tag, scope): - # pylint: disable=protected-access - val = _gen_logging_ops._scalar_summary(tags=tag, values=tensor, name=scope) + val = _gen_logging_ops.scalar_summary(tags=tag, values=tensor, name=scope) _summary_op_util.collect(val, collections, [_ops.GraphKeys.SUMMARIES]) return val +@tf_export('summary.image') def image(name, tensor, max_outputs=3, collections=None, family=None): """Outputs a `Summary` protocol buffer with images. @@ -149,13 +151,13 @@ def image(name, tensor, max_outputs=3, collections=None, family=None): """ with _summary_op_util.summary_scope( name, family, values=[tensor]) as (tag, scope): - # pylint: disable=protected-access - val = _gen_logging_ops._image_summary( + val = _gen_logging_ops.image_summary( tag=tag, tensor=tensor, max_images=max_outputs, name=scope) _summary_op_util.collect(val, collections, [_ops.GraphKeys.SUMMARIES]) return val +@tf_export('summary.histogram') def histogram(name, values, collections=None, family=None): # pylint: disable=line-too-long """Outputs a `Summary` protocol buffer with a histogram. @@ -188,13 +190,13 @@ def histogram(name, values, collections=None, family=None): with _summary_op_util.summary_scope( name, family, values=[values], default_name='HistogramSummary') as (tag, scope): - # pylint: disable=protected-access - val = _gen_logging_ops._histogram_summary( + val = _gen_logging_ops.histogram_summary( tag=tag, values=values, name=scope) _summary_op_util.collect(val, collections, [_ops.GraphKeys.SUMMARIES]) return val +@tf_export('summary.audio') def audio(name, tensor, sample_rate, max_outputs=3, collections=None, family=None): # pylint: disable=line-too-long @@ -232,16 +234,16 @@ def audio(name, tensor, sample_rate, max_outputs=3, collections=None, """ with _summary_op_util.summary_scope( name, family=family, values=[tensor]) as (tag, scope): - # pylint: disable=protected-access sample_rate = _ops.convert_to_tensor( sample_rate, dtype=_dtypes.float32, name='sample_rate') - val = _gen_logging_ops._audio_summary_v2( + val = _gen_logging_ops.audio_summary_v2( tag=tag, tensor=tensor, max_outputs=max_outputs, sample_rate=sample_rate, name=scope) _summary_op_util.collect(val, collections, [_ops.GraphKeys.SUMMARIES]) return val +@tf_export('summary.merge') def merge(inputs, collections=None, name=None): # pylint: disable=line-too-long """Merges summaries. @@ -274,18 +276,18 @@ def merge(inputs, collections=None, name=None): @end_compatbility """ # pylint: enable=line-too-long - if _context.in_eager_mode(): + if _context.executing_eagerly(): raise RuntimeError( 'Merging tf.summary.* ops is not compatible with eager execution. ' 'Use tf.contrib.summary instead.') name = _summary_op_util.clean_tag(name) with _ops.name_scope(name, 'Merge', inputs): - # pylint: disable=protected-access - val = _gen_logging_ops._merge_summary(inputs=inputs, name=name) + val = _gen_logging_ops.merge_summary(inputs=inputs, name=name) _summary_op_util.collect(val, collections, []) return val +@tf_export('summary.merge_all') def merge_all(key=_ops.GraphKeys.SUMMARIES, scope=None): """Merges all summaries collected in the default graph. @@ -307,7 +309,7 @@ def merge_all(key=_ops.GraphKeys.SUMMARIES, scope=None): summaries under eager execution, use `tf.contrib.summary` instead. @end_compatbility """ - if _context.in_eager_mode(): + if _context.executing_eagerly(): raise RuntimeError( 'Merging tf.summary.* ops is not compatible with eager execution. ' 'Use tf.contrib.summary instead.') @@ -318,6 +320,7 @@ def merge_all(key=_ops.GraphKeys.SUMMARIES, scope=None): return merge(summary_ops) +@tf_export('summary.get_summary_description') def get_summary_description(node_def): """Given a TensorSummary node_def, retrieve its SummaryDescription. diff --git a/tensorflow/python/summary/summary_iterator.py b/tensorflow/python/summary/summary_iterator.py index 6969c4cf1500bf4b1fda900336158e5af4395ea6..321b11ffb73487405428340df94010ed8ddbfcd4 100644 --- a/tensorflow/python/summary/summary_iterator.py +++ b/tensorflow/python/summary/summary_iterator.py @@ -21,8 +21,10 @@ from __future__ import print_function from tensorflow.core.util import event_pb2 from tensorflow.python.lib.io import tf_record +from tensorflow.python.util.tf_export import tf_export +@tf_export('train.summary_iterator') def summary_iterator(path): # pylint: disable=line-too-long """An iterator for reading `Event` protocol buffers from an event file. diff --git a/tensorflow/python/summary/text_summary.py b/tensorflow/python/summary/text_summary.py index 94a85d73e2f77388f9a29b1c135fc6046a8362d0..6418c847f3c819cf2491bb449921d15c39eae288 100644 --- a/tensorflow/python/summary/text_summary.py +++ b/tensorflow/python/summary/text_summary.py @@ -26,10 +26,12 @@ from __future__ import print_function from tensorflow.core.framework import summary_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.ops.summary_ops import tensor_summary +from tensorflow.python.util.tf_export import tf_export PLUGIN_NAME = "text" +@tf_export("summary.text") def text_summary(name, tensor, collections=None): """Summarizes textual data. diff --git a/tensorflow/python/summary/writer/writer.py b/tensorflow/python/summary/writer/writer.py index 12f120116f4439059f42c7212469ee835cc13ef4..57f78c156b1334a5486b29f2ddec957e49156e73 100644 --- a/tensorflow/python/summary/writer/writer.py +++ b/tensorflow/python/summary/writer/writer.py @@ -32,6 +32,7 @@ from tensorflow.python.platform import gfile from tensorflow.python.platform import tf_logging as logging from tensorflow.python.summary import plugin_asset from tensorflow.python.summary.writer.event_file_writer import EventFileWriter +from tensorflow.python.util.tf_export import tf_export _PLUGINS_DIR = "plugins" @@ -276,6 +277,7 @@ class SummaryToEventTransformer(object): self.event_writer.add_event(event) +@tf_export("summary.FileWriter") class FileWriter(SummaryToEventTransformer): """Writes `Summary` protocol buffers to event files. @@ -341,7 +343,7 @@ class FileWriter(SummaryToEventTransformer): summaries under eager execution, use `tf.contrib.summary` instead. @end_compatbility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "tf.summary.FileWriter is not compatible with eager execution. " "Use tf.contrib.summary instead.") diff --git a/tensorflow/python/summary/writer/writer_cache.py b/tensorflow/python/summary/writer/writer_cache.py index bad289303c0fd0de7836b03a6762d04505521a89..645fa28a37fb125b6b1224961251bc8879d5fe6d 100644 --- a/tensorflow/python/summary/writer/writer_cache.py +++ b/tensorflow/python/summary/writer/writer_cache.py @@ -22,8 +22,10 @@ import threading from tensorflow.python.framework import ops from tensorflow.python.summary.writer.writer import FileWriter +from tensorflow.python.util.tf_export import tf_export +@tf_export('summary.FileWriterCache') class FileWriterCache(object): """Cache for file writers. diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD index 63f16c53a29fd65c32077dd29e3b1823c11d457b..1de1adcfbc35e2b760f362cb9784dd415b9a4dc4 100644 --- a/tensorflow/python/tools/BUILD +++ b/tensorflow/python/tools/BUILD @@ -14,6 +14,7 @@ py_library( name = "tools_pip", deps = [ ":freeze_graph", + ":import_pb_to_tensorboard", ":inspect_checkpoint", ":optimize_for_inference", ":print_selective_registration_header", @@ -248,7 +249,10 @@ py_test( "//tensorflow/cc/saved_model:saved_model_half_plus_two", ], srcs_version = "PY2AND3", - tags = ["manual"], + tags = [ + "manual", + "no-internal-py3", + ], deps = [ ":saved_model_cli", "//tensorflow/core:protos_all_py", diff --git a/tensorflow/python/tools/freeze_graph.py b/tensorflow/python/tools/freeze_graph.py index fd78f44c999e41f404a6499a0addd2131da6d287..e9f1def48c462dcd8a5acf0e3d29d562cd1b3d58 100644 --- a/tensorflow/python/tools/freeze_graph.py +++ b/tensorflow/python/tools/freeze_graph.py @@ -56,8 +56,6 @@ from tensorflow.python.saved_model import tag_constants from tensorflow.python.tools import saved_model_utils from tensorflow.python.training import saver as saver_lib -FLAGS = None - def freeze_graph_with_def_protos(input_graph_def, input_saver_def, @@ -101,15 +99,15 @@ def freeze_graph_with_def_protos(input_graph_def, _ = importer.import_graph_def(input_graph_def, name="") with session.Session() as sess: if input_saver_def: - saver = saver_lib.Saver(saver_def=input_saver_def, - write_version=checkpoint_version) + saver = saver_lib.Saver( + saver_def=input_saver_def, write_version=checkpoint_version) saver.restore(sess, input_checkpoint) elif input_meta_graph_def: restorer = saver_lib.import_meta_graph( input_meta_graph_def, clear_devices=True) restorer.restore(sess, input_checkpoint) if initializer_nodes: - sess.run(initializer_nodes.split(",")) + sess.run(initializer_nodes.replace(" ", "").split(",")) elif input_saved_model_dir: if saved_model_tags is None: saved_model_tags = [] @@ -126,29 +124,31 @@ def freeze_graph_with_def_protos(input_graph_def, # 'global_step' or a similar housekeeping element) so skip it. continue var_list[key] = tensor - saver = saver_lib.Saver(var_list=var_list, - write_version=checkpoint_version) + saver = saver_lib.Saver( + var_list=var_list, write_version=checkpoint_version) saver.restore(sess, input_checkpoint) if initializer_nodes: - sess.run(initializer_nodes.split(",")) + sess.run(initializer_nodes.replace(" ", "").split(",")) - variable_names_whitelist = (variable_names_whitelist.split(",") - if variable_names_whitelist else None) - variable_names_blacklist = (variable_names_blacklist.split(",") - if variable_names_blacklist else None) + variable_names_whitelist = ( + variable_names_whitelist.replace(" ", "").split(",") + if variable_names_whitelist else None) + variable_names_blacklist = ( + variable_names_blacklist.replace(" ", "").split(",") + if variable_names_blacklist else None) if input_meta_graph_def: output_graph_def = graph_util.convert_variables_to_constants( sess, input_meta_graph_def.graph_def, - output_node_names.split(","), + output_node_names.replace(" ", "").split(","), variable_names_whitelist=variable_names_whitelist, variable_names_blacklist=variable_names_blacklist) else: output_graph_def = graph_util.convert_variables_to_constants( sess, input_graph_def, - output_node_names.split(","), + output_node_names.replace(" ", "").split(","), variable_names_whitelist=variable_names_whitelist, variable_names_blacklist=variable_names_blacklist) @@ -237,24 +237,41 @@ def freeze_graph(input_graph, if input_saver: input_saver_def = _parse_input_saver_proto(input_saver, input_binary) freeze_graph_with_def_protos( - input_graph_def, input_saver_def, input_checkpoint, output_node_names, - restore_op_name, filename_tensor_name, output_graph, clear_devices, - initializer_nodes, variable_names_whitelist, variable_names_blacklist, - input_meta_graph_def, input_saved_model_dir, - saved_model_tags.split(","), checkpoint_version=checkpoint_version) - + input_graph_def, + input_saver_def, + input_checkpoint, + output_node_names, + restore_op_name, + filename_tensor_name, + output_graph, + clear_devices, + initializer_nodes, + variable_names_whitelist, + variable_names_blacklist, + input_meta_graph_def, + input_saved_model_dir, + saved_model_tags.replace(" ", "").split(","), + checkpoint_version=checkpoint_version) -def main(unused_args): - freeze_graph(FLAGS.input_graph, FLAGS.input_saver, FLAGS.input_binary, - FLAGS.input_checkpoint, FLAGS.output_node_names, - FLAGS.restore_op_name, FLAGS.filename_tensor_name, - FLAGS.output_graph, FLAGS.clear_devices, FLAGS.initializer_nodes, - FLAGS.variable_names_whitelist, FLAGS.variable_names_blacklist, - FLAGS.input_meta_graph, FLAGS.input_saved_model_dir, - FLAGS.saved_model_tags, FLAGS.checkpoint_version) +def main(unused_args, flags): + if flags.checkpoint_version == 1: + checkpoint_version = saver_pb2.SaverDef.V1 + elif flags.checkpoint_version == 2: + checkpoint_version = saver_pb2.SaverDef.V2 + else: + print("Invalid checkpoint version (must be '1' or '2'): %d" % + flags.checkpoint_version) + return -1 + freeze_graph(flags.input_graph, flags.input_saver, flags.input_binary, + flags.input_checkpoint, flags.output_node_names, + flags.restore_op_name, flags.filename_tensor_name, + flags.output_graph, flags.clear_devices, flags.initializer_nodes, + flags.variable_names_whitelist, flags.variable_names_blacklist, + flags.input_meta_graph, flags.input_saved_model_dir, + flags.saved_model_tags, checkpoint_version) -if __name__ == "__main__": +def run_main(): parser = argparse.ArgumentParser() parser.register("type", "bool", lambda v: v.lower() == "true") parser.add_argument( @@ -275,7 +292,7 @@ if __name__ == "__main__": parser.add_argument( "--checkpoint_version", type=int, - default=saver_pb2.SaverDef.V2, + default=2, help="Tensorflow variable file format") parser.add_argument( "--output_graph", @@ -356,5 +373,10 @@ if __name__ == "__main__": separated by \',\'. For tag-set contains multiple tags, all tags \ must be passed in.\ """) - FLAGS, unparsed = parser.parse_known_args() - app.run(main=main, argv=[sys.argv[0]] + unparsed) + flags, unparsed = parser.parse_known_args() + + my_main = lambda unused_args: main(unused_args, flags) + app.run(main=my_main, argv=[sys.argv[0]] + unparsed) + +if __name__ == '__main__': + run_main() diff --git a/tensorflow/python/tools/freeze_graph_test.py b/tensorflow/python/tools/freeze_graph_test.py index 342732465d48f40a4ffeac97146fb1b6d564c568..91f0061ebccaebbdbb09f283d9d52d813459f493 100644 --- a/tensorflow/python/tools/freeze_graph_test.py +++ b/tensorflow/python/tools/freeze_graph_test.py @@ -84,9 +84,18 @@ class FreezeGraphTest(test_util.TensorFlowTestCase): input_meta_graph = checkpoint_meta_graph_file freeze_graph.freeze_graph( - input_graph_path, input_saver_def_path, input_binary, checkpoint_path, - output_node_names, restore_op_name, filename_tensor_name, - output_graph_path, clear_devices, "", "", input_meta_graph, + input_graph_path, + input_saver_def_path, + input_binary, + checkpoint_path, + output_node_names, + restore_op_name, + filename_tensor_name, + output_graph_path, + clear_devices, + "", + "", + input_meta_graph, checkpoint_version=saver_write_version) # Now we make sure the variable is now a constant, and that the graph still diff --git a/tensorflow/python/tools/inspect_checkpoint.py b/tensorflow/python/tools/inspect_checkpoint.py index dd876cbe7fcd64a8de70eb28f67996df9de1dd7d..6504fbc10755c5c543016b8d56d6d53f3311b249 100644 --- a/tensorflow/python/tools/inspect_checkpoint.py +++ b/tensorflow/python/tools/inspect_checkpoint.py @@ -30,7 +30,7 @@ FLAGS = None def print_tensors_in_checkpoint_file(file_name, tensor_name, all_tensors, - all_tensor_names): + all_tensor_names=False): """Prints tensors in a checkpoint file. If no `tensor_name` is provided, prints the tensor names and shapes @@ -139,7 +139,7 @@ if __name__ == "__main__": const=True, type="bool", default=False, - help="If True, print the values of all the tensors.") + help="If True, print the names and values of all the tensors.") parser.add_argument( "--all_tensor_names", nargs="?", diff --git a/tensorflow/python/tools/optimize_for_inference_lib.py b/tensorflow/python/tools/optimize_for_inference_lib.py index c2687bf557b03ff588fd369771077c92ba012a15..9c1927122252f45ddfa8092045c7589fa0f45532 100644 --- a/tensorflow/python/tools/optimize_for_inference_lib.py +++ b/tensorflow/python/tools/optimize_for_inference_lib.py @@ -349,6 +349,7 @@ def fold_batch_norms(input_graph_def): bias_add_op.op = "BiasAdd" bias_add_op.name = node.name bias_add_op.attr["T"].CopyFrom(conv_op.attr["T"]) + bias_add_op.attr["data_format"].CopyFrom(conv_op.attr["data_format"]) bias_add_op.input.extend([new_conv_op.name, offset_op.name]) new_ops.extend([scaled_weights_op, new_conv_op, offset_op, bias_add_op]) diff --git a/tensorflow/python/tools/optimize_for_inference_test.py b/tensorflow/python/tools/optimize_for_inference_test.py index 7686bb0f1462253f45e574eaeaf8aeb555f26106..084a4500f8e1eb7f75e1e01668fae655b5e06763 100644 --- a/tensorflow/python/tools/optimize_for_inference_test.py +++ b/tensorflow/python/tools/optimize_for_inference_test.py @@ -173,48 +173,56 @@ class OptimizeForInferenceTest(test.TestCase): self.assertNotEqual("BatchNormWithGlobalNormalization", node.op) def testFoldFusedBatchNorms(self): - with self.test_session() as sess: - inputs = [1, 4, 2, 5, 3, 6, -1, -4, -2, -5, -3, -6] - input_op = constant_op.constant( - np.array(inputs), shape=[1, 1, 6, 2], dtype=dtypes.float32) - weights = [1, 2, 3, 4, 0.1, 0.2, 0.3, 0.4] - weights_op = constant_op.constant( - np.array(weights), shape=[1, 2, 2, 2], dtype=dtypes.float32) - conv_op = nn_ops.conv2d( - input_op, weights_op, [1, 1, 1, 1], padding="SAME", name="conv_op") - mean_op = constant_op.constant( - np.array([10, 20]), shape=[2], dtype=dtypes.float32) - variance_op = constant_op.constant( - np.array([0.25, 0.5]), shape=[2], dtype=dtypes.float32) - beta_op = constant_op.constant( - np.array([0.1, 0.6]), shape=[2], dtype=dtypes.float32) - gamma_op = constant_op.constant( - np.array([1.0, 2.0]), shape=[2], dtype=dtypes.float32) - ops.get_default_graph().graph_def_versions.producer = 9 - gen_nn_ops._fused_batch_norm( - conv_op, - gamma_op, - beta_op, - mean_op, - variance_op, - 0.00001, - is_training=False, - name="output") - original_graph_def = sess.graph_def - original_result = sess.run(["output:0"]) - optimized_graph_def = optimize_for_inference_lib.fold_batch_norms( - original_graph_def) - - with self.test_session() as sess: - _ = importer.import_graph_def( - optimized_graph_def, input_map={}, name="optimized") - optimized_result = sess.run(["optimized/output:0"]) - - self.assertAllClose( - original_result, optimized_result, rtol=1e-04, atol=1e-06) - - for node in optimized_graph_def.node: - self.assertNotEqual("FusedBatchNorm", node.op) + for data_format, use_gpu in [("NHWC", False), ("NCHW", True)]: + with self.test_session(use_gpu=use_gpu) as sess: + inputs = [1, 4, 2, 5, 3, 6, -1, -4, -2, -5, -3, -6] + input_op = constant_op.constant( + np.array(inputs), + shape=[1, 1, 6, 2] if data_format == "NHWC" else [1, 2, 1, 6], + dtype=dtypes.float32) + weights = [1, 2, 3, 4, 0.1, 0.2, 0.3, 0.4] + weights_op = constant_op.constant( + np.array(weights), shape=[1, 2, 2, 2], dtype=dtypes.float32) + conv_op = nn_ops.conv2d( + input_op, + weights_op, [1, 1, 1, 1], + padding="SAME", + data_format=data_format, + name="conv_op") + mean_op = constant_op.constant( + np.array([10, 20]), shape=[2], dtype=dtypes.float32) + variance_op = constant_op.constant( + np.array([0.25, 0.5]), shape=[2], dtype=dtypes.float32) + beta_op = constant_op.constant( + np.array([0.1, 0.6]), shape=[2], dtype=dtypes.float32) + gamma_op = constant_op.constant( + np.array([1.0, 2.0]), shape=[2], dtype=dtypes.float32) + ops.get_default_graph().graph_def_versions.producer = 9 + gen_nn_ops._fused_batch_norm( + conv_op, + gamma_op, + beta_op, + mean_op, + variance_op, + 0.00001, + is_training=False, + data_format=data_format, + name="output") + original_graph_def = sess.graph_def + original_result = sess.run(["output:0"]) + optimized_graph_def = optimize_for_inference_lib.fold_batch_norms( + original_graph_def) + + with self.test_session(use_gpu=use_gpu) as sess: + _ = importer.import_graph_def( + optimized_graph_def, input_map={}, name="optimized") + optimized_result = sess.run(["optimized/output:0"]) + + self.assertAllClose( + original_result, optimized_result, rtol=1e-04, atol=1e-06) + + for node in optimized_graph_def.node: + self.assertNotEqual("FusedBatchNorm", node.op) def testFuseResizePadAndConv(self): with self.test_session() as sess: diff --git a/tensorflow/python/tools/print_selective_registration_header_test.py b/tensorflow/python/tools/print_selective_registration_header_test.py index 36978b0860a423569149cd0572629f9f1f280637..4b3d98242caf683693430f08bd8cb74483f4bc74 100644 --- a/tensorflow/python/tools/print_selective_registration_header_test.py +++ b/tensorflow/python/tools/print_selective_registration_header_test.py @@ -24,6 +24,7 @@ import sys from google.protobuf import text_format from tensorflow.core.framework import graph_pb2 +from tensorflow.python.framework import test_util from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.tools import selective_registration_header_lib @@ -93,11 +94,16 @@ class PrintOpFilegroupTest(test.TestCase): ops_and_kernels = selective_registration_header_lib.get_ops_and_kernels( 'rawproto', self.WriteGraphFiles(graphs), default_ops) + matmul_prefix = '' + if test_util.IsMklEnabled(): + matmul_prefix = 'Mkl' + self.assertListEqual( [ ('BiasAdd', 'BiasOp'), # - ('MatMul', 'MatMulOp'), # - ('MatMul', 'MatMulOp'), # + ('MatMul', + matmul_prefix + 'MatMulOp'), # + ('MatMul', matmul_prefix + 'MatMulOp'), # ('NoOp', 'NoOp'), # ('Reshape', 'ReshapeOp'), # ('_Recv', 'RecvOp'), # @@ -112,8 +118,9 @@ class PrintOpFilegroupTest(test.TestCase): self.assertListEqual( [ ('BiasAdd', 'BiasOp'), # - ('MatMul', 'MatMulOp'), # - ('MatMul', 'MatMulOp'), # + ('MatMul', + matmul_prefix + 'MatMulOp'), # + ('MatMul', matmul_prefix + 'MatMulOp'), # ('NoOp', 'NoOp'), # ('Reshape', 'ReshapeOp'), # ('_Recv', 'RecvOp'), # diff --git a/tensorflow/python/tools/saved_model_cli.py b/tensorflow/python/tools/saved_model_cli.py index 5b0a584c10ea33f345c09324cbd47eb1789466e0..b88be4ae04d5dc7a7641fb8dbd7e56e61035869f 100644 --- a/tensorflow/python/tools/saved_model_cli.py +++ b/tensorflow/python/tools/saved_model_cli.py @@ -38,11 +38,15 @@ from tensorflow.core.example import example_pb2 from tensorflow.core.framework import types_pb2 from tensorflow.python.client import session from tensorflow.python.debug.wrappers import local_cli_wrapper +from tensorflow.python.framework import meta_graph as meta_graph_lib from tensorflow.python.framework import ops as ops_lib -from tensorflow.python.platform import app +from tensorflow.python.platform import app # pylint: disable=unused-import from tensorflow.python.saved_model import loader from tensorflow.python.tools import saved_model_utils +# Set of ops to blacklist. +_OP_BLACKLIST = set(['WriteFile', 'ReadFile']) + def _show_tag_sets(saved_model_dir): """Prints the tag-sets stored in SavedModel directory. @@ -115,7 +119,7 @@ def _get_outputs_tensor_info_from_meta_graph_def(meta_graph_def, signature_def_key).outputs -def _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key): +def _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key, indent=0): """Prints input and output TensorInfos. Prints the details of input and output TensorInfos for the SignatureDef mapped @@ -126,6 +130,7 @@ def _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key): tag_set: Group of tag(s) of the MetaGraphDef, in string format, separated by ','. For tag-set contains multiple tags, all tags must be passed in. signature_def_key: A SignatureDef key string. + indent: How far (in increments of 2 spaces) to indent each line of output. """ meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir, tag_set) @@ -134,29 +139,39 @@ def _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key): outputs_tensor_info = _get_outputs_tensor_info_from_meta_graph_def( meta_graph_def, signature_def_key) - print('The given SavedModel SignatureDef contains the following input(s):') + indent_str = " " * indent + def in_print(s): + print(indent_str + s) + + in_print('The given SavedModel SignatureDef contains the following input(s):') for input_key, input_tensor in sorted(inputs_tensor_info.items()): - print('inputs[\'%s\'] tensor_info:' % input_key) - _print_tensor_info(input_tensor) + in_print(' inputs[\'%s\'] tensor_info:' % input_key) + _print_tensor_info(input_tensor, indent+1) - print('The given SavedModel SignatureDef contains the following output(s):') + in_print('The given SavedModel SignatureDef contains the following ' + 'output(s):') for output_key, output_tensor in sorted(outputs_tensor_info.items()): - print('outputs[\'%s\'] tensor_info:' % output_key) - _print_tensor_info(output_tensor) + in_print(' outputs[\'%s\'] tensor_info:' % output_key) + _print_tensor_info(output_tensor, indent+1) - print('Method name is: %s' % - meta_graph_def.signature_def[signature_def_key].method_name) + in_print('Method name is: %s' % + meta_graph_def.signature_def[signature_def_key].method_name) -def _print_tensor_info(tensor_info): +def _print_tensor_info(tensor_info, indent=0): """Prints details of the given tensor_info. Args: tensor_info: TensorInfo object to be printed. + indent: How far (in increments of 2 spaces) to indent each line output """ - print(' dtype: ' + - {value: key - for (key, value) in types_pb2.DataType.items()}[tensor_info.dtype]) + indent_str = " " * indent + def in_print(s): + print(indent_str + s) + + in_print(' dtype: ' + + {value: key + for (key, value) in types_pb2.DataType.items()}[tensor_info.dtype]) # Display shape as tuple. if tensor_info.tensor_shape.unknown_rank: shape = 'unknown_rank' @@ -164,8 +179,8 @@ def _print_tensor_info(tensor_info): dims = [str(dim.size) for dim in tensor_info.tensor_shape.dim] shape = ', '.join(dims) shape = '(' + shape + ')' - print(' shape: ' + shape) - print(' name: ' + tensor_info.name) + in_print(' shape: ' + shape) + in_print(' name: ' + tensor_info.name) def _show_all(saved_model_dir): @@ -186,7 +201,8 @@ def _show_all(saved_model_dir): signature_def_map = get_signature_def_map(saved_model_dir, tag_set) for signature_def_key in sorted(signature_def_map.keys()): print('\nsignature_def[\'' + signature_def_key + '\']:') - _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key) + _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key, + indent=1) def get_meta_graph_def(saved_model_dir, tag_set): @@ -230,6 +246,27 @@ def get_signature_def_map(saved_model_dir, tag_set): return meta_graph.signature_def +def scan_meta_graph_def(meta_graph_def): + """Scans meta_graph_def and reports if there are ops on blacklist. + + Print ops if they are on black list, or print success if no blacklisted ops + found. + + Args: + meta_graph_def: MetaGraphDef protocol buffer. + """ + all_ops_set = set( + meta_graph_lib.ops_used_by_graph_def(meta_graph_def.graph_def)) + blacklisted_ops = _OP_BLACKLIST & all_ops_set + if blacklisted_ops: + # TODO(yifeif): print more warnings + print('MetaGraph with tag set %s contains the following blacklisted ops:' % + meta_graph_def.meta_info_def.tags, blacklisted_ops) + else: + print('MetaGraph with tag set %s does not contain blacklisted ops.' % + meta_graph_def.meta_info_def.tags) + + def run_saved_model_with_feed_dict(saved_model_dir, tag_set, signature_def_key, input_tensor_key_feed_dict, outdir, overwrite_flag, tf_debug=False): @@ -597,6 +634,21 @@ def run(args): args.overwrite, tf_debug=args.tf_debug) +def scan(args): + """Function triggered by scan command. + + Args: + args: A namespace parsed from command line. + """ + if args.tag_set: + scan_meta_graph_def( + saved_model_utils.get_meta_graph_def(args.dir, args.tag_set)) + else: + saved_model = reader.read_saved_model(args.dir) + for meta_graph_def in saved_model.meta_graphs: + scan_meta_graph_def(meta_graph_def) + + def create_parser(): """Creates a parser that parse the command line arguments. @@ -614,19 +666,19 @@ def create_parser(): show_msg = ( 'Usage examples:\n' 'To show all tag-sets in a SavedModel:\n' - '$saved_model_cli show --dir /tmp/saved_model\n' + '$saved_model_cli show --dir /tmp/saved_model\n\n' 'To show all available SignatureDef keys in a ' 'MetaGraphDef specified by its tag-set:\n' - '$saved_model_cli show --dir /tmp/saved_model --tag_set serve\n' + '$saved_model_cli show --dir /tmp/saved_model --tag_set serve\n\n' 'For a MetaGraphDef with multiple tags in the tag-set, all tags must be ' 'passed in, separated by \';\':\n' '$saved_model_cli show --dir /tmp/saved_model --tag_set serve,gpu\n\n' 'To show all inputs and outputs TensorInfo for a specific' ' SignatureDef specified by the SignatureDef key in a' ' MetaGraph.\n' - '$saved_model_cli show --dir /tmp/saved_model --tag_set serve ' - '--signature_def serving_default\n\n' - 'To show all available information in the SavedModel\n:' + '$saved_model_cli show --dir /tmp/saved_model --tag_set serve' + ' --signature_def serving_default\n\n' + 'To show all available information in the SavedModel:\n' '$saved_model_cli show --dir /tmp/saved_model --all') parser_show = subparsers.add_parser( 'show', @@ -658,12 +710,14 @@ def create_parser(): run_msg = ('Usage example:\n' 'To run input tensors from files through a MetaGraphDef and save' ' the output tensors to files:\n' - '$saved_model_cli show --dir /tmp/saved_model --tag_set serve ' - '--signature_def serving_default ' - '--inputs input1_key=/tmp/124.npz[x],input2_key=/tmp/123.npy ' - '--input_exprs \'input3_key=np.ones(2)\' --input_examples ' - '\'input4_key=[{"id":[26],"weights":[0.5, 0.5]}]\' ' - '--outdir=/out\n\n' + '$saved_model_cli show --dir /tmp/saved_model --tag_set serve \\\n' + ' --signature_def serving_default \\\n' + ' --inputs input1_key=/tmp/124.npz[x],input2_key=/tmp/123.npy ' + '\\\n' + ' --input_exprs \'input3_key=np.ones(2)\' \\\n' + ' --input_examples ' + '\'input4_key=[{"id":[26],"weights":[0.5, 0.5]}]\' \\\n' + ' --outdir=/out\n\n' 'For more information about input file format, please see:\n' 'https://www.tensorflow.org/programmers_guide/saved_model_cli\n') parser_run = subparsers.add_parser( @@ -716,6 +770,26 @@ def create_parser(): 'SavedModel.') parser_run.set_defaults(func=run) + # scan command + scan_msg = ('Usage example:\n' + 'To scan for blacklisted ops in SavedModel:\n' + '$saved_model_cli scan --dir /tmp/saved_model\n' + 'To scan a specific MetaGraph, pass in --tag_set\n') + parser_scan = subparsers.add_parser( + 'scan', + description=scan_msg, + formatter_class=argparse.RawTextHelpFormatter) + parser_scan.add_argument( + '--dir', + type=str, + required=True, + help='directory containing the SavedModel to execute') + parser_scan.add_argument( + '--tag_set', + type=str, + help='tag-set of graph in SavedModel to scan, separated by \',\'') + parser_scan.set_defaults(func=scan) + return parser diff --git a/tensorflow/python/tools/saved_model_cli_test.py b/tensorflow/python/tools/saved_model_cli_test.py index d6cbc49ba1e08a6b808b228fb8d69fc14f36e3d2..eedc893a38d3d0857dd49c7ce03f3921da48fdbd 100644 --- a/tensorflow/python/tools/saved_model_cli_test.py +++ b/tensorflow/python/tools/saved_model_cli_test.py @@ -61,83 +61,84 @@ class SavedModelCLITestCase(test.TestCase): exp_out = """MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs: signature_def['classify_x2_to_y3']: -The given SavedModel SignatureDef contains the following input(s): -inputs['inputs'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: x2:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['scores'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y3:0 -Method name is: tensorflow/serving/classify + The given SavedModel SignatureDef contains the following input(s): + inputs['inputs'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: x2:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['scores'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y3:0 + Method name is: tensorflow/serving/classify signature_def['classify_x_to_y']: -The given SavedModel SignatureDef contains the following input(s): -inputs['inputs'] tensor_info: - dtype: DT_STRING - shape: unknown_rank - name: tf_example:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['scores'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y:0 -Method name is: tensorflow/serving/classify + The given SavedModel SignatureDef contains the following input(s): + inputs['inputs'] tensor_info: + dtype: DT_STRING + shape: unknown_rank + name: tf_example:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['scores'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y:0 + Method name is: tensorflow/serving/classify signature_def['regress_x2_to_y3']: -The given SavedModel SignatureDef contains the following input(s): -inputs['inputs'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: x2:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['outputs'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y3:0 -Method name is: tensorflow/serving/regress + The given SavedModel SignatureDef contains the following input(s): + inputs['inputs'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: x2:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['outputs'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y3:0 + Method name is: tensorflow/serving/regress signature_def['regress_x_to_y']: -The given SavedModel SignatureDef contains the following input(s): -inputs['inputs'] tensor_info: - dtype: DT_STRING - shape: unknown_rank - name: tf_example:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['outputs'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y:0 -Method name is: tensorflow/serving/regress + The given SavedModel SignatureDef contains the following input(s): + inputs['inputs'] tensor_info: + dtype: DT_STRING + shape: unknown_rank + name: tf_example:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['outputs'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y:0 + Method name is: tensorflow/serving/regress signature_def['regress_x_to_y2']: -The given SavedModel SignatureDef contains the following input(s): -inputs['inputs'] tensor_info: - dtype: DT_STRING - shape: unknown_rank - name: tf_example:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['outputs'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y2:0 -Method name is: tensorflow/serving/regress + The given SavedModel SignatureDef contains the following input(s): + inputs['inputs'] tensor_info: + dtype: DT_STRING + shape: unknown_rank + name: tf_example:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['outputs'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y2:0 + Method name is: tensorflow/serving/regress signature_def['serving_default']: -The given SavedModel SignatureDef contains the following input(s): -inputs['x'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: x:0 -The given SavedModel SignatureDef contains the following output(s): -outputs['y'] tensor_info: - dtype: DT_FLOAT - shape: (-1, 1) - name: y:0 -Method name is: tensorflow/serving/predict""" + The given SavedModel SignatureDef contains the following input(s): + inputs['x'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: x:0 + The given SavedModel SignatureDef contains the following output(s): + outputs['y'] tensor_info: + dtype: DT_FLOAT + shape: (-1, 1) + name: y:0 + Method name is: tensorflow/serving/predict""" # pylint: enable=line-too-long + self.maxDiff = None # Produce a useful error msg if the comparison fails self.assertMultiLineEqual(output, exp_out) self.assertEqual(err.getvalue().strip(), '') @@ -193,11 +194,11 @@ Method name is: tensorflow/serving/predict""" output = out.getvalue().strip() expected_output = ( 'The given SavedModel SignatureDef contains the following input(s):\n' - 'inputs[\'x\'] tensor_info:\n' - ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: x:0\n' + ' inputs[\'x\'] tensor_info:\n' + ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: x:0\n' 'The given SavedModel SignatureDef contains the following output(s):\n' - 'outputs[\'y\'] tensor_info:\n' - ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: y:0\n' + ' outputs[\'y\'] tensor_info:\n' + ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: y:0\n' 'Method name is: tensorflow/serving/predict') self.assertEqual(output, expected_output) self.assertEqual(err.getvalue().strip(), '') @@ -524,6 +525,28 @@ Method name is: tensorflow/serving/predict""" y_expected = np.array([[2.5], [3.0]]) self.assertAllClose(y_expected, y_actual) + def testScanCommand(self): + self.parser = saved_model_cli.create_parser() + base_path = test.test_src_dir_path(SAVED_MODEL_PATH) + args = self.parser.parse_args(['scan', '--dir', base_path]) + with captured_output() as (out, _): + saved_model_cli.scan(args) + output = out.getvalue().strip() + self.assertTrue('does not contain blacklisted ops' in output) + + def testScanCommandFoundBlacklistedOp(self): + self.parser = saved_model_cli.create_parser() + base_path = test.test_src_dir_path(SAVED_MODEL_PATH) + args = self.parser.parse_args( + ['scan', '--dir', base_path, '--tag_set', 'serve']) + op_blacklist = saved_model_cli._OP_BLACKLIST + saved_model_cli._OP_BLACKLIST = set(['VariableV2']) + with captured_output() as (out, _): + saved_model_cli.scan(args) + saved_model_cli._OP_BLACKLIST = op_blacklist + output = out.getvalue().strip() + self.assertTrue('\'VariableV2\'' in output) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/training/adadelta.py b/tensorflow/python/training/adadelta.py index 13c07cfd7bf4333fee3edc3c3ad9d2fb7bcbaad2..c08e3cca007dc17f1112d53bf729c1accf61b5df 100644 --- a/tensorflow/python/training/adadelta.py +++ b/tensorflow/python/training/adadelta.py @@ -22,8 +22,10 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.AdadeltaOptimizer") class AdadeltaOptimizer(optimizer.Optimizer): """Optimizer that implements the Adadelta algorithm. diff --git a/tensorflow/python/training/adagrad.py b/tensorflow/python/training/adagrad.py index afa192f7cc6e0ecd629fd94252d26961f1407183..deb4e6f546379eff330235dbc302a30c44193830 100644 --- a/tensorflow/python/training/adagrad.py +++ b/tensorflow/python/training/adagrad.py @@ -25,8 +25,10 @@ from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.AdagradOptimizer") class AdagradOptimizer(optimizer.Optimizer): """Optimizer that implements the Adagrad algorithm. diff --git a/tensorflow/python/training/adagrad_da.py b/tensorflow/python/training/adagrad_da.py index b3f9ea323c2bb4fd9ecee93863fbc7955b47a947..5ba403554f570d9df33a5d525a40de2eb0d11138 100644 --- a/tensorflow/python/training/adagrad_da.py +++ b/tensorflow/python/training/adagrad_da.py @@ -23,8 +23,10 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.AdagradDAOptimizer") class AdagradDAOptimizer(optimizer.Optimizer): """Adagrad Dual Averaging algorithm for sparse linear models. diff --git a/tensorflow/python/training/adam.py b/tensorflow/python/training/adam.py index 0c69f8bf3997452f0eeb71c93f4fcf98eb27d8f9..006e360389b404a8edd97c9a8bf4b8876c828004 100644 --- a/tensorflow/python/training/adam.py +++ b/tensorflow/python/training/adam.py @@ -26,8 +26,10 @@ from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.AdamOptimizer") class AdamOptimizer(optimizer.Optimizer): """Optimizer that implements the Adam algorithm. @@ -104,10 +106,10 @@ class AdamOptimizer(optimizer.Optimizer): self._updated_lr = None def _get_beta_accumulators(self): - if context.in_graph_mode(): - graph = ops.get_default_graph() - else: + if context.executing_eagerly(): graph = None + else: + graph = ops.get_default_graph() return (self._get_non_slot_variable("beta1_power", graph=graph), self._get_non_slot_variable("beta2_power", graph=graph)) diff --git a/tensorflow/python/training/adam_test.py b/tensorflow/python/training/adam_test.py index a521f1299e035424d1c3897a469655db732b0dcd..9be8b6aafefa33977511cde24dd2e87dd6c3b81a 100644 --- a/tensorflow/python/training/adam_test.py +++ b/tensorflow/python/training/adam_test.py @@ -184,7 +184,7 @@ class AdamOptimizerTest(test.TestCase): # Shouldn't return non-slot variables from other graphs. self.assertEqual(0, len(opt.variables())) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(variables.global_variables_initializer()) # Fetch params to validate initial values self.assertAllClose([1.0, 2.0], self.evaluate(var0)) @@ -194,7 +194,7 @@ class AdamOptimizerTest(test.TestCase): # Run 3 steps of Adam for t in range(1, 4): - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(update) elif t > 1: opt.apply_gradients(zip([grads0, grads1], [var0, var1])) @@ -319,6 +319,15 @@ class AdamOptimizerTest(test.TestCase): # fails. optimizer.apply_gradients([(grads0, var0)]) + def testSlotsUniqueEager(self): + with context.eager_mode(): + v1 = resource_variable_ops.ResourceVariable(1.) + v2 = resource_variable_ops.ResourceVariable(1.) + opt = adam.AdamOptimizer(1.) + opt.minimize(lambda: v1 + v2) + # There should be two non-slot variables, and two unique slot variables + # for v1 and v2 respectively. + self.assertEqual(6, len(set(opt.variables()))) if __name__ == "__main__": test.main() diff --git a/tensorflow/python/training/basic_loops.py b/tensorflow/python/training/basic_loops.py index 52b0f4210612bad4a2e838153ac9cbdb1023bf66..7af821c81928e67e0f258bc064d582a4186995c1 100644 --- a/tensorflow/python/training/basic_loops.py +++ b/tensorflow/python/training/basic_loops.py @@ -18,8 +18,10 @@ from __future__ import division from __future__ import print_function from tensorflow.python.framework import errors +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.basic_train_loop") def basic_train_loop(supervisor, train_step_fn, args=None, kwargs=None, master=""): """Basic loop to train a model. diff --git a/tensorflow/python/training/basic_session_run_hooks.py b/tensorflow/python/training/basic_session_run_hooks.py index 864c2e4406ebba0f35c3a1e9d783af43046325ce..aae757b99aa9abb2fca112dcc781fc31e367649d 100644 --- a/tensorflow/python/training/basic_session_run_hooks.py +++ b/tensorflow/python/training/basic_session_run_hooks.py @@ -47,6 +47,7 @@ from tensorflow.python.training import session_run_hook from tensorflow.python.training import training_util from tensorflow.python.training.session_run_hook import SessionRunArgs from tensorflow.python.training.summary_io import SummaryWriterCache +from tensorflow.python.util.tf_export import tf_export class _HookTimer(object): @@ -85,6 +86,7 @@ class _HookTimer(object): raise NotImplementedError +@tf_export("train.SecondOrStepTimer") class SecondOrStepTimer(_HookTimer): """Timer that triggers at most once every N seconds or once every N steps. """ @@ -164,6 +166,7 @@ class NeverTriggerTimer(_HookTimer): return None +@tf_export("train.LoggingTensorHook") class LoggingTensorHook(session_run_hook.SessionRunHook): """Prints the given tensors every N local steps, every N seconds, or at end. @@ -262,6 +265,7 @@ class LoggingTensorHook(session_run_hook.SessionRunHook): self._log_tensors(values) +@tf_export("train.StopAtStepHook") class StopAtStepHook(session_run_hook.SessionRunHook): """Hook that requests stop at a specified step.""" @@ -317,6 +321,7 @@ class StopAtStepHook(session_run_hook.SessionRunHook): run_context.request_stop() +@tf_export("train.CheckpointSaverListener") class CheckpointSaverListener(object): """Interface for listeners that take action before or after checkpoint save. @@ -375,6 +380,7 @@ class CheckpointSaverListener(object): pass +@tf_export("train.CheckpointSaverHook") class CheckpointSaverHook(session_run_hook.SessionRunHook): """Saves checkpoints every N steps or seconds.""" @@ -497,6 +503,7 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook): return savers[0] +@tf_export("train.StepCounterHook") class StepCounterHook(session_run_hook.SessionRunHook): """Hook that counts steps per second.""" @@ -575,12 +582,14 @@ class StepCounterHook(session_run_hook.SessionRunHook): self._last_global_step = stale_global_step +@tf_export("train.NanLossDuringTrainingError") class NanLossDuringTrainingError(RuntimeError): def __str__(self): return "NaN loss during training." +@tf_export("train.NanTensorHook") class NanTensorHook(session_run_hook.SessionRunHook): """Monitors the loss tensor and stops training if loss is NaN. @@ -612,6 +621,7 @@ class NanTensorHook(session_run_hook.SessionRunHook): run_context.request_stop() +@tf_export("train.SummarySaverHook") class SummarySaverHook(session_run_hook.SessionRunHook): """Saves summaries every N steps.""" @@ -720,6 +730,7 @@ class SummarySaverHook(session_run_hook.SessionRunHook): return summary_op +@tf_export("train.GlobalStepWaiterHook") class GlobalStepWaiterHook(session_run_hook.SessionRunHook): """Delays execution until global step reaches `wait_until_step`. @@ -767,6 +778,7 @@ class GlobalStepWaiterHook(session_run_hook.SessionRunHook): time.sleep(0.5) +@tf_export("train.FinalOpsHook") class FinalOpsHook(session_run_hook.SessionRunHook): """A hook which evaluates `Tensors` at the end of a session.""" @@ -793,6 +805,7 @@ class FinalOpsHook(session_run_hook.SessionRunHook): feed_dict=self._final_ops_feed_dict) +@tf_export("train.FeedFnHook") class FeedFnHook(session_run_hook.SessionRunHook): """Runs `feed_fn` and sets the `feed_dict` accordingly.""" @@ -810,6 +823,7 @@ class FeedFnHook(session_run_hook.SessionRunHook): fetches=None, feed_dict=self.feed_fn()) +@tf_export("train.ProfilerHook") class ProfilerHook(session_run_hook.SessionRunHook): """Captures CPU/GPU profiling information every N steps or seconds. diff --git a/tensorflow/python/training/checkpoint_ops.py b/tensorflow/python/training/checkpoint_ops.py index 7f92d94d2be369709608d36c109863b0ebfb7bbe..a6e9662b7305a00f1fcf03245685e93b756942d3 100644 --- a/tensorflow/python/training/checkpoint_ops.py +++ b/tensorflow/python/training/checkpoint_ops.py @@ -149,7 +149,7 @@ def _load_and_remap_matrix(ckpt_path, num_rows_present = num_rows_to_load if remap_rows: row_remapping, num_rows_present = ( - gen_checkpoint_ops._generate_vocab_remapping( # pylint: disable=protected-access + gen_checkpoint_ops.generate_vocab_remapping( new_vocab_file=new_row_vocab_file, old_vocab_file=old_row_vocab_file, new_vocab_offset=new_row_vocab_offset, @@ -168,7 +168,7 @@ def _load_and_remap_matrix(ckpt_path, num_cols_present = new_col_vocab_size if remap_cols: col_remapping, num_cols_present = ( - gen_checkpoint_ops._generate_vocab_remapping( # pylint: disable=protected-access + gen_checkpoint_ops.generate_vocab_remapping( new_vocab_file=new_col_vocab_file, old_vocab_file=old_col_vocab_file, new_vocab_offset=0, # Offset is unused for cols (no partitioning). @@ -178,7 +178,7 @@ def _load_and_remap_matrix(ckpt_path, num_rows_to_load * new_col_vocab_size - num_rows_present * num_cols_present, 1 ]) - return_tensor = gen_checkpoint_ops._load_and_remap_matrix( # pylint: disable=protected-access + return_tensor = gen_checkpoint_ops.load_and_remap_matrix( ckpt_path=ckpt_path, old_tensor_name=old_tensor_name, row_remapping=row_remapping, diff --git a/tensorflow/python/training/checkpoint_utils.py b/tensorflow/python/training/checkpoint_utils.py index 63235a14540b739fca94db51fb29cdf22275b9fd..e7f88de1d2290a49f3b7bdf47417016d7e7c9cea 100644 --- a/tensorflow/python/training/checkpoint_utils.py +++ b/tensorflow/python/training/checkpoint_utils.py @@ -23,12 +23,14 @@ import six from tensorflow.python import pywrap_tensorflow from tensorflow.python.framework import ops from tensorflow.python.ops import io_ops +from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope as vs from tensorflow.python.ops import variables from tensorflow.python.platform import gfile from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import saver +from tensorflow.python.util.tf_export import tf_export __all__ = [ @@ -36,6 +38,7 @@ __all__ = [ ] +@tf_export("train.load_checkpoint") def load_checkpoint(ckpt_dir_or_file): """Returns `CheckpointReader` for checkpoint found in `ckpt_dir_or_file`. @@ -60,6 +63,7 @@ def load_checkpoint(ckpt_dir_or_file): return pywrap_tensorflow.NewCheckpointReader(filename) +@tf_export("train.load_variable") def load_variable(ckpt_dir_or_file, name): """Returns the tensor value of the given variable in the checkpoint. @@ -77,6 +81,7 @@ def load_variable(ckpt_dir_or_file, name): return reader.get_tensor(name) +@tf_export("train.list_variables") def list_variables(ckpt_dir_or_file): """Returns list of all variables in the checkpoint. @@ -95,6 +100,7 @@ def list_variables(ckpt_dir_or_file): return result +@tf_export("train.init_from_checkpoint") def init_from_checkpoint(ckpt_dir_or_file, assignment_map): """Initializes current variables with tensors loaded from given checkpoint. @@ -284,10 +290,20 @@ def _set_checkpoint_initializer(variable, name: Name of the operation. """ base_type = variable.dtype.base_dtype - with ops.colocate_with(variable): + # Do not colocate with variable since RestoreV2 op only runs on CPU and + # colocation will force variable (and other ops that colocate with variable) + # to be on CPU as well. It is okay to place the variable's initializer op on + # CPU since it will only be run once at the start. + with ops.device(variable.device), ops.device("/cpu:0"): restore_op = io_ops.restore_v2( ckpt_file, [tensor_name], [slice_spec], [base_type], name=name)[0] - variable._initializer_op = state_ops.assign(variable, restore_op) # pylint:disable=protected-access + if isinstance(variable, resource_variable_ops.ResourceVariable): + init_op = variable.assign(restore_op, read_value=False) + else: + init_op = state_ops.assign(variable, restore_op) + variable._initializer_op = init_op # pylint:disable=protected-access + restore_op.set_shape(variable.shape) + variable._initial_value = restore_op # pylint:disable=protected-access def _set_variable_or_list_initializer(variable_or_list, ckpt_file, diff --git a/tensorflow/python/training/checkpoint_utils_test.py b/tensorflow/python/training/checkpoint_utils_test.py index cd17faa040d5b85263b54bc53100b18f736a12e0..4e08a1c859fbaac75e7cd09ad498d9fea14c6338 100644 --- a/tensorflow/python/training/checkpoint_utils_test.py +++ b/tensorflow/python/training/checkpoint_utils_test.py @@ -26,6 +26,7 @@ from tensorflow.python.framework import errors_impl from tensorflow.python.framework import ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import partitioned_variables +from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test @@ -145,6 +146,36 @@ class CheckpointsTest(test.TestCase): # Check that tensors are not explicitly in the graph. self.assertLess(len(str(session.graph.as_graph_def())), 29000) + def testInitialValueComesFromCheckpoint(self): + checkpoint_dir = self.get_temp_dir() + with self.test_session() as session: + v1, _, _, _ = _create_checkpoints(session, checkpoint_dir) + + # New graph and session. + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as session: + with variable_scope.variable_scope( + "some_scope", initializer=init_ops.zeros_initializer()): + my1 = variable_scope.get_variable("my1", [1, 10]) + + before = my1.initialized_value() + + checkpoint_utils.init_from_checkpoint(checkpoint_dir, {"var1": my1}) + + after = my1.initialized_value() + + self.assertAllEqual(session.run(before), [[0.0] * 10]) + self.assertAllEqual(session.run(after), v1) + + session.run(variables.global_variables_initializer()) + + self.assertAllEqual(session.run(my1), v1) + self.assertAllEqual(session.run(my1.initialized_value()), v1) + self.assertAllClose(session.run(before), v1) + self.assertAllClose(session.run(after), v1) + with self.assertRaises(AssertionError): + self.assertAllClose(v1, [[0.0] * 10]) + def testInitWithScopeDoesNotCaptureSuffixes(self): checkpoint_dir = self.get_temp_dir() with self.test_session() as session: @@ -176,7 +207,9 @@ class CheckpointsTest(test.TestCase): checkpoint_utils.init_from_checkpoint(checkpoint_dir, {"useful_scope/": "useful_scope/"}) - self.assertEqual(my4._initializer_op.op.inputs[1].device, "/job:ps") + # initializer runs on the same task but always on CPU. + self.assertEqual(my4._initializer_op.op.inputs[1].device, + "/job:ps/device:CPU:0") def testInitFromRootCheckpoint(self): checkpoint_dir = self.get_temp_dir() @@ -332,6 +365,31 @@ class CheckpointsTest(test.TestCase): checkpoint_utils.init_from_checkpoint(checkpoint_dir, {"useful_scope": "some_scope/"}) + def testNoAdditionalReadOpsForResourceVariables(self): + checkpoint_dir = self.get_temp_dir() + with self.test_session() as session: + v1, _, _, _ = _create_checkpoints(session, checkpoint_dir) + + # New graph and session. + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as session: + my1 = resource_variable_ops.ResourceVariable([[0.0] * 10], name="my1") + + with ops.name_scope("init_from_checkpoint"): + checkpoint_utils.init_from_checkpoint(checkpoint_dir, {"var1": my1}) + + # Basic sanity checks: + session.run(variables.global_variables_initializer()) + self.assertAllEqual(session.run(my1), v1) + + ops_in_init_from_checkpoint_scope = [ + op for op in g.get_operations() + if (op.name.startswith("init_from_checkpoint/") and + not op.name.startswith("init_from_checkpoint/checkpoint_initializer" + ) and op.type != "AssignVariableOp") + ] + self.assertEqual(ops_in_init_from_checkpoint_scope, []) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/training/checkpointable.py b/tensorflow/python/training/checkpointable.py new file mode 100644 index 0000000000000000000000000000000000000000..96e3c4821cd319df0300edce2c722eb72ec752a6 --- /dev/null +++ b/tensorflow/python/training/checkpointable.py @@ -0,0 +1,676 @@ +"""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 + +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 gen_io_ops as io_ops +from tensorflow.python.util import nest + +# A key indicating a variable's value in an object's checkpointed Tensors +# (Checkpointable._gather_saveables_for_checkpoint). If this is the only key and +# the object has no dependencies, then its value may be restored on object +# creation (avoiding double assignment when executing eagerly). +VARIABLE_VALUE_KEY = "VARIABLE_VALUE" + +CheckpointableReference = collections.namedtuple( + "CheckpointableReference", + [ + # The local name for this dependency. + "name", + # The Checkpointable object being referenced. + "ref" + ]) + + +class CheckpointInitialValue(ops.Tensor): + """Tensor wrapper for managing update UIDs in `Variables`. + + When supplied as an initial value, objects of this type let a `Variable` + (`Variable`, `ResourceVariable`, etc.) know the UID of the restore the initial + value came from. This allows deferred restorations to be sequenced in the + order the user specified them, and lets us fall back on assignment if an + initial value is not set (e.g. due to a custom getter interfering). + + See comments in _add_variable_with_custom_getter for more information about + how `CheckpointInitialValue` is used. + """ + + def __init__(self, checkpoint_position, shape=None): + self.wrapped_value = checkpoint_position.value_tensors()[ + VARIABLE_VALUE_KEY] + if shape: + # We need to set the static shape information on the initializer if + # possible so we don't get a variable with an unknown shape. + self.wrapped_value.set_shape(shape) + self._checkpoint_position = checkpoint_position + + @property + def __class__(self): + return (self.wrapped_value.__class__, CheckpointInitialValue) + + def __getattr__(self, attr): + try: + return getattr(self.wrapped_value, attr) + except AttributeError: + return self.__getattribute__(attr) + + @property + def checkpoint_position(self): + return self._checkpoint_position + + +class _CheckpointPosition(object): + """Indicates a position within a `_Checkpoint`.""" + + def __init__(self, checkpoint, proto_id): + """Specify an object within a checkpoint. + + Args: + checkpoint: A _Checkpoint object. + proto_id: The index of this object in CheckpointableObjectGraph.nodes. + """ + self._checkpoint = checkpoint + self._proto_id = proto_id + + def restore(self, checkpointable): + """Restore this value into `checkpointable`.""" + if self.bind_object(checkpointable): + # This object's correspondence with a checkpointed object is new, so + # process deferred restorations for it and its dependencies. + restore_ops = checkpointable._restore_from_checkpoint_position(self) # pylint: disable=protected-access + if restore_ops: + self._checkpoint.restore_ops.extend(restore_ops) + + def bind_object(self, checkpointable): + """Set a checkpoint<->object correspondence and process slot variables. + + Args: + checkpointable: The object to record a correspondence for. + Returns: + True if this is a new assignment, False if this object has already been + mapped to a checkpointed `Object` proto. + Raises: + AssertionError: If another object is already bound to the `Object` proto. + """ + checkpoint = self.checkpoint + current_assignment = checkpoint.object_by_proto_id.get(self._proto_id, None) + if current_assignment is None: + checkpoint.object_by_proto_id[self._proto_id] = checkpointable + for deferred_slot_restoration in ( + checkpoint.deferred_slot_restorations.pop(self._proto_id, ())): + checkpointable._create_or_restore_slot_variable( # pylint: disable=protected-access + slot_variable_position=_CheckpointPosition( + checkpoint=checkpoint, + proto_id=deferred_slot_restoration.slot_variable_id), + variable=deferred_slot_restoration.original_variable, + slot_name=deferred_slot_restoration.slot_name) + for slot_restoration in checkpoint.slot_restorations.pop( + self._proto_id, ()): + optimizer_object = checkpoint.object_by_proto_id.get( + slot_restoration.optimizer_id, None) + if optimizer_object is None: + # The optimizer has not yet been created or tracked. Record in the + # checkpoint that the slot variables need to be restored when it is. + checkpoint.deferred_slot_restorations.setdefault( + slot_restoration.optimizer_id, []).append( + _DeferredSlotVariableRestoration( + original_variable=checkpointable, + slot_variable_id=slot_restoration.slot_variable_id, + slot_name=slot_restoration.slot_name)) + else: + optimizer_object._create_or_restore_slot_variable( # pylint: disable=protected-access + slot_variable_position=_CheckpointPosition( + checkpoint=checkpoint, + proto_id=slot_restoration.slot_variable_id), + variable=checkpointable, + slot_name=slot_restoration.slot_name) + return True # New assignment + else: + # The object was already mapped for this checkpoint load, which means + # we don't need to do anything besides check that the mapping is + # consistent (if the dependency DAG is not a tree then there are + # multiple paths to the same object). + if current_assignment is not checkpointable: + raise AssertionError( + ("Unable to load the checkpoint into this object graph. Either " + "the Checkpointable object references in the Python program " + "have changed in an incompatible way, or the checkpoint was " + "generated in an incompatible program.\n\nTwo checkpoint " + "references resolved to different objects (%s and %s).") + % (current_assignment, checkpointable)) + return False # Not a new assignment + + def is_simple_variable(self): + """Determine whether this value is restorable with a Tensor initializer.""" + attributes = self.object_proto.attributes + return (len(attributes) == 1 + and attributes[0].name == VARIABLE_VALUE_KEY + and not self.object_proto.children) + + def value_tensors(self): + """Create value `Tensor`s for this object's attributes. + + Does not require that the Python object has been created. Used for + restore-on-create when executing eagerly. + + Returns: + A dictionary mapping from object attribute names to `Tensor`s. + """ + value_tensors = {} + for serialized_tensor in self.object_proto.attributes: + checkpoint_key = serialized_tensor.checkpoint_key + dtype = self._checkpoint.dtype_map[checkpoint_key] + base_type = dtype.base_dtype + with ops.init_scope(): + with ops.device("/cpu:0"): + # Run the restore itself on the CPU. + value, = io_ops.restore_v2( + prefix=self._checkpoint.save_path, + tensor_names=[checkpoint_key], + shape_and_slices=[""], + dtypes=[base_type], + name="%s_checkpoint_read" % (serialized_tensor.name,)) + # Copy the value to the current device if necessary. + value_tensors[serialized_tensor.name] = array_ops.identity(value) + return value_tensors + + def restore_ops(self): + """Create or fetch restore ops for this object's attributes. + + Requires that the `Checkpointable` Python object has been bound to an object + ID in the checkpoint. + + Returns: + A list of operations when graph building, or an empty list when executing + eagerly. + """ + saveables = self.checkpointable._gather_saveables_for_checkpoint() # pylint: disable=protected-access + # Name saveables based on the name this object had when it was checkpointed. + named_saveables = {} + restore_ops = [] + building_graph = not context.executing_eagerly() + for serialized_tensor in self.object_proto.attributes: + saveable_factory = saveables.get(serialized_tensor.name, None) + if saveable_factory is None: + # Purposefully does not throw an exception if attributes have been added + # or deleted. Stores unused attributes so an exception can be raised if + # the user decides to check that everything in the checkpoint was + # loaded. + self._checkpoint.unused_attributes.setdefault( + self.checkpointable, []).append(serialized_tensor.name) + continue + if building_graph: + existing_ops = self._checkpoint.restore_ops_by_name.get( + serialized_tensor.name, None) + else: + existing_ops = None + if existing_ops is None: + if callable(saveable_factory): + saveable = saveable_factory(name=serialized_tensor.checkpoint_key) + else: + saveable = saveable_factory + named_saveables[serialized_tensor.checkpoint_key] = saveable + if named_saveables: + validated_saveables = ( + self._checkpoint.builder._ValidateAndSliceInputs(named_saveables)) # pylint: disable=protected-access + validated_names = set(saveable.name for saveable in validated_saveables) + if set(named_saveables.keys()) != validated_names: + raise AssertionError( + ("Saveable keys changed when validating. Got back %s, was " + "expecting %s") % (named_saveables.keys(), validated_names)) + all_tensors = self._checkpoint.builder.bulk_restore( + filename_tensor=self._checkpoint.save_path, + saveables=validated_saveables, preferred_shard=-1, + restore_sequentially=False) + saveable_index = 0 + for saveable in validated_saveables: + num_specs = len(saveable.specs) + saveable_tensors = all_tensors[ + saveable_index:saveable_index + num_specs] + saveable_index += num_specs + restore_op = saveable.restore(saveable_tensors, restored_shapes=None) + if building_graph: + assert saveable.name not in self._checkpoint.restore_ops_by_name + self._checkpoint.restore_ops_by_name[saveable.name] = restore_op + restore_ops.append(restore_op) + return restore_ops + + @property + def checkpoint(self): + return self._checkpoint + + @property + def checkpointable(self): + return self._checkpoint.object_by_proto_id[self._proto_id] + + @property + def object_proto(self): + return self._checkpoint.object_graph_proto.nodes[self._proto_id] + + @property + def restore_uid(self): + return self._checkpoint.restore_uid + + def __repr__(self): + return repr(self.object_proto) + + +_DeferredSlotVariableRestoration = collections.namedtuple( + "_DeferredSlotVariableRestoration", + [ + "original_variable", + "slot_variable_id", + "slot_name", + ] +) + +_SlotVariableRestoration = collections.namedtuple( + "_SlotVariableRestoration", + [ + # The checkpoint proto id of the optimizer object. + "optimizer_id", + # The checkpoint proto id of the slot variable. + "slot_variable_id", + "slot_name", + ]) + + +class CheckpointableBase(object): + """Base class for `Checkpointable` objects without automatic dependencies. + + This class has no __setattr__ override for performance reasons. Dependencies + must be added explicitly. Unless attribute assignment is performance-critical, + use `Checkpointable` instead. Use `CheckpointableBase` for `isinstance` + checks. + """ + + def _maybe_initialize_checkpointable(self): + """Initialize dependency management. + + Not __init__, since most objects will forget to call it. + """ + if hasattr(self, "_unconditional_checkpoint_dependencies"): + # __init__ already called. This check means that we don't need + # Checkpointable.__init__() in the constructor of every TensorFlow object. + return + # A list of CheckpointableReference objects. Some classes implementing + # `Checkpointable`, notably `Optimizer`s, may override the + # _checkpoint_dependencies property with conditional dependencies + # (e.g. based on the current graph when saving). + self._unconditional_checkpoint_dependencies = [] + # Maps names -> Checkpointable objects + self._unconditional_dependency_names = {} + # Restorations for other Checkpointable objects on which this object may + # eventually depend. + self._deferred_dependencies = {} # local name -> _CheckpointPosition list + # The UID of the highest assignment to this object. Used to ensure that the + # last requested assignment determines the final value of an object. + if hasattr(self, "_update_uid"): + raise AssertionError( + "Internal error: the object had an update UID set before its " + "initialization code was run.") + self._update_uid = -1 + + @property + def _checkpoint_dependencies(self): + """All dependencies of this object. + + May be overridden to include conditional dependencies. + + Returns: + A list of `CheckpointableReference` objects indicating named + `Checkpointable` dependencies which should be saved along with this + object. + """ + return self._unconditional_checkpoint_dependencies + + def _lookup_dependency(self, name): + """Look up a dependency by name. + + May be overridden to include conditional dependencies. + + Args: + name: The local name of the dependency. + Returns: + A `Checkpointable` object, or `None` if no dependency by this name was + found. + """ + return self._unconditional_dependency_names.get(name, None) + + def _add_variable_with_custom_getter( + self, name, shape=None, dtype=dtypes.float32, + initializer=None, getter=None, overwrite=False, + **kwargs_for_getter): + """Restore-on-create for a variable 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 there is a deferred + restoration left over from a call to + `_restore_from_checkpoint_position`. + getter: The getter to wrap which actually fetches the variable. + overwrite: If True, disables unique name and type checks. + **kwargs_for_getter: Passed to the getter. + + Returns: + The new variable object. + + Raises: + ValueError: If the variable name is not unique. + """ + self._maybe_initialize_checkpointable() + if not overwrite and self._lookup_dependency(name) is not None: + 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 context.executing_eagerly(): + # If this is a variable with a single Tensor stored in the checkpoint, we + # can set that value as an initializer rather than initializing and then + # assigning (when executing eagerly). This call returns None if there is + # nothing to restore. + checkpoint_initializer = self._preload_simple_restoration( + name=name, shape=shape) + else: + checkpoint_initializer = None + if (checkpoint_initializer is not None + and not ( + isinstance(initializer, CheckpointInitialValue) + and initializer.restore_uid > checkpoint_initializer.restore_uid)): + # If multiple Checkpointable objects are "creating" the same variable via + # the magic of custom getters, the one with the highest restore UID (the + # one called last) has to make the final initializer. If another custom + # getter interrupts this process by overwriting the initializer, then + # we'll catch that when we call _track_checkpointable. So this is "best + # effort" to set the initializer with the highest restore UID. + initializer = checkpoint_initializer + shape = None + + new_variable = getter( + name=name, shape=shape, dtype=dtype, initializer=initializer, + **kwargs_for_getter) + + # If we set an initializer and the variable processed it, tracking will not + # assign again. It will add this variable to our dependencies, and if there + # is a non-trivial restoration queued, it will handle that. This also + # handles slot variables. + if not overwrite or isinstance(new_variable, CheckpointableBase): + return self._track_checkpointable(new_variable, name=name, + overwrite=overwrite) + else: + # TODO(allenl): Some variable types are not yet supported. Remove this + # fallback once all get_variable() return types are Checkpointable. + return new_variable + + def _preload_simple_restoration(self, name, shape): + """Return a dependency's value for restore-on-create. + + Note the restoration is not deleted; if for some reason preload is called + and then not assigned to the variable (for example because a custom getter + overrides the initializer), the assignment will still happen once the + variable is tracked (determined based on checkpoint.restore_uid). + + Args: + name: The object-local name of the dependency holding the variable's + value. + shape: The shape of the variable being loaded into. + Returns: + An callable for use as a variable's initializer/initial_value, or None if + one should not be set (either because there was no variable with this name + in the checkpoint or because it needs more complex deserialization). Any + non-trivial deserialization will happen when the variable object is + tracked. + """ + deferred_dependencies_list = self._deferred_dependencies.get(name, ()) + if not deferred_dependencies_list: + # Nothing to do; we don't have a restore for this dependency queued up. + return + for checkpoint_position in deferred_dependencies_list: + if not checkpoint_position.is_simple_variable(): + # If _any_ pending restoration is too complicated to fit in an + # initializer (because it has dependencies, or because there are + # multiple Tensors to restore), bail and let the general tracking code + # handle it. + return None + checkpoint_position = max( + deferred_dependencies_list, + key=lambda restore: restore.checkpoint.restore_uid) + return CheckpointInitialValue( + checkpoint_position=checkpoint_position, shape=shape) + + def _track_checkpointable(self, checkpointable, name, overwrite=False): + """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 the names + provided when the checkpoint was written. To avoid breaking existing + checkpoints when modifying a class, neither variable names nor dependency + names (the names passed to `_track_checkpointable`) may change. + + Args: + checkpointable: A `Checkpointable` which this object depends on. + name: A local name for `checkpointable`, used for loading checkpoints into + the correct objects. + overwrite: Boolean, whether silently replacing dependencies is OK. Used + for __setattr__, where throwing an error on attribute reassignment would + be inappropriate. + + Returns: + `checkpointable`, for convenience when declaring a dependency and + assigning to a member variable in one statement. + + Raises: + TypeError: If `checkpointable` does not inherit from `Checkpointable`. + ValueError: If another object is already tracked by this name. + """ + self._maybe_initialize_checkpointable() + if not isinstance(checkpointable, CheckpointableBase): + raise TypeError( + ("Checkpointable._track_checkpointable() passed type %s, not a " + "Checkpointable.") % (type(checkpointable),)) + new_reference = CheckpointableReference(name=name, ref=checkpointable) + current_object = self._lookup_dependency(name) + if (current_object is not None + and current_object is not checkpointable): + if not overwrite: + raise ValueError( + ("Called Checkpointable._track_checkpointable() with name='%s', " + "but a Checkpointable with this name is already declared as a " + "dependency. Names must be unique (or overwrite=True).") % (name,)) + # This is a weird thing to do, but we're not going to stop people from + # using __setattr__. + for index, (old_name, _) in enumerate( + self._unconditional_checkpoint_dependencies): + if name == old_name: + self._unconditional_checkpoint_dependencies[index] = new_reference + else: + self._unconditional_checkpoint_dependencies.append(new_reference) + + self._unconditional_dependency_names[name] = checkpointable + self._handle_deferred_dependencies(name=name, checkpointable=checkpointable) + return checkpointable + + def _handle_deferred_dependencies(self, name, checkpointable): + """Pop and load any deferred checkpoint restores into `checkpointable`. + + This method does not add a new dependency on `checkpointable`, but it does + check if any outstanding/deferred dependencies have been queued waiting for + this dependency to be added (matched based on `name`). If so, + `checkpointable` and its dependencies are restored. The restorations are + considered fulfilled and so are deleted. + + `_track_checkpointable` is more appropriate for adding a + normal/unconditional dependency, and includes handling for deferred + restorations. This method allows objects such as `Optimizer` to use the same + restoration logic while managing conditional dependencies themselves, by + overriding `_checkpoint_dependencies` and `_lookup_dependency` to change the + object's dependencies based on the context it is saved/restored in (a single + optimizer instance can have state associated with multiple graphs). + + Args: + name: The name of the dependency within this object (`self`), used to + match `checkpointable` with values saved in a checkpoint. + checkpointable: The Checkpointable object to restore (inheriting from + `CheckpointableBase`). + """ + deferred_dependencies_list = self._deferred_dependencies.pop(name, ()) + for checkpoint_position in sorted( + deferred_dependencies_list, + key=lambda restore: restore.checkpoint.restore_uid, + reverse=True): + checkpoint_position.restore(checkpointable) + + def _restore_from_checkpoint_position(self, checkpoint_position): + """Restore this object and its dependencies (may be deferred).""" + # Attempt a breadth-first traversal, since presumably the user has more + # control over shorter paths. If we don't have all of the dependencies at + # this point, the end result is not breadth-first (since other deferred + # traversals will happen later). + visit_queue = collections.deque([checkpoint_position]) + restore_ops = [] + while visit_queue: + current_position = visit_queue.popleft() + restore_ops.extend(nest.flatten( + current_position.checkpointable # pylint: disable=protected-access + ._single_restoration_from_checkpoint_position( + checkpoint_position=current_position, + visit_queue=visit_queue))) + return restore_ops + + def _single_restoration_from_checkpoint_position( + self, checkpoint_position, visit_queue): + """Restore this object, and either queue its dependencies or defer them.""" + self._maybe_initialize_checkpointable() + checkpoint = checkpoint_position.checkpoint + # If the UID of this restore is lower than our current update UID, we don't + # need to actually restore the object. However, we should pass the + # restoration on to our dependencies. + if checkpoint.restore_uid > self._update_uid: + restore_ops = checkpoint_position.restore_ops() + # TODO(allenl): Get a list of feeds for saving Python state + self._update_uid = checkpoint.restore_uid + else: + restore_ops = () + for child in checkpoint_position.object_proto.children: + child_position = _CheckpointPosition( + checkpoint=checkpoint, + proto_id=child.node_id) + local_object = self._lookup_dependency(child.local_name) + if local_object is None: + # We don't yet have a dependency registered with this name. Save it + # in case we do. + self._deferred_dependencies.setdefault(child.local_name, []).append( + child_position) + else: + if child_position.bind_object(checkpointable=local_object): + # This object's correspondence is new, so dependencies need to be + # visited. Delay doing it so that we get a breadth-first dependency + # resolution order (shallowest paths first). The caller is responsible + # for emptying visit_queue. + visit_queue.append(child_position) + return restore_ops + + def _gather_saveables_for_checkpoint(self): + """Returns a dictionary of values to checkpoint with this object. + + Keys in the returned dictionary are local to this object and in a separate + namespace from dependencies. Values may either be `SaveableObject` factories + or variables easily converted to `SaveableObject`s (as in `tf.train.Saver`'s + `var_list` constructor argument). + + `SaveableObjects` have a name set, which Checkpointable needs to generate + itself. So rather than returning `SaveableObjects` directly, this method + should return a dictionary of callables which take `name` arguments and + return `SaveableObjects` with that name. + + If this object may also be passed to the global-name-based `tf.train.Saver`, + the returned callables should have a default value for their name argument + (i.e. be callable with no arguments). + + Returned values must be saved only by this object; if any value may be + shared, it should instead be a dependency. For example, variable objects + save their own values with the key `VARIABLE_VALUE_KEY`, but objects which + reference variables simply add a dependency. + + Returns: + The dictionary mapping attribute names to `SaveableObject` factories + described above. For example: + {VARIABLE_VALUE_KEY: + lambda name="global_name_for_this_object": + SaveableObject(name=name, ...)} + """ + return {} + + +class Checkpointable(CheckpointableBase): + """Manages dependencies on other objects. + + `Checkpointable` objects may have dependencies: other `Checkpointable` objects + which should be saved if the object declaring the dependency is saved. A + correctly saveable program has a dependency graph such that if changing a + global variable affects an object (e.g. changes the behavior of any of its + methods) then there is a chain of dependencies from the influenced object to + the variable. + + Dependency edges have names, and are created implicitly when a + `Checkpointable` object is assigned to an attribute of another + `Checkpointable` object. For example: + + ``` + obj = Checkpointable() + obj.v = ResourceVariable(0.) + ``` + + The `Checkpointable` object `obj` now has a dependency named "v" on a + variable. + + `Checkpointable` objects may specify `Tensor`s to be saved and restored + directly (e.g. a `Variable` indicating how to save itself) rather than through + dependencies on other objects. See + `Checkpointable._gather_saveables_for_checkpoint` for details. + """ + + def __setattr__(self, name, value): + """Support self.foo = checkpointable syntax.""" + # Perform the attribute assignment, and potentially call other __setattr__ + # overrides such as that for tf.keras.Model. + super(Checkpointable, self).__setattr__(name, value) + if isinstance(value, CheckpointableBase): + self._track_checkpointable( + value, name=name, + # Allow the user to switch the Checkpointable which is tracked by this + # name, since assigning a new variable to an attribute has + # historically been fine (e.g. Adam did this). + # TODO(allenl): Should this be a warning once Checkpointable save/load + # is usable? + overwrite=True) diff --git a/tensorflow/python/training/checkpointable_test.py b/tensorflow/python/training/checkpointable_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e79acb49758b6a7d69dd084692d434bea808db64 --- /dev/null +++ b/tensorflow/python/training/checkpointable_test.py @@ -0,0 +1,39 @@ +# 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 + +from tensorflow.python.platform import test +from tensorflow.python.training import checkpointable + + +class InterfaceTests(test.TestCase): + + def testMultipleAssignment(self): + root = checkpointable.Checkpointable() + root.leaf = checkpointable.Checkpointable() + root.leaf = root.leaf + duplicate_name_dep = checkpointable.Checkpointable() + with self.assertRaises(ValueError): + root._track_checkpointable(duplicate_name_dep, name="leaf") + # No error; we're overriding __setattr__, so we can't really stop people + # from doing this while maintaining backward compatibility. + root.leaf = duplicate_name_dep + root._track_checkpointable(duplicate_name_dep, name="leaf", overwrite=True) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/training/checkpointable_utils.py b/tensorflow/python/training/checkpointable_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..32123f87ef2d12497077ab0e2f7d4d4cad1ec5dd --- /dev/null +++ b/tensorflow/python/training/checkpointable_utils.py @@ -0,0 +1,78 @@ +"""Utilities for saving/loading 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 weakref + +from tensorflow.python.framework import ops +from tensorflow.python.training import checkpointable +from tensorflow.python.training import saver as saver_lib + + +class _Checkpoint(object): + """Holds the status of an object-based checkpoint load.""" + + def __init__(self, object_graph_proto, save_path, dtype_map=None): + """Specify the checkpoint being loaded. + + Args: + object_graph_proto: The CheckpointableObjectGraph protocol buffer + associated with this checkpoint. + save_path: A string `Tensor`. The path to the checkpoint, as returned by + `tf.train.latest_checkpoint`. + dtype_map: When executing eagerly, specifies dtypes for creating slot + variables. None when graph building. + """ + self.builder = saver_lib.BulkSaverBuilder() + self.object_graph_proto = object_graph_proto + self.restore_uid = ops.uid() + # Maps from objects to lists of attributes which were in the checkpoint but + # not loaded into any object, for error checking. + self.unused_attributes = weakref.WeakKeyDictionary() + # Dictionary mapping from an id in the protocol buffer flat array to + # Checkpointable Python objects. This mapping may be deferred if a + # checkpoint is restored before all dependencies have been tracked. Uses + # weak references so that partial restorations don't create reference cycles + # (as objects with deferred dependencies will generally have references to + # this object). + self.object_by_proto_id = weakref.WeakValueDictionary() + self.save_path = save_path + self.dtype_map = dtype_map + # When graph building, contains a list of ops to run to restore objects from + # this checkpoint. + self.restore_ops = [] + self.restore_ops_by_name = {} + # A mapping from optimizer proto ids to lists of slot variables to be + # restored when the optimizer is tracked. Only includes slot variables whose + # regular variables have already been created, and only for optimizer + # objects which have not yet been created/tracked. + self.deferred_slot_restorations = {} + # A mapping from variable proto ids to lists of slot variables to be + # restored when the variable is created/tracked. These get shifted over to + # deferred_slot_restorations if the optimizer hasn't been created when that + # happens. + self.slot_restorations = {} + for node_index, node in enumerate(self.object_graph_proto.nodes): + for slot_reference in node.slot_variables: + # `node` refers to an `Optimizer`, since only these have slot variables. + self.slot_restorations.setdefault( + slot_reference.original_variable_node_id, []).append( + checkpointable._SlotVariableRestoration( # pylint: disable=protected-access + optimizer_id=node_index, + slot_variable_id=slot_reference.slot_variable_node_id, + slot_name=slot_reference.slot_name)) diff --git a/tensorflow/python/training/coordinator.py b/tensorflow/python/training/coordinator.py index 0e31255b74f64657cffc4a2f58798835513f0444..0ff97d85e37e6167f1200ba56940f4a663c259a2 100644 --- a/tensorflow/python/training/coordinator.py +++ b/tensorflow/python/training/coordinator.py @@ -27,8 +27,10 @@ import six from tensorflow.python.framework import errors from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.Coordinator") class Coordinator(object): """A coordinator for threads. @@ -406,6 +408,7 @@ class Coordinator(object): # Threads for the standard services. +@tf_export("train.LooperThread") class LooperThread(threading.Thread): """A thread that runs code repeatedly, optionally on a timer. diff --git a/tensorflow/python/training/device_setter.py b/tensorflow/python/training/device_setter.py index 37ab625779f788b1b8e270a15db3244ea6f1bef3..d31c375b4ce48dcb9bc2918514707636a647c675 100644 --- a/tensorflow/python/training/device_setter.py +++ b/tensorflow/python/training/device_setter.py @@ -23,6 +23,16 @@ from tensorflow.core.framework import node_def_pb2 from tensorflow.python.framework import device as pydev from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import server_lib +from tensorflow.python.util.tf_export import tf_export + +# This is a tuple of PS ops used by tf.estimator.Esitmator which should work in +# almost all of cases. +STANDARD_PS_OPS = ( + "Variable", "VariableV2", "AutoReloadVariable", "MutableHashTable", + "MutableHashTableV2", "MutableHashTableOfTensors", + "MutableHashTableOfTensorsV2", "MutableDenseHashTable", + "MutableDenseHashTableV2", "VarHandleOp" +) class _RoundRobinStrategy(object): @@ -121,6 +131,7 @@ class _ReplicaDeviceChooser(object): return worker_device.to_string() +@tf_export("train.replica_device_setter") def replica_device_setter(ps_tasks=0, ps_device="/job:ps", worker_device="/job:worker", merge_devices=True, cluster=None, ps_ops=None, ps_strategy=None): @@ -168,8 +179,7 @@ def replica_device_setter(ps_tasks=0, ps_device="/job:ps", than overriding them. cluster: `ClusterDef` proto or `ClusterSpec`. ps_ops: List of strings representing `Operation` types that need to be - placed on `ps` devices. If `None`, defaults to - `["Variable", "VariableV2", "VarHandleOp"]`. + placed on `ps` devices. If `None`, defaults to `STANDARD_PS_OPS`. ps_strategy: A callable invoked for every ps `Operation` (i.e. matched by `ps_ops`), that takes the `Operation` and returns the ps task index to use. If `None`, defaults to a round-robin strategy across all `ps` @@ -199,7 +209,7 @@ def replica_device_setter(ps_tasks=0, ps_device="/job:ps", if ps_ops is None: # TODO(sherrym): Variables in the LOCAL_VARIABLES collection should not be # placed in the parameter server. - ps_ops = ["Variable", "VariableV2", "VarHandleOp"] + ps_ops = list(STANDARD_PS_OPS) if not merge_devices: logging.warning( diff --git a/tensorflow/python/training/ftrl.py b/tensorflow/python/training/ftrl.py index c64a1b3f799e776c7bbbbcfb691bdd97e4a34466..4fa081fab72df62107cf4957d4ff68240ced9ee0 100644 --- a/tensorflow/python/training/ftrl.py +++ b/tensorflow/python/training/ftrl.py @@ -22,8 +22,10 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.FtrlOptimizer") class FtrlOptimizer(optimizer.Optimizer): """Optimizer that implements the FTRL algorithm. @@ -51,7 +53,7 @@ class FtrlOptimizer(optimizer.Optimizer): learning_rate: A float value or a constant float `Tensor`. learning_rate_power: A float value, must be less or equal to zero. initial_accumulator_value: The starting value for accumulators. - Only positive values are allowed. + Only zero or positive values are allowed. l1_regularization_strength: A float value, must be greater than or equal to zero. l2_regularization_strength: A float value, must be greater than or @@ -82,9 +84,10 @@ class FtrlOptimizer(optimizer.Optimizer): """ super(FtrlOptimizer, self).__init__(use_locking, name) - if initial_accumulator_value <= 0.0: - raise ValueError("initial_accumulator_value %f needs to be positive" % - initial_accumulator_value) + if initial_accumulator_value < 0.0: + raise ValueError( + "initial_accumulator_value %f needs to be be positive or zero" % + initial_accumulator_value) if learning_rate_power > 0.0: raise ValueError("learning_rate_power %f needs to be negative or zero" % learning_rate_power) @@ -265,4 +268,3 @@ class FtrlOptimizer(optimizer.Optimizer): grad.dtype), math_ops.cast(self._learning_rate_power_tensor, grad.dtype), use_locking=self._use_locking) - diff --git a/tensorflow/python/training/gradient_descent.py b/tensorflow/python/training/gradient_descent.py index 5a536e27297f054671e7e44a9e5d20a8b36580b7..6caf29d83af546f821314179e17f7bf1a693ff1a 100644 --- a/tensorflow/python/training/gradient_descent.py +++ b/tensorflow/python/training/gradient_descent.py @@ -18,13 +18,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.GradientDescentOptimizer") class GradientDescentOptimizer(optimizer.Optimizer): """Optimizer that implements the gradient descent algorithm. """ @@ -41,6 +44,7 @@ class GradientDescentOptimizer(optimizer.Optimizer): """ super(GradientDescentOptimizer, self).__init__(use_locking, name) self._learning_rate = learning_rate + self._learning_rate_tensor = None def _apply_dense(self, grad, var): return training_ops.apply_gradient_descent( @@ -67,5 +71,6 @@ class GradientDescentOptimizer(optimizer.Optimizer): return var.scatter_sub(delta, use_locking=self._use_locking) def _prepare(self): - self._learning_rate_tensor = ops.convert_to_tensor(self._learning_rate, - name="learning_rate") + if not context.executing_eagerly() or self._learning_rate_tensor is None: + self._learning_rate_tensor = ops.convert_to_tensor(self._learning_rate, + name="learning_rate") diff --git a/tensorflow/python/training/input.py b/tensorflow/python/training/input.py index 331a51e8bc848917967fed06632fe0d1c5bcad9c..44f00a96deff64012705c4c81b185a9c4fac2295 100644 --- a/tensorflow/python/training/input.py +++ b/tensorflow/python/training/input.py @@ -44,6 +44,7 @@ from tensorflow.python.ops import sparse_ops from tensorflow.python.ops import variable_scope as vs from tensorflow.python.summary import summary from tensorflow.python.training import queue_runner +from tensorflow.python.util.tf_export import tf_export # pylint: disable=protected-access @@ -53,9 +54,12 @@ _restore_sparse = sparse_ops._take_many_sparse_from_tensors_map # pylint: enable=protected-access +@tf_export("train.match_filenames_once") def match_filenames_once(pattern, name=None): """Save the list of files matching pattern, so it is only computed once. + NOTE: The order of the files returned can be non-deterministic. + Args: pattern: A file pattern (glob), or 1D tensor of file patterns. name: A name for the operations (optional). @@ -70,6 +74,7 @@ def match_filenames_once(pattern, name=None): collections=[ops.GraphKeys.LOCAL_VARIABLES]) +@tf_export("train.limit_epochs") def limit_epochs(tensor, num_epochs=None, name=None): """Returns tensor `num_epochs` times and then raises an `OutOfRange` error. @@ -102,6 +107,7 @@ def limit_epochs(tensor, num_epochs=None, name=None): return array_ops.identity(tensor, name=name) +@tf_export("train.input_producer") def input_producer(input_tensor, element_shape=None, num_epochs=None, @@ -153,7 +159,7 @@ def input_producer(input_tensor, enabled. Please use the `tf.data` API to ingest data under eager execution. @end_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "Input pipelines based on Queues are not supported when eager execution" " is enabled. Please use tf.data to ingest data into your model" @@ -184,6 +190,7 @@ def input_producer(input_tensor, return q +@tf_export("train.string_input_producer") def string_input_producer(string_tensor, num_epochs=None, shuffle=True, @@ -253,6 +260,7 @@ def string_input_producer(string_tensor, cancel_op=cancel_op) +@tf_export("train.range_input_producer") def range_input_producer(limit, num_epochs=None, shuffle=True, seed=None, capacity=32, shared_name=None, name=None): """Produces the integers from 0 to limit-1 in a queue. @@ -290,6 +298,7 @@ def range_input_producer(limit, num_epochs=None, shuffle=True, seed=None, shared_name, "fraction_of_%d_full" % capacity, name) +@tf_export("train.slice_input_producer") def slice_input_producer(tensor_list, num_epochs=None, shuffle=True, seed=None, capacity=32, shared_name=None, name=None): """Produces a slice of each `Tensor` in `tensor_list`. @@ -728,7 +737,7 @@ def _batch(tensors, batch_size, keep_input, num_threads=1, capacity=32, allow_smaller_final_batch=False, shared_name=None, name=None): """Helper function for `batch` and `maybe_batch`.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError( "Input pipelines based on Queues are not supported when eager execution" " is enabled. Please use tf.data to ingest data into your model" @@ -766,7 +775,7 @@ def _batch_join(tensors_list, batch_size, keep_input, capacity=32, enqueue_many=False, shapes=None, dynamic_pad=False, allow_smaller_final_batch=False, shared_name=None, name=None): """Helper function for `batch_join` and `maybe_batch_join`.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError( "Input pipelines based on Queues are not supported when eager execution" " is enabled. Please use tf.data to ingest data into your model" @@ -801,7 +810,7 @@ def _shuffle_batch(tensors, batch_size, capacity, min_after_dequeue, shapes=None, allow_smaller_final_batch=False, shared_name=None, name=None): """Helper function for `shuffle_batch` and `maybe_shuffle_batch`.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError( "Input pipelines based on Queues are not supported when eager execution" " is enabled. Please use tf.data to ingest data into your model" @@ -846,7 +855,7 @@ def _shuffle_batch_join(tensors_list, batch_size, capacity, allow_smaller_final_batch=False, shared_name=None, name=None): """Helper function for `shuffle_batch_join` and `maybe_shuffle_batch_join`.""" - if context.in_eager_mode(): + if context.executing_eagerly(): raise ValueError( "Input pipelines based on Queues are not supported when eager execution" " is enabled. Please use tf.data to ingest data into your model" @@ -885,6 +894,7 @@ def _shuffle_batch_join(tensors_list, batch_size, capacity, # Batching functions ---------------------------------------------------------- +@tf_export("train.batch") def batch(tensors, batch_size, num_threads=1, capacity=32, enqueue_many=False, shapes=None, dynamic_pad=False, allow_smaller_final_batch=False, shared_name=None, name=None): @@ -979,6 +989,7 @@ def batch(tensors, batch_size, num_threads=1, capacity=32, name=name) +@tf_export("train.maybe_batch") def maybe_batch(tensors, keep_input, batch_size, num_threads=1, capacity=32, enqueue_many=False, shapes=None, dynamic_pad=False, allow_smaller_final_batch=False, shared_name=None, name=None): @@ -1031,6 +1042,7 @@ def maybe_batch(tensors, keep_input, batch_size, num_threads=1, capacity=32, name=name) +@tf_export("train.batch_join") def batch_join(tensors_list, batch_size, capacity=32, enqueue_many=False, shapes=None, dynamic_pad=False, allow_smaller_final_batch=False, shared_name=None, name=None): @@ -1136,6 +1148,7 @@ def batch_join(tensors_list, batch_size, capacity=32, enqueue_many=False, name=name) +@tf_export("train.maybe_batch_join") def maybe_batch_join(tensors_list, keep_input, batch_size, capacity=32, enqueue_many=False, shapes=None, dynamic_pad=False, allow_smaller_final_batch=False, shared_name=None, @@ -1188,6 +1201,7 @@ def maybe_batch_join(tensors_list, keep_input, batch_size, capacity=32, name=name) +@tf_export("train.shuffle_batch") def shuffle_batch(tensors, batch_size, capacity, min_after_dequeue, num_threads=1, seed=None, enqueue_many=False, shapes=None, allow_smaller_final_batch=False, shared_name=None, name=None): @@ -1287,6 +1301,7 @@ def shuffle_batch(tensors, batch_size, capacity, min_after_dequeue, name=name) +@tf_export("train.maybe_shuffle_batch") def maybe_shuffle_batch(tensors, batch_size, capacity, min_after_dequeue, keep_input, num_threads=1, seed=None, enqueue_many=False, shapes=None, @@ -1346,6 +1361,7 @@ def maybe_shuffle_batch(tensors, batch_size, capacity, min_after_dequeue, name=name) +@tf_export("train.shuffle_batch_join") def shuffle_batch_join(tensors_list, batch_size, capacity, min_after_dequeue, seed=None, enqueue_many=False, shapes=None, allow_smaller_final_batch=False, @@ -1439,6 +1455,7 @@ def shuffle_batch_join(tensors_list, batch_size, capacity, name=name) +@tf_export("train.maybe_shuffle_batch_join") def maybe_shuffle_batch_join(tensors_list, batch_size, capacity, min_after_dequeue, keep_input, seed=None, enqueue_many=False, shapes=None, diff --git a/tensorflow/python/training/learning_rate_decay.py b/tensorflow/python/training/learning_rate_decay.py index 343a49cded01ec3665170c85767bcc812531df78..10ab4c1137ff226d88902143d4f2281ad77de531 100644 --- a/tensorflow/python/training/learning_rate_decay.py +++ b/tensorflow/python/training/learning_rate_decay.py @@ -25,8 +25,10 @@ from tensorflow.python.framework import 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.util.tf_export import tf_export +@tf_export("train.exponential_decay") def exponential_decay(learning_rate, global_step, decay_steps, @@ -103,6 +105,7 @@ def exponential_decay(learning_rate, learning_rate, math_ops.pow(decay_rate, p), name=name) +@tf_export("train.piecewise_constant") def piecewise_constant(x, boundaries, values, name=None): """Piecewise constant from boundaries and interval values. @@ -182,6 +185,7 @@ def piecewise_constant(x, boundaries, values, name=None): return control_flow_ops.case(pred_fn_pairs, default, exclusive=True) +@tf_export("train.polynomial_decay") def polynomial_decay(learning_rate, global_step, decay_steps, @@ -291,6 +295,7 @@ def polynomial_decay(learning_rate, name=name) +@tf_export("train.natural_exp_decay") def natural_exp_decay(learning_rate, global_step, decay_steps, @@ -362,6 +367,7 @@ def natural_exp_decay(learning_rate, return math_ops.multiply(learning_rate, exponent, name=name) +@tf_export("train.inverse_time_decay") def inverse_time_decay(learning_rate, global_step, decay_steps, @@ -444,6 +450,7 @@ def inverse_time_decay(learning_rate, return math_ops.div(learning_rate, denom, name=name) +@tf_export("train.cosine_decay") def cosine_decay(learning_rate, global_step, decay_steps, alpha=0.0, name=None): """Applies cosine decay to the learning rate. @@ -503,6 +510,7 @@ def cosine_decay(learning_rate, global_step, decay_steps, alpha=0.0, name=None): return math_ops.multiply(learning_rate, decayed) +@tf_export("train.cosine_decay_restarts") def cosine_decay_restarts(learning_rate, global_step, first_decay_steps, @@ -596,6 +604,7 @@ def cosine_decay_restarts(learning_rate, return math_ops.multiply(learning_rate, decayed, name=name) +@tf_export("train.linear_cosine_decay") def linear_cosine_decay(learning_rate, global_step, decay_steps, @@ -679,6 +688,7 @@ def linear_cosine_decay(learning_rate, return math_ops.multiply(learning_rate, linear_cosine_decayed, name=name) +@tf_export("train.noisy_linear_cosine_decay") def noisy_linear_cosine_decay(learning_rate, global_step, decay_steps, diff --git a/tensorflow/python/training/learning_rate_decay_test.py b/tensorflow/python/training/learning_rate_decay_test.py index 1ce8c156a0b126f680bad62267f90e31a23febed..60306e4f1239a759ea1f68492a1211d5f0858997 100644 --- a/tensorflow/python/training/learning_rate_decay_test.py +++ b/tensorflow/python/training/learning_rate_decay_test.py @@ -43,8 +43,8 @@ class LRDecayTest(test_util.TensorFlowTestCase): def testStaircase(self): with self.test_session(): - step = gen_state_ops._variable(shape=[], dtype=dtypes.int32, - name="step", container="", shared_name="") + step = gen_state_ops.variable(shape=[], dtype=dtypes.int32, + name="step", container="", shared_name="") assign_100 = state_ops.assign(step, 100) assign_1 = state_ops.assign(step, 1) assign_2 = state_ops.assign(step, 2) @@ -113,7 +113,7 @@ class LRDecayTest(test_util.TensorFlowTestCase): learning_rate_decay.piecewise_constant(x, boundaries, values) # Test that ref types are valid. - if context.in_graph_mode(): + if not context.executing_eagerly(): x = variables.Variable(0.0) x_ref = x.op.outputs[0] # float32_ref tensor should be accepted boundaries, values = [1.0, 2.0], [1, 2, 3] @@ -264,8 +264,8 @@ class ExponentialDecayTest(test_util.TensorFlowTestCase): initial_lr = 0.1 k = 10 decay_rate = 0.96 - step = gen_state_ops._variable(shape=[], dtype=dtypes.int32, - name="step", container="", shared_name="") + step = gen_state_ops.variable( + shape=[], dtype=dtypes.int32, name="step", container="", shared_name="") assign_step = state_ops.assign(step, 0) increment_step = state_ops.assign_add(step, 1) decayed_lr = learning_rate_decay.natural_exp_decay(initial_lr, step, @@ -281,8 +281,8 @@ class ExponentialDecayTest(test_util.TensorFlowTestCase): initial_lr = 0.1 k = 10 decay_rate = 0.96 - step = gen_state_ops._variable(shape=[], dtype=dtypes.int32, - name="step", container="", shared_name="") + step = gen_state_ops.variable( + shape=[], dtype=dtypes.int32, name="step", container="", shared_name="") assign_step = state_ops.assign(step, 0) increment_step = state_ops.assign_add(step, 1) decayed_lr = learning_rate_decay.natural_exp_decay(initial_lr, @@ -304,8 +304,8 @@ class InverseDecayTest(test_util.TensorFlowTestCase): initial_lr = 0.1 k = 10 decay_rate = 0.96 - step = gen_state_ops._variable(shape=[], dtype=dtypes.int32, - name="step", container="", shared_name="") + step = gen_state_ops.variable( + shape=[], dtype=dtypes.int32, name="step", container="", shared_name="") assign_step = state_ops.assign(step, 0) increment_step = state_ops.assign_add(step, 1) decayed_lr = learning_rate_decay.inverse_time_decay(initial_lr, @@ -323,8 +323,8 @@ class InverseDecayTest(test_util.TensorFlowTestCase): initial_lr = 0.1 k = 10 decay_rate = 0.96 - step = gen_state_ops._variable(shape=[], dtype=dtypes.int32, - name="step", container="", shared_name="") + step = gen_state_ops.variable( + shape=[], dtype=dtypes.int32, name="step", container="", shared_name="") assign_step = state_ops.assign(step, 0) increment_step = state_ops.assign_add(step, 1) decayed_lr = learning_rate_decay.inverse_time_decay(initial_lr, diff --git a/tensorflow/python/training/momentum.py b/tensorflow/python/training/momentum.py index cf9530d87c46783b517884610b644b076bef6807..bd9fa79d8feac68c149f787ee8501bdddb173d33 100644 --- a/tensorflow/python/training/momentum.py +++ b/tensorflow/python/training/momentum.py @@ -22,8 +22,10 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.MomentumOptimizer") class MomentumOptimizer(optimizer.Optimizer): """Optimizer that implements the Momentum algorithm. diff --git a/tensorflow/python/training/momentum_test.py b/tensorflow/python/training/momentum_test.py index 6865513b0e4aad18d77887770a11243642958e7a..297a8bbde5447cff9465be36c0bb71f2490c60fc 100644 --- a/tensorflow/python/training/momentum_test.py +++ b/tensorflow/python/training/momentum_test.py @@ -66,7 +66,7 @@ class MomentumOptimizerTest(test.TestCase): mom_update = mom_opt.apply_gradients( zip([grads0, grads1], [var0, var1])) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(variables.global_variables_initializer()) # Fetch params to validate initial values self.assertAllClose([1.0, 2.0], self.evaluate(var0)) @@ -78,13 +78,13 @@ class MomentumOptimizerTest(test.TestCase): self.assertEquals(slot0.get_shape(), var0.get_shape()) slot1 = mom_opt.get_slot(var1, "momentum") self.assertEquals(slot1.get_shape(), var1.get_shape()) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertFalse(slot0 in variables.trainable_variables()) self.assertFalse(slot1 in variables.trainable_variables()) # Step 1: the momentum accumulators where 0. So we should see a normal # update: v -= grad * learning_rate - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(mom_update) # Check that the momentum accumulators have been updated. self.assertAllCloseAccordingToType(np.array([0.1, 0.1]), @@ -99,10 +99,10 @@ class MomentumOptimizerTest(test.TestCase): np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]), self.evaluate(var1)) # Step 2: the momentum accumulators contain the previous update. - if context.in_graph_mode(): - self.evaluate(mom_update) - else: + if context.executing_eagerly(): mom_opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + else: + self.evaluate(mom_update) # Check that the momentum accumulators have been updated. self.assertAllCloseAccordingToType( np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]), @@ -142,7 +142,7 @@ class MomentumOptimizerTest(test.TestCase): [1.0, 2.0], dtype=dtypes.float32, name="var0") var1 = resource_variable_ops.ResourceVariable( [3.0, 4.0], dtype=dtypes.float32, name="var1") - if context.in_eager_mode(): + if context.executing_eagerly(): loss = lambda: math_ops.reduce_sum(var0 + var1) else: loss = math_ops.reduce_sum(var0 + var1) @@ -157,7 +157,7 @@ class MomentumOptimizerTest(test.TestCase): [1.0, 2.0], dtype=dtypes.float32, name="var2") var3 = resource_variable_ops.ResourceVariable( [3.0, 4.0], dtype=dtypes.float32, name="var3") - if context.in_eager_mode(): + if context.executing_eagerly(): loss = lambda: math_ops.reduce_sum(var2 + var3) else: loss = math_ops.reduce_sum(var2 + var3) @@ -247,7 +247,7 @@ class MomentumOptimizerTest(test.TestCase): # pylint: enable=cell-var-from-loop opt = momentum_lib.MomentumOptimizer(learning_rate=1.0, momentum=0.0) - sgd_op = opt.minimize(loss if context.in_eager_mode() else loss()) + sgd_op = opt.minimize(loss) self.evaluate(variables.global_variables_initializer()) # Run 1 step of sgd self.evaluate(sgd_op) @@ -262,7 +262,7 @@ class MomentumOptimizerTest(test.TestCase): return math_ops.reduce_sum(embedding_ops.embedding_lookup(var0, [[1]])) opt = momentum_lib.MomentumOptimizer(learning_rate=1.0, momentum=0.0) - sgd_op = opt.minimize(loss if context.in_eager_mode() else loss()) + sgd_op = opt.minimize(loss) self.evaluate(variables.global_variables_initializer()) self.evaluate(sgd_op) self.assertAllCloseAccordingToType([[1, 1], [0, 0]], self.evaluate(var0)) diff --git a/tensorflow/python/training/monitored_session.py b/tensorflow/python/training/monitored_session.py index fa3517db27be4581deb85f77f022406b8b30ec56..6c5c9e01a76d539b550420134b09090b89beed46 100644 --- a/tensorflow/python/training/monitored_session.py +++ b/tensorflow/python/training/monitored_session.py @@ -41,6 +41,7 @@ from tensorflow.python.training import queue_runner from tensorflow.python.training import saver as training_saver from tensorflow.python.training import session_manager as sm from tensorflow.python.training import session_run_hook +from tensorflow.python.util.tf_export import tf_export # The list of exceptions that we should recover from. Exceptions not in this @@ -52,6 +53,7 @@ _PREEMPTION_ERRORS = (errors.AbortedError, errors.UnavailableError) USE_DEFAULT = object() +@tf_export('train.Scaffold') class Scaffold(object): """Structure to create or gather pieces commonly needed to train a model. @@ -272,6 +274,7 @@ class Scaffold(object): resources.initialize_resources(resources.local_resources())) +@tf_export('train.MonitoredTrainingSession') def MonitoredTrainingSession(master='', # pylint: disable=invalid-name is_chief=True, checkpoint_dir=None, @@ -381,6 +384,7 @@ def MonitoredTrainingSession(master='', # pylint: disable=invalid-name stop_grace_period_secs=stop_grace_period_secs) +@tf_export('train.SessionCreator') class SessionCreator(object): """A factory for tf.Session.""" @@ -390,6 +394,7 @@ class SessionCreator(object): 'create_session is not implemented for {}.'.format(self)) +@tf_export('train.ChiefSessionCreator') class ChiefSessionCreator(SessionCreator): """Creates a tf.Session for a chief.""" @@ -441,6 +446,7 @@ class ChiefSessionCreator(SessionCreator): init_fn=self._scaffold.init_fn) +@tf_export('train.WorkerSessionCreator') class WorkerSessionCreator(SessionCreator): """Creates a tf.Session for a worker.""" @@ -706,6 +712,7 @@ class _MonitoredSession(object): return self._coordinated_creator.tf_sess +@tf_export('train.MonitoredSession') class MonitoredSession(_MonitoredSession): """Session-like object that handles initialization, recovery and hooks. @@ -788,6 +795,7 @@ class MonitoredSession(_MonitoredSession): stop_grace_period_secs=stop_grace_period_secs) +@tf_export('train.SingularMonitoredSession') class SingularMonitoredSession(_MonitoredSession): """Session-like object that handles initialization, restoring, and hooks. diff --git a/tensorflow/python/training/moving_averages.py b/tensorflow/python/training/moving_averages.py index 43ed1ac170d0d993bf7b5bcaff3ff6a8cbbde6b2..b9ecb27df19d051c28ec1c3fe3cd9fd86717a5ed 100644 --- a/tensorflow/python/training/moving_averages.py +++ b/tensorflow/python/training/moving_averages.py @@ -26,6 +26,7 @@ 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 slot_creator +from tensorflow.python.util.tf_export import tf_export # TODO(touts): switch to variables.Variable. @@ -230,6 +231,7 @@ def _zero_debias(unbiased_var, value, decay): return unbiased_ema_delta +@tf_export("train.ExponentialMovingAverage") class ExponentialMovingAverage(object): """Maintains moving averages of variables by employing an exponential decay. @@ -398,7 +400,9 @@ class ExponentialMovingAverage(object): avg = slot_creator.create_zeros_slot( var, self._name, - colocate_with_primary=(var.op.type in ["Variable", "VariableV2"])) + colocate_with_primary=(var.op.type in ["Variable", + "VariableV2", + "VarHandleOp"])) if self._zero_debias: zero_debias_true.add(avg) self._averages[var] = avg diff --git a/tensorflow/python/training/moving_averages_test.py b/tensorflow/python/training/moving_averages_test.py index 6efdeb286657e761a4c46634b9408121765a447b..6717811bbb0f05723a5ad0fbcbfba75249d0d43b 100644 --- a/tensorflow/python/training/moving_averages_test.py +++ b/tensorflow/python/training/moving_averages_test.py @@ -376,7 +376,7 @@ class ExponentialMovingAverageTest(test.TestCase): with ops.device("/job:dev_v0"): v0 = variables.Variable(10.0, name="v0") with ops.device("/job:dev_v1"): - v1 = gen_state_ops._variable( + v1 = gen_state_ops.variable( shape=[1], dtype=dtypes.float32, name="v1", diff --git a/tensorflow/python/training/optimizer.py b/tensorflow/python/training/optimizer.py index a06b3eada6cf2f0e4b5690ebe0c92e60f5d2ec0e..af9cc3491c2d80c95428c70e33c684c5a7d9e088 100644 --- a/tensorflow/python/training/optimizer.py +++ b/tensorflow/python/training/optimizer.py @@ -34,21 +34,10 @@ 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 checkpointable from tensorflow.python.training import slot_creator from tensorflow.python.util import nest - - -def _get_variable_for(v): - """Returns the ResourceVariable responsible for v, or v if not necessary.""" - if context.in_eager_mode(): - return v - if v.op.type == "VarHandleOp": - for var in variables.trainable_variables(): - if (isinstance(var, resource_variable_ops.ResourceVariable) - and var.handle.op is v.op): - return var - raise ValueError("Got %s but could not locate source variable." % (str(v))) - return v +from tensorflow.python.util.tf_export import tf_export def _deduplicate_indexed_slices(values, indices): @@ -71,8 +60,8 @@ def _deduplicate_indexed_slices(values, indices): def _var_key(var): - if context.in_eager_mode(): - return var._shared_name # pylint: disable=protected-access + if context.executing_eagerly(): + return var._unique_id # pylint: disable=protected-access return (var.op.graph, var.op.name) @@ -96,6 +85,9 @@ class _RefVariableProcessor(_OptimizableVariable): def __init__(self, v): self._v = v + def __str__(self): + return "<_RefVariableProcessor(%s)>" % self._v + def target(self): return self._v._ref() # pylint: disable=protected-access @@ -174,20 +166,48 @@ class _StreamingModelPortProcessor(_OptimizableVariable): return g +class _TensorProcessor(_OptimizableVariable): + """Processor for ordinary Tensors. + + Even though a Tensor can't really be updated, sometimes it is useful to + compute the gradients with respect to a Tensor using the optimizer. Updating + the Tensor is, of course, unsupported. + """ + + def __init__(self, v): + self._v = v + + def target(self): + return self._v + + def update_op(self, optimizer, g): + raise NotImplementedError("Trying to update a Tensor ", self._v) + + def _get_processor(v): """The processor of v.""" - if context.in_eager_mode(): - return _DenseResourceVariableProcessor(v) + if context.executing_eagerly(): + if isinstance(v, ops.Tensor): + return _TensorProcessor(v) + else: + return _DenseResourceVariableProcessor(v) if v.op.type == "VarHandleOp": return _DenseResourceVariableProcessor(v) if isinstance(v, variables.Variable): return _RefVariableProcessor(v) if v.op.type == "SubmodelPort": return _StreamingModelPortProcessor(v) + if isinstance(v, ops.Tensor): + return _TensorProcessor(v) raise NotImplementedError("Trying to optimize unsupported type ", v) -class Optimizer(object): +@tf_export("train.Optimizer") +class Optimizer( + # Optimizers inherit from CheckpointableBase rather than Checkpointable + # since they do most of their dependency management themselves (slot + # variables are special-cased, and non-slot variables are keyed to graphs). + checkpointable.CheckpointableBase): """Base class for optimizers. This class defines the API to add Ops to train a model. You never use this @@ -298,9 +318,18 @@ class Optimizer(object): self._use_locking = use_locking self._name = name # Dictionary of slots. - # {slot_name : { variable_to_train: slot_for_the_variable, ...}, ... } + # {slot_name : + # {_var_key(variable_to_train): slot_for_the_variable, ... }, + # ... } self._slots = {} self._non_slot_dict = {} + # For implementing Checkpointable. Stores information about how to restore + # slot variables which have not yet been created + # (checkpointable._CheckpointPosition objects). + # {slot_name : + # {_var_key(variable_to_train): [checkpoint_position, ... ], ... }, + # ... } + self._deferred_slot_restorations = {} def get_name(self): return self._name @@ -380,7 +409,9 @@ class Optimizer(object): given variable. Args: - loss: A Tensor containing the value to minimize. + loss: A Tensor containing the value to minimize or a callable taking + no arguments which returns the value to minimize. When eager execution + is enabled it must be a callable. var_list: Optional list or tuple of `tf.Variable` to update to minimize `loss`. Defaults to the list of variables collected in the graph under the key `GraphKeys.TRAINABLE_VARIABLES`. @@ -399,37 +430,27 @@ class Optimizer(object): Raises: TypeError: If `var_list` contains anything else than `Variable` objects. ValueError: If some arguments are invalid. - RuntimeError: If called with eager execution enabled and if `grad_loss` - is not `None` or `loss` is not callable. + RuntimeError: If called with eager execution enabled and `loss` is + not callable. @compatibility(eager) - When eager execution is enabled, `loss` should be a Python function that - takes elements of `var_list` as arguments and computes the value to be - minimized. If `var_list` is None, `loss` should take no arguments. - Gradient computation is done with respect to the elements of `var_list` if - not None, else with respect to any trainable variables created during the - execution of the `loss` function. - `gate_gradients`, `aggregation_method`, `colocate_gradients_with_ops` and - `grad_loss` are ignored when eager execution is enabled. + When eager execution is enabled, `gate_gradients`, `aggregation_method`, + and `colocate_gradients_with_ops` are ignored. @end_compatibility """ - if context.in_eager_mode(): - if grad_loss is not None: - raise RuntimeError( - "`grad_loss` argument to Optimizer.compute_gradients " - "not supported when eager execution is enabled.") - if not callable(loss): - raise RuntimeError( - "`loss` passed to Optimizer.compute_gradients should " - "be a function when eager execution is enabled.") - # TODO(agarwal): consider passing parameters to the `loss` function. + if callable(loss): + with backprop.GradientTape() as tape: + if var_list is not None: + tape.watch(var_list) + loss_value = loss() if var_list is None: - return backprop.implicit_grad(loss)() - else: - var_list = nest.flatten(var_list) - grads = backprop.gradients_function(loss)(*var_list) - grads_and_vars = list(zip(grads, var_list)) - return grads_and_vars + var_list = tape.watched_variables() + grads = tape.gradient(loss_value, var_list, grad_loss) + return list(zip(grads, var_list)) + if context.executing_eagerly(): + raise RuntimeError( + "`loss` passed to Optimizer.compute_gradients should " + "be a function when eager execution is enabled.") if gate_gradients not in [Optimizer.GATE_NONE, Optimizer.GATE_OP, Optimizer.GATE_GRAPH]: raise ValueError("gate_gradients must be one of: Optimizer.GATE_NONE, " @@ -515,7 +536,7 @@ class Optimizer(object): raise ValueError("No gradients provided for any variable: %s." % ([str(v) for _, _, v in converted_grads_and_vars],)) with ops.init_scope(): - self._create_slots([_get_variable_for(v) for v in var_list]) + self._create_slots(var_list) update_ops = [] with ops.name_scope(name, self._name) as name: self._prepare() @@ -525,7 +546,7 @@ class Optimizer(object): # We colocate all ops created in _apply_dense or _apply_sparse # on the same device as the variable. # TODO(apassos): figure out how to get the variable name here. - scope_name = var.op.name if context.in_graph_mode() else "" + scope_name = "" if context.executing_eagerly() else var.op.name with ops.name_scope("update_" + scope_name), ops.colocate_with(var): update_ops.append(processor.update_op(self, grad)) if global_step is None: @@ -543,7 +564,7 @@ class Optimizer(object): else: apply_updates = state_ops.assign_add(global_step, 1, name=name) - if context.in_graph_mode(): + if not context.executing_eagerly(): if isinstance(apply_updates, ops.Tensor): apply_updates = apply_updates.op train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP) @@ -593,14 +614,14 @@ class Optimizer(object): Returns: A list of variables. """ - executing_eagerly = context.in_eager_mode() + executing_eagerly = context.executing_eagerly() current_graph = ops.get_default_graph() def _from_current_graph(variable): if executing_eagerly: # No variable.op in eager mode. We don't expect lots of eager graphs, # but behavior should be consistent with graph mode. - return variable._container_prefix == current_graph._container_prefix # pylint: disable=protected-access + return variable._graph_key == current_graph._graph_key # pylint: disable=protected-access else: return variable.op.graph is current_graph @@ -615,20 +636,54 @@ class Optimizer(object): 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 + eager = context.executing_eagerly() + graph = None if eager else colocate_with.graph key = (name, graph) v = self._non_slot_dict.get(key, None) if v is None: + self._maybe_initialize_checkpointable() with ops.colocate_with(colocate_with): + if eager: + restored_initial_value = self._preload_simple_restoration( + name=name, shape=None) + if restored_initial_value is not None: + initial_value = restored_initial_value v = variable_scope.variable(initial_value, name=name, trainable=False) + # Restore this variable by name if necessary, but don't add a + # Checkpointable dependency. Optimizers return the current graph's + # non-slot variables from _checkpoint_dependencies explicitly rather + # than unconditionally adding dependencies (since there may be multiple + # non-slot variables with the same name in different graphs, trying to + # save all of them would result in errors). + self._handle_deferred_dependencies(name=name, checkpointable=v) self._non_slot_dict[key] = v return v + @property + def _checkpoint_dependencies(self): + """From Checkpointable. Gather graph-specific non-slot variables to save.""" + current_graph_non_slot_variables = [] + current_graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access + for (name, _), variable_object in sorted(self._non_slot_dict.items(), + # Avoid comparing graphs + key=lambda item: item[0][0]): + if variable_object._graph_key == current_graph_key: # pylint: disable=protected-access + current_graph_non_slot_variables.append( + checkpointable.CheckpointableReference( + name=name, ref=variable_object)) + return (super(Optimizer, self)._checkpoint_dependencies + + current_graph_non_slot_variables) + + def _lookup_dependency(self, name): + """From Checkpointable. Find a non-slot variable in the current graph.""" + unconditional = super(Optimizer, self)._lookup_dependency(name) + if unconditional is not None: + return unconditional + graph = None if context.executing_eagerly() else ops.get_default_graph() + return self._get_non_slot_variable(name, graph=graph) + def _get_non_slot_variable(self, name, graph=None): return self._non_slot_dict.get((name, graph), None) @@ -866,7 +921,11 @@ class Optimizer(object): """ named_slots = self._slot_dict(slot_name) if _var_key(var) not in named_slots: - named_slots[_var_key(var)] = slot_creator.create_slot(var, val, op_name) + new_slot_variable = slot_creator.create_slot(var, val, op_name) + self._restore_slot_variable( + slot_name=slot_name, variable=var, + slot_variable=new_slot_variable) + named_slots[_var_key(var)] = new_slot_variable return named_slots[_var_key(var)] def _get_or_make_slot_with_initializer(self, var, initializer, shape, dtype, @@ -887,8 +946,12 @@ class Optimizer(object): """ named_slots = self._slot_dict(slot_name) if _var_key(var) not in named_slots: - named_slots[_var_key(var)] = slot_creator.create_slot_with_initializer( + new_slot_variable = slot_creator.create_slot_with_initializer( var, initializer, shape, dtype, op_name) + self._restore_slot_variable( + slot_name=slot_name, variable=var, + slot_variable=new_slot_variable) + named_slots[_var_key(var)] = new_slot_variable return named_slots[_var_key(var)] def _zeros_slot(self, var, slot_name, op_name): @@ -905,5 +968,78 @@ class Optimizer(object): """ named_slots = self._slot_dict(slot_name) if _var_key(var) not in named_slots: - named_slots[_var_key(var)] = slot_creator.create_zeros_slot(var, op_name) + new_slot_variable = slot_creator.create_zeros_slot(var, op_name) + self._restore_slot_variable( + slot_name=slot_name, variable=var, + slot_variable=new_slot_variable) + named_slots[_var_key(var)] = new_slot_variable return named_slots[_var_key(var)] + + # -------------- + # For implementing the Checkpointable interface. + # -------------- + + def _restore_slot_variable(self, slot_name, variable, slot_variable): + """Restore a newly created slot variable's value.""" + variable_key = _var_key(variable) + deferred_restorations = self._deferred_slot_restorations.get( + slot_name, {}).pop(variable_key, []) + # Iterate over restores, highest restore UID first to minimize the number + # of assignments. + deferred_restorations.sort(key=lambda position: position.restore_uid, + reverse=True) + for checkpoint_position in deferred_restorations: + checkpoint_position.restore(slot_variable) + + def _create_or_restore_slot_variable( + self, slot_variable_position, slot_name, variable): + """Restore a slot variable's value, possibly creating it. + + Called when a variable which has an associated slot variable is created or + restored. When executing eagerly, we create the slot variable with a + restoring initializer. + + No new variables are created when graph building. Instead, + _restore_slot_variable catches these after normal creation and adds restore + ops to the graph. This method is nonetheless important when graph building + for the case when a slot variable has already been created but `variable` + has just been added to a dependency graph (causing us to realize that the + slot variable needs to be restored). + + Args: + slot_variable_position: A `checkpointable._CheckpointPosition` object + indicating the slot variable `Checkpointable` object to be restored. + slot_name: The name of this `Optimizer`'s slot to restore into. + variable: The variable object this slot is being created for. + """ + named_slots = self._slot_dict(slot_name) + variable_key = _var_key(variable) + slot_variable = named_slots.get(variable_key, None) + if (slot_variable is None and context.executing_eagerly() and + slot_variable_position.is_simple_variable()): + initializer = checkpointable.CheckpointInitialValue( + checkpoint_position=slot_variable_position) + slot_variable = self._get_or_make_slot( + var=variable, + val=initializer, + slot_name=slot_name, + op_name=self._name) + # Slot variables are not owned by any one object (because we don't want to + # save the slot variable if the optimizer is saved without the non-slot + # variable, or if the non-slot variable is saved without the optimizer; + # it's a dependency hypergraph with edges of the form (optimizer, non-slot + # variable, variable)). So we don't _track_ slot variables anywhere, and + # instead special-case this dependency and otherwise pretend it's a normal + # graph. + if slot_variable is not None: + # If we've either made this slot variable, or if we've pulled out an + # existing slot variable, we should restore it. + slot_variable_position.restore(slot_variable) + else: + # We didn't make the slot variable. Defer restoring until it gets created + # normally. We keep a list rather than the one with the highest restore + # UID in case slot variables have their own dependencies, in which case + # those could differ between restores. + self._deferred_slot_restorations.setdefault( + slot_name, {}).setdefault(variable_key, []).append( + slot_variable_position) diff --git a/tensorflow/python/training/optimizer_test.py b/tensorflow/python/training/optimizer_test.py index 6bdae39073d48e0bd8b757a2d5145480e92d185f..0cab6410e83ca1880a0a4a80d2cfa5c17517af95 100644 --- a/tensorflow/python/training/optimizer_test.py +++ b/tensorflow/python/training/optimizer_test.py @@ -18,7 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -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 @@ -44,11 +43,10 @@ class OptimizerTest(test.TestCase): name='a_%d' % i) var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype, name='b_%d' % i) - def loss(v0, v1): - return 5 * v0 + 3 * v1 + def loss(): + return 5 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop # Note that for eager execution, minimize expects a function instead of a # Tensor. - cost = loss if context.in_eager_mode() else loss(var0, var1) global_step = resource_variable_ops.ResourceVariable( array_ops.zeros([], dtypes.int64), name='global_step_%d' % i) sgd_op = gradient_descent.GradientDescentOptimizer(3.0) @@ -58,7 +56,7 @@ class OptimizerTest(test.TestCase): self.assertAllClose([1.0, 2.0], self.evaluate(var0)) self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 1 step of sgd through optimizer - opt_op = sgd_op.minimize(cost, global_step, [var0, var1]) + opt_op = sgd_op.minimize(loss, global_step, [var0, var1]) self.evaluate(opt_op) # Validate updated params self.assertAllClose([-14., -13.], self.evaluate(var0)) @@ -125,10 +123,9 @@ class OptimizerTest(test.TestCase): [3.0, 4.0], dtype=dtype, trainable=False, name='b') return 5 * var0 + var1 # pylint: enable=cell-var-from-loop - cost = loss if context.in_eager_mode() else loss() sgd_op = gradient_descent.GradientDescentOptimizer(3.0) with self.assertRaisesRegexp(ValueError, 'No.*variables'): - sgd_op.minimize(cost) + sgd_op.minimize(loss) @test_util.run_in_graph_and_eager_modes() def testNoGradients(self): @@ -140,14 +137,13 @@ class OptimizerTest(test.TestCase): var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype, name='b%d' % i) # pylint: disable=cell-var-from-loop - def loss(_): + def loss(): return 5 * var0 # pylint: enable=cell-var-from-loop - cost = loss if context.in_eager_mode() else loss(var1) sgd_op = gradient_descent.GradientDescentOptimizer(3.0) with self.assertRaisesRegexp(ValueError, 'No gradients'): # var1 has no gradient - sgd_op.minimize(cost, var_list=[var1]) + sgd_op.minimize(loss, var_list=[var1]) @test_util.run_in_graph_and_eager_modes() def testNoGradientsForAnyVariables_Minimize(self): @@ -158,13 +154,12 @@ class OptimizerTest(test.TestCase): name='a_%d' % i) var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype, name='b_%d' % i) - def loss(unused_v1, unused_v2): + def loss(): return constant_op.constant(5.0) - cost = loss if context.in_eager_mode() else loss(var0, var1) sgd_op = gradient_descent.GradientDescentOptimizer(3.0) with self.assertRaisesRegexp(ValueError, 'No gradients provided for any variable'): - sgd_op.minimize(cost, var_list=[var0, var1]) + sgd_op.minimize(loss, var_list=[var0, var1]) @test_util.run_in_graph_and_eager_modes() def testNoGradientsForAnyVariables_ApplyGradients(self): @@ -189,11 +184,10 @@ class OptimizerTest(test.TestCase): name='a%d' % i) var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype, name='b%d' % i) - def loss(v0, v1): - return 5 * v0 + 3 * v1 - cost = loss if context.in_eager_mode() else loss(var0, var1) + def loss(): + return 5 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop sgd_op = gradient_descent.GradientDescentOptimizer(3.0) - grads_and_vars = sgd_op.compute_gradients(cost, [var0, var1]) + grads_and_vars = sgd_op.compute_gradients(loss, [var0, var1]) # Convert gradients to tf.Variables converted_grads = [ resource_variable_ops.ResourceVariable(array_ops.zeros([2], dtype), @@ -221,6 +215,21 @@ class OptimizerTest(test.TestCase): self.assertAllClose([-14., -13.], self.evaluate(var0)) self.assertAllClose([-6., -5.], self.evaluate(var1)) + @test_util.run_in_graph_and_eager_modes() + def testComputeGradientsWithTensors(self): + x = ops.convert_to_tensor(1.0) + def f(): + return x * x + sgd_op = gradient_descent.GradientDescentOptimizer(3.0) + grads_and_vars = sgd_op.compute_gradients(f, [x]) + self.assertEqual(1, len(grads_and_vars)) + grad, x_as_var = grads_and_vars[0] + self.assertIs(x, x_as_var) + self.assertEqual(2.0, self.evaluate(grad)) + + with self.assertRaises(NotImplementedError): + sgd_op.apply_gradients(grads_and_vars) + def testTrainOp(self): with self.test_session(): var0 = variables.Variable([1.0, 2.0]) diff --git a/tensorflow/python/training/proximal_adagrad.py b/tensorflow/python/training/proximal_adagrad.py index da31ab325d5e45e1943f554c45717cceb4dc638f..9bd677b8efcd447f74ec2a3cbe94d63eeb9a4dd1 100644 --- a/tensorflow/python/training/proximal_adagrad.py +++ b/tensorflow/python/training/proximal_adagrad.py @@ -23,8 +23,10 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.ProximalAdagradOptimizer") class ProximalAdagradOptimizer(optimizer.Optimizer): # pylint: disable=line-too-long """Optimizer that implements the Proximal Adagrad algorithm. diff --git a/tensorflow/python/training/proximal_gradient_descent.py b/tensorflow/python/training/proximal_gradient_descent.py index 53e9dc2ef2c86a20070fdbdc690b39d2c0e9df06..369b6cbb50e5c621737c095a24eeb473f3870534 100644 --- a/tensorflow/python/training/proximal_gradient_descent.py +++ b/tensorflow/python/training/proximal_gradient_descent.py @@ -24,8 +24,10 @@ from tensorflow.python.ops import math_ops # pylint: enable=unused-import from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.ProximalGradientDescentOptimizer") class ProximalGradientDescentOptimizer(optimizer.Optimizer): # pylint: disable=line-too-long """Optimizer that implements the proximal gradient descent algorithm. diff --git a/tensorflow/python/training/queue_runner_impl.py b/tensorflow/python/training/queue_runner_impl.py index 4e7c81d7b2913d71a23dcaa3751db2aaffdc67cf..d38c5499c73e1217effbc907077236cb6c8e0ae8 100644 --- a/tensorflow/python/training/queue_runner_impl.py +++ b/tensorflow/python/training/queue_runner_impl.py @@ -27,8 +27,10 @@ from tensorflow.python.eager import context from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.queue_runner.QueueRunner", "train.QueueRunner") class QueueRunner(object): """Holds a list of enqueue operations for a queue, each to be run in a thread. @@ -87,7 +89,7 @@ class QueueRunner(object): restoring from `queue_runner_def`. RuntimeError: If eager execution is enabled. """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError( "QueueRunners are not supported when eager execution is enabled. " "Instead, please use tf.data to get data into your model.") @@ -384,6 +386,7 @@ class QueueRunner(object): import_scope=import_scope) +@tf_export("train.queue_runner.add_queue_runner", "train.add_queue_runner") def add_queue_runner(qr, collection=ops.GraphKeys.QUEUE_RUNNERS): """Adds a `QueueRunner` to a collection in the graph. @@ -402,6 +405,8 @@ def add_queue_runner(qr, collection=ops.GraphKeys.QUEUE_RUNNERS): ops.add_to_collection(collection, qr) +@tf_export("train.queue_runner.start_queue_runners", + "train.start_queue_runners") def start_queue_runners(sess=None, coord=None, daemon=True, start=True, collection=ops.GraphKeys.QUEUE_RUNNERS): """Starts all queue runners collected in the graph. @@ -436,7 +441,7 @@ def start_queue_runners(sess=None, coord=None, daemon=True, start=True, use the `tf.data` API instead. @end_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Queues are not compatible with eager execution.") if sess is None: sess = ops.get_default_session() diff --git a/tensorflow/python/training/rmsprop.py b/tensorflow/python/training/rmsprop.py index 745e61201823d81ac01f117be89b910b2810f80c..341b970c92e42b4fe392d91f57219d713d2513e5 100644 --- a/tensorflow/python/training/rmsprop.py +++ b/tensorflow/python/training/rmsprop.py @@ -42,12 +42,15 @@ 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 init_ops from tensorflow.python.ops import math_ops from tensorflow.python.training import optimizer from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.RMSPropOptimizer") class RMSPropOptimizer(optimizer.Optimizer): """Optimizer that implements the RMSProp algorithm. @@ -105,9 +108,13 @@ class RMSPropOptimizer(optimizer.Optimizer): def _create_slots(self, var_list): for v in var_list: - init_rms = init_ops.ones_initializer(dtype=v.dtype) + if v.get_shape().is_fully_defined(): + init_rms = init_ops.ones_initializer(dtype=v.dtype.base_dtype) + else: + init_rms = array_ops.ones_like(v) self._get_or_make_slot_with_initializer(v, init_rms, v.get_shape(), - v.dtype, "rms", self._name) + v.dtype.base_dtype, "rms", + self._name) if self._centered: self._zeros_slot(v, "mg", self._name) self._zeros_slot(v, "momentum", self._name) diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py index 554472e043c83b7c8cc08ac8087618317cf73e90..5f2e56894adb42b3d34851db1a364b30b1a6d67e 100644 --- a/tensorflow/python/training/saver.py +++ b/tensorflow/python/training/saver.py @@ -50,9 +50,11 @@ from tensorflow.python.ops import string_ops from tensorflow.python.ops import variables from tensorflow.python.platform import gfile from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpointable from tensorflow.python.training import training_util from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export # Op names which identify variable reads which should be saved. @@ -195,8 +197,8 @@ class BaseSaverBuilder(object): # Copy the restored tensor to the variable's device. with ops.device(self._var_device): restored_tensor = array_ops.identity(restored_tensor) - return resource_variable_ops.shape_safe_assign_variable_handle( - self.handle_op, self._var_shape, restored_tensor) + return resource_variable_ops.shape_safe_assign_variable_handle( + self.handle_op, self._var_shape, restored_tensor) def __init__(self, write_version=saver_pb2.SaverDef.V2): self._write_version = write_version @@ -309,8 +311,7 @@ class BaseSaverBuilder(object): Returns: A string tensor. """ - # pylint: disable=protected-access - return gen_io_ops._sharded_filename(filename_tensor, shard, num_shards) + return gen_io_ops.sharded_filename(filename_tensor, shard, num_shards) def _AddSaveOps(self, filename_tensor, saveables): """Add ops to save variables that are on the same shard. @@ -419,8 +420,7 @@ class BaseSaverBuilder(object): sharded_saves.append(self._AddSaveOps(sharded_filename, saveables)) # Return the sharded name for the save path. with ops.control_dependencies([x.op for x in sharded_saves]): - # pylint: disable=protected-access - return gen_io_ops._sharded_filespec(filename_tensor, num_shards_tensor) + return gen_io_ops.sharded_filespec(filename_tensor, num_shards_tensor) def _AddRestoreOps(self, filename_tensor, @@ -576,10 +576,33 @@ class BaseSaverBuilder(object): names_to_saveables[name].append(var) else: names_to_saveables[name] = [var] + elif (isinstance(var, checkpointable.CheckpointableBase) + and not isinstance(var, variables.Variable)): + checkpointable_saveables = [ + (factory() if callable(factory) else factory) + for factory in var._gather_saveables_for_checkpoint().values()] + names_to_saveables.update( + BaseSaverBuilder.OpListToDict(checkpointable_saveables)) else: - if context.in_graph_mode(): + if context.executing_eagerly(): + if not isinstance(var, resource_variable_ops.ResourceVariable): + raise ValueError( + "Can only save/restore ResourceVariables when eager execution " + "is enabled, type: %s." % type(var)) + set_var = names_to_saveables.setdefault(var._shared_name, var) + if set_var is not var: + raise ValueError( + ("Two different ResourceVariable objects with the same " + "shared_name '%s' were passed to the Saver. This likely means " + "that they were created in different Graphs or isolation " + "contexts, and may not be checkpointed together.") % + (var._shared_name,)) + else: if convert_variable_to_tensor: - var = ops.internal_convert_to_tensor(var, as_ref=True) + if isinstance(var, resource_variable_ops.ResourceVariable): + var = var._graph_element # pylint: disable=protected-access + else: + var = ops.internal_convert_to_tensor(var, as_ref=True) if not BaseSaverBuilder._IsVariable(var): raise TypeError("Variable to save is not a Variable: %s" % var) if var.op.type == "ReadVariableOp": @@ -590,18 +613,6 @@ class BaseSaverBuilder(object): raise ValueError("At least two variables have the same name: %s" % name) names_to_saveables[name] = var - else: - if not isinstance(var, resource_variable_ops.ResourceVariable): - raise ValueError("Can only save/restore ResourceVariable eager " - "mode is enabled, type: %s." % type(var)) - set_var = names_to_saveables.setdefault(var._shared_name, var) - if set_var is not var: - raise ValueError( - ("Two different ResourceVariable objects with the same " - "shared_name '%s' were passed to the Saver. This likely means " - "that they were created in different Graphs or isolation " - "contexts, and may not be checkpointed together.") % ( - var._shared_name,)) # pylint: enable=protected-access return names_to_saveables @@ -663,13 +674,16 @@ class BaseSaverBuilder(object): # pylint: enable=protected-access else: # A variable or tensor. - if context.in_eager_mode(): + if context.executing_eagerly(): if not isinstance(op, resource_variable_ops.ResourceVariable): raise ValueError("Can only save/restore ResourceVariable eager " "mode is enabled, type: %s." % type(op)) saveable = BaseSaverBuilder.ResourceVariableSaveable(op, "", name) else: - variable = ops.internal_convert_to_tensor(op, as_ref=True) + if isinstance(op, resource_variable_ops.ResourceVariable): + variable = op._graph_element # pylint: disable=protected-access + else: + variable = ops.internal_convert_to_tensor(op, as_ref=True) if not BaseSaverBuilder._IsVariable(variable): raise TypeError("names_to_saveables must be a dict mapping string " "names to Tensors/Variables. Not a variable: %s" % @@ -767,8 +781,10 @@ class BaseSaverBuilder(object): build_save=True, build_restore=True): """build() with option to only perform save and restore.""" - if context.in_graph_mode() and (not build_save or not build_restore): - raise ValueError("Graph mode needs to build save and restore together.") + if not context.executing_eagerly() and (not build_save or + not build_restore): + raise ValueError("save and restore operations need to be built together " + " when eager execution is not enabled.") saveables = self._ValidateAndSliceInputs(names_to_saveables) if max_to_keep is None: @@ -805,22 +821,22 @@ class BaseSaverBuilder(object): # such usage model makes sense. # # assert restore_op.name.endswith("restore_all"), restore_op.name - if context.in_graph_mode(): + if context.executing_eagerly(): + # Store the tensor values to the tensor_names. + save_tensor_name = save_tensor.numpy() if build_save else "" return saver_pb2.SaverDef( - filename_tensor_name=filename_tensor.name, - save_tensor_name=save_tensor.name, - restore_op_name=restore_op.name, + filename_tensor_name=filename_tensor.numpy(), + save_tensor_name=save_tensor_name, + restore_op_name="", max_to_keep=max_to_keep, sharded=sharded, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, version=self._write_version) else: - # Store the tensor values to the tensor_names. - save_tensor_name = save_tensor.numpy() if build_save else "" return saver_pb2.SaverDef( - filename_tensor_name=filename_tensor.numpy(), - save_tensor_name=save_tensor_name, - restore_op_name="", + filename_tensor_name=filename_tensor.name, + save_tensor_name=save_tensor.name, + restore_op_name=restore_op.name, max_to_keep=max_to_keep, sharded=sharded, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, @@ -889,6 +905,7 @@ def _GetCheckpointFilename(save_dir, latest_filename): return os.path.join(save_dir, latest_filename) +@tf_export("train.generate_checkpoint_state_proto") def generate_checkpoint_state_proto(save_dir, model_checkpoint_path, all_model_checkpoint_paths=None): @@ -933,6 +950,7 @@ def generate_checkpoint_state_proto(save_dir, return coord_checkpoint_proto +@tf_export("train.update_checkpoint_state") def update_checkpoint_state(save_dir, model_checkpoint_path, all_model_checkpoint_paths=None, @@ -1025,6 +1043,7 @@ def _update_checkpoint_state(save_dir, text_format.MessageToString(ckpt)) +@tf_export("train.get_checkpoint_state") def get_checkpoint_state(checkpoint_dir, latest_filename=None): """Returns CheckpointState proto from the "checkpoint" file. @@ -1082,6 +1101,7 @@ def get_checkpoint_state(checkpoint_dir, latest_filename=None): return ckpt +@tf_export("train.Saver") class Saver(object): """Saves and restores variables. @@ -1115,8 +1135,9 @@ class Saver(object): the proliferation of checkpoint files on disk: * `max_to_keep` indicates the maximum number of recent checkpoint files to - keep. As new files are created, older files are deleted. If None or 0, - all checkpoint files are kept. Defaults to 5 (that is, the 5 most recent + keep. As new files are created, older files are deleted. If None or 0, + no checkpoints are deleted from the filesystem but only the last one is + kept in the `checkpoint` file. Defaults to 5 (that is, the 5 most recent checkpoint files are kept.) * `keep_checkpoint_every_n_hours`: In addition to keeping the most recent @@ -1265,7 +1286,7 @@ class Saver(object): raise ValueError( "If `var_list` is provided then build cannot be deferred. " "Either set defer_build=False or var_list=None.") - if context.in_eager_mode() and var_list is None: + if context.executing_eagerly() and var_list is None: raise RuntimeError( "When eager execution is enabled, `var_list` must specify a list or " "dict of variables to save") @@ -1284,7 +1305,12 @@ class Saver(object): self._write_version = write_version self._pad_step_number = pad_step_number self._filename = filename - if not defer_build and context.in_graph_mode(): + self._last_checkpoints = [] + self._checkpoints_to_be_deleted = [] + if context.executing_eagerly(): + self._next_checkpoint_time = ( + time.time() + self._keep_checkpoint_every_n_hours * 3600) + elif not defer_build: self.build() if self.saver_def: self._check_saver_def() @@ -1292,7 +1318,7 @@ class Saver(object): self._save_relative_paths = save_relative_paths def build(self): - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Use save/restore instead of build in eager mode.") self._build(self._filename, build_save=True, build_restore=True) @@ -1302,12 +1328,12 @@ class Saver(object): def _build(self, checkpoint_path, build_save, build_restore): """Builds saver_def.""" - if context.in_graph_mode(): + if not context.executing_eagerly(): if self._is_built: return self._is_built = True - if not self.saver_def or context.in_eager_mode(): + if not self.saver_def or context.executing_eagerly(): if self._builder is None: self._builder = BulkSaverBuilder(self._write_version) @@ -1344,17 +1370,17 @@ class Saver(object): self.saver_def.restore_op_name, self._name) self._check_saver_def() - # Updates next checkpoint time. - self._next_checkpoint_time = ( - time.time() + self.saver_def.keep_checkpoint_every_n_hours * 3600) - self._last_checkpoints = [] - self._checkpoints_to_be_deleted = [] + if not context.executing_eagerly(): + # Updates next checkpoint time. + # Set in __init__ when executing eagerly. + self._next_checkpoint_time = ( + time.time() + self.saver_def.keep_checkpoint_every_n_hours * 3600) def _check_saver_def(self): if not isinstance(self.saver_def, saver_pb2.SaverDef): raise ValueError("saver_def must be a saver_pb2.SaverDef: %s" % self.saver_def) - if context.in_graph_mode(): + if not context.executing_eagerly(): if not self.saver_def.save_tensor_name: raise ValueError("saver_def must specify the save_tensor_name: %s" % str(self.saver_def)) @@ -1592,9 +1618,9 @@ class Saver(object): [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes). Returns: - A string: path prefix used for the checkpoint files. If checkpoint - format is V1 and the saver is sharded, this string ends with: - '-?????-of-nnnnn' where 'nnnnn' is the number of shards created. + A string: path prefix used for the checkpoint files. If the saver is + sharded, this string ends with: '-?????-of-nnnnn' where 'nnnnn' + is the number of shards created. If the saver is empty, returns None. Raises: @@ -1604,7 +1630,7 @@ class Saver(object): RuntimeError: If save and restore ops weren't built. """ # pylint: enable=line-too-long - if not self._is_built and context.in_graph_mode(): + if not self._is_built and not context.executing_eagerly(): raise RuntimeError( "`build()` should be called before save if defer_build==True") if latest_filename is None: @@ -1636,21 +1662,21 @@ class Saver(object): "'latest_filename' collides with 'save_path': '%s' and '%s'" % (latest_filename, save_path)) - if (context.in_graph_mode() and + if (not context.executing_eagerly() and not isinstance(sess, session.SessionInterface)): raise TypeError("'sess' must be a Session; %s" % sess) save_path_parent = os.path.dirname(save_path) if not self._is_empty: try: - if context.in_graph_mode(): - model_checkpoint_path = sess.run( - self.saver_def.save_tensor_name, - {self.saver_def.filename_tensor_name: checkpoint_file}) - else: + if context.executing_eagerly(): self._build_eager( checkpoint_file, build_save=True, build_restore=False) model_checkpoint_path = self.saver_def.save_tensor_name + else: + model_checkpoint_path = sess.run( + self.saver_def.save_tensor_name, + {self.saver_def.filename_tensor_name: checkpoint_file}) model_checkpoint_path = compat.as_str(model_checkpoint_path) if write_state: @@ -1672,7 +1698,7 @@ class Saver(object): if write_meta_graph: meta_graph_filename = self._MetaGraphFilename( checkpoint_file, meta_graph_suffix=meta_graph_suffix) - if context.in_graph_mode(): + if not context.executing_eagerly(): with sess.graph.as_default(): self.export_meta_graph( meta_graph_filename, strip_default_attrs=strip_default_attrs) @@ -1744,17 +1770,12 @@ class Saver(object): return if save_path is None: raise ValueError("Can't load save_path when it is None.") - if (os.path.isfile(save_path) and - self._write_version != saver_pb2.SaverDef.V1): - raise ValueError("The specified path: %s is a file." - " Please specify only the path prefix" - " to the checkpoint files." % save_path) logging.info("Restoring parameters from %s", save_path) - if context.in_graph_mode(): + if context.executing_eagerly(): + self._build_eager(save_path, build_save=False, build_restore=True) + else: sess.run(self.saver_def.restore_op_name, {self.saver_def.filename_tensor_name: save_path}) - else: - self._build_eager(save_path, build_save=False, build_restore=True) @staticmethod def _add_collection_def(meta_graph_def, key, export_scope=None): @@ -1788,6 +1809,7 @@ def _prefix_to_checkpoint_path(prefix, format_version): return prefix # Just the data file. +@tf_export("train.latest_checkpoint") def latest_checkpoint(checkpoint_dir, latest_filename=None): """Finds the filename of latest saved checkpoint file. @@ -1817,6 +1839,7 @@ def latest_checkpoint(checkpoint_dir, latest_filename=None): return None +@tf_export("train.import_meta_graph") def import_meta_graph(meta_graph_or_file, clear_devices=False, import_scope=None, **kwargs): """Recreates a Graph saved in a `MetaGraphDef` proto. @@ -1892,7 +1915,7 @@ def import_meta_graph(meta_graph_or_file, clear_devices=False, execution is enabled. @end_compatibility """ # pylint: disable=g-doc-exception - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Exporting/importing meta graphs is not supported when " "eager execution is enabled. No graph exists when eager " "execution is enabled.") @@ -1918,6 +1941,7 @@ def import_meta_graph(meta_graph_or_file, clear_devices=False, return None +@tf_export("train.export_meta_graph") def export_meta_graph(filename=None, meta_info_def=None, graph_def=None, @@ -1974,7 +1998,7 @@ def export_meta_graph(filename=None, @end_compatibility """ # pylint: enable=line-too-long - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Exporting/importing meta graphs is not supported when " "eager execution is enabled. No graph exists when eager " "execution is enabled.") @@ -1994,6 +2018,7 @@ def export_meta_graph(filename=None, return meta_graph_def +@tf_export("train.checkpoint_exists") def checkpoint_exists(checkpoint_prefix): """Checks whether a V1 or V2 checkpoint exists with the specified prefix. @@ -2018,6 +2043,7 @@ def checkpoint_exists(checkpoint_prefix): return False +@tf_export("train.get_checkpoint_mtimes") def get_checkpoint_mtimes(checkpoint_prefixes): """Returns the mtimes (modification timestamps) of the checkpoints. diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py index c5a6f49df599434ab3bc1a9fe3d85db6f824071e..787582ae70e6d22ee8fa94f1d82df03aa0e29c14 100644 --- a/tensorflow/python/training/saver_test.py +++ b/tensorflow/python/training/saver_test.py @@ -35,6 +35,7 @@ from google.protobuf import text_format from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf import meta_graph_pb2 from tensorflow.core.protobuf import queue_runner_pb2 +from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.core.protobuf import saver_pb2 from tensorflow.python import pywrap_tensorflow from tensorflow.python.client import session @@ -53,6 +54,7 @@ from tensorflow.python.lib.io import file_io 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 gradients_impl from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops from tensorflow.python.ops import partitioned_variables @@ -66,6 +68,7 @@ from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.summary import summary from tensorflow.python.training import adam +from tensorflow.python.training import checkpointable from tensorflow.python.training import gradient_descent from tensorflow.python.training import queue_runner_impl from tensorflow.python.training import saver as saver_module @@ -89,7 +92,7 @@ class SaverTest(test.TestCase): v2_init = v2.insert("k1", 30.0) # Initialize all variables - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate([variables.global_variables_initializer(), v2_init]) # Check that the parameter nodes have been initialized. @@ -117,7 +120,7 @@ class SaverTest(test.TestCase): v2 = saver_test_utils.CheckpointedOp(name="v2") # Assert that the variables are not initialized. - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual( len(variables.report_uninitialized_variables().eval()), 2) self.assertEqual(0, len(v2.keys().eval())) @@ -140,7 +143,7 @@ class SaverTest(test.TestCase): v2_init = v2_2.insert("k1000", 3000.0) # Check that the parameter nodes have been initialized. - if context.in_graph_mode(): + if not context.executing_eagerly(): init_all_op = [variables.global_variables_initializer(), v2_init] self.evaluate(init_all_op) # TODO(xpan): Why _mutable_hash_table_v2 doesn't create empty @@ -249,10 +252,10 @@ class SaverTest(test.TestCase): with self.test_session(graph=ops_lib.Graph()) as sess: v = resource_variable_ops.ResourceVariable([1], caching_device="/cpu:0", name="v") - if context.in_graph_mode(): - self.evaluate(variables.global_variables_initializer()) - else: + if context.executing_eagerly(): sess = None + else: + self.evaluate(variables.global_variables_initializer()) save = saver_module.Saver([v]) save.save(sess, save_path) @@ -260,6 +263,24 @@ class SaverTest(test.TestCase): save2.restore(sess, save_path) self.assertEquals(self.evaluate(v), [1]) + def testNoAdditionalOpsAddedBySaverForResourceVariablesOutsideSaveScope(self): + with ops_lib.Graph().as_default() as g: + v = resource_variable_ops.ResourceVariable(1.0, name="v") + with ops_lib.name_scope("saver1"): + saver_module.Saver() + with ops_lib.name_scope("saver2"): + saver_module.Saver({"name": v}) + ops_in_saver1_scope_but_not_save_scope = [ + op for op in g.get_operations() + if (op.name.startswith("saver1/") and + not op.name.startswith("saver1/save/"))] + self.assertEqual(ops_in_saver1_scope_but_not_save_scope, []) + ops_in_saver2_scope_but_not_save_scope = [ + op for op in g.get_operations() + if (op.name.startswith("saver2/") and + not op.name.startswith("saver2/save/"))] + self.assertEqual(ops_in_saver2_scope_but_not_save_scope, []) + def testSaveCopyRestoreWithSaveRelativePaths(self): """Save, copy checkpoint dir and restore from copied dir. @@ -497,7 +518,7 @@ class SaverTest(test.TestCase): with self.test_session(graph=ops_lib.Graph()) as sess: var = resource_variable_ops.ResourceVariable(var_value, name=var_name) save = saver_module.Saver({var_name: var}) - if context.in_graph_mode(): + if not context.executing_eagerly(): self.evaluate(var.initializer) val = save.save(sess, save_path) self.assertEqual(save_path, val) @@ -657,11 +678,11 @@ class SaverTest(test.TestCase): { var._shared_name: var }, pad_step_number=pad_step_number) - if context.in_graph_mode(): + if context.executing_eagerly(): + sess = None + else: self.evaluate(var.initializer) sess = ops_lib.get_default_session() - else: - sess = None if use_tensor: global_step = constant_op.constant(global_step_int) val = save.save(sess, save_path, global_step=global_step) @@ -1039,6 +1060,77 @@ class MaxToKeepTest(test.TestCase): self.assertEqual(checkpoint_state.all_model_checkpoint_paths, all_model_checkpoint_paths) + def testMaxToKeepEager(self): + with context.eager_mode(): + save_dir = self._get_test_dir("max_to_keep_non_sharded") + + v = variable_scope.variable(10.0, name="v") + save = saver_module.Saver({"v": v}, max_to_keep=2) + self.evaluate(variables.global_variables_initializer()) + if not context.executing_eagerly(): + self.assertEqual([], save.last_checkpoints) + + s1 = save.save(None, os.path.join(save_dir, "s1")) + self.assertEqual([s1], save.last_checkpoints) + self.assertTrue(saver_module.checkpoint_exists(s1)) + self.assertCheckpointState( + model_checkpoint_path=s1, + all_model_checkpoint_paths=[s1], + save_dir=save_dir) + + s2 = save.save(None, os.path.join(save_dir, "s2")) + self.assertEqual([s1, s2], save.last_checkpoints) + self.assertTrue(saver_module.checkpoint_exists(s1)) + self.assertTrue(saver_module.checkpoint_exists(s2)) + self.assertCheckpointState( + model_checkpoint_path=s2, + all_model_checkpoint_paths=[s1, s2], + save_dir=save_dir) + + s3 = save.save(None, os.path.join(save_dir, "s3")) + self.assertEqual([s2, s3], save.last_checkpoints) + self.assertFalse(saver_module.checkpoint_exists(s1)) + self.assertTrue(saver_module.checkpoint_exists(s2)) + self.assertTrue(saver_module.checkpoint_exists(s3)) + self.assertCheckpointState( + model_checkpoint_path=s3, + all_model_checkpoint_paths=[s2, s3], + save_dir=save_dir) + + # Create a second helper, identical to the first. + save2 = saver_module.Saver({"v": v}, max_to_keep=2) + save2.set_last_checkpoints(save.last_checkpoints) + + # Exercise the first helper. + + # Adding s2 again (old s2 is removed first, then new s2 appended) + s2 = save.save(None, os.path.join(save_dir, "s2")) + self.assertEqual([s3, s2], save.last_checkpoints) + self.assertFalse(saver_module.checkpoint_exists(s1)) + self.assertTrue(saver_module.checkpoint_exists(s3)) + self.assertTrue(saver_module.checkpoint_exists(s2)) + self.assertCheckpointState( + model_checkpoint_path=s2, + all_model_checkpoint_paths=[s3, s2], + save_dir=save_dir) + + # Adding s1 (s3 should now be deleted as oldest in list) + s1 = save.save(None, os.path.join(save_dir, "s1")) + self.assertEqual([s2, s1], save.last_checkpoints) + self.assertFalse(saver_module.checkpoint_exists(s3)) + self.assertTrue(saver_module.checkpoint_exists(s2)) + self.assertCheckpointState( + model_checkpoint_path=s1, + all_model_checkpoint_paths=[s2, s1], + save_dir=save_dir) + + s2 = save2.save(None, os.path.join(save_dir, "s2")) + self.assertEqual([s3, s2], save2.last_checkpoints) + # Created by the first helper. + self.assertTrue(saver_module.checkpoint_exists(s1)) + # Deleted by the first helper. + self.assertFalse(saver_module.checkpoint_exists(s3)) + def testNonSharded(self): save_dir = self._get_test_dir("max_to_keep_non_sharded") @@ -1301,15 +1393,16 @@ class KeepCheckpointEveryNHoursTest(test.TestCase): gfile.MakeDirs(test_dir) return test_dir + @test_util.run_in_graph_and_eager_modes() @test.mock.patch.object(saver_module, "time") def testNonSharded(self, mock_time): save_dir = self._get_test_dir("keep_checkpoint_every_n_hours") with self.test_session() as sess: - v = variables.Variable([10.0], name="v") + v = variable_scope.variable([10.0], name="v") # Run the initializer NOW to avoid the 0.5s overhead of the first Run() # call, which throws the test timing off in fastbuild mode. - variables.global_variables_initializer().run() + self.evaluate(variables.global_variables_initializer()) # Create a saver that will keep the last 2 checkpoints plus one every 0.7 # seconds. start_time = time.time() @@ -1387,7 +1480,7 @@ class SaveRestoreWithVariableNameMap(test.TestCase): v0 = variable_op(-1.0, name="v0") v1 = variable_op(-1.0, name="v1") - if context.in_graph_mode(): + if not context.executing_eagerly(): with self.assertRaisesOpError("uninitialized"): self.evaluate(v0) with self.assertRaisesOpError("uninitialized"): @@ -1397,7 +1490,7 @@ class SaveRestoreWithVariableNameMap(test.TestCase): save.restore(sess, save_path) # Check that the parameter nodes have been restored. - if context.in_graph_mode(): + if not context.executing_eagerly(): self.assertEqual(10.0, self.evaluate(v0)) self.assertEqual(20.0, self.evaluate(v1)) @@ -1407,7 +1500,7 @@ class SaveRestoreWithVariableNameMap(test.TestCase): v0 = variable_op(-1.0, name="restore_prefix/v0") v1 = variable_op(-1.0, name="restore_prefix/v1") - if context.in_graph_mode(): + if not context.executing_eagerly(): with self.assertRaisesOpError("uninitialized"): self.evaluate(v0) with self.assertRaisesOpError("uninitialized"): @@ -1646,7 +1739,8 @@ class CheckpointStateTest(test.TestCase): os.path.join(save_dir, "./model.ckpt-687529")) -@test_util.with_c_api +# TODO(skyewm): reenable when this works with _USE_C_SHAPES=False +# @test_util.with_c_api class MetaGraphTest(test.TestCase): def _get_test_dir(self, dirname): @@ -2039,6 +2133,113 @@ class MetaGraphTest(test.TestCase): self._testGraphExtensionRestore(test_dir) self._testRestoreFromTrainGraphWithControlContext(test_dir) + def _testGradientSerDes(self, graph_fn): + """Tests that gradients can be computed after exporting and importing. + + Builds a graph, exports it, and verifies that it can be imported and the + gradient can be built and run correctly. + + Args: + graph_fn: takes a single float Tensor argument as input, outputs a single + Tensor + """ + test_dir = self._get_test_dir("nested_control_flow") + filename = os.path.join(test_dir, "metafile") + saver_ckpt = os.path.join(test_dir, "saver.ckpt") + + # Create while loop using `outer_body_fn`. + with ops_lib.Graph().as_default(): + var = variables.Variable(0.0) + var_name = var.name + output = graph_fn(var) + output_name = output.name + init_op = variables.global_variables_initializer() + + # Generate a MetaGraphDef containing the while loop. + with session.Session() as sess: + sess.run(init_op) + sess.run(output) + saver = saver_module.Saver() + saver.save(sess, saver_ckpt) + saver.export_meta_graph(filename) + + # Build and run the gradients of the while loop. We use this below to + # verify that the gradients are correct with an imported MetaGraphDef. + grad = gradients_impl.gradients([output], [var]) + # Turn off constant folding to avoid breaking testNestedControlFlowSerDes. + # It appears that a missing control dependency in the gradient graph + # causes the fetch node to not be triggered. + no_constfold_config = config_pb2.ConfigProto() + no_constfold_config.graph_options.rewrite_options.constant_folding = ( + rewriter_config_pb2.RewriterConfig.OFF) + with session.Session(config=no_constfold_config) as sess: + sess.run(init_op) + expected_grad_value = sess.run(grad) + + # Restore the MetaGraphDef into a new Graph. + with ops_lib.Graph().as_default(): + with session.Session() as sess: + saver = saver_module.import_meta_graph(filename) + saver.restore(sess, saver_ckpt) + + # Make sure we can still build gradients and get the same result. + var = ops_lib.get_default_graph().get_tensor_by_name(var_name) + output = ops_lib.get_default_graph().get_tensor_by_name(output_name) + grad = gradients_impl.gradients([output], [var]) + + init_op = variables.global_variables_initializer() + + with session.Session(config=no_constfold_config) as sess: + sess.run(init_op) + actual_grad_value = sess.run(grad) + self.assertEqual(expected_grad_value, actual_grad_value) + + def _testWhileLoopAndGradientSerDes(self, outer_body_fn): + # Build a while loop with `outer_body_fn`, export it, and verify that it can + # be imported and the gradient can be built and run correctly. + # pylint: disable=g-long-lambda + return self._testGradientSerDes( + lambda x: control_flow_ops.while_loop( + lambda i, y: i < 5, outer_body_fn, [0, x])[1]) + # pylint: enable=g-long-lambda + + def testNestedWhileLoopsSerDes(self): + # Test two simple nested while loops. + def body(i, x): + _, r = control_flow_ops.while_loop(lambda j, y: j < 3, + lambda j, y: (j + 1, y + x), + [0, 0.0]) + return i + 1, x + r + self._testWhileLoopAndGradientSerDes(body) + + def testNestedControlFlowSerDes(self): + # Test while loop in a cond in a while loop. + # pylint: disable=g-long-lambda + def body(i, x): + cond_result = control_flow_ops.cond( + i > 0, + lambda: control_flow_ops.while_loop( + lambda j, y: j < 3, + lambda j, y: (j + 1, y + x), + [0, 0.0])[1], + lambda: x) + return i + 1, cond_result + # pylint: enable=g-long-lambda + self._testWhileLoopAndGradientSerDes(body) + + def testNestedCondsSerDes(self): + # Test conds in a cond. + # pylint: disable=g-long-lambda + self._testGradientSerDes(lambda x: control_flow_ops.cond( + x > 0, + lambda: control_flow_ops.cond(x > 3, + lambda: array_ops.identity(x), + lambda: math_ops.multiply(x, 2.0)), + lambda: control_flow_ops.cond(x < -3, + lambda: constant_op.constant(1.0), + lambda: math_ops.multiply(x, -1.0)))) + # pylint: enable=g-long-lambda + def testStrippedOpListDef(self): with self.test_session(): # Creates a graph. @@ -2660,5 +2861,94 @@ class ScopedGraphTest(test.TestCase): self.assertEqual(2.0, var_dict2["variable2:0"].eval()) +class _OwnsAVariableSimple(checkpointable.CheckpointableBase): + """A Checkpointable object which can be saved using a tf.train.Saver.""" + + def __init__(self): + self.non_dep_variable = variable_scope.get_variable( + name="non_dep_variable", initializer=6., use_resource=True) + + def _gather_saveables_for_checkpoint(self): + return {checkpointable.VARIABLE_VALUE_KEY: self.non_dep_variable} + + # The Saver sorts by name before parsing, so we need a name property. + @property + def name(self): + return self.non_dep_variable.name + + +class _MirroringSaveable( + saver_module.BaseSaverBuilder.ResourceVariableSaveable): + + def __init__(self, primary_variable, mirrored_variable, name): + self._primary_variable = primary_variable + self._mirrored_variable = mirrored_variable + super(_MirroringSaveable, self).__init__( + self._primary_variable, "", name) + + def restore(self, restored_tensors, restored_shapes): + """Restore the same value into both variables.""" + tensor, = restored_tensors + return control_flow_ops.group( + self._primary_variable.assign(tensor), + self._mirrored_variable.assign(tensor)) + + +class _OwnsMirroredVariables(checkpointable.CheckpointableBase): + """A Checkpointable object which returns a more complex SaveableObject.""" + + def __init__(self): + self.non_dep_variable = variable_scope.get_variable( + name="non_dep_variable", initializer=6., use_resource=True) + self.mirrored = variable_scope.get_variable( + name="mirrored", initializer=15., use_resource=True) + + def _gather_saveables_for_checkpoint(self): + def _saveable_factory(name=self.non_dep_variable.name): + return _MirroringSaveable( + primary_variable=self.non_dep_variable, + mirrored_variable=self.mirrored, + name=name) + return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} + + # The Saver sorts by name before parsing, so we need a name property. + @property + def name(self): + return self.non_dep_variable.name + + +@test_util.with_c_api +class CheckpointableCompatibilityTests(test.TestCase): + + # TODO(allenl): Track down python3 reference cycles in these tests. + @test_util.run_in_graph_and_eager_modes() + def testNotSaveableButIsCheckpointable(self): + v = _OwnsAVariableSimple() + saver = saver_module.Saver(var_list=[v]) + test_dir = self.get_temp_dir() + prefix = os.path.join(test_dir, "ckpt") + self.evaluate(v.non_dep_variable.assign(42.)) + with self.test_session() as sess: + save_path = saver.save(sess, prefix) + self.evaluate(v.non_dep_variable.assign(43.)) + saver.restore(sess, save_path) + self.assertEqual(42., self.evaluate(v.non_dep_variable)) + + @test_util.run_in_graph_and_eager_modes() + def testMoreComplexSaveableReturned(self): + v = _OwnsMirroredVariables() + saver = saver_module.Saver(var_list=[v]) + test_dir = self.get_temp_dir() + prefix = os.path.join(test_dir, "ckpt") + self.evaluate(v.non_dep_variable.assign(42.)) + with self.test_session() as sess: + save_path = saver.save(sess, prefix) + self.evaluate(v.non_dep_variable.assign(43.)) + self.evaluate(v.mirrored.assign(44.)) + saver.restore(sess, save_path) + self.assertEqual(42., self.evaluate(v.non_dep_variable)) + self.assertEqual(42., self.evaluate(v.mirrored)) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/training/saver_test_utils.py b/tensorflow/python/training/saver_test_utils.py index 44b06b357ecbe4c8e330a2ccc49e83ddd4bf8c7d..2bbe5b6d845c304c4dc79fb3619c57211ca0489e 100644 --- a/tensorflow/python/training/saver_test_utils.py +++ b/tensorflow/python/training/saver_test_utils.py @@ -35,12 +35,12 @@ class CheckpointedOp(object): # pylint: disable=protected-access def __init__(self, name, table_ref=None): if table_ref is None: - self.table_ref = gen_lookup_ops._mutable_hash_table_v2( + self.table_ref = gen_lookup_ops.mutable_hash_table_v2( key_dtype=dtypes.string, value_dtype=dtypes.float32, name=name) else: self.table_ref = table_ref self._name = name - if context.in_graph_mode(): + if not context.executing_eagerly(): self._saveable = CheckpointedOp.CustomSaveable(self, name) ops_lib.add_to_collection(ops_lib.GraphKeys.SAVEABLE_OBJECTS, self._saveable) @@ -51,16 +51,16 @@ class CheckpointedOp(object): @property def saveable(self): - if context.in_graph_mode(): - return self._saveable - else: + if context.executing_eagerly(): return CheckpointedOp.CustomSaveable(self, self.name) + else: + return self._saveable def insert(self, keys, values): - return gen_lookup_ops._lookup_table_insert_v2(self.table_ref, keys, values) + return gen_lookup_ops.lookup_table_insert_v2(self.table_ref, keys, values) def lookup(self, keys, default): - return gen_lookup_ops._lookup_table_find_v2(self.table_ref, keys, default) + return gen_lookup_ops.lookup_table_find_v2(self.table_ref, keys, default) def keys(self): return self._export()[0] @@ -69,8 +69,8 @@ class CheckpointedOp(object): return self._export()[1] def _export(self): - return gen_lookup_ops._lookup_table_export_v2(self.table_ref, dtypes.string, - dtypes.float32) + return gen_lookup_ops.lookup_table_export_v2(self.table_ref, dtypes.string, + dtypes.float32) class CustomSaveable(saver_module.BaseSaverBuilder.SaveableObject): """A custom saveable for CheckpointedOp.""" @@ -86,6 +86,6 @@ class CheckpointedOp(object): super(CheckpointedOp.CustomSaveable, self).__init__(table, specs, name) def restore(self, restore_tensors, shapes): - return gen_lookup_ops._lookup_table_import_v2( + return gen_lookup_ops.lookup_table_import_v2( self.op.table_ref, restore_tensors[0], restore_tensors[1]) # pylint: enable=protected-access diff --git a/tensorflow/python/training/server_lib.py b/tensorflow/python/training/server_lib.py index 29da67a30a58c1b8b8e172b2ccede340880fef58..2f421d1cc0a0190670082fabf4e25470c6a1723b 100644 --- a/tensorflow/python/training/server_lib.py +++ b/tensorflow/python/training/server_lib.py @@ -23,6 +23,7 @@ from tensorflow.core.protobuf import tensorflow_server_pb2 from tensorflow.python import pywrap_tensorflow from tensorflow.python.framework import errors from tensorflow.python.util import compat +from tensorflow.python.util.tf_export import tf_export def _make_server_def(server_or_cluster_def, job_name, task_index, protocol, @@ -92,6 +93,7 @@ def _make_server_def(server_or_cluster_def, job_name, task_index, protocol, return server_def +@tf_export("train.Server") class Server(object): """An in-process TensorFlow server, for use in distributed training. @@ -221,6 +223,7 @@ class Server(object): start=start) +@tf_export("train.ClusterSpec") class ClusterSpec(object): """Represents a cluster as a set of "tasks", organized into "jobs". diff --git a/tensorflow/python/training/session_manager.py b/tensorflow/python/training/session_manager.py index b396a1e7d0a06ec7b952ba2980e081e01e681d4d..360e02fb44c1062f71bb50449b9ef381510a9c69 100644 --- a/tensorflow/python/training/session_manager.py +++ b/tensorflow/python/training/session_manager.py @@ -25,6 +25,7 @@ from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import saver as saver_mod +from tensorflow.python.util.tf_export import tf_export def _maybe_name(obj): @@ -44,6 +45,7 @@ def _maybe_name(obj): return "" % type(obj) +@tf_export("train.SessionManager") class SessionManager(object): """Training helper that restores from checkpoint and creates session. diff --git a/tensorflow/python/training/session_run_hook.py b/tensorflow/python/training/session_run_hook.py index 5b023d8a2672af5d1fab1c2566b19fca738fd1f7..89f40300650f3b6cd1ae15d946640c9df91771e2 100644 --- a/tensorflow/python/training/session_run_hook.py +++ b/tensorflow/python/training/session_run_hook.py @@ -96,8 +96,10 @@ from __future__ import division from __future__ import print_function import collections +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.SessionRunHook") class SessionRunHook(object): """Hook to extend calls to MonitoredSession.run().""" @@ -189,6 +191,7 @@ class SessionRunHook(object): pass +@tf_export("train.SessionRunArgs") class SessionRunArgs( collections.namedtuple("SessionRunArgs", ["fetches", "feed_dict", "options"])): @@ -213,6 +216,7 @@ class SessionRunArgs( return super(SessionRunArgs, cls).__new__(cls, fetches, feed_dict, options) +@tf_export("train.SessionRunContext") class SessionRunContext(object): """Provides information about the `session.run()` call being made. @@ -264,6 +268,7 @@ class SessionRunContext(object): self._stop_requested = True +@tf_export("train.SessionRunValues") class SessionRunValues( collections.namedtuple("SessionRunValues", ["results", "options", "run_metadata"])): diff --git a/tensorflow/python/training/slot_creator.py b/tensorflow/python/training/slot_creator.py index ea28b5ddfc2dbbf65ec60e86d29ff2a9988d2b97..9ac52dd0715d7ed15e2e57ed286be973614b01e5 100644 --- a/tensorflow/python/training/slot_creator.py +++ b/tensorflow/python/training/slot_creator.py @@ -48,11 +48,6 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables -def _is_resource(v): - """Returns true if v is something you get from a resource variable.""" - return isinstance(v, resource_variable_ops.ResourceVariable) - - def _create_slot_var(primary, val, scope, validate_shape, shape, dtype): """Helper function for creating a slot variable.""" @@ -60,9 +55,12 @@ def _create_slot_var(primary, val, scope, validate_shape, shape, dtype): # scope. current_partitioner = variable_scope.get_variable_scope().partitioner variable_scope.get_variable_scope().set_partitioner(None) + # When init from val instead of callable initializer, the shape is expected to + # be None, not or any fully defined shape. + shape = shape if callable(val) else None slot = variable_scope.get_variable( scope, initializer=val, trainable=False, - use_resource=_is_resource(primary), + use_resource=resource_variable_ops.is_resource_variable(primary), shape=shape, dtype=dtype, validate_shape=validate_shape) variable_scope.get_variable_scope().set_partitioner(current_partitioner) @@ -108,7 +106,11 @@ def create_slot(primary, val, name, colocate_with_primary=True): # and the same name has been previously used, the scope name will add '_N' # as suffix for unique identifications. validate_shape = val.get_shape().is_fully_defined() - with variable_scope.variable_scope(None, primary.op.name + "/" + name): + if context.executing_eagerly(): + prefix = primary._shared_name # pylint: disable=protected-access + else: + prefix = primary.op.name + with variable_scope.variable_scope(None, prefix + "/" + name): if colocate_with_primary: with ops.colocate_with(primary): return _create_slot_var(primary, val, "", validate_shape, None, None) @@ -140,7 +142,10 @@ def create_slot_with_initializer(primary, initializer, shape, dtype, name, # and the same name has been previously used, the scope name will add '_N' # as suffix for unique identifications. validate_shape = shape.is_fully_defined() - prefix = primary.op.name if context.in_graph_mode() else primary._shared_name # pylint: disable=protected-access + if context.executing_eagerly(): + prefix = primary._shared_name # pylint: disable=protected-access + else: + prefix = primary.op.name with variable_scope.variable_scope(None, prefix + "/" + name): if colocate_with_primary: with ops.colocate_with(primary): diff --git a/tensorflow/python/training/supervisor.py b/tensorflow/python/training/supervisor.py index e4514aaea223b6b254a7a72e11e6b70b576fd54b..7389e344c7d8eef8e26c4d24c0985ff66276deea 100644 --- a/tensorflow/python/training/supervisor.py +++ b/tensorflow/python/training/supervisor.py @@ -37,13 +37,15 @@ from tensorflow.python.training import saver as saver_mod from tensorflow.python.training import session_manager as session_manager_mod from tensorflow.python.training import training_util from tensorflow.python.util import deprecation +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.Supervisor") class Supervisor(object): """A training helper that checkpoints models and computes summaries. This class is deprecated. Please use - ${tf.train.MonitoredTrainingSession} instead. + @{tf.train.MonitoredTrainingSession} instead. The Supervisor is a small wrapper around a `Coordinator`, a `Saver`, and a `SessionManager` that takes care of common needs of TensorFlow @@ -303,7 +305,7 @@ class Supervisor(object): `Supervisor`s are not supported when eager execution is enabled. @end_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Supervisors are compatible with eager execution.") # Set default values of arguments. if graph is None: @@ -760,7 +762,7 @@ class Supervisor(object): execution is enabled, use the `tf.data` API. @end_compatibility """ - if context.in_eager_mode(): + if context.executing_eagerly(): raise RuntimeError("Queues are not compatible with eager execution.") if queue_runners is None: queue_runners = self._graph.get_collection(ops.GraphKeys.QUEUE_RUNNERS) diff --git a/tensorflow/python/training/sync_replicas_optimizer.py b/tensorflow/python/training/sync_replicas_optimizer.py index 47702fdad05d13015e0cbf7768129b0c53b6c14c..0c6cf910d1a01dc20b15fb1cd5dbb249fbb60ef5 100644 --- a/tensorflow/python/training/sync_replicas_optimizer.py +++ b/tensorflow/python/training/sync_replicas_optimizer.py @@ -31,6 +31,7 @@ from tensorflow.python.training import optimizer from tensorflow.python.training import queue_runner from tensorflow.python.training import session_manager from tensorflow.python.training import session_run_hook +from tensorflow.python.util.tf_export import tf_export # Please note that the gradients from replicas are averaged instead of summed @@ -38,6 +39,7 @@ from tensorflow.python.training import session_run_hook # rate according to the number of replicas. This change is introduced to be # consistent with how gradients are aggregated (averaged) within a batch in a # replica. +@tf_export("train.SyncReplicasOptimizer") class SyncReplicasOptimizer(optimizer.Optimizer): """Class to synchronize, aggregate gradients and pass them to the optimizer. diff --git a/tensorflow/python/training/training.py b/tensorflow/python/training/training.py index 03811fa38dd021fd5ff222bfbe32234606d6c681..6880cfc4db9b05c02df56c4c1f26f2fea1a7098c 100644 --- a/tensorflow/python/training/training.py +++ b/tensorflow/python/training/training.py @@ -28,6 +28,7 @@ See the @{$python/train} guide. @@ProximalGradientDescentOptimizer @@ProximalAdagradOptimizer @@RMSPropOptimizer +@@custom_gradient @@gradients @@AggregationMethod @@stop_gradient @@ -94,6 +95,8 @@ See the @{$python/train} guide. @@load_variable @@list_variables @@init_from_checkpoint +@@warm_start +@@VocabInfo """ # Optimizers. @@ -135,7 +138,7 @@ from tensorflow.python.training.queue_runner import * # For the module level doc. from tensorflow.python.training import input as _input -from tensorflow.python.training.input import * +from tensorflow.python.training.input import * # pylint: disable=redefined-builtin # pylint: enable=wildcard-import from tensorflow.python.training.basic_session_run_hooks import SecondOrStepTimer @@ -187,8 +190,11 @@ from tensorflow.python.training.training_util import get_global_step from tensorflow.python.training.training_util import assert_global_step from tensorflow.python.training.training_util import create_global_step from tensorflow.python.training.training_util import get_or_create_global_step +from tensorflow.python.training.warm_starting_util import VocabInfo +from tensorflow.python.training.warm_starting_util import warm_start from tensorflow.python.pywrap_tensorflow import do_quantize_training_on_graphdef from tensorflow.python.pywrap_tensorflow import NewCheckpointReader +from tensorflow.python.util.tf_export import tf_export # pylint: disable=wildcard-import # Training data protos. @@ -239,6 +245,23 @@ _allowed_symbols = [ "SequenceExample", # from example_pb2. "ServerDef", ] + +# pylint: disable=undefined-variable +tf_export("train.BytesList")(BytesList) +tf_export("train.ClusterDef")(ClusterDef) +tf_export("train.Example")(Example) +tf_export("train.Feature")(Feature) +tf_export("train.Features")(Features) +tf_export("train.FeatureList")(FeatureList) +tf_export("train.FeatureLists")(FeatureLists) +tf_export("train.FloatList")(FloatList) +tf_export("train.Int64List")(Int64List) +tf_export("train.JobDef")(JobDef) +tf_export("train.SaverDef")(SaverDef) +tf_export("train.SequenceExample")(SequenceExample) +tf_export("train.ServerDef")(ServerDef) +# pylint: enable=undefined-variable + # Include extra modules for docstrings because: # * Input methods in tf.train are documented in io_ops. # * Saver methods in tf.train are documented in state_ops. diff --git a/tensorflow/python/training/training_ops.py b/tensorflow/python/training/training_ops.py index e98c32b614418224b1bc14081bc35f175d769965..d7133cfb500ef11e5b94c7c36905e039f9c0bf46 100644 --- a/tensorflow/python/training/training_ops.py +++ b/tensorflow/python/training/training_ops.py @@ -19,7 +19,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.training import gen_training_ops +from tensorflow.python.training import gen_training_ops # pylint: disable=unused-import # go/tf-wildcard-import # pylint: disable=wildcard-import from tensorflow.python.training.gen_training_ops import * diff --git a/tensorflow/python/training/training_util.py b/tensorflow/python/training/training_util.py index 89a9e129328fe38da2ce497a7f26dc11446ea032..4f1abccc96ff95b6fcf9fe82c7a5d45fc2fd1c0c 100644 --- a/tensorflow/python/training/training_util.py +++ b/tensorflow/python/training/training_util.py @@ -29,6 +29,7 @@ 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 tf_logging as logging +from tensorflow.python.util.tf_export import tf_export # Picked a long key value to minimize the chance of collision with user defined @@ -40,6 +41,7 @@ GLOBAL_STEP_READ_KEY = 'global_step_read_op_cache' write_graph = graph_io.write_graph +@tf_export('train.global_step') def global_step(sess, global_step_tensor): """Small helper to get the global step. @@ -62,11 +64,12 @@ def global_step(sess, global_step_tensor): Returns: The global step value. """ - if context.in_eager_mode(): + if context.executing_eagerly(): return int(global_step_tensor.numpy()) return int(sess.run(global_step_tensor)) +@tf_export('train.get_global_step') def get_global_step(graph=None): """Get the global step tensor. @@ -101,6 +104,7 @@ def get_global_step(graph=None): return global_step_tensor +@tf_export('train.create_global_step') def create_global_step(graph=None): """Create global step tensor in graph. @@ -119,7 +123,7 @@ def create_global_step(graph=None): raise ValueError('"global_step" already exists.') # Create in proper graph and base name_scope. with graph.as_default() as g, g.name_scope(None): - if context.in_eager_mode(): + if context.executing_eagerly(): with ops.device('cpu:0'): return variable_scope.get_variable( ops.GraphKeys.GLOBAL_STEP, @@ -139,6 +143,7 @@ def create_global_step(graph=None): ops.GraphKeys.GLOBAL_STEP]) +@tf_export('train.get_or_create_global_step') def get_or_create_global_step(graph=None): """Returns and create (if necessary) the global step tensor. @@ -156,6 +161,7 @@ def get_or_create_global_step(graph=None): return global_step_tensor +@tf_export('train.assert_global_step') def assert_global_step(global_step_tensor): """Asserts `global_step_tensor` is a scalar int `Variable` or `Tensor`. diff --git a/tensorflow/python/estimator/warm_starting_util.py b/tensorflow/python/training/warm_starting_util.py similarity index 66% rename from tensorflow/python/estimator/warm_starting_util.py rename to tensorflow/python/training/warm_starting_util.py index 48110ef57fcba897cc495323973b2f6761c3add4..4d4fb394c1272d2bf510bb594d70b9aa2edb3df2 100644 --- a/tensorflow/python/estimator/warm_starting_util.py +++ b/tensorflow/python/training/warm_starting_util.py @@ -30,8 +30,10 @@ from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import checkpoint_ops from tensorflow.python.training import checkpoint_utils from tensorflow.python.training import saver +from tensorflow.python.util.tf_export import tf_export +@tf_export("train.VocabInfo", "estimator.VocabInfo") class VocabInfo( collections.namedtuple("VocabInfo", [ "new_vocab", @@ -41,7 +43,7 @@ class VocabInfo( "old_vocab_size", "backup_initializer", ])): - """Vocabulary information for WarmStartSettings. + """Vocabulary information for warm-starting. See @{tf.estimator.WarmStartSettings$WarmStartSettings} for examples of using VocabInfo to warm-start. @@ -81,171 +83,6 @@ class VocabInfo( ) -class WarmStartSettings( - collections.namedtuple("WarmStartSettings", [ - "ckpt_to_initialize_from", - "vars_to_warm_start", - "var_name_to_vocab_info", - "var_name_to_prev_var_name", - ])): - """Settings for warm-starting in Estimators. - - Example Use with canned `DNNEstimator`: - - ``` - emb_vocab_file = tf.feature_column.embedding_column( - tf.feature_column.categorical_column_with_vocabulary_file( - "sc_vocab_file", "new_vocab.txt", vocab_size=100), - dimension=8) - emb_vocab_list = tf.feature_column.embedding_column( - tf.feature_column.categorical_column_with_vocabulary_list( - "sc_vocab_list", vocabulary_list=["a", "b"]), - dimension=8) - estimator = tf.estimator.DNNClassifier( - hidden_units=[128, 64], feature_columns=[emb_vocab_file, emb_vocab_list], - warm_start_from=ws) - ``` - - where `ws` could be defined as: - - Warm-start all weights in the model (input layer and hidden weights). - Either the directory or a specific checkpoint can be provided (in the case - of the former, the latest checkpoint will be used): - - ``` - ws = WarmStartSettings(ckpt_to_initialize_from="/tmp") - ws = WarmStartSettings(ckpt_to_initialize_from="/tmp/model-1000") - ``` - - Warm-start only the embeddings (input layer) and their accumulator variables: - - ``` - ws = WarmStartSettings(ckpt_to_initialize_from="/tmp", - vars_to_warm_start=".*input_layer.*") - ``` - - Warm-start everything except the optimizer accumulator variables - (DNN defaults to Adagrad): - - ``` - ws = WarmStartSettings(ckpt_to_initialize_from="/tmp", - vars_to_warm_start="^(?!.*(Adagrad))") - ``` - - Warm-start all weights but the embedding parameters corresponding to - `sc_vocab_file` have a different vocab from the one used in the current - model: - - ``` - vocab_info = ws_util.VocabInfo( - new_vocab=sc_vocab_file.vocabulary_file, - new_vocab_size=sc_vocab_file.vocabulary_size, - num_oov_buckets=sc_vocab_file.num_oov_buckets, - old_vocab="old_vocab.txt" - ) - ws = WarmStartSettings( - ckpt_to_initialize_from="/tmp", - var_name_to_vocab_info={ - "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info - }) - ``` - - Warm-start only `sc_vocab_file` embeddings (and no other variables), which - have a different vocab from the one used in the current model: - - ``` - vocab_info = ws_util.VocabInfo( - new_vocab=sc_vocab_file.vocabulary_file, - new_vocab_size=sc_vocab_file.vocabulary_size, - num_oov_buckets=sc_vocab_file.num_oov_buckets, - old_vocab="old_vocab.txt" - ) - ws = WarmStartSettings( - ckpt_to_initialize_from="/tmp", - vars_to_warm_start=None, - var_name_to_vocab_info={ - "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info - }) - ``` - - Warm-start all weights but the parameters corresponding to `sc_vocab_file` - have a different vocab from the one used in current checkpoint, and only - 100 of those entries were used: - - ``` - vocab_info = ws_util.VocabInfo( - new_vocab=sc_vocab_file.vocabulary_file, - new_vocab_size=sc_vocab_file.vocabulary_size, - num_oov_buckets=sc_vocab_file.num_oov_buckets, - old_vocab="old_vocab.txt", - old_vocab_size=100 - ) - ws = WarmStartSettings( - ckpt_to_initialize_from="/tmp", - var_name_to_vocab_info={ - "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info - }) - ``` - - Warm-start all weights but the parameters corresponding to `sc_vocab_file` - have a different vocab from the one used in current checkpoint and the - parameters corresponding to `sc_vocab_list` have a different name from the - current checkpoint: - - ``` - vocab_info = ws_util.VocabInfo( - new_vocab=sc_vocab_file.vocabulary_file, - new_vocab_size=sc_vocab_file.vocabulary_size, - num_oov_buckets=sc_vocab_file.num_oov_buckets, - old_vocab="old_vocab.txt", - old_vocab_size=100 - ) - ws = WarmStartSettings( - ckpt_to_initialize_from="/tmp", - var_name_to_vocab_info={ - "input_layer/sc_vocab_file_embedding/embedding_weights": vocab_info - }, - var_name_to_prev_var_name={ - "input_layer/sc_vocab_list_embedding/embedding_weights": - "old_tensor_name" - }) - ``` - - Attributes: - ckpt_to_initialize_from: [Required] A string specifying the directory with - checkpoint file(s) or path to checkpoint from which to warm-start the - model parameters. - vars_to_warm_start: [Optional] A regular expression that captures which - variables to warm-start (see tf.get_collection). Defaults to `'.*'`, - which warm-starts all variables. If `None` is explicitly given, only - variables specified in `var_name_to_vocab_info` will be warm-started. - var_name_to_vocab_info: [Optional] Dict of variable names (strings) to - VocabInfo. The variable names should be "full" variables, not the names - of the partitions. If not explicitly provided, the variable is assumed to - have no vocabulary. - var_name_to_prev_var_name: [Optional] Dict of variable names (strings) to - name of the previously-trained variable in `ckpt_to_initialize_from`. If - not explicitly provided, the name of the variable is assumed to be same - between previous checkpoint and current model. - """ - - def __new__(cls, - ckpt_to_initialize_from, - vars_to_warm_start=".*", - var_name_to_vocab_info=None, - var_name_to_prev_var_name=None): - if not ckpt_to_initialize_from: - raise ValueError( - "`ckpt_to_initialize_from` MUST be set in WarmStartSettings") - return super(WarmStartSettings, cls).__new__( - cls, - ckpt_to_initialize_from, - vars_to_warm_start, - var_name_to_vocab_info or {}, - var_name_to_prev_var_name or {}, - ) - - def _is_variable(x): return (isinstance(x, variables_lib.Variable) or isinstance(x, resource_variable_ops.ResourceVariable)) @@ -380,8 +217,7 @@ def _warm_start_var_with_vocab(var, full_shape=slice_info.full_shape, var_offset=slice_info.var_offset) - # TODO(eddz): Support WarmStartSettings where class vocabularies need - # remapping too. + # TODO(eddz): Support cases where class vocabularies need remapping too. init = checkpoint_ops._load_and_remap_matrix_initializer( ckpt_path=checkpoint_utils._get_checkpoint_filename(prev_ckpt), old_tensor_name=prev_tensor_name, @@ -401,30 +237,53 @@ def _warm_start_var_with_vocab(var, # pylint: enable=protected-access -def _warm_start(warm_start_settings): +@tf_export("train.warm_start") +def warm_start(ckpt_to_initialize_from, + vars_to_warm_start=".*", + var_name_to_vocab_info=None, + var_name_to_prev_var_name=None): """Warm-starts a model using the given settings. If you are using a tf.estimator.Estimator, this will automatically be called during training. Args: - warm_start_settings: An object of `WarmStartSettings`. + ckpt_to_initialize_from: [Required] A string specifying the directory with + checkpoint file(s) or path to checkpoint from which to warm-start the + model parameters. + vars_to_warm_start: [Optional] A regular expression that captures which + variables to warm-start (see tf.get_collection). Defaults to `'.*'`, + which warm-starts all variables. If `None` is explicitly given, only + variables specified in `var_name_to_vocab_info` will be warm-started. + var_name_to_vocab_info: [Optional] Dict of variable names (strings) to + VocabInfo. The variable names should be "full" variables, not the names + of the partitions. If not explicitly provided, the variable is assumed to + have no vocabulary. + var_name_to_prev_var_name: [Optional] Dict of variable names (strings) to + name of the previously-trained variable in `ckpt_to_initialize_from`. If + not explicitly provided, the name of the variable is assumed to be same + between previous checkpoint and current model. Raises: ValueError: If the WarmStartSettings contains prev_var_name or VocabInfo configuration for variable names that are not used. This is to ensure a stronger check for variable configuration than relying on users to examine the logs. """ - logging.info("Warm-starting from: %s", - (warm_start_settings.ckpt_to_initialize_from,)) + if var_name_to_vocab_info is None: + var_name_to_vocab_info = {} + if var_name_to_prev_var_name is None: + var_name_to_prev_var_name = {} + logging.info("Warm-starting from: %s", (ckpt_to_initialize_from,)) # We have to deal with partitioned variables, since get_collection flattens # out the list. grouped_variables = {} - # Both warm_start_settings.vars_to_warm_start = '.*' and - # warm_start_settings.vars_to_warm_start = None will match everything here. + # Both vars_to_warm_start = '.*' and + # vars_to_warm_start = None will match everything here. for v in ops.get_collection( + # TODO(eddz): Allow for different collections here (to support + # warm-starting accumulators). ops.GraphKeys.TRAINABLE_VARIABLES, - scope=warm_start_settings.vars_to_warm_start): + scope=vars_to_warm_start): if not isinstance(v, list): var_name = _infer_var_name([v]) else: @@ -440,10 +299,10 @@ def _warm_start(warm_start_settings): vocab_info_used = set() for var_name, variable in six.iteritems(grouped_variables): - prev_var_name = warm_start_settings.var_name_to_prev_var_name.get(var_name) + prev_var_name = var_name_to_prev_var_name.get(var_name) if prev_var_name: prev_var_name_used.add(var_name) - vocab_info = warm_start_settings.var_name_to_vocab_info.get(var_name) + vocab_info = var_name_to_vocab_info.get(var_name) if vocab_info: vocab_info_used.add(var_name) logging.info( @@ -463,16 +322,16 @@ def _warm_start(warm_start_settings): variable, current_vocab_path=vocab_info.new_vocab, current_vocab_size=vocab_info.new_vocab_size, - prev_ckpt=warm_start_settings.ckpt_to_initialize_from, + prev_ckpt=ckpt_to_initialize_from, prev_vocab_path=vocab_info.old_vocab, previous_vocab_size=vocab_info.old_vocab_size, current_oov_buckets=vocab_info.num_oov_buckets, prev_tensor_name=prev_var_name, initializer=vocab_info.backup_initializer) else: - # For the special value of warm_start_settings.vars_to_warm_start = None, + # For the special value of vars_to_warm_start = None, # we only warm-start variables with explicitly specified vocabularies. - if warm_start_settings.vars_to_warm_start: + if vars_to_warm_start: logging.info("Warm-starting variable: {}; prev_var_name: {}".format( var_name, prev_var_name or "Unchanged")) # Because we use a default empty list in grouped_variables, single @@ -480,48 +339,22 @@ def _warm_start(warm_start_settings): # for init_from_checkpoint logic to work correctly. if len(variable) == 1: variable = variable[0] - _warm_start_var(variable, warm_start_settings.ckpt_to_initialize_from, - prev_var_name) + _warm_start_var(variable, ckpt_to_initialize_from, prev_var_name) prev_var_name_not_used = set( - warm_start_settings.var_name_to_prev_var_name.keys()) - prev_var_name_used - vocab_info_not_used = set( - warm_start_settings.var_name_to_vocab_info.keys()) - vocab_info_used + var_name_to_prev_var_name.keys()) - prev_var_name_used + vocab_info_not_used = set(var_name_to_vocab_info.keys()) - vocab_info_used if prev_var_name_not_used: raise ValueError( "You provided the following variables in " - "warm_start_settings.var_name_to_prev_var_name that were not used: " + "var_name_to_prev_var_name that were not used: " "{0}. Perhaps you misspelled them? Here is the list of viable " "variable names: {1}".format(prev_var_name_not_used, grouped_variables.keys())) if vocab_info_not_used: raise ValueError( "You provided the following variables in " - "warm_start_settings.var_name_to_vocab_info that were not used: {0}. " + "var_name_to_vocab_info that were not used: {0}. " " Perhaps you misspelled them? Here is the list of viable variable " "names: {1}".format(vocab_info_not_used, grouped_variables.keys())) - - -def _get_default_warm_start_settings(warm_start_from): - """Returns default WarmStartSettings. - - Args: - warm_start_from: Either a string representing the filepath of a checkpoint - to initialize from, or an instance of WarmStartSettings. - - Returns: - Either None or an instance of WarmStartSettings. - - Raises: - ValueError: If warm_start_from is not None but is neither a string nor an - instance of WarmStartSettings. - """ - if warm_start_from is None: - return None - if isinstance(warm_start_from, six.string_types): - return WarmStartSettings(ckpt_to_initialize_from=warm_start_from) - elif isinstance(warm_start_from, WarmStartSettings): - return warm_start_from - else: - raise ValueError("warm_start_from must be a string or a WarmStartSettings") diff --git a/tensorflow/python/estimator/warm_starting_util_test.py b/tensorflow/python/training/warm_starting_util_test.py similarity index 94% rename from tensorflow/python/estimator/warm_starting_util_test.py rename to tensorflow/python/training/warm_starting_util_test.py index 3985d9ebd04e6963339fcf9999f6367fe4dadc1a..6e445d8bd14cc13010541c1ab0f737f96a4b1e03 100644 --- a/tensorflow/python/estimator/warm_starting_util_test.py +++ b/tensorflow/python/training/warm_starting_util_test.py @@ -22,7 +22,6 @@ import os import numpy as np import six -from tensorflow.python.estimator import warm_starting_util as ws_util from tensorflow.python.feature_column import feature_column as fc from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -32,6 +31,7 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test from tensorflow.python.training import saver as saver_lib +from tensorflow.python.training import warm_starting_util as ws_util ones = init_ops.ones_initializer norms = init_ops.truncated_normal_initializer @@ -330,9 +330,7 @@ class WarmStartingUtilTest(test.TestCase): with ops.Graph().as_default() as g: with self.test_session(graph=g) as sess: cols_to_vars = self._create_linear_model([sc_int], partitioner) - ws_util._warm_start( - ws_util.WarmStartSettings( - self.get_temp_dir(), vars_to_warm_start=".*sc_int.*")) + ws_util.warm_start(self.get_temp_dir(), vars_to_warm_start=".*sc_int.*") sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. self._assert_cols_to_vars(cols_to_vars, {sc_int: [prev_int_val]}, sess) @@ -361,9 +359,8 @@ class WarmStartingUtilTest(test.TestCase): with ops.Graph().as_default() as g: with self.test_session(graph=g) as sess: cols_to_vars = self._create_linear_model([sc_hash], partitioner) - ws_util._warm_start( - ws_util.WarmStartSettings( - self.get_temp_dir(), vars_to_warm_start=".*sc_hash.*")) + ws_util.warm_start( + self.get_temp_dir(), vars_to_warm_start=".*sc_hash.*") sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. self._assert_cols_to_vars(cols_to_vars, {sc_hash: [prev_hash_val]}, @@ -398,9 +395,8 @@ class WarmStartingUtilTest(test.TestCase): cols_to_vars = self._create_linear_model([sc_vocab], partitioner) # Since old vocab is not explicitly set in WarmStartSettings, the old # vocab is assumed to be same as new vocab. - ws_util._warm_start( - ws_util.WarmStartSettings( - self.get_temp_dir(), vars_to_warm_start=".*sc_vocab.*")) + ws_util.warm_start( + self.get_temp_dir(), vars_to_warm_start=".*sc_vocab.*") sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. self._assert_cols_to_vars(cols_to_vars, {sc_vocab: [prev_vocab_val]}, @@ -435,11 +431,10 @@ class WarmStartingUtilTest(test.TestCase): cols_to_vars = self._create_linear_model([sc_vocab], partitioner) # Since old vocab is not explicitly set in WarmStartSettings, the old # vocab is assumed to be same as new vocab. - ws_util._warm_start( - ws_util.WarmStartSettings( - # Explicitly provide the file prefix instead of just the dir. - os.path.join(self.get_temp_dir(), "model-0"), - vars_to_warm_start=".*sc_vocab.*")) + ws_util.warm_start( + # Explicitly provide the file prefix instead of just the dir. + os.path.join(self.get_temp_dir(), "model-0"), + vars_to_warm_start=".*sc_vocab.*") sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. self._assert_cols_to_vars(cols_to_vars, {sc_vocab: [prev_vocab_val]}, @@ -485,13 +480,12 @@ class WarmStartingUtilTest(test.TestCase): num_oov_buckets=sc_vocab.num_oov_buckets, old_vocab=old_vocab_path, old_vocab_size=old_vocab_size) - warm_start_settings = ws_util.WarmStartSettings( + ws_util.warm_start( ckpt_to_initialize_from=self.get_temp_dir(), vars_to_warm_start=".*sc_vocab.*", var_name_to_vocab_info={ "linear_model/sc_vocab/weights": vocab_info }) - ws_util._warm_start(warm_start_settings) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. 'banana' isn't in the # first two entries of the old vocabulary, so it's newly initialized. @@ -523,9 +517,8 @@ class WarmStartingUtilTest(test.TestCase): with ops.Graph().as_default() as g: with self.test_session(graph=g) as sess: cols_to_vars = self._create_linear_model([real_bucket], partitioner) - ws_util._warm_start( - ws_util.WarmStartSettings( - self.get_temp_dir(), vars_to_warm_start=".*real_bucketized.*")) + ws_util.warm_start( + self.get_temp_dir(), vars_to_warm_start=".*real_bucketized.*") sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. self._assert_cols_to_vars(cols_to_vars, @@ -606,12 +599,11 @@ class WarmStartingUtilTest(test.TestCase): new_vocab_size=sc_vocab.vocabulary_size, num_oov_buckets=sc_vocab.num_oov_buckets, old_vocab=vocab_path) - ws_util._warm_start( - ws_util.WarmStartSettings( - self.get_temp_dir(), - var_name_to_vocab_info={ - "linear_model/sc_vocab/weights": vocab_info - })) + ws_util.warm_start( + self.get_temp_dir(), + var_name_to_vocab_info={ + "linear_model/sc_vocab/weights": vocab_info + }) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. self._assert_cols_to_vars(cols_to_vars, { @@ -668,7 +660,7 @@ class WarmStartingUtilTest(test.TestCase): new_vocab_size=sc_vocab.vocabulary_size, num_oov_buckets=sc_vocab.num_oov_buckets, old_vocab=prev_vocab_path) - ws_settings = ws_util.WarmStartSettings( + ws_util.warm_start( self.get_temp_dir(), vars_to_warm_start=".*(sc_keys|sc_vocab).*", var_name_to_vocab_info={ @@ -678,7 +670,6 @@ class WarmStartingUtilTest(test.TestCase): ws_util._infer_var_name(cols_to_vars[sc_keys]): "some_other_name" }) - ws_util._warm_start(ws_settings) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. Var corresponding to # sc_hash should not be warm-started. Var corresponding to sc_vocab @@ -732,7 +723,7 @@ class WarmStartingUtilTest(test.TestCase): new_vocab_size=sc_vocab.vocabulary_size, num_oov_buckets=sc_vocab.num_oov_buckets, old_vocab=prev_vocab_path) - ws_settings = ws_util.WarmStartSettings( + ws_util.warm_start( self.get_temp_dir(), vars_to_warm_start=".*(sc_keys|sc_vocab).*", var_name_to_vocab_info={ @@ -742,7 +733,6 @@ class WarmStartingUtilTest(test.TestCase): ws_util._infer_var_name(cols_to_vars[sc_keys]): "some_other_name" }) - ws_util._warm_start(ws_settings) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. Var corresponding to # sc_hash should not be warm-started. Var corresponding to sc_vocab @@ -796,7 +786,7 @@ class WarmStartingUtilTest(test.TestCase): new_vocab_size=sc_vocab.vocabulary_size, num_oov_buckets=sc_vocab.num_oov_buckets, old_vocab=prev_vocab_path) - ws_settings = ws_util.WarmStartSettings( + ws_util.warm_start( self.get_temp_dir(), # The special value of None here will ensure that only the variable # specified in var_name_to_vocab_info (sc_vocab embedding) is @@ -812,7 +802,6 @@ class WarmStartingUtilTest(test.TestCase): ws_util._infer_var_name(cols_to_vars[sc_keys]): "some_other_name" }) - ws_util._warm_start(ws_settings) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. Var corresponding to # sc_vocab should be correctly warm-started after vocab remapping, @@ -874,13 +863,12 @@ class WarmStartingUtilTest(test.TestCase): # use a truncated normal initializer. backup_initializer=init_ops.random_uniform_initializer( minval=0.42, maxval=0.42)) - ws_settings = ws_util.WarmStartSettings( + ws_util.warm_start( self.get_temp_dir(), var_name_to_vocab_info={ ws_util._infer_var_name(cols_to_vars[emb_vocab_column]): vocab_info }) - ws_util._warm_start(ws_settings) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. Var corresponding to # emb_vocab_column should be correctly warm-started after vocab @@ -947,13 +935,12 @@ class WarmStartingUtilTest(test.TestCase): # use a truncated normal initializer. backup_initializer=init_ops.random_uniform_initializer( minval=0.42, maxval=0.42)) - ws_settings = ws_util.WarmStartSettings( + ws_util.warm_start( self.get_temp_dir(), vars_to_warm_start=".*sc_vocab.*", var_name_to_vocab_info={ "linear_model/sc_vocab_embedding/embedding_weights": vocab_info }) - ws_util._warm_start(ws_settings) sess.run(variables.global_variables_initializer()) # Verify weights were correctly warm-started. Var corresponding to # emb_vocab should be correctly warm-started after vocab remapping. @@ -973,7 +960,6 @@ class WarmStartingUtilTest(test.TestCase): }, sess) def testErrorConditions(self): - self.assertRaises(ValueError, ws_util.WarmStartSettings, None) x = variable_scope.get_variable( "x", shape=[4, 1], @@ -983,9 +969,6 @@ class WarmStartingUtilTest(test.TestCase): # List of PartitionedVariable is invalid type when warm-starting with vocab. self.assertRaises(TypeError, ws_util._warm_start_var_with_vocab, [x], "/tmp", 5, "/tmp", "/tmp") - # Keys of type other than FeatureColumn. - self.assertRaises(TypeError, ws_util._warm_start, {"StringType": x}, - ws_util.WarmStartSettings("/tmp")) # Unused variable names raises ValueError. with ops.Graph().as_default(): @@ -997,18 +980,16 @@ class WarmStartingUtilTest(test.TestCase): partitioner=lambda shape, dtype: [2, 1]) self._write_checkpoint(sess) - self.assertRaises(ValueError, ws_util._warm_start, - ws_util.WarmStartSettings( - self.get_temp_dir(), - var_name_to_vocab_info={ - "y": ws_util.VocabInfo("", 1, 0, "") - })) - self.assertRaises(ValueError, ws_util._warm_start, - ws_util.WarmStartSettings( - self.get_temp_dir(), - var_name_to_prev_var_name={ - "y": "y2" - })) + self.assertRaises( + ValueError, + ws_util.warm_start, + self.get_temp_dir(), + var_name_to_vocab_info={"y": ws_util.VocabInfo("", 1, 0, "")}) + self.assertRaises( + ValueError, + ws_util.warm_start, + self.get_temp_dir(), + var_name_to_prev_var_name={"y": "y2"}) if __name__ == "__main__": diff --git a/tensorflow/python/user_ops/user_ops.py b/tensorflow/python/user_ops/user_ops.py index 17dbab706c9243c5f119dc82cc4428f03b90a18d..6f9b5d92bb2ea662c9a5af279f0fcf71f0efccc5 100644 --- a/tensorflow/python/user_ops/user_ops.py +++ b/tensorflow/python/user_ops/user_ops.py @@ -27,4 +27,4 @@ from tensorflow.python.ops.gen_user_ops import * # pylint: disable=wildcard-imp def my_fact(): """Example of overriding the generated code for an Op.""" - return _gen_user_ops._fact() # pylint: disable=protected-access + return _gen_user_ops.fact() diff --git a/tensorflow/python/util/compat.py b/tensorflow/python/util/compat.py index 270d96a3c7c831d8c06dd86199cf2dc5dfc43421..4163fcac79e3d237c4c4c4303e1db2c39e5fe7c6 100644 --- a/tensorflow/python/util/compat.py +++ b/tensorflow/python/util/compat.py @@ -41,8 +41,11 @@ import numpy as _np import six as _six from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.tf_export import tf_export +from tensorflow.python.util.tf_export import tf_export +@tf_export('compat.as_bytes', 'compat.as_str') def as_bytes(bytes_or_text, encoding='utf-8'): """Converts either bytes or unicode to `bytes`, using utf-8 encoding for text. @@ -65,6 +68,7 @@ def as_bytes(bytes_or_text, encoding='utf-8'): (bytes_or_text,)) +@tf_export('compat.as_text') def as_text(bytes_or_text, encoding='utf-8'): """Returns the given argument as a unicode string. @@ -93,6 +97,7 @@ else: as_str = as_text +@tf_export('compat.as_str_any') def as_str_any(value): """Converts to `str` as `str(value)`, but use `as_str` for `bytes`. @@ -108,6 +113,7 @@ def as_str_any(value): return str(value) +@tf_export('compat.path_to_str') def path_to_str(path): """Returns the file system path representation of a `PathLike` object, else as it is. @@ -125,11 +131,16 @@ def path_to_str(path): # Numpy 1.8 scalars don't inherit from numbers.Integral in Python 3, so we # need to check them specifically. The same goes from Real and Complex. integral_types = (_numbers.Integral, _np.integer) +tf_export('compat.integral_types').export_constant(__name__, 'integral_types') real_types = (_numbers.Real, _np.integer, _np.floating) +tf_export('compat.real_types').export_constant(__name__, 'real_types') complex_types = (_numbers.Complex, _np.number) +tf_export('compat.complex_types').export_constant(__name__, 'complex_types') # Either bytes or text. bytes_or_text_types = (bytes, _six.text_type) +tf_export('compat.bytes_or_text_types').export_constant(__name__, + 'bytes_or_text_types') _allowed_symbols = [ 'as_str', diff --git a/tensorflow/python/util/compat_internal.py b/tensorflow/python/util/compat_internal.py index a299b2fc3c302705d9493904e8ac0f81e4b8d371..1905c3e3832550906c601bd4545e72b5bd135e2c 100644 --- a/tensorflow/python/util/compat_internal.py +++ b/tensorflow/python/util/compat_internal.py @@ -12,16 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Functions for Python 2 vs. 3 compatibility that are private to TensorFlow.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.util.compat import as_str_any + def path_to_str(path): - """Returns the file system path representation of a `PathLike` object, else as it is. + """Returns the file system path representation of a `PathLike` object, + else as it is. Args: path: An object that can be converted to path representation. diff --git a/tensorflow/python/util/decorator_utils.py b/tensorflow/python/util/decorator_utils.py index df259c7f7c29f9a4b674d3e980b33d6dcf323769..7b4363c0e40802779cf47c75c5a5e5a901da37e2 100644 --- a/tensorflow/python/util/decorator_utils.py +++ b/tensorflow/python/util/decorator_utils.py @@ -82,7 +82,7 @@ def add_notice_to_docstring( lines = _normalize_docstring(doc).splitlines() lines[0] += ' ' + suffix_str - notice = [''] + notice + [instructions] + notice = [''] + notice + ([instructions] if instructions else []) if len(lines) > 1: # Make sure that we keep our distance from the main body diff --git a/tensorflow/python/util/deprecation.py b/tensorflow/python/util/deprecation.py index fbec8fd2d8edfcd503a3677583a0853872ec08bc..376be39978fb11463ae8a870492a359c89a9f2ce 100644 --- a/tensorflow/python/util/deprecation.py +++ b/tensorflow/python/util/deprecation.py @@ -22,9 +22,9 @@ import collections import functools import re -from tensorflow.python.eager import context from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import decorator_utils +from tensorflow.python.util import is_in_graph_mode from tensorflow.python.util import tf_contextlib from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect @@ -400,7 +400,7 @@ def deprecated_args(date, instructions, *deprecated_arg_names_or_tuples, """Deprecation wrapper.""" # TODO(apassos) figure out a way to have reasonable performance with # deprecation warnings and eager mode. - if context.in_graph_mode() and _PRINT_DEPRECATION_WARNINGS: + if is_in_graph_mode.IS_IN_GRAPH_MODE() and _PRINT_DEPRECATION_WARNINGS: invalid_args = [] named_args = tf_inspect.getcallargs(func, *args, **kwargs) for arg_name, spec in iter(deprecated_positions.items()): diff --git a/tensorflow/contrib/ndlstm/python/__init__.py b/tensorflow/python/util/is_in_graph_mode.py similarity index 67% rename from tensorflow/contrib/ndlstm/python/__init__.py rename to tensorflow/python/util/is_in_graph_mode.py index 1aa51a6ec40c042ca3c26c6b08e5bdb8a42a12bd..9ae89ecb714c25787732f0d6c671d78144bec395 100644 --- a/tensorflow/contrib/ndlstm/python/__init__.py +++ b/tensorflow/python/util/is_in_graph_mode.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,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Init file, giving convenient access to all ndlstm ops.""" - +"""A function that tells you if the program is running in graph mode.""" 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 +# Call IS_IN_GRAPH_MODE() when you want to know whether the thread is in +# graph mode. By default, we always are. +IS_IN_GRAPH_MODE = lambda: True diff --git a/tensorflow/python/util/nest.py b/tensorflow/python/util/nest.py index 874df3d1087e157f8bfcec12ba3495e341c14b7b..23c2c48f4b5a165bd6e356a6243b234619af1c4c 100644 --- a/tensorflow/python/util/nest.py +++ b/tensorflow/python/util/nest.py @@ -497,7 +497,9 @@ def assert_shallow_structure(shallow_tree, input_tree, check_types=True): shallow_tree: an arbitrarily nested structure. input_tree: an arbitrarily nested structure. check_types: if `True` (default) the sequence types of `shallow_tree` and - `input_tree` have to be the same. + `input_tree` have to be the same. Note that even with check_types==True, + this function will consider two different namedtuple classes with the same + name and _fields attribute to be the same class. Raises: TypeError: If `shallow_tree` is a sequence but `input_tree` is not. @@ -513,10 +515,21 @@ def assert_shallow_structure(shallow_tree, input_tree, check_types=True): "Input has type: %s." % type(input_tree)) if check_types and not isinstance(input_tree, type(shallow_tree)): - raise TypeError( - "The two structures don't have the same sequence type. Input " - "structure has type %s, while shallow structure has type %s." - % (type(input_tree), type(shallow_tree))) + # Duck-typing means that nest should be fine with two different + # namedtuples with identical name and fields. + shallow_is_namedtuple = _is_namedtuple(shallow_tree, False) + input_is_namedtuple = _is_namedtuple(input_tree, False) + if shallow_is_namedtuple and input_is_namedtuple: + if not _same_namedtuples(shallow_tree, input_tree): + raise TypeError( + "The two namedtuples don't have the same sequence type. Input " + "structure has type %s, while shallow structure has type %s." + % (type(input_tree), type(shallow_tree))) + else: + raise TypeError( + "The two structures don't have the same sequence type. Input " + "structure has type %s, while shallow structure has type %s." + % (type(input_tree), type(shallow_tree))) if len(input_tree) != len(shallow_tree): raise ValueError( @@ -532,8 +545,8 @@ def assert_shallow_structure(shallow_tree, input_tree, check_types=True): (list(_six.iterkeys(input_tree)), list(_six.iterkeys(shallow_tree)))) - input_tree = list(_six.iteritems(input_tree)) - shallow_tree = list(_six.iteritems(shallow_tree)) + input_tree = list(sorted(_six.iteritems(input_tree))) + shallow_tree = list(sorted(_six.iteritems(shallow_tree))) for shallow_branch, input_branch in zip(shallow_tree, input_tree): assert_shallow_structure(shallow_branch, input_branch, diff --git a/tensorflow/python/util/nest_test.py b/tensorflow/python/util/nest_test.py index 6bec397db577c5be5847a701ccc92367dc008fc9..4439d6241ea9607b194cbb17304dbb77dc9f57a8 100644 --- a/tensorflow/python/util/nest_test.py +++ b/tensorflow/python/util/nest_test.py @@ -425,6 +425,19 @@ class NestTest(test.TestCase): with self.assertRaisesRegexp(ValueError, expected_message): nest.assert_shallow_structure(inp_ab2, inp_ab1) + inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))]) + inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)]) + nest.assert_shallow_structure(inp_ab, inp_ba) + + # This assertion is expected to pass: two namedtuples with the same + # name and field names are considered to be identical. + same_name_type_0 = collections.namedtuple("same_name", ("a", "b")) + same_name_type_1 = collections.namedtuple("same_name", ("a", "b")) + inp_shallow = same_name_type_0(1, 2) + inp_deep = same_name_type_1(1, [1, 2, 3]) + nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=False) + nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=True) + def testFlattenUpTo(self): # Shallow tree ends at scalar. input_tree = [[[2, 2], [3, 3]], [[4, 9], [5, 5]]] diff --git a/tensorflow/python/util/port.i b/tensorflow/python/util/port.i index cea4d8468afe8816d71da6581635b8a7ab0c2388..2f730732bee373a6e6ead97fe3320645f37ac220 100644 --- a/tensorflow/python/util/port.i +++ b/tensorflow/python/util/port.i @@ -23,5 +23,6 @@ limitations under the License. %unignore tensorflow; %unignore tensorflow::IsGoogleCudaEnabled; %unignore tensorflow::CudaSupportsHalfMatMulAndConv; +%unignore tensorflow::IsMklEnabled; %include "tensorflow/core/util/port.h" %unignoreall diff --git a/tensorflow/python/util/tf_inspect.py b/tensorflow/python/util/tf_inspect.py index c4168f7b1ac80976a957e96c79c72fe3b288d622..4ab8a72a83b466c38c50b1c76004e7a6fe942a04 100644 --- a/tensorflow/python/util/tf_inspect.py +++ b/tensorflow/python/util/tf_inspect.py @@ -46,8 +46,10 @@ def getargspec(object): # pylint: disable=redefined-builtin def getfullargspec(obj): # pylint: disable=redefined-builtin - """TFDecorator-aware replacement for inspect.getfullargspec and fallback to - inspect.getargspec in Python 2. + """TFDecorator-aware replacement for `inspect.getfullargspec`/`getargspec`. + + This wrapper uses `inspect.getfullargspec` if available and falls back to + `inspect.getargspec` in Python 2. Args: obj: A callable, possibly decorated. @@ -134,6 +136,11 @@ def getmembers(object, predicate=None): # pylint: disable=redefined-builtin return _inspect.getmembers(object, predicate) +def getmodule(object): # pylint: disable=redefined-builtin + """TFDecorator-aware replacement for inspect.getmodule.""" + return _inspect.getmodule(object) + + def getmro(cls): """TFDecorator-aware replacement for inspect.getmro.""" return _inspect.getmro(cls) @@ -144,6 +151,11 @@ def getsource(object): # pylint: disable=redefined-builtin return _inspect.getsource(tf_decorator.unwrap(object)[1]) +def isbuiltin(object): # pylint: disable=redefined-builtin + """TFDecorator-aware replacement for inspect.isbuiltin.""" + return _inspect.isbuiltin(tf_decorator.unwrap(object)[1]) + + def isclass(object): # pylint: disable=redefined-builtin """TFDecorator-aware replacement for inspect.isclass.""" return _inspect.isclass(tf_decorator.unwrap(object)[1]) diff --git a/tensorflow/python/util/tf_inspect_test.py b/tensorflow/python/util/tf_inspect_test.py index a9e8ffb30c3392251c2bf7076e02aafd2338696b..129408449ebb45ac3a322f163a13b705cbb31f0c 100644 --- a/tensorflow/python/util/tf_inspect_test.py +++ b/tensorflow/python/util/tf_inspect_test.py @@ -124,6 +124,17 @@ class TfInspectTest(test.TestCase): inspect.getmembers(TestDecoratedClass), tf_inspect.getmembers(TestDecoratedClass)) + def testGetModule(self): + self.assertEqual( + inspect.getmodule(TestDecoratedClass), + tf_inspect.getmodule(TestDecoratedClass)) + self.assertEqual( + inspect.getmodule(test_decorated_function), + tf_inspect.getmodule(test_decorated_function)) + self.assertEqual( + inspect.getmodule(test_undecorated_function), + tf_inspect.getmodule(test_undecorated_function)) + def testGetSource(self): expected = '''@test_decorator('decorator') def test_decorated_function_with_defaults(a, b=2, c='Hello'): @@ -133,6 +144,19 @@ def test_decorated_function_with_defaults(a, b=2, c='Hello'): self.assertEqual( expected, tf_inspect.getsource(test_decorated_function_with_defaults)) + def testIsBuiltin(self): + self.assertEqual( + tf_inspect.isbuiltin(TestDecoratedClass), + inspect.isbuiltin(TestDecoratedClass)) + self.assertEqual( + tf_inspect.isbuiltin(test_decorated_function), + inspect.isbuiltin(test_decorated_function)) + self.assertEqual( + tf_inspect.isbuiltin(test_undecorated_function), + inspect.isbuiltin(test_undecorated_function)) + self.assertEqual(tf_inspect.isbuiltin(range), inspect.isbuiltin(range)) + self.assertEqual(tf_inspect.isbuiltin(max), inspect.isbuiltin(max)) + def testIsClass(self): self.assertTrue(tf_inspect.isclass(TestDecoratedClass)) self.assertFalse(tf_inspect.isclass(test_decorated_function)) diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py index 37733152e8ec6d7b026bf74e69e33bfe8f9f4e89..28e49afa023904abed076373685bb38f2537b7d4 100644 --- a/tensorflow/python/util/tf_should_use.py +++ b/tensorflow/python/util/tf_should_use.py @@ -47,7 +47,7 @@ def _add_should_use_warning(x, fatal_error=False): if x is None or x == []: # pylint: disable=g-explicit-bool-comparison return x - if context.in_eager_mode(): + if context.executing_eagerly(): # Typically not needed when executing eagerly (the main use case is for ops # which need to be incorporated into the graph), and even the no-op wrapper # creates reference cycles which require garbage collection. diff --git a/tensorflow/stream_executor/blas.cc b/tensorflow/stream_executor/blas.cc index da09d84921e2dd94942b3a62fe7366211c60aed1..31724cf6c9b97e45975b9e053459f7b8f5918dfa 100644 --- a/tensorflow/stream_executor/blas.cc +++ b/tensorflow/stream_executor/blas.cc @@ -79,6 +79,8 @@ string ComputationTypeString(ComputationType ty) { return "f32"; case ComputationType::kF64: return "f64"; + case ComputationType::kI32: + return "i32"; case ComputationType::kComplexF32: return "complex f32"; case ComputationType::kComplexF64: @@ -88,6 +90,10 @@ string ComputationTypeString(ComputationType ty) { } } +std::ostream& operator<<(std::ostream& os, ComputationType ty) { + return os << ComputationTypeString(ty); +} + } // namespace blas } // namespace gputools } // namespace perftools diff --git a/tensorflow/stream_executor/blas.h b/tensorflow/stream_executor/blas.h index 072f08554688276a05d9be85718de8750bd874c2..c5f778a5c74519c0f35cea5d59aac3d0d4564c56 100644 --- a/tensorflow/stream_executor/blas.h +++ b/tensorflow/stream_executor/blas.h @@ -104,6 +104,8 @@ enum class ComputationType { // Converts a ComputationType to a string. string ComputationTypeString(ComputationType ty); +std::ostream &operator<<(std::ostream &os, ComputationType ty); + // Opaque identifier for an "algorithm" used by a blas routine. This functions // as a hint to the blas library. typedef int64 AlgorithmType; diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc index 44a3a745ad86dc24f632e4a36691fba06171c9fb..c563f8f931b0a5689268329386d1252f2a45bdd1 100644 --- a/tensorflow/stream_executor/cuda/cuda_blas.cc +++ b/tensorflow/stream_executor/cuda/cuda_blas.cc @@ -13,17 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// Include cuBLAS headers early, and then set EIGEN_HAS_CUDA_FP16 -// if we have new enough CUDA (which we will only know after including -// cuda.h). This ensures that Eigen's Half.h does not attempt to make its own -// __half typedef if CUDA has already defined one (and conversely, that we do -// not include after Half.h has made its typedef). -#include "cuda/include/cuda.h" #include "cuda/include/cublas_v2.h" - -#if CUDA_VERSION >= 7050 -#define EIGEN_HAS_CUDA_FP16 -#endif +#include "cuda/include/cuda.h" #if CUDA_VERSION >= 8000 #define SE_CUDA_DATA_HALF CUDA_R_16F @@ -33,6 +24,34 @@ limitations under the License. #include "tensorflow/stream_executor/cuda/cuda_blas.h" +// Both Eigen Half.h and CUDA cuda_fp16.h provide similar typedef for __half. As +// such, there are two ways to get the typedef for __half: +// +// (1) Includes cuda_fp16.h and defines EIGEN_HAS_CUDA_FP16. +// (2) Neither includes cuda_fp16.h nor defines EIGEN_HAS_CUDA_FP16. +// +// Due to issue b/73793421, when the first approach is used and NVCC is used to +// compile this file, NVCC will complain duplicated definition for +// EIGEN_HAS_CUDA_FP16. On the other hand, when the second approach is used and +// clang is used to compile this file, clang will not understand __half +// due to missing the definition and macro EIGEN_HAS_CUDA_FP16. +// +// Because this file may be compiled with CLANG but will never be compiled with +// NVCC, we choose the first approach for CUDA < 9.0. For CUDA >= 9.0, we have +// to use the second approach because the data member in the __half defined +// by CUDA > 9.0 is `__x` while Eigen expects it to be `x`. +// +// TODO(b/73793421): Remove the following code block to switch to the second +// approach when the issue is fixed. +#if CUDA_VERSION < 9000 +#include "cuda/include/cuda_fp16.h" +#if CUDA_VERSION >= 7050 +#define EIGEN_HAS_CUDA_FP16 +#endif +#endif + +#include "third_party/eigen3/Eigen/Core" + #include #include @@ -2256,6 +2275,14 @@ bool CUDABlas::DoBlasGemmWithAlgorithm( DeviceMemory *c, int ldc, blas::ComputationType computation_type, blas::AlgorithmType algorithm, blas::ProfileResult *output_profile_result) { + if (computation_type == blas::ComputationType::kF32) { + return DoBlasGemmWithAlgorithmImpl( + stream, transa, transb, m, n, k, static_cast(alpha), a, lda, b, + ldb, static_cast(beta), c, ldc, computation_type, algorithm, + output_profile_result); + } + + CHECK_EQ(computation_type, blas::ComputationType::kF16); return DoBlasGemmWithAlgorithmImpl( stream, transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc, computation_type, algorithm, output_profile_result); diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc index 384445e6c1629e5518459b5382aa9b92698fb6ff..03e3e0857f9f70dcd3062f0766b50f6f75b2fa5e 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -109,6 +109,24 @@ string ToString(cudnnStatus_t status) { } } +template +cudnnDataType_t GetCudnnDataType(); + +template <> +cudnnDataType_t GetCudnnDataType() { + return CUDNN_DATA_DOUBLE; +} + +template <> +cudnnDataType_t GetCudnnDataType() { + return CUDNN_DATA_FLOAT; +} + +template <> +cudnnDataType_t GetCudnnDataType() { + return CUDNN_DATA_HALF; +} + namespace wrap { static port::ThreadPool* InitCudnnThreadpool() { @@ -256,7 +274,8 @@ CUDNN_DNN_ROUTINE_EACH_R6(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) // clang-format off #if CUDNN_VERSION >= 7000 #define CUDNN_DNN_ROUTINE_EACH_R7(__macro) \ - __macro(cudnnSetConvolutionMathType) + __macro(cudnnSetConvolutionMathType) \ + __macro(cudnnSetRNNMatrixMathType) // clang-format on CUDNN_DNN_ROUTINE_EACH_R7(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) @@ -559,7 +578,7 @@ class ScopedFilterDescriptor { // A helper function to decide whether to enable the TENSOR_OP_MATH math type static bool TensorOpMathEnabled() { static bool is_enabled = [] { - bool is_disabled; + bool is_disabled = false; TF_CHECK_OK( tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_TENSOR_OP_MATH", /*default_val=*/false, &is_disabled)); @@ -568,6 +587,38 @@ static bool TensorOpMathEnabled() { return is_enabled; } +// A helper function to decide whether to enable the TENSOR_OP_MATH math type +// for RNNs. +static bool RnnTensorOpMathEnabled() { + static bool is_enabled = [] { + bool is_disabled = false; + TF_CHECK_OK( + tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH", + /*default_val=*/false, &is_disabled)); + return !is_disabled; + }(); + return is_enabled; +} + +// A helper function to decide whether to use CUDNN_BATCHNORM_SPATIAL_PERSISTENT +// in batchnorm. This mode can be faster in some tasks because an optimized path +// may be selected for CUDNN_DATA_FLOAT and CUDNN_DATA_HALF data types, compute +// capability 6.0 or higher. The reason we set it to false by default is that +// this mode may use scaled atomic integer reduction that may cause a numerical +// overflow for certain input data range. +// TODO(yangzihao): Use autotune to choose between this mode and +// CUDNN_BATCHNORM_SPATIAL mode. +static bool BatchnormSpatialPersistentEnabled() { + static bool is_enabled = [] { + bool is_enabled = false; + TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar( + "TF_USE_CUDNN_BATCHNORM_SPATIAL_PERSISTENT", + /*default_val=*/false, &is_enabled)); + return is_enabled; + }(); + return is_enabled; +} + // Turns a ConvolutionDescriptor structure into a cudnn convolution handle // within a scope. class ScopedConvolutionDescriptor { @@ -1087,6 +1138,9 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon { SetFailure(cudnn_params_desc_->Status()); return; } + if (data_type == CUDNN_DATA_HALF) { + set_use_tensor_op_math(true); + } } ~CudnnRnnDescriptor() override { if (rnn_desc_) { @@ -1095,6 +1149,20 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon { CUDNN_RETURN_IF_FAIL(status, "Unable to destroy RNN descriptor"); } } + void set_use_tensor_op_math(bool use_tensor_op_math) { +#if CUDNN_VERSION >= 7000 + cudnnMathType_t math_type = + (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH); + if (RnnTensorOpMathEnabled()) { + cudnnStatus_t status = + wrap::cudnnSetRNNMatrixMathType(parent_, rnn_desc_, math_type); + if (status != CUDNN_STATUS_SUCCESS) { + LOG(FATAL) << "could not set cudnn RNN math type: " + << ToString(status); + } + } +#endif + } cudnnRNNDescriptor_t handle() const { if (!ok()) return nullptr; return rnn_desc_; @@ -2106,7 +2174,6 @@ inline cudnnConvolutionFwdAlgo_t GetCudnnConvolutionForwardAlgo( dnn::AlgorithmDesc GetCudnnConvolutionForwardAlgorithm( Stream* stream, CUDAExecutor* parent, void* dnn_handle, - int cudnn_type, // Actually cudnnDataType_t. const dnn::AlgorithmConfig& algorithm_config, bool is_profiling, const ScopedTensorDescriptor& input_nd, const ScopedFilterDescriptor& filter, @@ -2245,7 +2312,6 @@ struct ConvDoFP32ComputationFP16Input { // A group of helper functions to return the internal compute type for // convolutions in cudnn. -// TODO(yangzihao): Add support for float64. template cudnnDataType_t GetConvComputeType() { return CUDNN_DATA_FLOAT; @@ -2260,12 +2326,17 @@ cudnnDataType_t GetConvComputeType() { } } +template <> +cudnnDataType_t GetConvComputeType() { + return CUDNN_DATA_DOUBLE; +} + } // namespace template bool CudnnSupport::DoConvolveImpl( - Stream* stream, int cudnn_type, // Actually cudnnDataType_t. - const BatchDescriptor& batch_descriptor, const DeviceMemory& input_data, + Stream* stream, const BatchDescriptor& batch_descriptor, + const DeviceMemory& input_data, const FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, const ConvolutionDescriptor& convolution_descriptor, @@ -2273,12 +2344,11 @@ bool CudnnSupport::DoConvolveImpl( ScratchAllocator* scratch_allocator, const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { - ScopedTensorDescriptor input_nd{parent_, batch_descriptor, - static_cast(cudnn_type)}; - ScopedTensorDescriptor output_nd{parent_, output_descriptor, - static_cast(cudnn_type)}; + cudnnDataType_t cudnn_type = GetCudnnDataType(); + ScopedTensorDescriptor input_nd{parent_, batch_descriptor, cudnn_type}; + ScopedTensorDescriptor output_nd{parent_, output_descriptor, cudnn_type}; ScopedFilterDescriptor filter{parent_, filter_descriptor, batch_descriptor, - static_cast(cudnn_type)}; + cudnn_type}; ScopedConvolutionDescriptor conv{parent_, convolution_descriptor, GetConvComputeType()}; @@ -2289,9 +2359,15 @@ bool CudnnSupport::DoConvolveImpl( LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } // Alpha is the scaling factor for input. - float alpha = 1.0; + float falpha = 1.0; + double dalpha = 1.0; + void* alpha = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast(&dalpha) + : static_cast(&falpha); // Beta is the scaling factor for output. - float beta = 0.0; + float fbeta = 0.0; + double dbeta = 0.0; + void* beta = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast(&dbeta) + : static_cast(&fbeta); const bool is_profiling = output_profile_result != nullptr; cudnnConvolutionFwdAlgo_t algo; @@ -2429,11 +2505,11 @@ bool CudnnSupport::DoConvolveImpl( } status = wrap::cudnnConvolutionForward( parent_, ToHandle(dnn_handle_), - /*alpha=*/&alpha, /*srcDesc=*/input_nd.handle(), + /*alpha=*/alpha, /*srcDesc=*/input_nd.handle(), /*srcData=*/input_data.opaque(), /*filterDesc=*/filter.handle(), /*filterData=*/filter_data.opaque(), /*convDesc=*/conv.handle(), /*algo=*/algo, /*workSpace=*/scratch.opaque(), - /*workSpaceSizeInBytes=*/scratch.size(), /*beta=*/&beta, + /*workSpaceSizeInBytes=*/scratch.size(), /*beta=*/beta, /*destDesc=*/output_nd.handle(), /*destData=*/output_data->opaque()); if (is_profiling) { @@ -2505,9 +2581,8 @@ bool CudnnSupport::DoFusedConvolveImpl( const bool is_profiling = output_profile_result != nullptr; DeviceMemory scratch; dnn::AlgorithmDesc algotype = GetCudnnConvolutionForwardAlgorithm( - stream, parent_, dnn_handle_, cudnn_data_type, algorithm_config, - is_profiling, conv_input_nd, filter, conv, output_nd, scratch_allocator, - &scratch); + stream, parent_, dnn_handle_, algorithm_config, is_profiling, + conv_input_nd, filter, conv, output_nd, scratch_allocator, &scratch); if (algotype.is_default()) { if (!is_profiling) { LOG(ERROR) << "No suitable algorithm found"; @@ -2758,6 +2833,11 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl( ScopedTensorDescriptor scale_offset_descriptor{ parent_, scale_offset_desc, ToCudnnDataType(scale_data_type)}; cudnnBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL; +#if CUDNN_VERSION >= 7000 + if (BatchnormSpatialPersistentEnabled() && is_training) { + mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; + } +#endif float one = 1.0; float zero = 0.0; @@ -2859,6 +2939,11 @@ bool CudnnSupport::DoBatchNormalizationBackwardImpl( parent_, scale_offset_desc, static_cast(cudnn_scale_type)}; cudnnBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL; +#if CUDNN_VERSION >= 7000 + if (BatchnormSpatialPersistentEnabled()) { + mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; + } +#endif float one = 1.0; float zero = 0.0; @@ -2888,9 +2973,9 @@ bool CudnnSupport::DoConvolve( const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { return DoConvolveImpl( - stream, CUDNN_DATA_FLOAT, batch_descriptor, input_data, filter_descriptor, - filter_data, convolution_descriptor, output_descriptor, output_data, - scratch_allocator, algorithm_config, output_profile_result); + stream, batch_descriptor, input_data, filter_descriptor, filter_data, + convolution_descriptor, output_descriptor, output_data, scratch_allocator, + algorithm_config, output_profile_result); } bool CudnnSupport::DoConvolve( @@ -2899,10 +2984,14 @@ bool CudnnSupport::DoConvolve( const FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, const ConvolutionDescriptor& convolution_descriptor, - const BatchDescriptor& output_descriptor, - DeviceMemory* output_data) { - LOG(ERROR) << "double-based DNN not yet implemented"; - return false; + const BatchDescriptor& output_descriptor, DeviceMemory* output_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) { + return DoConvolveImpl( + stream, batch_descriptor, input_data, filter_descriptor, filter_data, + convolution_descriptor, output_descriptor, output_data, scratch_allocator, + algorithm_config, output_profile_result); } bool CudnnSupport::DoConvolve( @@ -2916,9 +3005,9 @@ bool CudnnSupport::DoConvolve( const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { return DoConvolveImpl( - stream, CUDNN_DATA_HALF, batch_descriptor, input_data, filter_descriptor, - filter_data, convolution_descriptor, output_descriptor, output_data, - scratch_allocator, algorithm_config, output_profile_result); + stream, batch_descriptor, input_data, filter_descriptor, filter_data, + convolution_descriptor, output_descriptor, output_data, scratch_allocator, + algorithm_config, output_profile_result); } bool CudnnSupport::DoFusedConvolve( @@ -3027,7 +3116,6 @@ bool CudnnSupport::DoFusedConvolve( template DeviceMemory CudnnSupport::MaybeTransformLayout( Stream* stream, - int cudnn_type, // Actually cudnnDataType_t. BatchDescriptor* output_descriptor, DeviceMemory backward_output_data, std::unique_ptr>* transform_scratch) { @@ -3041,11 +3129,11 @@ DeviceMemory CudnnSupport::MaybeTransformLayout( BatchDescriptor transformed_output_descriptor; transformed_output_descriptor.CloneFrom(*output_descriptor); transformed_output_descriptor.set_layout(dnn::DataLayout::kBatchDepthYX); - ScopedTensorDescriptor orig_out_back_nd{ - parent_, *output_descriptor, static_cast(cudnn_type)}; + cudnnDataType_t cudnn_type = GetCudnnDataType(); + ScopedTensorDescriptor orig_out_back_nd{parent_, *output_descriptor, + cudnn_type}; ScopedTensorDescriptor transformed_out_back_nd{ - parent_, transformed_output_descriptor, - static_cast(cudnn_type)}; + parent_, transformed_output_descriptor, cudnn_type}; float alpha = 1.0f; float beta = 0.0f; @@ -3092,7 +3180,6 @@ bool CudnnSupport::DoTransformTensor(Stream* stream, template bool CudnnSupport::DoConvolveBackwardDataImpl( Stream* stream, - int cudnn_type, // Actually cudnnDataType_t. const FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, const BatchDescriptor& output_descriptor_in, @@ -3109,25 +3196,29 @@ bool CudnnSupport::DoConvolveBackwardDataImpl( LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } + cudnnDataType_t cudnn_type = GetCudnnDataType(); // Alpha is the scaling factor for input. - float alpha = 1.0; + float falpha = 1.0; + double dalpha = 1.0; + void* alpha = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast(&dalpha) + : static_cast(&falpha); // Beta is the scaling factor for output. - float beta = 0.0; + float fbeta = 0.0; + double dbeta = 0.0; + void* beta = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast(&dbeta) + : static_cast(&fbeta); // TBD(keveman): remove once cuDNN supports kBatchYXDepth for backward pass. BatchDescriptor output_descriptor; output_descriptor.CloneFrom(output_descriptor_in); std::unique_ptr> transform_scratch; backward_output_data = MaybeTransformLayout( - stream, cudnn_type, &output_descriptor, backward_output_data, - &transform_scratch); + stream, &output_descriptor, backward_output_data, &transform_scratch); - ScopedTensorDescriptor out_back_nd{parent_, output_descriptor, - static_cast(cudnn_type)}; - ScopedTensorDescriptor in_back_nd{parent_, input_descriptor, - static_cast(cudnn_type)}; + ScopedTensorDescriptor out_back_nd{parent_, output_descriptor, cudnn_type}; + ScopedTensorDescriptor in_back_nd{parent_, input_descriptor, cudnn_type}; ScopedFilterDescriptor filter{parent_, filter_descriptor, input_descriptor, - static_cast(cudnn_type)}; + cudnn_type}; ScopedConvolutionDescriptor conv{parent_, convolution_descriptor, GetConvComputeType()}; @@ -3270,7 +3361,7 @@ bool CudnnSupport::DoConvolveBackwardDataImpl( status = wrap::cudnnConvolutionBackwardData_v3( #endif parent_, ToHandle(dnn_handle_), - /*alpha=*/&alpha, + /*alpha=*/alpha, /*filterDesc=*/filter.handle(), /*filterData=*/filter_data.opaque(), /*diffDesc=*/out_back_nd.handle(), @@ -3279,7 +3370,7 @@ bool CudnnSupport::DoConvolveBackwardDataImpl( /*algo=*/algo, /*workSpace=*/scratch.opaque(), /*workSpaceSizeInBytes=*/scratch.size(), - /*beta=*/&beta, + /*beta=*/beta, /*gradDesc=*/in_back_nd.handle(), /*gradData=*/backward_input_data->opaque()); if (is_profiling) { @@ -3304,10 +3395,28 @@ bool CudnnSupport::DoConvolveBackwardDataImpl( return true; } +bool CudnnSupport::DoConvolveBackwardData( + Stream* stream, const FilterDescriptor& filter_descriptor, + const DeviceMemory& filter_data, + const BatchDescriptor& output_descriptor, + DeviceMemory backward_output_data, + const ConvolutionDescriptor& convolution_descriptor, + const BatchDescriptor& input_descriptor, + DeviceMemory* backward_input_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) { + return DoConvolveBackwardDataImpl(stream, filter_descriptor, filter_data, + output_descriptor, backward_output_data, + convolution_descriptor, input_descriptor, + backward_input_data, scratch_allocator, + algorithm_config, output_profile_result); +} + bool CudnnSupport::DoConvolveBackwardData( Stream* stream, const FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, - const BatchDescriptor& output_descriptor_in, + const BatchDescriptor& output_descriptor, DeviceMemory backward_output_data, const ConvolutionDescriptor& convolution_descriptor, const BatchDescriptor& input_descriptor, @@ -3315,17 +3424,17 @@ bool CudnnSupport::DoConvolveBackwardData( ScratchAllocator* scratch_allocator, const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { - return DoConvolveBackwardDataImpl( - stream, CUDNN_DATA_FLOAT, filter_descriptor, filter_data, - output_descriptor_in, backward_output_data, convolution_descriptor, - input_descriptor, backward_input_data, scratch_allocator, - algorithm_config, output_profile_result); + return DoConvolveBackwardDataImpl(stream, filter_descriptor, filter_data, + output_descriptor, backward_output_data, + convolution_descriptor, input_descriptor, + backward_input_data, scratch_allocator, + algorithm_config, output_profile_result); } bool CudnnSupport::DoConvolveBackwardData( Stream* stream, const FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, - const BatchDescriptor& output_descriptor_in, + const BatchDescriptor& output_descriptor, DeviceMemory backward_output_data, const ConvolutionDescriptor& convolution_descriptor, const BatchDescriptor& input_descriptor, @@ -3333,17 +3442,16 @@ bool CudnnSupport::DoConvolveBackwardData( ScratchAllocator* scratch_allocator, const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { - return DoConvolveBackwardDataImpl( - stream, CUDNN_DATA_HALF, filter_descriptor, filter_data, - output_descriptor_in, backward_output_data, convolution_descriptor, - input_descriptor, backward_input_data, scratch_allocator, - algorithm_config, output_profile_result); + return DoConvolveBackwardDataImpl(stream, filter_descriptor, filter_data, + output_descriptor, backward_output_data, + convolution_descriptor, input_descriptor, + backward_input_data, scratch_allocator, + algorithm_config, output_profile_result); } template bool CudnnSupport::DoConvolveBackwardFilterImpl( - Stream* stream, int cudnn_type, // Actually cudnnDataType_t. - const dnn::BatchDescriptor& input_descriptor, + Stream* stream, const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, const dnn::BatchDescriptor& output_descriptor_in, DeviceMemory backward_output_data, @@ -3359,26 +3467,29 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl( LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } + cudnnDataType_t cudnn_type = GetCudnnDataType(); // Alpha is the scaling factor for input. - float alpha = 1.0; + float falpha = 1.0; + double dalpha = 1.0; + void* alpha = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast(&dalpha) + : static_cast(&falpha); // Beta is the scaling factor for output. - float beta = 0.0; + float fbeta = 0.0; + double dbeta = 0.0; + void* beta = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast(&dbeta) + : static_cast(&fbeta); // TBD(keveman): remove once cuDNN supports kBatchYXDepth for backward pass. BatchDescriptor output_descriptor; output_descriptor.CloneFrom(output_descriptor_in); std::unique_ptr> transform_scratch; backward_output_data = MaybeTransformLayout( - stream, static_cast(cudnn_type), - &output_descriptor, backward_output_data, - &transform_scratch); - - ScopedTensorDescriptor out_back_nd{parent_, output_descriptor, - static_cast(cudnn_type)}; - ScopedTensorDescriptor input_nd{parent_, input_descriptor, - static_cast(cudnn_type)}; + stream, &output_descriptor, backward_output_data, &transform_scratch); + + ScopedTensorDescriptor out_back_nd{parent_, output_descriptor, cudnn_type}; + ScopedTensorDescriptor input_nd{parent_, input_descriptor, cudnn_type}; ScopedFilterDescriptor filter{parent_, filter_descriptor, input_descriptor, - static_cast(cudnn_type)}; + cudnn_type}; ScopedConvolutionDescriptor conv{parent_, convolution_descriptor, GetConvComputeType()}; @@ -3521,7 +3632,7 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl( #else status = wrap::cudnnConvolutionBackwardFilter_v3( #endif - parent_, ToHandle(dnn_handle_), /*alpha=*/&alpha, + parent_, ToHandle(dnn_handle_), /*alpha=*/alpha, /*srcDesc=*/input_nd.handle(), /*srcData=*/input_data.opaque(), /*diffDesc=*/out_back_nd.handle(), @@ -3530,7 +3641,7 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl( /*algo=*/algo, /*workSpace=*/scratch.opaque(), /*workSpaceSizeInBytes=*/scratch.size(), - /*beta=*/&beta, + /*beta=*/beta, /*gradDesc=*/filter.handle(), /*gradData=*/backward_filter_data->opaque()); @@ -3556,10 +3667,28 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl( return true; } +bool CudnnSupport::DoConvolveBackwardFilter( + Stream* stream, const dnn::BatchDescriptor& input_descriptor, + const DeviceMemory& input_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + const dnn::FilterDescriptor& filter_descriptor, + DeviceMemory* backward_filter_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) { + return DoConvolveBackwardFilterImpl(stream, input_descriptor, input_data, + output_descriptor, backward_output_data, + convolution_descriptor, filter_descriptor, + backward_filter_data, scratch_allocator, + algorithm_config, output_profile_result); +} + bool CudnnSupport::DoConvolveBackwardFilter( Stream* stream, const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, - const dnn::BatchDescriptor& output_descriptor_in, + const dnn::BatchDescriptor& output_descriptor, DeviceMemory backward_output_data, const dnn::ConvolutionDescriptor& convolution_descriptor, const dnn::FilterDescriptor& filter_descriptor, @@ -3567,17 +3696,17 @@ bool CudnnSupport::DoConvolveBackwardFilter( ScratchAllocator* scratch_allocator, const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { - return DoConvolveBackwardFilterImpl( - stream, CUDNN_DATA_FLOAT, input_descriptor, input_data, - output_descriptor_in, backward_output_data, convolution_descriptor, - filter_descriptor, backward_filter_data, scratch_allocator, - algorithm_config, output_profile_result); + return DoConvolveBackwardFilterImpl(stream, input_descriptor, input_data, + output_descriptor, backward_output_data, + convolution_descriptor, filter_descriptor, + backward_filter_data, scratch_allocator, + algorithm_config, output_profile_result); } bool CudnnSupport::DoConvolveBackwardFilter( Stream* stream, const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, - const dnn::BatchDescriptor& output_descriptor_in, + const dnn::BatchDescriptor& output_descriptor, DeviceMemory backward_output_data, const dnn::ConvolutionDescriptor& convolution_descriptor, const dnn::FilterDescriptor& filter_descriptor, @@ -3585,17 +3714,16 @@ bool CudnnSupport::DoConvolveBackwardFilter( ScratchAllocator* scratch_allocator, const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { - return DoConvolveBackwardFilterImpl( - stream, CUDNN_DATA_HALF, input_descriptor, input_data, - output_descriptor_in, backward_output_data, convolution_descriptor, - filter_descriptor, backward_filter_data, scratch_allocator, - algorithm_config, output_profile_result); + return DoConvolveBackwardFilterImpl(stream, input_descriptor, input_data, + output_descriptor, backward_output_data, + convolution_descriptor, filter_descriptor, + backward_filter_data, scratch_allocator, + algorithm_config, output_profile_result); } template bool CudnnSupport::DoConvolveBackwardBiasImpl( - Stream* stream, int cudnn_type, // Actually cudnnDataType_t. - const dnn::BatchDescriptor& input_descriptor, + Stream* stream, const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, const dnn::BatchDescriptor& bias_descriptor, DeviceMemory* backward_bias_data) { @@ -3606,10 +3734,9 @@ bool CudnnSupport::DoConvolveBackwardBiasImpl( LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } - ScopedTensorDescriptor input_nd{parent_, input_descriptor, - static_cast(cudnn_type)}; - ScopedTensorDescriptor bias_nd{parent_, bias_descriptor, - static_cast(cudnn_type)}; + cudnnDataType_t cudnn_type = GetCudnnDataType(); + ScopedTensorDescriptor input_nd{parent_, input_descriptor, cudnn_type}; + ScopedTensorDescriptor bias_nd{parent_, bias_descriptor, cudnn_type}; // Alpha is the scaling factor for input. float alpha = 1.0; @@ -3633,9 +3760,8 @@ bool CudnnSupport::DoConvolveBackwardBias( const DeviceMemory& input_data, const BatchDescriptor& bias_descriptor, DeviceMemory* backward_bias_data) { - return DoConvolveBackwardBiasImpl(stream, CUDNN_DATA_DOUBLE, input_descriptor, - input_data, bias_descriptor, - backward_bias_data); + return DoConvolveBackwardBiasImpl(stream, input_descriptor, input_data, + bias_descriptor, backward_bias_data); } bool CudnnSupport::DoConvolveBackwardBias( @@ -3643,9 +3769,8 @@ bool CudnnSupport::DoConvolveBackwardBias( const DeviceMemory& input_data, const BatchDescriptor& bias_descriptor, DeviceMemory* backward_bias_data) { - return DoConvolveBackwardBiasImpl(stream, CUDNN_DATA_FLOAT, input_descriptor, - input_data, bias_descriptor, - backward_bias_data); + return DoConvolveBackwardBiasImpl(stream, input_descriptor, input_data, + bias_descriptor, backward_bias_data); } bool CudnnSupport::DoConvolveBackwardBias( @@ -3653,9 +3778,8 @@ bool CudnnSupport::DoConvolveBackwardBias( const DeviceMemory& input_data, const BatchDescriptor& bias_descriptor, DeviceMemory* backward_bias_data) { - return DoConvolveBackwardBiasImpl(stream, CUDNN_DATA_HALF, input_descriptor, - input_data, bias_descriptor, - backward_bias_data); + return DoConvolveBackwardBiasImpl(stream, input_descriptor, input_data, + bias_descriptor, backward_bias_data); } bool CudnnSupport::DoMatMul(Stream* stream, diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.h b/tensorflow/stream_executor/cuda/cuda_dnn.h index ee28c0bf57a51a63be7ebbce5c8f80e09737bb16..48d56f71e3195a897b6216ab9f5709326d1b86d3 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.h +++ b/tensorflow/stream_executor/cuda/cuda_dnn.h @@ -259,7 +259,10 @@ class CudnnSupport : public dnn::DnnSupport { const DeviceMemory& filter_data, const dnn::ConvolutionDescriptor& convolution_descriptor, const dnn::BatchDescriptor& output_descriptor, - DeviceMemory* output_data) override; + DeviceMemory* output_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) override; bool DoConvolve(Stream* stream, const dnn::BatchDescriptor& batch_descriptor, const DeviceMemory& input_data, @@ -371,6 +374,18 @@ class CudnnSupport : public dnn::DnnSupport { return false; } + bool DoConvolveBackwardData( + Stream* stream, const dnn::FilterDescriptor& filter_descriptor, + const DeviceMemory& filter_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + const dnn::BatchDescriptor& input_descriptor, + DeviceMemory* backward_input_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) override; + bool DoConvolveBackwardData( Stream* stream, const dnn::FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, @@ -395,6 +410,18 @@ class CudnnSupport : public dnn::DnnSupport { const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) override; + bool DoConvolveBackwardFilter( + Stream* stream, const dnn::BatchDescriptor& input_descriptor, + const DeviceMemory& input_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + const dnn::FilterDescriptor& filter_descriptor, + DeviceMemory* backward_filter_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) override; + bool DoConvolveBackwardFilter( Stream* stream, const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, @@ -611,7 +638,6 @@ class CudnnSupport : public dnn::DnnSupport { template DeviceMemory MaybeTransformLayout( Stream* stream, - int cudnn_type, // Actually cudnnDataType_t. dnn::BatchDescriptor* output_descriptor, DeviceMemory backward_output_data, std::unique_ptr>* transform_scratch) @@ -644,7 +670,6 @@ class CudnnSupport : public dnn::DnnSupport { template bool DoConvolveImpl(Stream* stream, - int cudnn_type, // Actually cudnnDataType_t. const dnn::BatchDescriptor& batch_descriptor, const DeviceMemory& input_data, const dnn::FilterDescriptor& filter_descriptor, @@ -675,7 +700,6 @@ class CudnnSupport : public dnn::DnnSupport { template bool DoConvolveBackwardDataImpl( Stream* stream, - int cudnn_type, // Actually cudnnDataType_t. const dnn::FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, const dnn::BatchDescriptor& output_descriptor, @@ -688,8 +712,7 @@ class CudnnSupport : public dnn::DnnSupport { template bool DoConvolveBackwardFilterImpl( - Stream* stream, int cudnn_type, // Actually cudnnDataType_t. - const dnn::BatchDescriptor& input_descriptor, + Stream* stream, const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, const dnn::BatchDescriptor& output_descriptor_in, DeviceMemory backward_output_data, @@ -702,7 +725,6 @@ class CudnnSupport : public dnn::DnnSupport { template bool DoConvolveBackwardBiasImpl(Stream* stream, - int cudnn_type, // Actually cudnnDataType_t. const dnn::BatchDescriptor& input_descriptor, const DeviceMemory& input_data, const dnn::BatchDescriptor& bias_descriptor, diff --git a/tensorflow/stream_executor/cuda/cuda_driver.cc b/tensorflow/stream_executor/cuda/cuda_driver.cc index a017ff64d4c69b6952b442464877dc26a800ad37..58e1e58c593a3d938d97baff2356bce2c215a7a1 100644 --- a/tensorflow/stream_executor/cuda/cuda_driver.cc +++ b/tensorflow/stream_executor/cuda/cuda_driver.cc @@ -1503,6 +1503,19 @@ static port::StatusOr GetSimpleAttribute(CUdevice device, return true; } +/* static */ port::StatusOr CUDADriver::GetDeviceAttribute( + CUdevice_attribute attribute, CUdevice device) { + int val; + CUresult res = cuDeviceGetAttribute(&val, attribute, device); + if (res != CUDA_SUCCESS) { + return port::Status{ + port::error::INTERNAL, + port::Printf("failed to get device attribute %d for device %d: %s", + attribute, device, ToString(res).c_str())}; + } + return val; +} + /* static */ bool CUDADriver::IsEccEnabled(CUdevice device, bool *result) { int value = -1; CUresult res = diff --git a/tensorflow/stream_executor/cuda/cuda_driver.h b/tensorflow/stream_executor/cuda/cuda_driver.h index 4002ba2021d1a2e2c36bd1786a3084ee8c08bb78..fa9172b3f008d3083309126bbfa4a1ab961030e1 100644 --- a/tensorflow/stream_executor/cuda/cuda_driver.h +++ b/tensorflow/stream_executor/cuda/cuda_driver.h @@ -400,12 +400,20 @@ class CUDADriver { // Returns a grab-bag of device properties in a caller-owned device_properties // structure for device_ordinal via cuDeviceGetProperties. - // This call is deprecated in the NVIDIA driver API. + // + // This call is deprecated in the NVIDIA driver API; its replacement is + // GetDeviceAttribute // // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html#group__CUDA__DEVICE__DEPRECATED_1g65a5b4e25186bd257df80b98c98cffe6 static bool GetDeviceProperties(CUdevprop *device_properties, int device_ordinal); + // Gets a specific integer-valued property about the given device. + // + // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE.html#group__CUDA__DEVICE_1g9c3e1414f0ad901d3278a4d6645fc266 + static port::StatusOr GetDeviceAttribute(CUdevice_attribute attribute, + CUdevice device); + // Returns whether ECC is enabled for the given CUdevice via // cuDeviceGetattribute with CU_DEVICE_ATTRIBUTE_ECC_ENABLED. // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE.html#group__CUDA__DEVICE_1g9c3e1414f0ad901d3278a4d6645fc266 diff --git a/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc b/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc index 4bbd531e14f18fc24d87b4fa655fe72e9f56b129..5ecaf46b8cae3c1e1f312816e7e5aec8ff8ce306 100644 --- a/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc +++ b/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc @@ -1103,6 +1103,18 @@ DeviceDescription *CUDAExecutor::PopulateDeviceDescription() const { builder.set_device_memory_size(device_memory_size); } + port::StatusOr mem_clock_khz = CUDADriver::GetDeviceAttribute( + CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE, device_ordinal_); + port::StatusOr mem_bus_width_bits = CUDADriver::GetDeviceAttribute( + CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH, device_ordinal_); + if (mem_clock_khz.ok() && mem_bus_width_bits.ok()) { + // Times 2 because HBM is DDR memory; it gets two data bits per each data + // lane. + builder.set_memory_bandwidth(2 * int64_t{mem_clock_khz.ValueOrDie()} * + 1000 * + int64_t{mem_bus_width_bits.ValueOrDie()} / 8); + } + { BlockDim block_dim_limit; FillBlockDimLimit(&block_dim_limit); diff --git a/tensorflow/stream_executor/device_description.cc b/tensorflow/stream_executor/device_description.cc index a98143e34bbb42c3aee76c27e1648c49397a0e44..52f5319a3b16c771ce89843a963841b25df5467e 100644 --- a/tensorflow/stream_executor/device_description.cc +++ b/tensorflow/stream_executor/device_description.cc @@ -50,6 +50,7 @@ DeviceDescription::DeviceDescription() shared_memory_alloc_granularity_(1), device_address_bits_(kUninitializedUint64), device_memory_size_(kUninitializedUint64), + memory_bandwidth_(kUninitializedUint64), shared_memory_per_core_(kUninitializedUint64), shared_memory_per_block_(kUninitializedUint64), clock_rate_ghz_(-1.0), @@ -85,6 +86,8 @@ std::unique_ptr> DeviceDescription::ToMap() const { result["Device Address Bits"] = port::StrCat(device_address_bits()); result["Device Memory Size"] = port::HumanReadableNumBytes::ToString(device_memory_size()); + result["Memory Bandwidth"] = port::StrCat( + port::HumanReadableNumBytes::ToString(memory_bandwidth_), "/s"); result["Shared Memory Per Core"] = port::HumanReadableNumBytes::ToString(shared_memory_per_core_); diff --git a/tensorflow/stream_executor/device_description.h b/tensorflow/stream_executor/device_description.h index f2b35bcb4345a37f72541979564cbbb7944595c2..fcf0928096ed1f1bdf0499efb92af2bc9cb0eaa2 100644 --- a/tensorflow/stream_executor/device_description.h +++ b/tensorflow/stream_executor/device_description.h @@ -140,6 +140,11 @@ class DeviceDescription { // Returns the device memory size in bytes. uint64 device_memory_size() const { return device_memory_size_; } + // Returns the device's memory bandwidth in bytes/sec. (This is for + // reads/writes to/from the device's own memory, not for transfers between the + // host and device.) + uint64 memory_bandwidth() const { return memory_bandwidth_; } + // Returns the device's core clock rate in GHz. float clock_rate_ghz() const { return clock_rate_ghz_; } @@ -212,6 +217,7 @@ class DeviceDescription { uint64 device_address_bits_; uint64 device_memory_size_; + uint64 memory_bandwidth_; // Shared memory limits on a given device. uint64 shared_memory_per_core_; @@ -305,6 +311,9 @@ class DeviceDescriptionBuilder { void set_device_memory_size(uint64 value) { device_description_->device_memory_size_ = value; } + void set_memory_bandwidth(uint64 value) { + device_description_->memory_bandwidth_ = value; + } void set_shared_memory_per_core(int64 value) { device_description_->shared_memory_per_core_ = value; diff --git a/tensorflow/stream_executor/dnn.h b/tensorflow/stream_executor/dnn.h index f4162b096299ca9405e1f3045e370d0da1acf8da..b41536e638873412a31a0cdbbd3ba3a818dd9cf2 100644 --- a/tensorflow/stream_executor/dnn.h +++ b/tensorflow/stream_executor/dnn.h @@ -896,7 +896,7 @@ class DnnSupport { // offset: offset parameters. // estimated_mean: population mean estimated during training. // Used for inference only; empty for training. - // estimated_variance: population variance estimated during traning, + // estimated_variance: population variance estimated during training, // used for inference only; empty for training. // x_desc: dimensions of the input data, which is the same as the dimensions // of the output. @@ -1172,7 +1172,9 @@ class DnnSupport { const DeviceMemory& filter_data, const dnn::ConvolutionDescriptor& convolution_descriptor, const dnn::BatchDescriptor& output_descriptor, - DeviceMemory* output_data) = 0; + DeviceMemory* output_data, ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + dnn::ProfileResult* output_profile_result) = 0; // Enqueues a half-precision convolution operation onto the stream. // See DoConvolve above for argument details. @@ -1273,6 +1275,18 @@ class DnnSupport { bool with_winograd_nonfused, int cc_major, int cc_minor, std::vector* out_algorithms); + virtual bool DoConvolveBackwardData( + Stream* stream, const FilterDescriptor& filter_descriptor, + const DeviceMemory& filter_data, + const BatchDescriptor& output_descriptor, + DeviceMemory backward_output_data, + const ConvolutionDescriptor& convolution_descriptor, + const BatchDescriptor& input_descriptor, + DeviceMemory* backward_input_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + ProfileResult* output_profile_result) = 0; + virtual bool DoConvolveBackwardData( Stream* stream, const FilterDescriptor& filter_descriptor, const DeviceMemory& filter_data, @@ -1322,6 +1336,18 @@ class DnnSupport { bool with_winograd_nonfused, int cc_major, int cc_minor, std::vector* out_algorithms); + virtual bool DoConvolveBackwardFilter( + Stream* stream, const BatchDescriptor& input_descriptor, + const DeviceMemory& input_data, + const BatchDescriptor& output_descriptor, + DeviceMemory backward_output_data, + const ConvolutionDescriptor& convolution_descriptor, + const FilterDescriptor& filter_descriptor, + DeviceMemory* backward_filter_data, + ScratchAllocator* scratch_allocator, + const dnn::AlgorithmConfig& algorithm_config, + ProfileResult* output_profile_result) = 0; + virtual bool DoConvolveBackwardFilter( Stream* stream, const BatchDescriptor& input_descriptor, const DeviceMemory& input_data, diff --git a/tensorflow/stream_executor/dso_loader.cc b/tensorflow/stream_executor/dso_loader.cc index d71938634d6e6fe092d9a1e0861215bb101e824f..95168836278add5d6592ff0c3d0f7245e6f6bc5b 100644 --- a/tensorflow/stream_executor/dso_loader.cc +++ b/tensorflow/stream_executor/dso_loader.cc @@ -33,6 +33,10 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/platform/port.h" +#if !defined(PLATFORM_GOOGLE) +#include "cuda/cuda_config.h" +#endif + namespace perftools { namespace gputools { namespace internal { @@ -97,11 +101,12 @@ string GetCudnnVersion() { return TF_CUDNN_VERSION; } /* static */ port::Status DsoLoader::GetLibcuptiDsoHandle(void** dso_handle) { #if defined(ANDROID_TEGRA) - // On Android devices the CUDA version number is not added to the library name. - return GetDsoHandle(FindDsoPath(port::Env::Default()->FormatLibraryFileName( - "cupti", ""), - GetCudaCuptiLibraryPath()), - dso_handle); + // On Android devices the CUDA version number is not added to the library + // name. + return GetDsoHandle( + FindDsoPath(port::Env::Default()->FormatLibraryFileName("cupti", ""), + GetCudaCuptiLibraryPath()), + dso_handle); #else return GetDsoHandle(FindDsoPath(port::Env::Default()->FormatLibraryFileName( "cupti", GetCudaVersion()), diff --git a/tensorflow/stream_executor/dso_loader.h b/tensorflow/stream_executor/dso_loader.h index 9495f7253a1d475f0b5321b71419febd086832af..354c7b50b8209755991827b3c36afac790cb952b 100644 --- a/tensorflow/stream_executor/dso_loader.h +++ b/tensorflow/stream_executor/dso_loader.h @@ -28,10 +28,6 @@ limitations under the License. #include "tensorflow/stream_executor/platform.h" #include "tensorflow/stream_executor/platform/mutex.h" -#if !defined(PLATFORM_GOOGLE) -#include "cuda/cuda_config.h" -#endif - namespace perftools { namespace gputools { namespace internal { diff --git a/tensorflow/stream_executor/multi_platform_manager.cc b/tensorflow/stream_executor/multi_platform_manager.cc index f23224ae772b9c5915426feaef1155fc9711f075..f9f3737a06dad3f146ef9fc8e2ec50160b3a01b5 100644 --- a/tensorflow/stream_executor/multi_platform_manager.cc +++ b/tensorflow/stream_executor/multi_platform_manager.cc @@ -23,11 +23,37 @@ limitations under the License. namespace perftools { namespace gputools { +/* static */ mutex MultiPlatformManager::platforms_mutex_{LINKER_INITIALIZED}; + +/* static */ port::StatusOr MultiPlatformManager::LookupByNameLocked( + const string& target) { + PlatformMap* platform_map = GetPlatformMap(); + auto it = platform_map->find(port::Lowercase(target)); + if (it == platform_map->end()) { + return port::Status( + port::error::NOT_FOUND, + "could not find registered platform with name: \"" + target + "\""); + } + return it->second; +} + +/* static */ port::StatusOr MultiPlatformManager::LookupByIdLocked( + const Platform::Id& id) { + PlatformIdMap* platform_map = GetPlatformByIdMap(); + auto it = platform_map->find(id); + if (it == platform_map->end()) { + return port::Status( + port::error::NOT_FOUND, + port::Printf("could not find registered platform with id: 0x%p", id)); + } + return it->second; +} + /* static */ port::Status MultiPlatformManager::RegisterPlatform( std::unique_ptr platform) { CHECK(platform != nullptr); string key = port::Lowercase(platform->Name()); - mutex_lock lock(GetPlatformsMutex()); + mutex_lock lock(platforms_mutex_); if (GetPlatformMap()->find(key) != GetPlatformMap()->end()) { return port::Status(port::error::INTERNAL, "platform is already registered with name: \"" + @@ -45,33 +71,63 @@ namespace gputools { /* static */ port::StatusOr MultiPlatformManager::PlatformWithName( const string& target) { - tf_shared_lock lock(GetPlatformsMutex()); - auto it = GetPlatformMap()->find(port::Lowercase(target)); + mutex_lock lock(platforms_mutex_); - if (it == GetPlatformMap()->end()) { - return port::Status( - port::error::NOT_FOUND, - "could not find registered platform with name: \"" + target + "\""); + SE_ASSIGN_OR_RETURN(Platform * platform, LookupByNameLocked(target)); + if (!platform->Initialized()) { + SE_RETURN_IF_ERROR(platform->Initialize({})); } - return it->second; + return platform; } /* static */ port::StatusOr MultiPlatformManager::PlatformWithId( const Platform::Id& id) { - tf_shared_lock lock(GetPlatformsMutex()); - auto it = GetPlatformByIdMap()->find(id); - if (it == GetPlatformByIdMap()->end()) { + mutex_lock lock(platforms_mutex_); + + SE_ASSIGN_OR_RETURN(Platform * platform, LookupByIdLocked(id)); + if (!platform->Initialized()) { + SE_RETURN_IF_ERROR(platform->Initialize({})); + } + + return platform; +} + +/* static */ port::StatusOr +MultiPlatformManager::InitializePlatformWithName( + const string& target, const std::map& options) { + mutex_lock lock(platforms_mutex_); + + SE_ASSIGN_OR_RETURN(Platform * platform, LookupByNameLocked(target)); + if (platform->Initialized()) { + return port::Status(port::error::FAILED_PRECONDITION, + "platform \"" + target + "\" is already initialized"); + } + + SE_RETURN_IF_ERROR(platform->Initialize(options)); + + return platform; +} + +/* static */ port::StatusOr +MultiPlatformManager::InitializePlatformWithId( + const Platform::Id& id, const std::map& options) { + mutex_lock lock(platforms_mutex_); + + SE_ASSIGN_OR_RETURN(Platform * platform, LookupByIdLocked(id)); + if (platform->Initialized()) { return port::Status( - port::error::NOT_FOUND, - port::Printf("could not find registered platform with id: 0x%p", id)); + port::error::FAILED_PRECONDITION, + port::Printf("platform with id 0x%p is already initialized", id)); } - return it->second; + SE_RETURN_IF_ERROR(platform->Initialize(options)); + + return platform; } /* static */ void MultiPlatformManager::ClearPlatformRegistry() { - mutex_lock lock(GetPlatformsMutex()); + mutex_lock lock(platforms_mutex_); GetPlatformMap()->clear(); GetPlatformByIdMap()->clear(); } diff --git a/tensorflow/stream_executor/multi_platform_manager.h b/tensorflow/stream_executor/multi_platform_manager.h index ea6155b4826439b98262530e70e6463eb1fda237..438653ee20bdb1fd83cd9e75c4bcd35af277cc28 100644 --- a/tensorflow/stream_executor/multi_platform_manager.h +++ b/tensorflow/stream_executor/multi_platform_manager.h @@ -67,13 +67,13 @@ limitations under the License. #include #include #include -#include "tensorflow/stream_executor/platform/port.h" #include "tensorflow/stream_executor/lib/status.h" #include "tensorflow/stream_executor/lib/statusor.h" #include "tensorflow/stream_executor/platform.h" #include "tensorflow/stream_executor/platform/mutex.h" #include "tensorflow/stream_executor/platform/port.h" +#include "tensorflow/stream_executor/platform/thread_annotations.h" namespace perftools { namespace gputools { @@ -85,26 +85,43 @@ class MultiPlatformManager { // already registered. The associated listener, if not null, will be used to // trace events for ALL executors for that platform. // Takes ownership of listener. - static port::Status RegisterPlatform(std::unique_ptr platform); + static port::Status RegisterPlatform(std::unique_ptr platform) + LOCKS_EXCLUDED(platforms_mutex_); - // Retrieves the platform registered with the given platform name; e.g. - // "CUDA", "OpenCL", ... + // Retrieves the platform registered with the given platform name (e.g. + // "CUDA", "OpenCL", ...) or id (an opaque, comparable value provided by the + // Platform's Id() method). + // + // If the platform has not already been initialized, it will be initialized + // with a default set of parameters. // // If the requested platform is not registered, an error status is returned. // Ownership of the platform is NOT transferred to the caller -- // the MultiPlatformManager owns the platforms in a singleton-like fashion. - static port::StatusOr PlatformWithName(const string& target); - - // Retrieves the platform registered with the given platform ID, which - // is an opaque (but comparable) value. + static port::StatusOr PlatformWithName(const string& target) + LOCKS_EXCLUDED(platforms_mutex_); + static port::StatusOr PlatformWithId(const Platform::Id& id) + LOCKS_EXCLUDED(platforms_mutex_); + + // Retrieves the platform registered with the given platform name (e.g. + // "CUDA", "OpenCL", ...) or id (an opaque, comparable value provided by the + // Platform's Id() method). + // + // The platform will be initialized with the given options. If the platform + // was already initialized, an error will be returned. // // If the requested platform is not registered, an error status is returned. // Ownership of the platform is NOT transferred to the caller -- // the MultiPlatformManager owns the platforms in a singleton-like fashion. - static port::StatusOr PlatformWithId(const Platform::Id& id); + static port::StatusOr InitializePlatformWithName( + const string& target, const std::map& options) + LOCKS_EXCLUDED(platforms_mutex_); + static port::StatusOr InitializePlatformWithId( + const Platform::Id& id, const std::map& options) + LOCKS_EXCLUDED(platforms_mutex_); // Clears the set of registered platforms, primarily used for testing. - static void ClearPlatformRegistry(); + static void ClearPlatformRegistry() LOCKS_EXCLUDED(platforms_mutex_); // Although the MultiPlatformManager "owns" its platforms, it holds them as // undecorated pointers to prevent races during program exit (between this @@ -122,17 +139,16 @@ class MultiPlatformManager { // Provides access to the available set of platforms under a lock. static port::Status WithPlatforms( - std::function callback) { - mutex_lock lock(GetPlatformsMutex()); + std::function callback) + LOCKS_EXCLUDED(platforms_mutex_) { + mutex_lock lock(platforms_mutex_); return callback(GetPlatformMap()); } private: - // mutex that guards the platform map. - static mutex& GetPlatformsMutex() { - static mutex* platforms_mutex = new mutex; - return *platforms_mutex; - } + using PlatformIdMap = std::map; + + static mutex platforms_mutex_; // TODO(b/22689637): Clean up these two maps; make sure they coexist nicely. // TODO(b/22689637): Move this (whatever the final/"official" map is) to @@ -147,12 +163,21 @@ class MultiPlatformManager { // Holds a Platform::Id-to-object mapping. // Unlike platforms_ above, this map does not own its contents. - static std::map* GetPlatformByIdMap() { - using PlatformIdMap = std::map; + static PlatformIdMap* GetPlatformByIdMap() { static PlatformIdMap* instance = new PlatformIdMap; return instance; } + // Looks up the platform object with the given name. Assumes the Platforms + // mutex is held. + static port::StatusOr LookupByNameLocked(const string& target) + EXCLUSIVE_LOCKS_REQUIRED(platforms_mutex_); + + // Looks up the platform object with the given id. Assumes the Platforms + // mutex is held. + static port::StatusOr LookupByIdLocked(const Platform::Id& id) + EXCLUSIVE_LOCKS_REQUIRED(platforms_mutex_); + SE_DISALLOW_COPY_AND_ASSIGN(MultiPlatformManager); }; diff --git a/tensorflow/stream_executor/platform.cc b/tensorflow/stream_executor/platform.cc index 93f08d06dae862f24b5b533395af63139f344f77..4cdc22bd16a3ea66037696f6a9d70bcb86ef5ebb 100644 --- a/tensorflow/stream_executor/platform.cc +++ b/tensorflow/stream_executor/platform.cc @@ -85,6 +85,17 @@ StreamExecutorConfig::StreamExecutorConfig(int ordinal_in) Platform::~Platform() {} +bool Platform::Initialized() const { return true; } + +port::Status Platform::Initialize( + const std::map &platform_options) { + if (!platform_options.empty()) { + return port::Status(port::error::UNIMPLEMENTED, + "this platform does not support custom initialization"); + } + return port::Status::OK(); +} + port::Status Platform::ForceExecutorShutdown() { return port::Status(port::error::UNIMPLEMENTED, "executor shutdown is not supported on this platform"); diff --git a/tensorflow/stream_executor/platform.h b/tensorflow/stream_executor/platform.h index f0a0e60e02f951018b39ef831cd2f7dd3256f87d..54f8aa86c269ff0d32648e1d4629179cafd5be76 100644 --- a/tensorflow/stream_executor/platform.h +++ b/tensorflow/stream_executor/platform.h @@ -111,6 +111,9 @@ class Platform { // Returns a key uniquely identifying this platform. virtual Id id() const = 0; + // Name of this platform. + virtual const string& Name() const = 0; + // Returns the number of devices accessible on this platform. // // Note that, though these devices are visible, if there is only one userspace @@ -118,8 +121,17 @@ class Platform { // device, a call to ExecutorForDevice may return an error status. virtual int VisibleDeviceCount() const = 0; - // Name of this platform. - virtual const string& Name() const = 0; + // Returns true iff the platform has been initialized. + virtual bool Initialized() const; + + // Initializes the platform with a custom set of options. The platform must be + // initialized before obtaining StreamExecutor objects. The interpretation of + // the platform_options argument is implementation specific. This method may + // return an error if unrecognized options are provided. If using + // MultiPlatformManager, this method will be called automatically by + // InitializePlatformWithId/InitializePlatformWithName. + virtual port::Status Initialize( + const std::map& platform_options); // Returns a device with the given ordinal on this platform with a default // plugin configuration or, if none can be found with the given ordinal or @@ -156,6 +168,8 @@ class Platform { // This is only useful on platforms which bind a device to a single process // that has obtained the device context. May return UNIMPLEMENTED on platforms // that have no reason to destroy device contexts. + // + // The platform must be reinitialized after this is called. virtual port::Status ForceExecutorShutdown(); // Registers a TraceListener to listen to all StreamExecutors for this diff --git a/tensorflow/stream_executor/stream.cc b/tensorflow/stream_executor/stream.cc index ba5001e273632c893b05eea64542f1b156e28c47..6bbb5f0b2e728a855b64beb33d6e28b1b3b399c3 100644 --- a/tensorflow/stream_executor/stream.cc +++ b/tensorflow/stream_executor/stream.cc @@ -681,6 +681,37 @@ Stream &Stream::ThenFusedConvolveWithAlgorithm( return *this; } +Stream &Stream::ThenConvolveWithAlgorithm( + const dnn::BatchDescriptor &input_descriptor, + const DeviceMemory &input_data, + const dnn::FilterDescriptor &filter_descriptor, + const DeviceMemory &filter_data, + const dnn::ConvolutionDescriptor &convolution_descriptor, + const dnn::BatchDescriptor &output_descriptor, DeviceMemory *output, + ScratchAllocator *scratch_allocator, + const dnn::AlgorithmConfig &algorithm_config, + dnn::ProfileResult *output_profile_result) { + VLOG_CALL(PARAM(input_descriptor), PARAM(input_data), + PARAM(filter_descriptor), PARAM(filter_data), + PARAM(convolution_descriptor), PARAM(output_descriptor), + PARAM(output), PARAM(algorithm_config)); + + if (ok()) { + if (dnn::DnnSupport *dnn = parent_->AsDnn()) { + auto status = dnn->DoConvolve( + this, input_descriptor, input_data, filter_descriptor, filter_data, + convolution_descriptor, output_descriptor, output, scratch_allocator, + algorithm_config, output_profile_result); + if (!status && !output_profile_result) { + SetError(); + } + } else { + SetErrorAndLogNoDnnSupport(); + } + } + return *this; +} + Stream &Stream::ThenConvolveWithAlgorithm( const dnn::BatchDescriptor &input_descriptor, const DeviceMemory &input_data, @@ -890,6 +921,39 @@ Stream &Stream::ThenConvolveBackwardDataWithScratch( return *this; } +Stream &Stream::ThenConvolveBackwardDataWithAlgorithm( + const dnn::FilterDescriptor &filter_descriptor, + const DeviceMemory &filter_data, + const dnn::BatchDescriptor &output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor &convolution_descriptor, + const dnn::BatchDescriptor &input_descriptor, + DeviceMemory *backward_input_data, + ScratchAllocator *scratch_allocator, + const dnn::AlgorithmConfig &algorithm_config, + dnn::ProfileResult *output_profile_result) { + VLOG_CALL(PARAM(filter_descriptor), PARAM(filter_data), + PARAM(output_descriptor), PARAM(backward_output_data), + PARAM(convolution_descriptor), PARAM(input_descriptor), + PARAM(backward_input_data)); + + if (ok()) { + if (dnn::DnnSupport *dnn = parent_->AsDnn()) { + auto status = dnn->DoConvolveBackwardData( + this, filter_descriptor, filter_data, output_descriptor, + backward_output_data, convolution_descriptor, input_descriptor, + backward_input_data, scratch_allocator, algorithm_config, + output_profile_result); + if (!status && !output_profile_result) { + SetError(); + } + } else { + SetErrorAndLogNoDnnSupport(); + } + } + return *this; +} + Stream &Stream::ThenConvolveBackwardDataWithAlgorithm( const dnn::FilterDescriptor &filter_descriptor, const DeviceMemory &filter_data, @@ -1026,6 +1090,39 @@ Stream &Stream::ThenConvolveBackwardFilterWithScratch( return *this; } +Stream &Stream::ThenConvolveBackwardFilterWithAlgorithm( + const dnn::BatchDescriptor &input_descriptor, + const DeviceMemory &input_data, + const dnn::BatchDescriptor &output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor &convolution_descriptor, + const dnn::FilterDescriptor &filter_descriptor, + DeviceMemory *backward_filter_data, + ScratchAllocator *scratch_allocator, + const dnn::AlgorithmConfig &algorithm_config, + dnn::ProfileResult *output_profile_result) { + VLOG_CALL(PARAM(input_descriptor), PARAM(input_data), + PARAM(output_descriptor), PARAM(backward_output_data), + PARAM(convolution_descriptor), PARAM(filter_descriptor), + PARAM(backward_filter_data)); + + if (ok()) { + if (dnn::DnnSupport *dnn = parent_->AsDnn()) { + auto status = dnn->DoConvolveBackwardFilter( + this, input_descriptor, input_data, output_descriptor, + backward_output_data, convolution_descriptor, filter_descriptor, + backward_filter_data, scratch_allocator, algorithm_config, + output_profile_result); + if (!status && !output_profile_result) { + SetError(); + } + } else { + SetErrorAndLogNoDnnSupport(); + } + } + return *this; +} + Stream &Stream::ThenConvolveBackwardFilterWithAlgorithm( const dnn::BatchDescriptor &input_descriptor, const DeviceMemory &input_data, @@ -4923,12 +5020,6 @@ Stream &Stream::ThenTransformTensor(const dnn::BatchDescriptor &input_desc, return *this; } -Stream &Stream::ThenDoHostCallbackForTest(std::function callback) { - VLOG_CALL(PARAM(callback)); - - return ThenDoHostCallback(callback); -} - Stream &Stream::ThenDoHostCallback(std::function callback) { VLOG_CALL(PARAM(callback)); diff --git a/tensorflow/stream_executor/stream.h b/tensorflow/stream_executor/stream.h index a2fb2ea2375d0f245ae3bf3ccb04803d01663def..d7d11315699b85cae4d479b79bc8fc2717b2d8fb 100644 --- a/tensorflow/stream_executor/stream.h +++ b/tensorflow/stream_executor/stream.h @@ -358,6 +358,17 @@ class Stream { const dnn::BatchDescriptor &output_descriptor, DeviceMemory *output, ScratchAllocator *scratch_allocator); + Stream &ThenConvolveWithAlgorithm( + const dnn::BatchDescriptor &input_descriptor, + const DeviceMemory &input_data, + const dnn::FilterDescriptor &filter_descriptor, + const DeviceMemory &filter_data, + const dnn::ConvolutionDescriptor &convolution_descriptor, + const dnn::BatchDescriptor &output_descriptor, + DeviceMemory *output, ScratchAllocator *scratch_allocator, + const dnn::AlgorithmConfig &algorithm_config, + dnn::ProfileResult *output_profile_result); + Stream &ThenConvolveWithAlgorithm( const dnn::BatchDescriptor &input_descriptor, const DeviceMemory &input_data, @@ -476,6 +487,18 @@ class Stream { DeviceMemory *backward_input_data, ScratchAllocator *scratch_allocator); + Stream &ThenConvolveBackwardDataWithAlgorithm( + const dnn::FilterDescriptor &filter_descriptor, + const DeviceMemory &filter_data, + const dnn::BatchDescriptor &output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor &convolution_descriptor, + const dnn::BatchDescriptor &input_descriptor, + DeviceMemory *backward_input_data, + ScratchAllocator *scratch_allocator, + const dnn::AlgorithmConfig &algorithm_config, + dnn::ProfileResult *output_profile_result); + Stream &ThenConvolveBackwardDataWithAlgorithm( const dnn::FilterDescriptor &filter_descriptor, const DeviceMemory &filter_data, @@ -529,6 +552,18 @@ class Stream { DeviceMemory *backward_filter_data, ScratchAllocator *scratch_allocator); + Stream &ThenConvolveBackwardFilterWithAlgorithm( + const dnn::BatchDescriptor &input_descriptor, + const DeviceMemory &input_data, + const dnn::BatchDescriptor &output_descriptor, + DeviceMemory backward_output_data, + const dnn::ConvolutionDescriptor &convolution_descriptor, + const dnn::FilterDescriptor &filter_descriptor, + DeviceMemory *backward_filter_data, + ScratchAllocator *scratch_allocator, + const dnn::AlgorithmConfig &algorithm_config, + dnn::ProfileResult *output_profile_result); + Stream &ThenConvolveBackwardFilterWithAlgorithm( const dnn::BatchDescriptor &input_descriptor, const DeviceMemory &input_data, @@ -1933,16 +1968,15 @@ class Stream { // Entrains onto the stream a callback to the host (from the device). // Host callbacks block/occupy the stream just as device functions // (execute one at a time, block later stream operations). + // // Behavior is undefined when synchronizing using OpenCL user events. // Behavior is undefined if host callbacks call device routines or insert // them into any stream. + // // On certain platforms, ThenDoHostCallback is expected to have significant // negative effects on performance. Stream &ThenDoHostCallback(std::function callback); - // Identical to ThenDoHostCallback; only exposed for testing purposes. - Stream &ThenDoHostCallbackForTest(std::function callback); - // Returns the StreamExecutor (parent object) associated with this stream. StreamExecutor *parent() const { CHECK(parent_ != nullptr); diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 23d11c88ed687d919bd605af2f83f1ef77209370..c8d175dcf2cdd73bf49821cee1619145539f7042 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -219,6 +219,13 @@ def tf_copts(android_optimization_level_override="-O2", is_external=False): "//conditions:default": ["-pthread"] })) + +def tfe_xla_copts(): + return select({ + "//tensorflow:with_xla_support": ["-DTENSORFLOW_EAGER_USE_XLA"], + "//conditions:default": [], + }) + def tf_opts_nortti_if_android(): return if_android([ "-fno-rtti", @@ -491,6 +498,9 @@ def tf_gen_op_wrappers_cc(name, # is invalid to specify both "hidden" and "op_whitelist". # cc_linkopts: Optional linkopts to be added to tf_cc_binary that contains the # specified ops. +# gen_locally: if True, the genrule to generate the Python library will be run +# without sandboxing. This would help when the genrule depends on symlinks +# which may not be supported in the sandbox. def tf_gen_op_wrapper_py(name, out=None, hidden=None, @@ -501,7 +511,8 @@ def tf_gen_op_wrapper_py(name, generated_target_name=None, op_whitelist=[], cc_linkopts=[], - api_def_srcs=[]): + api_def_srcs=[], + gen_locally=False): if (hidden or hidden_file) and op_whitelist: fail('Cannot pass specify both hidden and op_whitelist.') @@ -556,6 +567,7 @@ def tf_gen_op_wrapper_py(name, outs=[out], srcs=api_def_srcs + [hidden_file], tools=[tool_name] + tf_binary_additional_srcs(), + local = (1 if gen_locally else 0), cmd=("$(location " + tool_name + ") " + api_def_args_str + " @$(location " + hidden_file + ") " + ("1" if require_shape_functions else "0") + " > $@")) @@ -565,6 +577,7 @@ def tf_gen_op_wrapper_py(name, outs=[out], srcs=api_def_srcs, tools=[tool_name] + tf_binary_additional_srcs(), + local = (1 if gen_locally else 0), cmd=("$(location " + tool_name + ") " + api_def_args_str + " " + op_list_arg + " " + ("1" if require_shape_functions else "0") + " " + @@ -605,7 +618,7 @@ def tf_cc_test(name, srcs=srcs + tf_binary_additional_srcs(), copts=tf_copts() + extra_copts, linkopts=select({ - "//tensorflow:android": [ + clean_dep("//tensorflow:android"): [ "-pie", ], clean_dep("//tensorflow:windows"): [], @@ -671,6 +684,7 @@ def tf_cuda_cc_test(name, tags=[], data=[], size="medium", + extra_copts=[], linkstatic=0, args=[], linkopts=[]): @@ -681,6 +695,7 @@ def tf_cuda_cc_test(name, tags=tags + ["manual"], data=data, size=size, + extra_copts=extra_copts, linkstatic=linkstatic, linkopts=linkopts, args=args) @@ -701,6 +716,7 @@ def tf_cuda_cc_test(name, tags=tags + tf_cuda_tests_tags(), data=data, size=size, + extra_copts=extra_copts, linkopts=linkopts, args=args) @@ -889,6 +905,14 @@ def tf_cuda_library(deps=None, cuda_deps=None, copts=tf_copts(), **kwargs): if not cuda_deps: cuda_deps = [] + if 'linkstatic' not in kwargs or kwargs['linkstatic'] != 1: + enable_text_relocation_linkopt = select({ + clean_dep("//tensorflow:darwin"): [], + "//conditions:default": ['-Wl,-z,notext'],}) + if 'linkopts' in kwargs: + kwargs['linkopts'] += enable_text_relocation_linkopt + else: + kwargs['linkopts'] = enable_text_relocation_linkopt native.cc_library( deps=deps + if_cuda(cuda_deps + [ clean_dep("//tensorflow/core:cuda"), @@ -1142,22 +1166,6 @@ def transitive_hdrs(name, deps=[], **kwargs): # the libraries in deps. def cc_header_only_library(name, deps=[], includes=[], **kwargs): _transitive_hdrs(name=name + "_gather", deps=deps) - - # We could generalize the following, but rather than complicate things - # here, we'll do the minimal use case for now, and hope bazel comes up - # with a better solution before too long. We'd expect it to compute - # the right include path by itself, but it doesn't, possibly because - # _transitive_hdrs lost some information about the include path. - if "@nsync//:nsync_headers" in deps: - # Buiding tensorflow from @org_tensorflow finds this two up. - nsynch = "../../external/nsync/public" - # Building tensorflow from elsewhere finds it four up. - # Note that native.repository_name() is not yet available in TF's Kokoro. - if REPOSITORY_NAME != "@": - nsynch = "../../" + nsynch - includes = includes[:] - includes.append(nsynch) - native.cc_library(name=name, hdrs=[":" + name + "_gather"], includes=includes, @@ -1166,9 +1174,24 @@ def cc_header_only_library(name, deps=[], includes=[], **kwargs): def tf_custom_op_library_additional_deps(): return [ "@protobuf_archive//:protobuf_headers", - "@nsync//:nsync_headers", clean_dep("//third_party/eigen3"), clean_dep("//tensorflow/core:framework_headers_lib"), + ] + if_windows(["//tensorflow/python:pywrap_tensorflow_import_lib"]) + +# A list of targets that contains the implemenation of +# tf_custom_op_library_additional_deps. It's used to generate a DEF file for +# exporting symbols from _pywrap_tensorflow.dll on Windows. +def tf_custom_op_library_additional_deps_impl(): + return [ + # for @protobuf_archive//:protobuf_headers + "@protobuf_archive//:protobuf", + # for @nsync//:nsync_headers + "@nsync//:nsync_cpp", + # for //third_party/eigen3 + clean_dep("//third_party/eigen3"), + # for //tensorflow/core:framework_headers_lib + clean_dep("//tensorflow/core:framework"), + clean_dep("//tensorflow/core:reader_base"), ] # Traverse the dependency graph along the "deps" attribute of the @@ -1255,6 +1278,7 @@ def tf_custom_op_library(name, srcs=[], gpu_srcs=[], deps=[], linkopts=[]): deps=deps + if_cuda(cuda_deps), data=[name + "_check_deps"], copts=tf_copts(is_external=True), + features = ["windows_export_all_symbols"], linkopts=linkopts + select({ "//conditions:default": [ "-lm", @@ -1296,6 +1320,46 @@ def tf_extension_linkopts(): def tf_extension_copts(): return [] # No extension c opts +# In tf_py_wrap_cc generated libraries +# module init functions are not exported unless +# they contain one of the keywords in the version file +# this prevents custom python modules. +# This function attempts to append init_module_name to list of +# exported functions in version script +def _append_init_to_versionscript_impl(ctx): + mod_name = ctx.attr.module_name + if ctx.attr.is_version_script: + ctx.actions.expand_template( + template=ctx.file.template_file, + output=ctx.outputs.versionscript, + substitutions={ + "global:":"global:\n init_%s;\n PyInit_*;"%(mod_name), + }, + is_executable=False, + ) + else: + ctx.actions.expand_template( + template=ctx.file.template_file, + output=ctx.outputs.versionscript, + substitutions={ + "*tensorflow*":"*tensorflow*\ninit_%s\nPyInit_*\n"%(mod_name), + }, + is_executable=False, + ) + + +_append_init_to_versionscript= rule( + implementation=_append_init_to_versionscript_impl, + attrs={ + "module_name":attr.string(mandatory=True), + "template_file":attr.label(allow_files=True,single_file=True,mandatory=True), + "is_version_script":attr.bool(default=True, + doc='whether target is a ld version script or exported symbol list', + mandatory=False), + }, + outputs={"versionscript":"%{name}.lds"}, +) + def tf_py_wrap_cc(name, srcs, swig_includes=[], @@ -1317,26 +1381,39 @@ def tf_py_wrap_cc(name, toolchain_deps=["//tools/defaults:crosstool"], module_name=module_name, py_module_name=name) + vscriptname=name+"_versionscript" + _append_init_to_versionscript( + name=vscriptname, + module_name=module_name, + is_version_script=select({ + "@local_config_cuda//cuda:darwin":False, + "//conditions:default":True, + }), + template_file=select({ + "@local_config_cuda//cuda:darwin":clean_dep("//tensorflow:tf_exported_symbols.lds"), + "//conditions:default":clean_dep("//tensorflow:tf_version_script.lds") + }) + ) extra_linkopts = select({ "@local_config_cuda//cuda:darwin": [ "-Wl,-exported_symbols_list", - clean_dep("//tensorflow:tf_exported_symbols.lds") + "%s.lds"%vscriptname, ], clean_dep("//tensorflow:windows"): [], clean_dep("//tensorflow:windows_msvc"): [], "//conditions:default": [ "-Wl,--version-script", - clean_dep("//tensorflow:tf_version_script.lds") + "%s.lds"%vscriptname, ] }) extra_deps += select({ "@local_config_cuda//cuda:darwin": [ - clean_dep("//tensorflow:tf_exported_symbols.lds") + "%s.lds"%vscriptname, ], clean_dep("//tensorflow:windows"): [], clean_dep("//tensorflow:windows_msvc"): [], "//conditions:default": [ - clean_dep("//tensorflow:tf_version_script.lds") + "%s.lds"%vscriptname, ] }) @@ -1348,7 +1425,8 @@ def tf_py_wrap_cc(name, ]) + tf_extension_copts()), linkopts=tf_extension_linkopts() + extra_linkopts, linkstatic=1, - deps=deps + extra_deps) + deps=deps + extra_deps, + **kwargs) native.genrule( name="gen_" + cc_library_pyd_name, srcs=[":" + cc_library_name], diff --git a/tensorflow/tools/api/generator/BUILD b/tensorflow/tools/api/generator/BUILD index d11031639592aa1d3e6ce1c7f09c2f0679b29854..14ce8dbeb3614b177bf4424c3e00ae7dd17f72f0 100644 --- a/tensorflow/tools/api/generator/BUILD +++ b/tensorflow/tools/api/generator/BUILD @@ -1,5 +1,6 @@ # Description: # Scripts used to generate TensorFlow Python API. + licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) @@ -21,7 +22,7 @@ py_binary( srcs = ["create_python_api.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow:tensorflow_py", + "//tensorflow/python", ], ) @@ -41,42 +42,91 @@ genrule( # every module exported using tf_export. For e.g. if an op is decorated with # @tf_export('module1.module2', 'module3'). Then, outs should include # api/module1/module2/__init__.py and api/module3/__init__.py. + # keep sorted outs = [ "api/__init__.py", + "api/app/__init__.py", "api/bitwise/__init__.py", + "api/compat/__init__.py", "api/contrib/__init__.py", "api/contrib/stat_summarizer/__init__.py", + "api/data/__init__.py", "api/distributions/__init__.py", "api/distributions/bijectors/__init__.py", "api/errors/__init__.py", - "api/image/__init__.py", - "api/linalg/__init__.py", - "api/nn/__init__.py", - "api/spectral/__init__.py", - "api/train/__init__.py", - "api/app/__init__.py", + "api/estimator/__init__.py", + "api/estimator/export/__init__.py", + "api/estimator/inputs/__init__.py", + "api/feature_column/__init__.py", "api/gfile/__init__.py", "api/graph_util/__init__.py", + "api/image/__init__.py", + "api/initializers/__init__.py", "api/keras/__init__.py", + "api/keras/activations/__init__.py", + "api/keras/applications/__init__.py", + "api/keras/applications/densenet/__init__.py", + "api/keras/applications/inception_resnet_v2/__init__.py", + "api/keras/applications/inception_v3/__init__.py", + "api/keras/applications/mobilenet/__init__.py", + "api/keras/applications/nasnet/__init__.py", + "api/keras/applications/resnet50/__init__.py", + "api/keras/applications/vgg16/__init__.py", + "api/keras/applications/vgg19/__init__.py", + "api/keras/applications/xception/__init__.py", "api/keras/backend/__init__.py", + "api/keras/callbacks/__init__.py", + "api/keras/constraints/__init__.py", "api/keras/datasets/__init__.py", "api/keras/datasets/boston_housing/__init__.py", "api/keras/datasets/cifar10/__init__.py", "api/keras/datasets/cifar100/__init__.py", + "api/keras/datasets/fashion_mnist/__init__.py", "api/keras/datasets/imdb/__init__.py", "api/keras/datasets/mnist/__init__.py", "api/keras/datasets/reuters/__init__.py", + "api/keras/estimator/__init__.py", + "api/keras/initializers/__init__.py", + "api/keras/layers/__init__.py", + "api/keras/losses/__init__.py", + "api/keras/metrics/__init__.py", + "api/keras/models/__init__.py", + "api/keras/optimizers/__init__.py", + "api/keras/preprocessing/__init__.py", + "api/keras/preprocessing/image/__init__.py", + "api/keras/preprocessing/sequence/__init__.py", + "api/keras/preprocessing/text/__init__.py", + "api/keras/regularizers/__init__.py", "api/keras/utils/__init__.py", + "api/keras/wrappers/__init__.py", + "api/keras/wrappers/scikit_learn/__init__.py", + "api/layers/__init__.py", + "api/linalg/__init__.py", "api/logging/__init__.py", - "api/resource_loader/__init__.py", - "api/sysconfig/__init__.py", - "api/test/__init__.py", - "api/initializers/__init__.py", - "api/keras/initializers/__init__.py", + "api/losses/__init__.py", + "api/manip/__init__.py", "api/metrics/__init__.py", + "api/nn/__init__.py", "api/nn/rnn_cell/__init__.py", + "api/profiler/__init__.py", + "api/python_io/__init__.py", + "api/resource_loader/__init__.py", + "api/saved_model/__init__.py", + "api/saved_model/builder/__init__.py", + "api/saved_model/constants/__init__.py", + "api/saved_model/loader/__init__.py", + "api/saved_model/main_op/__init__.py", + "api/saved_model/signature_constants/__init__.py", + "api/saved_model/signature_def_utils/__init__.py", + "api/saved_model/tag_constants/__init__.py", + "api/saved_model/utils/__init__.py", "api/sets/__init__.py", + "api/spectral/__init__.py", "api/summary/__init__.py", + "api/sysconfig/__init__.py", + "api/test/__init__.py", + "api/train/__init__.py", + "api/train/queue_runner/__init__.py", ], cmd = "$(location create_python_api) $(OUTS)", tools = ["create_python_api"], @@ -86,7 +136,9 @@ py_library( name = "python_api", srcs = [":python_api_gen"], srcs_version = "PY2AND3", + visibility = ["//tensorflow:__subpackages__"], deps = [ "//tensorflow/contrib:contrib_py", # keep + "//tensorflow/python", # keep ], ) diff --git a/tensorflow/tools/api/generator/create_python_api.py b/tensorflow/tools/api/generator/create_python_api.py index 1557314939bd85c0467426216f90aa3891ca0ac0..bb7c3e77a3a82b63188fc5a4a21706da7a276fde 100644 --- a/tensorflow/tools/api/generator/create_python_api.py +++ b/tensorflow/tools/api/generator/create_python_api.py @@ -23,15 +23,14 @@ import collections import os import sys -# This import is needed so that we can traverse over TensorFlow modules. -import tensorflow as tf # pylint: disable=unused-import +from tensorflow import python as tf from tensorflow.python.util import tf_decorator _API_CONSTANTS_ATTR = '_tf_api_constants' _API_NAMES_ATTR = '_tf_api_names' _API_DIR = '/api/' -_CONTRIB_IMPORT = 'from tensorflow import contrib' +_OUTPUT_MODULE = 'tensorflow.tools.api.generator.api' _GENERATED_FILE_HEADER = """\"\"\"Imports for Python API. This file is MACHINE GENERATED! Do not edit. @@ -92,7 +91,7 @@ def get_api_imports(): if module_contents_name == _API_CONSTANTS_ATTR: for exports, value in attr: for export in exports: - names = ['tf'] + export.split('.') + names = export.split('.') dest_module = '.'.join(names[:-1]) import_str = format_import(module.__name__, value, names[-1]) module_imports[dest_module].append(import_str) @@ -104,29 +103,43 @@ def get_api_imports(): if hasattr(attr, '__dict__') and _API_NAMES_ATTR in attr.__dict__: # The same op might be accessible from multiple modules. # We only want to consider location where function was defined. - if attr.__module__ != module.__name__: + # Here we check if the op is defined in another TensorFlow module in + # sys.modules. + if (hasattr(attr, '__module__') and + attr.__module__.startswith(tf.__name__) and + attr.__module__ != module.__name__ and + attr.__module__ in sys.modules and + module_contents_name in dir(sys.modules[attr.__module__])): continue for export in attr._tf_api_names: # pylint: disable=protected-access - names = ['tf'] + export.split('.') + names = export.split('.') dest_module = '.'.join(names[:-1]) import_str = format_import( module.__name__, module_contents_name, names[-1]) module_imports[dest_module].append(import_str) # Import all required modules in their parent modules. - # For e.g. if we import 'tf.foo.bar.Value'. Then, we also - # import 'bar' in 'tf.foo'. - dest_modules = set(module_imports.keys()) - for dest_module in dest_modules: - dest_module_split = dest_module.split('.') - for dest_submodule_index in range(1, len(dest_module_split)): - dest_submodule = '.'.join(dest_module_split[:dest_submodule_index]) + # For e.g. if we import 'foo.bar.Value'. Then, we also + # import 'bar' in 'foo'. + imported_modules = set(module_imports.keys()) + for module in imported_modules: + if not module: + continue + module_split = module.split('.') + parent_module = '' # we import submodules in their parent_module + + for submodule_index in range(len(module_split)): + import_from = _OUTPUT_MODULE + if submodule_index > 0: + parent_module += ('.' + module_split[submodule_index-1] if parent_module + else module_split[submodule_index-1]) + import_from += '.' + parent_module submodule_import = format_import( - '', dest_module_split[dest_submodule_index], - dest_module_split[dest_submodule_index]) - if submodule_import not in module_imports[dest_submodule]: - module_imports[dest_submodule].append(submodule_import) + import_from, module_split[submodule_index], + module_split[submodule_index]) + if submodule_import not in module_imports[parent_module]: + module_imports[parent_module].append(submodule_import) return module_imports @@ -151,8 +164,8 @@ def create_api_files(output_files): # First get module directory under _API_DIR. module_dir = os.path.dirname( output_file[output_file.rfind(_API_DIR)+len(_API_DIR):]) - # Convert / to . and prefix with tf. - module_name = '.'.join(['tf', module_dir.replace('/', '.')]).strip('.') + # Convert / to . + module_name = module_dir.replace('/', '.').strip('.') module_name_to_file_path[module_name] = output_file # Create file for each expected output in genrule. @@ -162,16 +175,14 @@ def create_api_files(output_files): open(file_path, 'a').close() module_imports = get_api_imports() - module_imports['tf'].append(_CONTRIB_IMPORT) # Include all of contrib. # Add imports to output files. missing_output_files = [] for module, exports in module_imports.items(): # Make sure genrule output file list is in sync with API exports. if module not in module_name_to_file_path: - module_without_tf = module[len('tf.'):] module_file_path = '"api/%s/__init__.py"' % ( - module_without_tf.replace('.', '/')) + module.replace('.', '/')) missing_output_files.append(module_file_path) continue with open(module_name_to_file_path[module], 'w') as fp: diff --git a/tensorflow/tools/api/golden/tensorflow.-graph.pbtxt b/tensorflow/tools/api/golden/tensorflow.-graph.pbtxt index 75361803a3991f380d6be2485cfd3d05fd1572e1..cdaeb55e30865e082054085f47d6a071ebf3affd 100644 --- a/tensorflow/tools/api/golden/tensorflow.-graph.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.-graph.pbtxt @@ -130,6 +130,10 @@ tf_class { name: "prevent_fetching" argspec: "args=[\'self\', \'op\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "switch_to_thread_local" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "unique_name" argspec: "args=[\'self\', \'name\', \'mark_as_used\'], varargs=None, keywords=None, defaults=[\'True\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.-variable.pbtxt b/tensorflow/tools/api/golden/tensorflow.-variable.pbtxt index bc7cf7267f7d23121402e63903f01ddc6caa2e04..5a02bb2175e2d6ad71722799143090f2735c1a37 100644 --- a/tensorflow/tools/api/golden/tensorflow.-variable.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.-variable.pbtxt @@ -1,6 +1,7 @@ path: "tensorflow.Variable" tf_class { is_instance: "" + is_instance: "" is_instance: "" member { name: "SaveSliceInfo" diff --git a/tensorflow/tools/api/golden/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/tensorflow.data.-dataset.pbtxt index 42de5c0c80023ad5bd7f33a564780060998307c1..0900adaf762df1415c8db63c3879ca2fabc28d9f 100644 --- a/tensorflow/tools/api/golden/tensorflow.data.-dataset.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.data.-dataset.pbtxt @@ -64,7 +64,7 @@ tf_class { } member_method { name: "list_files" - argspec: "args=[\'file_pattern\'], varargs=None, keywords=None, defaults=None" + argspec: "args=[\'file_pattern\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " } member_method { name: "make_initializable_iterator" diff --git a/tensorflow/tools/api/golden/tensorflow.data.-fixed-length-record-dataset.pbtxt b/tensorflow/tools/api/golden/tensorflow.data.-fixed-length-record-dataset.pbtxt index e2fc8d6cb1d318cc50828f22e8e575cc28c7aaad..7b16ac90c925beb25e065d26e73ee2a54b06d9dc 100644 --- a/tensorflow/tools/api/golden/tensorflow.data.-fixed-length-record-dataset.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.data.-fixed-length-record-dataset.pbtxt @@ -65,7 +65,7 @@ tf_class { } member_method { name: "list_files" - argspec: "args=[\'file_pattern\'], varargs=None, keywords=None, defaults=None" + argspec: "args=[\'file_pattern\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " } member_method { name: "make_initializable_iterator" diff --git a/tensorflow/tools/api/golden/tensorflow.data.-t-f-record-dataset.pbtxt b/tensorflow/tools/api/golden/tensorflow.data.-t-f-record-dataset.pbtxt index 9770389e5ef1e29a80ae1da2725d9862f6521ff9..9cf5f2ae2057ab4a16131527cf2ef2fa6ada28e5 100644 --- a/tensorflow/tools/api/golden/tensorflow.data.-t-f-record-dataset.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.data.-t-f-record-dataset.pbtxt @@ -17,7 +17,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'filenames\', \'compression_type\', \'buffer_size\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + argspec: "args=[\'self\', \'filenames\', \'compression_type\', \'buffer_size\', \'num_parallel_reads\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -65,7 +65,7 @@ tf_class { } member_method { name: "list_files" - argspec: "args=[\'file_pattern\'], varargs=None, keywords=None, defaults=None" + argspec: "args=[\'file_pattern\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " } member_method { name: "make_initializable_iterator" diff --git a/tensorflow/tools/api/golden/tensorflow.data.-text-line-dataset.pbtxt b/tensorflow/tools/api/golden/tensorflow.data.-text-line-dataset.pbtxt index 7263230c1c7182bb812cb2e433aedd415bcd16c7..8c3d6691439e619c906996a3ddaea4317c4a9597 100644 --- a/tensorflow/tools/api/golden/tensorflow.data.-text-line-dataset.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.data.-text-line-dataset.pbtxt @@ -65,7 +65,7 @@ tf_class { } member_method { name: "list_files" - argspec: "args=[\'file_pattern\'], varargs=None, keywords=None, defaults=None" + argspec: "args=[\'file_pattern\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " } member_method { name: "make_initializable_iterator" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-classifier.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-classifier.pbtxt index 874a73f661d782ff5637b751f104fd2209734599..be9ba4ce85bd5b9905a39e3f45873c534594e15f 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-classifier.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-classifier.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-regressor.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-regressor.pbtxt index 8da2a2b6867a3f9a3d82fcdb76ac4a62d5cee825..91fca67b6b5b1187b61f398a152793362c0c6e30 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-regressor.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-baseline-regressor.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-classifier.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-classifier.pbtxt index efc441ae2f2a00f663c11f84c1155bece0c8e08a..cd4f72fcf839fa89f25c7ed115ee6c61294283c3 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-classifier.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-classifier.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-classifier.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-classifier.pbtxt index 20ce87987060d9013bd071d6fc9f1f4f33467121..303fd74a64d0c7f5a0292a4eaabec63455c29381 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-classifier.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-classifier.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-regressor.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-regressor.pbtxt index 73211aaf8ba5f925982afe3d17c4b8f009250cb8..c97ea7969eff3e6952a604e72ce140b49d304461 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-regressor.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-linear-combined-regressor.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-regressor.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-regressor.pbtxt index 27a159639d2098aace2e69718d9ac4e38a29fdc3..4b5b5bf0e3599a81e2e853ae8ba34ef12cc63097 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-regressor.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-d-n-n-regressor.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-estimator.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-estimator.pbtxt index 76f527f796e95f342eb144ae3de87ff234338021..42a0d595216ad28363727b9d7c066fc37fddd02c 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-estimator.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-estimator.pbtxt @@ -44,7 +44,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-linear-classifier.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-linear-classifier.pbtxt index c45318b98a034255d32c326179813de14cf1d4c8..2de52d6c57cc70b562c3c10b7f23cd15b63e25f8 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-linear-classifier.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-linear-classifier.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-linear-regressor.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-linear-regressor.pbtxt index 04a2aa080d0704a8b7ec98f8eafda4bd1944e567..e552f33720bb939b8a98d34ef3de78bda7db976c 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-linear-regressor.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-linear-regressor.pbtxt @@ -45,7 +45,7 @@ tf_class { } member_method { name: "predict" - argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], " } member_method { name: "train" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-vocab-info.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-vocab-info.pbtxt index a16e3aedae96e7289e73c49ac7890550dd5ddb08..5301b94eb361251a1cb4d02a5d8168f7c8191045 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-vocab-info.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-vocab-info.pbtxt @@ -1,7 +1,7 @@ path: "tensorflow.estimator.VocabInfo" tf_class { - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" member { name: "backup_initializer" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-warm-start-settings.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-warm-start-settings.pbtxt index afdd6bb058353594415cd1abe726070f84ae46b6..43f5343359aff3b856a2b3708e4cda7cec29e146 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-warm-start-settings.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-warm-start-settings.pbtxt @@ -1,7 +1,7 @@ path: "tensorflow.estimator.WarmStartSettings" tf_class { - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" member { name: "ckpt_to_initialize_from" diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.export.-tensor-serving-input-receiver.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.export.-tensor-serving-input-receiver.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..4fe92643bf9867765499d7bf475b9cdd1686aec5 --- /dev/null +++ b/tensorflow/tools/api/golden/tensorflow.estimator.export.-tensor-serving-input-receiver.pbtxt @@ -0,0 +1,27 @@ +path: "tensorflow.estimator.export.TensorServingInputReceiver" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "features" + mtype: "" + } + member { + name: "receiver_tensors" + mtype: "" + } + member { + name: "receiver_tensors_alternatives" + mtype: "" + } + member_method { + name: "__init__" + } + member_method { + name: "count" + } + member_method { + name: "index" + } +} diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.export.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.export.pbtxt index 4d0dddb3bc0305a28fab0c95c31e4869f5db0aa8..bd72f6cd79f7dffb9f0a7f8ae43751c4ecba939d 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.export.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.export.pbtxt @@ -20,6 +20,10 @@ tf_module { name: "ServingInputReceiver" mtype: "" } + member { + name: "TensorServingInputReceiver" + mtype: "" + } member_method { name: "build_parsing_serving_input_receiver_fn" argspec: "args=[\'feature_spec\', \'default_batch_size\'], varargs=None, keywords=None, defaults=[\'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.image.pbtxt b/tensorflow/tools/api/golden/tensorflow.image.pbtxt index baedf596e8fbce921ed7e0570542b8a11655dba4..3fc64dae888012169af3ea7695154b73f24d90c8 100644 --- a/tensorflow/tools/api/golden/tensorflow.image.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.image.pbtxt @@ -100,6 +100,14 @@ tf_module { name: "hsv_to_rgb" argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " } + member_method { + name: "image_gradients" + argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "is_jpeg" + argspec: "args=[\'contents\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "non_max_suppression" argspec: "args=[\'boxes\', \'scores\', \'max_output_size\', \'iou_threshold\', \'name\'], varargs=None, keywords=None, defaults=[\'0.5\', \'None\'], " @@ -112,6 +120,10 @@ tf_module { name: "per_image_standardization" argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "psnr" + argspec: "args=[\'a\', \'b\', \'max_val\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "random_brightness" argspec: "args=[\'image\', \'max_delta\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], " @@ -184,6 +196,18 @@ tf_module { name: "sample_distorted_bounding_box" argspec: "args=[\'image_size\', \'bounding_boxes\', \'seed\', \'seed2\', \'min_object_covered\', \'aspect_ratio_range\', \'area_range\', \'max_attempts\', \'use_image_if_no_bounding_boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'0.1\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } + member_method { + name: "sobel_edges" + argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "ssim" + argspec: "args=[\'img1\', \'img2\', \'max_val\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "ssim_multiscale" + argspec: "args=[\'img1\', \'img2\', \'max_val\', \'power_factors\'], varargs=None, keywords=None, defaults=[\'(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)\'], " + } member_method { name: "total_variation" argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.initializers.pbtxt b/tensorflow/tools/api/golden/tensorflow.initializers.pbtxt index 21a0f84d22fc2d06e551c9a709f3963e812333b8..eaf0036cacfadce335a84bcf61f47f9d360be7e2 100644 --- a/tensorflow/tools/api/golden/tensorflow.initializers.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.initializers.pbtxt @@ -1,17 +1,9 @@ path: "tensorflow.initializers" tf_module { - member { - name: "absolute_import" - mtype: "" - } member { name: "constant" mtype: "" } - member { - name: "division" - mtype: "" - } member { name: "identity" mtype: "" @@ -24,10 +16,6 @@ tf_module { name: "orthogonal" mtype: "" } - member { - name: "print_function" - mtype: "" - } member { name: "random_normal" mtype: "" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt index 2bf584fa2936990b467b2da9c48620a31814691a..7be2f4f61f6b9637f372591e49efc0c93c7a8c0a 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt @@ -1,10 +1,10 @@ path: "tensorflow.keras.Model" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -38,6 +38,10 @@ tf_class { name: "input_spec" mtype: "" } + member { + name: "layers" + mtype: "" + } member { name: "losses" mtype: "" @@ -108,11 +112,11 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'inputs\', \'outputs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_loss" - argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_update" @@ -120,7 +124,7 @@ tf_class { } member_method { name: "add_variable" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " } member_method { name: "add_weight" @@ -136,11 +140,11 @@ tf_class { } member_method { name: "call" - argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\', \'training\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "compile" - argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "compute_mask" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt index 0a6096813155d59eb1c7920f2bcd250ed9730982..0f2428d77a537959cf2c46dfa350208abea8cb36 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt @@ -1,11 +1,11 @@ path: "tensorflow.keras.Sequential" tf_class { - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -39,6 +39,10 @@ tf_class { name: "input_spec" mtype: "" } + member { + name: "layers" + mtype: "" + } member { name: "losses" mtype: "" @@ -125,7 +129,7 @@ tf_class { } member_method { name: "add_loss" - argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_update" @@ -133,7 +137,7 @@ tf_class { } member_method { name: "add_variable" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " } member_method { name: "add_weight" @@ -149,7 +153,7 @@ tf_class { } member_method { name: "call" - argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } member_method { name: "compile" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-base-logger.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-base-logger.pbtxt index ea4d5143540611f0585b67910cb319454b8560dc..454823fd23e72c6aa6bf6aa608707fa3b893b986 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-base-logger.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-base-logger.pbtxt @@ -5,7 +5,7 @@ tf_class { is_instance: "" member_method { name: "__init__" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + argspec: "args=[\'self\', \'stateful_metrics\'], varargs=None, keywords=None, defaults=[\'None\'], " } member_method { name: "on_batch_begin" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-progbar-logger.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-progbar-logger.pbtxt index 0e6901f28affdfc73092c2b9f3af07d17db61a9f..543de0ad48b86502fc83374e5e6d82822485f331 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-progbar-logger.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.callbacks.-progbar-logger.pbtxt @@ -5,7 +5,7 @@ tf_class { is_instance: "" member_method { name: "__init__" - argspec: "args=[\'self\', \'count_mode\'], varargs=None, keywords=None, defaults=[\'samples\'], " + argspec: "args=[\'self\', \'count_mode\', \'stateful_metrics\'], varargs=None, keywords=None, defaults=[\'samples\', \'None\'], " } member_method { name: "on_batch_begin" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.datasets.fashion_mnist.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.datasets.fashion_mnist.pbtxt index 791cfda23345fea7df1cfb107ae5dec06354bd48..a0e14356fa5e91bc81bd89f6eb8c07087956c392 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.datasets.fashion_mnist.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.datasets.fashion_mnist.pbtxt @@ -1,3 +1,7 @@ path: "tensorflow.keras.datasets.fashion_mnist" tf_module { + member_method { + name: "load_data" + argspec: "args=[], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt index f4ab075959906cdf350ec5d49dc86f928b7eb7ae..db8f626b98b70fd99f38e696aa16c72e74e86e25 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Activation" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activity-regularization.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activity-regularization.pbtxt index eb558cddafc3972127786353072767f0d53bf174..809b3a5430449176a0d7423ec7f4499ceb620890 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activity-regularization.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activity-regularization.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.ActivityRegularization" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt index 770a107b664d7ab0a8aedf292a34d4258a201859..68d41bb6cc258ca87d4664ac0fb9d5649f89ebaf 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Add" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-alpha-dropout.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-alpha-dropout.pbtxt index 0ce42b706ec20a8ea1cc83ec95cb64d9be2e5710..970b777e514194db4ac49fe58bea737b35436217 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-alpha-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-alpha-dropout.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.AlphaDropout" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling1-d.pbtxt index d6c98fa225ce924bc8e20f8531516eaed4d32ffb..529c64ab293d596012aefd42e0695bd1eb7e44d1 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling2-d.pbtxt index 754fd310c6d8ddb994db0590342b29f8cb7abd71..7e7c330d74fe3b71ecd0eb87e34719e47ae70784 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling3-d.pbtxt index 9b62880c7931d151fb98cc1dc3149dcbd4dd103d..ada8466d7473072b1878861ab36ec40b07fa1914 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average-pooling3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt index b371ad148cee16dd243869d929e0c1c002794682..2a5c1cd530a7a532f6cdd3c184f4ee7eb88d23d3 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Average" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool1-d.pbtxt index 3e2aba55fd63326bb0e232fdce06f32884db7a0a..9a2cb29815d59f3761ea25e9ea36ff6489c85b88 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool2-d.pbtxt index fb37308cce0124538648c3837e1e802794d7f1ae..f5e991ea42e5ee2723b64574d4598dc8463f1c8c 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool3-d.pbtxt index 813470ffc7c87727eb0b958e54806f530399806a..31732214a62524017e39776cdfb9ab629746e8ae 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-avg-pool3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-batch-normalization.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-batch-normalization.pbtxt index e251ac18e511b58a49816126d9941b98e4f91088..422eddf10db6763e10405dba5537ca161d1b8994 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-batch-normalization.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-batch-normalization.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.BatchNormalization" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt index 2f5e65a0c58eb82c43b013f9c2fbffa4e813c4d2..9053a37916314198842bc21b0608a9b69a64c264 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Bidirectional" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -73,6 +74,10 @@ tf_class { name: "scope_name" mtype: "" } + member { + name: "trainable" + mtype: "" + } member { name: "trainable_variables" mtype: "" @@ -159,7 +164,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" @@ -175,7 +180,7 @@ tf_class { } member_method { name: "get_updates_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_weights" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt index ff08def0a08e5201bc01d61be3f2d66d712c384b..3d536d2182fc4480a2ee5fba177543ca21fbd5ac 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Concatenate" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt index 6db22ca0320519fd9c101456c9c9c0e26a9a11e0..6a7da1aef8db64ad11bb5a5ba357f33eeb99170b 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv1-d.pbtxt index 577f206e3510a9995d5d383ac440b4f68ea39fe5..801a0339720919f8b3f6beee0f045d58b2c0a371 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d-transpose.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d-transpose.pbtxt index 72924c32b43e5edb39938cc0cd909cffefa61be1..13352e264a5305190717bb973a3f2bce4d7f4fff 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d-transpose.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d.pbtxt index 16be08d9b2bae8fe1faecf34c4d87ac9b9baf142..f400e4a15c362037e85ac375cee98bb5f6358669 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d-transpose.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d-transpose.pbtxt index 11e05f884d781166616a9c9a61dacbc8fdae6ae3..b3a9f573b8ba652d2544b21f36f65fe81a6ebb50 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d-transpose.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d.pbtxt index 72b72d6b3b1e410dda0b0a529449f0135203fc1b..a9be09c0abd19aeb4df30116ef2befc3948bfbf4 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution1-d.pbtxt index ee93247f63ed700dc6058041bd0ea4ff5c879078..be1ef5eb928d16cc6bf78c289aa20d815c728b23 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d-transpose.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d-transpose.pbtxt index e5023287e5f38553f3553a37b5a908790072b5c7..30034f7eaf6d9073695353e5c8d9ead0cc8de7cc 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d-transpose.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d.pbtxt index ba38cb7121c9d312e7ba9d7147bdc67673d1ad2e..189b38054c004facfeeff8ad2ae87848b89040f2 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d-transpose.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d-transpose.pbtxt index 58724a1e1661609ef3c000c7ca1dfe9b3235acff..a76d85c629c1fe620dafd62a0f0e05e9009109e2 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d-transpose.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d.pbtxt index 98d52c430c659d0fc3e9299f7bede9190dad2fcf..782195d4ad5883d8c0ea6a657cc10258f2080a55 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping1-d.pbtxt index 33b6ebe1af731f66f88a9493502f69049ab34b42..2cb7a39ea595e1ff699b96554cb135377d20a488 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping1-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Cropping1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping2-d.pbtxt index 4b241ebb0f68c270a9448b02138d44f82211f418..80803306992bba3b601824a93cb3086ef3947369 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping2-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Cropping2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping3-d.pbtxt index 1856a9ee21347ed6ca3dd592517eb644e205a5b7..678f40bbc23db15ff7c1138169478fb4412a449d 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping3-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Cropping3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt index a8c37af31f649d28ca2ab7614178f2dee58c13fc..fac826109b6a32305ece86c4990f08afe2236ce8 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Dense" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt index 07d3f023e54105c606b198c05750ffa78ee5d0c8..285d544af2d69d564afdec748598b39b6b95670f 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Dot" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt index e2e21b5f123f63fa38cb0e344be9a12fc091f20b..b77976974cccb96fc2373c093d2bdf279560c46f 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Dropout" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-e-l-u.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-e-l-u.pbtxt index 92b9760d53e35d3e5066a730bb5cbda45492cc64..b07714d3f2d158496e0482f8611e55ea0fb0fd51 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-e-l-u.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-e-l-u.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.ELU" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt index 83c528b40117222ac2b3e85ad338459948d0aa8c..e67d4ddfc47077d62319ab097e5333a373cbfc80 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Embedding" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt index 73609752886c8c57a78f6bc02cc46d2c7ff6e996..b2a668e5a88d312656f48ddd0e9f7aa9f6306991 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Flatten" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u-cell.pbtxt index b329f1c46bb07ab7684dec6aaf45a20b98c27ed9..1fd3febad26df16576dedca1df7560bf230c08ec 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u-cell.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.GRUCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u.pbtxt index d0f6d2a14f936c47ef78f7bf819c13facea2efcb..f5f41d879dcb840551c00a7272bbcfbe51dbee89 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-g-r-u.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GRU" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activation" @@ -227,7 +228,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-dropout.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-dropout.pbtxt index 57596badf1881950270fa6d3c074afb65daaa8eb..f4f1a5d51c5d5689918af4facf907f79d9ca71ec 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-dropout.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.GaussianDropout" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-noise.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-noise.pbtxt index 3829353cc3c195a750ad862707c5c8563e203fba..e502df5e177d422403d0643c18a9588afb9d9713 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-noise.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-noise.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.GaussianNoise" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling1-d.pbtxt index e53e78a977b32eaf2e31867044aedd39ab2dd34f..9c8d5bfcd8966384230e7d5cdcc1cac53a0eab9a 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling1-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalAveragePooling1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling2-d.pbtxt index 48fcd1044e06b2fe61aadb6c3675ce82197ff003..8dd65f1f248daaf120780f19050c45d297b7902e 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling2-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalAveragePooling2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling3-d.pbtxt index 66c06ed47289eb2d83d97778a7b13dab821722d2..5e30571cc730ee23767a044036b590460deec00b 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-average-pooling3-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalAveragePooling3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool1-d.pbtxt index 4f2420f74ab3069952e4a44bf61e5e12b3e80ea3..ba90fa454696d1cb4e77d80a2dc77ff65def4714 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool1-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalAvgPool1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool2-d.pbtxt index 7912a6d933b851521358e0246d04688da410b909..8823857758307c208527b144c0cc73b566f2f115 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool2-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalAvgPool2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool3-d.pbtxt index d5b2d2c274ad97071497045271c0a595f8e0e062..500ced852ba6b19502769ba9052f2e364af7e283 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-avg-pool3-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalAvgPool3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool1-d.pbtxt index d88ff17eb6df7bbba7d3af4344fc8ddc367ae44c..cf2717ed46b56e639fb774c1e922648e1653ec0d 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool1-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalMaxPool1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool2-d.pbtxt index c8cc5a0ddfdd54cbb47de922591a9842abf63396..a86ff1a46997f19b11e6ef03be432b45687a2df2 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool2-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalMaxPool2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool3-d.pbtxt index 7956c5a340d963cfd5976e8af56da222848a164a..e01cc7c1b09ad6a40380613d54b771c6a1c89c1c 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pool3-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalMaxPool3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling1-d.pbtxt index 0a7e16413dfbd80d448eb1bad5771915475d96b2..259c1fb37c787f5318570b7aca6935d2f0ed997f 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling1-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalMaxPooling1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling2-d.pbtxt index 6c8a58a996f5313ea48e395e7e443a7c21f198ee..0c41bf97f763f1e40e8fac714709ccac1483a00b 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling2-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalMaxPooling2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling3-d.pbtxt index 7678ce8aab63fcfa76c0ac61346a723c1dfe1ee7..bec8817aa393ba2d8a6410408938402366cbb01d 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-global-max-pooling3-d.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.GlobalMaxPooling3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-layer.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-layer.pbtxt index d46fd41a3f33002a9bbe755851278c9729ccd1d1..17be86222901c0f5a9a18c0e5f1c5bcac6c06a17 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-layer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-layer.pbtxt @@ -1,9 +1,9 @@ path: "tensorflow.keras.layers.InputLayer" tf_class { - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m-cell.pbtxt index 3b171b137af699c9608494a17c5651b439fe4545..6d2a8c56196d9b3c80f570c7f1d3ac803253fff6 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m-cell.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.LSTMCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m.pbtxt index 0036d6805bb67fe20a1373927f15d8f564bdbe1d..490b5b618c65e28f1ae2e01e8d35e7f3973cc180 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-l-s-t-m.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.LSTM" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activation" @@ -231,7 +232,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt index ca0144929942f7024a4e8bac5552bf0547ceb56d..21a65b838af35e2f540eacab823513e7bf54b434 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Lambda" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt index c52ad727545c0bf4f199714d71180eac3f1bf62a..127b04738e70c11b2dc1071cf174cf5de23c5133 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt @@ -1,7 +1,8 @@ path: "tensorflow.keras.layers.Layer" tf_class { - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-leaky-re-l-u.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-leaky-re-l-u.pbtxt index 8134fb738683b79764662d9ea7f721fe04751162..87e49f2ed5b5d73aee5e9aa2511485b1f3f4bcd9 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-leaky-re-l-u.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-leaky-re-l-u.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.LeakyReLU" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected1-d.pbtxt index c5d452300947d7f74e7458e2a04bfdfabb1c1da2..1aa3aad3246b83931a47e69a4aa76fdf2b5aee22 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected1-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.LocallyConnected1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected2-d.pbtxt index bcbed9241b525a953c8b499197facaefebe8cc44..5e9dc7d4774c651a186a4e320d0cfd088e87b6b3 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-locally-connected2-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.LocallyConnected2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt index 244e79b4ffe60ddd6aa56d2780d80dfd66c494a9..0d101e5b68cdb2cdf24ed472c724cfc885e3d95d 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Masking" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool1-d.pbtxt index 56cbf5df785ef0e2614ea7e9e6cfe1335e148eec..c85cd49ac8ce2c1fc0759671865b7174cd1c1480 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool2-d.pbtxt index 33c2d30e86f9cdc3fb9f4f498bfc2c94497fe2dd..4f59e330c92f96101c65a9a24f66196e84587ccb 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool3-d.pbtxt index 94f91059b7a1e291c38fe0045accc6c03f226603..c0ea0eb0505d20e70d641f2a646a060d7dbfabda 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pool3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling1-d.pbtxt index 247230a6d68b8ea93a30a2f5846d8baaa78cb13e..ca37ae51314516ae67c7725eb2ccd3d25154e2ac 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling2-d.pbtxt index 8d61b67e7ce9564d31b0bd904a58540d19c89172..3ede2378347f5eddb0e8fae775a0200ea484d3f8 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling3-d.pbtxt index ad2e30802006e934730e5c75247e958329f7121c..d87e25a7ba8e7cce615431723b53a0106c2b5279 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-max-pooling3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt index ff0db15f190675d533c50c277eb1cb60e0b95e55..e4df7b48ae6b41400375920a48ef8577bb69376e 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Maximum" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt index 1d3f33f04516345ee32f16befe0d7200d2cdad00..6bf7c77743c31b6d74df35d827e9d5bc9a25d303 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.Multiply" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-p-re-l-u.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-p-re-l-u.pbtxt index c86bc49b22a8cc3e004a77f4a21594aacb2c665a..c14be132b7e406c99841576be8d8fa9ab99aa816 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-p-re-l-u.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-p-re-l-u.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.PReLU" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt index 2043e1a1263f0f0745b7c6446cc670fd6b0f0000..72ffbceae01da900778dba1ec14e646aa17b39e5 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Permute" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-r-n-n.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-r-n-n.pbtxt index b29f65d79dc0ffca176a8f2905df9474841f572e..d3e780c8b22ed580f61ffc3d9b2bad7278391402 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-r-n-n.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-r-n-n.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.RNN" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -162,7 +163,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-repeat-vector.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-repeat-vector.pbtxt index 4b0e98520a0dd86c085fa7345af445e1ae253d3b..a27980a9d17397e558a4b732e3dc332a0c1e8432 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-repeat-vector.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-repeat-vector.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.RepeatVector" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt index 34bc71af8a26ff6e4d7c81a3877751df5209906f..67f991276c6908ff54fd516e84533542a5f60528 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Reshape" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv1-d.pbtxt index dd67b76523cc50409516e29f963f59d039455bfd..fccea5e8af5ab81e712669ff1b2567d8bde8607e 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv1-d.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv2-d.pbtxt index 5d898fb2bd86b39cb8fab755382bb96cce231fa6..d20663bdb0bc2eea323d35b1e3d4d27122f50472 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-conv2-d.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution1-d.pbtxt index bf62c095e7cc3fbeac95919a0f9fdc545efd3d25..889fa0a1b58bbd3babd293b7b1b45915a9ee3ca4 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution1-d.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution2-d.pbtxt index c758d87993b3acba88a13c7bc9eaeee929a22652..c850f3fedc814b20f0f95cc3cf4fd5c973446b5b 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-separable-convolution2-d.pbtxt @@ -4,8 +4,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n-cell.pbtxt index 6e3cde3e3eaba4f9985411d66a220f7cdd4ee7ad..526d88ccba60eb25c68432e5baa03fd3a878f718 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n-cell.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.SimpleRNNCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n.pbtxt index b875898a8196e0359297f77cb10c1d8177f05d45..7fddae34472411f49d42b4d65d12034d056ec818 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-simple-r-n-n.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.SimpleRNN" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activation" @@ -219,7 +220,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt index ee4b2fa39ed34a544ee800e9370e4f34c4a17041..5b9b62fc970238e49e6d4849285606d0a7908b23 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Softmax" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout1-d.pbtxt index e4727072e375b9fc4dc99a1536eaaf3df5415369..769da30999993fad05ae0f7c04e256e6cf01a774 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout1-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout2-d.pbtxt index c5ff7043115ccdd3bc4a1147790b20feda410f65..fca2e42a1519fcf3a9f0ec996c50b148b2df05fd 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout2-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout3-d.pbtxt index 476a7f362cf88e234e964f6f6645ee4ed0cbaff8..36e8de09a967c5940bf8078234f5980a78ec8009 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-spatial-dropout3-d.pbtxt @@ -3,8 +3,9 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-stacked-r-n-n-cells.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-stacked-r-n-n-cells.pbtxt index db9f90caef7d6192ec7b68ddffd784a83ec0ac69..a96f16fae99af9c30959d228202055e9aebfaf58 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-stacked-r-n-n-cells.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-stacked-r-n-n-cells.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.StackedRNNCells" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -122,7 +123,7 @@ tf_class { } member_method { name: "call" - argspec: "args=[\'self\', \'inputs\', \'states\'], varargs=None, keywords=kwargs, defaults=None" + argspec: "args=[\'self\', \'inputs\', \'states\', \'constants\'], varargs=None, keywords=kwargs, defaults=[\'None\'], " } member_method { name: "compute_mask" @@ -158,7 +159,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-thresholded-re-l-u.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-thresholded-re-l-u.pbtxt index ef31c5443efa0c0e5a7a2e0a422d2a9c9c49baaf..e1cbd0e150ed890ae57c1725249d1340fc2cb663 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-thresholded-re-l-u.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-thresholded-re-l-u.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.ThresholdedReLU" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-time-distributed.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-time-distributed.pbtxt index 2a7059d9aa7ac12d8130c30622bc5f190562695c..f0d35728fb1c42d563ff0598dd84da51a766a764 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-time-distributed.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-time-distributed.pbtxt @@ -2,8 +2,9 @@ path: "tensorflow.keras.layers.TimeDistributed" tf_class { is_instance: "" is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -69,6 +70,10 @@ tf_class { name: "scope_name" mtype: "" } + member { + name: "trainable" + mtype: "" + } member { name: "trainable_variables" mtype: "" @@ -155,7 +160,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" @@ -171,7 +176,7 @@ tf_class { } member_method { name: "get_updates_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_weights" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling1-d.pbtxt index a81b83be49e0073f242efc6890e419b4fe172ab2..74efaea6ddb22ec2fe9d41558978c183b0e06671 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling1-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.UpSampling1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling2-d.pbtxt index 5403279d45ec7b93bae7907b891c659a043e96d0..dc5bd5fd5319f9bbd601a3c4083ae566b47e1aaa 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling2-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.UpSampling2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling3-d.pbtxt index 96c337caf28d43fabd0b90df016f4e8ab0c408db..e01ccfb74aead591f1018cdcbb1c888767ecdb20 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-up-sampling3-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.UpSampling3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt index 58bffa087521517fe7f0b5dcd6cae0a8b39a4e25..7e6f90f7623677244865ac285c134dc79f7b9b69 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.Wrapper" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -68,6 +69,10 @@ tf_class { name: "scope_name" mtype: "" } + member { + name: "trainable" + mtype: "" + } member { name: "trainable_variables" mtype: "" @@ -154,7 +159,7 @@ tf_class { } member_method { name: "get_losses_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_output_at" @@ -170,7 +175,7 @@ tf_class { } member_method { name: "get_updates_for" - argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } member_method { name: "get_weights" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding1-d.pbtxt index b81a4b1c50b22f13eacb521cfc8bc288bd40c81f..4d0d402dad442ccf52267f5ce40b05400afbfbc7 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding1-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.ZeroPadding1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding2-d.pbtxt index 1a26f2f3c9bbaa2aa567e76e1aafe14805ecff38..b353a529bcf8e543d334fee57fca26ebc83036a4 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding2-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.ZeroPadding2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding3-d.pbtxt index 310277fe67433fd870ae3d907984f402576925b2..9fe1256e616dbca4f35101df160dc55bc68bfa8a 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-zero-padding3-d.pbtxt @@ -1,8 +1,9 @@ path: "tensorflow.keras.layers.ZeroPadding3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.metrics.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.metrics.pbtxt index de285c1aab197ea5cae9c94048a5131f8463ebde..42729e4237685638d38301cece6e93383ddfffba 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.metrics.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.metrics.pbtxt @@ -22,7 +22,7 @@ tf_module { } member_method { name: "deserialize" - argspec: "args=[\'name\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], " } member_method { name: "get" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt index 0b816b58631d12471c2e9db96fc5395796d96ddf..8ccf15f9ab0fcfa59907ff05a962a84d3d86ccb4 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt @@ -1,10 +1,10 @@ path: "tensorflow.keras.models.Model" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -38,6 +38,10 @@ tf_class { name: "input_spec" mtype: "" } + member { + name: "layers" + mtype: "" + } member { name: "losses" mtype: "" @@ -108,11 +112,11 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'inputs\', \'outputs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_loss" - argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_update" @@ -120,7 +124,7 @@ tf_class { } member_method { name: "add_variable" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " } member_method { name: "add_weight" @@ -136,11 +140,11 @@ tf_class { } member_method { name: "call" - argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\', \'training\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "compile" - argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "compute_mask" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt index 7c1bfcb22558ec3a64c63ebbf0466f9114ef68ee..102eb3220334516e0051f952353920f229f4ff20 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt @@ -1,11 +1,11 @@ path: "tensorflow.keras.models.Sequential" tf_class { - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" @@ -39,6 +39,10 @@ tf_class { name: "input_spec" mtype: "" } + member { + name: "layers" + mtype: "" + } member { name: "losses" mtype: "" @@ -125,7 +129,7 @@ tf_class { } member_method { name: "add_loss" - argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_update" @@ -133,7 +137,7 @@ tf_class { } member_method { name: "add_variable" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " } member_method { name: "add_weight" @@ -149,7 +153,7 @@ tf_class { } member_method { name: "call" - argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } member_method { name: "compile" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.utils.-progbar.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.utils.-progbar.pbtxt index 3adc6b6faa6f62330f9ac3d621f29adfc380a09d..16e1cbe650e1662f8694fd7137ad20a48a90675b 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.utils.-progbar.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.utils.-progbar.pbtxt @@ -4,7 +4,7 @@ tf_class { is_instance: "" member_method { name: "__init__" - argspec: "args=[\'self\', \'target\', \'width\', \'verbose\', \'interval\'], varargs=None, keywords=None, defaults=[\'30\', \'1\', \'0.05\'], " + argspec: "args=[\'self\', \'target\', \'width\', \'verbose\', \'interval\', \'stateful_metrics\'], varargs=None, keywords=None, defaults=[\'30\', \'1\', \'0.05\', \'None\'], " } member_method { name: "add" @@ -12,6 +12,6 @@ tf_class { } member_method { name: "update" - argspec: "args=[\'self\', \'current\', \'values\', \'force\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], " + argspec: "args=[\'self\', \'current\', \'values\'], varargs=None, keywords=None, defaults=[\'None\'], " } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling1-d.pbtxt index de81206bc8b25046cd48c79ff8f154041c0e0cb0..1c4f550d7f05b8be33326cb39d7a5f3bf663f5e6 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling1-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling2-d.pbtxt index 72d5496464210efd9e423996dfb274dd9564f761..d2db0952693f2989e6a9e8748a254eb4db483206 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling2-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling3-d.pbtxt index 595e77ff9f8b64b6606fb075f3edf2281b4c3c1f..34d9a9df281c09a2e2030daf74a2ceb8066085bb 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling3-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt index 0c4aa2ff2612269727026141574726ad6df5cdbd..21ad0efecf88c42a3a679910ddfe095585a7933a 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.layers.BatchNormalization" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt index 5f576d0189309442dc4cea3d3617ab3144420165..ed38747c7671a267bb640ecb96a4c5fcc46c5edf 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d-transpose.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d-transpose.pbtxt index 675a7c76e569d3163ecd2c547841b4c36078b21d..ff453c6059477c20528fc768d93c65d208cdfc4a 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d-transpose.pbtxt @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt index eaabbf6aab172aea5c51f8071076890bb6b5bcf7..5583bd22dce18b0a0593b73bde509818b63b3f29 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d-transpose.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d-transpose.pbtxt index 838e070d79d2d7cfbd631f1a5e9960412cfdae5a..63f0c32a7c8f7e530c76c64fa619102bc12f9ad9 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d-transpose.pbtxt @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt index 4bd8cfc1a48cd839e2ffa54d0d0ca863060406d8..b77726252ccca30a7c6555fb569eb65b69e34998 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt index 57eccb03ffeb90652b019b5ce8a519797e4a3a3d..92db9f6dcd2f77c4253eb77df4a26fb632b2a766 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.layers.Dense" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt index a1ec00eeeaa98a6199e29b187b0760ddc92db09d..80fa846a24c9162d8521bdb4f098b9cd8e34aedb 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.layers.Dropout" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt index a06943d51a52f1951056136445b0d5786d801b5b..f63213b3dde40aa54b165c1c269c26fd2cd9e3b4 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.layers.Flatten" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-layer.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-layer.pbtxt index 24fda0c87ed0aeabd0fd4a16bb2efab444f8cd8a..4e45b2d513bb72bb47433d72c310d6a34fbc0c01 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-layer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-layer.pbtxt @@ -1,6 +1,7 @@ path: "tensorflow.layers.Layer" tf_class { is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling1-d.pbtxt index 4c3d00e0e1ddfe95c56f9ebc7c5d609c79dd44d4..19ec33fce775caa634e71e2295ac945a6f70ade9 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling1-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling2-d.pbtxt index f7e2017b0c9438130f1cfb2431eb73ca4d3103c5..76180c333a21c592a3b53bb445df9b12d3596552 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling2-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling3-d.pbtxt index 84780926a38ff811a5ab35fadfac690a6dbbbbe2..ded75c8ff09efc6746ddd2284f53d2c021cc473c 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling3-d.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv1-d.pbtxt index 05799ecfc9fdb9ff44620a67dcdbdc4426fddced..3dbfa5453f8e0ebb02429df9c4cbdf98de6b8ced 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv1-d.pbtxt @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv2-d.pbtxt index c2aeb35c4648bcce22ca73c838a85803a6b9cedf..ab171df1d1650e19836018f3316e6919f6d36def 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv2-d.pbtxt @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.pbtxt index 59134f84891ad5518dcb5331ce04475482c8b59e..df74c32e1f10cc7540ef105adef6be681e93d089 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.pbtxt @@ -76,10 +76,6 @@ tf_module { name: "SeparableConv2D" mtype: "" } - member_method { - name: "Input" - argspec: "args=[\'shape\', \'batch_size\', \'name\', \'dtype\', \'sparse\', \'tensor\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"\", \'False\', \'None\'], " - } member_method { name: "average_pooling1d" argspec: "args=[\'inputs\', \'pool_size\', \'strides\', \'padding\', \'data_format\', \'name\'], varargs=None, keywords=None, defaults=[\'valid\', \'channels_last\', \'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-l-s-t-m-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-l-s-t-m-cell.pbtxt index a2e728f94b41341b1a7c2a06d2c92d490f6eeb87..9c71a24d0500e2091e0ae94cc4dd7ed6b788a54f 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-l-s-t-m-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-l-s-t-m-cell.pbtxt @@ -1,9 +1,10 @@ path: "tensorflow.nn.rnn_cell.BasicLSTMCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt index 4211faa1ec615da8938d9a858a19a9e9a76378cd..9e19f96b7452616956fb7fd3ca62d8f4b25a2122 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt @@ -1,9 +1,10 @@ path: "tensorflow.nn.rnn_cell.BasicRNNCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-device-wrapper.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-device-wrapper.pbtxt index 0d253e5dd233d6d2b6ad0070a463c283a8769dab..7540aa62861895a7c41840476d4edb79785a77a9 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-device-wrapper.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-device-wrapper.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt index 97edf245f6fbed393a6fb8dbf1e83649e9ac4b4e..fc1ff386690f9c7acb11d4cc0770e394f78350ad 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-g-r-u-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-g-r-u-cell.pbtxt index 06fdc638c82b0d19b03857e33f083a94b7fd133b..751122cfff3bf9c55dd9fa264fdf2e1960940724 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-g-r-u-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-g-r-u-cell.pbtxt @@ -1,9 +1,10 @@ path: "tensorflow.nn.rnn_cell.GRUCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-l-s-t-m-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-l-s-t-m-cell.pbtxt index ef48cff0c329a7af5009d31fda429cf649c24261..4b6313f395fd8fd4ec2af78365117620263e7a55 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-l-s-t-m-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-l-s-t-m-cell.pbtxt @@ -1,9 +1,10 @@ path: "tensorflow.nn.rnn_cell.LSTMCell" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-multi-r-n-n-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-multi-r-n-n-cell.pbtxt index 9a6c73a079884b8ab92be1c9e89b2a9f34aad851..00e8c71140596ecea237ce05a09feff1fbb49001 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-multi-r-n-n-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-multi-r-n-n-cell.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-r-n-n-cell.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-r-n-n-cell.pbtxt index 27488f8e73f20456fae911511ecd2e41a60da351..3852f90dd6c4a254e20e789bdeb7796d61cef6bc 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-r-n-n-cell.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-r-n-n-cell.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.nn.rnn_cell.RNNCell" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-residual-wrapper.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-residual-wrapper.pbtxt index 3310836ed26387718115c2454300b9edfe930451..8f3f0f7506ef49014b31cd4bc04f1cb1e0d696fc 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-residual-wrapper.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-residual-wrapper.pbtxt @@ -3,6 +3,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "activity_regularizer" diff --git a/tensorflow/tools/api/golden/tensorflow.pbtxt b/tensorflow/tools/api/golden/tensorflow.pbtxt index e8890e9cc0a3c659b3f5f377136a2ca616d55993..99e09c375941b0463d610cd3ea46de6a887ba1b1 100644 --- a/tensorflow/tools/api/golden/tensorflow.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.pbtxt @@ -596,6 +596,10 @@ tf_module { name: "add_to_collection" argspec: "args=[\'name\', \'value\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "add_to_collections" + argspec: "args=[\'names\', \'value\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "all_variables" argspec: "args=[], varargs=None, keywords=None, defaults=None" @@ -892,6 +896,10 @@ tf_module { name: "cumsum" argspec: "args=[\'x\', \'axis\', \'exclusive\', \'reverse\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'False\', \'False\', \'None\'], " } + member_method { + name: "custom_gradient" + argspec: "args=[\'f\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "decode_base64" argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " @@ -964,6 +972,10 @@ tf_module { name: "einsum" argspec: "args=[\'equation\'], varargs=inputs, keywords=kwargs, defaults=None" } + member_method { + name: "enable_eager_execution" + argspec: "args=[\'config\', \'device_policy\', \'execution_mode\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + } member_method { name: "encode_base64" argspec: "args=[\'input\', \'pad\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], " @@ -980,6 +992,10 @@ tf_module { name: "erfc" argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " } + member_method { + name: "executing_eagerly" + argspec: "args=[], varargs=None, keywords=None, defaults=None" + } member_method { name: "exp" argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " @@ -1094,7 +1110,7 @@ tf_module { } member_method { name: "get_local_variable" - argspec: "args=[], varargs=args, keywords=kwargs, defaults=None" + argspec: "args=[\'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'collections\', \'caching_device\', \'partitioner\', \'validate_shape\', \'use_resource\', \'custom_getter\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'False\', \'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "get_seed" @@ -1600,6 +1616,10 @@ tf_module { name: "reduce_sum" argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], " } + member_method { + name: "regex_replace" + argspec: "args=[\'input\', \'pattern\', \'rewrite\', \'replace_global\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], " + } member_method { name: "register_tensor_conversion_function" argspec: "args=[\'base_type\', \'conversion_func\', \'priority\'], varargs=None, keywords=None, defaults=[\'100\'], " @@ -1988,10 +2008,22 @@ tf_module { name: "tile" argspec: "args=[\'input\', \'multiples\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " } + member_method { + name: "timestamp" + argspec: "args=[\'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "to_bfloat16" argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'ToBFloat16\'], " } + member_method { + name: "to_complex128" + argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'ToComplex128\'], " + } + member_method { + name: "to_complex64" + argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'ToComplex64\'], " + } member_method { name: "to_double" argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'ToDouble\'], " @@ -2056,6 +2088,22 @@ tf_module { name: "unsorted_segment_max" argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " } + member_method { + name: "unsorted_segment_mean" + argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "unsorted_segment_min" + argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "unsorted_segment_prod" + argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "unsorted_segment_sqrt_n" + argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "unsorted_segment_sum" argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.train.-adadelta-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-adadelta-optimizer.pbtxt index 863beaea4cf05a67e572c97b556bc1eb598d9ced..16bfbf20d5227d6308248bebcb62f32a2df8ef41 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-adadelta-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-adadelta-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.AdadeltaOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-adagrad-d-a-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-adagrad-d-a-optimizer.pbtxt index 0a7aa9b6bc14c95e74ab05a3aeb71b770a918f60..61cde9181c2367153b7b289b41bd932482bb92fd 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-adagrad-d-a-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-adagrad-d-a-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.AdagradDAOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-adagrad-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-adagrad-optimizer.pbtxt index 83724fea55d005e9476801feb1bf58cb004aa141..0a998c1afe4fff6e215360bc1cf8fc135754223c 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-adagrad-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-adagrad-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.AdagradOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-adam-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-adam-optimizer.pbtxt index e285b27a0531e00d27941fe451570a5056995c17..cc5954152577796ee7a5a6e1cedc873647d64f7c 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-adam-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-adam-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.AdamOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-ftrl-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-ftrl-optimizer.pbtxt index fc28577d6ed1328ae85970cf22cc458b7cf54344..1add3a902122341a706c38b19ea6ff5882c26445 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-ftrl-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-ftrl-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.FtrlOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-gradient-descent-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-gradient-descent-optimizer.pbtxt index bf3c1d81f877e3a8a7e24d5455e9c5bf6a41f764..ef5bbd6ace29abb5c73516176fcc7594a58d493a 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-gradient-descent-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-gradient-descent-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.GradientDescentOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-momentum-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-momentum-optimizer.pbtxt index a640c8d2c6366951cbba6a15d2000d9369cbbdbf..3d6e87f5eb44de9d6ce1bdd25a54b8df9020cc03 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-momentum-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-momentum-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.MomentumOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-optimizer.pbtxt index 6b33c236a35f09422a42a17b3ffddf5ba7b1595f..e73861ff7cb2d90d8efac72cdd7de3b27395f29e 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-optimizer.pbtxt @@ -1,6 +1,7 @@ path: "tensorflow.train.Optimizer" tf_class { is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-proximal-adagrad-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-proximal-adagrad-optimizer.pbtxt index d23fcaed7b4cee397dcf9c51eb3b521e5461c9e5..301b35b199c87890a0aef4139eb06253592ce0c4 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-proximal-adagrad-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-proximal-adagrad-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.ProximalAdagradOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-proximal-gradient-descent-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-proximal-gradient-descent-optimizer.pbtxt index b6c03e71d9ffb50bd6377b489fcc444453bd9752..8815befa936a85522011111a4a6270d22cbc25ae 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-proximal-gradient-descent-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-proximal-gradient-descent-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.ProximalGradientDescentOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-r-m-s-prop-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-r-m-s-prop-optimizer.pbtxt index 4a82db11cb8d85bd0c44135ecaf507c62fae41a1..e9819683ba5ec1bcacb3cdbcb2d787e866a77b6f 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-r-m-s-prop-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-r-m-s-prop-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.RMSPropOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-sync-replicas-optimizer.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-sync-replicas-optimizer.pbtxt index e9131bf544f2e7f08928f46d2be06a00259690be..3db96aff876b88b80b647570cf68b1ebc0b2da3b 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-sync-replicas-optimizer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-sync-replicas-optimizer.pbtxt @@ -2,6 +2,7 @@ path: "tensorflow.train.SyncReplicasOptimizer" tf_class { is_instance: "" is_instance: "" + is_instance: "" is_instance: "" member { name: "GATE_GRAPH" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-vocab-info.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-vocab-info.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..4ce7cb111163e103a1cebe30d5c6f3eeb4234693 --- /dev/null +++ b/tensorflow/tools/api/golden/tensorflow.train.-vocab-info.pbtxt @@ -0,0 +1,39 @@ +path: "tensorflow.train.VocabInfo" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "backup_initializer" + mtype: "" + } + member { + name: "new_vocab" + mtype: "" + } + member { + name: "new_vocab_size" + mtype: "" + } + member { + name: "num_oov_buckets" + mtype: "" + } + member { + name: "old_vocab" + mtype: "" + } + member { + name: "old_vocab_size" + mtype: "" + } + member_method { + name: "__init__" + } + member_method { + name: "count" + } + member_method { + name: "index" + } +} diff --git a/tensorflow/tools/api/golden/tensorflow.train.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.pbtxt index e49c719a334455d1f8f39fa67332be8bb81f2bc2..c75ee474aa471524e4b5c8a7e2dd4a9da4b08eae 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.pbtxt @@ -224,6 +224,10 @@ tf_module { name: "SyncReplicasOptimizer" mtype: "" } + member { + name: "VocabInfo" + mtype: "" + } member { name: "WorkerSessionCreator" mtype: "" @@ -402,7 +406,7 @@ tf_module { } member_method { name: "sdca_optimizer" - argspec: "args=[\'sparse_example_indices\', \'sparse_feature_indices\', \'sparse_feature_values\', \'dense_features\', \'example_weights\', \'example_labels\', \'sparse_indices\', \'sparse_weights\', \'dense_weights\', \'example_state_data\', \'loss_type\', \'l1\', \'l2\', \'num_loss_partitions\', \'num_inner_iterations\', \'adaptative\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], " + argspec: "args=[\'sparse_example_indices\', \'sparse_feature_indices\', \'sparse_feature_values\', \'dense_features\', \'example_weights\', \'example_labels\', \'sparse_indices\', \'sparse_weights\', \'dense_weights\', \'example_state_data\', \'loss_type\', \'l1\', \'l2\', \'num_loss_partitions\', \'num_inner_iterations\', \'adaptative\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], " } member_method { name: "sdca_shrink_l1" @@ -436,6 +440,10 @@ tf_module { name: "update_checkpoint_state" argspec: "args=[\'save_dir\', \'model_checkpoint_path\', \'all_model_checkpoint_paths\', \'latest_filename\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } + member_method { + name: "warm_start" + argspec: "args=[\'ckpt_to_initialize_from\', \'vars_to_warm_start\', \'var_name_to_vocab_info\', \'var_name_to_prev_var_name\'], varargs=None, keywords=None, defaults=[\'.*\', \'None\', \'None\'], " + } member_method { name: "write_graph" argspec: "args=[\'graph_or_graph_def\', \'logdir\', \'name\', \'as_text\'], varargs=None, keywords=None, defaults=[\'True\'], " diff --git a/tensorflow/tools/api/tests/BUILD b/tensorflow/tools/api/tests/BUILD index 8fb6b1cdfd8981e427062e186f6ac26b24231b8b..15bf1abb5f8f541c435be77b1a3c2f13382f2438 100644 --- a/tensorflow/tools/api/tests/BUILD +++ b/tensorflow/tools/api/tests/BUILD @@ -17,19 +17,15 @@ py_test( name = "api_compatibility_test", srcs = ["api_compatibility_test.py"], data = [ - ":convert_from_multiline", - "//tensorflow/core/api_def:base_api_def", - "//tensorflow/core/api_def:python_api_def", - "//tensorflow/python:hidden_ops", "//tensorflow/tools/api/golden:api_golden", "//tensorflow/tools/api/tests:API_UPDATE_WARNING.txt", "//tensorflow/tools/api/tests:README.txt", ], srcs_version = "PY2AND3", deps = [ + "//tensorflow:experimental_tensorflow_py", "//tensorflow:tensorflow_py", "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", "//tensorflow/python:lib", "//tensorflow/python:platform", "//tensorflow/tools/api/lib:python_object_to_proto_visitor", diff --git a/tensorflow/tools/api/tests/api_compatibility_test.py b/tensorflow/tools/api/tests/api_compatibility_test.py index afcbf50944cc47b3ae3086b17279f2ce2fdc6ee7..baa7a0889de2534604f94f34f6852a8fdf8819fa 100644 --- a/tensorflow/tools/api/tests/api_compatibility_test.py +++ b/tensorflow/tools/api/tests/api_compatibility_test.py @@ -28,18 +28,16 @@ from __future__ import division from __future__ import print_function import argparse -from collections import defaultdict import os import re -import subprocess import sys import unittest import tensorflow as tf +from tensorflow import experimental_api as api from google.protobuf import text_format -from tensorflow.core.framework import api_def_pb2 from tensorflow.python.lib.io import file_io from tensorflow.python.platform import resource_loader from tensorflow.python.platform import test @@ -49,6 +47,9 @@ from tensorflow.tools.api.lib import python_object_to_proto_visitor from tensorflow.tools.common import public_api from tensorflow.tools.common import traverse +if hasattr(tf, 'experimental_api'): + del tf.experimental_api + # FLAGS defined at the bottom: FLAGS = None # DEFINE_boolean, update_goldens, default False: @@ -67,11 +68,6 @@ _API_GOLDEN_FOLDER = 'tensorflow/tools/api/golden' _TEST_README_FILE = 'tensorflow/tools/api/tests/README.txt' _UPDATE_WARNING_FILE = 'tensorflow/tools/api/tests/API_UPDATE_WARNING.txt' -_CONVERT_FROM_MULTILINE_SCRIPT = 'tensorflow/tools/api/tests/convert_from_multiline' -_BASE_API_DIR = 'tensorflow/core/api_def/base_api' -_PYTHON_API_DIR = 'tensorflow/core/api_def/python_api' -_HIDDEN_OPS_FILE = 'tensorflow/python/ops/hidden_ops.txt' - def _KeyToFilePath(key): """From a given key, construct a filepath.""" @@ -96,55 +92,6 @@ def _FileNameToKey(filename): return api_object_key -def _GetSymbol(symbol_id): - """Get TensorFlow symbol based on the given identifier. - - Args: - symbol_id: Symbol identifier in the form module1.module2. ... .sym. - - Returns: - Symbol corresponding to the given id. - """ - # Ignore first module which should be tensorflow - symbol_id_split = symbol_id.split('.')[1:] - symbol = tf - for sym in symbol_id_split: - symbol = getattr(symbol, sym) - return symbol - - -def _IsGenModule(module_name): - if not module_name: - return False - module_name_split = module_name.split('.') - return module_name_split[-1].startswith('gen_') - - -def _GetHiddenOps(): - hidden_ops_file = file_io.FileIO(_HIDDEN_OPS_FILE, 'r') - hidden_ops = set() - for line in hidden_ops_file: - line = line.strip() - if not line: - continue - if line[0] == '#': # comment line - continue - # If line is of the form "op_name # comment", only keep the op_name. - line_split = line.split('#') - hidden_ops.add(line_split[0].strip()) - return hidden_ops - - -def _GetGoldenApiDefs(): - old_api_def_files = file_io.get_matching_files(_GetApiDefFilePath('*')) - return {file_path: file_io.read_file_to_string(file_path) - for file_path in old_api_def_files} - - -def _GetApiDefFilePath(graph_op_name): - return os.path.join(_PYTHON_API_DIR, 'api_def_%s.pbtxt' % graph_op_name) - - class ApiCompatibilityTest(test.TestCase): def __init__(self, *args, **kwargs): @@ -166,7 +113,8 @@ class ApiCompatibilityTest(test.TestCase): expected_dict, actual_dict, verbose=False, - update_goldens=False): + update_goldens=False, + additional_missing_object_message=''): """Diff given dicts of protobufs and report differences a readable way. Args: @@ -177,6 +125,8 @@ class ApiCompatibilityTest(test.TestCase): verbose: Whether to log the full diffs, or simply report which files were different. update_goldens: Whether to update goldens when there are diffs found. + additional_missing_object_message: Message to print when a symbol is + missing. """ diffs = [] verbose_diffs = [] @@ -195,7 +145,8 @@ class ApiCompatibilityTest(test.TestCase): verbose_diff_message = '' # First check if the key is not found in one or the other. if key in only_in_expected: - diff_message = 'Object %s expected but not found (removed).' % key + diff_message = 'Object %s expected but not found (removed). %s' % ( + key, additional_missing_object_message) verbose_diff_message = diff_message elif key in only_in_actual: diff_message = 'New object %s found (added).' % key @@ -222,7 +173,7 @@ class ApiCompatibilityTest(test.TestCase): logging.error('%d differences found between API and golden.', diff_count) messages = verbose_diffs if verbose else diffs for i in range(diff_count): - logging.error('Issue %d\t: %s', i + 1, messages[i]) + print('Issue %d\t: %s' % (i + 1, messages[i]), file=sys.stderr) if update_goldens: # Write files if requested. @@ -286,187 +237,59 @@ class ApiCompatibilityTest(test.TestCase): verbose=FLAGS.verbose_diffs, update_goldens=FLAGS.update_goldens) + @unittest.skipUnless( + sys.version_info.major == 2, + 'API compabitility test goldens are generated using python2.') + def testNewAPIBackwardsCompatibility(self): + # Extract all API stuff. + visitor = python_object_to_proto_visitor.PythonObjectToProtoVisitor() -class ApiDefTest(test.TestCase): - - def __init__(self, *args, **kwargs): - super(ApiDefTest, self).__init__(*args, **kwargs) - self._first_cap_pattern = re.compile('(.)([A-Z][a-z]+)') - self._all_cap_pattern = re.compile('([a-z0-9])([A-Z])') - - def _GenerateLowerCaseOpName(self, op_name): - lower_case_name = self._first_cap_pattern.sub(r'\1_\2', op_name) - return self._all_cap_pattern.sub(r'\1_\2', lower_case_name).lower() - - def _CreatePythonApiDef(self, base_api_def, endpoint_names): - """Creates Python ApiDef that overrides base_api_def if needed. - - Args: - base_api_def: (api_def_pb2.ApiDef) base ApiDef instance. - endpoint_names: List of Python endpoint names. - - Returns: - api_def_pb2.ApiDef instance with overrides for base_api_def - if module.name endpoint is different from any existing - endpoints in base_api_def. Otherwise, returns None. - """ - endpoint_names_set = set(endpoint_names) - - # If the only endpoint is equal to graph_op_name then - # it is equivalent to having no endpoints. - if (not base_api_def.endpoint and len(endpoint_names) == 1 - and endpoint_names[0] == - self._GenerateLowerCaseOpName(base_api_def.graph_op_name)): - return None - - base_endpoint_names_set = { - self._GenerateLowerCaseOpName(endpoint.name) - for endpoint in base_api_def.endpoint} - - if endpoint_names_set == base_endpoint_names_set: - return None # All endpoints are the same - - api_def = api_def_pb2.ApiDef() - api_def.graph_op_name = base_api_def.graph_op_name - - for endpoint_name in sorted(endpoint_names): - new_endpoint = api_def.endpoint.add() - new_endpoint.name = endpoint_name - - return api_def - - def _GetBaseApiMap(self): - """Get a map from graph op name to its base ApiDef. + public_api_visitor = public_api.PublicAPIVisitor(visitor) + public_api_visitor.do_not_descend_map['tf'].append('contrib') + public_api_visitor.do_not_descend_map['tf.GPUOptions'] = ['Experimental'] + # TODO(annarev): Make slide_dataset available in API. + public_api_visitor.private_map['tf'] = ['slide_dataset'] + traverse.traverse(api, public_api_visitor) - Returns: - Dictionary mapping graph op name to corresponding ApiDef. - """ - # Convert base ApiDef in Multiline format to Proto format. - converted_base_api_dir = os.path.join( - test.get_temp_dir(), 'temp_base_api_defs') - subprocess.check_call( - [os.path.join(resource_loader.get_root_dir_with_all_resources(), - _CONVERT_FROM_MULTILINE_SCRIPT), - _BASE_API_DIR, converted_base_api_dir]) - - name_to_base_api_def = {} - base_api_files = file_io.get_matching_files( - os.path.join(converted_base_api_dir, 'api_def_*.pbtxt')) - for base_api_file in base_api_files: - if file_io.file_exists(base_api_file): - api_defs = api_def_pb2.ApiDefs() - text_format.Merge( - file_io.read_file_to_string(base_api_file), api_defs) - for api_def in api_defs.op: - name_to_base_api_def[api_def.graph_op_name] = api_def - return name_to_base_api_def - - def _AddHiddenOpOverrides(self, name_to_base_api_def, api_def_map): - """Adds ApiDef overrides to api_def_map for hidden Python ops. + proto_dict = visitor.GetProtos() - Args: - name_to_base_api_def: Map from op name to base api_def_pb2.ApiDef. - api_def_map: Map from file path to api_def_pb2.ApiDefs for Python API - overrides. - """ - hidden_ops = _GetHiddenOps() - for hidden_op in hidden_ops: - if hidden_op not in name_to_base_api_def: - logging.warning('Unexpected hidden op name: %s' % hidden_op) - continue + # Read all golden files. + expression = os.path.join( + resource_loader.get_root_dir_with_all_resources(), + _KeyToFilePath('*')) + golden_file_list = file_io.get_matching_files(expression) - base_api_def = name_to_base_api_def[hidden_op] - if base_api_def.visibility != api_def_pb2.ApiDef.HIDDEN: - api_def = api_def_pb2.ApiDef() - api_def.graph_op_name = base_api_def.graph_op_name - api_def.visibility = api_def_pb2.ApiDef.HIDDEN + def _ReadFileToProto(filename): + """Read a filename, create a protobuf from its contents.""" + ret_val = api_objects_pb2.TFAPIObject() + text_format.Merge(file_io.read_file_to_string(filename), ret_val) + return ret_val - file_path = _GetApiDefFilePath(base_api_def.graph_op_name) - api_def_map[file_path].op.extend([api_def]) + golden_proto_dict = { + _FileNameToKey(filename): _ReadFileToProto(filename) + for filename in golden_file_list + } - @unittest.skipUnless( - sys.version_info.major == 2 and os.uname()[0] == 'Linux', - 'API compabitility test goldens are generated using python2 on Linux.') - def testAPIDefCompatibility(self): - # Get base ApiDef - name_to_base_api_def = self._GetBaseApiMap() - snake_to_camel_graph_op_names = { - self._GenerateLowerCaseOpName(name): name - for name in name_to_base_api_def.keys()} - # Extract Python API - visitor = python_object_to_proto_visitor.PythonObjectToProtoVisitor() - public_api_visitor = public_api.PublicAPIVisitor(visitor) - public_api_visitor.do_not_descend_map['tf'].append('contrib') - traverse.traverse(tf, public_api_visitor) - proto_dict = visitor.GetProtos() + # user_ops is an empty module. It is currently available in TensorFlow API + # but we don't keep empty modules in the new API. + # We delete user_ops from golden_proto_dict to make sure assert passes + # when diffing new API against goldens. + # TODO(annarev): remove user_ops from goldens once we switch to new API. + tf_module = golden_proto_dict['tensorflow'].tf_module + for i in range(len(tf_module.member)): + if tf_module.member[i].name == 'user_ops': + del tf_module.member[i] + break - # Map from file path to Python ApiDefs. - new_api_defs_map = defaultdict(api_def_pb2.ApiDefs) - # We need to override all endpoints even if 1 endpoint differs from base - # ApiDef. So, we first create a map from an op to all its endpoints. - op_to_endpoint_name = defaultdict(list) - - # Generate map from generated python op to endpoint names. - for public_module, value in proto_dict.items(): - module_obj = _GetSymbol(public_module) - for sym in value.tf_module.member_method: - obj = getattr(module_obj, sym.name) - - # Check if object is defined in gen_* module. That is, - # the object has been generated from OpDef. - if hasattr(obj, '__module__') and _IsGenModule(obj.__module__): - if obj.__name__ not in snake_to_camel_graph_op_names: - # Symbol might be defined only in Python and not generated from - # C++ api. - continue - relative_public_module = public_module[len('tensorflow.'):] - full_name = (relative_public_module + '.' + sym.name - if relative_public_module else sym.name) - op_to_endpoint_name[obj].append(full_name) - - # Generate Python ApiDef overrides. - for op, endpoint_names in op_to_endpoint_name.items(): - graph_op_name = snake_to_camel_graph_op_names[op.__name__] - api_def = self._CreatePythonApiDef( - name_to_base_api_def[graph_op_name], endpoint_names) - - if api_def: - file_path = _GetApiDefFilePath(graph_op_name) - api_defs = new_api_defs_map[file_path] - api_defs.op.extend([api_def]) - - self._AddHiddenOpOverrides(name_to_base_api_def, new_api_defs_map) - - old_api_defs_map = _GetGoldenApiDefs() - for file_path, new_api_defs in new_api_defs_map.items(): - # Get new ApiDef string. - new_api_defs_str = str(new_api_defs) - - # Get current ApiDef for the given file. - old_api_defs_str = ( - old_api_defs_map[file_path] if file_path in old_api_defs_map else '') - - if old_api_defs_str == new_api_defs_str: - continue - - if FLAGS.update_goldens: - logging.info('Updating %s...' % file_path) - file_io.write_string_to_file(file_path, new_api_defs_str) - else: - self.assertMultiLineEqual( - old_api_defs_str, new_api_defs_str, - 'To update golden API files, run api_compatibility_test locally ' - 'with --update_goldens=True flag.') - - for file_path in set(old_api_defs_map) - set(new_api_defs_map): - if FLAGS.update_goldens: - logging.info('Deleting %s...' % file_path) - file_io.delete_file(file_path) - else: - self.fail( - '%s file is no longer needed and should be removed.' - 'To update golden API files, run api_compatibility_test locally ' - 'with --update_goldens=True flag.' % file_path) + # Diff them. Do not fail if called with update. + # If the test is run to update goldens, only report diffs but do not fail. + self._AssertProtoDictEquals( + golden_proto_dict, + proto_dict, + verbose=FLAGS.verbose_diffs, + update_goldens=False, + additional_missing_object_message= + 'Check if tf_export decorator/call is missing for this symbol.') if __name__ == '__main__': diff --git a/tensorflow/tools/benchmark/benchmark_model.cc b/tensorflow/tools/benchmark/benchmark_model.cc index ecab6f8769ae2d0126f63580030ed6ff756015d0..15523028c726fefa13641a1369cf4274bcfb9973 100644 --- a/tensorflow/tools/benchmark/benchmark_model.cc +++ b/tensorflow/tools/benchmark/benchmark_model.cc @@ -48,33 +48,14 @@ limitations under the License. namespace tensorflow { namespace benchmark_model { -Status InitializeSession(int num_threads, const string& graph, - std::unique_ptr* session, - std::unique_ptr* graph_def) { - LOG(INFO) << "Loading TensorFlow."; +namespace { - tensorflow::SessionOptions options; - tensorflow::ConfigProto& config = options.config; - if (num_threads > 0) { - config.set_intra_op_parallelism_threads(num_threads); +Status InitializeVariables(Session* session, + const std::vector& init_ops) { + LOG(INFO) << "Initializing graph variables"; + for (const string& init_op : init_ops) { + TF_RETURN_IF_ERROR(session->Run({}, {}, {init_op}, nullptr)); } - LOG(INFO) << "Got config, " << config.device_count_size() << " devices"; - - session->reset(tensorflow::NewSession(options)); - graph_def->reset(new GraphDef()); - tensorflow::GraphDef tensorflow_graph; - Status s = ReadBinaryProto(Env::Default(), graph, graph_def->get()); - if (!s.ok()) { - LOG(ERROR) << "Could not create TensorFlow Graph: " << s; - return s; - } - - s = (*session)->Create(*(graph_def->get())); - if (!s.ok()) { - LOG(ERROR) << "Could not create TensorFlow Session: " << s; - return s; - } - return Status::OK(); } @@ -247,8 +228,56 @@ void RecordBenchmarkEntry(const string& output_prefix, TF_QCHECK_OK(node_reporter.Close()); } +void SleepSeconds(double sleep_seconds) { + if (sleep_seconds <= 0.0) { + return; + } +#ifdef PLATFORM_WINDOWS + Sleep(sleep_seconds * 1000); +#else + // Convert the inference_delay string into a timespec. + timespec req; + req.tv_sec = static_cast(sleep_seconds); + req.tv_nsec = (sleep_seconds - req.tv_sec) * 1000000000; + nanosleep(&req, nullptr); +#endif +} + +} // namespace + +Status InitializeSession(int num_threads, const string& graph, + std::unique_ptr* session, + std::unique_ptr* graph_def) { + LOG(INFO) << "Loading TensorFlow."; + + tensorflow::SessionOptions options; + tensorflow::ConfigProto& config = options.config; + if (num_threads > 0) { + config.set_intra_op_parallelism_threads(num_threads); + } + LOG(INFO) << "Got config, " << config.device_count_size() << " devices"; + + session->reset(tensorflow::NewSession(options)); + graph_def->reset(new GraphDef()); + tensorflow::GraphDef tensorflow_graph; + Status s = ReadBinaryProto(Env::Default(), graph, graph_def->get()); + if (!s.ok()) { + LOG(ERROR) << "Could not create TensorFlow Graph: " << s; + return s; + } + + s = (*session)->Create(*(graph_def->get())); + if (!s.ok()) { + LOG(ERROR) << "Could not create TensorFlow Session: " << s; + return s; + } + + return Status::OK(); +} + Status RunBenchmark(const std::vector& inputs, - const std::vector& outputs, Session* session, + const std::vector& outputs, + const std::vector& targets, Session* session, StatSummarizer* stats, int64* inference_time_us) { std::vector > input_tensors; CreateTensorsFromInputInfo(inputs, &input_tensors); @@ -264,8 +293,8 @@ Status RunBenchmark(const std::vector& inputs, RunMetadata run_metadata; const int64 start_time = Env::Default()->NowMicros(); - s = session->Run(run_options, input_tensors, outputs, {}, &output_tensors, - &run_metadata); + s = session->Run(run_options, input_tensors, outputs, targets, + &output_tensors, &run_metadata); const int64 end_time = Env::Default()->NowMicros(); *inference_time_us = end_time - start_time; @@ -283,24 +312,10 @@ Status RunBenchmark(const std::vector& inputs, return s; } -void SleepSeconds(double sleep_seconds) { - if (sleep_seconds <= 0.0) { - return; - } -#ifdef PLATFORM_WINDOWS - Sleep(sleep_seconds * 1000); -#else - // Convert the inference_delay string into a timespec. - timespec req; - req.tv_sec = static_cast(sleep_seconds); - req.tv_nsec = (sleep_seconds - req.tv_sec) * 1000000000; - nanosleep(&req, nullptr); -#endif -} - Status TimeMultipleRuns(double sleep_seconds, int num_runs, double max_time_s, const std::vector& inputs, - const std::vector& outputs, Session* session, + const std::vector& outputs, + const std::vector& targets, Session* session, StatSummarizer* stats, int64* total_time_us, int64* actual_num_runs) { *total_time_us = 0; @@ -315,7 +330,8 @@ Status TimeMultipleRuns(double sleep_seconds, int num_runs, double max_time_s, const bool until_max_time = num_runs <= 0; for (int i = 0; until_max_time || i < num_runs; ++i) { int64 time; - Status run_status = RunBenchmark(inputs, outputs, session, stats, &time); + Status run_status = + RunBenchmark(inputs, outputs, targets, session, stats, &time); stat.UpdateStat(time); (*total_time_us) += time; ++(*actual_num_runs); @@ -345,11 +361,13 @@ Status TimeMultipleRuns(double sleep_seconds, int num_runs, double max_time_s, int Main(int argc, char** argv) { string graph = "/data/local/tmp/tensorflow_inception_graph.pb"; + string init_ops_string = ""; string input_layer_string = "input:0"; string input_layer_shape_string = "1,224,224,3"; string input_layer_type_string = "float"; string input_layer_values_string = ""; string output_layer_string = "output:0"; + string target_layer_string = ""; int max_num_runs = 1000; string max_time = "10.0"; string inference_delay = "-1.0"; @@ -371,12 +389,14 @@ int Main(int argc, char** argv) { std::vector flag_list = { Flag("graph", &graph, "graph file name"), + Flag("init_ops", &init_ops_string, "init ops"), Flag("input_layer", &input_layer_string, "input layer names"), Flag("input_layer_shape", &input_layer_shape_string, "input layer shape"), Flag("input_layer_type", &input_layer_type_string, "input layer type"), Flag("input_layer_values", &input_layer_values_string, "values to initialize the inputs with"), Flag("output_layer", &output_layer_string, "output layer name"), + Flag("target_layer", &target_layer_string, "target layer name"), Flag("max_num_runs", &max_num_runs, "number of runs max"), Flag("max_time", &max_time, "length to run max"), Flag("inference_delay", &inference_delay, @@ -410,6 +430,7 @@ int Main(int argc, char** argv) { return -1; } + std::vector init_ops = str_util::Split(init_ops_string, ','); std::vector input_layers = str_util::Split(input_layer_string, ','); std::vector input_layer_shapes = str_util::Split(input_layer_shape_string, ':'); @@ -418,6 +439,7 @@ int Main(int argc, char** argv) { std::vector input_layer_values = str_util::Split(input_layer_values_string, ':'); std::vector output_layers = str_util::Split(output_layer_string, ','); + std::vector target_layers = str_util::Split(target_layer_string, ','); if ((input_layers.size() != input_layer_shapes.size()) || (input_layers.size() != input_layer_types.size())) { LOG(ERROR) << "There must be the same number of items in --input_layer," @@ -441,10 +463,12 @@ int Main(int argc, char** argv) { } LOG(INFO) << "Graph: [" << graph << "]"; + LOG(INFO) << "Init ops:" << init_ops_string; LOG(INFO) << "Input layers: [" << input_layer_string << "]"; LOG(INFO) << "Input shapes: [" << input_layer_shape_string << "]"; LOG(INFO) << "Input types: [" << input_layer_type_string << "]"; LOG(INFO) << "Output layers: [" << output_layer_string << "]"; + LOG(INFO) << "Target layers: [" << target_layer_string << "]"; LOG(INFO) << "Num runs: [" << max_num_runs << "]"; LOG(INFO) << "Inter-inference delay (seconds): [" << inference_delay << "]"; LOG(INFO) << "Inter-benchmark delay (seconds): [" << inter_benchmark_delay @@ -470,6 +494,16 @@ int Main(int argc, char** argv) { return -1; } + if (!init_ops.empty()) { + Status initialize_variables_status = + InitializeVariables(session.get(), init_ops); + if (!initialize_variables_status.ok()) { + LOG(ERROR) << "Graph variables initialization failed with " + << initialize_variables_status; + return -1; + } + } + StatSummarizerOptions stats_options; stats_options.show_run_order = show_run_order; stats_options.run_order_limit = run_order_limit; @@ -520,9 +554,10 @@ int Main(int argc, char** argv) { int64 warmup_time_us = 0; int64 num_warmup_runs = 0; if (warmup_runs > 0) { - Status warmup_time_status = TimeMultipleRuns( - inter_inference_sleep_seconds, warmup_runs, -1.0, inputs, output_layers, - session.get(), nullptr, &warmup_time_us, &num_warmup_runs); + Status warmup_time_status = + TimeMultipleRuns(inter_inference_sleep_seconds, warmup_runs, -1.0, + inputs, output_layers, target_layers, session.get(), + nullptr, &warmup_time_us, &num_warmup_runs); if (!warmup_time_status.ok()) { LOG(ERROR) << "Timing failed with " << warmup_time_status; return -1; @@ -536,8 +571,8 @@ int Main(int argc, char** argv) { int64 no_stat_num_runs = 0; Status no_stat_time_status = TimeMultipleRuns( inter_inference_sleep_seconds, max_num_runs, max_benchmark_time_seconds, - inputs, output_layers, session.get(), nullptr, &no_stat_time_us, - &no_stat_num_runs); + inputs, output_layers, target_layers, session.get(), nullptr, + &no_stat_time_us, &no_stat_num_runs); const double no_stat_wall_time = no_stat_time_us / 1000000.0; if (!no_stat_time_status.ok()) { LOG(ERROR) << "Timing failed with " << no_stat_time_status; @@ -551,8 +586,8 @@ int Main(int argc, char** argv) { int64 stat_num_runs = 0; Status stat_time_status = TimeMultipleRuns( inter_inference_sleep_seconds, max_num_runs, max_benchmark_time_seconds, - inputs, output_layers, session.get(), stats.get(), &stat_time_us, - &stat_num_runs); + inputs, output_layers, target_layers, session.get(), stats.get(), + &stat_time_us, &stat_num_runs); if (!stat_time_status.ok()) { LOG(ERROR) << "Timing failed with " << stat_time_status; return -1; diff --git a/tensorflow/tools/benchmark/benchmark_model.h b/tensorflow/tools/benchmark/benchmark_model.h index dff62c5b5d518da8f9034295626e46db783f343d..dc5f0080374e70edad52965cc0a95f99751baa48 100644 --- a/tensorflow/tools/benchmark/benchmark_model.h +++ b/tensorflow/tools/benchmark/benchmark_model.h @@ -37,13 +37,15 @@ Status InitializeSession(int num_threads, const string& graph, // Does a single run of the model that's been loaded into the given session. Status RunBenchmark(const std::vector& inputs, - const std::vector& outputs, Session* session, + const std::vector& outputs, + const std::vector& targets, Session* session, StatSummarizer* stats, int64* inference_time_us); // Runs the model multiple time, keeping track of timing information. Status TimeMultipleRuns(double sleep_seconds, int num_runs, double max_time_s, const std::vector& inputs, - const std::vector& outputs, Session* session, + const std::vector& outputs, + const std::vector& targets, Session* session, StatSummarizer* stats, int64* total_time_us, int64* actual_num_runs); diff --git a/tensorflow/tools/benchmark/benchmark_model_test.cc b/tensorflow/tools/benchmark/benchmark_model_test.cc index bb4eb5352039b01a6692621906eff005187cfa36..16ab2ff66e763a0ca5130f075f988bade9c8abd1 100644 --- a/tensorflow/tools/benchmark/benchmark_model_test.cc +++ b/tensorflow/tools/benchmark/benchmark_model_test.cc @@ -64,8 +64,8 @@ TEST(BenchmarkModelTest, InitializeAndRun) { int64 time; int64 num_runs = 0; TF_ASSERT_OK(benchmark_model::TimeMultipleRuns( - 0.0, 10, 0.0, {input}, {output_name}, session.get(), stats.get(), &time, - &num_runs)); + 0.0, 10, 0.0, {input}, {output_name}, {}, session.get(), stats.get(), + &time, &num_runs)); ASSERT_EQ(num_runs, 10); } diff --git a/tensorflow/tools/ci_build/Dockerfile.cmake b/tensorflow/tools/ci_build/Dockerfile.cmake index ec90c83aacd068e8f9c16e5be8eb6e1cef098ea6..d5dea4f3e41841aed5aeac02fcca850dbfdfaeb3 100644 --- a/tensorflow/tools/ci_build/Dockerfile.cmake +++ b/tensorflow/tools/ci_build/Dockerfile.cmake @@ -23,11 +23,12 @@ RUN /install/install_deb_packages.sh RUN apt-get update RUN apt-get install -y --no-install-recommends python-pip +RUN pip install --upgrade wheel RUN pip install --upgrade astor RUN pip install --upgrade gast RUN pip install --upgrade numpy RUN pip install --upgrade termcolor # Install golang -RUN add-apt-repository -y ppa:ubuntu-lxc/lxd-stable -RUN apt-get install -y golang +RUN apt-get install -t xenial-backports -y golang-1.9 +ENV PATH=${PATH}:/usr/lib/go-1.9/bin diff --git a/tensorflow/tools/ci_build/Dockerfile.rbe.cpu b/tensorflow/tools/ci_build/Dockerfile.rbe.cpu new file mode 100644 index 0000000000000000000000000000000000000000..6f0798b1afc34bc08df6f3f8f467a329fcf0fe9b --- /dev/null +++ b/tensorflow/tools/ci_build/Dockerfile.rbe.cpu @@ -0,0 +1,14 @@ +FROM launcher.gcr.io/google/rbe-debian8:r322167 +LABEL maintainer="Yu Yi " + +# Copy install scripts +COPY install/*.sh /install/ + +# Setup envvars +ENV CC /usr/local/bin/clang +ENV CXX /usr/local/bin/clang++ +ENV AR /usr/bin/ar + +# Run pip install script for RBE Debian8 container. +RUN /install/install_pip_packages_remote.sh +RUN /install/install_pip_packages.sh diff --git a/tensorflow/tools/ci_build/Dockerfile.rbe.gpu b/tensorflow/tools/ci_build/Dockerfile.rbe.gpu new file mode 100644 index 0000000000000000000000000000000000000000..24ff4765a619701cd614414d2b06f7fa4ce7d8c0 --- /dev/null +++ b/tensorflow/tools/ci_build/Dockerfile.rbe.gpu @@ -0,0 +1,26 @@ +FROM nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04 + +LABEL maintainer="Nick Lopez " + +# In the Ubuntu 16.04 images, cudnn is placed in system paths. Move them to +# /usr/local/cuda +RUN cp -P /usr/include/cudnn.h /usr/local/cuda/include +RUN cp -P /usr/lib/x86_64-linux-gnu/libcudnn* /usr/local/cuda/lib64 + +# Copy and run the install scripts. +COPY install/*.sh /install/ +ARG DEBIAN_FRONTEND=noninteractive +RUN /install/install_bootstrap_deb_packages.sh +RUN add-apt-repository -y ppa:openjdk-r/ppa && \ + add-apt-repository -y ppa:george-edison55/cmake-3.x +RUN /install/install_deb_packages.sh +RUN /install/install_pip_packages.sh +RUN /install/install_golang.sh + +# Install clang from pre-built package +RUN cd /tmp && \ + wget https://storage.googleapis.com/clang-builds-stable/clang-ubuntu16_04/clang_r323528.tar.gz && \ + echo "26752d9f5785df07193fac8316ba5d5ba3bec36d970c29a1577360848818ac74 clang_r323528.tar.gz" | sha256sum -c && \ + tar -C /usr/local -xf clang_r323528.tar.gz && \ + rm clang_r323528.tar.gz + diff --git a/tensorflow/tools/ci_build/builds/test_tutorials.sh b/tensorflow/tools/ci_build/builds/test_tutorials.sh index 67e5af556405a5c659000a07a79a6bd9a1d1e542..db335f14ca4f88ade7a540ffab7ed9de67f1248e 100755 --- a/tensorflow/tools/ci_build/builds/test_tutorials.sh +++ b/tensorflow/tools/ci_build/builds/test_tutorials.sh @@ -277,17 +277,6 @@ test_ptb_word_lm() { fi } - -# ----------------------------------------------------------- -# translate_test -test_translate_test() { - LOG_FILE=$1 - - run_in_directory "${TEST_DIR}" "${LOG_FILE}" \ - "${TF_MODELS_DIR}/tutorials/rnn/translate/translate.py" --self_test=True -} - - # Run the tutorial tests test_runner "tutorial test-on-install" \ "${TUT_TESTS}" "${TF_BUILD_TUT_TEST_BLACKLIST}" "${LOGS_DIR}" diff --git a/tensorflow/tools/ci_build/builds/with_the_same_user b/tensorflow/tools/ci_build/builds/with_the_same_user index 5817716c8dec37dfdfd50defb4b20b1deafced70..d4bf546d401d058bd205a70c147615c8efc4f4ba 100755 --- a/tensorflow/tools/ci_build/builds/with_the_same_user +++ b/tensorflow/tools/ci_build/builds/with_the_same_user @@ -36,8 +36,13 @@ else rm /this_is_writable_file_system fi +if [ -n "${CI_BUILD_USER_FORCE_BADNAME}" ]; then + ADDUSER_OPTS="--force-badname" +fi + getent group "${CI_BUILD_GID}" || addgroup --gid "${CI_BUILD_GID}" "${CI_BUILD_GROUP}" -getent passwd "${CI_BUILD_UID}" || adduser --gid "${CI_BUILD_GID}" --uid "${CI_BUILD_UID}" \ +getent passwd "${CI_BUILD_UID}" || adduser ${ADDUSER_OPTS} \ + --gid "${CI_BUILD_GID}" --uid "${CI_BUILD_UID}" \ --gecos "${CI_BUILD_USER} (generated by with_the_same_user script)" \ --disabled-password --home "${CI_BUILD_HOME}" --quiet "${CI_BUILD_USER}" usermod -a -G sudo "${CI_BUILD_USER}" diff --git a/tensorflow/tools/ci_build/remote/remote_docker_build.sh b/tensorflow/tools/ci_build/ci_rbe_docker_build.sh similarity index 58% rename from tensorflow/tools/ci_build/remote/remote_docker_build.sh rename to tensorflow/tools/ci_build/ci_rbe_docker_build.sh index e00a66aabaf1068c772aabce2391616518be44d4..cd811de6bdf9275b799a608381c76713a6c7a65b 100755 --- a/tensorflow/tools/ci_build/remote/remote_docker_build.sh +++ b/tensorflow/tools/ci_build/ci_rbe_docker_build.sh @@ -16,25 +16,19 @@ # Build TensorFlow Docker images for remote build # # Usage: -# remote_docker_build.sh -c # docker image for cpu build -# remote_docker_build.sh -g # docker image for gpu build - +# ci_rbe_docker_build.sh -c # docker image for cpu build +# ci_rbe_docker_build.sh -g # docker image for gpu build function main { - publish=true cpu_build=false gpu_build=false - publish=true + publish=false script_dir=$(dirname "$(readlink -f "$0")") cd $script_dir - trap cleanup_on_finish EXIT - set_script_flags $@ - build_base_image - build_tf_image if [ "$publish" = true ] ; then @@ -50,17 +44,14 @@ function set_script_flags { c) cpu_build=true ;; - f) - base_image_build_script=$OPTARG - ;; g) gpu_build=true ;; h) print_usage ;; - n) - publish=false + p) + publish=true ;; *) print_usage "ERROR: unknown option" @@ -76,7 +67,6 @@ function print_usage { echo "Usage: $(basename $0) -c | -g [options]" echo " -c build image for CPU build (base image debian8-clang)" echo " -g build image for GPU build (base image nvidia-clang)" - echo " -f the script which build the {debian8,nvidia}-clang base image" echo "[option] is one of" echo " -n not publish the locally-built image to GCR;" echo " the build process will publish image to GCR by default" @@ -87,54 +77,22 @@ function print_usage { exit 1 } - -# Build nvidia-cuba-clang base image for GPU image. -# For CPU the `clang-debian8` from Cloud Launcher will be used directly: -# https://console.cloud.google.com/launcher/details/google/clang-debian8?filter=category:developer-tools&q=clang -function build_base_image { - if [ "$gpu_build" = true ] ; then - base_image="nvidia-cuda" - # Run a 2-stage build for clang base image, see - # https://github.com/llvm-mirror/llvm/blob/master/docs/Docker.rst - $base_image_build_script \ - --source $base_image \ - --branch branches/google/stable \ - --docker-repository ${base_image}-clang --docker-tag "latest" \ - -p clang -i stage2-install-clang -i stage2-install-clang-headers \ - -- \ - -DLLVM_TARGETS_TO_BUILD=Native -DCMAKE_BUILD_TYPE=Release \ - -DBOOTSTRAP_CMAKE_BUILD_TYPE=Release \ - -DCLANG_ENABLE_BOOTSTRAP=ON \ - -DCLANG_BOOTSTRAP_TARGETS="install-clang;install-clang-headers" - fi -} - - function build_tf_image { if [ "$cpu_build" = true ] ; then - dockerfile="Dockerfile.cpu" - tf_image="tensorflow-remote" + dockerfile="Dockerfile.rbe.cpu" + tf_image="tensorflow-rbe-cpu" else - dockerfile="Dockerfile.gpu" - tf_image="tensorflow-remote-gpu" + dockerfile="Dockerfile.rbe.gpu" + tf_image="tensorflow-rbe-gpu" fi docker build -f $dockerfile -t $tf_image . } - function publish_tf_image { gcr_tf_image="gcr.io/tensorflow/${tf_image}" docker tag $tf_image $gcr_tf_image gcloud docker -- push $gcr_tf_image } - -function cleanup_on_finish { - cd $script_dir - rm -rf $llvm_docker_src - docker rmi -f ${base_image}-clang ${base_image}-clang-build -} - - main $@ diff --git a/tensorflow/tools/ci_build/ci_sanity.sh b/tensorflow/tools/ci_build/ci_sanity.sh index 106ea19d46cd7034bbe1b0248cd0d20568ad37e5..aeac085d30aef746366192361f249eb01f95e8da 100755 --- a/tensorflow/tools/ci_build/ci_sanity.sh +++ b/tensorflow/tools/ci_build/ci_sanity.sh @@ -185,7 +185,9 @@ do_pylint() { # C0330 bad-continuation # C0301 line-too-long # C0326 bad-whitespace - grep -E '(\[E|\[W0311|\[W0312|\[C0330|\[C0301|\[C0326)' ${OUTPUT_FILE} > ${ERRORS_FILE} + # W0611 unused-import + # W0622 redefined-builtin + grep -E '(\[E|\[W0311|\[W0312|\[C0330|\[C0301|\[C0326|\[W0611|\[W0622)' ${OUTPUT_FILE} > ${ERRORS_FILE} N_ERRORS=0 while read -r LINE; do @@ -321,7 +323,7 @@ do_external_licenses_check(){ EXTRA_LICENSES_FILE="$(mktemp)_extra_licenses.log" echo "Getting external dependencies for ${BUILD_TARGET}" - bazel query "attr('licenses', 'notice', deps(${BUILD_TARGET}))" --no_implicit_deps --no_host_deps --keep_going \ + bazel query "attr('licenses', 'notice', deps(${BUILD_TARGET}))" --keep_going \ | grep -E -v "^//tensorflow" \ | sed -e 's|:.*||' \ | sort \ @@ -330,7 +332,7 @@ do_external_licenses_check(){ echo echo "Getting list of external licenses mentioned in ${LICENSES_TARGET}." - bazel query "deps(${LICENSES_TARGET})" --no_implicit_deps --no_host_deps --keep_going \ + bazel query "deps(${LICENSES_TARGET})" --keep_going \ | grep -E -v "^//tensorflow" \ | sed -e 's|:.*||' \ | sort \ @@ -344,6 +346,18 @@ do_external_licenses_check(){ EXTERNAL_LICENSES_CHECK_END_TIME=$(date +'%s') + # Blacklist + echo ${MISSING_LICENSES_FILE} + grep -e "@bazel_tools//third_party/" -e "@com_google_absl//absl" -e "@org_tensorflow//" -v ${MISSING_LICENSES_FILE} > temp.txt + mv temp.txt ${MISSING_LICENSES_FILE} + + # Whitelist + echo ${EXTRA_LICENSE_FILE} + grep -e "@bazel_tools//src" -e "@bazel_tools//tools/" -e "@com_google_absl//" -e "//external" -e "@local" -v ${EXTRA_LICENSES_FILE} > temp.txt + mv temp.txt ${EXTRA_LICENSES_FILE} + + + echo echo "do_external_licenses_check took $((EXTERNAL_LICENSES_CHECK_END_TIME - EXTERNAL_LICENSES_CHECK_START_TIME)) s" echo @@ -517,9 +531,14 @@ do_check_futures_test() { python check_futures_test.py } +do_check_file_name_test() { + cd "$ROOT_DIR/tensorflow/tools/test" + python file_name_test.py +} + # Supply all sanity step commands and descriptions -SANITY_STEPS=("do_pylint PYTHON2" "do_pylint PYTHON3" "do_check_futures_test" "do_buildifier" "do_bazel_nobuild" "do_pip_package_licenses_check" "do_lib_package_licenses_check" "do_java_package_licenses_check" "do_pip_smoke_test" "do_check_load_py_test" "do_code_link_check" "do_cmake_python_sanity") -SANITY_STEPS_DESC=("Python 2 pylint" "Python 3 pylint" "Check that python files have certain __future__ imports" "buildifier check" "bazel nobuild" "pip: license check for external dependencies" "C library: license check for external dependencies" "Java Native Library: license check for external dependencies" "Pip Smoke Test: Checking py_test dependencies exist in pip package" "Check load py_test: Check that BUILD files with py_test target properly load py_test" "Code Link Check: Check there are no broken links" "Test entries in /tensorflow/contrib/cmake/python_{modules|protos|protos_cc}.txt for validity and consistency") +SANITY_STEPS=("do_pylint PYTHON2" "do_pylint PYTHON3" "do_check_futures_test" "do_buildifier" "do_bazel_nobuild" "do_pip_package_licenses_check" "do_lib_package_licenses_check" "do_java_package_licenses_check" "do_pip_smoke_test" "do_check_load_py_test" "do_code_link_check" "do_cmake_python_sanity" "do_check_file_name_test") +SANITY_STEPS_DESC=("Python 2 pylint" "Python 3 pylint" "Check that python files have certain __future__ imports" "buildifier check" "bazel nobuild" "pip: license check for external dependencies" "C library: license check for external dependencies" "Java Native Library: license check for external dependencies" "Pip Smoke Test: Checking py_test dependencies exist in pip package" "Check load py_test: Check that BUILD files with py_test target properly load py_test" "Code Link Check: Check there are no broken links" "Test entries in /tensorflow/contrib/cmake/python_{modules|protos|protos_cc}.txt for validity and consistency" "Check file names for cases") INCREMENTAL_FLAG="" DEFAULT_BAZEL_CONFIGS="--config=hdfs --config=gcp" diff --git a/tensorflow/tools/ci_build/copy_binary.py b/tensorflow/tools/ci_build/copy_binary.py index 90fd6a6e71f19649406234bc93025c15e4a5063c..ff26b052f3ea0b4866e12a663c569a3fa39ad698 100755 --- a/tensorflow/tools/ci_build/copy_binary.py +++ b/tensorflow/tools/ci_build/copy_binary.py @@ -43,7 +43,7 @@ BINARY_STRING_TEMPLATE = "%s-%s-%s.whl" def check_existence(filename): """Check the existence of file or dir.""" if not os.path.exists(filename): - raise RuntimeError("%s not found.") + raise RuntimeError("%s not found." % filename) def copy_binary(directory, origin_tag, new_tag, version, gpu=False): diff --git a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh index cfeaebdbf57c01fef7cd81dae76217429336d0ff..d0816c92b7308a1079579e605ee9af491a0533fb 100755 --- a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh +++ b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh @@ -54,3 +54,6 @@ for i in `seq 0 $((TF_GPU_COUNT-1))`; do fi done +echo "Cannot find a free GPU to run the test $* on, exiting with failure..." +exit 1 + diff --git a/tensorflow/tools/ci_build/install/install_bazel.sh b/tensorflow/tools/ci_build/install/install_bazel.sh index cf8737c2d8c746b6ad6c436745193290e31326ea..3e27a94cf2bf3110ac181d6ef5a57366be17255f 100755 --- a/tensorflow/tools/ci_build/install/install_bazel.sh +++ b/tensorflow/tools/ci_build/install/install_bazel.sh @@ -15,7 +15,7 @@ # ============================================================================== # Select bazel version. -BAZEL_VERSION="0.8.0" +BAZEL_VERSION="0.11.0" set +e local_bazel_ver=$(bazel version 2>&1 | grep -i label | awk '{print $3}') diff --git a/tensorflow/tools/ci_build/install/install_pip_packages.sh b/tensorflow/tools/ci_build/install/install_pip_packages.sh index 71744c04f2f432bc76eadfac406233ad8241a52a..d406b83a6246d18c335fb52cea1256d7809fa61a 100755 --- a/tensorflow/tools/ci_build/install/install_pip_packages.sh +++ b/tensorflow/tools/ci_build/install/install_pip_packages.sh @@ -43,8 +43,8 @@ pip2 install --upgrade werkzeug==0.11.10 pip3 install --upgrade werkzeug==0.11.10 # Install bleach. html5lib will be picked up as a dependency. -pip2 install --upgrade bleach==1.5.0 -pip3 install --upgrade bleach==1.5.0 +pip2 install --upgrade bleach==2.0.0 +pip3 install --upgrade bleach==2.0.0 # Install markdown. pip2 install --upgrade markdown==2.6.8 diff --git a/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh b/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh new file mode 100755 index 0000000000000000000000000000000000000000..39a6d557d185d8564a79315fc738a054325aa0bc --- /dev/null +++ b/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# 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. +# ============================================================================== + +set -e + +if [ ! -f /usr/bin/x86_64-linux-gnu-gcc ]; then + ln -s /usr/local/bin/clang /usr/bin/x86_64-linux-gnu-gcc +fi + +pip2 install -U pip +pip3 install -U pip +pip2 install -U setuptools +pip3 install -U setuptools + +# The rest of the pip packages will be installed in +# `install_pip_packages.sh` diff --git a/tensorflow/tools/ci_build/osx/cpu/run_contrib.sh b/tensorflow/tools/ci_build/osx/cpu/run_contrib.sh index 509ee38ec4fd584037f8e43726c01391430c1817..5c5a36139f50e85e70ce4bff5ca8054f7570b0f5 100755 --- a/tensorflow/tools/ci_build/osx/cpu/run_contrib.sh +++ b/tensorflow/tools/ci_build/osx/cpu/run_contrib.sh @@ -31,7 +31,7 @@ export CC_OPT_FLAGS='-mavx' export PYTHON_BIN_PATH=$(which python2) yes "" | $PYTHON_BIN_PATH configure.py which bazel -bazel test --test_tag_filters=-no_oss,-gpu,-benchmark-test,-nomac \ +bazel test --test_tag_filters=-no_oss,-gpu,-benchmark-test,-nomac,-no_mac \ --test_timeout 300,450,1200,3600 \ --test_size_filters=small,medium --config=opt \ --jobs=${N_JOBS} --build_tests_only --test_output=errors -k -- \ diff --git a/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh b/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh index 05547136704394ed9262f566a2bfb4160b73c7fd..338066131b5d4511ae9f0646a1269b182cf8e1fa 100755 --- a/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh +++ b/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh @@ -31,7 +31,7 @@ export CC_OPT_FLAGS='-mavx' export PYTHON_BIN_PATH=$(which python2) yes "" | $PYTHON_BIN_PATH configure.py which bazel -bazel test --test_tag_filters=-no_oss,-gpu,-benchmark-test,-nomac \ +bazel test --test_tag_filters=-no_oss,-gpu,-benchmark-test,-nomac,-no_mac \ --test_timeout 300,450,1200,3600 --config=opt \ --test_size_filters=small,medium \ --jobs=${N_JOBS} --build_tests_only --test_output=errors -k -- \ diff --git a/tensorflow/tools/ci_build/osx/cpu/run_py3_cc_core.sh b/tensorflow/tools/ci_build/osx/cpu/run_py3_cc_core.sh index 8f839ca110e5bbeba6fb7f0baaeab2fe6f126319..920a261ae3c8d68ec0b0d311fd361e3843eebd86 100755 --- a/tensorflow/tools/ci_build/osx/cpu/run_py3_cc_core.sh +++ b/tensorflow/tools/ci_build/osx/cpu/run_py3_cc_core.sh @@ -30,7 +30,7 @@ export TF_NEED_CUDA=0 export PYTHON_BIN_PATH=$(which python3) yes "" | $PYTHON_BIN_PATH configure.py which bazel -bazel test --test_tag_filters=-no_oss,-gpu,-benchmark-test,-nomac \ +bazel test --test_tag_filters=-no_oss,-gpu,-benchmark-test,-nomac,-no_mac \ --test_timeout 300,450,1200,3600 \ --test_size_filters=small,medium \ --jobs=${N_JOBS} --build_tests_only --test_output=errors -k -- \ diff --git a/tensorflow/tools/ci_build/remote/Dockerfile.cpu b/tensorflow/tools/ci_build/remote/Dockerfile.cpu deleted file mode 100644 index 7b01d8320d26f38c92ad8f404da3188809a6d400..0000000000000000000000000000000000000000 --- a/tensorflow/tools/ci_build/remote/Dockerfile.cpu +++ /dev/null @@ -1,27 +0,0 @@ -FROM launcher.gcr.io/google/clang-debian8:latest - -RUN apt-get update && apt-get --no-install-recommends install -y \ - binutils \ - binutils-gold \ - curl \ - libstdc++-4.9-dev \ - python \ - python-dev \ - python-numpy \ - python-pip \ - unzip \ - zip && \ - rm -rf /var/lib/apt/lists/* - -RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \ - python get-pip.py && \ - rm get-pip.py - -# Set up grpc -RUN pip install --upgrade enum34 futures mock numpy six backports.weakref portpicker && \ - pip install --pre 'protobuf>=3.0.0a3' && \ - pip install 'grpcio>=1.1.3' - -# TODO: Set up golang which is compatible with clang - -WORKDIR /botexec diff --git a/tensorflow/tools/ci_build/remote/Dockerfile.gpu b/tensorflow/tools/ci_build/remote/Dockerfile.gpu deleted file mode 100644 index 47ffd44163dd3e4b99f06689e1aa6f19f84cc2ca..0000000000000000000000000000000000000000 --- a/tensorflow/tools/ci_build/remote/Dockerfile.gpu +++ /dev/null @@ -1,27 +0,0 @@ -FROM nvidia-cuda-clang:latest - -RUN apt-get update && apt-get --no-install-recommends install -y \ - binutils \ - binutils-gold \ - curl \ - libstdc++-4.9-dev \ - python \ - python-dev \ - python-numpy \ - python-pip \ - unzip \ - zip && \ - rm -rf /var/lib/apt/lists/* - -RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \ - python get-pip.py && \ - rm get-pip.py - -# Set up grpc -RUN pip install --upgrade \ - enum34 futures astor gast mock numpy six \ - backports.weakref termcolor && \ - pip install --pre 'protobuf>=3.0.0a3' && \ - pip install 'grpcio>=1.1.3' - -WORKDIR /botexec diff --git a/tensorflow/tools/ci_build/update_version.py b/tensorflow/tools/ci_build/update_version.py index 347d0769a92cc767f2e263fce0e21d7d0bc8e586..52a0da9a14847e863d92fee9ef7e63e4af0cf068 100755 --- a/tensorflow/tools/ci_build/update_version.py +++ b/tensorflow/tools/ci_build/update_version.py @@ -261,14 +261,12 @@ def major_minor_change(old_version, new_version): def update_dockerfiles(old_version, new_version): """Update dockerfiles if there was a major change.""" if major_minor_change(old_version, new_version): - old_r_major_minor = r"r%s\.%s" % (old_version.major, old_version.minor) - old_r_major_minor_string = old_r_major_minor.replace("\\", "") - r_major_minor = r"r%s\.%s" % (new_version.major, new_version.minor) - r_major_minor_string = r_major_minor.replace("\\", "") + old_r_major_minor = "r%s.%s" % (old_version.major, old_version.minor) + r_major_minor = "r%s.%s" % (new_version.major, new_version.minor) print("Detected Major.Minor change.") print("Updating pattern %s to %s in additional files" - % (old_r_major_minor_string, r_major_minor_string)) + % (old_r_major_minor, r_major_minor)) # Update dockerfiles replace_string_in_line(old_r_major_minor, r_major_minor, DEVEL_DOCKERFILE) diff --git a/tensorflow/tools/ci_build/windows/bazel/common_env.sh b/tensorflow/tools/ci_build/windows/bazel/common_env.sh index 1c35d74af72ad0a72b0016356888c8cf77e20e56..7d4cc7ac3005f7ff9a79d18228e86d6b74e1e8b0 100644 --- a/tensorflow/tools/ci_build/windows/bazel/common_env.sh +++ b/tensorflow/tools/ci_build/windows/bazel/common_env.sh @@ -34,6 +34,9 @@ export BAZEL_SH=${BAZEL_SH:-"C:/tools/msys64/usr/bin/bash"} export PYTHON_BASE_PATH="${PYTHON_DIRECTORY:-Program Files/Anaconda3}" +# Set the path to find bazel. +export PATH="/c/tools/bazel/:$PATH" + # Set Python path for ./configure export PYTHON_BIN_PATH="C:/${PYTHON_BASE_PATH}/python.exe" export PYTHON_LIB_PATH="C:/${PYTHON_BASE_PATH}/lib/site-packages" diff --git a/tensorflow/tools/ci_build/windows/cpu/cmake/run_build.bat b/tensorflow/tools/ci_build/windows/cpu/cmake/run_build.bat index 957729bb37db3ae49800c277f4090a52117c699d..c1bc71850754c5b4b42a6eb50be465ba8f98c218 100644 --- a/tensorflow/tools/ci_build/windows/cpu/cmake/run_build.bat +++ b/tensorflow/tools/ci_build/windows/cpu/cmake/run_build.bat @@ -36,7 +36,7 @@ SET CMAKE_DIR=%REPO_ROOT%\tensorflow\contrib\cmake SET MSBUILD_EXE="C:\Program Files (x86)\MSBuild\14.0\Bin\msbuild.exe" :: Run cmake to create Visual Studio Project files. -%CMAKE_EXE% %CMAKE_DIR% -A x64 -DSWIG_EXECUTABLE=%SWIG_EXE% -DPYTHON_EXECUTABLE=%PY_EXE% -DCMAKE_BUILD_TYPE=Release -DPYTHON_LIBRARIES=%PY_LIB% -Dtensorflow_BUILD_PYTHON_TESTS=%BUILD_PYTHON_TESTS% -Dtensorflow_BUILD_CC_TESTS=%BUILD_CC_TESTS% -Dtensorflow_TF_NIGHTLY=%TF_NIGHTLY% -Dtensorflow_DISABLE_EIGEN_FORCEINLINE=%DISABLE_FORCEINLINE% +%CMAKE_EXE% %CMAKE_DIR% -A x64 -DSWIG_EXECUTABLE=%SWIG_EXE% -DPYTHON_EXECUTABLE=%PY_EXE% -DCMAKE_BUILD_TYPE=Release -DPYTHON_LIBRARIES=%PY_LIB% -Dtensorflow_BUILD_PYTHON_TESTS=%BUILD_PYTHON_TESTS% -Dtensorflow_BUILD_CC_TESTS=%BUILD_CC_TESTS% -Dtensorflow_TF_NIGHTLY=%TF_NIGHTLY% -Dtensorflow_DISABLE_EIGEN_FORCEINLINE=%DISABLE_FORCEINLINE% -Dtensorflow_WIN_CPU_SIMD_OPTIONS=/arch:AVX :: Run msbuild in the resulting VS project files to build a pip package. %MSBUILD_EXE% /p:Configuration=Release /maxcpucount:32 tf_python_build_pip_package.vcxproj diff --git a/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh b/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh index 8b8ba31a0dda88ad3c43330e0208a9fa6a7d0276..40189a6d1b1388092dc397251f4b581726995a77 100644 --- a/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh +++ b/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh @@ -65,4 +65,5 @@ bazel test -c opt $BUILD_OPTS -k --test_output=errors \ --define=no_tensorflow_py_deps=true --test_lang_filters=py \ --test_tag_filters=-no_pip,-no_windows,-no_oss \ --build_tag_filters=-no_pip,-no_windows,-no_oss --build_tests_only \ - //${PY_TEST_DIR}/tensorflow/python/... + //${PY_TEST_DIR}/tensorflow/python/... \ + //${PY_TEST_DIR}/tensorflow/contrib/... diff --git a/tensorflow/tools/ci_build/windows/gpu/cmake/run_build.bat b/tensorflow/tools/ci_build/windows/gpu/cmake/run_build.bat index 5a362de3992156fea8a5fc6ab4c70ba67ab47f89..4656afe0256d03540fed6912677c8e93f9cf9eb6 100644 --- a/tensorflow/tools/ci_build/windows/gpu/cmake/run_build.bat +++ b/tensorflow/tools/ci_build/windows/gpu/cmake/run_build.bat @@ -37,7 +37,7 @@ SET CMAKE_DIR=%REPO_ROOT%\tensorflow\contrib\cmake SET MSBUILD_EXE="C:\Program Files (x86)\MSBuild\14.0\Bin\msbuild.exe" :: Run cmake to create Visual Studio Project files. -%CMAKE_EXE% %CMAKE_DIR% -A x64 -DSWIG_EXECUTABLE=%SWIG_EXE% -DPYTHON_EXECUTABLE=%PY_EXE% -DCMAKE_BUILD_TYPE=Release -DPYTHON_LIBRARIES=%PY_LIB% -Dtensorflow_BUILD_PYTHON_TESTS=%BUILD_PYTHON_TESTS% -Dtensorflow_BUILD_CC_TESTS=%BUILD_CC_TESTS% -Dtensorflow_ENABLE_GPU=ON -DCUDNN_HOME=%CUDNN_HOME% -Dtensorflow_TF_NIGHTLY=%TF_NIGHTLY% -Dtensorflow_DISABLE_EIGEN_FORCEINLINE=%DISABLE_FORCEINLINE% +%CMAKE_EXE% %CMAKE_DIR% -A x64 -DSWIG_EXECUTABLE=%SWIG_EXE% -DPYTHON_EXECUTABLE=%PY_EXE% -DCMAKE_BUILD_TYPE=Release -DPYTHON_LIBRARIES=%PY_LIB% -Dtensorflow_BUILD_PYTHON_TESTS=%BUILD_PYTHON_TESTS% -Dtensorflow_BUILD_CC_TESTS=%BUILD_CC_TESTS% -Dtensorflow_ENABLE_GPU=ON -DCUDNN_HOME=%CUDNN_HOME% -Dtensorflow_TF_NIGHTLY=%TF_NIGHTLY% -Dtensorflow_DISABLE_EIGEN_FORCEINLINE=%DISABLE_FORCEINLINE% -Dtensorflow_WIN_CPU_SIMD_OPTIONS=/arch:AVX -G"Visual Studio 14" :: Run msbuild in the resulting VS project files to build a pip package. %MSBUILD_EXE% /p:Configuration=Release /maxcpucount:32 tf_python_build_pip_package.vcxproj diff --git a/tensorflow/tools/ci_build/windows/libtensorflow_cpu.sh b/tensorflow/tools/ci_build/windows/libtensorflow_cpu.sh index fa28e3d79ca4ee5f429a41dd3e871248d5c047ca..583d1d5f09527861015458c636af2259b34d45f8 100755 --- a/tensorflow/tools/ci_build/windows/libtensorflow_cpu.sh +++ b/tensorflow/tools/ci_build/windows/libtensorflow_cpu.sh @@ -41,7 +41,7 @@ run_configure_for_cpu_build # build_libtensorflow_tarball in ../builds/libtensorflow.sh # cannot be used on Windows since it relies on pkg_tar rules. # So we do something special here -bazel build -c opt \ +bazel build -c opt --copt=/arch:AVX \ tensorflow:libtensorflow.so \ tensorflow/tools/lib_package:clicenses_generate \ tensorflow/java:libtensorflow_jni.so \ diff --git a/tensorflow/tools/ci_build/windows/libtensorflow_gpu.sh b/tensorflow/tools/ci_build/windows/libtensorflow_gpu.sh index 573c926203fc76b787ba08b10bd71c8effda29b6..94276c6c5c9ce897ca24f03efe3d93e1ea1e00c9 100644 --- a/tensorflow/tools/ci_build/windows/libtensorflow_gpu.sh +++ b/tensorflow/tools/ci_build/windows/libtensorflow_gpu.sh @@ -41,7 +41,7 @@ run_configure_for_gpu_build # build_libtensorflow_tarball in ../builds/libtensorflow.sh # cannot be used on Windows since it relies on pkg_tar rules. # So we do something special here -bazel build -c opt \ +bazel build -c opt --copt=/arch:AVX \ tensorflow:libtensorflow.so \ tensorflow/tools/lib_package:clicenses_generate \ tensorflow/java:libtensorflow_jni.so \ diff --git a/tensorflow/tools/compatibility/tf_upgrade.py b/tensorflow/tools/compatibility/tf_upgrade.py index 6e90b286c99f894ddd25268afc69043759571c36..1f8833582af4c922115e637117e775e619439786 100644 --- a/tensorflow/tools/compatibility/tf_upgrade.py +++ b/tensorflow/tools/compatibility/tf_upgrade.py @@ -662,9 +662,9 @@ class TFAPIChangeSpec(APIChangeSpec): def _reverse_handler(file_edit_recorder, node): # TODO(aselle): Could check for a literal list of bools and try to convert # them to indices. - comment = ("ERROR: tf.reverse has had its argument semantics changed\n" - "significantly the converter cannot detect this reliably, so you" - "need to inspect this usage manually.\n") + comment = ("ERROR: tf.reverse has had its argument semantics changed " + "significantly the converter cannot detect this reliably, so " + "you need to inspect this usage manually.\n") file_edit_recorder.add( comment, node.lineno, diff --git a/tensorflow/tools/compatibility/tf_upgrade_test.py b/tensorflow/tools/compatibility/tf_upgrade_test.py index a495f9883b284869d043441d1cfecca01296eda3..3d02eacba6e7a91e6d3c88e8297306de9782f4bf 100644 --- a/tensorflow/tools/compatibility/tf_upgrade_test.py +++ b/tensorflow/tools/compatibility/tf_upgrade_test.py @@ -114,7 +114,7 @@ class TestUpgrade(test_util.TensorFlowTestCase): self.assertEqual(errors, ["test.py:1: tf.reverse requires manual check."]) def testListComprehension(self): - def _test(input, output): + def _test(input, output): # pylint: disable=redefined-builtin _, unused_report, errors, new_text = self._upgrade(input) self.assertEqual(new_text, output) _test("tf.concat(0, \t[x for x in y])\n", diff --git a/tensorflow/tools/def_file_filter/BUILD b/tensorflow/tools/def_file_filter/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tensorflow/tools/def_file_filter/BUILD.tpl b/tensorflow/tools/def_file_filter/BUILD.tpl new file mode 100644 index 0000000000000000000000000000000000000000..3cb72f49797d8014e4df3f44155c967ab0e2f9b3 --- /dev/null +++ b/tensorflow/tools/def_file_filter/BUILD.tpl @@ -0,0 +1,15 @@ +# Description: +# Tools for filtering DEF file for TensorFlow on Windows +# +# On Windows, we use a DEF file generated by Bazel to export +# symbols from the tensorflow dynamic library(_pywrap_tensorflow.dll). +# The maximum number of symbols that can be exported per DLL is 64K, +# so we have to filter some useless symbols through this python script. + +package(default_visibility = ["//visibility:public"]) + +py_binary( + name = "def_file_filter", + srcs = ["def_file_filter.py"], + srcs_version = "PY2AND3", +) diff --git a/tensorflow/tools/def_file_filter/def_file_filter.py.tpl b/tensorflow/tools/def_file_filter/def_file_filter.py.tpl new file mode 100644 index 0000000000000000000000000000000000000000..8bdc03eb0f19fd6daae826727f429bc1255f0eca --- /dev/null +++ b/tensorflow/tools/def_file_filter/def_file_filter.py.tpl @@ -0,0 +1,168 @@ +# 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. +# ============================================================================== + +"""def_file_filter.py - tool to filter a windows def file. + +The def file can be used to export symbols from the tensorflow dll to enable +tf.load_library(). + +Because the linker allows only 64K symbols to be exported per dll +we filter the symbols down to the essentials. The regular expressions +we use for this are specific to tensorflow. + +TODO: this works fine but there is an issue with exporting +'const char * const' and importing it from a user_ops. The problem is +on the importing end and using __declspec(dllimport) works around it. +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import argparse +import io +import os +import re +import subprocess +import sys +import tempfile + +# External tools we use that come with visual studio sdk +UNDNAME = "%{undname_bin_path}" + +# Exclude if matched +EXCLUDE_RE = re.compile(r"RTTI|deleting destructor|::internal::") + +# Include if matched before exclude +INCLUDEPRE_RE = re.compile(r"google::protobuf::internal::ExplicitlyConstructed|" + r"google::protobuf::internal::ArenaImpl::AllocateAligned|" # for contrib/data/_prefetching_ops + r"google::protobuf::internal::ArenaImpl::AddCleanup|" # for contrib/data/_prefetching_ops + r"google::protobuf::Arena::OnArenaAllocation|" # for contrib/data/_prefetching_ops + r"tensorflow::internal::LogMessage|" + r"tensorflow::internal::LogString|" + r"tensorflow::internal::CheckOpMessageBuilder|" + r"tensorflow::internal::MakeCheckOpValueString|" + r"tensorflow::internal::PickUnusedPortOrDie|" + r"tensorflow::internal::ValidateDevice|" + r"tensorflow::ops::internal::Enter|" + r"tensorflow::strings::internal::AppendPieces|" + r"tensorflow::strings::internal::CatPieces|" + r"tensorflow::io::internal::JoinPathImpl") + +# Include if matched after exclude +INCLUDE_RE = re.compile(r"^(TF_\w*)$|" + r"^(TFE_\w*)$|" + r"nsync::|" + r"tensorflow::|" + r"functor::|" + r"perftools::gputools") + +# We want to identify data members explicitly in the DEF file, so that no one +# can implicitly link against the DLL if they use one of the variables exported +# from the DLL and the header they use does not decorate the symbol with +# __declspec(dllimport). It is easier to detect what a data symbol does +# NOT look like, so doing it with the below regex. +DATA_EXCLUDE_RE = re.compile(r"[)(]|" + r"vftable|" + r"vbtable|" + r"vcall|" + r"RTTI|" + r"protobuf::internal::ExplicitlyConstructed") + +def get_args(): + """Parse command line.""" + filename_list = lambda x: x.split(";") + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=filename_list, + help="paths to input def file", + required=True) + parser.add_argument("--output", help="output deffile", required=True) + parser.add_argument("--target", help="name of the target", required=True) + args = parser.parse_args() + return args + + +def main(): + """main.""" + args = get_args() + + # Pipe dumpbin to extract all linkable symbols from libs. + # Good symbols are collected in candidates and also written to + # a temp file. + candidates = [] + tmpfile = tempfile.NamedTemporaryFile(mode="w", delete=False) + for def_file_path in args.input: + def_file = open(def_file_path, 'r') + for line in def_file: + cols = line.split() + sym = cols[0] + tmpfile.file.write(sym + "\n") + candidates.append(sym) + tmpfile.file.close() + + # Run the symbols through undname to get their undecorated name + # so we can filter on something readable. + with open(args.output, "w") as def_fp: + # track dupes + taken = set() + + # Header for the def file. + def_fp.write("LIBRARY " + args.target + "\n") + def_fp.write("EXPORTS\n") + def_fp.write("\t ??1OpDef@tensorflow@@UEAA@XZ\n") + + # Each symbols returned by undname matches the same position in candidates. + # 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")): + decorated = candidates[idx] + if decorated in taken: + # Symbol is already in output, done. + dupes += 1 + continue + + if not INCLUDEPRE_RE.search(line): + if EXCLUDE_RE.search(line): + continue + if not INCLUDE_RE.search(line): + continue + + if "deleting destructor" in line: + # Some of the symbols convered by INCLUDEPRE_RE export deleting + # destructor symbols, which is a bad idea. + # So we filter out such symbols here. + continue + + if DATA_EXCLUDE_RE.search(line): + def_fp.write("\t" + decorated + "\n") + else: + def_fp.write("\t" + decorated + " DATA\n") + taken.add(decorated) + def_fp.close() + + exit_code = proc.wait() + if exit_code != 0: + print("{} failed, exit={}".format(UNDNAME, exit_code)) + return exit_code + + os.unlink(tmpfile.name) + + print("symbols={}, taken={}, dupes={}" + .format(len(candidates), len(taken), dupes)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tensorflow/tools/def_file_filter/def_file_filter_configure.bzl b/tensorflow/tools/def_file_filter/def_file_filter_configure.bzl new file mode 100644 index 0000000000000000000000000000000000000000..47539b2423e602bb9771541ae5b01ba76c79f56f --- /dev/null +++ b/tensorflow/tools/def_file_filter/def_file_filter_configure.bzl @@ -0,0 +1,56 @@ +"""Repository rule for def file filter autoconfiguration. + +This repository reuses Bazel's VC detect mechanism to find undname.exe, +which is a tool used in def_file_filter.py. + +def_file_filter.py is for filtering the DEF file for TensorFlow on Windows. +On Windows, we use a DEF file generated by Bazel to export symbols from the +tensorflow dynamic library(_pywrap_tensorflow.dll). The maximum number of +symbols that can be exported per DLL is 64K, so we have to filter some useless +symbols through this python script. + +`def_file_filter_config` depends on the following environment variables: + * `BAZEL_VC` + * `BAZEL_VS` + * `VS90COMNTOOLS` + * `VS100COMNTOOLS` + * `VS110COMNTOOLS` + * `VS120COMNTOOLS` + * `VS140COMNTOOLS` +""" + +load("@bazel_tools//tools/cpp:windows_cc_configure.bzl", "find_vc_path") +load("@bazel_tools//tools/cpp:windows_cc_configure.bzl", "find_msvc_tool") +load("@bazel_tools//tools/cpp:lib_cc_configure.bzl", "auto_configure_fail") + +def _def_file_filter_configure_impl(repository_ctx): + if repository_ctx.os.name.lower().find("windows") == -1: + repository_ctx.symlink(Label("//tensorflow/tools/def_file_filter:BUILD.tpl"), "BUILD") + repository_ctx.file("def_file_filter.py", "") + return + vc_path = find_vc_path(repository_ctx) + if vc_path == "visual-studio-not-found": + auto_configure_fail("Visual C++ build tools not found on your machine") + undname_bin_path = find_msvc_tool(repository_ctx, vc_path, "undname.exe").replace("\\", "\\\\") + + repository_ctx.template( + "def_file_filter.py", + Label("//tensorflow/tools/def_file_filter:def_file_filter.py.tpl"), + { + "%{undname_bin_path}": undname_bin_path, + }) + repository_ctx.symlink(Label("//tensorflow/tools/def_file_filter:BUILD.tpl"), "BUILD") + + +def_file_filter_configure = repository_rule( + implementation = _def_file_filter_configure_impl, + environ = [ + "BAZEL_VC", + "BAZEL_VS", + "VS90COMNTOOLS", + "VS100COMNTOOLS", + "VS110COMNTOOLS", + "VS120COMNTOOLS", + "VS140COMNTOOLS" + ], +) diff --git a/tensorflow/tools/dist_test/README.md b/tensorflow/tools/dist_test/README.md index c1b1f79bbd4b657768b9bbcab93efa3354774915..228d5ee35d1839c60b51a85bd606c1ba86e46886 100644 --- a/tensorflow/tools/dist_test/README.md +++ b/tensorflow/tools/dist_test/README.md @@ -17,6 +17,14 @@ cesnsu model: ./local_test.sh --model_name CENSUS_WIDENDEEP +You can test specify version of TensorFlow: + +```shell +./local_test.sh ${whl_file_url} +``` + +For example, you can find these TensorFlow python package URLs from [here](https://www.tensorflow.org/install/install_linux#the_url_of_the_tensorflow_python_package) for Ubuntu. + **2) Launch a remote k8s cluster on Google Kubernetes Engine (GKE) and run the test suite on it** diff --git a/tensorflow/tools/dist_test/build_server.sh b/tensorflow/tools/dist_test/build_server.sh index 878fabd248f3c1dd5cb79983df5220ebf5893026..225c0347416ec8c8fef855946d18e838bd767690 100755 --- a/tensorflow/tools/dist_test/build_server.sh +++ b/tensorflow/tools/dist_test/build_server.sh @@ -16,14 +16,15 @@ # # Builds the test server for distributed (GRPC) TensorFlow # -# Usage: build_server.sh [--test] +# Usage: build_server.sh [--test] # # Arguments: # docker_image_name: Name of the docker image to build. # E.g.: tensorflow/tf_grpc_test_server:0.11.0rc1 # -# whl_url: URL from which the TensorFlow whl file will be downloaded. +# whl_file_location: URL from which the TensorFlow whl file will be downloaded. # E.g.: https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-0.11.0rc1-cp27-none-linux_x86_64.whl +# E.g.: /path/to/folder/tensorflow-0.11.0rc1-cp27-none-linux_x86_64.whl # # The optional flag --test lets the script to use the Dockerfile for the # testing GRPC server. Without the flag, the script will build the non-test @@ -41,11 +42,11 @@ die() { # Check arguments if [[ $# -lt 2 ]]; then - die "Usage: $0 [--test]" + die "Usage: $0 [--test]" fi DOCKER_IMG_NAME=$1 -WHL_URL=$2 +WHL_FILE_LOCATION=$2 shift 2 # Current script directory @@ -53,7 +54,7 @@ DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" BUILD_DIR=$(mktemp -d) echo "" -echo "Using whl file URL: ${WHL_URL}" +echo "Using whl file URL: ${WHL_FILE_LOCATION}" echo "Building in temporary directory: ${BUILD_DIR}" cp -r ${DIR}/* "${BUILD_DIR}"/ || \ @@ -65,9 +66,15 @@ if [[ $1 == "--test" ]]; then fi echo "Using Docker file: ${DOCKER_FILE}" +if [[ $WHL_FILE_LOCATION =~ 'http://' || $WHL_FILE_LOCATION =~ 'https://' ]]; then + # Download whl file into the build context directory. + wget -P "${BUILD_DIR}" "${WHL_FILE_LOCATION}" || \ + die "Failed to download tensorflow whl file from URL: ${WHL_FILE_LOCATION}" +else + cp "${WHL_FILE_LOCATION}" "${BUILD_DIR}" +fi + # Download whl file into the build context directory. -wget -P "${BUILD_DIR}" ${WHL_URL} || \ - die "Failed to download tensorflow whl file from URL: ${WHL_URL}" if [[ ! -f "${DOCKER_FILE}" ]]; then die "ERROR: Unable to find dockerfile: ${DOCKER_FILE}" diff --git a/tensorflow/tools/dist_test/local_test.sh b/tensorflow/tools/dist_test/local_test.sh index 7d7f92d246e1ca0b519ac3bf30fde673621ff755..caae7fd5305af9846628eaf00348dd08df4e827f 100755 --- a/tensorflow/tools/dist_test/local_test.sh +++ b/tensorflow/tools/dist_test/local_test.sh @@ -16,27 +16,27 @@ # # Tests distributed TensorFlow on a locally running TF GRPC cluster. # -# This script peforms the following steps: -# 1) Build the docker-in-docker (dind) image capable of running docker and -# Kubernetes (k8s) cluster inside. +# This script performs the following steps: +# 1) Build the docker image capable of running distributed TensorFlow in docker. # 2) Run a container from the aforementioned image and start docker service # in it -# 3) Call a script to launch a k8s TensorFlow GRPC cluster inside the container +# 3) Call a script to launch a distributed TensorFlow GRPC cluster inside the container # and run the distributed test suite. # -# Usage: local_test.sh +# Usage: local_test.sh # [--leave_container_running] # [--model_name ] # [--num_workers ] # [--num_parameter_servers ] # [--sync_replicas] # -# E.g., local_test.sh --model_name CENSUS_WIDENDEEP -# local_test.sh --num_workers 3 --num_parameter_servers 3 +# E.g., local_test.sh --model_name CENSUS_WIDENDEEP +# local_test.sh --num_workers 3 --num_parameter_servers 3 # # Arguments: -# -# Specify custom TensorFlow whl file URL to install in the test Docker image. +# whl_file_location: URL from which the TensorFlow whl file will be acquired. +# E.g.: https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-0.11.0rc1-cp27-none-linux_x86_64.whl +# E.g.: /path/to/folder/tensorflow-0.11.0rc1-cp27-none-linux_x86_64.whl # # --leave_container_running: Do not stop the docker-in-docker container after # the termination of the tests, e.g., for debugging @@ -63,15 +63,9 @@ die() { # Configurations DOCKER_IMG_NAME="tensorflow/tf-dist-test-local-cluster" -LOCAL_K8S_CACHE=${HOME}/kubernetes -# Helper function -get_container_id_by_image_name() { - # Get the id of a container by image name - # Usage: get_docker_container_id_by_image_name - - docker ps | grep $1 | awk '{print $1}' -} +# Use TensorFlow v1.5.0 for Python 2.7 and CPU only as we set num_gpus to 0 in the below +DEFAULT_WHL_FILE_LOCATION="https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0-cp27-none-linux_x86_64.whl" # Parse input arguments LEAVE_CONTAINER_RUNNING=0 @@ -81,9 +75,10 @@ NUM_WORKERS=2 NUM_PARAMETER_SERVERS=2 SYNC_REPLICAS_FLAG="" -WHL_URL=${1} -if [[ -z "${WHL_URL}" ]]; then - die "whl file URL is not specified" +WHL_FILE_LOCATION=${1} +if [[ -z "${WHL_FILE_LOCATION}" ]]; then + WHL_FILE_LOCATION=${DEFAULT_WHL_FILE_LOCATION} + echo "use default whl file location" fi while true; do @@ -98,8 +93,8 @@ while true; do NUM_PARAMETER_SERVERS=$2 elif [[ $1 == "--sync_replicas" ]]; then SYNC_REPLICAS_FLAG="--sync_replicas" - elif [[ $1 == "--whl_url" ]]; then - WHL_URL=$2 + elif [[ $1 == "--WHL_FILE_LOCATION" ]]; then + WHL_FILE_LOCATION=$2 fi shift @@ -120,7 +115,7 @@ DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # Get utility functions source ${DIR}/scripts/utils.sh -# Build docker-in-docker image for local k8s cluster. +# Build docker image for local distributed TensorFlow cluster. NO_CACHE_FLAG="" if [[ ! -z "${TF_DIST_DOCKER_NO_CACHE}" ]] && [[ "${TF_DIST_DOCKER_NO_CACHE}" != "0" ]]; then @@ -130,15 +125,19 @@ fi # Create docker build context directory. BUILD_DIR=$(mktemp -d) echo "" -echo "Using whl file URL: ${WHL_URL}" +echo "Using whl file location: ${WHL_FILE_LOCATION}" echo "Building in temporary directory: ${BUILD_DIR}" cp -r ${DIR}/* "${BUILD_DIR}"/ || \ die "Failed to copy files to ${BUILD_DIR}" -# Download whl file into the build context directory. -wget -P "${BUILD_DIR}" ${WHL_URL} || \ - die "Failed to download tensorflow whl file from URL: ${WHL_URL}" +if [[ $WHL_FILE_LOCATION =~ 'http://' || $WHL_FILE_LOCATION =~ 'https://' ]]; then + # Download whl file into the build context directory. + wget -P "${BUILD_DIR}" "${WHL_FILE_LOCATION}" || \ + die "Failed to download tensorflow whl file from URL: ${WHL_FILE_LOCATION}" +else + cp "${WHL_FILE_LOCATION}" "${BUILD_DIR}" +fi # Build docker image for test. docker build ${NO_CACHE_FLAG} -t ${DOCKER_IMG_NAME} \ diff --git a/tensorflow/tools/dist_test/python/mnist_replica.py b/tensorflow/tools/dist_test/python/mnist_replica.py index a2d12442c44553a287637029843021b7541fa3fa..d6e7f317dd0b52203e354676425dbbbcd53e1973 100644 --- a/tensorflow/tools/dist_test/python/mnist_replica.py +++ b/tensorflow/tools/dist_test/python/mnist_replica.py @@ -56,7 +56,7 @@ flags.DEFINE_integer("task_index", None, flags.DEFINE_integer("num_gpus", 1, "Total number of gpus for each machine." "If you don't use GPU, please set it to '0'") flags.DEFINE_integer("replicas_to_aggregate", None, - "Number of replicas to aggregate before parameter update" + "Number of replicas to aggregate before parameter update " "is applied (For sync_replicas mode only; default: " "num_workers)") flags.DEFINE_integer("hidden_units", 100, diff --git a/tensorflow/tools/docker/Dockerfile.devel b/tensorflow/tools/docker/Dockerfile.devel index 5dc4a053fd2cae7d83739507fea31e7afc92d77c..22c73c3fe13f2cb763295fa25b43e2f82c0e8962 100644 --- a/tensorflow/tools/docker/Dockerfile.devel +++ b/tensorflow/tools/docker/Dockerfile.devel @@ -57,7 +57,7 @@ RUN echo "startup --batch" >>/etc/bazel.bazelrc RUN echo "build --spawn_strategy=standalone --genrule_strategy=standalone" \ >>/etc/bazel.bazelrc # Install the most recent bazel release. -ENV BAZEL_VERSION 0.8.0 +ENV BAZEL_VERSION 0.11.0 WORKDIR / RUN mkdir /bazel && \ cd /bazel && \ @@ -70,7 +70,7 @@ RUN mkdir /bazel && \ # Download and build TensorFlow. WORKDIR /tensorflow -RUN git clone --branch=r1.5 --depth=1 https://github.com/tensorflow/tensorflow.git . +RUN git clone --branch=r1.6 --depth=1 https://github.com/tensorflow/tensorflow.git . # TODO(craigcitro): Don't install the pip package, since it makes it # more difficult to experiment with local changes. Instead, just add diff --git a/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl b/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl index 96b260ad3aeb78622dd1ad276f7d524dd598e3bf..3690e7dfe57a4682276a90b10cb84c9a329b3f5e 100644 --- a/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl +++ b/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl @@ -3,7 +3,7 @@ FROM tensorflow/tensorflow:latest-devel LABEL maintainer="Clayne Robison" # These arguments are parameterized. Use --build-args to override. -ARG TF_BRANCH=r1.5 +ARG TF_BRANCH=r1.6 ARG WHL_DIR=/whl RUN apt-get update && apt-get install -y --no-install-recommends \ diff --git a/tensorflow/tools/docker/Dockerfile.devel-gpu b/tensorflow/tools/docker/Dockerfile.devel-gpu index 07ffd3839a32ef194100322e54b9133412e4b664..69ba340f9201266fd2c2f86571e83f6acdcda950 100644 --- a/tensorflow/tools/docker/Dockerfile.devel-gpu +++ b/tensorflow/tools/docker/Dockerfile.devel-gpu @@ -66,7 +66,7 @@ RUN echo "startup --batch" >>/etc/bazel.bazelrc RUN echo "build --spawn_strategy=standalone --genrule_strategy=standalone" \ >>/etc/bazel.bazelrc # Install the most recent bazel release. -ENV BAZEL_VERSION 0.8.0 +ENV BAZEL_VERSION 0.11.0 WORKDIR / RUN mkdir /bazel && \ cd /bazel && \ @@ -79,7 +79,7 @@ RUN mkdir /bazel && \ # Download and build TensorFlow. WORKDIR /tensorflow -RUN git clone --branch=r1.5 --depth=1 https://github.com/tensorflow/tensorflow.git . +RUN git clone --branch=r1.6 --depth=1 https://github.com/tensorflow/tensorflow.git . # Configure the build for our CUDA configuration. ENV CI_BUILD_PYTHON python diff --git a/tensorflow/tools/docker/Dockerfile.gpu b/tensorflow/tools/docker/Dockerfile.gpu index b6682cd68163ec870ed815b45ac4fdd9233f88c6..625321e1235202f78a2d5e1a5b2d9d05e1e3f9ba 100644 --- a/tensorflow/tools/docker/Dockerfile.gpu +++ b/tensorflow/tools/docker/Dockerfile.gpu @@ -1,11 +1,18 @@ -FROM nvidia/cuda:9.0-cudnn7-runtime-ubuntu16.04 +FROM nvidia/cuda:9.0-base-ubuntu16.04 LABEL maintainer="Craig Citro " # Pick up some TF dependencies RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ + cuda-command-line-tools-9-0 \ + cuda-cublas-9-0 \ + cuda-cufft-9-0 \ + cuda-curand-9-0 \ + cuda-cusolver-9-0 \ + cuda-cusparse-9-0 \ curl \ + libcudnn7=7.0.5.15-1+cuda9.0 \ libfreetype6-dev \ libpng12-dev \ libzmq3-dev \ diff --git a/tensorflow/tools/docs/generate_1_0.py b/tensorflow/tools/docs/generate_1_0.py index cdc03fdcacf44f7be49e739962b63ba84cf94896..f4384e0ced77718c80d4d146a2d72072588a0541 100644 --- a/tensorflow/tools/docs/generate_1_0.py +++ b/tensorflow/tools/docs/generate_1_0.py @@ -53,7 +53,6 @@ if __name__ == '__main__': 'factorization', 'grid_rnn', 'labeled_tensor', - 'ndlstm', 'quantization', 'session_bundle', 'slim', diff --git a/tensorflow/tools/docs/generate_lib.py b/tensorflow/tools/docs/generate_lib.py index 003f972070cb05aa6f34a3748d47f019744de058..34dd419f15676babfa9a36c2c0960b01248b6f69 100644 --- a/tensorflow/tools/docs/generate_lib.py +++ b/tensorflow/tools/docs/generate_lib.py @@ -215,7 +215,6 @@ def _get_default_do_not_descend_map(): # Block contrib.keras to de-clutter the docs 'keras', 'labeled_tensor', - 'ndlstm', 'quantization', 'session_bundle', 'slim', diff --git a/tensorflow/tools/docs/parser.py b/tensorflow/tools/docs/parser.py index 3db164c2b5b78dbcb3c408ce89c067d33c2a2af4..e758229535e7b10994a39cbafb37e116fd2a465c 100644 --- a/tensorflow/tools/docs/parser.py +++ b/tensorflow/tools/docs/parser.py @@ -111,8 +111,8 @@ SYMBOL_REFERENCE_RE = re.compile( r""" # Start with a literal "@{". @\{ - # Group at least 1 symbol: not "}" or "\n". - ([^}\n]+) + # Group at least 1 symbol, not "}". + ([^}]+) # Followed by a closing "}" \} """, diff --git a/tensorflow/tools/docs/parser_test.py b/tensorflow/tools/docs/parser_test.py index 8a0e9af5216c881326449b3e85b94c0be331fa37..fca5436ca5fadd1fb5da07d7523bb51c871164b5 100644 --- a/tensorflow/tools/docs/parser_test.py +++ b/tensorflow/tools/docs/parser_test.py @@ -76,8 +76,9 @@ class ParserTest(googletest.TestCase): pass string = ( - 'A @{tf.reference}, another @{tf.reference}, a member ' - '@{tf.reference.foo}, and a @{tf.third$link `text` with `code` in it}.') + 'A @{tf.reference}, another @{tf.reference$with\nnewline}, a member ' + '@{tf.reference.foo}, and a @{tf.third$link `text` with `code` in ' + 'it}.') duplicate_of = {'tf.third': 'tf.fourth'} index = {'tf.reference': HasOneMember, 'tf.reference.foo': HasOneMember.foo, @@ -93,7 +94,7 @@ class ParserTest(googletest.TestCase): self.assertEqual('A ' 'tf.reference, ' 'another ' - 'tf.reference, ' + 'with\nnewline, ' 'a member ' 'tf.reference.foo, ' 'and a link ' diff --git a/tensorflow/tools/docs/pretty_docs.py b/tensorflow/tools/docs/pretty_docs.py index b5df633800ae5a3ce67cf03910d472b9908d6249..543b5fa6fefcd8e8dca99ad7eac7cca76781ccd3 100644 --- a/tensorflow/tools/docs/pretty_docs.py +++ b/tensorflow/tools/docs/pretty_docs.py @@ -162,7 +162,7 @@ def _build_class_page(page_info): parts.append(h3.format(**method_info.__dict__)) if method_info.signature is not None: - parts.append(_build_signature(method_info)) + parts.append(_build_signature(method_info, use_full_name=False)) parts.append(method_info.doc.docstring) parts.append(_build_function_details(method_info.doc.function_details)) @@ -259,14 +259,14 @@ def _build_module_page(page_info): return ''.join(parts) -def _build_signature(obj_info): +def _build_signature(obj_info, use_full_name=True): """Returns a md code block showing the function signature.""" # Special case tf.range, since it has an optional first argument if obj_info.full_name == 'tf.range': return ( '``` python\n' - "range(limit, delta=1, dtype=None, name='range')\n" - "range(start, limit, delta=1, dtype=None, name='range')\n" + "tf.range(limit, delta=1, dtype=None, name='range')\n" + "tf.range(start, limit, delta=1, dtype=None, name='range')\n" '```\n\n') parts = ['``` python'] @@ -281,7 +281,11 @@ def _build_signature(obj_info): sig = ',\n'.join(' %s' % sig_item for sig_item in obj_info.signature) sig = '\n'+sig+'\n' - parts.append(signature_template.format(name=obj_info.short_name, sig=sig)) + if use_full_name: + obj_name = obj_info.full_name + else: + obj_name = obj_info.short_name + parts.append(signature_template.format(name=obj_name, sig=sig)) parts.append('```\n\n') return '\n'.join(parts) diff --git a/tensorflow/tools/git/gen_git_source.py b/tensorflow/tools/git/gen_git_source.py index 3630dbd740e981971bdc9ff45b756b45095d437d..cbcdbf5b807a585865e2e3f19291e55388d55cb1 100755 --- a/tensorflow/tools/git/gen_git_source.py +++ b/tensorflow/tools/git/gen_git_source.py @@ -114,6 +114,13 @@ def configure(src_base_path, gen_path, debug=False): for target, src in link_map.items(): if src is None: open(os.path.join(gen_path, target), "w").write("") + elif not os.path.exists(src): + # Git repo is configured in a way we don't support such as having + # packed refs. Even though in a git repo, tf.__git_version__ will not + # be accurate. + # TODO(mikecase): Support grabbing git info when using packed refs. + open(os.path.join(gen_path, target), "w").write("") + spec["git"] = False else: try: # In python 3.5, symlink function exists even on Windows. But requires diff --git a/tensorflow/tools/graph_transforms/BUILD b/tensorflow/tools/graph_transforms/BUILD index b5465b7fb32856833fc2a12c8dfea58c2e8e79dd..b7d7fac315367824c6fdb379eaa6b0812d89b28b 100644 --- a/tensorflow/tools/graph_transforms/BUILD +++ b/tensorflow/tools/graph_transforms/BUILD @@ -91,7 +91,6 @@ cc_library( srcs = [ "add_default_attributes.cc", "backports.cc", - "fake_quantize_training.cc", "flatten_atrous.cc", "fold_batch_norms.cc", "fold_constants_lib.cc", @@ -99,22 +98,21 @@ cc_library( "freeze_requantization_ranges.cc", "fuse_convolutions.cc", "insert_logging.cc", - "remove_ema.cc", "obfuscate_names.cc", + "quantize_nodes.cc", + "quantize_weights.cc", "remove_attribute.cc", + "remove_control_dependencies.cc", "remove_device.cc", "remove_nodes.cc", "rename_attribute.cc", "rename_op.cc", + "round_weights.cc", "set_device.cc", "sort_by_execution_order.cc", "sparsify_gather.cc", "strip_unused_nodes.cc", - ] + if_not_windows([ - "quantize_nodes.cc", - "quantize_weights.cc", - "round_weights.cc", - ]), + ], hdrs = [ "fold_constants_lib.h", ], @@ -134,8 +132,8 @@ cc_library( "//tensorflow/core:tensorflow", "//tensorflow/contrib/rnn:gru_ops_op_lib", "//tensorflow/contrib/rnn:lstm_ops_op_lib", + "//tensorflow/core/kernels:quantization_utils", ] + if_not_windows([ - "//tensorflow/core/kernels:quantized_ops", "//tensorflow/core/kernels:remote_fused_graph_rewriter_transform", "//tensorflow/core/kernels/hexagon:hexagon_rewriter_transform", ]), @@ -148,7 +146,6 @@ tf_cc_test( srcs = [ "add_default_attributes_test.cc", "backports_test.cc", - "fake_quantize_training_test.cc", "flatten_atrous_test.cc", "fold_batch_norms_test.cc", "fold_constants_test.cc", @@ -161,7 +158,6 @@ tf_cc_test( "quantize_weights_test.cc", "remove_attribute_test.cc", "remove_device_test.cc", - "remove_ema_test.cc", "remove_nodes_test.cc", "rename_attribute_test.cc", "rename_op_test.cc", diff --git a/tensorflow/tools/graph_transforms/README.md b/tensorflow/tools/graph_transforms/README.md index 345d9eadb858cadebe03ecb3297aea52ba54bd37..67badb4869029b684cae05130d3c1e190dfb40d2 100644 --- a/tensorflow/tools/graph_transforms/README.md +++ b/tensorflow/tools/graph_transforms/README.md @@ -639,6 +639,13 @@ specified devices may not be available. In order to work with graphs like these, you can run this transform to wipe the slate clean and delete the device specifier from all ops. +### remove_control_dependencies + +Args: None \ +Prerequisites: None + +Removes all control dependencies from the graph. + ### remove_nodes Args: diff --git a/tensorflow/tools/graph_transforms/fake_quantize_training.cc b/tensorflow/tools/graph_transforms/fake_quantize_training.cc deleted file mode 100644 index 61aecc6e16d817d245421f18fa39c70aa45b2bef..0000000000000000000000000000000000000000 --- a/tensorflow/tools/graph_transforms/fake_quantize_training.cc +++ /dev/null @@ -1,51 +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. -==============================================================================*/ - -#define EIGEN_USE_THREADS - -#include "tensorflow/core/graph/quantize_training.h" -#include "tensorflow/tools/graph_transforms/transform_utils.h" - -namespace tensorflow { -namespace graph_transforms { - -// EXPERIMENTAL: This can change without warning. -// Rewrites the GraphDef for quantized training. -// Rewrites the forward pass to include the precision loss with quantization so -// the model can learn to deal with such loss and achieve better accuracy when -// it is quantized later for inference. -// Quantization range information is collected in FakeQuantizeWithMinMaxVars -// ops. -// -// TODO(suharshs): Provide instructions on converting the resulting graph for -// inference. -// TODO(suharshs): Implement this using the GTT rather than calling the old -// prototype function. -Status FakeQuantizeTraining(const GraphDef& input_graph_def, - const TransformFuncContext& context, - GraphDef* output_graph_def) { - // TODO(suharshs): Make num_bits a parameter. - const int32 num_bits = 8; - // TODO(suharshs): Make quantization op a parameter? - const string quant_op_type = "FakeQuantWithMinMaxVars"; - - return DoQuantizeTrainingOnGraphDef(input_graph_def, num_bits, quant_op_type, - output_graph_def); -} - -REGISTER_GRAPH_TRANSFORM("fake_quantize_training", FakeQuantizeTraining); - -} // namespace graph_transforms -} // namespace tensorflow diff --git a/tensorflow/tools/graph_transforms/fake_quantize_training_test.cc b/tensorflow/tools/graph_transforms/fake_quantize_training_test.cc deleted file mode 100644 index 5e4ab209e97808c3f42ecf73fb763ef9d7ab1cfe..0000000000000000000000000000000000000000 --- a/tensorflow/tools/graph_transforms/fake_quantize_training_test.cc +++ /dev/null @@ -1,63 +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/cc/ops/const_op.h" -#include "tensorflow/cc/ops/math_ops.h" -#include "tensorflow/core/framework/tensor_testutil.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/tools/graph_transforms/transform_utils.h" - -namespace tensorflow { -namespace graph_transforms { - -// Declare here, so we don't need a public header. -Status FakeQuantizeTraining(const GraphDef& input_graph_def, - const TransformFuncContext& context, - GraphDef* output_graph_def); - -class FakeQuantizeTrainingTest : public ::testing::Test {}; - -// For now, since the fake_quantize_training transform just calls the -// quantize_training rewrite from tensorflow/core/graph/quantize_training.h, -// we just test that the graph has been changed by the transform. -// TODO(suharshs): Once we implement the fake_quantize_training transform -// using the GTT, write proper tests of the transform here. -TEST_F(FakeQuantizeTrainingTest, TransformOccurred) { - auto root = tensorflow::Scope::DisabledShapeInferenceScope(); - using namespace ::tensorflow::ops; // NOLINT(build/namespaces) - - Tensor a_data(DT_FLOAT, TensorShape()); - test::FillIota(&a_data, 1.0f); - Output a_const = Const(root.WithOpName("a"), Input::Initializer(a_data)); - - Tensor b_data(DT_FLOAT, TensorShape()); - test::FillIota(&b_data, 1.0f); - Output b_const = Const(root.WithOpName("b"), Input::Initializer(b_data)); - - Output matmul = MatMul(root.WithOpName("matmul"), a_const, b_const); - GraphDef graph_def; - TF_ASSERT_OK(root.ToGraphDef(&graph_def)); - - GraphDef result; - TransformFuncContext context; - TF_ASSERT_OK(FakeQuantizeTraining(graph_def, context, &result)); - - // Test that the transformation resulted in a graph with more nodes. - EXPECT_GT(result.node_size(), graph_def.node_size()); -} - -} // namespace graph_transforms -} // namespace tensorflow diff --git a/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc b/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc index d89afe85c72883323cec3c14342fd60adebd024d..d86f65325be1c3f5151ab8d0a0c3c64afa3dcf0f 100644 --- a/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc +++ b/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc @@ -182,6 +182,36 @@ Status FuseBatchNormWithConv(const NodeMatch& match, return Status::OK(); } +Status FuseBatchNormWithBatchToSpace(const NodeMatch& match, + std::vector* new_nodes) { + // Calculate the scale and offset values to apply. + std::vector scale_values; + std::vector offset_values; + TF_RETURN_IF_ERROR( + GetScaleAndOffsetValues(match, &scale_values, &offset_values)); + + // Fuse conv weights, and set the final output node name as batch_norm_node. + const NodeDef& batch_norm_node = match.node; + const NodeMatch& batch_to_space_node_match = match.inputs[0]; + const NodeMatch& conv_node_match = batch_to_space_node_match.inputs[0]; + const NodeDef& batch_to_space_node = batch_to_space_node_match.node; + const NodeDef& conv_node = conv_node_match.node; + + string biasadd_name = conv_node.name() + "/biasadd"; + TF_RETURN_IF_ERROR( + FuseScaleOffsetToConvWeights(scale_values, offset_values, conv_node_match, + biasadd_name , new_nodes)); + + NodeDef new_batch_to_space_node = batch_to_space_node; + // reuse batch_norm node name + new_batch_to_space_node.set_name(batch_norm_node.name()); + new_batch_to_space_node.set_input(0, biasadd_name); + new_nodes->push_back(batch_to_space_node_match.inputs[1].node); + new_nodes->push_back(batch_to_space_node_match.inputs[2].node); + new_nodes->push_back(new_batch_to_space_node); + return Status::OK(); +} + Status FuseBatchNormWithConvConcat(const NodeMatch& match, std::vector* new_nodes) { // Calculate the scale and offset values to apply. @@ -284,6 +314,43 @@ Status FoldOldBatchNorms(const GraphDef& input_graph_def, current_graph_def = replaced_graph_def; } while (did_graph_change); + do { + did_graph_change = false; + GraphDef replaced_graph_def; + TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes( + current_graph_def, // clang-format off + {"BatchNormWithGlobalNormalization|FusedBatchNorm", // batch_norm_node + { + {"BatchToSpaceND", // batch_to_space_node + { + {"Conv2D", // conv_node + { + {"*"}, // input_node + {"Const"}, // weights_node + } + }, + {"Const"}, // block_shape + {"Const"}, // crops + } + }, + {"Const"}, // mean_node + {"Const"}, // variance_node + {"Const"}, // beta_node + {"Const"}, // gamma_node + } + }, // clang-format on + [&did_graph_change](const NodeMatch& match, + const std::set& input_nodes, + const std::set& output_nodes, + std::vector* new_nodes) { + TF_RETURN_IF_ERROR(FuseBatchNormWithBatchToSpace(match, new_nodes)); + did_graph_change = true; + return Status::OK(); + }, + {}, &replaced_graph_def)); + current_graph_def = replaced_graph_def; + } while (did_graph_change); + do { did_graph_change = false; GraphDef replaced_graph_def; diff --git a/tensorflow/tools/graph_transforms/fold_old_batch_norms_test.cc b/tensorflow/tools/graph_transforms/fold_old_batch_norms_test.cc index b30ba9ac8b92db68eb3374c51a7f31b69cd1e3cf..7651a03fe51012678d6d6fc495fd82e497aa512b 100644 --- a/tensorflow/tools/graph_transforms/fold_old_batch_norms_test.cc +++ b/tensorflow/tools/graph_transforms/fold_old_batch_norms_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/cc/ops/const_op.h" #include "tensorflow/cc/ops/image_ops.h" #include "tensorflow/cc/ops/nn_ops.h" +#include "tensorflow/cc/ops/array_ops.h" #include "tensorflow/cc/ops/sendrecv_ops.h" #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/core/framework/tensor_testutil.h" @@ -298,6 +299,96 @@ class FoldOldBatchNormsTest : public ::testing::Test { } }; +void TestFoldFusedBatchNormsWithBatchToSpace() { + auto root = tensorflow::Scope::NewRootScope(); + using namespace ::tensorflow::ops; // NOLINT(build/namespaces) + + Tensor input_data(DT_FLOAT, TensorShape({2, 1, 3, 2})); + test::FillValues( + &input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f, + -5.0f, -3.0f, -6.0f}); + Output input_op = + Const(root.WithOpName("input_op"), Input::Initializer(input_data)); + + Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2})); + test::FillValues(&weights_data, + {1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f}); + Output weights_op = + Const(root.WithOpName("weights_op"), Input::Initializer(weights_data)); + + Output conv_op = Conv2D(root.WithOpName("conv_op"), input_op, weights_op, + {1, 1, 1, 1}, "VALID"); + + Tensor block_shape_data(DT_INT32, TensorShape({2})); + test::FillValues(&block_shape_data, {1, 2}); + Output block_shape_op = + Const(root.WithOpName("block_shape_op"), Input::Initializer(block_shape_data)); + + Tensor crops_data(DT_INT32, TensorShape({2, 2})); + test::FillValues(&crops_data, {0, 0, 0, 1}); + Output crops_op = + Const(root.WithOpName("crops_op"), Input::Initializer(crops_data)); + + Output batch_to_space_op = BatchToSpaceND(root.WithOpName("batch_to_space_op"), + conv_op, block_shape_op, crops_data); + + Tensor mean_data(DT_FLOAT, TensorShape({2})); + test::FillValues(&mean_data, {10.0f, 20.0f}); + Output mean_op = + Const(root.WithOpName("mean_op"), Input::Initializer(mean_data)); + + Tensor variance_data(DT_FLOAT, TensorShape({2})); + test::FillValues(&variance_data, {0.25f, 0.5f}); + Output variance_op = Const(root.WithOpName("variance_op"), + Input::Initializer(variance_data)); + + Tensor beta_data(DT_FLOAT, TensorShape({2})); + test::FillValues(&beta_data, {0.1f, 0.6f}); + Output beta_op = + Const(root.WithOpName("beta_op"), Input::Initializer(beta_data)); + + Tensor gamma_data(DT_FLOAT, TensorShape({2})); + test::FillValues(&gamma_data, {1.0f, 2.0f}); + Output gamma_op = + Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data)); + + GraphDef original_graph_def; + TF_ASSERT_OK(root.ToGraphDef(&original_graph_def)); + + NodeDef batch_norm_node; + batch_norm_node.set_op("FusedBatchNorm"); + batch_norm_node.set_name("output"); + AddNodeInput("batch_to_space_op", &batch_norm_node); + AddNodeInput("gamma_op", &batch_norm_node); + AddNodeInput("beta_op", &batch_norm_node); + AddNodeInput("mean_op", &batch_norm_node); + AddNodeInput("variance_op", &batch_norm_node); + SetNodeAttr("T", DT_FLOAT, &batch_norm_node); + SetNodeAttr("epsilon", 0.00001f, &batch_norm_node); + SetNodeAttr("is_training", false, &batch_norm_node); + *(original_graph_def.mutable_node()->Add()) = batch_norm_node; + + std::unique_ptr original_session(NewSession(SessionOptions())); + TF_ASSERT_OK(original_session->Create(original_graph_def)); + std::vector original_outputs; + TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs)); + + GraphDef fused_graph_def; + TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}}, + &fused_graph_def)); + + std::unique_ptr fused_session(NewSession(SessionOptions())); + TF_ASSERT_OK(fused_session->Create(fused_graph_def)); + std::vector fused_outputs; + TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs)); + + test::ExpectTensorNear(original_outputs[0], fused_outputs[0], 1e-5); + + for (const NodeDef& node : fused_graph_def.node()) { + EXPECT_NE("FusedBatchNormWithBatchToSpace", node.op()); + } +} + TEST_F(FoldOldBatchNormsTest, TestFoldOldBatchNorms) { TestFoldOldBatchNorms(); } @@ -307,7 +398,7 @@ TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNorms) { } TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNormsWithConcat) { - // Test axis is not 3, so all weigths and offsets are fused to each of inputs + // Test axis is not 3, so all weights and offsets are fused to each of inputs // of conv2d. TestFoldFusedBatchNormsWithConcat(/*split=*/true); // Test axis = 3, BatchNorm weights and offsets will be split before fused @@ -315,5 +406,9 @@ TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNormsWithConcat) { TestFoldFusedBatchNormsWithConcat(/*split=*/false); } +TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNormsWithBatchToSpace) { + TestFoldFusedBatchNormsWithBatchToSpace(); +} + } // namespace graph_transforms } // namespace tensorflow diff --git a/tensorflow/tools/graph_transforms/remove_control_dependencies.cc b/tensorflow/tools/graph_transforms/remove_control_dependencies.cc new file mode 100644 index 0000000000000000000000000000000000000000..cba6b78fc5c43ca17f4f30930eb74efdb12940a1 --- /dev/null +++ b/tensorflow/tools/graph_transforms/remove_control_dependencies.cc @@ -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. +==============================================================================*/ +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/tools/graph_transforms/transform_utils.h" + +namespace tensorflow { +namespace graph_transforms { + +// Remove control depdencies in preparation for inference. +// In the tensorflow graph, control dependencies are represented as extra +// inputs which are referenced with "^tensor_name". +// See node_def.proto for more details. +Status RemoveControlDependencies(const GraphDef& input_graph_def, + const TransformFuncContext& context, + GraphDef* output_graph_def) { + output_graph_def->Clear(); + for (const NodeDef& node : input_graph_def.node()) { + NodeDef* new_node = output_graph_def->mutable_node()->Add(); + *new_node = node; + new_node->clear_input(); + for (const auto& input : node.input()) { + if (input[0] != '^') { + new_node->add_input(input); + } + } + } + return Status::OK(); +} + +REGISTER_GRAPH_TRANSFORM("remove_control_dependencies", RemoveControlDependencies); + +} // namespace graph_transforms +} // namespace tensorflow diff --git a/tensorflow/tools/graph_transforms/remove_ema.cc b/tensorflow/tools/graph_transforms/remove_ema.cc deleted file mode 100644 index 22e26267025c3fed4f44ffbc09d55d8d355cc448..0000000000000000000000000000000000000000 --- a/tensorflow/tools/graph_transforms/remove_ema.cc +++ /dev/null @@ -1,146 +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. -==============================================================================*/ - -#define EIGEN_USE_THREADS - -#include "tensorflow/tools/graph_transforms/transform_utils.h" - -namespace tensorflow { -namespace graph_transforms { - -// EXPERIMENTAL: This can change without warning. -// Given a graph that has gone through the FakeQuantizeTraining transform and -// has been frozen afterwards, RemoveEMA simplifies the FakeQuantize estimated -// moving average subgraphs to make it compatible with the QuantizeNodes -// transform. -Status RemoveEMA(const GraphDef& input_graph_def, - const TransformFuncContext& context, - GraphDef* output_graph_def) { - TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes( - input_graph_def, // clang-format off - {"FakeQuantWithMinMaxVars", - { - {"*"}, - {"Assign", - { - {"Const"}, - {"Merge", - { - {"Switch", - { - {"Min", - { - {"*"}, - {"Range", - { - {"*"}, - {"*"}, - {"*"}, - } - } - } - }, - {"IsVariableInitialized"} - } - }, - {"Sub", - { - {"Const"}, - {"Mul", - { - {"Sub"}, - {"Sub", - { - {"Const"}, - {"Const"} - } - } - } - } - } - } - } - } - } - }, - {"Assign", - { - {"Const"}, - {"Merge", - { - {"Switch", - { - {"Max"}, - {"IsVariableInitialized"} - } - }, - {"Sub", - { - {"Const"}, - {"Mul", - { - {"Sub"}, - {"Sub", - { - {"Const"}, - {"Const"} - } - } - } - } - } - } - } - } - } - }, - } - }, // clang-format on - [](const NodeMatch& match, const std::set& input_nodes, - const std::set& output_nodes, - std::vector* new_nodes) { - const NodeDef& fake_quant_node = match.node; - const NodeDef& input_node = match.inputs[0].node; - const NodeDef& min_var_node = match.inputs[1].inputs[0].node; - const NodeDef& max_var_node = match.inputs[2].inputs[0].node; - - // Make a new FakeQuantizeWithMinMaxVars operation that uses constants - // for its min/max arguments rather than an entire EMA subgraph. - NodeDef new_fake_quant_node; - new_fake_quant_node.set_op(fake_quant_node.op()); - new_fake_quant_node.set_name(fake_quant_node.name()); - AddNodeInput(input_node.name(), &new_fake_quant_node); - AddNodeInput(min_var_node.name(), &new_fake_quant_node); - AddNodeInput(max_var_node.name(), &new_fake_quant_node); - CopyNodeAttr(fake_quant_node, "narrow_range", "narrow_range", - &new_fake_quant_node); - CopyNodeAttr(fake_quant_node, "num_bits", "num_bits", - &new_fake_quant_node); - - new_nodes->push_back(new_fake_quant_node); - new_nodes->push_back(input_node); - new_nodes->push_back(min_var_node); - new_nodes->push_back(max_var_node); - - return Status::OK(); - }, - {}, output_graph_def)); - return Status::OK(); -} - -REGISTER_GRAPH_TRANSFORM("remove_ema", RemoveEMA); - -} // namespace graph_transforms -} // namespace tensorflow diff --git a/tensorflow/tools/graph_transforms/remove_ema_test.cc b/tensorflow/tools/graph_transforms/remove_ema_test.cc deleted file mode 100644 index 27db90e2729487f89324622f7a63aca1c5a58fe7..0000000000000000000000000000000000000000 --- a/tensorflow/tools/graph_transforms/remove_ema_test.cc +++ /dev/null @@ -1,121 +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/cc/ops/const_op.h" -#include "tensorflow/cc/ops/math_ops.h" -#include "tensorflow/core/framework/tensor_testutil.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/core/public/session.h" -#include "tensorflow/tools/graph_transforms/transform_utils.h" - -namespace tensorflow { -namespace graph_transforms { - -// Declare transformations here, so we don't need a public header. -Status FakeQuantizeTraining(const GraphDef& input_graph_def, - const TransformFuncContext& context, - GraphDef* output_graph_def); - -Status RemoveEMA(const GraphDef& input_graph_def, - const TransformFuncContext& context, - GraphDef* output_graph_def); - -Status QuantizeNodes(const GraphDef& input_graph_def, - const TransformFuncContext& context, - GraphDef* output_graph_def); - -class RemoveEMATest : public ::testing::Test {}; - -TEST_F(RemoveEMATest, FakeQuant_RemoveEMA_QuantizeTraining) { - // Build a small graph. - auto root = tensorflow::Scope::NewRootScope(); - using namespace ::tensorflow::ops; // NOLINT(build/namespaces) - - Tensor a_data(DT_FLOAT, TensorShape({1, 1})); - test::FillIota(&a_data, 1.0f); - Output a_const = Const(root.WithOpName("a"), Input::Initializer(a_data)); - - Tensor b_data(DT_FLOAT, TensorShape({1, 1})); - test::FillIota(&b_data, 1.0f); - Output b_const = Const(root.WithOpName("b"), Input::Initializer(b_data)); - - Output matmul = MatMul(root.WithOpName("matmul"), a_const, b_const); - GraphDef graph_def; - TF_ASSERT_OK(root.ToGraphDef(&graph_def)); - - // (1) FakeQuantize the graph. - GraphDef fake_quantized_graph_def; - TransformFuncContext context; - TF_ASSERT_OK( - FakeQuantizeTraining(graph_def, context, &fake_quantized_graph_def)); - - // Test that the transformation resulted in a graph with more nodes. - EXPECT_GT(fake_quantized_graph_def.node_size(), graph_def.node_size()); - - // (2) Run the graph to initialize the newly added variables. - std::unique_ptr session(NewSession(SessionOptions())); - TF_ASSERT_OK(session->Create(fake_quantized_graph_def)); - std::vector outputs; - TF_ASSERT_OK(session->Run({}, {"matmul"}, {}, &outputs)); - - // (3) Freeze the graph. Create a "frozen graph" that matches what we would - // expect if we actually froze the above graph. - // TODO(suharshs): Use a c++ freeze graph alternative, when one is available. - GraphDef frozen_graph_def; - for (const NodeDef& node : fake_quantized_graph_def.node()) { - if (node.op() == "Variable" || node.op() == "VariableV2") { - NodeDef const_node; - const_node.set_op("Const"); - const_node.set_name(node.name()); - SetNodeAttr("dtype", DT_FLOAT, &const_node); - Tensor tensor(DT_FLOAT, {}); - tensor.flat()(0) = 1.0f; - SetNodeTensorAttr("value", tensor, &const_node); - *(frozen_graph_def.mutable_node()->Add()) = const_node; - } else { - *(frozen_graph_def.mutable_node()->Add()) = node; - } - } - - // Test that freezing the graph resulted in a graph with the same number of - // nodes. - EXPECT_EQ(frozen_graph_def.node_size(), fake_quantized_graph_def.node_size()); - - // (4) RemoveEMA on the graph to make it compatible with QuantizeNodes. - GraphDef removed_ema_graph_def; - TF_ASSERT_OK(RemoveEMA(frozen_graph_def, context, &removed_ema_graph_def)); - - // Test that the transformation resulted in a graph with less nodes. - EXPECT_LT(removed_ema_graph_def.node_size(), frozen_graph_def.node_size()); - - // (5) QuantizeNodes and inspect the final graph. - // TODO(suharshs): Add a more thorough inspection of the structure of - // the output graph. - GraphDef quantized_graph_def; - TF_ASSERT_OK( - QuantizeNodes(removed_ema_graph_def, context, &quantized_graph_def)); - - // Test that the transformation resulted in a graph with more nodes. - EXPECT_GT(quantized_graph_def.node_size(), removed_ema_graph_def.node_size()); - - // Make sure that the FakeQuantizeWithMinMaxVars op has been removed. - for (const NodeDef& node : quantized_graph_def.node()) { - EXPECT_NE(node.op(), "FakeQuantWithMinMaxVars"); - } -} - -} // namespace graph_transforms -} // namespace tensorflow diff --git a/tensorflow/tools/graph_transforms/remove_nodes.cc b/tensorflow/tools/graph_transforms/remove_nodes.cc index 119b44d6a4a4d066b734ae8a0e655c771087d0db..05f036a86a09b2a6a94e9c1a1220803eabc64da5 100644 --- a/tensorflow/tools/graph_transforms/remove_nodes.cc +++ b/tensorflow/tools/graph_transforms/remove_nodes.cc @@ -81,7 +81,17 @@ Status RemoveNodes(const GraphDef& input_graph_def, return Status::OK(); } const NodeDef& input_node = match.inputs[0].node; - inputs_to_rename[replace_node.name()] = input_node.name(); + string target_name = input_node.name(); + for (const string& input : replace_node.input()) { + if (!input.compare(0, target_name.size(), target_name)) { + if (input.size() == target_name.size() || + input[target_name.size()] == ':') { + target_name = input; + break; + } + } + } + inputs_to_rename[replace_node.name()] = target_name; inputs_to_rename["^" + replace_node.name()] = "^" + input_node.name(); new_nodes->push_back(input_node); diff --git a/tensorflow/tools/graph_transforms/sparsify_gather.cc b/tensorflow/tools/graph_transforms/sparsify_gather.cc index 593c654f9fbbabe7e89c1ff2a43e56d30e8919d6..701e350fc39d083665f5420e6b73510c182e12ce 100644 --- a/tensorflow/tools/graph_transforms/sparsify_gather.cc +++ b/tensorflow/tools/graph_transforms/sparsify_gather.cc @@ -86,8 +86,17 @@ void CreateConstNode(const Tensor& tensor, const string& name, SetNodeTensorAttr("value", tensor, node_def); } +string GetMonolithicTensorKey(const string& tensor_slice_name) { + std::vector names = Split(tensor_slice_name, "/"); + if (StringPiece(names[names.size() - 1]).starts_with("part_")) { + CHECK_GE(names.size(), 2); + names.pop_back(); + } + return Join(names, "/"); +} + Status ObtainTensorSlice(const GraphDef& input_graph_def, - const string& tensor_name, + const string& target_name, string* shape_slice_string) { string restore_node_name; for (const auto& node : input_graph_def.node()) { @@ -95,39 +104,53 @@ Status ObtainTensorSlice(const GraphDef& input_graph_def, if (node_name_parts.size() == 2 && StringPiece(node_name_parts[0]).starts_with("save") && StringPiece(node_name_parts[1]).starts_with("Assign") && - node.input(0) == tensor_name) { + node.input(0) == target_name) { restore_node_name = node.input(1); break; } } + + std::vector restore_node_parts = Split(restore_node_name, ":"); + CHECK_LE(restore_node_parts.size(), 2); + string tensor_names_node; string shape_and_slices_node; for (const auto& node : input_graph_def.node()) { - if ((node.name() == restore_node_name) && (node.op() == "RestoreV2")) { + if ((node.name() == restore_node_parts[0]) && (node.op() == "RestoreV2")) { + tensor_names_node = node.input(1); shape_and_slices_node = node.input(2); break; } } + + int offset = -1; + for (const auto& node : input_graph_def.node()) { + if (node.name() == tensor_names_node) { + Tensor tensor_names_tensor; + TF_RETURN_IF_ERROR(GetNodeAttr(node, "value", &tensor_names_tensor)); + const auto& tensor_names_value = tensor_names_tensor.flat(); + for (int i = 0; i < tensor_names_value.size(); i++) { + if (tensor_names_value(i) == GetMonolithicTensorKey(target_name)) { + offset = i; + break; + } + } + } + } + if (offset == -1) { + return errors::Internal("Unable to find RestoreV2 entry for variable: ", + target_name); + } for (const auto& node : input_graph_def.node()) { if (node.name() == shape_and_slices_node) { Tensor shape_and_slices_tensor; TF_RETURN_IF_ERROR(GetNodeAttr(node, "value", &shape_and_slices_tensor)); const auto& shape_and_slices_value = shape_and_slices_tensor.flat(); - *shape_slice_string = shape_and_slices_value(0); + *shape_slice_string = shape_and_slices_value(offset); return Status::OK(); } } - return errors::Internal("Unable to find slice for variable: ", tensor_name); -} - -string GetMonolithicTensorKey(const string& tensor_slice_name) { - std::vector names = Split(tensor_slice_name, "/"); - CHECK_GE(names.size(), 2); - CHECK(StringPiece(names[names.size() - 1]).starts_with("part_")); - - // Remove the "part_x" suffix - names.pop_back(); - return Join(names, "/"); + return errors::Internal("Unable to find slice for variable: ", target_name); } Status ReadTensorFromCheckpoint( @@ -181,6 +204,22 @@ Status ObtainVariableInfo( return Status::OK(); } +Status RemoveInputAtIndex(NodeDef* n, int index) { + for (int i = index; i < n->input_size() - 1; i++) { + n->mutable_input()->SwapElements(i, i + 1); + } + n->mutable_input()->RemoveLast(); + return Status::OK(); +} + +Status RemoveNodeAtIndex(GraphDef* g, int index) { + for (int i = index; i < g->node_size() - 1; i++) { + g->mutable_node()->SwapElements(i, i + 1); + } + g->mutable_node()->RemoveLast(); + return Status::OK(); +} + Status SparsifyGatherInternal( const GraphDef& input_graph_def, const std::unique_ptr >& @@ -301,13 +340,13 @@ Status SparsifyGatherInternal( TF_RETURN_IF_ERROR(ReadTensorFromCheckpoint( weights_node.name(), ckpt_reader, (*shapes_and_slices)[weights_node.name()], &weight)); - // Add both both weight and identity node names. - removed_node_names.push_back(weights_node.name()); - removed_node_names.push_back(match.inputs[0].node.name()); - for (auto input_node : match.inputs[0].node.input()) { - auto parsed_input = StringReplace(input_node, "^", "", true); - refs[parsed_input]--; - } + } + // Add both both weight and identity node names. + removed_node_names.push_back(weights_node.name()); + removed_node_names.push_back(match.inputs[0].node.name()); + for (auto input_node : match.inputs[0].node.input()) { + auto parsed_input = StringReplace(input_node, "^", "", true); + refs[parsed_input]--; } Tensor indices_tensor; Tensor values_tensor; @@ -462,32 +501,53 @@ Status SparsifyGatherInternal( removed_node_names.push_back(parsed_input); } } - replaced_graph_def.mutable_node()->SwapElements( - i, replaced_graph_def.node_size() - 1); - replaced_graph_def.mutable_node()->RemoveLast(); + TF_RETURN_IF_ERROR(RemoveNodeAtIndex(&replaced_graph_def, i)); continue; } int j = 0; + bool deleted_inputs = false; while (j < replaced_graph_def.node(i).input_size()) { if (replaced_graph_def.node(i).input(j) == name || replaced_graph_def.node(i).input(j) == ("^" + name)) { - replaced_graph_def.mutable_node(i)->mutable_input()->SwapElements( - j, replaced_graph_def.node(i).input_size() - 1); - replaced_graph_def.mutable_node(i)->mutable_input()->RemoveLast(); + TF_RETURN_IF_ERROR( + RemoveInputAtIndex(replaced_graph_def.mutable_node(i), j)); + deleted_inputs = true; continue; } j++; } - if (!replaced_graph_def.node(i).input_size()) { - if ((refs.find(replaced_graph_def.node(i).name()) != refs.end()) && - (refs[replaced_graph_def.node(i).name()] == 0)) { + if (deleted_inputs) { + if (replaced_graph_def.node(i).op() == "ConcatV2") { + if (replaced_graph_def.node(i).input_size() > 2) { + SetNodeAttr("N", replaced_graph_def.node(i).input_size() - 1, + replaced_graph_def.mutable_node(i)); + } else if (replaced_graph_def.node(i).input_size() == 2) { + if (refs[replaced_graph_def.node(i).input(1)] != 1) { + return errors::Internal( + "Expect axis tensor of ConcatV2 node to only be referenced " + "once."); + } + refs[replaced_graph_def.node(i).input(1)] -= 1; + removed_node_names.push_back(replaced_graph_def.node(i).input(1)); + replaced_graph_def.mutable_node(i)->mutable_input()->RemoveLast(); + replaced_graph_def.mutable_node(i)->mutable_attr()->erase("N"); + replaced_graph_def.mutable_node(i)->set_op("Identity"); + } else { + return errors::Internal( + "ConcatV2 should have at least two elements"); + } + } + if ((replaced_graph_def.node(i).op() == "Assign" || + replaced_graph_def.node(i).op() == "Reshape" || + replaced_graph_def.node(i).op() == "Equal" || + replaced_graph_def.node(i).op() == "Mean" || + replaced_graph_def.node(i).op() == "ScalarSummary") && + replaced_graph_def.node(i).input_size() == 1) { + removed_node_names.push_back(replaced_graph_def.node(i).name()); + } + if (!replaced_graph_def.node(i).input_size()) { removed_node_names.push_back(replaced_graph_def.node(i).name()); } - } - - if (replaced_graph_def.node(i).op() == "Assign" && - replaced_graph_def.node(i).input_size() == 1) { - removed_node_names.push_back(replaced_graph_def.node(i).name()); } i++; } @@ -528,17 +588,22 @@ Status SparsifyGather(const GraphDef& input_graph_def, }; // clang-format on + GraphDef cleaned_input_graph_def; + RemoveAttributes(input_graph_def, {"_output_shapes"}, + &cleaned_input_graph_def); + GraphDef temp_output; std::unique_ptr ckpt_reader; TF_RETURN_IF_ERROR(InitializeCheckpointReader(context, &ckpt_reader)); std::unique_ptr > shapes_and_slices; - TF_RETURN_IF_ERROR(ObtainVariableInfo(input_graph_def, &shapes_and_slices)); + TF_RETURN_IF_ERROR( + ObtainVariableInfo(cleaned_input_graph_def, &shapes_and_slices)); - TF_RETURN_IF_ERROR(SparsifyGatherInternal(input_graph_def, shapes_and_slices, - context, gather_pattern, - ckpt_reader, &temp_output)); + TF_RETURN_IF_ERROR(SparsifyGatherInternal( + cleaned_input_graph_def, shapes_and_slices, context, gather_pattern, + ckpt_reader, &temp_output)); TF_RETURN_IF_ERROR(SparsifyGatherInternal(temp_output, shapes_and_slices, context, gather_v2_pattern, diff --git a/tensorflow/tools/graph_transforms/sparsify_gather_test.cc b/tensorflow/tools/graph_transforms/sparsify_gather_test.cc index 6627df1331a6eaf49857c3ecba4a4d55859cad7c..d41321c9a6df755eed099ec453f162e2132cfb57 100644 --- a/tensorflow/tools/graph_transforms/sparsify_gather_test.cc +++ b/tensorflow/tools/graph_transforms/sparsify_gather_test.cc @@ -71,7 +71,7 @@ class SparsifyGatherTest : public ::testing::Test { } void TestSinglePartition(bool gather_v2, bool include_shared_init, - bool test_variable, + bool test_variable, bool test_kept_concat, const string& shared_init_name = "group_deps") { GraphDef graph_def; @@ -106,11 +106,15 @@ class SparsifyGatherTest : public ::testing::Test { NodeDef* save_const_node = CreateNode("save/Const", "Const", {}, &graph_def); + Tensor tensor_names_values(DT_STRING, TensorShape({1})); + test::FillValues(&tensor_names_values, {"w"}); NodeDef* tensor_names_node = CreateNode("save/RestoreV2/tensor_names", "Const", {}, &graph_def); + SetNodeTensorAttr("value", tensor_names_values, + tensor_names_node); + NodeDef* tensor_shapes_slices_node = CreateNode( "save/RestoreV2/shape_and_slices", "Const", {}, &graph_def); - Tensor shapes_slices_val(DT_STRING, TensorShape({1})); shapes_slices_val.flat()(0) = "4 1 0,4:0,1"; SetNodeTensorAttr("value", shapes_slices_val, @@ -139,6 +143,26 @@ class SparsifyGatherTest : public ::testing::Test { } } + NodeDef* concat_axis_node = + CreateNode("linear/concat/axis", "Const", {}, &graph_def); + NodeDef* concat_input_node = + CreateNode("concat/input/node", "Const", {}, &graph_def); + NodeDef* concat_node = nullptr; + if (!test_kept_concat) { + concat_node = CreateNode( + "concat/node", "ConcatV2", + {identity_node, concat_input_node, concat_axis_node}, &graph_def); + SetNodeAttr("N", 2, concat_node); + } else { + NodeDef* concat_input_node_2 = + CreateNode("concat/input/node_2", "Const", {}, &graph_def); + concat_node = CreateNode("concat/node", "ConcatV2", + {identity_node, concat_input_node, + concat_input_node_2, concat_axis_node}, + &graph_def); + SetNodeAttr("N", 3, concat_node); + } + // Run the op. GraphDef result; TransformFuncContext context; @@ -166,6 +190,23 @@ class SparsifyGatherTest : public ::testing::Test { EXPECT_EQ(1, node_lookup.count("ids")); EXPECT_EQ("Const", node_lookup.at("ids")->op()); + EXPECT_EQ(1, node_lookup.count("concat/node")); + + if (!test_kept_concat) { + EXPECT_EQ(0, node_lookup.count("linear/concat/axis")); + EXPECT_EQ("Identity", node_lookup.at("concat/node")->op()); + EXPECT_EQ(1, node_lookup.at("concat/node")->input_size()); + EXPECT_EQ("concat/input/node", node_lookup.at("concat/node")->input(0)); + } else { + EXPECT_EQ(1, node_lookup.count("linear/concat/axis")); + EXPECT_EQ("ConcatV2", node_lookup.at("concat/node")->op()); + EXPECT_EQ(3, node_lookup.at("concat/node")->input_size()); + EXPECT_EQ("concat/input/node", node_lookup.at("concat/node")->input(0)); + EXPECT_EQ("concat/input/node_2", node_lookup.at("concat/node")->input(1)); + EXPECT_EQ("linear/concat/axis", node_lookup.at("concat/node")->input(2)); + EXPECT_EQ(2, node_lookup.at("concat/node")->attr().at("N").i()); + } + EXPECT_EQ(1, node_lookup.count("w/part_1/indices")); EXPECT_EQ("Const", node_lookup.at("w/part_1/indices")->op()); Tensor expected_indices_tensor(DT_INT64, TensorShape({3})); @@ -273,6 +314,29 @@ class SparsifyGatherTest : public ::testing::Test { SetNodeTensorAttr("value", weights, w_node1); SetNodeTensorAttr("value", weights, w_node2); } else { + NodeDef* save_const_node = + CreateNode("save/Const", "Const", {}, &graph_def); + + NodeDef* tensor_names_node = + CreateNode("save/RestoreV2/tensor_names", "Const", {}, &graph_def); + Tensor tensor_names_values(DT_STRING, TensorShape({2})); + test::FillValues(&tensor_names_values, {"w1", "w2"}); + SetNodeTensorAttr("value", tensor_names_values, + tensor_names_node); + + NodeDef* tensor_shapes_slices_node = CreateNode( + "save/RestoreV2/shape_and_slices", "Const", {}, &graph_def); + Tensor shapes_slices_val(DT_STRING, TensorShape({2})); + shapes_slices_val.flat()(0) = "4 1 0,4:0,1"; + shapes_slices_val.flat()(1) = "4 1 0,4:0,1"; + SetNodeTensorAttr("value", shapes_slices_val, + tensor_shapes_slices_node); + + NodeDef* restore_node = CreateNode( + "save/RestoreV2", "RestoreV2", + {save_const_node, tensor_names_node, tensor_shapes_slices_node}, + &graph_def); + w_node1 = CreateNode("w1/part_1", "VariableV2", {}, &graph_def); zeros_shape1 = CreateNode("w1/part_1/Initializer/zeros/shape_as_tensor", @@ -284,23 +348,7 @@ class SparsifyGatherTest : public ::testing::Test { assign_node1 = CreateNode("w1/part_1/Assign", "Assign", {w_node1, zeros_node1}, &graph_def); - NodeDef* save_const_node = - CreateNode("save/Const", "Const", {}, &graph_def); - NodeDef* tensor_names_node1 = - CreateNode("save/RestoreV2/tensor_names", "Const", {}, &graph_def); - NodeDef* tensor_shapes_slices_node1 = CreateNode( - "save/RestoreV2/shape_and_slices", "Const", {}, &graph_def); - - Tensor shapes_slices_val1(DT_STRING, TensorShape({1})); - shapes_slices_val1.flat()(0) = "4 1 0,4:0,1"; - SetNodeTensorAttr("value", shapes_slices_val1, - tensor_shapes_slices_node1); - - NodeDef* restore_node1 = CreateNode( - "save/RestoreV2", "RestoreV2", - {save_const_node, tensor_names_node1, tensor_shapes_slices_node1}, - &graph_def); - CreateNode("save/Assign", "Assign", {w_node1, restore_node1}, &graph_def); + CreateNode("save/Assign", "Assign", {w_node1, restore_node}, &graph_def); w_node2 = CreateNode("w2/part_1", "VariableV2", {}, &graph_def); zeros_shape2 = CreateNode("w2/part_1/Initializer/zeros/shape_as_tensor", @@ -312,21 +360,7 @@ class SparsifyGatherTest : public ::testing::Test { assign_node2 = CreateNode("w2/part_1/Assign", "Assign", {w_node2, zeros_node2}, &graph_def); - NodeDef* tensor_names_node2 = - CreateNode("save/RestoreV2_1/tensor_names", "Const", {}, &graph_def); - NodeDef* tensor_shapes_slices_node2 = CreateNode( - "save/RestoreV2_1/shape_and_slices", "Const", {}, &graph_def); - - Tensor shapes_slices_val2(DT_STRING, TensorShape({1})); - shapes_slices_val2.flat()(0) = "4 1 0,4:0,1"; - SetNodeTensorAttr("value", shapes_slices_val2, - tensor_shapes_slices_node2); - - NodeDef* restore_node2 = CreateNode( - "save/RestoreV2_1", "RestoreV2", - {save_const_node, tensor_names_node2, tensor_shapes_slices_node2}, - &graph_def); - CreateNode("save/Assign_1", "Assign", {w_node2, restore_node2}, + CreateNode("save/Assign_1", "Assign", {w_node2, restore_node}, &graph_def); BundleWriter writer(Env::Default(), checkpoint_path); @@ -344,6 +378,13 @@ class SparsifyGatherTest : public ::testing::Test { MakeGather("gather1", gather_v2, identity_node1, input_node, &graph_def); MakeGather("gather2", gather_v2, identity_node2, input_node, &graph_def); + NodeDef* concat_axis_node = + CreateNode("linear/concat/axis", "Const", {}, &graph_def); + NodeDef* concat_node = CreateNode( + "concat/node", "ConcatV2", + {identity_node1, identity_node2, concat_axis_node}, &graph_def); + SetNodeAttr("N", 2, concat_node); + // Shared init node if (include_shared_init) { if (!test_variable) { @@ -515,6 +556,9 @@ class SparsifyGatherTest : public ::testing::Test { node_lookup.at("gather2/LookupTableFind")->input(2)); EXPECT_EQ("gather2/LookupTableFind", node_lookup.at("gather2")->input(0)); + EXPECT_EQ(0, node_lookup.count("linear/concat/axis")); + EXPECT_EQ(0, node_lookup.count("concat/node")); + // Check control deps. EXPECT_EQ(2, node_lookup.at(shared_init_name)->input_size()); EXPECT_NE(std::find(node_lookup.at(shared_init_name)->input().begin(), @@ -550,18 +594,31 @@ class SparsifyGatherTest : public ::testing::Test { }; TEST_F(SparsifyGatherTest, TestSinglePartition) { - TestSinglePartition(false, false, false); - TestSinglePartition(false, true, false); - TestSinglePartition(true, false, false); - TestSinglePartition(true, true, false); - TestSinglePartition(false, false, true); - TestSinglePartition(false, true, true); - TestSinglePartition(true, false, true); - TestSinglePartition(true, true, true); - TestSinglePartition(false, true, false, "shared_inits"); - TestSinglePartition(true, true, false, "shared_inits"); - TestSinglePartition(false, true, true, "shared_inits"); - TestSinglePartition(true, true, true, "shared_inits"); + TestSinglePartition(false, false, false, false); + TestSinglePartition(false, true, false, false); + TestSinglePartition(true, false, false, false); + TestSinglePartition(true, true, false, false); + TestSinglePartition(false, false, true, false); + TestSinglePartition(false, true, true, false); + TestSinglePartition(true, false, true, false); + TestSinglePartition(true, true, true, false); + TestSinglePartition(false, true, false, false, "shared_inits"); + TestSinglePartition(true, true, false, false, "shared_inits"); + TestSinglePartition(false, true, true, false, "shared_inits"); + TestSinglePartition(true, true, true, false, "shared_inits"); + + TestSinglePartition(false, false, false, true); + TestSinglePartition(false, true, false, true); + TestSinglePartition(true, false, false, true); + TestSinglePartition(true, true, false, true); + TestSinglePartition(false, false, true, true); + TestSinglePartition(false, true, true, true); + TestSinglePartition(true, false, true, true); + TestSinglePartition(true, true, true, true); + TestSinglePartition(false, true, false, true, "shared_inits"); + TestSinglePartition(true, true, false, true, "shared_inits"); + TestSinglePartition(false, true, true, true, "shared_inits"); + TestSinglePartition(true, true, true, true, "shared_inits"); } TEST_F(SparsifyGatherTest, TestMultiPartition) { diff --git a/tensorflow/tools/lib_package/BUILD b/tensorflow/tools/lib_package/BUILD index dbc81599de8539ce58933f9d40bf99fcae8f8e67..3fbdb5cacd1fd0039deaae5ac330b6c2ca006a68 100644 --- a/tensorflow/tools/lib_package/BUILD +++ b/tensorflow/tools/lib_package/BUILD @@ -27,6 +27,7 @@ pkg_tar( ":cheaders", ":clib", ":clicenses", + ":eager_cheaders", ], ) @@ -57,7 +58,6 @@ pkg_tar( name = "cheaders", files = [ "//tensorflow/c:headers", - "//tensorflow/c/eager:headers", ], package_dir = "include/tensorflow/c", # Mark as "manual" till @@ -68,6 +68,20 @@ pkg_tar( tags = ["manual"], ) +pkg_tar( + name = "eager_cheaders", + files = [ + "//tensorflow/c/eager:headers", + ], + package_dir = "include/tensorflow/c/eager", + # Mark as "manual" till + # https://github.com/bazelbuild/bazel/issues/2352 + # and https://github.com/bazelbuild/bazel/issues/1580 + # are resolved, otherwise these rules break when built + # with Python 3. + tags = ["manual"], +) + pkg_tar( name = "clib", files = ["//tensorflow:libtensorflow.so"], @@ -99,6 +113,7 @@ genrule( "//third_party/hadoop:LICENSE.txt", "//third_party/eigen3:LICENSE", "//third_party/fft2d:LICENSE", + "@aws//:LICENSE", "@boringssl//:LICENSE", "@com_googlesource_code_re2//:LICENSE", "@cub_archive//:LICENSE.TXT", @@ -112,8 +127,10 @@ genrule( "@jemalloc//:COPYING", "@jpeg//:LICENSE.md", "@libxsmm_archive//:LICENSE", + "@llvm//:LICENSE.TXT", "@lmdb//:LICENSE", "@local_config_sycl//sycl:LICENSE.text", + "@nasm//:LICENSE", "@nsync//:LICENSE", "@png_archive//:LICENSE", "@protobuf_archive//:LICENSE", @@ -134,6 +151,7 @@ genrule( "//third_party/hadoop:LICENSE.txt", "//third_party/eigen3:LICENSE", "//third_party/fft2d:LICENSE", + "@aws//:LICENSE", "@boringssl//:LICENSE", "@com_googlesource_code_re2//:LICENSE", "@cub_archive//:LICENSE.TXT", @@ -147,8 +165,10 @@ genrule( "@jemalloc//:COPYING", "@jpeg//:LICENSE.md", "@libxsmm_archive//:LICENSE", + "@llvm//:LICENSE.TXT", "@lmdb//:LICENSE", "@local_config_sycl//sycl:LICENSE.text", + "@nasm//:LICENSE", "@nsync//:LICENSE", "@png_archive//:LICENSE", "@protobuf_archive//:LICENSE", diff --git a/tensorflow/tools/pip_package/BUILD b/tensorflow/tools/pip_package/BUILD index 598080ed2753b862056ebcc76c4c572ae45b46e6..2607b9d7049828db95e7cbdddad30541a26942b0 100644 --- a/tensorflow/tools/pip_package/BUILD +++ b/tensorflow/tools/pip_package/BUILD @@ -11,6 +11,7 @@ load( ) load("//third_party/mkl:build_defs.bzl", "if_mkl") load("//tensorflow:tensorflow.bzl", "if_cuda") +load("@local_config_tensorrt//:build_defs.bzl", "if_tensorrt") load("//tensorflow/core:platform/default/build_config_root.bzl", "tf_additional_license_deps") # This returns a list of headers of all public header libraries (e.g., @@ -47,37 +48,65 @@ py_binary( deps = ["//tensorflow:tensorflow_py"], ) +COMMON_PIP_DEPS = [ + ":licenses", + "MANIFEST.in", + "README", + "setup.py", + ":included_headers", + "//tensorflow:tensorflow_py", + "//tensorflow/contrib/boosted_trees:boosted_trees_pip", + "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip", + "//tensorflow/contrib/data/python/kernel_tests:dataset_serialization_test", + "//tensorflow/contrib/data/python/ops:contrib_op_loader", + "//tensorflow/contrib/eager/python/examples:examples_pip", + "//tensorflow/contrib/eager/python:checkpointable_utils", + "//tensorflow/contrib/eager/python:evaluator", + "//tensorflow/contrib/gan:gan", + "//tensorflow/contrib/graph_editor:graph_editor_pip", + "//tensorflow/contrib/keras:keras", + "//tensorflow/contrib/labeled_tensor:labeled_tensor_pip", + "//tensorflow/contrib/nn:nn_py", + "//tensorflow/contrib/predictor:predictor_pip", + "//tensorflow/contrib/py2tf:py2tf", + "//tensorflow/contrib/py2tf/converters:converters", + "//tensorflow/contrib/py2tf/converters:test_lib", + "//tensorflow/contrib/py2tf/impl:impl", + "//tensorflow/contrib/py2tf/pyct:pyct", + "//tensorflow/contrib/py2tf/pyct/static_analysis:static_analysis", + "//tensorflow/contrib/receptive_field:receptive_field_pip", + "//tensorflow/contrib/session_bundle:session_bundle_pip", + "//tensorflow/contrib/signal:signal_py", + "//tensorflow/contrib/signal:test_util", + "//tensorflow/contrib/slim:slim", + "//tensorflow/contrib/slim/python/slim/data:data_pip", + "//tensorflow/contrib/slim/python/slim/nets:nets_pip", + "//tensorflow/contrib/specs:specs", + "//tensorflow/contrib/summary:summary_test_util", + "//tensorflow/contrib/tensor_forest:init_py", + "//tensorflow/contrib/tensor_forest/hybrid:hybrid_pip", + "//tensorflow/contrib/timeseries:timeseries_pip", + "//tensorflow/contrib/tpu", + "//tensorflow/examples/tutorials/mnist:package", + "//tensorflow/python:distributed_framework_test_lib", + "//tensorflow/python:meta_graph_testdata", + "//tensorflow/python:spectral_ops_test_util", + "//tensorflow/python:util_example_parser_configuration", + "//tensorflow/python/debug:debug_pip", + "//tensorflow/python/eager:eager_pip", + "//tensorflow/python/saved_model:saved_model", + "//tensorflow/python/tools:tools_pip", + "//tensorflow/python:test_ops", + "//tensorflow/tools/dist_test/server:grpc_tensorflow_server", +] + # On Windows, python binary is a zip file of runfiles tree. # Add everything to its data dependency for generating a runfiles tree # for building the pip package on Windows. py_binary( name = "simple_console_for_windows", srcs = ["simple_console_for_windows.py"], - data = [ - "MANIFEST.in", - "README", - "setup.py", - ":included_headers", - "//tensorflow/contrib/nn:nn_py", - "//tensorflow/contrib/session_bundle:session_bundle_pip", - "//tensorflow/contrib/signal:signal_py", - "//tensorflow/contrib/slim/python/slim/data:data_pip", - "//tensorflow/python:util_example_parser_configuration", - "//tensorflow/python/debug:debug_pip", - "//tensorflow/python/saved_model", - "//tensorflow/python:spectral_ops_test_util", - "//tensorflow/python/tools:tools_pip", - "//tensorflow/python/eager:eager_pip", - "//tensorflow/contrib/summary:summary_test_util", - # These targets don't build on Windows yet. Exclude them for now. - # "//tensorflow/contrib/ndlstm", - # "//tensorflow/contrib/slim", - # "//tensorflow/contrib/slim/python/slim/nets:nets_pip", - # "//tensorflow/contrib/specs", - # "//tensorflow/contrib/tensor_forest:init_py", - # "//tensorflow/contrib/tensor_forest/hybrid:hybrid_pip", - # "//tensorflow/examples/tutorials/mnist:package", - ], + data = COMMON_PIP_DEPS, srcs_version = "PY2AND3", deps = ["//tensorflow:tensorflow_py"], ) @@ -88,28 +117,40 @@ filegroup( "//third_party/eigen3:LICENSE", "//third_party/fft2d:LICENSE", "//third_party/hadoop:LICENSE.txt", + "@absl_py//absl/flags:LICENSE", + "@arm_neon_2_x86_sse//:LICENSE", + "@astor_archive//:LICENSE", + "@aws//:LICENSE", "@boringssl//:LICENSE", + "@com_google_absl//:LICENSE", "@com_googlesource_code_re2//:LICENSE", "@cub_archive//:LICENSE.TXT", "@curl//:COPYING", "@eigen_archive//:COPYING.MPL2", "@farmhash_archive//:COPYING", "@fft2d//:fft/readme.txt", + "@flatbuffers//:LICENSE.txt", + "@gast_archive//:PKG-INFO", "@gemmlowp//:LICENSE", "@gif_archive//:COPYING", "@grpc//:LICENSE", "@highwayhash//:LICENSE", "@jemalloc//:COPYING", "@jpeg//:LICENSE.md", + "@kafka//:LICENSE", "@libxsmm_archive//:LICENSE", "@lmdb//:LICENSE", "@local_config_sycl//sycl:LICENSE.text", "@grpc//third_party/nanopb:LICENSE.txt", + "@nasm//:LICENSE", "@nsync//:LICENSE", + "@pcre//:LICENCE", "@png_archive//:LICENSE", "@protobuf_archive//:LICENSE", "@six_archive//:LICENSE", "@snappy//:COPYING", + "@swig//:LICENSE", + "@termcolor_archive//:COPYING.txt", "@zlib_archive//:zlib.h", "@org_python_pypi_backports_weakref//:LICENSE", ] + if_mkl([ @@ -126,62 +167,16 @@ sh_binary( data = select({ "//tensorflow:windows": [":simple_console_for_windows"], "//tensorflow:windows_msvc": [":simple_console_for_windows"], - "//conditions:default": [ - ":licenses", - "MANIFEST.in", - "README", - "setup.py", - ":included_headers", + "//conditions:default": COMMON_PIP_DEPS + [ ":simple_console", - "//tensorflow:tensorflow_py", - "//tensorflow/contrib/boosted_trees:boosted_trees_pip", - "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip", - "//tensorflow/contrib/data/python/kernel_tests:dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:prefetching_py", - "//tensorflow/contrib/eager/python/examples:examples_pip", - "//tensorflow/contrib/eager/python:checkpointable", - "//tensorflow/contrib/eager/python:evaluator", - "//tensorflow/contrib/gan:gan", - "//tensorflow/contrib/graph_editor:graph_editor_pip", - "//tensorflow/contrib/keras:keras", - "//tensorflow/contrib/labeled_tensor:labeled_tensor_pip", + "//tensorflow/contrib/lite/python:interpreter_test_data", "//tensorflow/contrib/lite/toco:toco", "//tensorflow/contrib/lite/toco/python:toco_wrapper", "//tensorflow/contrib/lite/toco/python:toco_from_protos", - "//tensorflow/contrib/ndlstm:ndlstm", - "//tensorflow/contrib/nn:nn_py", - "//tensorflow/contrib/predictor:predictor_pip", - "//tensorflow/contrib/py2tf:py2tf_internal", - "//tensorflow/contrib/py2tf/converters:converters", - "//tensorflow/contrib/py2tf/converters:test_lib", - "//tensorflow/contrib/py2tf/pyct:pyct", - "//tensorflow/contrib/py2tf/pyct/static_analysis:static_analysis", - "//tensorflow/contrib/receptive_field:receptive_field_pip", - "//tensorflow/contrib/session_bundle:session_bundle_pip", - "//tensorflow/contrib/signal:signal_py", - "//tensorflow/contrib/signal:test_util", - "//tensorflow/contrib/slim:slim", - "//tensorflow/contrib/slim/python/slim/data:data_pip", - "//tensorflow/contrib/slim/python/slim/nets:nets_pip", - "//tensorflow/contrib/specs:specs", - "//tensorflow/contrib/summary:summary_test_util", - "//tensorflow/contrib/tensor_forest:init_py", - "//tensorflow/contrib/tensor_forest/hybrid:hybrid_pip", - "//tensorflow/contrib/timeseries:timeseries_pip", - "//tensorflow/contrib/tpu", - "//tensorflow/examples/tutorials/mnist:package", - "//tensorflow/python:distributed_framework_test_lib", - "//tensorflow/python:meta_graph_testdata", - "//tensorflow/python:spectral_ops_test_util", - "//tensorflow/python:util_example_parser_configuration", - "//tensorflow/python/debug:debug_pip", - "//tensorflow/python/eager:eager_pip", - "//tensorflow/python/saved_model:saved_model", - "//tensorflow/python/tools:tools_pip", - "//tensorflow/python:test_ops", - "//tensorflow/tools/dist_test/server:grpc_tensorflow_server", ], - }) + if_mkl(["//third_party/mkl:intel_binary_blob"]), + }) + if_mkl(["//third_party/mkl:intel_binary_blob"]) + if_tensorrt([ + "//tensorflow/contrib/tensorrt:init_py", + ]), ) # A genrule for generating a marker file for the pip package on Windows diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py index 38a900738786e2413f5b1dd914caaebeafc92e21..b66c45ec13067e20fbd3c333cf562001945a63ea 100644 --- a/tensorflow/tools/pip_package/pip_smoke_test.py +++ b/tensorflow/tools/pip_package/pip_smoke_test.py @@ -58,6 +58,10 @@ BLACKLIST = [ # contrib "//tensorflow/contrib/session_bundle:session_bundle_half_plus_two", "//tensorflow/contrib/keras:testing_utils", + "//tensorflow/contrib/lite/python:interpreter", + "//tensorflow/contrib/lite/python:interpreter_test", + "//tensorflow/contrib/lite/python:interpreter.py", + "//tensorflow/contrib/lite/python:interpreter_test.py", "//tensorflow/contrib/ffmpeg:test_data", "//tensorflow/contrib/factorization/examples:mnist", "//tensorflow/contrib/factorization/examples:mnist.py", @@ -65,7 +69,6 @@ BLACKLIST = [ "//tensorflow/contrib/framework:checkpoint_ops_testdata", "//tensorflow/contrib/bayesflow:reinforce_simple_example", "//tensorflow/contrib/bayesflow:examples/reinforce_simple/reinforce_simple_example.py", # pylint:disable=line-too-long - "//tensorflow/contrib/py2tf:py2tf_internal", "//tensorflow/contrib/timeseries/examples:predict", "//tensorflow/contrib/timeseries/examples:multivariate", "//tensorflow/contrib/timeseries/examples:known_anomaly", diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py index a456300fae81afbeef4362d0ba1f588de2053870..e0152da4dff36212d736cd9c0e6d9be3b2324e31 100644 --- a/tensorflow/tools/pip_package/setup.py +++ b/tensorflow/tools/pip_package/setup.py @@ -29,16 +29,17 @@ from setuptools.dist import Distribution # This version string is semver compatible, but incompatible with pip. # For pip, we will remove all '-' characters from this string, and use the # result for pip. -_VERSION = '1.5.0' +_VERSION = '1.6.0' REQUIRED_PACKAGES = [ 'absl-py >= 0.1.6', 'astor >= 0.6.0', 'gast >= 0.2.0', - 'numpy >= 1.12.1', + 'grpcio >= 1.8.6', + 'numpy >= 1.13.3', 'six >= 1.10.0', 'protobuf >= 3.4.0', - 'tensorflow-tensorboard >= 1.5.0, < 1.6.0', + 'tensorboard >= 1.6.0, < 1.7.0', 'termcolor >= 1.1.0', ] @@ -61,7 +62,7 @@ else: if 'tf_nightly' in project_name: for i, pkg in enumerate(REQUIRED_PACKAGES): if 'tensorboard' in pkg: - REQUIRED_PACKAGES[i] = 'tb-nightly >= 1.5.0a0, < 1.6.0a0' + REQUIRED_PACKAGES[i] = 'tb-nightly >= 1.7.0a0, < 1.8.0a0' break # weakref.finalize and enum were introduced in Python 3.4 @@ -71,7 +72,7 @@ if sys.version_info < (3, 4): # pylint: disable=line-too-long CONSOLE_SCRIPTS = [ - 'freeze_graph = tensorflow.python.tools.freeze_graph:main', + 'freeze_graph = tensorflow.python.tools.freeze_graph:run_main', 'toco_from_protos = tensorflow.contrib.lite.toco.python.toco_from_protos:main', 'toco = tensorflow.contrib.lite.toco.python.toco_wrapper:main', 'saved_model_cli = tensorflow.python.tools.saved_model_cli:main', @@ -79,13 +80,13 @@ CONSOLE_SCRIPTS = [ # is now declared by the tensorboard pip package. If we remove the # TensorBoard command, pip will inappropriately remove it during install, # even though the command is not removed, just moved to a different wheel. - 'tensorboard = tensorboard.main:main', + 'tensorboard = tensorboard.main:run_main', ] # pylint: enable=line-too-long # remove the tensorboard console script if building tf_nightly if 'tf_nightly' in project_name: - CONSOLE_SCRIPTS.remove('tensorboard = tensorboard.main:main') + CONSOLE_SCRIPTS.remove('tensorboard = tensorboard.main:run_main') TEST_PACKAGES = [ 'scipy >= 0.15.1', @@ -180,9 +181,10 @@ def find_files(pattern, root): matches = ['../' + x for x in find_files('*', 'external') if '.py' not in x] -so_lib_paths = [i for i in os.listdir('.') - if os.path.isdir(i) - and fnmatch.fnmatch(i, '_solib_*')] +so_lib_paths = [ + i for i in os.listdir('.') + if os.path.isdir(i) and fnmatch.fnmatch(i, '_solib_*') +] for path in so_lib_paths: matches.extend( @@ -198,8 +200,7 @@ headers = (list(find_files('*.h', 'tensorflow/core')) + list(find_files('*.h', 'tensorflow/stream_executor')) + list(find_files('*.h', 'google/protobuf_archive/src')) + list(find_files('*', 'third_party/eigen3')) + - list(find_files('*', 'external/eigen_archive')) + - list(find_files('*.h', 'external/nsync/public'))) + list(find_files('*', 'external/eigen_archive'))) setup( name=project_name, diff --git a/tensorflow/tools/test/file_name_test.py b/tensorflow/tools/test/file_name_test.py new file mode 100644 index 0000000000000000000000000000000000000000..16fb8a822d09ed136cf79dd2473fc202ca632d83 --- /dev/null +++ b/tensorflow/tools/test/file_name_test.py @@ -0,0 +1,48 @@ +#!/usr/bin/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. +# ============================================================================== +# +# Test that checks if we have any issues with case insensitive filesystems. + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')) +ERROR_MESSAGE = """ +Files with same name but different case detected in directory: {} +""" + + +def main(): + # Make sure BASE_DIR ends with tensorflow. If it doesn't, we probably + # computed the wrong directory. + if os.path.split(BASE_DIR)[-1] != 'tensorflow': + raise AssertionError( + "BASE_DIR = '%s' doesn't end with tensorflow" % BASE_DIR) + + for dirpath, dirnames, filenames in os.walk(BASE_DIR, followlinks=True): + lowercase_directories = [x.lower() for x in dirnames] + lowercase_files = [x.lower() for x in filenames] + + lowercase_dir_contents = lowercase_directories + lowercase_files + if len(lowercase_dir_contents) != len(set(lowercase_dir_contents)): + raise AssertionError(ERROR_MESSAGE.format(dirpath)) + + +if __name__ == '__main__': + main() diff --git a/tensorflow/tools/test/run_and_gather_logs_lib.py b/tensorflow/tools/test/run_and_gather_logs_lib.py index a953ed1b53d13504f92d2ffeb4c1ac6bcb0b8477..3b4921bb983a72223b092d99eb3fb59332fc6345 100644 --- a/tensorflow/tools/test/run_and_gather_logs_lib.py +++ b/tensorflow/tools/test/run_and_gather_logs_lib.py @@ -136,7 +136,7 @@ def run_and_gather_logs(name, test_name, test_args, gpu_config = gpu_info_lib.gather_gpu_devices() if gpu_config: gpu_name = gpu_config[0].model - gpu_short_name_match = re.search(r"Tesla (K40|K80|P100)", gpu_name) + gpu_short_name_match = re.search(r"Tesla (K40|K80|P100|V100)", gpu_name) if gpu_short_name_match: gpu_short_name = gpu_short_name_match.group(0) test_adjusted_name = name + "|" + gpu_short_name.replace(" ", "_") diff --git a/tensorflow/tools/test/upload_test_benchmarks.py b/tensorflow/tools/test/upload_test_benchmarks.py index 77cc9f75f7725438918f681833d58e9ecb4a2f70..edd093510e5483ef946e1177641371f00c92b50b 100644 --- a/tensorflow/tools/test/upload_test_benchmarks.py +++ b/tensorflow/tools/test/upload_test_benchmarks.py @@ -88,6 +88,7 @@ import os import shutil from google.cloud import datastore +from six import text_type def is_real_file(dirpath, fname): @@ -150,7 +151,7 @@ def upload_benchmark_data(client, data): """ test_result = json.loads(data) - test_name = unicode(test_result["name"]) + test_name = text_type(test_result["name"]) start_time = datetime.datetime.utcfromtimestamp( float(test_result["startTime"])) batch = [] @@ -162,7 +163,7 @@ def upload_benchmark_data(client, data): t_val.update({ "test": test_name, "start": start_time, - "info": unicode(data) + "info": text_type(data) }) batch.append(t_val) @@ -170,7 +171,7 @@ def upload_benchmark_data(client, data): # the attribute to be fetched and displayed. The full entry information is # also stored as a non-indexed JSON blob. for ent in test_result["entries"].get("entry", []): - ent_name = unicode(ent["name"]) + ent_name = text_type(ent["name"]) e_key = client.key("Entry") e_val = datastore.Entity(e_key, exclude_from_indexes=["info"]) e_val.update({ @@ -178,7 +179,7 @@ def upload_benchmark_data(client, data): "start": start_time, "entry": ent_name, "timing": ent["wallTime"], - "info": unicode(json.dumps(ent)) + "info": text_type(json.dumps(ent)) }) batch.append(e_val) diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index 660520d36a7e10295178f64f9f4657a07af3eada..a922808a70f8bd082f2aab7d0f01c404b304e121 100644 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -5,12 +5,15 @@ load("//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure") load("//third_party/mkl:build_defs.bzl", "mkl_repository") load("//third_party/git:git_configure.bzl", "git_configure") load("//third_party/py:python_configure.bzl", "python_configure") + load("//third_party/sycl:sycl_configure.bzl", "sycl_configure") load("//third_party/toolchains/clang6:repo.bzl", "clang6_configure") load("//third_party/toolchains/cpus/arm:arm_compiler_configure.bzl", "arm_compiler_configure") load("//third_party:repo.bzl", "tf_http_archive") load("@io_bazel_rules_closure//closure/private:java_import_external.bzl", "java_import_external") load("@io_bazel_rules_closure//closure:defs.bzl", "filegroup_external") +load("//tensorflow/tools/def_file_filter:def_file_filter_configure.bzl", + "def_file_filter_configure") def _extract_version_number(bazel_version): """Extracts the semantic version number from a version string @@ -66,7 +69,7 @@ def tf_workspace(path_prefix="", tf_repo_name=""): # We must check the bazel version before trying to parse any other BUILD # files, in case the parsing of those build files depends on the bazel # version we require here. - check_bazel_version_at_least("0.5.4") + check_bazel_version_at_least("0.10.0") clang6_configure(name="local_config_clang6") cuda_configure(name="local_config_cuda") tensorrt_configure(name="local_config_tensorrt") @@ -74,6 +77,10 @@ def tf_workspace(path_prefix="", tf_repo_name=""): sycl_configure(name="local_config_sycl") python_configure(name="local_config_python") + # For windows bazel build + # TODO: Remove def file filter when TensorFlow can export symbols properly on Windows. + def_file_filter_configure(name = "local_config_def_file_filter") + # Point //external/local_config_arm_compiler to //external/arm_compiler arm_compiler_configure( name="local_config_arm_compiler", @@ -114,17 +121,19 @@ def tf_workspace(path_prefix="", tf_repo_name=""): ], sha256 = "5996380e3e8b981f55d1c8d58e709c00dbb4806ba367be75d0925a68cc2f6478", strip_prefix = "abseil-cpp-720c017e30339fd1786ce4aac68bc8559736e53f", + build_file = str(Label("//third_party:com_google_absl.BUILD")), ) tf_http_archive( name = "eigen_archive", urls = [ - "https://mirror.bazel.build/bitbucket.org/eigen/eigen/get/14e1418fcf12.tar.gz", - "https://bitbucket.org/eigen/eigen/get/14e1418fcf12.tar.gz", + "https://mirror.bazel.build/bitbucket.org/eigen/eigen/get/2355b229ea4c.tar.gz", + "https://bitbucket.org/eigen/eigen/get/2355b229ea4c.tar.gz", ], - sha256 = "2b526c6888639025323fd4f2600533c0f982d304ea48e4f1663e8066bd9f6368", - strip_prefix = "eigen-eigen-14e1418fcf12", + sha256 = "0cadb31a35b514bf2dfd6b5d38205da94ef326ec6908fc3fd7c269948467214f", + strip_prefix = "eigen-eigen-2355b229ea4c", build_file = str(Label("//third_party:eigen.BUILD")), + patch_file = str(Label("//third_party:eigen_fix_cuda_compilation.patch")) ) tf_http_archive( @@ -178,11 +187,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "gemmlowp", urls = [ - "https://mirror.bazel.build/github.com/google/gemmlowp/archive/010bb3e71a26ca1d0884a167081d092b43563996.zip", - "https://github.com/google/gemmlowp/archive/010bb3e71a26ca1d0884a167081d092b43563996.zip", + "https://mirror.bazel.build/github.com/google/gemmlowp/archive/7c7c744640ddc3d0af18fb245b4d23228813a71b.zip", + "https://github.com/google/gemmlowp/archive/7c7c744640ddc3d0af18fb245b4d23228813a71b.zip", ], - sha256 = "dd2557072bde12141419cb8320a9c25e6ec41a8ae53c2ac78c076a347bb46d9d", - strip_prefix = "gemmlowp-010bb3e71a26ca1d0884a167081d092b43563996", + sha256 = "b852cc90259a7357c8a323f108f2cec6e85979fc3b18b5590b99e0130044b2cf", + strip_prefix = "gemmlowp-7c7c744640ddc3d0af18fb245b4d23228813a71b", ) tf_http_archive( @@ -352,11 +361,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "protobuf_archive", urls = [ - "https://mirror.bazel.build/github.com/google/protobuf/archive/b04e5cba356212e4e8c66c61bbe0c3a20537c5b9.tar.gz", - "https://github.com/google/protobuf/archive/b04e5cba356212e4e8c66c61bbe0c3a20537c5b9.tar.gz", + "https://mirror.bazel.build/github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz", + "https://github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz", ], - sha256 = "e178a25c52efcb6b05988bdbeace4c0d3f2d2fe5b46696d1d9898875c3803d6a", - strip_prefix = "protobuf-b04e5cba356212e4e8c66c61bbe0c3a20537c5b9", + sha256 = "846d907acf472ae233ec0882ef3a2d24edbbe834b80c305e867ac65a1f2c59e3", + strip_prefix = "protobuf-396336eb961b75f03b25824fe86cf6490fb75e3a", ) # We need to import the protobuf library under the names com_google_protobuf @@ -365,21 +374,21 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "com_google_protobuf", urls = [ - "https://mirror.bazel.build/github.com/google/protobuf/archive/b04e5cba356212e4e8c66c61bbe0c3a20537c5b9.tar.gz", - "https://github.com/google/protobuf/archive/b04e5cba356212e4e8c66c61bbe0c3a20537c5b9.tar.gz", + "https://mirror.bazel.build/github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz", + "https://github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz", ], - sha256 = "e178a25c52efcb6b05988bdbeace4c0d3f2d2fe5b46696d1d9898875c3803d6a", - strip_prefix = "protobuf-b04e5cba356212e4e8c66c61bbe0c3a20537c5b9", + sha256 = "846d907acf472ae233ec0882ef3a2d24edbbe834b80c305e867ac65a1f2c59e3", + strip_prefix = "protobuf-396336eb961b75f03b25824fe86cf6490fb75e3a", ) tf_http_archive( name = "com_google_protobuf_cc", urls = [ - "https://mirror.bazel.build/github.com/google/protobuf/archive/b04e5cba356212e4e8c66c61bbe0c3a20537c5b9.tar.gz", - "https://github.com/google/protobuf/archive/b04e5cba356212e4e8c66c61bbe0c3a20537c5b9.tar.gz", + "https://mirror.bazel.build/github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz", + "https://github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz", ], - sha256 = "e178a25c52efcb6b05988bdbeace4c0d3f2d2fe5b46696d1d9898875c3803d6a", - strip_prefix = "protobuf-b04e5cba356212e4e8c66c61bbe0c3a20537c5b9", + sha256 = "846d907acf472ae233ec0882ef3a2d24edbbe834b80c305e867ac65a1f2c59e3", + strip_prefix = "protobuf-396336eb961b75f03b25824fe86cf6490fb75e3a", ) tf_http_archive( @@ -472,11 +481,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "llvm", urls = [ - "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/36a30fc7c9ee6fdfe5157190ad15c1801b1ab2de.tar.gz", - "https://github.com/llvm-mirror/llvm/archive/36a30fc7c9ee6fdfe5157190ad15c1801b1ab2de.tar.gz", + "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/738ee045416377e8c2094f7f61508ac1c178ff37.tar.gz", + "https://github.com/llvm-mirror/llvm/archive/738ee045416377e8c2094f7f61508ac1c178ff37.tar.gz", ], - sha256 = "3b74ecd8f59c712b4daf715a4da15c43ebdd40edcd4c30737bffef62f6a2bc9d", - strip_prefix = "llvm-36a30fc7c9ee6fdfe5157190ad15c1801b1ab2de", + sha256 = "4442ed6a05c13752338036b1b9f16b09264de24b6c0bf62325fb9ff75a09340f", + strip_prefix = "llvm-738ee045416377e8c2094f7f61508ac1c178ff37", build_file = str(Label("//third_party/llvm:llvm.BUILD")), ) @@ -640,11 +649,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): ) java_import_external( - name = "javax_validation", - jar_sha256 = "e459f313ebc6db2483f8ceaad39af07086361b474fa92e40f442e8de5d9895dc", + name = "org_checkerframework_qual", + jar_sha256 = "a17501717ef7c8dda4dba73ded50c0d7cde440fd721acfeacbf19786ceac1ed6", jar_urls = [ - "http://mirror.bazel.build/repo1.maven.org/maven2/javax/validation/validation-api/1.0.0.GA/validation-api-1.0.0.GA.jar", - "http://repo1.maven.org/maven2/javax/validation/validation-api/1.0.0.GA/validation-api-1.0.0.GA.jar", + "http://mirror.bazel.build/repo1.maven.org/maven2/org/checkerframework/checker-qual/2.4.0/checker-qual-2.4.0.jar", + "http://repo1.maven.org/maven2/org/checkerframework/checker-qual/2.4.0/checker-qual-2.4.0.jar", ], licenses = ["notice"], # Apache 2.0 ) @@ -663,11 +672,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "cub_archive", urls = [ - "https://mirror.bazel.build/github.com/NVlabs/cub/archive/1.7.4.zip", - "https://github.com/NVlabs/cub/archive/1.7.4.zip", + "https://mirror.bazel.build/github.com/NVlabs/cub/archive/1.8.0.zip", + "https://github.com/NVlabs/cub/archive/1.8.0.zip", ], - sha256 = "20a1a39fd97e5da7f40f5f2e7fd73fd2ea59f9dc4bb8a6c5f228aa543e727e31", - strip_prefix = "cub-1.7.4", + sha256 = "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3", + strip_prefix = "cub-1.8.0", build_file = str(Label("//third_party:cub.BUILD")), ) @@ -686,11 +695,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "bazel_toolchains", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/bazel-toolchains/archive/f3b09700fae5d7b6e659d7cefe0dcc6e8498504c.tar.gz", - "https://github.com/bazelbuild/bazel-toolchains/archive/f3b09700fae5d7b6e659d7cefe0dcc6e8498504c.tar.gz", + "https://mirror.bazel.build/github.com/bazelbuild/bazel-toolchains/archive/44200e0c026d86c53470d107b3697a3e46469c43.tar.gz", + "https://github.com/bazelbuild/bazel-toolchains/archive/44200e0c026d86c53470d107b3697a3e46469c43.tar.gz", ], - sha256 = "ed829b5eea8af1f405f4cc3d6ecfc3b1365bb7843171036030a31b5127002311", - strip_prefix = "bazel-toolchains-f3b09700fae5d7b6e659d7cefe0dcc6e8498504c", + strip_prefix = "bazel-toolchains-44200e0c026d86c53470d107b3697a3e46469c43", + sha256 = "699b55a6916c687f4b7dc092dbbf5f64672cde0dc965f79717735ec4e5416556", ) tf_http_archive( diff --git a/third_party/com_google_absl.BUILD b/third_party/com_google_absl.BUILD new file mode 100644 index 0000000000000000000000000000000000000000..8fca145f751eacfa3e5a0af046dcc8c19e6a85d4 --- /dev/null +++ b/third_party/com_google_absl.BUILD @@ -0,0 +1,5 @@ +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) # Apache + +exports_files(["LICENSE"]) diff --git a/third_party/eigen_fix_cuda_compilation.patch b/third_party/eigen_fix_cuda_compilation.patch new file mode 100644 index 0000000000000000000000000000000000000000..b921a7c31d5c96c79cd3033b13c60a8f7e63ba75 --- /dev/null +++ b/third_party/eigen_fix_cuda_compilation.patch @@ -0,0 +1,38 @@ +diff --git a/Eigen/src/Core/ProductEvaluators.h b/Eigen/src/Core/ProductEvaluators.h +--- a/Eigen/src/Core/ProductEvaluators.h ++++ b/Eigen/src/Core/ProductEvaluators.h +@@ -137,7 +137,7 @@ struct Assignment::type> + { + typedef Product SrcXprType; +- static EIGEN_STRONG_INLINE ++ static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE + void run(DstXprType &dst, const SrcXprType &src, const internal::assign_op &) + { + Index dstRows = src.rows(); +@@ -390,7 +390,7 @@ struct generic_product_impl::Scalar Scalar; + + template +- static EIGEN_STRONG_INLINE void evalTo(Dst& dst, const Lhs& lhs, const Rhs& rhs) ++ static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalTo(Dst& dst, const Lhs& lhs, const Rhs& rhs) + { + // Same as: dst.noalias() = lhs.lazyProduct(rhs); + // but easier on the compiler side +@@ -398,14 +398,14 @@ struct generic_product_impl +- static EIGEN_STRONG_INLINE void addTo(Dst& dst, const Lhs& lhs, const Rhs& rhs) ++ static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void addTo(Dst& dst, const Lhs& lhs, const Rhs& rhs) + { + // dst.noalias() += lhs.lazyProduct(rhs); + call_assignment_no_alias(dst, lhs.lazyProduct(rhs), internal::add_assign_op()); + } + + template +- static EIGEN_STRONG_INLINE void subTo(Dst& dst, const Lhs& lhs, const Rhs& rhs) ++ static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void subTo(Dst& dst, const Lhs& lhs, const Rhs& rhs) + { + // dst.noalias() -= lhs.lazyProduct(rhs); + call_assignment_no_alias(dst, lhs.lazyProduct(rhs), internal::sub_assign_op()); diff --git a/third_party/examples/eager/spinn/README.md b/third_party/examples/eager/spinn/README.md index 6bd3d53e56d01e15491ecd383dcc763a19d75b88..7f477d19208257474d0481ca04c04679f589b751 100644 --- a/third_party/examples/eager/spinn/README.md +++ b/third_party/examples/eager/spinn/README.md @@ -66,3 +66,44 @@ Other eager execution examples can be found under [tensorflow/contrib/eager/pyth ```bash tensorboard --logdir /tmp/spinn-logs ``` + +- After training, you may use the model to perform inference on input data in + the SNLI data format. The premise and hypotheses sentences are specified with + the command-line flags `--inference_premise` and `--inference_hypothesis`, + respecitvely. Each sentence should include the words, as well as parentheses + representing a binary parsing of the sentence. The words and parentheses + should all be separated by spaces. For instance, + + ```bash + python spinn.py --data_root /tmp/spinn-data --logdir /tmp/spinn-logs \ + --inference_premise '( ( The dog ) ( ( is running ) . ) )' \ + --inference_hypothesis '( ( The dog ) ( moves . ) )' + ``` + + which will generate an output like the following, due to the semantic + consistency of the two sentences. + + ```none + Inference logits: + entailment: 1.101249 (winner) + contradiction: -2.374171 + neutral: -0.296733 + ``` + + By contrast, the following sentence pair: + + ```bash + python spinn.py --data_root /tmp/spinn-data --logdir /tmp/spinn-logs \ + --inference_premise '( ( The dog ) ( ( is running ) . ) )' \ + --inference_hypothesis '( ( The dog ) ( rests . ) )' + ``` + + will give you an output like the following, due to the semantic + contradiction of the two sentences. + + ```none + Inference logits: + entailment: -1.070098 + contradiction: 2.798695 (winner) + neutral: -1.402287 + ``` diff --git a/third_party/examples/eager/spinn/spinn.py b/third_party/examples/eager/spinn/spinn.py index a2fa18eeb1077c8a1ccd4ab0bcd178f952e17270..8a1c7db2ea14365be53a796a79fce77900e668e1 100644 --- a/third_party/examples/eager/spinn/spinn.py +++ b/third_party/examples/eager/spinn/spinn.py @@ -471,6 +471,15 @@ class SNLIClassifierTrainer(object): def learning_rate(self): return self._learning_rate + @property + def model(self): + return self._model + + @property + def variables(self): + return (self._model.variables + [self.learning_rate] + + self._optimizer.variables()) + def _batch_n_correct(logits, label): """Calculate number of correct predictions in a batch. @@ -488,13 +497,12 @@ def _batch_n_correct(logits, label): tf.argmax(logits, axis=1), label)), tf.float32)).numpy() -def _evaluate_on_dataset(snli_data, batch_size, model, trainer, use_gpu): +def _evaluate_on_dataset(snli_data, batch_size, trainer, use_gpu): """Run evaluation on a dataset. Args: snli_data: The `data.SnliData` to use in this evaluation. batch_size: The batch size to use during this evaluation. - model: An instance of `SNLIClassifier` to evaluate. trainer: An instance of `SNLIClassifierTrainer to use for this evaluation. use_gpu: Whether GPU is being used. @@ -509,7 +517,7 @@ def _evaluate_on_dataset(snli_data, batch_size, model, trainer, use_gpu): snli_data, batch_size): if use_gpu: label, prem, hypo = label.gpu(), prem.gpu(), hypo.gpu() - logits = model(prem, prem_trans, hypo, hypo_trans, training=False) + logits = trainer.model(prem, prem_trans, hypo, hypo_trans, training=False) loss_val = trainer.loss(label, logits) batch_size = tf.shape(label)[0] mean_loss(loss_val, weights=batch_size.gpu() if use_gpu else batch_size) @@ -536,13 +544,19 @@ def _get_dataset_iterator(snli_data, batch_size): return tfe.Iterator(dataset) -def train_spinn(embed, train_data, dev_data, test_data, config): - """Train a SPINN model. +def train_or_infer_spinn(embed, + word2index, + train_data, + dev_data, + test_data, + config): + """Perform Training or Inference on a SPINN model. Args: embed: The embedding matrix as a float32 numpy array with shape [vocabulary_size, word_vector_len]. word_vector_len is the length of a word embedding vector. + word2index: A `dict` mapping word to word index. train_data: An instance of `data.SnliData`, for the train split. dev_data: Same as above, for the dev split. test_data: Same as above, for the test split. @@ -550,13 +564,35 @@ def train_spinn(embed, train_data, dev_data, test_data, config): details. Returns: - 1. Final loss value on the test split. - 2. Final fraction of correct classifications on the test split. + If `config.inference_premise ` and `config.inference_hypothesis` are not + `None`, i.e., inference mode: the logits for the possible labels of the + SNLI data set, as a `Tensor` of three floats. + else: + The trainer object. + Raises: + ValueError: if only one of config.inference_premise and + config.inference_hypothesis is specified. """ + # TODO(cais): Refactor this function into separate one for training and + # inference. use_gpu = tfe.num_gpus() > 0 and not config.force_cpu device = "gpu:0" if use_gpu else "cpu:0" print("Using device: %s" % device) + if ((config.inference_premise and not config.inference_hypothesis) or + (not config.inference_premise and config.inference_hypothesis)): + raise ValueError( + "--inference_premise and --inference_hypothesis must be both " + "specified or both unspecified, but only one is specified.") + + if config.inference_premise: + # Inference mode. + inference_sentence_pair = [ + data.encode_sentence(config.inference_premise, word2index), + data.encode_sentence(config.inference_hypothesis, word2index)] + else: + inference_sentence_pair = None + log_header = ( " Time Epoch Iteration Progress (%Epoch) Loss Dev/Loss" " Accuracy Dev/Accuracy") @@ -569,16 +605,36 @@ def train_spinn(embed, train_data, dev_data, test_data, config): summary_writer = tf.contrib.summary.create_file_writer( config.logdir, flush_millis=10000) - train_len = train_data.num_batches(config.batch_size) + with tf.device(device), \ - tfe.restore_variables_on_create( - tf.train.latest_checkpoint(config.logdir)), \ summary_writer.as_default(), \ tf.contrib.summary.always_record_summaries(): - model = SNLIClassifier(config, embed) - global_step = tf.train.get_or_create_global_step() - trainer = SNLIClassifierTrainer(model, config.lr) - + with tfe.restore_variables_on_create( + tf.train.latest_checkpoint(config.logdir)): + model = SNLIClassifier(config, embed) + global_step = tf.train.get_or_create_global_step() + trainer = SNLIClassifierTrainer(model, config.lr) + + if inference_sentence_pair: + # Inference mode. + with tfe.restore_variables_on_create( + tf.train.latest_checkpoint(config.logdir)): + prem, prem_trans = inference_sentence_pair[0] + hypo, hypo_trans = inference_sentence_pair[1] + hypo_trans = inference_sentence_pair[1][1] + inference_logits = model( # pylint: disable=not-callable + tf.constant(prem), tf.constant(prem_trans), + tf.constant(hypo), tf.constant(hypo_trans), training=False) + inference_logits = inference_logits[0][1:] + max_index = tf.argmax(inference_logits) + print("\nInference logits:") + for i, (label, logit) in enumerate( + zip(data.POSSIBLE_LABELS, inference_logits)): + winner_tag = " (winner)" if max_index == i else "" + print(" {0:<16}{1:.6f}{2}".format(label + ":", logit, winner_tag)) + return inference_logits + + train_len = train_data.num_batches(config.batch_size) start = time.time() iterations = 0 mean_loss = tfe.metrics.Mean() @@ -594,23 +650,24 @@ def train_spinn(embed, train_data, dev_data, test_data, config): # remain on CPU. Same in _evaluate_on_dataset(). iterations += 1 - batch_train_loss, batch_train_logits = trainer.train_batch( - label, prem, prem_trans, hypo, hypo_trans) + with tfe.restore_variables_on_create( + tf.train.latest_checkpoint(config.logdir)): + batch_train_loss, batch_train_logits = trainer.train_batch( + label, prem, prem_trans, hypo, hypo_trans) batch_size = tf.shape(label)[0] mean_loss(batch_train_loss.numpy(), weights=batch_size.gpu() if use_gpu else batch_size) accuracy(tf.argmax(batch_train_logits, axis=1), label) if iterations % config.save_every == 0: - all_variables = ( - model.variables + [trainer.learning_rate] + [global_step]) + all_variables = trainer.variables + [global_step] saver = tfe.Saver(all_variables) saver.save(os.path.join(config.logdir, "ckpt"), global_step=global_step) if iterations % config.dev_every == 0: dev_loss, dev_frac_correct = _evaluate_on_dataset( - dev_data, config.batch_size, model, trainer, use_gpu) + dev_data, config.batch_size, trainer, use_gpu) print(dev_log_template.format( time.time() - start, epoch, iterations, 1 + batch_idx, train_len, @@ -638,10 +695,12 @@ def train_spinn(embed, train_data, dev_data, test_data, config): trainer.decay_learning_rate(config.lr_decay_by) test_loss, test_frac_correct = _evaluate_on_dataset( - test_data, config.batch_size, model, trainer, use_gpu) + test_data, config.batch_size, trainer, use_gpu) print("Final test loss: %g; accuracy: %g%%" % (test_loss, test_frac_correct * 100.0)) + return trainer + def main(_): config = FLAGS @@ -650,18 +709,24 @@ def main(_): vocab = data.load_vocabulary(FLAGS.data_root) word2index, embed = data.load_word_vectors(FLAGS.data_root, vocab) - print("Loading train, dev and test data...") - train_data = data.SnliData( - os.path.join(FLAGS.data_root, "snli/snli_1.0/snli_1.0_train.txt"), - word2index, sentence_len_limit=FLAGS.sentence_len_limit) - dev_data = data.SnliData( - os.path.join(FLAGS.data_root, "snli/snli_1.0/snli_1.0_dev.txt"), - word2index, sentence_len_limit=FLAGS.sentence_len_limit) - test_data = data.SnliData( - os.path.join(FLAGS.data_root, "snli/snli_1.0/snli_1.0_test.txt"), - word2index, sentence_len_limit=FLAGS.sentence_len_limit) - - train_spinn(embed, train_data, dev_data, test_data, config) + if not (config.inference_premise or config.inference_hypothesis): + print("Loading train, dev and test data...") + train_data = data.SnliData( + os.path.join(FLAGS.data_root, "snli/snli_1.0/snli_1.0_train.txt"), + word2index, sentence_len_limit=FLAGS.sentence_len_limit) + dev_data = data.SnliData( + os.path.join(FLAGS.data_root, "snli/snli_1.0/snli_1.0_dev.txt"), + word2index, sentence_len_limit=FLAGS.sentence_len_limit) + test_data = data.SnliData( + os.path.join(FLAGS.data_root, "snli/snli_1.0/snli_1.0_test.txt"), + word2index, sentence_len_limit=FLAGS.sentence_len_limit) + else: + train_data = None + dev_data = None + test_data = None + + train_or_infer_spinn( + embed, word2index, train_data, dev_data, test_data, config) if __name__ == "__main__": @@ -678,6 +743,15 @@ if __name__ == "__main__": parser.add_argument("--logdir", type=str, default="/tmp/spinn-logs", help="Directory in which summaries will be written for " "TensorBoard.") + parser.add_argument("--inference_premise", type=str, default=None, + help="Premise sentence for inference. Must be " + "accompanied by --inference_hypothesis. If specified, " + "will override all training parameters and perform " + "inference.") + parser.add_argument("--inference_hypothesis", type=str, default=None, + help="Hypothesis sentence for inference. Must be " + "accompanied by --inference_premise. If specified, will " + "override all training parameters and perform inference.") parser.add_argument("--epochs", type=int, default=50, help="Number of epochs to train.") parser.add_argument("--batch_size", type=int, default=128, diff --git a/third_party/flatbuffers/flatbuffers.BUILD b/third_party/flatbuffers/flatbuffers.BUILD index f6b8e6ddb05e67a4bb4833a3bba6db3cbd4c79e0..824c97be60e7ef148a363b964ed330ba3c5fcb0c 100644 --- a/third_party/flatbuffers/flatbuffers.BUILD +++ b/third_party/flatbuffers/flatbuffers.BUILD @@ -4,6 +4,8 @@ package( licenses(["notice"]) # Apache 2.0 +exports_files(["LICENSE.txt"]) + config_setting( name = "freebsd", values = {"cpu": "freebsd"}, diff --git a/third_party/gast.BUILD b/third_party/gast.BUILD index 06db528ada27e2f26f6de48c1ce6e9348ce09173..4866982e1fda6d6f19e575c8b0c0273cb9de154b 100644 --- a/third_party/gast.BUILD +++ b/third_party/gast.BUILD @@ -3,7 +3,7 @@ licenses(["notice"]) # BSD 3-clause -exports_files(["LICENSE"]) +exports_files(["PKG-INFO"]) py_library( name = "gast", diff --git a/third_party/gpus/crosstool/CROSSTOOL_clang.tpl b/third_party/gpus/crosstool/CROSSTOOL_clang.tpl index e4363d604577de09241d635b6990c9dd6429efe0..2f09473ee2ddf9a38ca0c7aa11094690607b532f 100644 --- a/third_party/gpus/crosstool/CROSSTOOL_clang.tpl +++ b/third_party/gpus/crosstool/CROSSTOOL_clang.tpl @@ -49,6 +49,7 @@ toolchain { flag_set { action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag: "-lstdc++" } @@ -75,6 +76,7 @@ toolchain { name: "alwayslink" flag_set { action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" action: "c++-link-executable" flag_group { flag: "-Wl,-no-as-needed" @@ -116,6 +118,7 @@ toolchain { } flag_set { action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag: "-Wl,-z,relro,-z,now" } @@ -161,6 +164,7 @@ toolchain { flag_set { action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { # Stamp the binary with a unique identifier. flag: "-Wl,--build-id=md5" @@ -176,6 +180,7 @@ toolchain { action: "c++-compile" action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag:"-no-canonical-prefixes" } @@ -199,6 +204,7 @@ toolchain { flag_set { action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag: "-B/usr/bin/" } @@ -246,6 +252,7 @@ toolchain { } flag_set { action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" action: "c++-link-executable" flag_group { flag: "-Wl,--gc-sections" diff --git a/third_party/gpus/cuda/remote.BUILD.tpl b/third_party/gpus/cuda/remote.BUILD.tpl index d88d512b90c352e6a301ed6efe8266d8dd6bf744..f774def5e6cec25e4920ecce0076340a31c70386 100644 --- a/third_party/gpus/cuda/remote.BUILD.tpl +++ b/third_party/gpus/cuda/remote.BUILD.tpl @@ -41,65 +41,65 @@ config_setting( alias( name = "cuda_headers", - actual = "%{remote_cuda_repo}cuda:cuda_headers", + actual = "%{remote_cuda_repo}/cuda:cuda_headers", ) alias( name = "cudart_static", - actual = "%{remote_cuda_repo}cuda:cudart_static", + actual = "%{remote_cuda_repo}/cuda:cudart_static", ) alias( name = "cuda_driver", - actual = "%{remote_cuda_repo}cuda:cuda_driver", + actual = "%{remote_cuda_repo}/cuda:cuda_driver", ) alias( name = "cudart", - actual = "%{remote_cuda_repo}cuda:cudart", + actual = "%{remote_cuda_repo}/cuda:cudart", ) alias( name = "cublas", - actual = "%{remote_cuda_repo}cuda:cublas", + actual = "%{remote_cuda_repo}/cuda:cublas", ) alias( name = "cusolver", - actual = "%{remote_cuda_repo}cuda:cusolver", + actual = "%{remote_cuda_repo}/cuda:cusolver", ) alias( name = "cudnn", - actual = "%{remote_cuda_repo}cuda:cudnn", + actual = "%{remote_cuda_repo}/cuda:cudnn", ) alias( name = "cufft", - actual = "%{remote_cuda_repo}cuda:cufft", + actual = "%{remote_cuda_repo}/cuda:cufft", ) alias( name = "curand", - actual = "%{remote_cuda_repo}cuda:curand", + actual = "%{remote_cuda_repo}/cuda:curand", ) alias( name = "cuda", - actual = "%{remote_cuda_repo}cuda:cuda", + actual = "%{remote_cuda_repo}/cuda:cuda", ) alias( name = "cupti_headers", - actual = "%{remote_cuda_repo}cuda:cupti_headers", + actual = "%{remote_cuda_repo}/cuda:cupti_headers", ) alias( name = "cupti_dsos", - actual = "%{remote_cuda_repo}cuda:cupti_dsos", + actual = "%{remote_cuda_repo}/cuda:cupti_dsos", ) alias( name = "libdevice_root", - actual = "%{remote_cuda_repo}cuda:libdevice_root", + actual = "%{remote_cuda_repo}/cuda:libdevice_root", ) diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl index 8e1dd8a54f53f58a367af21ef1c17f2695431bad..6c9c128db635ecb23b7aa930dca1d20bae17f71c 100644 --- a/third_party/gpus/cuda_configure.bzl +++ b/third_party/gpus/cuda_configure.bzl @@ -38,6 +38,64 @@ _DEFAULT_CUDA_TOOLKIT_PATH = "/usr/local/cuda" _DEFAULT_CUDNN_INSTALL_PATH = "/usr/local/cuda" _DEFAULT_CUDA_COMPUTE_CAPABILITIES = ["3.5", "5.2"] +# Lookup paths for CUDA / cuDNN libraries, relative to the install directories. +# +# Paths will be tried out in the order listed below. The first successful path +# will be used. For example, when looking for the cudart libraries, the first +# attempt will be lib64/cudart inside the CUDA toolkit. +CUDA_LIB_PATHS = [ + "lib64/", + "lib64/stubs/", + "lib/x86_64-linux-gnu/", + "lib/x64/", + "lib/", + "", +] + +# Lookup paths for cupti.h, relative to the CUDA toolkit directory. +# +# On most systems, the cupti library is not installed in the same directory as +# the other CUDA libraries but rather in a special extras/CUPTI directory. +CUPTI_HEADER_PATHS = [ + "extras/CUPTI/include/", + "include/cuda/CUPTI/", +] + +# Lookup paths for the cupti library, relative to the +# +# On most systems, the cupti library is not installed in the same directory as +# the other CUDA libraries but rather in a special extras/CUPTI directory. +CUPTI_LIB_PATHS = [ + "extras/CUPTI/lib64/", + "lib/x86_64-linux-gnu", + "lib64/", + "extras/CUPTI/libx64/", + "extras/CUPTI/lib/", + "lib/", +] + +# Lookup paths for CUDA headers (cuda.h) relative to the CUDA toolkit directory. +CUDA_INCLUDE_PATHS = [ + "include/", + "include/cuda/" +] + +# Lookup paths for cudnn.h relative to the CUDNN install directory. +CUDNN_INCLUDE_PATHS = [ + "", + "include/", + "include/cuda/", +] + +# Lookup paths for NVVM libdevice relative to the CUDA directory toolkit. +# +# libdevice implements mathematical functions for GPU kernels, and is provided +# in NVVM bitcode (a subset of LLVM bitcode). +NVVM_LIBDEVICE_PATHS = [ + "nvvm/libdevice/", + "share/cuda/", +] + load(":download_clang.bzl", "download_clang") # TODO(dzc): Once these functions have been factored out of Bazel's @@ -367,11 +425,20 @@ def find_cuda_define(repository_ctx, header_dir, header_file, define): if result.stdout.find(define) == -1: auto_configure_fail("Cannot find line containing '%s' in %s" % (define, h_path)) - version = result.stdout - # Remove the new line and '\' character if any. - version = version.replace("\\", " ") - version = version.replace("\n", " ") - version = version.replace(define, "").lstrip() + # Split results to lines + lines = result.stdout.split('\n') + num_lines = len(lines) + for l in range(num_lines): + line = lines[l] + if define in line: # Find the line with define + version = line + if l != num_lines-1 and line[-1] == '\\': # Add next line, if multiline + version = version[:-1] + lines[l+1] + break + # Remove any comments + version = version.split("//")[0] + # Remove define name + version = version.replace(define, "").strip() # Remove the code after the version number. version_end = version.find(" ") if version_end != -1: @@ -513,31 +580,31 @@ def _find_cuda_lib(lib, repository_ctx, cpu_value, basedir, version="", path: The full path to the library. """ file_name = _lib_name(lib, cpu_value, version, static) - if cpu_value == "Linux": - path = repository_ctx.path("%s/lib64/%s" % (basedir, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - path = repository_ctx.path("%s/lib64/stubs/%s" % (basedir, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - path = repository_ctx.path( - "%s/lib/x86_64-linux-gnu/%s" % (basedir, file_name)) + for relative_path in CUDA_LIB_PATHS: + path = repository_ctx.path("%s/%s%s" % (basedir, relative_path, file_name)) if path.exists: return struct(file_name=file_name, path=str(path.realpath)) + auto_configure_fail("Cannot find cuda library %s" % file_name) - elif cpu_value == "Windows": - path = repository_ctx.path("%s/lib/x64/%s" % (basedir, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - path = repository_ctx.path("%s/lib/%s" % (basedir, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - path = repository_ctx.path("%s/%s" % (basedir, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) +def _find_cupti_header_dir(repository_ctx, cuda_config): + """Returns the path to the directory containing cupti.h - auto_configure_fail("Cannot find cuda library %s" % file_name) + On most systems, the cupti library is not installed in the same directory as + the other CUDA libraries but rather in a special extras/CUPTI directory. + + Args: + repository_ctx: The repository context. + cuda_config: The CUDA config as returned by _get_cuda_config + + Returns: + The path of the directory containing the cupti header. + """ + cuda_toolkit_path = cuda_config.cuda_toolkit_path + for relative_path in CUPTI_HEADER_PATHS: + if repository_ctx.path("%s/%scupti.h" % (cuda_toolkit_path, relative_path)).exists: + return ("%s/%s" % (cuda_toolkit_path, relative_path))[:-1] + auto_configure_fail("Cannot find cupti.h under %s" % cuda_toolkit_path) def _find_cupti_lib(repository_ctx, cuda_config): @@ -557,35 +624,13 @@ def _find_cupti_lib(repository_ctx, cuda_config): """ file_name = _lib_name("cupti", cuda_config.cpu_value, cuda_config.cuda_version) - if cuda_config.cpu_value == "Linux": - path = repository_ctx.path( - "%s/extras/CUPTI/lib64/%s" % (cuda_config.cuda_toolkit_path, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - - path = repository_ctx.path( - "%s/lib/x86_64-linux-gnu/%s" % (cuda_config.cuda_toolkit_path, - file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - - elif cuda_config.cpu_value == "Windows": + cuda_toolkit_path = cuda_config.cuda_toolkit_path + for relative_path in CUPTI_LIB_PATHS: path = repository_ctx.path( - "%s/extras/CUPTI/libx64/%s" % - (cuda_config.cuda_toolkit_path, file_name)) + "%s/%s%s" % (cuda_toolkit_path, relative_path, file_name)) if path.exists: return struct(file_name=file_name, path=str(path.realpath)) - path = repository_ctx.path( - "%s/extras/CUPTI/lib/%s" % (cuda_config.cuda_toolkit_path, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - - path = repository_ctx.path( - "%s/lib/%s" % (cuda_config.cuda_toolkit_path, file_name)) - if path.exists: - return struct(file_name=file_name, path=str(path.realpath)) - auto_configure_fail("Cannot find cupti library %s" % file_name) def _find_libs(repository_ctx, cuda_config): @@ -626,6 +671,23 @@ def _find_libs(repository_ctx, cuda_config): } +def _find_cuda_include_path(repository_ctx, cuda_config): + """Returns the path to the directory containing cuda.h + + Args: + repository_ctx: The repository context. + cuda_config: The CUDA config as returned by _get_cuda_config + + Returns: + The path of the directory containing the CUDA headers. + """ + cuda_toolkit_path = cuda_config.cuda_toolkit_path + for relative_path in CUDA_INCLUDE_PATHS: + if repository_ctx.path("%s/%scuda.h" % (cuda_toolkit_path, relative_path)).exists: + return ("%s/%s" % (cuda_toolkit_path, relative_path))[:-1] + auto_configure_fail("Cannot find cuda.h under %s" % cuda_toolkit_path) + + def _find_cudnn_header_dir(repository_ctx, cudnn_install_basedir): """Returns the path to the directory containing cudnn.h @@ -637,15 +699,31 @@ def _find_cudnn_header_dir(repository_ctx, cudnn_install_basedir): Returns: The path of the directory containing the cudnn header. """ - if repository_ctx.path(cudnn_install_basedir + "/cudnn.h").exists: - return cudnn_install_basedir - if repository_ctx.path(cudnn_install_basedir + "/include/cudnn.h").exists: - return cudnn_install_basedir + "/include" + for relative_path in CUDA_INCLUDE_PATHS: + if repository_ctx.path("%s/%scudnn.h" % (cudnn_install_basedir, relative_path)).exists: + return ("%s/%s" % (cudnn_install_basedir, relative_path))[:-1] if repository_ctx.path("/usr/include/cudnn.h").exists: return "/usr/include" auto_configure_fail("Cannot find cudnn.h under %s" % cudnn_install_basedir) +def _find_nvvm_libdevice_dir(repository_ctx, cuda_config): + """Returns the path to the directory containing libdevice in bitcode format. + + Args: + repository_ctx: The repository context. + cuda_config: The CUDA config as returned by _get_cuda_config + + Returns: + The path of the directory containing the CUDA headers. + """ + cuda_toolkit_path = cuda_config.cuda_toolkit_path + for relative_path in NVVM_LIBDEVICE_PATHS: + if repository_ctx.path("%s/%slibdevice.10.bc" % (cuda_toolkit_path, relative_path)).exists: + return ("%s/%s" % (cuda_toolkit_path, relative_path))[:-1] + auto_configure_fail("Cannot find libdevice.10.bc under %s" % cuda_toolkit_path) + + def _cudart_static_linkopt(cpu_value): """Returns additional platform-specific linkopts for cudart.""" return "" if cpu_value == "Darwin" else "\"-lrt\"," @@ -826,7 +904,7 @@ def symlink_genrule_for_dir(repository_ctx, src_dir, dest_dir, genrule_name, if src_dir != None: src_dir = _norm_path(src_dir) dest_dir = _norm_path(dest_dir) - files = _read_dir(repository_ctx, src_dir) + files = '\n'.join(sorted(_read_dir(repository_ctx, src_dir).splitlines())) # Create a list with the src_dir stripped to use for outputs. dest_files = files.replace(src_dir, '').splitlines() src_files = files.splitlines() @@ -916,21 +994,22 @@ def _create_local_cuda_repository(repository_ctx): """Creates the repository containing files set up to build with CUDA.""" cuda_config = _get_cuda_config(repository_ctx) + cuda_include_path = _find_cuda_include_path(repository_ctx, cuda_config) cudnn_header_dir = _find_cudnn_header_dir(repository_ctx, cuda_config.cudnn_install_basedir) + cupti_header_dir = _find_cupti_header_dir(repository_ctx, cuda_config) + nvvm_libdevice_dir = _find_nvvm_libdevice_dir(repository_ctx, cuda_config) # Set up symbolic links for the cuda toolkit by creating genrules to do # symlinking. We create one genrule for each directory we want to track under # cuda_toolkit_path cuda_toolkit_path = cuda_config.cuda_toolkit_path - cuda_include_path = cuda_toolkit_path + "/include" genrules = [symlink_genrule_for_dir(repository_ctx, cuda_include_path, "cuda/include", "cuda-include")] genrules.append(symlink_genrule_for_dir(repository_ctx, - cuda_toolkit_path + "/nvvm", "cuda/nvvm", "cuda-nvvm")) + nvvm_libdevice_dir, "cuda/nvvm/libdevice", "cuda-nvvm")) genrules.append(symlink_genrule_for_dir(repository_ctx, - cuda_toolkit_path + "/extras/CUPTI/include", - "cuda/extras/CUPTI/include", "cuda-extras")) + cupti_header_dir, "cuda/extras/CUPTI/include", "cuda-extras")) cuda_libs = _find_libs(repository_ctx, cuda_config) cuda_lib_src = [] @@ -1077,6 +1156,7 @@ cuda_configure = repository_rule( _TF_CUDNN_VERSION, _TF_CUDA_COMPUTE_CAPABILITIES, _TF_CUDA_CONFIG_REPO, + "NVVMIR_LIBRARY_DIR", ], ) diff --git a/third_party/jpeg/jpeg.BUILD b/third_party/jpeg/jpeg.BUILD index ca2d38d6878cee81b29e949175c5133f492cf06b..4418ac32fc4b08713ff1d1f0d78042803153c886 100644 --- a/third_party/jpeg/jpeg.BUILD +++ b/third_party/jpeg/jpeg.BUILD @@ -145,9 +145,9 @@ cc_library( "jpeglib.h", "jsimd.h", "jsimddct.h", - "simd/jsimd.h", "simd/jccolor-altivec.c", "simd/jcgray-altivec.c", + "simd/jcsample.h", "simd/jcsample-altivec.c", "simd/jdcolor-altivec.c", "simd/jdmerge-altivec.c", @@ -157,15 +157,15 @@ cc_library( "simd/jidctfst-altivec.c", "simd/jidctint-altivec.c", "simd/jquanti-altivec.c", - "simd/jsimd_powerpc.c", + "simd/jsimd.h", "simd/jsimd_altivec.h", - "simd/jcsample.h", + "simd/jsimd_powerpc.c", ], hdrs = [ - "simd/jdmrgext-altivec.c", # should have been named .inc - "simd/jccolext-altivec.c", # should have been named .inc - "simd/jcgryext-altivec.c", # should have been named .inc - "simd/jdcolext-altivec.c", # should have been named .inc + "simd/jccolext-altivec.c", # should have been named .inc + "simd/jcgryext-altivec.c", # should have been named .inc + "simd/jdcolext-altivec.c", # should have been named .inc + "simd/jdmrgext-altivec.c", # should have been named .inc ], copts = libjpegturbo_copts, nocopts = libjpegturbo_nocopts, @@ -526,12 +526,12 @@ config_setting( config_setting( name = "armeabi-v7a", - values = {"android_cpu": "armeabi-v7a"}, + values = {"cpu": "armeabi-v7a"}, ) config_setting( name = "arm64-v8a", - values = {"android_cpu": "arm64-v8a"}, + values = {"cpu": "arm64-v8a"}, ) config_setting( @@ -545,7 +545,6 @@ config_setting( ) config_setting( - name = "linux_ppc64le", - values = {"cpu": "ppc"}, - + name = "linux_ppc64le", + values = {"cpu": "ppc"}, ) diff --git a/third_party/kafka/BUILD b/third_party/kafka/BUILD index a61a9e1f6c2b29ad3b992e810c0cab463dfd7feb..a839ca717e695f35fac684b510f0a022010e0710 100644 --- a/third_party/kafka/BUILD +++ b/third_party/kafka/BUILD @@ -130,12 +130,16 @@ cc_library( ], hdrs = [ "config.h", + "src-cpp/rdkafkacpp.h", + "src-cpp/rdkafkacpp_int.h", + "src/lz4.c", + "src/snappy_compat.h", ], - defines = [ + copts = [ + "-Iexternal/kafka/src", + "-Iexternal/kafka/src-cpp", ], - includes = [ - "src", - "src-cpp", + defines = [ ], linkopts = [ "-lpthread", @@ -143,5 +147,6 @@ cc_library( visibility = ["//visibility:public"], deps = [ "@boringssl//:ssl", + "@zlib_archive//:zlib", ], ) diff --git a/third_party/llvm/llvm.BUILD b/third_party/llvm/llvm.BUILD index a9e1341a03c2e725e96bd7c8cbd7b09853bb8af4..28293a36593d8fa67a2d85631a0769e03d508354 100644 --- a/third_party/llvm/llvm.BUILD +++ b/third_party/llvm/llvm.BUILD @@ -1024,6 +1024,7 @@ cc_library( deps = [ ":arm_desc", ":arm_info", + ":arm_utils", ":config", ":mc_disassembler", ":support", diff --git a/third_party/pcre.BUILD b/third_party/pcre.BUILD index e2cdec40295d369548ff26e3493b5d2300041916..3a8e7a10b43debb5eeca690a64d5795de998a3ac 100644 --- a/third_party/pcre.BUILD +++ b/third_party/pcre.BUILD @@ -1,6 +1,6 @@ licenses(["notice"]) # BSD -exports_files(["COPYING"]) +exports_files(["LICENCE"]) cc_library( name = "pcre", diff --git a/third_party/py/BUILD.tpl b/third_party/py/BUILD.tpl index de06ad5f27e7c08aade4a8f51ab60ba52d012b7b..1dd8ab433a37a127b98ae7069bffcbfd4f6d8bd1 100644 --- a/third_party/py/BUILD.tpl +++ b/third_party/py/BUILD.tpl @@ -2,20 +2,26 @@ licenses(["restricted"]) package(default_visibility = ["//visibility:public"]) +# To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib +# See https://docs.python.org/3/extending/windows.html +cc_import( + name = "python_lib", + interface_library = select({ + ":windows": ":python_import_lib", + # A placeholder for Unix platforms which makes --no_build happy. + "//conditions:default": "not-existing.lib", + }), + system_provided = 1, +) + cc_library( name = "python_headers", hdrs = [":python_include"], - data = select({ - ":windows": [":python_import_lib"], + deps = select({ + ":windows": [":python_lib"], "//conditions:default": [], }), includes = ["python_include"], - linkopts = select({ - # TODO(pcloudy): Ideally, this should just go into deps after resolving - # https://github.com/bazelbuild/bazel/issues/3237, - ":windows": ["$(locations :python_import_lib)"], - "//conditions:default": [], - }), ) cc_library( diff --git a/third_party/py/python_configure.bzl b/third_party/py/python_configure.bzl index c16eb3a12a86f3c2eb3813f5c8c7631fec8e97c6..954f21f5f8fe8029c869f8870464a750cfc8a3db 100644 --- a/third_party/py/python_configure.bzl +++ b/third_party/py/python_configure.bzl @@ -118,7 +118,7 @@ def _symlink_genrule_for_dir(repository_ctx, src_dir, dest_dir, genrule_name, if src_dir != None: src_dir = _norm_path(src_dir) dest_dir = _norm_path(dest_dir) - files = _read_dir(repository_ctx, src_dir) + files = '\n'.join(sorted(_read_dir(repository_ctx, src_dir).splitlines())) # Create a list with the src_dir stripped to use for outputs. dest_files = files.replace(src_dir, '').splitlines() src_files = files.splitlines() diff --git a/third_party/repo.bzl b/third_party/repo.bzl index 11e9c842d2f0c8deb123d7b13d85865b089d73d7..aa178fa8cab92d9d299e5ed09927d8572816a0af 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -27,7 +27,7 @@ def _wrap_bash_cmd(ctx, cmd): bazel_sh = _get_env_var(ctx, "BAZEL_SH") if not bazel_sh: fail("BAZEL_SH environment variable is not set") - cmd = [bazel_sh, "-c", " ".join(cmd)] + cmd = [bazel_sh, "-l", "-c", " ".join(cmd)] return cmd def _get_env_var(ctx, name): diff --git a/third_party/tensorrt/BUILD.tpl b/third_party/tensorrt/BUILD.tpl index feaeb0bea63a982139d34f7c5f37a3f4c66d2af5..57682e8735013544d76b14fc2c41dfff3d50f691 100644 --- a/third_party/tensorrt/BUILD.tpl +++ b/third_party/tensorrt/BUILD.tpl @@ -3,6 +3,8 @@ licenses(["notice"]) +exports_files(["LICENSE"]) + load("@local_config_cuda//cuda:build_defs.bzl", "cuda_default_copts") package(default_visibility = ["//visibility:public"]) @@ -32,36 +34,6 @@ cc_library( visibility = ["//visibility:public"], ) -cc_library( - name = "nv_infer_plugin", - srcs = [%{nv_infer_plugin}], - data = [%{nv_infer_plugin}], - includes = [ - "include", - ], - copts= cuda_default_copts(), - deps = [ - "@local_config_cuda//cuda:cuda", - ":nv_infer", - ":tensorrt_headers", - ], - linkstatic = 1, - visibility = ["//visibility:public"], -) - -cc_library( - name = "nv_parsers", - srcs = [%{nv_parsers}], - data = [%{nv_parsers}], - includes = [ - "include", - ], - copts= cuda_default_copts(), - deps = [ - ":tensorrt_headers", - ], - linkstatic = 1, - visibility = ["//visibility:public"], -) %{tensorrt_genrules} + diff --git a/third_party/tensorrt/LICENSE b/third_party/tensorrt/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..146d9b765c5db44c2f5bea8fa5010eef5ec0c68f --- /dev/null +++ b/third_party/tensorrt/LICENSE @@ -0,0 +1,203 @@ +Copyright 2018 The TensorFlow Authors. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2018, The TensorFlow Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/third_party/tensorrt/tensorrt_configure.bzl b/third_party/tensorrt/tensorrt_configure.bzl index 8aa0f28f39d4dd8e24d5f162bf6535edbb030ee6..9b946505a615372aa7de317c8ee390a2cd4b60e9 100644 --- a/third_party/tensorrt/tensorrt_configure.bzl +++ b/third_party/tensorrt/tensorrt_configure.bzl @@ -19,11 +19,8 @@ load( _TENSORRT_INSTALL_PATH = "TENSORRT_INSTALL_PATH" _TF_TENSORRT_VERSION = "TF_TENSORRT_VERSION" -_TF_TENSORRT_LIBS = ["nvinfer", "nvinfer_plugin", "nvparsers"] -_TF_TENSORRT_HEADERS = [ - "NvInfer.h", "NvInferPlugin.h", "NvCaffeParser.h", "NvUffParser.h", - "NvUtils.h" -] +_TF_TENSORRT_LIBS = ["nvinfer"] +_TF_TENSORRT_HEADERS = ["NvInfer.h", "NvUtils.h"] _DEFINE_TENSORRT_SONAME_MAJOR = "#define NV_TENSORRT_SONAME_MAJOR" _DEFINE_TENSORRT_SONAME_MINOR = "#define NV_TENSORRT_SONAME_MINOR" @@ -60,6 +57,10 @@ def _find_trt_header_dir(repository_ctx, trt_install_path): path = "/usr/include/x86_64-linux-gnu" if _headers_exist(repository_ctx, path): return path + if trt_install_path == "/usr/lib/aarch64-linux-gnu": + path = "/usr/include/aarch64-linux-gnu" + if _headers_exist(repository_ctx, path): + return path path = str(repository_ctx.path("%s/../include" % trt_install_path).realpath) if _headers_exist(repository_ctx, path): return path diff --git a/third_party/termcolor.BUILD b/third_party/termcolor.BUILD index 6000e3289deff8183193883a9b796da9384365b8..655d7cb85e584027d12014c53718a15e2522b4ae 100644 --- a/third_party/termcolor.BUILD +++ b/third_party/termcolor.BUILD @@ -3,7 +3,7 @@ licenses(["notice"]) # MIT -exports_files(["LICENSE"]) +exports_files(["COPYING.txt"]) py_library( name = "termcolor", diff --git a/third_party/toolchains/gpus/crosstool/BUILD b/third_party/toolchains/gpus/crosstool/BUILD index a8c6b0f0291363f3a7576a70e78b3428fb984957..1f9065007ca884a46bfa391d1ee8a8f0333da235 100644 --- a/third_party/toolchains/gpus/crosstool/BUILD +++ b/third_party/toolchains/gpus/crosstool/BUILD @@ -50,3 +50,8 @@ filegroup( name = "empty", srcs = [], ) + +filegroup( + name = "crosstool_wrapper_driver_is_not_gcc", + srcs = ["clang/bin/crosstool_wrapper_driver_is_not_gcc"], +) diff --git a/third_party/toolchains/gpus/crosstool/CROSSTOOL b/third_party/toolchains/gpus/crosstool/CROSSTOOL index a47e0c7cd74edcea777d76854c2d7e97d69897fa..d6ee7e38c414dd59b76c7b2b4c95c55831bb30a8 100644 --- a/third_party/toolchains/gpus/crosstool/CROSSTOOL +++ b/third_party/toolchains/gpus/crosstool/CROSSTOOL @@ -53,6 +53,7 @@ toolchain { flag_set { action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag: "-lstdc++" } @@ -79,6 +80,7 @@ toolchain { name: "alwayslink" flag_set { action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" action: "c++-link-executable" flag_group { flag: "-Wl,-no-as-needed" @@ -120,6 +122,7 @@ toolchain { } flag_set { action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag: "-Wl,-z,relro,-z,now" } @@ -141,8 +144,8 @@ toolchain { flag_group { # All warnings are enabled. Maybe enable -Werror as well? flag: "-Wall" - # TODO(ngiraldo): Some parts of the codebase set -Werror and hit this - # warning, so switch it off for now. + # Some parts of the codebase set -Werror and hit this warning, so + # switch it off for now. flag: "-Wno-invalid-partial-specialization" } } @@ -165,6 +168,7 @@ toolchain { flag_set { action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { # Stamp the binary with a unique identifier. flag: "-Wl,--build-id=md5" @@ -180,6 +184,7 @@ toolchain { action: "c++-compile" action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag:"-no-canonical-prefixes" } @@ -203,6 +208,7 @@ toolchain { flag_set { action: "c++-link-executable" action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" flag_group { flag: "-B/usr/bin/" } @@ -250,6 +256,7 @@ toolchain { } flag_set { action: "c++-link-dynamic-library" + action: "c++-link-nodeps-dynamic-library" action: "c++-link-executable" flag_group { flag: "-Wl,--gc-sections" @@ -296,7 +303,7 @@ toolchain { cxx_builtin_include_directory: "/usr/include/x86_64-linux-gnu/c++/5.4.0" cxx_builtin_include_directory: "/usr/include/c++/5.4.0/backward" cxx_builtin_include_directory: "/usr/local/include" - cxx_builtin_include_directory: "/usr/local/lib/clang/6.0.0/include" + cxx_builtin_include_directory: "/usr/local/lib/clang/7.0.0/include" cxx_builtin_include_directory: "/usr/include/x86_64-linux-gnu" cxx_builtin_include_directory: "/usr/include" } diff --git a/third_party/toolchains/gpus/cuda/BUILD b/third_party/toolchains/gpus/cuda/BUILD index 39136de99c901d6d6a9dafefe3163972511ec122..4cb83809383afa52d5a1d98777f8e5bb2d266286 100644 --- a/third_party/toolchains/gpus/cuda/BUILD +++ b/third_party/toolchains/gpus/cuda/BUILD @@ -51,6 +51,7 @@ cc_library( includes = [ ".", "cuda/include", + "cuda/include/crt", ], visibility = ["//visibility:public"], ) @@ -84,8 +85,8 @@ cc_library( cc_library( name = "cudart", - srcs = ["cuda/lib/libcudart.so.8.0"], - data = ["cuda/lib/libcudart.so.8.0"], + srcs = ["cuda/lib/libcudart.so.9.0"], + data = ["cuda/lib/libcudart.so.9.0"], includes = [ ".", "cuda/include", @@ -96,8 +97,8 @@ cc_library( cc_library( name = "cublas", - srcs = ["cuda/lib/libcublas.so.8.0"], - data = ["cuda/lib/libcublas.so.8.0"], + srcs = ["cuda/lib/libcublas.so.9.0"], + data = ["cuda/lib/libcublas.so.9.0"], includes = [ ".", "cuda/include", @@ -108,8 +109,8 @@ cc_library( cc_library( name = "cusolver", - srcs = ["cuda/lib/libcusolver.so.8.0"], - data = ["cuda/lib/libcusolver.so.8.0"], + srcs = ["cuda/lib/libcusolver.so.9.0"], + data = ["cuda/lib/libcusolver.so.9.0"], includes = [ ".", "cuda/include", @@ -121,8 +122,8 @@ cc_library( cc_library( name = "cudnn", - srcs = ["cuda/lib/libcudnn.so.6"], - data = ["cuda/lib/libcudnn.so.6"], + srcs = ["cuda/lib/libcudnn.so.7"], + data = ["cuda/lib/libcudnn.so.7"], includes = [ ".", "cuda/include", @@ -133,8 +134,8 @@ cc_library( cc_library( name = "cufft", - srcs = ["cuda/lib/libcufft.so.8.0"], - data = ["cuda/lib/libcufft.so.8.0"], + srcs = ["cuda/lib/libcufft.so.9.0"], + data = ["cuda/lib/libcufft.so.9.0"], includes = [ ".", "cuda/include", @@ -145,8 +146,8 @@ cc_library( cc_library( name = "curand", - srcs = ["cuda/lib/libcurand.so.8.0"], - data = ["cuda/lib/libcurand.so.8.0"], + srcs = ["cuda/lib/libcurand.so.9.0"], + data = ["cuda/lib/libcurand.so.9.0"], includes = [ ".", "cuda/include", @@ -183,7 +184,7 @@ cc_library( cc_library( name = "cupti_dsos", - data = ["cuda/lib/libcupti.so.8.0"], + data = ["cuda/lib/libcupti.so.9.0"], includes = [ ".", "cuda/include", @@ -200,1063 +201,990 @@ cc_library( genrule( name = "cuda-include", outs = [ - "cuda/include/math_functions.hpp", - "cuda/include/cufft.h", - "cuda/include/nvgraph.h", - "cuda/include/curand_normal.h", - "cuda/include/curand_uniform.h", - "cuda/include/nppi_data_exchange_and_initialization.h", - "cuda/include/cuda_gl_interop.h", - "cuda/include/nppi_compression_functions.h", - "cuda/include/npp.h", + "cuda/include/CL/cl.h", + "cuda/include/CL/cl.hpp", + "cuda/include/CL/cl_egl.h", + "cuda/include/CL/cl_ext.h", + "cuda/include/CL/cl_gl.h", + "cuda/include/CL/cl_gl_ext.h", + "cuda/include/CL/cl_platform.h", + "cuda/include/CL/opencl.h", + "cuda/include/builtin_types.h", + "cuda/include/channel_descriptor.h", + "cuda/include/common_functions.h", + "cuda/include/cooperative_groups.h", + "cuda/include/cooperative_groups_helpers.h", + "cuda/include/crt/common_functions.h", + "cuda/include/crt/device_double_functions.h", + "cuda/include/crt/device_double_functions.hpp", + "cuda/include/crt/device_functions.h", + "cuda/include/crt/device_functions.hpp", + "cuda/include/crt/func_macro.h", + "cuda/include/crt/host_config.h", + "cuda/include/crt/host_defines.h", + "cuda/include/crt/host_runtime.h", + "cuda/include/crt/math_functions.h", + "cuda/include/crt/math_functions.hpp", + "cuda/include/crt/mma.h", + "cuda/include/crt/mma.hpp", + "cuda/include/crt/nvfunctional", + "cuda/include/crt/sm_70_rt.h", + "cuda/include/crt/sm_70_rt.hpp", + "cuda/include/crt/storage_class.h", + "cuda/include/cuComplex.h", + "cuda/include/cublas.h", + "cuda/include/cublasXt.h", + "cuda/include/cublas_api.h", + "cuda/include/cublas_v2.h", "cuda/include/cuda.h", - "cuda/include/nppi_statistics_functions.h", - "cuda/include/vector_functions.hpp", - "cuda/include/sm_32_intrinsics.hpp", - "cuda/include/sm_32_intrinsics.h", - "cuda/include/curand_discrete.h", + "cuda/include/cudaEGL.h", + "cuda/include/cudaGL.h", + "cuda/include/cudaProfiler.h", + "cuda/include/cudaVDPAU.h", + "cuda/include/cuda_device_runtime_api.h", + "cuda/include/cuda_fp16.h", + "cuda/include/cuda_fp16.hpp", + "cuda/include/cuda_gl_interop.h", + "cuda/include/cuda_occupancy.h", + "cuda/include/cuda_profiler_api.h", "cuda/include/cuda_runtime.h", + "cuda/include/cuda_runtime_api.h", + "cuda/include/cuda_surface_types.h", + "cuda/include/cuda_texture_types.h", + "cuda/include/cuda_vdpau_interop.h", + "cuda/include/cudalibxt.h", + "cuda/include/cudnn.h", + "cuda/include/cufft.h", "cuda/include/cufftXt.h", - "cuda/include/sm_61_intrinsics.h", - "cuda/include/texture_fetch_functions.h", + "cuda/include/cufftw.h", + "cuda/include/curand.h", + "cuda/include/curand_discrete.h", + "cuda/include/curand_discrete2.h", + "cuda/include/curand_globals.h", + "cuda/include/curand_kernel.h", + "cuda/include/curand_lognormal.h", "cuda/include/curand_mrg32k3a.h", - "cuda/include/host_defines.h", - "cuda/include/common_functions.h", - "cuda/include/nppi_support_functions.h", - "cuda/include/nppi_linear_transforms.h", - "cuda/include/device_double_functions.hpp", - "cuda/include/math_constants.h", - "cuda/include/nvToolsExtSync.h", - "cuda/include/npps_initialization.h", + "cuda/include/curand_mtgp32.h", + "cuda/include/curand_mtgp32_host.h", + "cuda/include/curand_mtgp32_kernel.h", + "cuda/include/curand_mtgp32dc_p_11213.h", + "cuda/include/curand_normal.h", + "cuda/include/curand_normal_static.h", + "cuda/include/curand_philox4x32_x.h", + "cuda/include/curand_poisson.h", + "cuda/include/curand_precalc.h", + "cuda/include/curand_uniform.h", + "cuda/include/cusolverDn.h", + "cuda/include/cusolverRf.h", + "cuda/include/cusolverSp.h", "cuda/include/cusolverSp_LOWLEVEL_PREVIEW.h", - "cuda/include/texture_indirect_functions.hpp", - "cuda/include/cudaProfiler.h", - "cuda/include/npps_filtering_functions.h", + "cuda/include/cusolver_common.h", + "cuda/include/cusparse.h", "cuda/include/cusparse_v2.h", - "cuda/include/nppi.h", - "cuda/include/surface_indirect_functions.h", - "cuda/include/sm_30_intrinsics.h", + "cuda/include/device_atomic_functions.h", + "cuda/include/device_atomic_functions.hpp", "cuda/include/device_double_functions.h", - "cuda/include/sm_35_intrinsics.h", - "cuda/include/cusolverSp.h", - "cuda/include/library_types.h", - "cuda/include/surface_indirect_functions.hpp", - "cuda/include/cudalibxt.h", - "cuda/include/channel_descriptor.h", + "cuda/include/device_double_functions.hpp", + "cuda/include/device_functions.h", + "cuda/include/device_functions.hpp", "cuda/include/device_functions_decls.h", - "cuda/include/curand_kernel.h", - "cuda/include/curand_mtgp32_host.h", - "cuda/include/nvToolsExtCuda.h", - "cuda/include/nvToolsExt.h", - "cuda/include/cuComplex.h", - "cuda/include/sm_32_atomic_functions.h", - "cuda/include/texture_indirect_functions.h", - "cuda/include/sm_32_atomic_functions.hpp", - "cuda/include/sm_20_intrinsics.hpp", "cuda/include/device_launch_parameters.h", - "cuda/include/curand_mtgp32.h", - "cuda/include/texture_fetch_functions.hpp", - "cuda/include/cuda_occupancy.h", - "cuda/include/CL/opencl.h", - "cuda/include/CL/cl_platform.h", - "cuda/include/CL/cl_egl.h", - "cuda/include/CL/cl_gl.h", - "cuda/include/CL/cl.h", - "cuda/include/CL/cl_gl_ext.h", - "cuda/include/CL/cl_ext.h", - "cuda/include/CL/cl.hpp", + "cuda/include/device_types.h", + "cuda/include/driver_functions.h", + "cuda/include/driver_types.h", + "cuda/include/dynlink_cuda.h", + "cuda/include/dynlink_cuda_cuda.h", + "cuda/include/dynlink_cuviddec.h", + "cuda/include/dynlink_nvcuvid.h", + "cuda/include/fatBinaryCtl.h", + "cuda/include/fatbinary.h", "cuda/include/host_config.h", - "cuda/include/cuda_surface_types.h", + "cuda/include/host_defines.h", + "cuda/include/library_types.h", + "cuda/include/math_constants.h", "cuda/include/math_functions.h", + "cuda/include/math_functions.hpp", + "cuda/include/math_functions_dbl_ptx3.h", + "cuda/include/math_functions_dbl_ptx3.hpp", + "cuda/include/mma.h", + "cuda/include/npp.h", + "cuda/include/nppcore.h", + "cuda/include/nppdefs.h", + "cuda/include/nppi.h", + "cuda/include/nppi_arithmetic_and_logical_operations.h", + "cuda/include/nppi_color_conversion.h", + "cuda/include/nppi_compression_functions.h", + "cuda/include/nppi_computer_vision.h", + "cuda/include/nppi_data_exchange_and_initialization.h", + "cuda/include/nppi_filtering_functions.h", + "cuda/include/nppi_geometry_transforms.h", + "cuda/include/nppi_linear_transforms.h", + "cuda/include/nppi_morphological_operations.h", + "cuda/include/nppi_statistics_functions.h", + "cuda/include/nppi_support_functions.h", + "cuda/include/nppi_threshold_and_compare_operations.h", + "cuda/include/npps.h", + "cuda/include/npps_arithmetic_and_logical_operations.h", + "cuda/include/npps_conversion_functions.h", + "cuda/include/npps_filtering_functions.h", + "cuda/include/npps_initialization.h", + "cuda/include/npps_statistics_functions.h", + "cuda/include/npps_support_functions.h", + "cuda/include/nppversion.h", + "cuda/include/nvToolsExt.h", + "cuda/include/nvToolsExtCuda.h", + "cuda/include/nvToolsExtCudaRt.h", "cuda/include/nvToolsExtMeta.h", + "cuda/include/nvToolsExtSync.h", + "cuda/include/nvblas.h", + "cuda/include/nvfunctional", + "cuda/include/nvgraph.h", + "cuda/include/nvml.h", + "cuda/include/nvrtc.h", + "cuda/include/sm_20_atomic_functions.h", "cuda/include/sm_20_atomic_functions.hpp", - "cuda/include/device_functions.h", - "cuda/include/device_types.h", - "cuda/include/npps_conversion_functions.h", - "cuda/include/curand_precalc.h", - "cuda/include/cusolverRf.h", + "cuda/include/sm_20_intrinsics.h", + "cuda/include/sm_20_intrinsics.hpp", + "cuda/include/sm_30_intrinsics.h", + "cuda/include/sm_30_intrinsics.hpp", + "cuda/include/sm_32_atomic_functions.h", + "cuda/include/sm_32_atomic_functions.hpp", + "cuda/include/sm_32_intrinsics.h", + "cuda/include/sm_32_intrinsics.hpp", + "cuda/include/sm_35_atomic_functions.h", + "cuda/include/sm_35_intrinsics.h", + "cuda/include/sm_60_atomic_functions.h", "cuda/include/sm_60_atomic_functions.hpp", - "cuda/include/cuviddec.h", - "cuda/include/curand_discrete2.h", - "cuda/include/device_functions.hpp", - "cuda/include/thrust/transform_scan.h", - "cuda/include/thrust/system_error.h", - "cuda/include/thrust/device_malloc.h", - "cuda/include/thrust/partition.h", - "cuda/include/thrust/unique.h", - "cuda/include/thrust/device_delete.h", - "cuda/include/thrust/execution_policy.h", + "cuda/include/sm_61_intrinsics.h", + "cuda/include/sm_61_intrinsics.hpp", + "cuda/include/sobol_direction_vectors.h", + "cuda/include/surface_functions.h", + "cuda/include/surface_functions.hpp", + "cuda/include/surface_indirect_functions.h", + "cuda/include/surface_indirect_functions.hpp", + "cuda/include/surface_types.h", + "cuda/include/texture_fetch_functions.h", + "cuda/include/texture_fetch_functions.hpp", + "cuda/include/texture_indirect_functions.h", + "cuda/include/texture_indirect_functions.hpp", + "cuda/include/texture_types.h", "cuda/include/thrust/adjacent_difference.h", - "cuda/include/thrust/sequence.h", - "cuda/include/thrust/merge.h", - "cuda/include/thrust/device_new.h", - "cuda/include/thrust/transform_reduce.h", - "cuda/include/thrust/device_vector.h", - "cuda/include/thrust/gather.h", - "cuda/include/thrust/sort.h", - "cuda/include/thrust/scan.h", - "cuda/include/thrust/detail/temporary_array.h", - "cuda/include/thrust/detail/util/align.h", - "cuda/include/thrust/detail/util/blocking.h", - "cuda/include/thrust/detail/transform.inl", - "cuda/include/thrust/detail/device_vector.inl", + "cuda/include/thrust/advance.h", + "cuda/include/thrust/binary_search.h", + "cuda/include/thrust/complex.h", + "cuda/include/thrust/copy.h", + "cuda/include/thrust/count.h", + "cuda/include/thrust/detail/adjacent_difference.inl", + "cuda/include/thrust/detail/advance.inl", + "cuda/include/thrust/detail/allocator/allocator_traits.h", + "cuda/include/thrust/detail/allocator/allocator_traits.inl", + "cuda/include/thrust/detail/allocator/copy_construct_range.h", + "cuda/include/thrust/detail/allocator/copy_construct_range.inl", + "cuda/include/thrust/detail/allocator/default_construct_range.h", + "cuda/include/thrust/detail/allocator/default_construct_range.inl", + "cuda/include/thrust/detail/allocator/destroy_range.h", + "cuda/include/thrust/detail/allocator/destroy_range.inl", + "cuda/include/thrust/detail/allocator/fill_construct_range.h", + "cuda/include/thrust/detail/allocator/fill_construct_range.inl", + "cuda/include/thrust/detail/allocator/malloc_allocator.h", + "cuda/include/thrust/detail/allocator/malloc_allocator.inl", + "cuda/include/thrust/detail/allocator/no_throw_allocator.h", + "cuda/include/thrust/detail/allocator/tagged_allocator.h", + "cuda/include/thrust/detail/allocator/tagged_allocator.inl", + "cuda/include/thrust/detail/allocator/temporary_allocator.h", + "cuda/include/thrust/detail/allocator/temporary_allocator.inl", "cuda/include/thrust/detail/binary_search.inl", - "cuda/include/thrust/detail/overlapped_copy.h", - "cuda/include/thrust/detail/vector_base.inl", - "cuda/include/thrust/detail/device_reference.inl", - "cuda/include/thrust/detail/functional/actor.h", - "cuda/include/thrust/detail/functional/value.h", - "cuda/include/thrust/detail/functional/operators.h", - "cuda/include/thrust/detail/functional/operators/logical_operators.h", - "cuda/include/thrust/detail/functional/operators/relational_operators.h", - "cuda/include/thrust/detail/functional/operators/assignment_operator.h", - "cuda/include/thrust/detail/functional/operators/bitwise_operators.h", - "cuda/include/thrust/detail/functional/operators/operator_adaptors.h", - "cuda/include/thrust/detail/functional/operators/arithmetic_operators.h", - "cuda/include/thrust/detail/functional/operators/compound_assignment_operators.h", - "cuda/include/thrust/detail/functional/argument.h", - "cuda/include/thrust/detail/functional/placeholder.h", - "cuda/include/thrust/detail/functional/actor.inl", - "cuda/include/thrust/detail/functional/composite.h", - "cuda/include/thrust/detail/static_map.h", - "cuda/include/thrust/detail/type_traits/has_nested_type.h", - "cuda/include/thrust/detail/type_traits/is_call_possible.h", - "cuda/include/thrust/detail/type_traits/function_traits.h", - "cuda/include/thrust/detail/type_traits/pointer_traits.h", - "cuda/include/thrust/detail/type_traits/has_member_function.h", - "cuda/include/thrust/detail/type_traits/algorithm/intermediate_type_from_function_and_iterators.h", - "cuda/include/thrust/detail/type_traits/minimum_type.h", - "cuda/include/thrust/detail/type_traits/has_trivial_assign.h", - "cuda/include/thrust/detail/type_traits/is_metafunction_defined.h", - "cuda/include/thrust/detail/type_traits/iterator/is_discard_iterator.h", - "cuda/include/thrust/detail/type_traits/iterator/is_output_iterator.h", - "cuda/include/thrust/detail/type_traits/result_of_adaptable_function.h", - "cuda/include/thrust/detail/reference.h", - "cuda/include/thrust/detail/inner_product.inl", - "cuda/include/thrust/detail/use_default.h", - "cuda/include/thrust/detail/sequence.inl", - "cuda/include/thrust/detail/sort.inl", - "cuda/include/thrust/detail/equal.inl", - "cuda/include/thrust/detail/execution_policy.h", - "cuda/include/thrust/detail/integer_traits.h", - "cuda/include/thrust/detail/type_traits.h", - "cuda/include/thrust/detail/reverse.inl", - "cuda/include/thrust/detail/tabulate.inl", - "cuda/include/thrust/detail/unique.inl", - "cuda/include/thrust/detail/scatter.inl", - "cuda/include/thrust/detail/set_operations.inl", - "cuda/include/thrust/detail/device_malloc.inl", - "cuda/include/thrust/detail/copy_if.inl", - "cuda/include/thrust/detail/fill.inl", - "cuda/include/thrust/detail/temporary_array.inl", - "cuda/include/thrust/detail/transform_scan.inl", - "cuda/include/thrust/detail/minmax.h", - "cuda/include/thrust/detail/swap.inl", - "cuda/include/thrust/detail/pointer.inl", - "cuda/include/thrust/detail/transform_reduce.inl", - "cuda/include/thrust/detail/config.h", - "cuda/include/thrust/detail/distance.inl", - "cuda/include/thrust/detail/pair.inl", - "cuda/include/thrust/detail/allocator/temporary_allocator.h", - "cuda/include/thrust/detail/allocator/tagged_allocator.h", - "cuda/include/thrust/detail/allocator/destroy_range.inl", - "cuda/include/thrust/detail/allocator/destroy_range.h", - "cuda/include/thrust/detail/allocator/no_throw_allocator.h", - "cuda/include/thrust/detail/allocator/default_construct_range.inl", - "cuda/include/thrust/detail/allocator/fill_construct_range.inl", - "cuda/include/thrust/detail/allocator/tagged_allocator.inl", - "cuda/include/thrust/detail/allocator/malloc_allocator.h", - "cuda/include/thrust/detail/allocator/allocator_traits.h", - "cuda/include/thrust/detail/allocator/copy_construct_range.h", - "cuda/include/thrust/detail/allocator/allocator_traits.inl", - "cuda/include/thrust/detail/allocator/default_construct_range.h", - "cuda/include/thrust/detail/allocator/copy_construct_range.inl", - "cuda/include/thrust/detail/allocator/malloc_allocator.inl", - "cuda/include/thrust/detail/allocator/temporary_allocator.inl", - "cuda/include/thrust/detail/allocator/fill_construct_range.h", - "cuda/include/thrust/detail/temporary_buffer.h", - "cuda/include/thrust/detail/reduce.inl", - "cuda/include/thrust/detail/device_new.inl", - "cuda/include/thrust/detail/pointer.h", - "cuda/include/thrust/detail/for_each.inl", - "cuda/include/thrust/detail/generate.inl", - "cuda/include/thrust/detail/dispatch/is_trivial_copy.h", - "cuda/include/thrust/detail/adjacent_difference.inl", - "cuda/include/thrust/detail/tuple_meta_transform.h", - "cuda/include/thrust/detail/functional.inl", - "cuda/include/thrust/detail/remove.inl", - "cuda/include/thrust/detail/tuple_transform.h", - "cuda/include/thrust/detail/merge.inl", - "cuda/include/thrust/detail/extrema.inl", - "cuda/include/thrust/detail/trivial_sequence.h", - "cuda/include/thrust/detail/vector_base.h", - "cuda/include/thrust/detail/count.inl", - "cuda/include/thrust/detail/uninitialized_copy.inl", - "cuda/include/thrust/detail/function.h", - "cuda/include/thrust/detail/swap_ranges.inl", - "cuda/include/thrust/detail/device_delete.inl", - "cuda/include/thrust/detail/static_assert.h", - "cuda/include/thrust/detail/logical.inl", - "cuda/include/thrust/detail/seq.h", - "cuda/include/thrust/detail/mpl/math.h", - "cuda/include/thrust/detail/mismatch.inl", - "cuda/include/thrust/detail/internal_functional.h", - "cuda/include/thrust/detail/get_iterator_value.h", - "cuda/include/thrust/detail/copy.inl", - "cuda/include/thrust/detail/copy.h", + "cuda/include/thrust/detail/complex/arithmetic.h", + "cuda/include/thrust/detail/complex/c99math.h", + "cuda/include/thrust/detail/complex/catrig.h", "cuda/include/thrust/detail/complex/catrigf.h", - "cuda/include/thrust/detail/complex/cpowf.h", - "cuda/include/thrust/detail/complex/csqrtf.h", + "cuda/include/thrust/detail/complex/ccosh.h", "cuda/include/thrust/detail/complex/ccoshf.h", - "cuda/include/thrust/detail/complex/csinhf.h", + "cuda/include/thrust/detail/complex/cexp.h", + "cuda/include/thrust/detail/complex/cexpf.h", + "cuda/include/thrust/detail/complex/clog.h", "cuda/include/thrust/detail/complex/clogf.h", - "cuda/include/thrust/detail/complex/ccosh.h", - "cuda/include/thrust/detail/complex/arithmetic.h", - "cuda/include/thrust/detail/complex/csqrt.h", - "cuda/include/thrust/detail/complex/cpow.h", "cuda/include/thrust/detail/complex/complex.inl", - "cuda/include/thrust/detail/complex/math_private.h", - "cuda/include/thrust/detail/complex/c99math.h", + "cuda/include/thrust/detail/complex/cpow.h", + "cuda/include/thrust/detail/complex/cpowf.h", "cuda/include/thrust/detail/complex/cproj.h", - "cuda/include/thrust/detail/complex/catrig.h", - "cuda/include/thrust/detail/complex/ctanhf.h", - "cuda/include/thrust/detail/complex/cexpf.h", "cuda/include/thrust/detail/complex/csinh.h", - "cuda/include/thrust/detail/complex/stream.h", + "cuda/include/thrust/detail/complex/csinhf.h", + "cuda/include/thrust/detail/complex/csqrt.h", + "cuda/include/thrust/detail/complex/csqrtf.h", "cuda/include/thrust/detail/complex/ctanh.h", - "cuda/include/thrust/detail/complex/cexp.h", - "cuda/include/thrust/detail/complex/clog.h", - "cuda/include/thrust/detail/range/head_flags.h", - "cuda/include/thrust/detail/range/tail_flags.h", - "cuda/include/thrust/detail/execute_with_allocator.h", - "cuda/include/thrust/detail/integer_math.h", - "cuda/include/thrust/detail/swap.h", - "cuda/include/thrust/detail/uninitialized_fill.inl", - "cuda/include/thrust/detail/scan.inl", - "cuda/include/thrust/detail/gather.inl", - "cuda/include/thrust/detail/reference_forward_declaration.h", - "cuda/include/thrust/detail/numeric_traits.h", - "cuda/include/thrust/detail/reference.inl", - "cuda/include/thrust/detail/cstdint.h", - "cuda/include/thrust/detail/device_free.inl", - "cuda/include/thrust/detail/copy_if.h", - "cuda/include/thrust/detail/partition.inl", - "cuda/include/thrust/detail/find.inl", - "cuda/include/thrust/detail/config/forceinline.h", - "cuda/include/thrust/detail/config/debug.h", - "cuda/include/thrust/detail/config/config.h", - "cuda/include/thrust/detail/config/host_device.h", - "cuda/include/thrust/detail/config/host_system.h", + "cuda/include/thrust/detail/complex/ctanhf.h", + "cuda/include/thrust/detail/complex/math_private.h", + "cuda/include/thrust/detail/complex/stream.h", + "cuda/include/thrust/detail/config.h", "cuda/include/thrust/detail/config/compiler.h", - "cuda/include/thrust/detail/config/device_system.h", "cuda/include/thrust/detail/config/compiler_fence.h", + "cuda/include/thrust/detail/config/config.h", + "cuda/include/thrust/detail/config/debug.h", + "cuda/include/thrust/detail/config/device_system.h", "cuda/include/thrust/detail/config/exec_check_disable.h", - "cuda/include/thrust/detail/config/simple_defines.h", + "cuda/include/thrust/detail/config/forceinline.h", "cuda/include/thrust/detail/config/global_workarounds.h", - "cuda/include/thrust/detail/replace.inl", + "cuda/include/thrust/detail/config/host_device.h", + "cuda/include/thrust/detail/config/host_system.h", + "cuda/include/thrust/detail/config/simple_defines.h", + "cuda/include/thrust/detail/contiguous_storage.h", + "cuda/include/thrust/detail/contiguous_storage.inl", + "cuda/include/thrust/detail/copy.h", + "cuda/include/thrust/detail/copy.inl", + "cuda/include/thrust/detail/copy_if.h", + "cuda/include/thrust/detail/copy_if.inl", + "cuda/include/thrust/detail/count.inl", + "cuda/include/thrust/detail/cstdint.h", + "cuda/include/thrust/detail/device_delete.inl", + "cuda/include/thrust/detail/device_free.inl", + "cuda/include/thrust/detail/device_malloc.inl", + "cuda/include/thrust/detail/device_new.inl", "cuda/include/thrust/detail/device_ptr.inl", - "cuda/include/thrust/detail/tuple.inl", - "cuda/include/thrust/detail/malloc_and_free.h", + "cuda/include/thrust/detail/device_reference.inl", + "cuda/include/thrust/detail/device_vector.inl", + "cuda/include/thrust/detail/dispatch/is_trivial_copy.h", + "cuda/include/thrust/detail/distance.inl", + "cuda/include/thrust/detail/equal.inl", + "cuda/include/thrust/detail/execute_with_allocator.h", + "cuda/include/thrust/detail/execution_policy.h", + "cuda/include/thrust/detail/extrema.inl", + "cuda/include/thrust/detail/fill.inl", + "cuda/include/thrust/detail/find.inl", + "cuda/include/thrust/detail/for_each.inl", + "cuda/include/thrust/detail/function.h", + "cuda/include/thrust/detail/functional.inl", + "cuda/include/thrust/detail/functional/actor.h", + "cuda/include/thrust/detail/functional/actor.inl", + "cuda/include/thrust/detail/functional/argument.h", + "cuda/include/thrust/detail/functional/composite.h", + "cuda/include/thrust/detail/functional/operators.h", + "cuda/include/thrust/detail/functional/operators/arithmetic_operators.h", + "cuda/include/thrust/detail/functional/operators/assignment_operator.h", + "cuda/include/thrust/detail/functional/operators/bitwise_operators.h", + "cuda/include/thrust/detail/functional/operators/compound_assignment_operators.h", + "cuda/include/thrust/detail/functional/operators/logical_operators.h", + "cuda/include/thrust/detail/functional/operators/operator_adaptors.h", + "cuda/include/thrust/detail/functional/operators/relational_operators.h", + "cuda/include/thrust/detail/functional/placeholder.h", + "cuda/include/thrust/detail/functional/value.h", + "cuda/include/thrust/detail/gather.inl", + "cuda/include/thrust/detail/generate.inl", + "cuda/include/thrust/detail/get_iterator_value.h", "cuda/include/thrust/detail/host_vector.inl", + "cuda/include/thrust/detail/inner_product.inl", + "cuda/include/thrust/detail/integer_math.h", + "cuda/include/thrust/detail/integer_traits.h", + "cuda/include/thrust/detail/internal_functional.h", + "cuda/include/thrust/detail/logical.inl", + "cuda/include/thrust/detail/malloc_and_free.h", + "cuda/include/thrust/detail/merge.inl", + "cuda/include/thrust/detail/minmax.h", + "cuda/include/thrust/detail/mismatch.inl", + "cuda/include/thrust/detail/mpl/math.h", + "cuda/include/thrust/detail/numeric_traits.h", + "cuda/include/thrust/detail/overlapped_copy.h", + "cuda/include/thrust/detail/pair.inl", + "cuda/include/thrust/detail/partition.inl", + "cuda/include/thrust/detail/pointer.h", + "cuda/include/thrust/detail/pointer.inl", + "cuda/include/thrust/detail/range/head_flags.h", + "cuda/include/thrust/detail/range/tail_flags.h", "cuda/include/thrust/detail/raw_pointer_cast.h", - "cuda/include/thrust/detail/advance.inl", - "cuda/include/thrust/detail/contiguous_storage.h", "cuda/include/thrust/detail/raw_reference_cast.h", - "cuda/include/thrust/detail/contiguous_storage.inl", - "cuda/include/thrust/reverse.h", - "cuda/include/thrust/device_malloc_allocator.h", - "cuda/include/thrust/scatter.h", - "cuda/include/thrust/pair.h", - "cuda/include/thrust/advance.h", - "cuda/include/thrust/find.h", - "cuda/include/thrust/device_ptr.h", - "cuda/include/thrust/generate.h", - "cuda/include/thrust/uninitialized_fill.h", - "cuda/include/thrust/system/system_error.h", - "cuda/include/thrust/system/detail/bad_alloc.h", - "cuda/include/thrust/system/detail/adl/transform_scan.h", - "cuda/include/thrust/system/detail/adl/unique_by_key.h", - "cuda/include/thrust/system/detail/adl/partition.h", - "cuda/include/thrust/system/detail/adl/unique.h", - "cuda/include/thrust/system/detail/adl/adjacent_difference.h", - "cuda/include/thrust/system/detail/adl/sequence.h", - "cuda/include/thrust/system/detail/adl/merge.h", - "cuda/include/thrust/system/detail/adl/transform_reduce.h", - "cuda/include/thrust/system/detail/adl/gather.h", - "cuda/include/thrust/system/detail/adl/sort.h", - "cuda/include/thrust/system/detail/adl/scan.h", - "cuda/include/thrust/system/detail/adl/temporary_buffer.h", - "cuda/include/thrust/system/detail/adl/scan_by_key.h", - "cuda/include/thrust/system/detail/adl/reverse.h", - "cuda/include/thrust/system/detail/adl/assign_value.h", - "cuda/include/thrust/system/detail/adl/scatter.h", - "cuda/include/thrust/system/detail/adl/find.h", - "cuda/include/thrust/system/detail/adl/generate.h", - "cuda/include/thrust/system/detail/adl/uninitialized_fill.h", - "cuda/include/thrust/system/detail/adl/remove.h", - "cuda/include/thrust/system/detail/adl/tabulate.h", - "cuda/include/thrust/system/detail/adl/for_each.h", - "cuda/include/thrust/system/detail/adl/reduce_by_key.h", - "cuda/include/thrust/system/detail/adl/reduce.h", - "cuda/include/thrust/system/detail/adl/equal.h", - "cuda/include/thrust/system/detail/adl/copy.h", - "cuda/include/thrust/system/detail/adl/swap_ranges.h", - "cuda/include/thrust/system/detail/adl/uninitialized_copy.h", - "cuda/include/thrust/system/detail/adl/binary_search.h", - "cuda/include/thrust/system/detail/adl/set_operations.h", - "cuda/include/thrust/system/detail/adl/mismatch.h", - "cuda/include/thrust/system/detail/adl/extrema.h", - "cuda/include/thrust/system/detail/adl/count.h", - "cuda/include/thrust/system/detail/adl/replace.h", + "cuda/include/thrust/detail/reduce.inl", + "cuda/include/thrust/detail/reference.h", + "cuda/include/thrust/detail/reference.inl", + "cuda/include/thrust/detail/reference_forward_declaration.h", + "cuda/include/thrust/detail/remove.inl", + "cuda/include/thrust/detail/replace.inl", + "cuda/include/thrust/detail/reverse.inl", + "cuda/include/thrust/detail/scan.inl", + "cuda/include/thrust/detail/scatter.inl", + "cuda/include/thrust/detail/seq.h", + "cuda/include/thrust/detail/sequence.inl", + "cuda/include/thrust/detail/set_operations.inl", + "cuda/include/thrust/detail/sort.inl", + "cuda/include/thrust/detail/static_assert.h", + "cuda/include/thrust/detail/static_map.h", + "cuda/include/thrust/detail/swap.h", + "cuda/include/thrust/detail/swap.inl", + "cuda/include/thrust/detail/swap_ranges.inl", + "cuda/include/thrust/detail/tabulate.inl", + "cuda/include/thrust/detail/temporary_array.h", + "cuda/include/thrust/detail/temporary_array.inl", + "cuda/include/thrust/detail/temporary_buffer.h", + "cuda/include/thrust/detail/transform.inl", + "cuda/include/thrust/detail/transform_reduce.inl", + "cuda/include/thrust/detail/transform_scan.inl", + "cuda/include/thrust/detail/trivial_sequence.h", + "cuda/include/thrust/detail/tuple.inl", + "cuda/include/thrust/detail/tuple_meta_transform.h", + "cuda/include/thrust/detail/tuple_transform.h", + "cuda/include/thrust/detail/type_traits.h", + "cuda/include/thrust/detail/type_traits/algorithm/intermediate_type_from_function_and_iterators.h", + "cuda/include/thrust/detail/type_traits/function_traits.h", + "cuda/include/thrust/detail/type_traits/has_member_function.h", + "cuda/include/thrust/detail/type_traits/has_nested_type.h", + "cuda/include/thrust/detail/type_traits/has_trivial_assign.h", + "cuda/include/thrust/detail/type_traits/is_call_possible.h", + "cuda/include/thrust/detail/type_traits/is_metafunction_defined.h", + "cuda/include/thrust/detail/type_traits/iterator/is_discard_iterator.h", + "cuda/include/thrust/detail/type_traits/iterator/is_output_iterator.h", + "cuda/include/thrust/detail/type_traits/minimum_type.h", + "cuda/include/thrust/detail/type_traits/pointer_traits.h", + "cuda/include/thrust/detail/type_traits/result_of_adaptable_function.h", + "cuda/include/thrust/detail/uninitialized_copy.inl", + "cuda/include/thrust/detail/uninitialized_fill.inl", + "cuda/include/thrust/detail/unique.inl", + "cuda/include/thrust/detail/use_default.h", + "cuda/include/thrust/detail/util/align.h", + "cuda/include/thrust/detail/util/blocking.h", + "cuda/include/thrust/detail/vector_base.h", + "cuda/include/thrust/detail/vector_base.inl", + "cuda/include/thrust/device_allocator.h", + "cuda/include/thrust/device_delete.h", + "cuda/include/thrust/device_free.h", + "cuda/include/thrust/device_malloc.h", + "cuda/include/thrust/device_malloc_allocator.h", + "cuda/include/thrust/device_new.h", + "cuda/include/thrust/device_new_allocator.h", + "cuda/include/thrust/device_ptr.h", + "cuda/include/thrust/device_reference.h", + "cuda/include/thrust/device_vector.h", + "cuda/include/thrust/distance.h", + "cuda/include/thrust/equal.h", + "cuda/include/thrust/execution_policy.h", + "cuda/include/thrust/extrema.h", + "cuda/include/thrust/fill.h", + "cuda/include/thrust/find.h", + "cuda/include/thrust/for_each.h", + "cuda/include/thrust/functional.h", + "cuda/include/thrust/gather.h", + "cuda/include/thrust/generate.h", + "cuda/include/thrust/host_vector.h", + "cuda/include/thrust/inner_product.h", + "cuda/include/thrust/iterator/constant_iterator.h", + "cuda/include/thrust/iterator/counting_iterator.h", + "cuda/include/thrust/iterator/detail/any_assign.h", + "cuda/include/thrust/iterator/detail/any_system_tag.h", + "cuda/include/thrust/iterator/detail/constant_iterator_base.h", + "cuda/include/thrust/iterator/detail/counting_iterator.inl", + "cuda/include/thrust/iterator/detail/device_system_tag.h", + "cuda/include/thrust/iterator/detail/discard_iterator_base.h", + "cuda/include/thrust/iterator/detail/distance_from_result.h", + "cuda/include/thrust/iterator/detail/host_system_tag.h", + "cuda/include/thrust/iterator/detail/is_iterator_category.h", + "cuda/include/thrust/iterator/detail/is_trivial_iterator.h", + "cuda/include/thrust/iterator/detail/iterator_adaptor_base.h", + "cuda/include/thrust/iterator/detail/iterator_category_to_system.h", + "cuda/include/thrust/iterator/detail/iterator_category_to_traversal.h", + "cuda/include/thrust/iterator/detail/iterator_category_with_system_and_traversal.h", + "cuda/include/thrust/iterator/detail/iterator_facade_category.h", + "cuda/include/thrust/iterator/detail/iterator_traits.inl", + "cuda/include/thrust/iterator/detail/iterator_traversal_tags.h", + "cuda/include/thrust/iterator/detail/join_iterator.h", + "cuda/include/thrust/iterator/detail/minimum_category.h", + "cuda/include/thrust/iterator/detail/minimum_system.h", + "cuda/include/thrust/iterator/detail/normal_iterator.h", + "cuda/include/thrust/iterator/detail/permutation_iterator_base.h", + "cuda/include/thrust/iterator/detail/retag.h", + "cuda/include/thrust/iterator/detail/reverse_iterator.inl", + "cuda/include/thrust/iterator/detail/reverse_iterator_base.h", + "cuda/include/thrust/iterator/detail/tagged_iterator.h", + "cuda/include/thrust/iterator/detail/transform_iterator.inl", + "cuda/include/thrust/iterator/detail/transform_output_iterator.inl", + "cuda/include/thrust/iterator/detail/tuple_of_iterator_references.h", + "cuda/include/thrust/iterator/detail/universal_categories.h", + "cuda/include/thrust/iterator/detail/zip_iterator.inl", + "cuda/include/thrust/iterator/detail/zip_iterator_base.h", + "cuda/include/thrust/iterator/discard_iterator.h", + "cuda/include/thrust/iterator/iterator_adaptor.h", + "cuda/include/thrust/iterator/iterator_categories.h", + "cuda/include/thrust/iterator/iterator_facade.h", + "cuda/include/thrust/iterator/iterator_traits.h", + "cuda/include/thrust/iterator/permutation_iterator.h", + "cuda/include/thrust/iterator/retag.h", + "cuda/include/thrust/iterator/reverse_iterator.h", + "cuda/include/thrust/iterator/transform_iterator.h", + "cuda/include/thrust/iterator/transform_output_iterator.h", + "cuda/include/thrust/iterator/zip_iterator.h", + "cuda/include/thrust/logical.h", + "cuda/include/thrust/memory.h", + "cuda/include/thrust/merge.h", + "cuda/include/thrust/mismatch.h", + "cuda/include/thrust/pair.h", + "cuda/include/thrust/partition.h", + "cuda/include/thrust/random.h", + "cuda/include/thrust/random/detail/discard_block_engine.inl", + "cuda/include/thrust/random/detail/linear_congruential_engine.inl", + "cuda/include/thrust/random/detail/linear_congruential_engine_discard.h", + "cuda/include/thrust/random/detail/linear_feedback_shift_engine.inl", + "cuda/include/thrust/random/detail/linear_feedback_shift_engine_wordmask.h", + "cuda/include/thrust/random/detail/mod.h", + "cuda/include/thrust/random/detail/normal_distribution.inl", + "cuda/include/thrust/random/detail/normal_distribution_base.h", + "cuda/include/thrust/random/detail/random_core_access.h", + "cuda/include/thrust/random/detail/subtract_with_carry_engine.inl", + "cuda/include/thrust/random/detail/uniform_int_distribution.inl", + "cuda/include/thrust/random/detail/uniform_real_distribution.inl", + "cuda/include/thrust/random/detail/xor_combine_engine.inl", + "cuda/include/thrust/random/detail/xor_combine_engine_max.h", + "cuda/include/thrust/random/discard_block_engine.h", + "cuda/include/thrust/random/linear_congruential_engine.h", + "cuda/include/thrust/random/linear_feedback_shift_engine.h", + "cuda/include/thrust/random/normal_distribution.h", + "cuda/include/thrust/random/subtract_with_carry_engine.h", + "cuda/include/thrust/random/uniform_int_distribution.h", + "cuda/include/thrust/random/uniform_real_distribution.h", + "cuda/include/thrust/random/xor_combine_engine.h", + "cuda/include/thrust/reduce.h", + "cuda/include/thrust/remove.h", + "cuda/include/thrust/replace.h", + "cuda/include/thrust/reverse.h", + "cuda/include/thrust/scan.h", + "cuda/include/thrust/scatter.h", + "cuda/include/thrust/sequence.h", + "cuda/include/thrust/set_operations.h", + "cuda/include/thrust/sort.h", + "cuda/include/thrust/swap.h", + "cuda/include/thrust/system/cpp/detail/adjacent_difference.h", + "cuda/include/thrust/system/cpp/detail/assign_value.h", + "cuda/include/thrust/system/cpp/detail/binary_search.h", + "cuda/include/thrust/system/cpp/detail/copy.h", + "cuda/include/thrust/system/cpp/detail/copy_if.h", + "cuda/include/thrust/system/cpp/detail/count.h", + "cuda/include/thrust/system/cpp/detail/equal.h", + "cuda/include/thrust/system/cpp/detail/execution_policy.h", + "cuda/include/thrust/system/cpp/detail/extrema.h", + "cuda/include/thrust/system/cpp/detail/fill.h", + "cuda/include/thrust/system/cpp/detail/find.h", + "cuda/include/thrust/system/cpp/detail/for_each.h", + "cuda/include/thrust/system/cpp/detail/gather.h", + "cuda/include/thrust/system/cpp/detail/generate.h", + "cuda/include/thrust/system/cpp/detail/get_value.h", + "cuda/include/thrust/system/cpp/detail/inner_product.h", + "cuda/include/thrust/system/cpp/detail/iter_swap.h", + "cuda/include/thrust/system/cpp/detail/logical.h", + "cuda/include/thrust/system/cpp/detail/malloc_and_free.h", + "cuda/include/thrust/system/cpp/detail/memory.inl", + "cuda/include/thrust/system/cpp/detail/merge.h", + "cuda/include/thrust/system/cpp/detail/mismatch.h", + "cuda/include/thrust/system/cpp/detail/par.h", + "cuda/include/thrust/system/cpp/detail/partition.h", + "cuda/include/thrust/system/cpp/detail/reduce.h", + "cuda/include/thrust/system/cpp/detail/reduce_by_key.h", + "cuda/include/thrust/system/cpp/detail/remove.h", + "cuda/include/thrust/system/cpp/detail/replace.h", + "cuda/include/thrust/system/cpp/detail/reverse.h", + "cuda/include/thrust/system/cpp/detail/scan.h", + "cuda/include/thrust/system/cpp/detail/scan_by_key.h", + "cuda/include/thrust/system/cpp/detail/scatter.h", + "cuda/include/thrust/system/cpp/detail/sequence.h", + "cuda/include/thrust/system/cpp/detail/set_operations.h", + "cuda/include/thrust/system/cpp/detail/sort.h", + "cuda/include/thrust/system/cpp/detail/swap_ranges.h", + "cuda/include/thrust/system/cpp/detail/tabulate.h", + "cuda/include/thrust/system/cpp/detail/temporary_buffer.h", + "cuda/include/thrust/system/cpp/detail/transform.h", + "cuda/include/thrust/system/cpp/detail/transform_reduce.h", + "cuda/include/thrust/system/cpp/detail/transform_scan.h", + "cuda/include/thrust/system/cpp/detail/uninitialized_copy.h", + "cuda/include/thrust/system/cpp/detail/uninitialized_fill.h", + "cuda/include/thrust/system/cpp/detail/unique.h", + "cuda/include/thrust/system/cpp/detail/unique_by_key.h", + "cuda/include/thrust/system/cpp/detail/vector.inl", + "cuda/include/thrust/system/cpp/execution_policy.h", + "cuda/include/thrust/system/cpp/memory.h", + "cuda/include/thrust/system/cpp/vector.h", + "cuda/include/thrust/system/cuda/config.h", + "cuda/include/thrust/system/cuda/detail/adjacent_difference.h", + "cuda/include/thrust/system/cuda/detail/assign_value.h", + "cuda/include/thrust/system/cuda/detail/binary_search.h", + "cuda/include/thrust/system/cuda/detail/copy.h", + "cuda/include/thrust/system/cuda/detail/copy_if.h", + "cuda/include/thrust/system/cuda/detail/core/agent_launcher.h", + "cuda/include/thrust/system/cuda/detail/core/alignment.h", + "cuda/include/thrust/system/cuda/detail/core/triple_chevron_launch.h", + "cuda/include/thrust/system/cuda/detail/core/util.h", + "cuda/include/thrust/system/cuda/detail/count.h", + "cuda/include/thrust/system/cuda/detail/cross_system.h", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_histogram.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_radix_sort_downsweep.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_radix_sort_upsweep.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_reduce_by_key.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_rle.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_scan.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_segment_fixup.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_select_if.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_spmv_csrt.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_spmv_orig.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/agent_spmv_row_based.cuh", + "cuda/include/thrust/system/cuda/detail/cub/agent/single_pass_scan_operators.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_adjacent_difference.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_discontinuity.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_exchange.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_histogram.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_load.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_radix_rank.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_radix_sort.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_raking_layout.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_scan.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_shuffle.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/block_store.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_atomic.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_sort.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking_commutative_only.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_warp_reductions.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_raking.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans2.cuh", + "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans3.cuh", + "cuda/include/thrust/system/cuda/detail/cub/cub.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_histogram.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_partition.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_radix_sort.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_run_length_encode.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_scan.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_segmented_radix_sort.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_segmented_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_select.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/device_spmv.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_histogram.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_radix_sort.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_reduce_by_key.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_rle.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_scan.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_select_if.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_csrt.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_orig.cuh", + "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_row_based.cuh", + "cuda/include/thrust/system/cuda/detail/cub/grid/grid_barrier.cuh", + "cuda/include/thrust/system/cuda/detail/cub/grid/grid_even_share.cuh", + "cuda/include/thrust/system/cuda/detail/cub/grid/grid_mapping.cuh", + "cuda/include/thrust/system/cuda/detail/cub/grid/grid_queue.cuh", + "cuda/include/thrust/system/cuda/detail/cub/host/mutex.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/arg_index_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_output_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/constant_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/discard_output_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/tex_obj_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/tex_ref_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/iterator/transform_input_iterator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/thread/thread_load.cuh", + "cuda/include/thrust/system/cuda/detail/cub/thread/thread_operators.cuh", + "cuda/include/thrust/system/cuda/detail/cub/thread/thread_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/thread/thread_scan.cuh", + "cuda/include/thrust/system/cuda/detail/cub/thread/thread_search.cuh", + "cuda/include/thrust/system/cuda/detail/cub/thread/thread_store.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_allocator.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_arch.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_debug.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_device.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_macro.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_namespace.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_ptx.cuh", + "cuda/include/thrust/system/cuda/detail/cub/util_type.cuh", + "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_shfl.cuh", + "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_smem.cuh", + "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_shfl.cuh", + "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_smem.cuh", + "cuda/include/thrust/system/cuda/detail/cub/warp/warp_reduce.cuh", + "cuda/include/thrust/system/cuda/detail/cub/warp/warp_scan.cuh", + "cuda/include/thrust/system/cuda/detail/equal.h", + "cuda/include/thrust/system/cuda/detail/error.inl", + "cuda/include/thrust/system/cuda/detail/execution_policy.h", + "cuda/include/thrust/system/cuda/detail/extrema.h", + "cuda/include/thrust/system/cuda/detail/fill.h", + "cuda/include/thrust/system/cuda/detail/find.h", + "cuda/include/thrust/system/cuda/detail/for_each.h", + "cuda/include/thrust/system/cuda/detail/gather.h", + "cuda/include/thrust/system/cuda/detail/generate.h", + "cuda/include/thrust/system/cuda/detail/get_value.h", + "cuda/include/thrust/system/cuda/detail/guarded_cuda_runtime_api.h", + "cuda/include/thrust/system/cuda/detail/guarded_driver_types.h", + "cuda/include/thrust/system/cuda/detail/inner_product.h", + "cuda/include/thrust/system/cuda/detail/internal/copy_cross_system.h", + "cuda/include/thrust/system/cuda/detail/internal/copy_device_to_device.h", + "cuda/include/thrust/system/cuda/detail/iter_swap.h", + "cuda/include/thrust/system/cuda/detail/logical.h", + "cuda/include/thrust/system/cuda/detail/malloc_and_free.h", + "cuda/include/thrust/system/cuda/detail/memory.inl", + "cuda/include/thrust/system/cuda/detail/memory_buffer.h", + "cuda/include/thrust/system/cuda/detail/merge.h", + "cuda/include/thrust/system/cuda/detail/mismatch.h", + "cuda/include/thrust/system/cuda/detail/par.h", + "cuda/include/thrust/system/cuda/detail/par_to_seq.h", + "cuda/include/thrust/system/cuda/detail/parallel_for.h", + "cuda/include/thrust/system/cuda/detail/partition.h", + "cuda/include/thrust/system/cuda/detail/reduce.h", + "cuda/include/thrust/system/cuda/detail/reduce_by_key.h", + "cuda/include/thrust/system/cuda/detail/remove.h", + "cuda/include/thrust/system/cuda/detail/replace.h", + "cuda/include/thrust/system/cuda/detail/reverse.h", + "cuda/include/thrust/system/cuda/detail/scan.h", + "cuda/include/thrust/system/cuda/detail/scan_by_key.h", + "cuda/include/thrust/system/cuda/detail/scatter.h", + "cuda/include/thrust/system/cuda/detail/sequence.h", + "cuda/include/thrust/system/cuda/detail/set_operations.h", + "cuda/include/thrust/system/cuda/detail/sort.h", + "cuda/include/thrust/system/cuda/detail/swap_ranges.h", + "cuda/include/thrust/system/cuda/detail/tabulate.h", + "cuda/include/thrust/system/cuda/detail/temporary_buffer.h", + "cuda/include/thrust/system/cuda/detail/terminate.h", + "cuda/include/thrust/system/cuda/detail/transform.h", + "cuda/include/thrust/system/cuda/detail/transform_reduce.h", + "cuda/include/thrust/system/cuda/detail/transform_scan.h", + "cuda/include/thrust/system/cuda/detail/uninitialized_copy.h", + "cuda/include/thrust/system/cuda/detail/uninitialized_fill.h", + "cuda/include/thrust/system/cuda/detail/unique.h", + "cuda/include/thrust/system/cuda/detail/unique_by_key.h", + "cuda/include/thrust/system/cuda/detail/util.h", + "cuda/include/thrust/system/cuda/detail/vector.inl", + "cuda/include/thrust/system/cuda/error.h", + "cuda/include/thrust/system/cuda/execution_policy.h", + "cuda/include/thrust/system/cuda/experimental/pinned_allocator.h", + "cuda/include/thrust/system/cuda/memory.h", + "cuda/include/thrust/system/cuda/vector.h", + "cuda/include/thrust/system/detail/adl/adjacent_difference.h", + "cuda/include/thrust/system/detail/adl/assign_value.h", + "cuda/include/thrust/system/detail/adl/binary_search.h", + "cuda/include/thrust/system/detail/adl/copy.h", + "cuda/include/thrust/system/detail/adl/copy_if.h", + "cuda/include/thrust/system/detail/adl/count.h", + "cuda/include/thrust/system/detail/adl/equal.h", + "cuda/include/thrust/system/detail/adl/extrema.h", + "cuda/include/thrust/system/detail/adl/fill.h", + "cuda/include/thrust/system/detail/adl/find.h", + "cuda/include/thrust/system/detail/adl/for_each.h", + "cuda/include/thrust/system/detail/adl/gather.h", + "cuda/include/thrust/system/detail/adl/generate.h", "cuda/include/thrust/system/detail/adl/get_value.h", "cuda/include/thrust/system/detail/adl/inner_product.h", - "cuda/include/thrust/system/detail/adl/copy_if.h", - "cuda/include/thrust/system/detail/adl/logical.h", "cuda/include/thrust/system/detail/adl/iter_swap.h", + "cuda/include/thrust/system/detail/adl/logical.h", "cuda/include/thrust/system/detail/adl/malloc_and_free.h", - "cuda/include/thrust/system/detail/adl/fill.h", + "cuda/include/thrust/system/detail/adl/merge.h", + "cuda/include/thrust/system/detail/adl/mismatch.h", + "cuda/include/thrust/system/detail/adl/partition.h", + "cuda/include/thrust/system/detail/adl/reduce.h", + "cuda/include/thrust/system/detail/adl/reduce_by_key.h", + "cuda/include/thrust/system/detail/adl/remove.h", + "cuda/include/thrust/system/detail/adl/replace.h", + "cuda/include/thrust/system/detail/adl/reverse.h", + "cuda/include/thrust/system/detail/adl/scan.h", + "cuda/include/thrust/system/detail/adl/scan_by_key.h", + "cuda/include/thrust/system/detail/adl/scatter.h", + "cuda/include/thrust/system/detail/adl/sequence.h", + "cuda/include/thrust/system/detail/adl/set_operations.h", + "cuda/include/thrust/system/detail/adl/sort.h", + "cuda/include/thrust/system/detail/adl/swap_ranges.h", + "cuda/include/thrust/system/detail/adl/tabulate.h", + "cuda/include/thrust/system/detail/adl/temporary_buffer.h", "cuda/include/thrust/system/detail/adl/transform.h", + "cuda/include/thrust/system/detail/adl/transform_reduce.h", + "cuda/include/thrust/system/detail/adl/transform_scan.h", + "cuda/include/thrust/system/detail/adl/uninitialized_copy.h", + "cuda/include/thrust/system/detail/adl/uninitialized_fill.h", + "cuda/include/thrust/system/detail/adl/unique.h", + "cuda/include/thrust/system/detail/adl/unique_by_key.h", + "cuda/include/thrust/system/detail/bad_alloc.h", "cuda/include/thrust/system/detail/errno.h", "cuda/include/thrust/system/detail/error_category.inl", - "cuda/include/thrust/system/detail/sequential/transform_scan.h", - "cuda/include/thrust/system/detail/sequential/unique_by_key.h", - "cuda/include/thrust/system/detail/sequential/stable_primitive_sort.h", - "cuda/include/thrust/system/detail/sequential/stable_primitive_sort.inl", - "cuda/include/thrust/system/detail/sequential/stable_merge_sort.h", - "cuda/include/thrust/system/detail/sequential/sort.inl", - "cuda/include/thrust/system/detail/sequential/partition.h", - "cuda/include/thrust/system/detail/sequential/unique.h", - "cuda/include/thrust/system/detail/sequential/execution_policy.h", - "cuda/include/thrust/system/detail/sequential/adjacent_difference.h", - "cuda/include/thrust/system/detail/sequential/sequence.h", - "cuda/include/thrust/system/detail/sequential/merge.h", - "cuda/include/thrust/system/detail/sequential/transform_reduce.h", - "cuda/include/thrust/system/detail/sequential/gather.h", - "cuda/include/thrust/system/detail/sequential/sort.h", - "cuda/include/thrust/system/detail/sequential/copy_backward.h", - "cuda/include/thrust/system/detail/sequential/stable_radix_sort.inl", - "cuda/include/thrust/system/detail/sequential/scan.h", - "cuda/include/thrust/system/detail/sequential/temporary_buffer.h", - "cuda/include/thrust/system/detail/sequential/scan_by_key.h", - "cuda/include/thrust/system/detail/sequential/reverse.h", - "cuda/include/thrust/system/detail/sequential/assign_value.h", - "cuda/include/thrust/system/detail/sequential/scatter.h", - "cuda/include/thrust/system/detail/sequential/find.h", - "cuda/include/thrust/system/detail/sequential/stable_merge_sort.inl", - "cuda/include/thrust/system/detail/sequential/merge.inl", - "cuda/include/thrust/system/detail/sequential/generate.h", - "cuda/include/thrust/system/detail/sequential/uninitialized_fill.h", - "cuda/include/thrust/system/detail/sequential/general_copy.h", - "cuda/include/thrust/system/detail/sequential/insertion_sort.h", - "cuda/include/thrust/system/detail/sequential/remove.h", - "cuda/include/thrust/system/detail/sequential/tabulate.h", - "cuda/include/thrust/system/detail/sequential/for_each.h", - "cuda/include/thrust/system/detail/sequential/reduce_by_key.h", - "cuda/include/thrust/system/detail/sequential/reduce.h", - "cuda/include/thrust/system/detail/sequential/equal.h", - "cuda/include/thrust/system/detail/sequential/stable_radix_sort.h", - "cuda/include/thrust/system/detail/sequential/copy.inl", - "cuda/include/thrust/system/detail/sequential/copy.h", - "cuda/include/thrust/system/detail/sequential/swap_ranges.h", - "cuda/include/thrust/system/detail/sequential/uninitialized_copy.h", - "cuda/include/thrust/system/detail/sequential/binary_search.h", - "cuda/include/thrust/system/detail/sequential/set_operations.h", - "cuda/include/thrust/system/detail/sequential/mismatch.h", - "cuda/include/thrust/system/detail/sequential/extrema.h", - "cuda/include/thrust/system/detail/sequential/count.h", - "cuda/include/thrust/system/detail/sequential/trivial_copy.h", - "cuda/include/thrust/system/detail/sequential/replace.h", - "cuda/include/thrust/system/detail/sequential/get_value.h", - "cuda/include/thrust/system/detail/sequential/inner_product.h", - "cuda/include/thrust/system/detail/sequential/copy_if.h", - "cuda/include/thrust/system/detail/sequential/logical.h", - "cuda/include/thrust/system/detail/sequential/iter_swap.h", - "cuda/include/thrust/system/detail/sequential/malloc_and_free.h", - "cuda/include/thrust/system/detail/sequential/fill.h", - "cuda/include/thrust/system/detail/sequential/transform.h", - "cuda/include/thrust/system/detail/error_condition.inl", - "cuda/include/thrust/system/detail/internal/decompose.h", "cuda/include/thrust/system/detail/error_code.inl", - "cuda/include/thrust/system/detail/generic/transform_scan.h", - "cuda/include/thrust/system/detail/generic/memory.inl", - "cuda/include/thrust/system/detail/generic/transform.inl", - "cuda/include/thrust/system/detail/generic/binary_search.inl", - "cuda/include/thrust/system/detail/generic/scan_by_key.inl", - "cuda/include/thrust/system/detail/generic/unique_by_key.h", - "cuda/include/thrust/system/detail/generic/inner_product.inl", - "cuda/include/thrust/system/detail/generic/select_system.h", - "cuda/include/thrust/system/detail/generic/sequence.inl", - "cuda/include/thrust/system/detail/generic/sort.inl", - "cuda/include/thrust/system/detail/generic/equal.inl", - "cuda/include/thrust/system/detail/generic/partition.h", - "cuda/include/thrust/system/detail/generic/unique.h", + "cuda/include/thrust/system/detail/error_condition.inl", "cuda/include/thrust/system/detail/generic/adjacent_difference.h", - "cuda/include/thrust/system/detail/generic/tag.h", - "cuda/include/thrust/system/detail/generic/unique_by_key.inl", - "cuda/include/thrust/system/detail/generic/sequence.h", - "cuda/include/thrust/system/detail/generic/type_traits.h", - "cuda/include/thrust/system/detail/generic/merge.h", - "cuda/include/thrust/system/detail/generic/reverse.inl", - "cuda/include/thrust/system/detail/generic/tabulate.inl", - "cuda/include/thrust/system/detail/generic/unique.inl", - "cuda/include/thrust/system/detail/generic/scatter.inl", - "cuda/include/thrust/system/detail/generic/set_operations.inl", - "cuda/include/thrust/system/detail/generic/copy_if.inl", - "cuda/include/thrust/system/detail/generic/transform_reduce.h", - "cuda/include/thrust/system/detail/generic/transform_scan.inl", - "cuda/include/thrust/system/detail/generic/gather.h", - "cuda/include/thrust/system/detail/generic/reduce_by_key.inl", - "cuda/include/thrust/system/detail/generic/transform_reduce.inl", - "cuda/include/thrust/system/detail/generic/sort.h", - "cuda/include/thrust/system/detail/generic/distance.inl", - "cuda/include/thrust/system/detail/generic/scan.h", - "cuda/include/thrust/system/detail/generic/temporary_buffer.h", - "cuda/include/thrust/system/detail/generic/reduce.inl", - "cuda/include/thrust/system/detail/generic/scan_by_key.h", - "cuda/include/thrust/system/detail/generic/reverse.h", - "cuda/include/thrust/system/detail/generic/temporary_buffer.inl", - "cuda/include/thrust/system/detail/generic/scatter.h", - "cuda/include/thrust/system/detail/generic/generate.inl", "cuda/include/thrust/system/detail/generic/adjacent_difference.inl", - "cuda/include/thrust/system/detail/generic/remove.inl", "cuda/include/thrust/system/detail/generic/advance.h", - "cuda/include/thrust/system/detail/generic/find.h", - "cuda/include/thrust/system/detail/generic/merge.inl", - "cuda/include/thrust/system/detail/generic/scalar/binary_search.inl", - "cuda/include/thrust/system/detail/generic/scalar/binary_search.h", - "cuda/include/thrust/system/detail/generic/extrema.inl", - "cuda/include/thrust/system/detail/generic/generate.h", - "cuda/include/thrust/system/detail/generic/uninitialized_fill.h", + "cuda/include/thrust/system/detail/generic/advance.inl", + "cuda/include/thrust/system/detail/generic/binary_search.h", + "cuda/include/thrust/system/detail/generic/binary_search.inl", + "cuda/include/thrust/system/detail/generic/copy.h", + "cuda/include/thrust/system/detail/generic/copy.inl", + "cuda/include/thrust/system/detail/generic/copy_if.h", + "cuda/include/thrust/system/detail/generic/copy_if.inl", + "cuda/include/thrust/system/detail/generic/count.h", "cuda/include/thrust/system/detail/generic/count.inl", - "cuda/include/thrust/system/detail/generic/remove.h", - "cuda/include/thrust/system/detail/generic/uninitialized_copy.inl", - "cuda/include/thrust/system/detail/generic/tabulate.h", - "cuda/include/thrust/system/detail/generic/for_each.h", "cuda/include/thrust/system/detail/generic/distance.h", - "cuda/include/thrust/system/detail/generic/swap_ranges.inl", - "cuda/include/thrust/system/detail/generic/reduce_by_key.h", - "cuda/include/thrust/system/detail/generic/reduce.h", + "cuda/include/thrust/system/detail/generic/distance.inl", "cuda/include/thrust/system/detail/generic/equal.h", - "cuda/include/thrust/system/detail/generic/mismatch.inl", - "cuda/include/thrust/system/detail/generic/copy.inl", - "cuda/include/thrust/system/detail/generic/copy.h", - "cuda/include/thrust/system/detail/generic/swap_ranges.h", - "cuda/include/thrust/system/detail/generic/uninitialized_copy.h", - "cuda/include/thrust/system/detail/generic/binary_search.h", - "cuda/include/thrust/system/detail/generic/set_operations.h", - "cuda/include/thrust/system/detail/generic/uninitialized_fill.inl", - "cuda/include/thrust/system/detail/generic/mismatch.h", - "cuda/include/thrust/system/detail/generic/scan.inl", - "cuda/include/thrust/system/detail/generic/gather.inl", + "cuda/include/thrust/system/detail/generic/equal.inl", "cuda/include/thrust/system/detail/generic/extrema.h", - "cuda/include/thrust/system/detail/generic/count.h", - "cuda/include/thrust/system/detail/generic/replace.h", + "cuda/include/thrust/system/detail/generic/extrema.inl", + "cuda/include/thrust/system/detail/generic/fill.h", + "cuda/include/thrust/system/detail/generic/find.h", + "cuda/include/thrust/system/detail/generic/find.inl", + "cuda/include/thrust/system/detail/generic/for_each.h", + "cuda/include/thrust/system/detail/generic/gather.h", + "cuda/include/thrust/system/detail/generic/gather.inl", + "cuda/include/thrust/system/detail/generic/generate.h", + "cuda/include/thrust/system/detail/generic/generate.inl", "cuda/include/thrust/system/detail/generic/inner_product.h", - "cuda/include/thrust/system/detail/generic/copy_if.h", + "cuda/include/thrust/system/detail/generic/inner_product.inl", "cuda/include/thrust/system/detail/generic/logical.h", - "cuda/include/thrust/system/detail/generic/partition.inl", "cuda/include/thrust/system/detail/generic/memory.h", - "cuda/include/thrust/system/detail/generic/find.inl", + "cuda/include/thrust/system/detail/generic/memory.inl", + "cuda/include/thrust/system/detail/generic/merge.h", + "cuda/include/thrust/system/detail/generic/merge.inl", + "cuda/include/thrust/system/detail/generic/mismatch.h", + "cuda/include/thrust/system/detail/generic/mismatch.inl", + "cuda/include/thrust/system/detail/generic/partition.h", + "cuda/include/thrust/system/detail/generic/partition.inl", + "cuda/include/thrust/system/detail/generic/reduce.h", + "cuda/include/thrust/system/detail/generic/reduce.inl", + "cuda/include/thrust/system/detail/generic/reduce_by_key.h", + "cuda/include/thrust/system/detail/generic/reduce_by_key.inl", + "cuda/include/thrust/system/detail/generic/remove.h", + "cuda/include/thrust/system/detail/generic/remove.inl", + "cuda/include/thrust/system/detail/generic/replace.h", "cuda/include/thrust/system/detail/generic/replace.inl", - "cuda/include/thrust/system/detail/generic/advance.inl", - "cuda/include/thrust/system/detail/generic/fill.h", + "cuda/include/thrust/system/detail/generic/reverse.h", + "cuda/include/thrust/system/detail/generic/reverse.inl", + "cuda/include/thrust/system/detail/generic/scalar/binary_search.h", + "cuda/include/thrust/system/detail/generic/scalar/binary_search.inl", + "cuda/include/thrust/system/detail/generic/scan.h", + "cuda/include/thrust/system/detail/generic/scan.inl", + "cuda/include/thrust/system/detail/generic/scan_by_key.h", + "cuda/include/thrust/system/detail/generic/scan_by_key.inl", + "cuda/include/thrust/system/detail/generic/scatter.h", + "cuda/include/thrust/system/detail/generic/scatter.inl", + "cuda/include/thrust/system/detail/generic/select_system.h", + "cuda/include/thrust/system/detail/generic/sequence.h", + "cuda/include/thrust/system/detail/generic/sequence.inl", + "cuda/include/thrust/system/detail/generic/set_operations.h", + "cuda/include/thrust/system/detail/generic/set_operations.inl", + "cuda/include/thrust/system/detail/generic/sort.h", + "cuda/include/thrust/system/detail/generic/sort.inl", + "cuda/include/thrust/system/detail/generic/swap_ranges.h", + "cuda/include/thrust/system/detail/generic/swap_ranges.inl", + "cuda/include/thrust/system/detail/generic/tabulate.h", + "cuda/include/thrust/system/detail/generic/tabulate.inl", + "cuda/include/thrust/system/detail/generic/tag.h", + "cuda/include/thrust/system/detail/generic/temporary_buffer.h", + "cuda/include/thrust/system/detail/generic/temporary_buffer.inl", "cuda/include/thrust/system/detail/generic/transform.h", + "cuda/include/thrust/system/detail/generic/transform.inl", + "cuda/include/thrust/system/detail/generic/transform_reduce.h", + "cuda/include/thrust/system/detail/generic/transform_reduce.inl", + "cuda/include/thrust/system/detail/generic/transform_scan.h", + "cuda/include/thrust/system/detail/generic/transform_scan.inl", + "cuda/include/thrust/system/detail/generic/type_traits.h", + "cuda/include/thrust/system/detail/generic/uninitialized_copy.h", + "cuda/include/thrust/system/detail/generic/uninitialized_copy.inl", + "cuda/include/thrust/system/detail/generic/uninitialized_fill.h", + "cuda/include/thrust/system/detail/generic/uninitialized_fill.inl", + "cuda/include/thrust/system/detail/generic/unique.h", + "cuda/include/thrust/system/detail/generic/unique.inl", + "cuda/include/thrust/system/detail/generic/unique_by_key.h", + "cuda/include/thrust/system/detail/generic/unique_by_key.inl", + "cuda/include/thrust/system/detail/internal/decompose.h", + "cuda/include/thrust/system/detail/sequential/adjacent_difference.h", + "cuda/include/thrust/system/detail/sequential/assign_value.h", + "cuda/include/thrust/system/detail/sequential/binary_search.h", + "cuda/include/thrust/system/detail/sequential/copy.h", + "cuda/include/thrust/system/detail/sequential/copy.inl", + "cuda/include/thrust/system/detail/sequential/copy_backward.h", + "cuda/include/thrust/system/detail/sequential/copy_if.h", + "cuda/include/thrust/system/detail/sequential/count.h", + "cuda/include/thrust/system/detail/sequential/equal.h", + "cuda/include/thrust/system/detail/sequential/execution_policy.h", + "cuda/include/thrust/system/detail/sequential/extrema.h", + "cuda/include/thrust/system/detail/sequential/fill.h", + "cuda/include/thrust/system/detail/sequential/find.h", + "cuda/include/thrust/system/detail/sequential/for_each.h", + "cuda/include/thrust/system/detail/sequential/gather.h", + "cuda/include/thrust/system/detail/sequential/general_copy.h", + "cuda/include/thrust/system/detail/sequential/generate.h", + "cuda/include/thrust/system/detail/sequential/get_value.h", + "cuda/include/thrust/system/detail/sequential/inner_product.h", + "cuda/include/thrust/system/detail/sequential/insertion_sort.h", + "cuda/include/thrust/system/detail/sequential/iter_swap.h", + "cuda/include/thrust/system/detail/sequential/logical.h", + "cuda/include/thrust/system/detail/sequential/malloc_and_free.h", + "cuda/include/thrust/system/detail/sequential/merge.h", + "cuda/include/thrust/system/detail/sequential/merge.inl", + "cuda/include/thrust/system/detail/sequential/mismatch.h", + "cuda/include/thrust/system/detail/sequential/partition.h", + "cuda/include/thrust/system/detail/sequential/reduce.h", + "cuda/include/thrust/system/detail/sequential/reduce_by_key.h", + "cuda/include/thrust/system/detail/sequential/remove.h", + "cuda/include/thrust/system/detail/sequential/replace.h", + "cuda/include/thrust/system/detail/sequential/reverse.h", + "cuda/include/thrust/system/detail/sequential/scan.h", + "cuda/include/thrust/system/detail/sequential/scan_by_key.h", + "cuda/include/thrust/system/detail/sequential/scatter.h", + "cuda/include/thrust/system/detail/sequential/sequence.h", + "cuda/include/thrust/system/detail/sequential/set_operations.h", + "cuda/include/thrust/system/detail/sequential/sort.h", + "cuda/include/thrust/system/detail/sequential/sort.inl", + "cuda/include/thrust/system/detail/sequential/stable_merge_sort.h", + "cuda/include/thrust/system/detail/sequential/stable_merge_sort.inl", + "cuda/include/thrust/system/detail/sequential/stable_primitive_sort.h", + "cuda/include/thrust/system/detail/sequential/stable_primitive_sort.inl", + "cuda/include/thrust/system/detail/sequential/stable_radix_sort.h", + "cuda/include/thrust/system/detail/sequential/stable_radix_sort.inl", + "cuda/include/thrust/system/detail/sequential/swap_ranges.h", + "cuda/include/thrust/system/detail/sequential/tabulate.h", + "cuda/include/thrust/system/detail/sequential/temporary_buffer.h", + "cuda/include/thrust/system/detail/sequential/transform.h", + "cuda/include/thrust/system/detail/sequential/transform_reduce.h", + "cuda/include/thrust/system/detail/sequential/transform_scan.h", + "cuda/include/thrust/system/detail/sequential/trivial_copy.h", + "cuda/include/thrust/system/detail/sequential/uninitialized_copy.h", + "cuda/include/thrust/system/detail/sequential/uninitialized_fill.h", + "cuda/include/thrust/system/detail/sequential/unique.h", + "cuda/include/thrust/system/detail/sequential/unique_by_key.h", "cuda/include/thrust/system/detail/system_error.inl", - "cuda/include/thrust/system/omp/execution_policy.h", - "cuda/include/thrust/system/omp/vector.h", - "cuda/include/thrust/system/omp/detail/transform_scan.h", - "cuda/include/thrust/system/omp/detail/memory.inl", - "cuda/include/thrust/system/omp/detail/reduce_intervals.inl", - "cuda/include/thrust/system/omp/detail/unique_by_key.h", - "cuda/include/thrust/system/omp/detail/sort.inl", - "cuda/include/thrust/system/omp/detail/partition.h", - "cuda/include/thrust/system/omp/detail/unique.h", - "cuda/include/thrust/system/omp/detail/execution_policy.h", + "cuda/include/thrust/system/error_code.h", "cuda/include/thrust/system/omp/detail/adjacent_difference.h", - "cuda/include/thrust/system/omp/detail/unique_by_key.inl", - "cuda/include/thrust/system/omp/detail/sequence.h", - "cuda/include/thrust/system/omp/detail/merge.h", - "cuda/include/thrust/system/omp/detail/unique.inl", + "cuda/include/thrust/system/omp/detail/assign_value.h", + "cuda/include/thrust/system/omp/detail/binary_search.h", + "cuda/include/thrust/system/omp/detail/copy.h", + "cuda/include/thrust/system/omp/detail/copy.inl", + "cuda/include/thrust/system/omp/detail/copy_if.h", "cuda/include/thrust/system/omp/detail/copy_if.inl", - "cuda/include/thrust/system/omp/detail/transform_reduce.h", - "cuda/include/thrust/system/omp/detail/gather.h", - "cuda/include/thrust/system/omp/detail/reduce_by_key.inl", - "cuda/include/thrust/system/omp/detail/sort.h", - "cuda/include/thrust/system/omp/detail/scan.h", - "cuda/include/thrust/system/omp/detail/temporary_buffer.h", + "cuda/include/thrust/system/omp/detail/count.h", "cuda/include/thrust/system/omp/detail/default_decomposition.h", - "cuda/include/thrust/system/omp/detail/reduce.inl", - "cuda/include/thrust/system/omp/detail/scan_by_key.h", - "cuda/include/thrust/system/omp/detail/reverse.h", - "cuda/include/thrust/system/omp/detail/assign_value.h", - "cuda/include/thrust/system/omp/detail/scatter.h", - "cuda/include/thrust/system/omp/detail/for_each.inl", "cuda/include/thrust/system/omp/detail/default_decomposition.inl", - "cuda/include/thrust/system/omp/detail/remove.inl", - "cuda/include/thrust/system/omp/detail/vector.inl", - "cuda/include/thrust/system/omp/detail/find.h", - "cuda/include/thrust/system/omp/detail/generate.h", - "cuda/include/thrust/system/omp/detail/uninitialized_fill.h", - "cuda/include/thrust/system/omp/detail/remove.h", - "cuda/include/thrust/system/omp/detail/tabulate.h", - "cuda/include/thrust/system/omp/detail/for_each.h", - "cuda/include/thrust/system/omp/detail/reduce_by_key.h", - "cuda/include/thrust/system/omp/detail/reduce.h", "cuda/include/thrust/system/omp/detail/equal.h", - "cuda/include/thrust/system/omp/detail/copy.inl", - "cuda/include/thrust/system/omp/detail/copy.h", - "cuda/include/thrust/system/omp/detail/swap_ranges.h", - "cuda/include/thrust/system/omp/detail/uninitialized_copy.h", - "cuda/include/thrust/system/omp/detail/binary_search.h", - "cuda/include/thrust/system/omp/detail/set_operations.h", - "cuda/include/thrust/system/omp/detail/mismatch.h", + "cuda/include/thrust/system/omp/detail/execution_policy.h", "cuda/include/thrust/system/omp/detail/extrema.h", - "cuda/include/thrust/system/omp/detail/count.h", - "cuda/include/thrust/system/omp/detail/replace.h", + "cuda/include/thrust/system/omp/detail/fill.h", + "cuda/include/thrust/system/omp/detail/find.h", + "cuda/include/thrust/system/omp/detail/for_each.h", + "cuda/include/thrust/system/omp/detail/for_each.inl", + "cuda/include/thrust/system/omp/detail/gather.h", + "cuda/include/thrust/system/omp/detail/generate.h", "cuda/include/thrust/system/omp/detail/get_value.h", "cuda/include/thrust/system/omp/detail/inner_product.h", - "cuda/include/thrust/system/omp/detail/copy_if.h", - "cuda/include/thrust/system/omp/detail/logical.h", - "cuda/include/thrust/system/omp/detail/partition.inl", "cuda/include/thrust/system/omp/detail/iter_swap.h", + "cuda/include/thrust/system/omp/detail/logical.h", + "cuda/include/thrust/system/omp/detail/malloc_and_free.h", + "cuda/include/thrust/system/omp/detail/memory.inl", + "cuda/include/thrust/system/omp/detail/merge.h", + "cuda/include/thrust/system/omp/detail/mismatch.h", "cuda/include/thrust/system/omp/detail/par.h", + "cuda/include/thrust/system/omp/detail/partition.h", + "cuda/include/thrust/system/omp/detail/partition.inl", + "cuda/include/thrust/system/omp/detail/reduce.h", + "cuda/include/thrust/system/omp/detail/reduce.inl", + "cuda/include/thrust/system/omp/detail/reduce_by_key.h", + "cuda/include/thrust/system/omp/detail/reduce_by_key.inl", "cuda/include/thrust/system/omp/detail/reduce_intervals.h", - "cuda/include/thrust/system/omp/detail/malloc_and_free.h", - "cuda/include/thrust/system/omp/detail/fill.h", + "cuda/include/thrust/system/omp/detail/reduce_intervals.inl", + "cuda/include/thrust/system/omp/detail/remove.h", + "cuda/include/thrust/system/omp/detail/remove.inl", + "cuda/include/thrust/system/omp/detail/replace.h", + "cuda/include/thrust/system/omp/detail/reverse.h", + "cuda/include/thrust/system/omp/detail/scan.h", + "cuda/include/thrust/system/omp/detail/scan_by_key.h", + "cuda/include/thrust/system/omp/detail/scatter.h", + "cuda/include/thrust/system/omp/detail/sequence.h", + "cuda/include/thrust/system/omp/detail/set_operations.h", + "cuda/include/thrust/system/omp/detail/sort.h", + "cuda/include/thrust/system/omp/detail/sort.inl", + "cuda/include/thrust/system/omp/detail/swap_ranges.h", + "cuda/include/thrust/system/omp/detail/tabulate.h", + "cuda/include/thrust/system/omp/detail/temporary_buffer.h", "cuda/include/thrust/system/omp/detail/transform.h", - "cuda/include/thrust/system/omp/memory.h", - "cuda/include/thrust/system/tbb/execution_policy.h", - "cuda/include/thrust/system/tbb/vector.h", - "cuda/include/thrust/system/tbb/detail/transform_scan.h", - "cuda/include/thrust/system/tbb/detail/memory.inl", - "cuda/include/thrust/system/tbb/detail/unique_by_key.h", - "cuda/include/thrust/system/tbb/detail/sort.inl", - "cuda/include/thrust/system/tbb/detail/partition.h", - "cuda/include/thrust/system/tbb/detail/unique.h", - "cuda/include/thrust/system/tbb/detail/execution_policy.h", + "cuda/include/thrust/system/omp/detail/transform_reduce.h", + "cuda/include/thrust/system/omp/detail/transform_scan.h", + "cuda/include/thrust/system/omp/detail/uninitialized_copy.h", + "cuda/include/thrust/system/omp/detail/uninitialized_fill.h", + "cuda/include/thrust/system/omp/detail/unique.h", + "cuda/include/thrust/system/omp/detail/unique.inl", + "cuda/include/thrust/system/omp/detail/unique_by_key.h", + "cuda/include/thrust/system/omp/detail/unique_by_key.inl", + "cuda/include/thrust/system/omp/detail/vector.inl", + "cuda/include/thrust/system/omp/execution_policy.h", + "cuda/include/thrust/system/omp/memory.h", + "cuda/include/thrust/system/omp/vector.h", + "cuda/include/thrust/system/system_error.h", "cuda/include/thrust/system/tbb/detail/adjacent_difference.h", - "cuda/include/thrust/system/tbb/detail/unique_by_key.inl", - "cuda/include/thrust/system/tbb/detail/sequence.h", - "cuda/include/thrust/system/tbb/detail/merge.h", - "cuda/include/thrust/system/tbb/detail/unique.inl", - "cuda/include/thrust/system/tbb/detail/copy_if.inl", - "cuda/include/thrust/system/tbb/detail/transform_reduce.h", - "cuda/include/thrust/system/tbb/detail/gather.h", - "cuda/include/thrust/system/tbb/detail/reduce_by_key.inl", - "cuda/include/thrust/system/tbb/detail/sort.h", - "cuda/include/thrust/system/tbb/detail/scan.h", - "cuda/include/thrust/system/tbb/detail/temporary_buffer.h", - "cuda/include/thrust/system/tbb/detail/reduce.inl", - "cuda/include/thrust/system/tbb/detail/scan_by_key.h", - "cuda/include/thrust/system/tbb/detail/reverse.h", "cuda/include/thrust/system/tbb/detail/assign_value.h", - "cuda/include/thrust/system/tbb/detail/scatter.h", - "cuda/include/thrust/system/tbb/detail/for_each.inl", - "cuda/include/thrust/system/tbb/detail/remove.inl", - "cuda/include/thrust/system/tbb/detail/vector.inl", - "cuda/include/thrust/system/tbb/detail/find.h", - "cuda/include/thrust/system/tbb/detail/merge.inl", - "cuda/include/thrust/system/tbb/detail/generate.h", - "cuda/include/thrust/system/tbb/detail/uninitialized_fill.h", - "cuda/include/thrust/system/tbb/detail/remove.h", - "cuda/include/thrust/system/tbb/detail/tabulate.h", - "cuda/include/thrust/system/tbb/detail/for_each.h", - "cuda/include/thrust/system/tbb/detail/reduce_by_key.h", - "cuda/include/thrust/system/tbb/detail/reduce.h", - "cuda/include/thrust/system/tbb/detail/equal.h", - "cuda/include/thrust/system/tbb/detail/copy.inl", - "cuda/include/thrust/system/tbb/detail/copy.h", - "cuda/include/thrust/system/tbb/detail/swap_ranges.h", - "cuda/include/thrust/system/tbb/detail/uninitialized_copy.h", "cuda/include/thrust/system/tbb/detail/binary_search.h", - "cuda/include/thrust/system/tbb/detail/set_operations.h", - "cuda/include/thrust/system/tbb/detail/mismatch.h", - "cuda/include/thrust/system/tbb/detail/scan.inl", - "cuda/include/thrust/system/tbb/detail/extrema.h", + "cuda/include/thrust/system/tbb/detail/copy.h", + "cuda/include/thrust/system/tbb/detail/copy.inl", + "cuda/include/thrust/system/tbb/detail/copy_if.h", + "cuda/include/thrust/system/tbb/detail/copy_if.inl", "cuda/include/thrust/system/tbb/detail/count.h", - "cuda/include/thrust/system/tbb/detail/replace.h", + "cuda/include/thrust/system/tbb/detail/equal.h", + "cuda/include/thrust/system/tbb/detail/execution_policy.h", + "cuda/include/thrust/system/tbb/detail/extrema.h", + "cuda/include/thrust/system/tbb/detail/fill.h", + "cuda/include/thrust/system/tbb/detail/find.h", + "cuda/include/thrust/system/tbb/detail/for_each.h", + "cuda/include/thrust/system/tbb/detail/for_each.inl", + "cuda/include/thrust/system/tbb/detail/gather.h", + "cuda/include/thrust/system/tbb/detail/generate.h", "cuda/include/thrust/system/tbb/detail/get_value.h", "cuda/include/thrust/system/tbb/detail/inner_product.h", - "cuda/include/thrust/system/tbb/detail/copy_if.h", - "cuda/include/thrust/system/tbb/detail/logical.h", - "cuda/include/thrust/system/tbb/detail/partition.inl", "cuda/include/thrust/system/tbb/detail/iter_swap.h", + "cuda/include/thrust/system/tbb/detail/logical.h", + "cuda/include/thrust/system/tbb/detail/malloc_and_free.h", + "cuda/include/thrust/system/tbb/detail/memory.inl", + "cuda/include/thrust/system/tbb/detail/merge.h", + "cuda/include/thrust/system/tbb/detail/merge.inl", + "cuda/include/thrust/system/tbb/detail/mismatch.h", "cuda/include/thrust/system/tbb/detail/par.h", + "cuda/include/thrust/system/tbb/detail/partition.h", + "cuda/include/thrust/system/tbb/detail/partition.inl", + "cuda/include/thrust/system/tbb/detail/reduce.h", + "cuda/include/thrust/system/tbb/detail/reduce.inl", + "cuda/include/thrust/system/tbb/detail/reduce_by_key.h", + "cuda/include/thrust/system/tbb/detail/reduce_by_key.inl", "cuda/include/thrust/system/tbb/detail/reduce_intervals.h", - "cuda/include/thrust/system/tbb/detail/malloc_and_free.h", - "cuda/include/thrust/system/tbb/detail/fill.h", + "cuda/include/thrust/system/tbb/detail/remove.h", + "cuda/include/thrust/system/tbb/detail/remove.inl", + "cuda/include/thrust/system/tbb/detail/replace.h", + "cuda/include/thrust/system/tbb/detail/reverse.h", + "cuda/include/thrust/system/tbb/detail/scan.h", + "cuda/include/thrust/system/tbb/detail/scan.inl", + "cuda/include/thrust/system/tbb/detail/scan_by_key.h", + "cuda/include/thrust/system/tbb/detail/scatter.h", + "cuda/include/thrust/system/tbb/detail/sequence.h", + "cuda/include/thrust/system/tbb/detail/set_operations.h", + "cuda/include/thrust/system/tbb/detail/sort.h", + "cuda/include/thrust/system/tbb/detail/sort.inl", + "cuda/include/thrust/system/tbb/detail/swap_ranges.h", + "cuda/include/thrust/system/tbb/detail/tabulate.h", + "cuda/include/thrust/system/tbb/detail/temporary_buffer.h", "cuda/include/thrust/system/tbb/detail/transform.h", - "cuda/include/thrust/system/tbb/memory.h", - "cuda/include/thrust/system/error_code.h", - "cuda/include/thrust/system/cpp/execution_policy.h", - "cuda/include/thrust/system/cpp/vector.h", - "cuda/include/thrust/system/cpp/detail/transform_scan.h", - "cuda/include/thrust/system/cpp/detail/memory.inl", - "cuda/include/thrust/system/cpp/detail/unique_by_key.h", - "cuda/include/thrust/system/cpp/detail/partition.h", - "cuda/include/thrust/system/cpp/detail/unique.h", - "cuda/include/thrust/system/cpp/detail/execution_policy.h", - "cuda/include/thrust/system/cpp/detail/adjacent_difference.h", - "cuda/include/thrust/system/cpp/detail/sequence.h", - "cuda/include/thrust/system/cpp/detail/merge.h", - "cuda/include/thrust/system/cpp/detail/transform_reduce.h", - "cuda/include/thrust/system/cpp/detail/gather.h", - "cuda/include/thrust/system/cpp/detail/sort.h", - "cuda/include/thrust/system/cpp/detail/scan.h", - "cuda/include/thrust/system/cpp/detail/temporary_buffer.h", - "cuda/include/thrust/system/cpp/detail/scan_by_key.h", - "cuda/include/thrust/system/cpp/detail/reverse.h", - "cuda/include/thrust/system/cpp/detail/assign_value.h", - "cuda/include/thrust/system/cpp/detail/scatter.h", - "cuda/include/thrust/system/cpp/detail/vector.inl", - "cuda/include/thrust/system/cpp/detail/find.h", - "cuda/include/thrust/system/cpp/detail/generate.h", - "cuda/include/thrust/system/cpp/detail/uninitialized_fill.h", - "cuda/include/thrust/system/cpp/detail/remove.h", - "cuda/include/thrust/system/cpp/detail/tabulate.h", - "cuda/include/thrust/system/cpp/detail/for_each.h", - "cuda/include/thrust/system/cpp/detail/reduce_by_key.h", - "cuda/include/thrust/system/cpp/detail/reduce.h", - "cuda/include/thrust/system/cpp/detail/equal.h", - "cuda/include/thrust/system/cpp/detail/copy.h", - "cuda/include/thrust/system/cpp/detail/swap_ranges.h", - "cuda/include/thrust/system/cpp/detail/uninitialized_copy.h", - "cuda/include/thrust/system/cpp/detail/binary_search.h", - "cuda/include/thrust/system/cpp/detail/set_operations.h", - "cuda/include/thrust/system/cpp/detail/mismatch.h", - "cuda/include/thrust/system/cpp/detail/extrema.h", - "cuda/include/thrust/system/cpp/detail/count.h", - "cuda/include/thrust/system/cpp/detail/replace.h", - "cuda/include/thrust/system/cpp/detail/get_value.h", - "cuda/include/thrust/system/cpp/detail/inner_product.h", - "cuda/include/thrust/system/cpp/detail/copy_if.h", - "cuda/include/thrust/system/cpp/detail/logical.h", - "cuda/include/thrust/system/cpp/detail/iter_swap.h", - "cuda/include/thrust/system/cpp/detail/par.h", - "cuda/include/thrust/system/cpp/detail/malloc_and_free.h", - "cuda/include/thrust/system/cpp/detail/fill.h", - "cuda/include/thrust/system/cpp/detail/transform.h", - "cuda/include/thrust/system/cpp/memory.h", - "cuda/include/thrust/system/cuda/execution_policy.h", - "cuda/include/thrust/system/cuda/vector.h", - "cuda/include/thrust/system/cuda/error.h", - "cuda/include/thrust/system/cuda/detail/copy_device_to_device.h", - "cuda/include/thrust/system/cuda/detail/transform_scan.h", - "cuda/include/thrust/system/cuda/detail/memory.inl", - "cuda/include/thrust/system/cuda/detail/cub/util_allocator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/grid/grid_mapping.cuh", - "cuda/include/thrust/system/cuda/detail/cub/grid/grid_barrier.cuh", - "cuda/include/thrust/system/cuda/detail/cub/grid/grid_even_share.cuh", - "cuda/include/thrust/system/cuda/detail/cub/grid/grid_queue.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_device.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_run_length_encode.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_partition.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_radix_sort.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_rle_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_histogram_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_reduce_by_key_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_scan_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_select_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_reduce_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_radix_sort_dispatch.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_scan.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_select.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_reduce.cuh", - "cuda/include/thrust/system/cuda/detail/cub/device/device_histogram.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_reduce.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_histo.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_scan.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_radix_sort_downsweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_radix_sort_upsweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_satomic.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_sort.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_gatomic.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_select.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_scan_prefix_operators.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_reduce_by_key.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_macro.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_namespace.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_radix_sort_upsweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_histogram_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_rle_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_select_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_scan_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_reduce_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_satomic_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_sort_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_gatomic_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_radix_sort_downsweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_reduce_by_key_sweep.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_scan_prefix_operators.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_type.cuh", - "cuda/include/thrust/system/cuda/detail/cub/host/spinlock.cuh", - "cuda/include/thrust/system/cuda/detail/cub/warp/warp_reduce.cuh", - "cuda/include/thrust/system/cuda/detail/cub/warp/warp_scan.cuh", - "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_shfl.cuh", - "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_smem.cuh", - "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_shfl.cuh", - "cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_smem.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_ptx.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_debug.cuh", - "cuda/include/thrust/system/cuda/detail/cub/cub.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/transform_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/tex_obj_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/tex_ref_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_output_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/arg_index_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/iterator/constant_input_iterator.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_scan.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_load.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_discontinuity.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_radix_rank.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_shift.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_store.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_reduce.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_exchange.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_radix_sort.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_histogram.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/block_raking_layout.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_warp_reductions.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking_commutative_only.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_atomic.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_raking.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_sort.cuh", - "cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking.cuh", - "cuda/include/thrust/system/cuda/detail/cub/thread/thread_load.cuh", - "cuda/include/thrust/system/cuda/detail/cub/thread/thread_store.cuh", - "cuda/include/thrust/system/cuda/detail/cub/thread/thread_scan.cuh", - "cuda/include/thrust/system/cuda/detail/cub/thread/thread_operators.cuh", - "cuda/include/thrust/system/cuda/detail/cub/thread/thread_reduce.cuh", - "cuda/include/thrust/system/cuda/detail/cub/util_arch.cuh", - "cuda/include/thrust/system/cuda/detail/reduce_intervals.inl", - "cuda/include/thrust/system/cuda/detail/copy_cross_system.inl", - "cuda/include/thrust/system/cuda/detail/unique_by_key.h", - "cuda/include/thrust/system/cuda/detail/bulk.h", - "cuda/include/thrust/system/cuda/detail/sort.inl", - "cuda/include/thrust/system/cuda/detail/partition.h", - "cuda/include/thrust/system/cuda/detail/unique.h", - "cuda/include/thrust/system/cuda/detail/execution_policy.h", - "cuda/include/thrust/system/cuda/detail/cuda_launch_config.h", - "cuda/include/thrust/system/cuda/detail/cub.h", - "cuda/include/thrust/system/cuda/detail/adjacent_difference.h", - "cuda/include/thrust/system/cuda/detail/sequence.h", - "cuda/include/thrust/system/cuda/detail/merge.h", - "cuda/include/thrust/system/cuda/detail/set_symmetric_difference.inl", - "cuda/include/thrust/system/cuda/detail/copy_if.inl", - "cuda/include/thrust/system/cuda/detail/transform_reduce.h", - "cuda/include/thrust/system/cuda/detail/error.inl", - "cuda/include/thrust/system/cuda/detail/gather.h", - "cuda/include/thrust/system/cuda/detail/reduce_by_key.inl", - "cuda/include/thrust/system/cuda/detail/sort.h", - "cuda/include/thrust/system/cuda/detail/synchronize.h", - "cuda/include/thrust/system/cuda/detail/scan.h", - "cuda/include/thrust/system/cuda/detail/temporary_indirect_permutation.h", - "cuda/include/thrust/system/cuda/detail/extern_shared_ptr.h", - "cuda/include/thrust/system/cuda/detail/detail/set_operation.inl", - "cuda/include/thrust/system/cuda/detail/detail/balanced_path.h", - "cuda/include/thrust/system/cuda/detail/detail/virtualized_smem_closure.h", - "cuda/include/thrust/system/cuda/detail/detail/stable_primitive_sort.h", - "cuda/include/thrust/system/cuda/detail/detail/set_operation.h", - "cuda/include/thrust/system/cuda/detail/detail/stable_primitive_sort.inl", - "cuda/include/thrust/system/cuda/detail/detail/stable_merge_sort.h", - "cuda/include/thrust/system/cuda/detail/detail/launch_closure.inl", - "cuda/include/thrust/system/cuda/detail/detail/merge.h", - "cuda/include/thrust/system/cuda/detail/detail/alignment.h", - "cuda/include/thrust/system/cuda/detail/detail/stable_radix_sort.inl", - "cuda/include/thrust/system/cuda/detail/detail/stable_sort_each.h", - "cuda/include/thrust/system/cuda/detail/detail/launch_calculator.inl", - "cuda/include/thrust/system/cuda/detail/detail/stable_merge_sort.inl", - "cuda/include/thrust/system/cuda/detail/detail/launch_closure.h", - "cuda/include/thrust/system/cuda/detail/detail/stable_radix_sort.h", - "cuda/include/thrust/system/cuda/detail/detail/uninitialized.h", - "cuda/include/thrust/system/cuda/detail/detail/cached_temporary_allocator.h", - "cuda/include/thrust/system/cuda/detail/detail/launch_calculator.h", - "cuda/include/thrust/system/cuda/detail/detail/stable_sort_each.inl", - "cuda/include/thrust/system/cuda/detail/temporary_buffer.h", - "cuda/include/thrust/system/cuda/detail/default_decomposition.h", - "cuda/include/thrust/system/cuda/detail/reduce.inl", - "cuda/include/thrust/system/cuda/detail/scan_by_key.h", - "cuda/include/thrust/system/cuda/detail/reverse.h", - "cuda/include/thrust/system/cuda/detail/assign_value.h", - "cuda/include/thrust/system/cuda/detail/scatter.h", - "cuda/include/thrust/system/cuda/detail/reduce_intervals.hpp", - "cuda/include/thrust/system/cuda/detail/for_each.inl", - "cuda/include/thrust/system/cuda/detail/default_decomposition.inl", - "cuda/include/thrust/system/cuda/detail/guarded_cuda_runtime_api.h", - "cuda/include/thrust/system/cuda/detail/adjacent_difference.inl", - "cuda/include/thrust/system/cuda/detail/vector.inl", - "cuda/include/thrust/system/cuda/detail/throw_on_error.h", - "cuda/include/thrust/system/cuda/detail/find.h", - "cuda/include/thrust/system/cuda/detail/terminate.h", - "cuda/include/thrust/system/cuda/detail/merge.inl", - "cuda/include/thrust/system/cuda/detail/trivial_copy.inl", - "cuda/include/thrust/system/cuda/detail/generate.h", - "cuda/include/thrust/system/cuda/detail/execute_on_stream.h", - "cuda/include/thrust/system/cuda/detail/uninitialized_fill.h", - "cuda/include/thrust/system/cuda/detail/remove.h", - "cuda/include/thrust/system/cuda/detail/tabulate.h", - "cuda/include/thrust/system/cuda/detail/for_each.h", - "cuda/include/thrust/system/cuda/detail/reduce_by_key.h", - "cuda/include/thrust/system/cuda/detail/decomposition.h", - "cuda/include/thrust/system/cuda/detail/reduce.h", - "cuda/include/thrust/system/cuda/detail/equal.h", - "cuda/include/thrust/system/cuda/detail/runtime_introspection.h", - "cuda/include/thrust/system/cuda/detail/copy.inl", - "cuda/include/thrust/system/cuda/detail/copy.h", - "cuda/include/thrust/system/cuda/detail/swap_ranges.h", - "cuda/include/thrust/system/cuda/detail/uninitialized_copy.h", - "cuda/include/thrust/system/cuda/detail/binary_search.h", - "cuda/include/thrust/system/cuda/detail/runtime_introspection.inl", - "cuda/include/thrust/system/cuda/detail/set_operations.h", - "cuda/include/thrust/system/cuda/detail/mismatch.h", - "cuda/include/thrust/system/cuda/detail/scan.inl", - "cuda/include/thrust/system/cuda/detail/synchronize.inl", - "cuda/include/thrust/system/cuda/detail/extrema.h", - "cuda/include/thrust/system/cuda/detail/set_union.inl", - "cuda/include/thrust/system/cuda/detail/set_intersection.inl", - "cuda/include/thrust/system/cuda/detail/count.h", - "cuda/include/thrust/system/cuda/detail/trivial_copy.h", - "cuda/include/thrust/system/cuda/detail/copy_device_to_device.inl", - "cuda/include/thrust/system/cuda/detail/replace.h", - "cuda/include/thrust/system/cuda/detail/bulk/malloc.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/config.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/closure.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/tail_flags.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/terminate.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/alignment.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/guarded_cuda_runtime_api.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/choose_sizes.inl", - "cuda/include/thrust/system/cuda/detail/bulk/detail/tuple_meta_transform.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_task.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/head_flags.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/synchronize.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/throw_on_error.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/parameter_ptr.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/cuda_launcher.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/triple_chevron_launcher.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/runtime_introspection.inl", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/cuda_launch_config.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/runtime_introspection.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/async.inl", - "cuda/include/thrust/system/cuda/detail/bulk/detail/tuple_transform.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/pointer_traits.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/apply_from_tuple.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/detail/is_contiguous_iterator.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/iterator.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/choose_sizes.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/copy.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/merge.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/accumulate.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/scan.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/detail/stable_merge_sort.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/gather.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/sort.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/reduce.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/scatter.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/adjacent_difference.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/reduce_by_key.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/algorithm/for_each.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/bulk.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/execution_policy.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/iterator/strided_iterator.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/uninitialized.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/async.hpp", - "cuda/include/thrust/system/cuda/detail/bulk/future.hpp", - "cuda/include/thrust/system/cuda/detail/guarded_driver_types.h", - "cuda/include/thrust/system/cuda/detail/get_value.h", - "cuda/include/thrust/system/cuda/detail/inner_product.h", - "cuda/include/thrust/system/cuda/detail/copy_if.h", - "cuda/include/thrust/system/cuda/detail/logical.h", - "cuda/include/thrust/system/cuda/detail/iter_swap.h", - "cuda/include/thrust/system/cuda/detail/block/merge.h", - "cuda/include/thrust/system/cuda/detail/block/inclusive_scan.h", - "cuda/include/thrust/system/cuda/detail/block/merge.inl", - "cuda/include/thrust/system/cuda/detail/block/merging_sort.h", - "cuda/include/thrust/system/cuda/detail/block/exclusive_scan.h", - "cuda/include/thrust/system/cuda/detail/block/reduce.h", - "cuda/include/thrust/system/cuda/detail/block/copy.h", - "cuda/include/thrust/system/cuda/detail/block/odd_even_sort.h", - "cuda/include/thrust/system/cuda/detail/par.h", - "cuda/include/thrust/system/cuda/detail/copy_cross_system.h", - "cuda/include/thrust/system/cuda/detail/reduce_intervals.h", - "cuda/include/thrust/system/cuda/detail/malloc_and_free.h", - "cuda/include/thrust/system/cuda/detail/fill.h", - "cuda/include/thrust/system/cuda/detail/set_difference.inl", - "cuda/include/thrust/system/cuda/detail/transform.h", - "cuda/include/thrust/system/cuda/experimental/pinned_allocator.h", - "cuda/include/thrust/system/cuda/memory.h", - "cuda/include/thrust/remove.h", + "cuda/include/thrust/system/tbb/detail/transform_reduce.h", + "cuda/include/thrust/system/tbb/detail/transform_scan.h", + "cuda/include/thrust/system/tbb/detail/uninitialized_copy.h", + "cuda/include/thrust/system/tbb/detail/uninitialized_fill.h", + "cuda/include/thrust/system/tbb/detail/unique.h", + "cuda/include/thrust/system/tbb/detail/unique.inl", + "cuda/include/thrust/system/tbb/detail/unique_by_key.h", + "cuda/include/thrust/system/tbb/detail/unique_by_key.inl", + "cuda/include/thrust/system/tbb/detail/vector.inl", + "cuda/include/thrust/system/tbb/execution_policy.h", + "cuda/include/thrust/system/tbb/memory.h", + "cuda/include/thrust/system/tbb/vector.h", + "cuda/include/thrust/system_error.h", "cuda/include/thrust/tabulate.h", - "cuda/include/thrust/for_each.h", - "cuda/include/thrust/distance.h", - "cuda/include/thrust/reduce.h", - "cuda/include/thrust/equal.h", - "cuda/include/thrust/complex.h", - "cuda/include/thrust/device_allocator.h", - "cuda/include/thrust/copy.h", + "cuda/include/thrust/transform.h", + "cuda/include/thrust/transform_reduce.h", + "cuda/include/thrust/transform_scan.h", + "cuda/include/thrust/tuple.h", "cuda/include/thrust/uninitialized_copy.h", - "cuda/include/thrust/device_reference.h", - "cuda/include/thrust/binary_search.h", - "cuda/include/thrust/set_operations.h", - "cuda/include/thrust/swap.h", - "cuda/include/thrust/mismatch.h", - "cuda/include/thrust/extrema.h", - "cuda/include/thrust/count.h", - "cuda/include/thrust/device_free.h", - "cuda/include/thrust/random/discard_block_engine.h", - "cuda/include/thrust/random/normal_distribution.h", - "cuda/include/thrust/random/detail/linear_feedback_shift_engine_wordmask.h", - "cuda/include/thrust/random/detail/subtract_with_carry_engine.inl", - "cuda/include/thrust/random/detail/xor_combine_engine_max.h", - "cuda/include/thrust/random/detail/linear_congruential_engine_discard.h", - "cuda/include/thrust/random/detail/uniform_int_distribution.inl", - "cuda/include/thrust/random/detail/discard_block_engine.inl", - "cuda/include/thrust/random/detail/uniform_real_distribution.inl", - "cuda/include/thrust/random/detail/random_core_access.h", - "cuda/include/thrust/random/detail/mod.h", - "cuda/include/thrust/random/detail/linear_feedback_shift_engine.inl", - "cuda/include/thrust/random/detail/linear_congruential_engine.inl", - "cuda/include/thrust/random/detail/xor_combine_engine.inl", - "cuda/include/thrust/random/detail/normal_distribution.inl", - "cuda/include/thrust/random/detail/normal_distribution_base.h", - "cuda/include/thrust/random/uniform_int_distribution.h", - "cuda/include/thrust/random/linear_feedback_shift_engine.h", - "cuda/include/thrust/random/xor_combine_engine.h", - "cuda/include/thrust/random/subtract_with_carry_engine.h", - "cuda/include/thrust/random/linear_congruential_engine.h", - "cuda/include/thrust/random/uniform_real_distribution.h", - "cuda/include/thrust/functional.h", - "cuda/include/thrust/replace.h", - "cuda/include/thrust/device_new_allocator.h", - "cuda/include/thrust/host_vector.h", + "cuda/include/thrust/uninitialized_fill.h", + "cuda/include/thrust/unique.h", "cuda/include/thrust/version.h", - "cuda/include/thrust/inner_product.h", - "cuda/include/thrust/iterator/iterator_traits.h", - "cuda/include/thrust/iterator/discard_iterator.h", - "cuda/include/thrust/iterator/retag.h", - "cuda/include/thrust/iterator/permutation_iterator.h", - "cuda/include/thrust/iterator/transform_iterator.h", - "cuda/include/thrust/iterator/detail/reverse_iterator.inl", - "cuda/include/thrust/iterator/detail/zip_iterator.inl", - "cuda/include/thrust/iterator/detail/counting_iterator.inl", - "cuda/include/thrust/iterator/detail/distance_from_result.h", - "cuda/include/thrust/iterator/detail/host_system_tag.h", - "cuda/include/thrust/iterator/detail/iterator_traversal_tags.h", - "cuda/include/thrust/iterator/detail/retag.h", - "cuda/include/thrust/iterator/detail/tagged_iterator.h", - "cuda/include/thrust/iterator/detail/iterator_traits.inl", - "cuda/include/thrust/iterator/detail/minimum_category.h", - "cuda/include/thrust/iterator/detail/discard_iterator_base.h", - "cuda/include/thrust/iterator/detail/iterator_category_to_traversal.h", - "cuda/include/thrust/iterator/detail/zip_iterator_base.h", - "cuda/include/thrust/iterator/detail/normal_iterator.h", - "cuda/include/thrust/iterator/detail/join_iterator.h", - "cuda/include/thrust/iterator/detail/device_system_tag.h", - "cuda/include/thrust/iterator/detail/universal_categories.h", - "cuda/include/thrust/iterator/detail/reverse_iterator_base.h", - "cuda/include/thrust/iterator/detail/minimum_system.h", - "cuda/include/thrust/iterator/detail/tuple_of_iterator_references.h", - "cuda/include/thrust/iterator/detail/is_iterator_category.h", - "cuda/include/thrust/iterator/detail/permutation_iterator_base.h", - "cuda/include/thrust/iterator/detail/any_assign.h", - "cuda/include/thrust/iterator/detail/any_system_tag.h", - "cuda/include/thrust/iterator/detail/is_trivial_iterator.h", - "cuda/include/thrust/iterator/detail/iterator_category_to_system.h", - "cuda/include/thrust/iterator/detail/iterator_adaptor_base.h", - "cuda/include/thrust/iterator/detail/constant_iterator_base.h", - "cuda/include/thrust/iterator/detail/transform_iterator.inl", - "cuda/include/thrust/iterator/detail/iterator_facade_category.h", - "cuda/include/thrust/iterator/detail/iterator_category_with_system_and_traversal.h", - "cuda/include/thrust/iterator/constant_iterator.h", - "cuda/include/thrust/iterator/counting_iterator.h", - "cuda/include/thrust/iterator/iterator_adaptor.h", - "cuda/include/thrust/iterator/iterator_facade.h", - "cuda/include/thrust/iterator/iterator_categories.h", - "cuda/include/thrust/iterator/reverse_iterator.h", - "cuda/include/thrust/iterator/zip_iterator.h", - "cuda/include/thrust/logical.h", - "cuda/include/thrust/tuple.h", - "cuda/include/thrust/memory.h", - "cuda/include/thrust/random.h", - "cuda/include/thrust/fill.h", - "cuda/include/thrust/transform.h", - "cuda/include/texture_types.h", - "cuda/include/nppversion.h", - "cuda/include/cuda_texture_types.h", - "cuda/include/fatbinary.h", - "cuda/include/cublasXt.h", - "cuda/include/cuda_fp16.h", "cuda/include/vector_functions.h", - "cuda/include/cusparse.h", - "cuda/include/nppi_filtering_functions.h", - "cuda/include/nppi_morphological_operations.h", - "cuda/include/sobol_direction_vectors.h", - "cuda/include/nvblas.h", - "cuda/include/curand_mtgp32dc_p_11213.h", - "cuda/include/nvcuvid.h", - "cuda/include/cuda_runtime_api.h", - "cuda/include/curand_mtgp32_kernel.h", - "cuda/include/cublas_v2.h", - "cuda/include/builtin_types.h", - "cuda/include/nppi_geometry_transforms.h", - "cuda/include/npps_support_functions.h", - "cuda/include/cufftw.h", - "cuda/include/cuda_device_runtime_api.h", - "cuda/include/sm_30_intrinsics.hpp", + "cuda/include/vector_functions.hpp", "cuda/include/vector_types.h", - "cuda/include/sm_35_atomic_functions.h", - "cuda/include/sm_20_intrinsics.h", - "cuda/include/driver_types.h", - "cuda/include/nvToolsExtCudaRt.h", - "cuda/include/curand_globals.h", - "cuda/include/device_atomic_functions.h", - "cuda/include/surface_types.h", - "cuda/include/nvrtc.h", - "cuda/include/nppdefs.h", - "cuda/include/sm_60_atomic_functions.h", - "cuda/include/driver_functions.h", - "cuda/include/cusolver_common.h", - "cuda/include/cublas.h", - "cuda/include/curand_lognormal.h", - "cuda/include/device_atomic_functions.hpp", - "cuda/include/crt/device_runtime.h", - "cuda/include/crt/storage_class.h", - "cuda/include/crt/func_macro.h", - "cuda/include/crt/host_runtime.h", - "cuda/include/nppi_arithmetic_and_logical_operations.h", - "cuda/include/npps_arithmetic_and_logical_operations.h", - "cuda/include/nppi_computer_vision.h", - "cuda/include/surface_functions.hpp", - "cuda/include/surface_functions.h", - "cuda/include/curand_normal_static.h", - "cuda/include/curand.h", - "cuda/include/math_functions_dbl_ptx3.h", - "cuda/include/curand_philox4x32_x.h", - "cuda/include/nppi_threshold_and_compare_operations.h", - "cuda/include/nvml.h", - "cuda/include/npps.h", - "cuda/include/cuda_vdpau_interop.h", - "cuda/include/sm_61_intrinsics.hpp", - "cuda/include/cublas_api.h", - "cuda/include/nppi_color_conversion.h", - "cuda/include/math_functions_dbl_ptx3.hpp", - "cuda/include/nppcore.h", - "cuda/include/cudaGL.h", - "cuda/include/fatBinaryCtl.h", - "cuda/include/npps_statistics_functions.h", - "cuda/include/cudaVDPAU.h", - "cuda/include/curand_poisson.h", - "cuda/include/cusolverDn.h", - "cuda/include/cuda_profiler_api.h", - "cuda/include/sm_20_atomic_functions.h", - "cuda/include/nvfunctional", ], cmd = """ -if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-8.0/include/math_functions.hpp" "$(@D)/cuda/include/math_functions.hpp" && cp "/usr/local/cuda-8.0/include/cufft.h" "$(@D)/cuda/include/cufft.h" && cp "/usr/local/cuda-8.0/include/nvgraph.h" "$(@D)/cuda/include/nvgraph.h" && cp "/usr/local/cuda-8.0/include/curand_normal.h" "$(@D)/cuda/include/curand_normal.h" && cp "/usr/local/cuda-8.0/include/curand_uniform.h" "$(@D)/cuda/include/curand_uniform.h" && cp "/usr/local/cuda-8.0/include/nppi_data_exchange_and_initialization.h" "$(@D)/cuda/include/nppi_data_exchange_and_initialization.h" && cp "/usr/local/cuda-8.0/include/cuda_gl_interop.h" "$(@D)/cuda/include/cuda_gl_interop.h" && cp "/usr/local/cuda-8.0/include/nppi_compression_functions.h" "$(@D)/cuda/include/nppi_compression_functions.h" && cp "/usr/local/cuda-8.0/include/npp.h" "$(@D)/cuda/include/npp.h" && cp "/usr/local/cuda-8.0/include/cuda.h" "$(@D)/cuda/include/cuda.h" && cp "/usr/local/cuda-8.0/include/nppi_statistics_functions.h" "$(@D)/cuda/include/nppi_statistics_functions.h" && cp "/usr/local/cuda-8.0/include/vector_functions.hpp" "$(@D)/cuda/include/vector_functions.hpp" && cp "/usr/local/cuda-8.0/include/sm_32_intrinsics.hpp" "$(@D)/cuda/include/sm_32_intrinsics.hpp" && cp "/usr/local/cuda-8.0/include/sm_32_intrinsics.h" "$(@D)/cuda/include/sm_32_intrinsics.h" && cp "/usr/local/cuda-8.0/include/curand_discrete.h" "$(@D)/cuda/include/curand_discrete.h" && cp "/usr/local/cuda-8.0/include/cuda_runtime.h" "$(@D)/cuda/include/cuda_runtime.h" && cp "/usr/local/cuda-8.0/include/cufftXt.h" "$(@D)/cuda/include/cufftXt.h" && cp "/usr/local/cuda-8.0/include/sm_61_intrinsics.h" "$(@D)/cuda/include/sm_61_intrinsics.h" && cp "/usr/local/cuda-8.0/include/texture_fetch_functions.h" "$(@D)/cuda/include/texture_fetch_functions.h" && cp "/usr/local/cuda-8.0/include/curand_mrg32k3a.h" "$(@D)/cuda/include/curand_mrg32k3a.h" && cp "/usr/local/cuda-8.0/include/host_defines.h" "$(@D)/cuda/include/host_defines.h" && cp "/usr/local/cuda-8.0/include/common_functions.h" "$(@D)/cuda/include/common_functions.h" && cp "/usr/local/cuda-8.0/include/nppi_support_functions.h" "$(@D)/cuda/include/nppi_support_functions.h" && cp "/usr/local/cuda-8.0/include/nppi_linear_transforms.h" "$(@D)/cuda/include/nppi_linear_transforms.h" && cp "/usr/local/cuda-8.0/include/device_double_functions.hpp" "$(@D)/cuda/include/device_double_functions.hpp" && cp "/usr/local/cuda-8.0/include/math_constants.h" "$(@D)/cuda/include/math_constants.h" && cp "/usr/local/cuda-8.0/include/nvToolsExtSync.h" "$(@D)/cuda/include/nvToolsExtSync.h" && cp "/usr/local/cuda-8.0/include/npps_initialization.h" "$(@D)/cuda/include/npps_initialization.h" && cp "/usr/local/cuda-8.0/include/cusolverSp_LOWLEVEL_PREVIEW.h" "$(@D)/cuda/include/cusolverSp_LOWLEVEL_PREVIEW.h" && cp "/usr/local/cuda-8.0/include/texture_indirect_functions.hpp" "$(@D)/cuda/include/texture_indirect_functions.hpp" && cp "/usr/local/cuda-8.0/include/cudaProfiler.h" "$(@D)/cuda/include/cudaProfiler.h" && cp "/usr/local/cuda-8.0/include/npps_filtering_functions.h" "$(@D)/cuda/include/npps_filtering_functions.h" && cp "/usr/local/cuda-8.0/include/cusparse_v2.h" "$(@D)/cuda/include/cusparse_v2.h" && cp "/usr/local/cuda-8.0/include/nppi.h" "$(@D)/cuda/include/nppi.h" && cp "/usr/local/cuda-8.0/include/surface_indirect_functions.h" "$(@D)/cuda/include/surface_indirect_functions.h" && cp "/usr/local/cuda-8.0/include/sm_30_intrinsics.h" "$(@D)/cuda/include/sm_30_intrinsics.h" && cp "/usr/local/cuda-8.0/include/device_double_functions.h" "$(@D)/cuda/include/device_double_functions.h" && cp "/usr/local/cuda-8.0/include/sm_35_intrinsics.h" "$(@D)/cuda/include/sm_35_intrinsics.h" && cp "/usr/local/cuda-8.0/include/cusolverSp.h" "$(@D)/cuda/include/cusolverSp.h" && cp "/usr/local/cuda-8.0/include/library_types.h" "$(@D)/cuda/include/library_types.h" && cp "/usr/local/cuda-8.0/include/surface_indirect_functions.hpp" "$(@D)/cuda/include/surface_indirect_functions.hpp" && cp "/usr/local/cuda-8.0/include/cudalibxt.h" "$(@D)/cuda/include/cudalibxt.h" && cp "/usr/local/cuda-8.0/include/channel_descriptor.h" "$(@D)/cuda/include/channel_descriptor.h" && cp "/usr/local/cuda-8.0/include/device_functions_decls.h" "$(@D)/cuda/include/device_functions_decls.h" && cp "/usr/local/cuda-8.0/include/curand_kernel.h" "$(@D)/cuda/include/curand_kernel.h" && cp "/usr/local/cuda-8.0/include/curand_mtgp32_host.h" "$(@D)/cuda/include/curand_mtgp32_host.h" && cp "/usr/local/cuda-8.0/include/nvToolsExtCuda.h" "$(@D)/cuda/include/nvToolsExtCuda.h" && cp "/usr/local/cuda-8.0/include/nvToolsExt.h" "$(@D)/cuda/include/nvToolsExt.h" && cp "/usr/local/cuda-8.0/include/cuComplex.h" "$(@D)/cuda/include/cuComplex.h" && cp "/usr/local/cuda-8.0/include/sm_32_atomic_functions.h" "$(@D)/cuda/include/sm_32_atomic_functions.h" && cp "/usr/local/cuda-8.0/include/texture_indirect_functions.h" "$(@D)/cuda/include/texture_indirect_functions.h" && cp "/usr/local/cuda-8.0/include/sm_32_atomic_functions.hpp" "$(@D)/cuda/include/sm_32_atomic_functions.hpp" && cp "/usr/local/cuda-8.0/include/sm_20_intrinsics.hpp" "$(@D)/cuda/include/sm_20_intrinsics.hpp" && cp "/usr/local/cuda-8.0/include/device_launch_parameters.h" "$(@D)/cuda/include/device_launch_parameters.h" && cp "/usr/local/cuda-8.0/include/curand_mtgp32.h" "$(@D)/cuda/include/curand_mtgp32.h" && cp "/usr/local/cuda-8.0/include/texture_fetch_functions.hpp" "$(@D)/cuda/include/texture_fetch_functions.hpp" && cp "/usr/local/cuda-8.0/include/cuda_occupancy.h" "$(@D)/cuda/include/cuda_occupancy.h" && cp "/usr/local/cuda-8.0/include/CL/opencl.h" "$(@D)/cuda/include/CL/opencl.h" && cp "/usr/local/cuda-8.0/include/CL/cl_platform.h" "$(@D)/cuda/include/CL/cl_platform.h" && cp "/usr/local/cuda-8.0/include/CL/cl_egl.h" "$(@D)/cuda/include/CL/cl_egl.h" && cp "/usr/local/cuda-8.0/include/CL/cl_gl.h" "$(@D)/cuda/include/CL/cl_gl.h" && cp "/usr/local/cuda-8.0/include/CL/cl.h" "$(@D)/cuda/include/CL/cl.h" && cp "/usr/local/cuda-8.0/include/CL/cl_gl_ext.h" "$(@D)/cuda/include/CL/cl_gl_ext.h" && cp "/usr/local/cuda-8.0/include/CL/cl_ext.h" "$(@D)/cuda/include/CL/cl_ext.h" && cp "/usr/local/cuda-8.0/include/CL/cl.hpp" "$(@D)/cuda/include/CL/cl.hpp" && cp "/usr/local/cuda-8.0/include/host_config.h" "$(@D)/cuda/include/host_config.h" && cp "/usr/local/cuda-8.0/include/cuda_surface_types.h" "$(@D)/cuda/include/cuda_surface_types.h" && cp "/usr/local/cuda-8.0/include/math_functions.h" "$(@D)/cuda/include/math_functions.h" && cp "/usr/local/cuda-8.0/include/nvToolsExtMeta.h" "$(@D)/cuda/include/nvToolsExtMeta.h" && cp "/usr/local/cuda-8.0/include/sm_20_atomic_functions.hpp" "$(@D)/cuda/include/sm_20_atomic_functions.hpp" && cp "/usr/local/cuda-8.0/include/device_functions.h" "$(@D)/cuda/include/device_functions.h" && cp "/usr/local/cuda-8.0/include/device_types.h" "$(@D)/cuda/include/device_types.h" && cp "/usr/local/cuda-8.0/include/npps_conversion_functions.h" "$(@D)/cuda/include/npps_conversion_functions.h" && cp "/usr/local/cuda-8.0/include/curand_precalc.h" "$(@D)/cuda/include/curand_precalc.h" && cp "/usr/local/cuda-8.0/include/cusolverRf.h" "$(@D)/cuda/include/cusolverRf.h" && cp "/usr/local/cuda-8.0/include/sm_60_atomic_functions.hpp" "$(@D)/cuda/include/sm_60_atomic_functions.hpp" && cp "/usr/local/cuda-8.0/include/cuviddec.h" "$(@D)/cuda/include/cuviddec.h" && cp "/usr/local/cuda-8.0/include/curand_discrete2.h" "$(@D)/cuda/include/curand_discrete2.h" && cp "/usr/local/cuda-8.0/include/device_functions.hpp" "$(@D)/cuda/include/device_functions.hpp" && cp "/usr/local/cuda-8.0/include/thrust/transform_scan.h" "$(@D)/cuda/include/thrust/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system_error.h" "$(@D)/cuda/include/thrust/system_error.h" && cp "/usr/local/cuda-8.0/include/thrust/device_malloc.h" "$(@D)/cuda/include/thrust/device_malloc.h" && cp "/usr/local/cuda-8.0/include/thrust/partition.h" "$(@D)/cuda/include/thrust/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/unique.h" "$(@D)/cuda/include/thrust/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/device_delete.h" "$(@D)/cuda/include/thrust/device_delete.h" && cp "/usr/local/cuda-8.0/include/thrust/execution_policy.h" "$(@D)/cuda/include/thrust/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/adjacent_difference.h" "$(@D)/cuda/include/thrust/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/sequence.h" "$(@D)/cuda/include/thrust/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/merge.h" "$(@D)/cuda/include/thrust/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/device_new.h" "$(@D)/cuda/include/thrust/device_new.h" && cp "/usr/local/cuda-8.0/include/thrust/transform_reduce.h" "$(@D)/cuda/include/thrust/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/device_vector.h" "$(@D)/cuda/include/thrust/device_vector.h" && cp "/usr/local/cuda-8.0/include/thrust/gather.h" "$(@D)/cuda/include/thrust/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/sort.h" "$(@D)/cuda/include/thrust/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/scan.h" "$(@D)/cuda/include/thrust/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/temporary_array.h" "$(@D)/cuda/include/thrust/detail/temporary_array.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/util/align.h" "$(@D)/cuda/include/thrust/detail/util/align.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/util/blocking.h" "$(@D)/cuda/include/thrust/detail/util/blocking.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/transform.inl" "$(@D)/cuda/include/thrust/detail/transform.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_vector.inl" "$(@D)/cuda/include/thrust/detail/device_vector.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/binary_search.inl" "$(@D)/cuda/include/thrust/detail/binary_search.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/overlapped_copy.h" "$(@D)/cuda/include/thrust/detail/overlapped_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/vector_base.inl" "$(@D)/cuda/include/thrust/detail/vector_base.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_reference.inl" "$(@D)/cuda/include/thrust/detail/device_reference.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/actor.h" "$(@D)/cuda/include/thrust/detail/functional/actor.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/value.h" "$(@D)/cuda/include/thrust/detail/functional/value.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/logical_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/logical_operators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/relational_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/relational_operators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/assignment_operator.h" "$(@D)/cuda/include/thrust/detail/functional/operators/assignment_operator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/bitwise_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/bitwise_operators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/operator_adaptors.h" "$(@D)/cuda/include/thrust/detail/functional/operators/operator_adaptors.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/arithmetic_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/arithmetic_operators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/operators/compound_assignment_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/compound_assignment_operators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/argument.h" "$(@D)/cuda/include/thrust/detail/functional/argument.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/placeholder.h" "$(@D)/cuda/include/thrust/detail/functional/placeholder.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/actor.inl" "$(@D)/cuda/include/thrust/detail/functional/actor.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional/composite.h" "$(@D)/cuda/include/thrust/detail/functional/composite.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/static_map.h" "$(@D)/cuda/include/thrust/detail/static_map.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/has_nested_type.h" "$(@D)/cuda/include/thrust/detail/type_traits/has_nested_type.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/is_call_possible.h" "$(@D)/cuda/include/thrust/detail/type_traits/is_call_possible.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/function_traits.h" "$(@D)/cuda/include/thrust/detail/type_traits/function_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/pointer_traits.h" "$(@D)/cuda/include/thrust/detail/type_traits/pointer_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/has_member_function.h" "$(@D)/cuda/include/thrust/detail/type_traits/has_member_function.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/algorithm/intermediate_type_from_function_and_iterators.h" "$(@D)/cuda/include/thrust/detail/type_traits/algorithm/intermediate_type_from_function_and_iterators.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/minimum_type.h" "$(@D)/cuda/include/thrust/detail/type_traits/minimum_type.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/has_trivial_assign.h" "$(@D)/cuda/include/thrust/detail/type_traits/has_trivial_assign.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/is_metafunction_defined.h" "$(@D)/cuda/include/thrust/detail/type_traits/is_metafunction_defined.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/iterator/is_discard_iterator.h" "$(@D)/cuda/include/thrust/detail/type_traits/iterator/is_discard_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/iterator/is_output_iterator.h" "$(@D)/cuda/include/thrust/detail/type_traits/iterator/is_output_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits/result_of_adaptable_function.h" "$(@D)/cuda/include/thrust/detail/type_traits/result_of_adaptable_function.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/reference.h" "$(@D)/cuda/include/thrust/detail/reference.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/inner_product.inl" "$(@D)/cuda/include/thrust/detail/inner_product.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/use_default.h" "$(@D)/cuda/include/thrust/detail/use_default.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/sequence.inl" "$(@D)/cuda/include/thrust/detail/sequence.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/sort.inl" "$(@D)/cuda/include/thrust/detail/sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/equal.inl" "$(@D)/cuda/include/thrust/detail/equal.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/execution_policy.h" "$(@D)/cuda/include/thrust/detail/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/integer_traits.h" "$(@D)/cuda/include/thrust/detail/integer_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/type_traits.h" "$(@D)/cuda/include/thrust/detail/type_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/reverse.inl" "$(@D)/cuda/include/thrust/detail/reverse.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/tabulate.inl" "$(@D)/cuda/include/thrust/detail/tabulate.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/unique.inl" "$(@D)/cuda/include/thrust/detail/unique.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/scatter.inl" "$(@D)/cuda/include/thrust/detail/scatter.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/set_operations.inl" "$(@D)/cuda/include/thrust/detail/set_operations.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_malloc.inl" "$(@D)/cuda/include/thrust/detail/device_malloc.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/copy_if.inl" "$(@D)/cuda/include/thrust/detail/copy_if.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/fill.inl" "$(@D)/cuda/include/thrust/detail/fill.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/temporary_array.inl" "$(@D)/cuda/include/thrust/detail/temporary_array.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/transform_scan.inl" "$(@D)/cuda/include/thrust/detail/transform_scan.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/minmax.h" "$(@D)/cuda/include/thrust/detail/minmax.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/swap.inl" "$(@D)/cuda/include/thrust/detail/swap.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/pointer.inl" "$(@D)/cuda/include/thrust/detail/pointer.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/transform_reduce.inl" "$(@D)/cuda/include/thrust/detail/transform_reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/config.h" "$(@D)/cuda/include/thrust/detail/config.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/distance.inl" "$(@D)/cuda/include/thrust/detail/distance.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/pair.inl" "$(@D)/cuda/include/thrust/detail/pair.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/temporary_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/temporary_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/tagged_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/tagged_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/destroy_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/destroy_range.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/destroy_range.h" "$(@D)/cuda/include/thrust/detail/allocator/destroy_range.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/no_throw_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/no_throw_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/default_construct_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/default_construct_range.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/fill_construct_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/fill_construct_range.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/tagged_allocator.inl" "$(@D)/cuda/include/thrust/detail/allocator/tagged_allocator.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/malloc_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/malloc_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/allocator_traits.h" "$(@D)/cuda/include/thrust/detail/allocator/allocator_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/copy_construct_range.h" "$(@D)/cuda/include/thrust/detail/allocator/copy_construct_range.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/allocator_traits.inl" "$(@D)/cuda/include/thrust/detail/allocator/allocator_traits.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/default_construct_range.h" "$(@D)/cuda/include/thrust/detail/allocator/default_construct_range.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/copy_construct_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/copy_construct_range.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/malloc_allocator.inl" "$(@D)/cuda/include/thrust/detail/allocator/malloc_allocator.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/temporary_allocator.inl" "$(@D)/cuda/include/thrust/detail/allocator/temporary_allocator.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/allocator/fill_construct_range.h" "$(@D)/cuda/include/thrust/detail/allocator/fill_construct_range.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/detail/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/reduce.inl" "$(@D)/cuda/include/thrust/detail/reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_new.inl" "$(@D)/cuda/include/thrust/detail/device_new.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/pointer.h" "$(@D)/cuda/include/thrust/detail/pointer.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/for_each.inl" "$(@D)/cuda/include/thrust/detail/for_each.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/generate.inl" "$(@D)/cuda/include/thrust/detail/generate.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/dispatch/is_trivial_copy.h" "$(@D)/cuda/include/thrust/detail/dispatch/is_trivial_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/adjacent_difference.inl" "$(@D)/cuda/include/thrust/detail/adjacent_difference.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/tuple_meta_transform.h" "$(@D)/cuda/include/thrust/detail/tuple_meta_transform.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/functional.inl" "$(@D)/cuda/include/thrust/detail/functional.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/remove.inl" "$(@D)/cuda/include/thrust/detail/remove.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/tuple_transform.h" "$(@D)/cuda/include/thrust/detail/tuple_transform.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/merge.inl" "$(@D)/cuda/include/thrust/detail/merge.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/extrema.inl" "$(@D)/cuda/include/thrust/detail/extrema.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/trivial_sequence.h" "$(@D)/cuda/include/thrust/detail/trivial_sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/vector_base.h" "$(@D)/cuda/include/thrust/detail/vector_base.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/count.inl" "$(@D)/cuda/include/thrust/detail/count.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/uninitialized_copy.inl" "$(@D)/cuda/include/thrust/detail/uninitialized_copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/function.h" "$(@D)/cuda/include/thrust/detail/function.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/swap_ranges.inl" "$(@D)/cuda/include/thrust/detail/swap_ranges.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_delete.inl" "$(@D)/cuda/include/thrust/detail/device_delete.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/static_assert.h" "$(@D)/cuda/include/thrust/detail/static_assert.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/logical.inl" "$(@D)/cuda/include/thrust/detail/logical.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/seq.h" "$(@D)/cuda/include/thrust/detail/seq.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/mpl/math.h" "$(@D)/cuda/include/thrust/detail/mpl/math.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/mismatch.inl" "$(@D)/cuda/include/thrust/detail/mismatch.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/internal_functional.h" "$(@D)/cuda/include/thrust/detail/internal_functional.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/get_iterator_value.h" "$(@D)/cuda/include/thrust/detail/get_iterator_value.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/copy.inl" "$(@D)/cuda/include/thrust/detail/copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/copy.h" "$(@D)/cuda/include/thrust/detail/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/catrigf.h" "$(@D)/cuda/include/thrust/detail/complex/catrigf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/cpowf.h" "$(@D)/cuda/include/thrust/detail/complex/cpowf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/csqrtf.h" "$(@D)/cuda/include/thrust/detail/complex/csqrtf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/ccoshf.h" "$(@D)/cuda/include/thrust/detail/complex/ccoshf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/csinhf.h" "$(@D)/cuda/include/thrust/detail/complex/csinhf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/clogf.h" "$(@D)/cuda/include/thrust/detail/complex/clogf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/ccosh.h" "$(@D)/cuda/include/thrust/detail/complex/ccosh.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/arithmetic.h" "$(@D)/cuda/include/thrust/detail/complex/arithmetic.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/csqrt.h" "$(@D)/cuda/include/thrust/detail/complex/csqrt.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/cpow.h" "$(@D)/cuda/include/thrust/detail/complex/cpow.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/complex.inl" "$(@D)/cuda/include/thrust/detail/complex/complex.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/math_private.h" "$(@D)/cuda/include/thrust/detail/complex/math_private.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/c99math.h" "$(@D)/cuda/include/thrust/detail/complex/c99math.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/cproj.h" "$(@D)/cuda/include/thrust/detail/complex/cproj.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/catrig.h" "$(@D)/cuda/include/thrust/detail/complex/catrig.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/ctanhf.h" "$(@D)/cuda/include/thrust/detail/complex/ctanhf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/cexpf.h" "$(@D)/cuda/include/thrust/detail/complex/cexpf.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/csinh.h" "$(@D)/cuda/include/thrust/detail/complex/csinh.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/stream.h" "$(@D)/cuda/include/thrust/detail/complex/stream.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/ctanh.h" "$(@D)/cuda/include/thrust/detail/complex/ctanh.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/cexp.h" "$(@D)/cuda/include/thrust/detail/complex/cexp.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/complex/clog.h" "$(@D)/cuda/include/thrust/detail/complex/clog.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/range/head_flags.h" "$(@D)/cuda/include/thrust/detail/range/head_flags.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/range/tail_flags.h" "$(@D)/cuda/include/thrust/detail/range/tail_flags.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/execute_with_allocator.h" "$(@D)/cuda/include/thrust/detail/execute_with_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/integer_math.h" "$(@D)/cuda/include/thrust/detail/integer_math.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/swap.h" "$(@D)/cuda/include/thrust/detail/swap.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/uninitialized_fill.inl" "$(@D)/cuda/include/thrust/detail/uninitialized_fill.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/scan.inl" "$(@D)/cuda/include/thrust/detail/scan.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/gather.inl" "$(@D)/cuda/include/thrust/detail/gather.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/reference_forward_declaration.h" "$(@D)/cuda/include/thrust/detail/reference_forward_declaration.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/numeric_traits.h" "$(@D)/cuda/include/thrust/detail/numeric_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/reference.inl" "$(@D)/cuda/include/thrust/detail/reference.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/cstdint.h" "$(@D)/cuda/include/thrust/detail/cstdint.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_free.inl" "$(@D)/cuda/include/thrust/detail/device_free.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/copy_if.h" "$(@D)/cuda/include/thrust/detail/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/partition.inl" "$(@D)/cuda/include/thrust/detail/partition.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/find.inl" "$(@D)/cuda/include/thrust/detail/find.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/forceinline.h" "$(@D)/cuda/include/thrust/detail/config/forceinline.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/debug.h" "$(@D)/cuda/include/thrust/detail/config/debug.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/config.h" "$(@D)/cuda/include/thrust/detail/config/config.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/host_device.h" "$(@D)/cuda/include/thrust/detail/config/host_device.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/host_system.h" "$(@D)/cuda/include/thrust/detail/config/host_system.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/compiler.h" "$(@D)/cuda/include/thrust/detail/config/compiler.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/device_system.h" "$(@D)/cuda/include/thrust/detail/config/device_system.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/compiler_fence.h" "$(@D)/cuda/include/thrust/detail/config/compiler_fence.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/exec_check_disable.h" "$(@D)/cuda/include/thrust/detail/config/exec_check_disable.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/simple_defines.h" "$(@D)/cuda/include/thrust/detail/config/simple_defines.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/config/global_workarounds.h" "$(@D)/cuda/include/thrust/detail/config/global_workarounds.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/replace.inl" "$(@D)/cuda/include/thrust/detail/replace.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/device_ptr.inl" "$(@D)/cuda/include/thrust/detail/device_ptr.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/tuple.inl" "$(@D)/cuda/include/thrust/detail/tuple.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/detail/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/host_vector.inl" "$(@D)/cuda/include/thrust/detail/host_vector.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/raw_pointer_cast.h" "$(@D)/cuda/include/thrust/detail/raw_pointer_cast.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/advance.inl" "$(@D)/cuda/include/thrust/detail/advance.inl" && cp "/usr/local/cuda-8.0/include/thrust/detail/contiguous_storage.h" "$(@D)/cuda/include/thrust/detail/contiguous_storage.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/raw_reference_cast.h" "$(@D)/cuda/include/thrust/detail/raw_reference_cast.h" && cp "/usr/local/cuda-8.0/include/thrust/detail/contiguous_storage.inl" "$(@D)/cuda/include/thrust/detail/contiguous_storage.inl" && cp "/usr/local/cuda-8.0/include/thrust/reverse.h" "$(@D)/cuda/include/thrust/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/device_malloc_allocator.h" "$(@D)/cuda/include/thrust/device_malloc_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/scatter.h" "$(@D)/cuda/include/thrust/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/pair.h" "$(@D)/cuda/include/thrust/pair.h" && cp "/usr/local/cuda-8.0/include/thrust/advance.h" "$(@D)/cuda/include/thrust/advance.h" && cp "/usr/local/cuda-8.0/include/thrust/find.h" "$(@D)/cuda/include/thrust/find.h" && cp "/usr/local/cuda-8.0/include/thrust/device_ptr.h" "$(@D)/cuda/include/thrust/device_ptr.h" && cp "/usr/local/cuda-8.0/include/thrust/generate.h" "$(@D)/cuda/include/thrust/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/uninitialized_fill.h" "$(@D)/cuda/include/thrust/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/system_error.h" "$(@D)/cuda/include/thrust/system/system_error.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/bad_alloc.h" "$(@D)/cuda/include/thrust/system/detail/bad_alloc.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/transform_scan.h" "$(@D)/cuda/include/thrust/system/detail/adl/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/unique_by_key.h" "$(@D)/cuda/include/thrust/system/detail/adl/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/partition.h" "$(@D)/cuda/include/thrust/system/detail/adl/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/unique.h" "$(@D)/cuda/include/thrust/system/detail/adl/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/detail/adl/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/sequence.h" "$(@D)/cuda/include/thrust/system/detail/adl/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/merge.h" "$(@D)/cuda/include/thrust/system/detail/adl/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/transform_reduce.h" "$(@D)/cuda/include/thrust/system/detail/adl/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/gather.h" "$(@D)/cuda/include/thrust/system/detail/adl/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/sort.h" "$(@D)/cuda/include/thrust/system/detail/adl/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/scan.h" "$(@D)/cuda/include/thrust/system/detail/adl/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/detail/adl/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/scan_by_key.h" "$(@D)/cuda/include/thrust/system/detail/adl/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/reverse.h" "$(@D)/cuda/include/thrust/system/detail/adl/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/assign_value.h" "$(@D)/cuda/include/thrust/system/detail/adl/assign_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/scatter.h" "$(@D)/cuda/include/thrust/system/detail/adl/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/find.h" "$(@D)/cuda/include/thrust/system/detail/adl/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/generate.h" "$(@D)/cuda/include/thrust/system/detail/adl/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/detail/adl/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/remove.h" "$(@D)/cuda/include/thrust/system/detail/adl/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/tabulate.h" "$(@D)/cuda/include/thrust/system/detail/adl/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/for_each.h" "$(@D)/cuda/include/thrust/system/detail/adl/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/detail/adl/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/reduce.h" "$(@D)/cuda/include/thrust/system/detail/adl/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/equal.h" "$(@D)/cuda/include/thrust/system/detail/adl/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/copy.h" "$(@D)/cuda/include/thrust/system/detail/adl/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/swap_ranges.h" "$(@D)/cuda/include/thrust/system/detail/adl/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/detail/adl/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/adl/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/set_operations.h" "$(@D)/cuda/include/thrust/system/detail/adl/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/mismatch.h" "$(@D)/cuda/include/thrust/system/detail/adl/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/extrema.h" "$(@D)/cuda/include/thrust/system/detail/adl/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/count.h" "$(@D)/cuda/include/thrust/system/detail/adl/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/replace.h" "$(@D)/cuda/include/thrust/system/detail/adl/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/get_value.h" "$(@D)/cuda/include/thrust/system/detail/adl/get_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/inner_product.h" "$(@D)/cuda/include/thrust/system/detail/adl/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/copy_if.h" "$(@D)/cuda/include/thrust/system/detail/adl/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/logical.h" "$(@D)/cuda/include/thrust/system/detail/adl/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/iter_swap.h" "$(@D)/cuda/include/thrust/system/detail/adl/iter_swap.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/detail/adl/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/fill.h" "$(@D)/cuda/include/thrust/system/detail/adl/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/adl/transform.h" "$(@D)/cuda/include/thrust/system/detail/adl/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/errno.h" "$(@D)/cuda/include/thrust/system/detail/errno.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/error_category.inl" "$(@D)/cuda/include/thrust/system/detail/error_category.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/transform_scan.h" "$(@D)/cuda/include/thrust/system/detail/sequential/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/unique_by_key.h" "$(@D)/cuda/include/thrust/system/detail/sequential/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/stable_primitive_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_primitive_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/stable_primitive_sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_primitive_sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/stable_merge_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_merge_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/partition.h" "$(@D)/cuda/include/thrust/system/detail/sequential/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/unique.h" "$(@D)/cuda/include/thrust/system/detail/sequential/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/execution_policy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/detail/sequential/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/sequence.h" "$(@D)/cuda/include/thrust/system/detail/sequential/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/merge.h" "$(@D)/cuda/include/thrust/system/detail/sequential/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/transform_reduce.h" "$(@D)/cuda/include/thrust/system/detail/sequential/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/gather.h" "$(@D)/cuda/include/thrust/system/detail/sequential/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/copy_backward.h" "$(@D)/cuda/include/thrust/system/detail/sequential/copy_backward.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/stable_radix_sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_radix_sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/scan.h" "$(@D)/cuda/include/thrust/system/detail/sequential/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/detail/sequential/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/scan_by_key.h" "$(@D)/cuda/include/thrust/system/detail/sequential/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/reverse.h" "$(@D)/cuda/include/thrust/system/detail/sequential/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/assign_value.h" "$(@D)/cuda/include/thrust/system/detail/sequential/assign_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/scatter.h" "$(@D)/cuda/include/thrust/system/detail/sequential/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/find.h" "$(@D)/cuda/include/thrust/system/detail/sequential/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/stable_merge_sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_merge_sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/merge.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/merge.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/generate.h" "$(@D)/cuda/include/thrust/system/detail/sequential/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/detail/sequential/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/general_copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/general_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/insertion_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/insertion_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/remove.h" "$(@D)/cuda/include/thrust/system/detail/sequential/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/tabulate.h" "$(@D)/cuda/include/thrust/system/detail/sequential/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/for_each.h" "$(@D)/cuda/include/thrust/system/detail/sequential/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/detail/sequential/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/reduce.h" "$(@D)/cuda/include/thrust/system/detail/sequential/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/equal.h" "$(@D)/cuda/include/thrust/system/detail/sequential/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/stable_radix_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_radix_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/copy.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/swap_ranges.h" "$(@D)/cuda/include/thrust/system/detail/sequential/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/sequential/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/set_operations.h" "$(@D)/cuda/include/thrust/system/detail/sequential/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/mismatch.h" "$(@D)/cuda/include/thrust/system/detail/sequential/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/extrema.h" "$(@D)/cuda/include/thrust/system/detail/sequential/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/count.h" "$(@D)/cuda/include/thrust/system/detail/sequential/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/trivial_copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/trivial_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/replace.h" "$(@D)/cuda/include/thrust/system/detail/sequential/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/get_value.h" "$(@D)/cuda/include/thrust/system/detail/sequential/get_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/inner_product.h" "$(@D)/cuda/include/thrust/system/detail/sequential/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/copy_if.h" "$(@D)/cuda/include/thrust/system/detail/sequential/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/logical.h" "$(@D)/cuda/include/thrust/system/detail/sequential/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/iter_swap.h" "$(@D)/cuda/include/thrust/system/detail/sequential/iter_swap.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/detail/sequential/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/fill.h" "$(@D)/cuda/include/thrust/system/detail/sequential/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/sequential/transform.h" "$(@D)/cuda/include/thrust/system/detail/sequential/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/error_condition.inl" "$(@D)/cuda/include/thrust/system/detail/error_condition.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/internal/decompose.h" "$(@D)/cuda/include/thrust/system/detail/internal/decompose.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/error_code.inl" "$(@D)/cuda/include/thrust/system/detail/error_code.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/transform_scan.h" "$(@D)/cuda/include/thrust/system/detail/generic/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/memory.inl" "$(@D)/cuda/include/thrust/system/detail/generic/memory.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/transform.inl" "$(@D)/cuda/include/thrust/system/detail/generic/transform.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/binary_search.inl" "$(@D)/cuda/include/thrust/system/detail/generic/binary_search.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scan_by_key.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scan_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/unique_by_key.h" "$(@D)/cuda/include/thrust/system/detail/generic/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/inner_product.inl" "$(@D)/cuda/include/thrust/system/detail/generic/inner_product.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/select_system.h" "$(@D)/cuda/include/thrust/system/detail/generic/select_system.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/sequence.inl" "$(@D)/cuda/include/thrust/system/detail/generic/sequence.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/sort.inl" "$(@D)/cuda/include/thrust/system/detail/generic/sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/equal.inl" "$(@D)/cuda/include/thrust/system/detail/generic/equal.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/partition.h" "$(@D)/cuda/include/thrust/system/detail/generic/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/unique.h" "$(@D)/cuda/include/thrust/system/detail/generic/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/detail/generic/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/tag.h" "$(@D)/cuda/include/thrust/system/detail/generic/tag.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/unique_by_key.inl" "$(@D)/cuda/include/thrust/system/detail/generic/unique_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/sequence.h" "$(@D)/cuda/include/thrust/system/detail/generic/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/type_traits.h" "$(@D)/cuda/include/thrust/system/detail/generic/type_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/merge.h" "$(@D)/cuda/include/thrust/system/detail/generic/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/reverse.inl" "$(@D)/cuda/include/thrust/system/detail/generic/reverse.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/tabulate.inl" "$(@D)/cuda/include/thrust/system/detail/generic/tabulate.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/unique.inl" "$(@D)/cuda/include/thrust/system/detail/generic/unique.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scatter.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scatter.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/set_operations.inl" "$(@D)/cuda/include/thrust/system/detail/generic/set_operations.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/copy_if.inl" "$(@D)/cuda/include/thrust/system/detail/generic/copy_if.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/transform_reduce.h" "$(@D)/cuda/include/thrust/system/detail/generic/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/transform_scan.inl" "$(@D)/cuda/include/thrust/system/detail/generic/transform_scan.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/gather.h" "$(@D)/cuda/include/thrust/system/detail/generic/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/detail/generic/reduce_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/transform_reduce.inl" "$(@D)/cuda/include/thrust/system/detail/generic/transform_reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/sort.h" "$(@D)/cuda/include/thrust/system/detail/generic/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/distance.inl" "$(@D)/cuda/include/thrust/system/detail/generic/distance.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scan.h" "$(@D)/cuda/include/thrust/system/detail/generic/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/detail/generic/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/reduce.inl" "$(@D)/cuda/include/thrust/system/detail/generic/reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scan_by_key.h" "$(@D)/cuda/include/thrust/system/detail/generic/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/reverse.h" "$(@D)/cuda/include/thrust/system/detail/generic/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/temporary_buffer.inl" "$(@D)/cuda/include/thrust/system/detail/generic/temporary_buffer.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scatter.h" "$(@D)/cuda/include/thrust/system/detail/generic/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/generate.inl" "$(@D)/cuda/include/thrust/system/detail/generic/generate.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/adjacent_difference.inl" "$(@D)/cuda/include/thrust/system/detail/generic/adjacent_difference.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/remove.inl" "$(@D)/cuda/include/thrust/system/detail/generic/remove.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/advance.h" "$(@D)/cuda/include/thrust/system/detail/generic/advance.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/find.h" "$(@D)/cuda/include/thrust/system/detail/generic/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/merge.inl" "$(@D)/cuda/include/thrust/system/detail/generic/merge.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scalar/binary_search.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scalar/binary_search.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scalar/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/generic/scalar/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/extrema.inl" "$(@D)/cuda/include/thrust/system/detail/generic/extrema.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/generate.h" "$(@D)/cuda/include/thrust/system/detail/generic/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/count.inl" "$(@D)/cuda/include/thrust/system/detail/generic/count.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/remove.h" "$(@D)/cuda/include/thrust/system/detail/generic/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/uninitialized_copy.inl" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/tabulate.h" "$(@D)/cuda/include/thrust/system/detail/generic/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/for_each.h" "$(@D)/cuda/include/thrust/system/detail/generic/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/distance.h" "$(@D)/cuda/include/thrust/system/detail/generic/distance.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/swap_ranges.inl" "$(@D)/cuda/include/thrust/system/detail/generic/swap_ranges.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/detail/generic/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/reduce.h" "$(@D)/cuda/include/thrust/system/detail/generic/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/equal.h" "$(@D)/cuda/include/thrust/system/detail/generic/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/mismatch.inl" "$(@D)/cuda/include/thrust/system/detail/generic/mismatch.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/copy.inl" "$(@D)/cuda/include/thrust/system/detail/generic/copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/copy.h" "$(@D)/cuda/include/thrust/system/detail/generic/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/swap_ranges.h" "$(@D)/cuda/include/thrust/system/detail/generic/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/generic/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/set_operations.h" "$(@D)/cuda/include/thrust/system/detail/generic/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/uninitialized_fill.inl" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_fill.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/mismatch.h" "$(@D)/cuda/include/thrust/system/detail/generic/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/scan.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scan.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/gather.inl" "$(@D)/cuda/include/thrust/system/detail/generic/gather.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/extrema.h" "$(@D)/cuda/include/thrust/system/detail/generic/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/count.h" "$(@D)/cuda/include/thrust/system/detail/generic/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/replace.h" "$(@D)/cuda/include/thrust/system/detail/generic/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/inner_product.h" "$(@D)/cuda/include/thrust/system/detail/generic/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/copy_if.h" "$(@D)/cuda/include/thrust/system/detail/generic/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/logical.h" "$(@D)/cuda/include/thrust/system/detail/generic/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/partition.inl" "$(@D)/cuda/include/thrust/system/detail/generic/partition.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/memory.h" "$(@D)/cuda/include/thrust/system/detail/generic/memory.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/find.inl" "$(@D)/cuda/include/thrust/system/detail/generic/find.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/replace.inl" "$(@D)/cuda/include/thrust/system/detail/generic/replace.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/advance.inl" "$(@D)/cuda/include/thrust/system/detail/generic/advance.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/fill.h" "$(@D)/cuda/include/thrust/system/detail/generic/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/generic/transform.h" "$(@D)/cuda/include/thrust/system/detail/generic/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/detail/system_error.inl" "$(@D)/cuda/include/thrust/system/detail/system_error.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/execution_policy.h" "$(@D)/cuda/include/thrust/system/omp/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/vector.h" "$(@D)/cuda/include/thrust/system/omp/vector.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/omp/detail/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/memory.inl" "$(@D)/cuda/include/thrust/system/omp/detail/memory.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reduce_intervals.inl" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_intervals.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/omp/detail/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/sort.inl" "$(@D)/cuda/include/thrust/system/omp/detail/sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/partition.h" "$(@D)/cuda/include/thrust/system/omp/detail/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/unique.h" "$(@D)/cuda/include/thrust/system/omp/detail/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/omp/detail/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/omp/detail/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/unique_by_key.inl" "$(@D)/cuda/include/thrust/system/omp/detail/unique_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/sequence.h" "$(@D)/cuda/include/thrust/system/omp/detail/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/merge.h" "$(@D)/cuda/include/thrust/system/omp/detail/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/unique.inl" "$(@D)/cuda/include/thrust/system/omp/detail/unique.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/copy_if.inl" "$(@D)/cuda/include/thrust/system/omp/detail/copy_if.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/omp/detail/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/gather.h" "$(@D)/cuda/include/thrust/system/omp/detail/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/sort.h" "$(@D)/cuda/include/thrust/system/omp/detail/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/scan.h" "$(@D)/cuda/include/thrust/system/omp/detail/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/omp/detail/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/default_decomposition.h" "$(@D)/cuda/include/thrust/system/omp/detail/default_decomposition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reduce.inl" "$(@D)/cuda/include/thrust/system/omp/detail/reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/omp/detail/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reverse.h" "$(@D)/cuda/include/thrust/system/omp/detail/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/omp/detail/assign_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/scatter.h" "$(@D)/cuda/include/thrust/system/omp/detail/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/for_each.inl" "$(@D)/cuda/include/thrust/system/omp/detail/for_each.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/default_decomposition.inl" "$(@D)/cuda/include/thrust/system/omp/detail/default_decomposition.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/remove.inl" "$(@D)/cuda/include/thrust/system/omp/detail/remove.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/vector.inl" "$(@D)/cuda/include/thrust/system/omp/detail/vector.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/find.h" "$(@D)/cuda/include/thrust/system/omp/detail/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/generate.h" "$(@D)/cuda/include/thrust/system/omp/detail/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/omp/detail/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/remove.h" "$(@D)/cuda/include/thrust/system/omp/detail/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/omp/detail/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/for_each.h" "$(@D)/cuda/include/thrust/system/omp/detail/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reduce.h" "$(@D)/cuda/include/thrust/system/omp/detail/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/equal.h" "$(@D)/cuda/include/thrust/system/omp/detail/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/copy.inl" "$(@D)/cuda/include/thrust/system/omp/detail/copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/copy.h" "$(@D)/cuda/include/thrust/system/omp/detail/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/omp/detail/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/omp/detail/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/omp/detail/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/omp/detail/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/omp/detail/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/extrema.h" "$(@D)/cuda/include/thrust/system/omp/detail/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/count.h" "$(@D)/cuda/include/thrust/system/omp/detail/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/replace.h" "$(@D)/cuda/include/thrust/system/omp/detail/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/get_value.h" "$(@D)/cuda/include/thrust/system/omp/detail/get_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/omp/detail/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/omp/detail/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/logical.h" "$(@D)/cuda/include/thrust/system/omp/detail/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/partition.inl" "$(@D)/cuda/include/thrust/system/omp/detail/partition.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/omp/detail/iter_swap.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/par.h" "$(@D)/cuda/include/thrust/system/omp/detail/par.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/reduce_intervals.h" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_intervals.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/omp/detail/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/fill.h" "$(@D)/cuda/include/thrust/system/omp/detail/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/detail/transform.h" "$(@D)/cuda/include/thrust/system/omp/detail/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/omp/memory.h" "$(@D)/cuda/include/thrust/system/omp/memory.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/execution_policy.h" "$(@D)/cuda/include/thrust/system/tbb/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/vector.h" "$(@D)/cuda/include/thrust/system/tbb/vector.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/tbb/detail/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/memory.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/memory.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/tbb/detail/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/sort.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/partition.h" "$(@D)/cuda/include/thrust/system/tbb/detail/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/unique.h" "$(@D)/cuda/include/thrust/system/tbb/detail/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/tbb/detail/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/tbb/detail/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/unique_by_key.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/unique_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/sequence.h" "$(@D)/cuda/include/thrust/system/tbb/detail/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/merge.h" "$(@D)/cuda/include/thrust/system/tbb/detail/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/unique.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/unique.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/copy_if.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/copy_if.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/tbb/detail/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/gather.h" "$(@D)/cuda/include/thrust/system/tbb/detail/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/sort.h" "$(@D)/cuda/include/thrust/system/tbb/detail/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/scan.h" "$(@D)/cuda/include/thrust/system/tbb/detail/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/tbb/detail/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/reduce.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/tbb/detail/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/reverse.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/tbb/detail/assign_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/scatter.h" "$(@D)/cuda/include/thrust/system/tbb/detail/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/for_each.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/for_each.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/remove.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/remove.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/vector.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/vector.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/find.h" "$(@D)/cuda/include/thrust/system/tbb/detail/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/merge.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/merge.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/generate.h" "$(@D)/cuda/include/thrust/system/tbb/detail/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/tbb/detail/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/remove.h" "$(@D)/cuda/include/thrust/system/tbb/detail/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/tbb/detail/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/for_each.h" "$(@D)/cuda/include/thrust/system/tbb/detail/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/reduce.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/equal.h" "$(@D)/cuda/include/thrust/system/tbb/detail/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/copy.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/copy.h" "$(@D)/cuda/include/thrust/system/tbb/detail/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/tbb/detail/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/tbb/detail/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/tbb/detail/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/tbb/detail/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/tbb/detail/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/scan.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/scan.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/extrema.h" "$(@D)/cuda/include/thrust/system/tbb/detail/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/count.h" "$(@D)/cuda/include/thrust/system/tbb/detail/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/replace.h" "$(@D)/cuda/include/thrust/system/tbb/detail/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/get_value.h" "$(@D)/cuda/include/thrust/system/tbb/detail/get_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/tbb/detail/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/tbb/detail/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/logical.h" "$(@D)/cuda/include/thrust/system/tbb/detail/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/partition.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/partition.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/tbb/detail/iter_swap.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/par.h" "$(@D)/cuda/include/thrust/system/tbb/detail/par.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/reduce_intervals.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce_intervals.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/tbb/detail/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/fill.h" "$(@D)/cuda/include/thrust/system/tbb/detail/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/detail/transform.h" "$(@D)/cuda/include/thrust/system/tbb/detail/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/tbb/memory.h" "$(@D)/cuda/include/thrust/system/tbb/memory.h" && cp "/usr/local/cuda-8.0/include/thrust/system/error_code.h" "$(@D)/cuda/include/thrust/system/error_code.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/execution_policy.h" "$(@D)/cuda/include/thrust/system/cpp/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/vector.h" "$(@D)/cuda/include/thrust/system/cpp/vector.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/cpp/detail/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/memory.inl" "$(@D)/cuda/include/thrust/system/cpp/detail/memory.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/cpp/detail/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/partition.h" "$(@D)/cuda/include/thrust/system/cpp/detail/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/unique.h" "$(@D)/cuda/include/thrust/system/cpp/detail/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/cpp/detail/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/cpp/detail/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/sequence.h" "$(@D)/cuda/include/thrust/system/cpp/detail/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/merge.h" "$(@D)/cuda/include/thrust/system/cpp/detail/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/cpp/detail/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/gather.h" "$(@D)/cuda/include/thrust/system/cpp/detail/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/sort.h" "$(@D)/cuda/include/thrust/system/cpp/detail/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/scan.h" "$(@D)/cuda/include/thrust/system/cpp/detail/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/cpp/detail/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/cpp/detail/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/reverse.h" "$(@D)/cuda/include/thrust/system/cpp/detail/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/cpp/detail/assign_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/scatter.h" "$(@D)/cuda/include/thrust/system/cpp/detail/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/vector.inl" "$(@D)/cuda/include/thrust/system/cpp/detail/vector.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/find.h" "$(@D)/cuda/include/thrust/system/cpp/detail/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/generate.h" "$(@D)/cuda/include/thrust/system/cpp/detail/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/cpp/detail/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/remove.h" "$(@D)/cuda/include/thrust/system/cpp/detail/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/cpp/detail/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/for_each.h" "$(@D)/cuda/include/thrust/system/cpp/detail/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/cpp/detail/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/reduce.h" "$(@D)/cuda/include/thrust/system/cpp/detail/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/equal.h" "$(@D)/cuda/include/thrust/system/cpp/detail/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/copy.h" "$(@D)/cuda/include/thrust/system/cpp/detail/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/cpp/detail/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/cpp/detail/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/cpp/detail/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/cpp/detail/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/cpp/detail/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/extrema.h" "$(@D)/cuda/include/thrust/system/cpp/detail/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/count.h" "$(@D)/cuda/include/thrust/system/cpp/detail/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/replace.h" "$(@D)/cuda/include/thrust/system/cpp/detail/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/get_value.h" "$(@D)/cuda/include/thrust/system/cpp/detail/get_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/cpp/detail/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/cpp/detail/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/logical.h" "$(@D)/cuda/include/thrust/system/cpp/detail/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/cpp/detail/iter_swap.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/par.h" "$(@D)/cuda/include/thrust/system/cpp/detail/par.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/cpp/detail/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/fill.h" "$(@D)/cuda/include/thrust/system/cpp/detail/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/detail/transform.h" "$(@D)/cuda/include/thrust/system/cpp/detail/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cpp/memory.h" "$(@D)/cuda/include/thrust/system/cpp/memory.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/execution_policy.h" "$(@D)/cuda/include/thrust/system/cuda/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/vector.h" "$(@D)/cuda/include/thrust/system/cuda/vector.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/error.h" "$(@D)/cuda/include/thrust/system/cuda/error.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy_device_to_device.h" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_device_to_device.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/cuda/detail/transform_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/memory.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/memory.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_allocator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_allocator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/grid/grid_mapping.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_mapping.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/grid/grid_barrier.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_barrier.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/grid/grid_even_share.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_even_share.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/grid/grid_queue.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_queue.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_device.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_device.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_run_length_encode.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_run_length_encode.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_partition.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_partition.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_radix_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_radix_sort.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_rle_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_rle_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_histogram_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_histogram_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_reduce_by_key_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_reduce_by_key_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_scan_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_scan_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_select_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_select_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_reduce_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_reduce_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/dispatch/device_radix_sort_dispatch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/device_radix_sort_dispatch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_scan.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_select.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_select.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_reduce.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/device/device_histogram.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_histogram.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_reduce.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_histo.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_histo.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_scan.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_radix_sort_downsweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_radix_sort_downsweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_radix_sort_upsweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_radix_sort_upsweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_satomic.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_satomic.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_sort.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_gatomic.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/specializations/block_range_histo_gatomic.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_select.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_select.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_scan_prefix_operators.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_scan_prefix_operators.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_range/block_range_reduce_by_key.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_range/block_range_reduce_by_key.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_macro.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_macro.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_namespace.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_namespace.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_radix_sort_upsweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_radix_sort_upsweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_histogram_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_histogram_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_rle_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_rle_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_select_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_select_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_scan_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_scan_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_reduce_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_reduce_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_satomic_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_satomic_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_sort_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_sort_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_gatomic_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/specializations/block_histogram_gatomic_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_radix_sort_downsweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_radix_sort_downsweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_reduce_by_key_sweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_reduce_by_key_sweep.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block_sweep/block_scan_prefix_operators.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block_sweep/block_scan_prefix_operators.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_type.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_type.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/host/spinlock.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/host/spinlock.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/warp/warp_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/warp_reduce.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/warp/warp_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/warp_scan.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_shfl.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_shfl.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_smem.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_smem.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_shfl.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_shfl.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_smem.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_smem.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_ptx.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_ptx.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_debug.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_debug.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/cub.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/cub.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/transform_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/transform_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/tex_obj_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/tex_obj_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/tex_ref_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/tex_ref_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/cache_modified_output_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_output_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/cache_modified_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/arg_index_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/arg_index_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/iterator/constant_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/constant_input_iterator.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_scan.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_load.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_load.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_discontinuity.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_discontinuity.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_radix_rank.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_radix_rank.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_shift.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_shift.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_store.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_store.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_reduce.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_exchange.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_exchange.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_radix_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_radix_sort.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_histogram.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_histogram.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/block_raking_layout.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_raking_layout.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_warp_reductions.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_warp_reductions.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking_commutative_only.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking_commutative_only.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_atomic.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_atomic.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_raking.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_raking.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_sort.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/thread/thread_load.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_load.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/thread/thread_store.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_store.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/thread/thread_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_scan.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/thread/thread_operators.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_operators.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/thread/thread_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_reduce.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub/util_arch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_arch.cuh" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce_intervals.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce_intervals.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy_cross_system.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_cross_system.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/cuda/detail/unique_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk.h" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/sort.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/partition.h" "$(@D)/cuda/include/thrust/system/cuda/detail/partition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/unique.h" "$(@D)/cuda/include/thrust/system/cuda/detail/unique.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/execution_policy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cuda_launch_config.h" "$(@D)/cuda/include/thrust/system/cuda/detail/cuda_launch_config.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/cub.h" "$(@D)/cuda/include/thrust/system/cuda/detail/cub.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/cuda/detail/adjacent_difference.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/sequence.h" "$(@D)/cuda/include/thrust/system/cuda/detail/sequence.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/merge.h" "$(@D)/cuda/include/thrust/system/cuda/detail/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/set_symmetric_difference.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/set_symmetric_difference.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy_if.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_if.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/cuda/detail/transform_reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/error.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/error.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/gather.h" "$(@D)/cuda/include/thrust/system/cuda/detail/gather.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce_by_key.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/synchronize.h" "$(@D)/cuda/include/thrust/system/cuda/detail/synchronize.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/scan.h" "$(@D)/cuda/include/thrust/system/cuda/detail/scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/temporary_indirect_permutation.h" "$(@D)/cuda/include/thrust/system/cuda/detail/temporary_indirect_permutation.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/extern_shared_ptr.h" "$(@D)/cuda/include/thrust/system/cuda/detail/extern_shared_ptr.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/set_operation.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/set_operation.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/balanced_path.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/balanced_path.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/virtualized_smem_closure.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/virtualized_smem_closure.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_primitive_sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_primitive_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/set_operation.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/set_operation.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_primitive_sort.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_primitive_sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_merge_sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_merge_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/launch_closure.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/launch_closure.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/merge.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/alignment.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/alignment.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_radix_sort.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_radix_sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_sort_each.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_sort_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/launch_calculator.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/launch_calculator.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_merge_sort.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_merge_sort.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/launch_closure.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/launch_closure.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_radix_sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_radix_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/uninitialized.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/uninitialized.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/cached_temporary_allocator.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/cached_temporary_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/launch_calculator.h" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/launch_calculator.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/detail/stable_sort_each.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/detail/stable_sort_each.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/cuda/detail/temporary_buffer.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/default_decomposition.h" "$(@D)/cuda/include/thrust/system/cuda/detail/default_decomposition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/cuda/detail/scan_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reverse.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reverse.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/cuda/detail/assign_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/scatter.h" "$(@D)/cuda/include/thrust/system/cuda/detail/scatter.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce_intervals.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce_intervals.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/for_each.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/for_each.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/default_decomposition.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/default_decomposition.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/guarded_cuda_runtime_api.h" "$(@D)/cuda/include/thrust/system/cuda/detail/guarded_cuda_runtime_api.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/adjacent_difference.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/adjacent_difference.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/vector.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/vector.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/throw_on_error.h" "$(@D)/cuda/include/thrust/system/cuda/detail/throw_on_error.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/find.h" "$(@D)/cuda/include/thrust/system/cuda/detail/find.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/terminate.h" "$(@D)/cuda/include/thrust/system/cuda/detail/terminate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/merge.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/merge.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/trivial_copy.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/trivial_copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/generate.h" "$(@D)/cuda/include/thrust/system/cuda/detail/generate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/execute_on_stream.h" "$(@D)/cuda/include/thrust/system/cuda/detail/execute_on_stream.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/cuda/detail/uninitialized_fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/remove.h" "$(@D)/cuda/include/thrust/system/cuda/detail/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/cuda/detail/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/for_each.h" "$(@D)/cuda/include/thrust/system/cuda/detail/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce_by_key.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/decomposition.h" "$(@D)/cuda/include/thrust/system/cuda/detail/decomposition.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/equal.h" "$(@D)/cuda/include/thrust/system/cuda/detail/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/runtime_introspection.h" "$(@D)/cuda/include/thrust/system/cuda/detail/runtime_introspection.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/copy.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/cuda/detail/swap_ranges.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/cuda/detail/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/runtime_introspection.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/runtime_introspection.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/cuda/detail/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/cuda/detail/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/scan.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/scan.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/synchronize.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/synchronize.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/extrema.h" "$(@D)/cuda/include/thrust/system/cuda/detail/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/set_union.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/set_union.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/set_intersection.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/set_intersection.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/count.h" "$(@D)/cuda/include/thrust/system/cuda/detail/count.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/trivial_copy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/trivial_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy_device_to_device.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_device_to_device.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/replace.h" "$(@D)/cuda/include/thrust/system/cuda/detail/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/malloc.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/malloc.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/config.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/config.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/closure.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/closure.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/tail_flags.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/tail_flags.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/terminate.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/terminate.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/alignment.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/alignment.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/guarded_cuda_runtime_api.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/guarded_cuda_runtime_api.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/choose_sizes.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/choose_sizes.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/tuple_meta_transform.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/tuple_meta_transform.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_task.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_task.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/head_flags.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/head_flags.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/synchronize.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/synchronize.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/throw_on_error.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/throw_on_error.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/parameter_ptr.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/parameter_ptr.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/cuda_launcher.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/cuda_launcher.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/triple_chevron_launcher.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/triple_chevron_launcher.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/runtime_introspection.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/runtime_introspection.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/cuda_launch_config.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/cuda_launch_config.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/runtime_introspection.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/cuda_launcher/runtime_introspection.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/async.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/async.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/tuple_transform.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/tuple_transform.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/pointer_traits.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/pointer_traits.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/apply_from_tuple.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/apply_from_tuple.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/detail/is_contiguous_iterator.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/detail/is_contiguous_iterator.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/iterator.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/iterator.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/choose_sizes.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/choose_sizes.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/copy.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/copy.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/merge.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/merge.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/accumulate.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/accumulate.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/scan.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/scan.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/detail/stable_merge_sort.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/detail/stable_merge_sort.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/gather.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/gather.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/sort.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/sort.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/reduce.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/reduce.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/scatter.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/scatter.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/adjacent_difference.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/adjacent_difference.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/reduce_by_key.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/reduce_by_key.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/algorithm/for_each.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/algorithm/for_each.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/bulk.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/bulk.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/execution_policy.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/execution_policy.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/iterator/strided_iterator.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/iterator/strided_iterator.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/uninitialized.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/uninitialized.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/async.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/async.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/bulk/future.hpp" "$(@D)/cuda/include/thrust/system/cuda/detail/bulk/future.hpp" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/guarded_driver_types.h" "$(@D)/cuda/include/thrust/system/cuda/detail/guarded_driver_types.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/get_value.h" "$(@D)/cuda/include/thrust/system/cuda/detail/get_value.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/cuda/detail/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_if.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/logical.h" "$(@D)/cuda/include/thrust/system/cuda/detail/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/cuda/detail/iter_swap.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/merge.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/merge.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/inclusive_scan.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/inclusive_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/merge.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/block/merge.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/merging_sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/merging_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/exclusive_scan.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/exclusive_scan.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/reduce.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/copy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/block/odd_even_sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/block/odd_even_sort.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/par.h" "$(@D)/cuda/include/thrust/system/cuda/detail/par.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/copy_cross_system.h" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_cross_system.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/reduce_intervals.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce_intervals.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/cuda/detail/malloc_and_free.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/fill.h" "$(@D)/cuda/include/thrust/system/cuda/detail/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/set_difference.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/set_difference.inl" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/detail/transform.h" "$(@D)/cuda/include/thrust/system/cuda/detail/transform.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/experimental/pinned_allocator.h" "$(@D)/cuda/include/thrust/system/cuda/experimental/pinned_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/system/cuda/memory.h" "$(@D)/cuda/include/thrust/system/cuda/memory.h" && cp "/usr/local/cuda-8.0/include/thrust/remove.h" "$(@D)/cuda/include/thrust/remove.h" && cp "/usr/local/cuda-8.0/include/thrust/tabulate.h" "$(@D)/cuda/include/thrust/tabulate.h" && cp "/usr/local/cuda-8.0/include/thrust/for_each.h" "$(@D)/cuda/include/thrust/for_each.h" && cp "/usr/local/cuda-8.0/include/thrust/distance.h" "$(@D)/cuda/include/thrust/distance.h" && cp "/usr/local/cuda-8.0/include/thrust/reduce.h" "$(@D)/cuda/include/thrust/reduce.h" && cp "/usr/local/cuda-8.0/include/thrust/equal.h" "$(@D)/cuda/include/thrust/equal.h" && cp "/usr/local/cuda-8.0/include/thrust/complex.h" "$(@D)/cuda/include/thrust/complex.h" && cp "/usr/local/cuda-8.0/include/thrust/device_allocator.h" "$(@D)/cuda/include/thrust/device_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/copy.h" "$(@D)/cuda/include/thrust/copy.h" && cp "/usr/local/cuda-8.0/include/thrust/uninitialized_copy.h" "$(@D)/cuda/include/thrust/uninitialized_copy.h" && cp "/usr/local/cuda-8.0/include/thrust/device_reference.h" "$(@D)/cuda/include/thrust/device_reference.h" && cp "/usr/local/cuda-8.0/include/thrust/binary_search.h" "$(@D)/cuda/include/thrust/binary_search.h" && cp "/usr/local/cuda-8.0/include/thrust/set_operations.h" "$(@D)/cuda/include/thrust/set_operations.h" && cp "/usr/local/cuda-8.0/include/thrust/swap.h" "$(@D)/cuda/include/thrust/swap.h" && cp "/usr/local/cuda-8.0/include/thrust/mismatch.h" "$(@D)/cuda/include/thrust/mismatch.h" && cp "/usr/local/cuda-8.0/include/thrust/extrema.h" "$(@D)/cuda/include/thrust/extrema.h" && cp "/usr/local/cuda-8.0/include/thrust/count.h" "$(@D)/cuda/include/thrust/count.h" && cp "/usr/local/cuda-8.0/include/thrust/device_free.h" "$(@D)/cuda/include/thrust/device_free.h" && cp "/usr/local/cuda-8.0/include/thrust/random/discard_block_engine.h" "$(@D)/cuda/include/thrust/random/discard_block_engine.h" && cp "/usr/local/cuda-8.0/include/thrust/random/normal_distribution.h" "$(@D)/cuda/include/thrust/random/normal_distribution.h" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/linear_feedback_shift_engine_wordmask.h" "$(@D)/cuda/include/thrust/random/detail/linear_feedback_shift_engine_wordmask.h" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/subtract_with_carry_engine.inl" "$(@D)/cuda/include/thrust/random/detail/subtract_with_carry_engine.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/xor_combine_engine_max.h" "$(@D)/cuda/include/thrust/random/detail/xor_combine_engine_max.h" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/linear_congruential_engine_discard.h" "$(@D)/cuda/include/thrust/random/detail/linear_congruential_engine_discard.h" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/uniform_int_distribution.inl" "$(@D)/cuda/include/thrust/random/detail/uniform_int_distribution.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/discard_block_engine.inl" "$(@D)/cuda/include/thrust/random/detail/discard_block_engine.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/uniform_real_distribution.inl" "$(@D)/cuda/include/thrust/random/detail/uniform_real_distribution.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/random_core_access.h" "$(@D)/cuda/include/thrust/random/detail/random_core_access.h" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/mod.h" "$(@D)/cuda/include/thrust/random/detail/mod.h" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/linear_feedback_shift_engine.inl" "$(@D)/cuda/include/thrust/random/detail/linear_feedback_shift_engine.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/linear_congruential_engine.inl" "$(@D)/cuda/include/thrust/random/detail/linear_congruential_engine.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/xor_combine_engine.inl" "$(@D)/cuda/include/thrust/random/detail/xor_combine_engine.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/normal_distribution.inl" "$(@D)/cuda/include/thrust/random/detail/normal_distribution.inl" && cp "/usr/local/cuda-8.0/include/thrust/random/detail/normal_distribution_base.h" "$(@D)/cuda/include/thrust/random/detail/normal_distribution_base.h" && cp "/usr/local/cuda-8.0/include/thrust/random/uniform_int_distribution.h" "$(@D)/cuda/include/thrust/random/uniform_int_distribution.h" && cp "/usr/local/cuda-8.0/include/thrust/random/linear_feedback_shift_engine.h" "$(@D)/cuda/include/thrust/random/linear_feedback_shift_engine.h" && cp "/usr/local/cuda-8.0/include/thrust/random/xor_combine_engine.h" "$(@D)/cuda/include/thrust/random/xor_combine_engine.h" && cp "/usr/local/cuda-8.0/include/thrust/random/subtract_with_carry_engine.h" "$(@D)/cuda/include/thrust/random/subtract_with_carry_engine.h" && cp "/usr/local/cuda-8.0/include/thrust/random/linear_congruential_engine.h" "$(@D)/cuda/include/thrust/random/linear_congruential_engine.h" && cp "/usr/local/cuda-8.0/include/thrust/random/uniform_real_distribution.h" "$(@D)/cuda/include/thrust/random/uniform_real_distribution.h" && cp "/usr/local/cuda-8.0/include/thrust/functional.h" "$(@D)/cuda/include/thrust/functional.h" && cp "/usr/local/cuda-8.0/include/thrust/replace.h" "$(@D)/cuda/include/thrust/replace.h" && cp "/usr/local/cuda-8.0/include/thrust/device_new_allocator.h" "$(@D)/cuda/include/thrust/device_new_allocator.h" && cp "/usr/local/cuda-8.0/include/thrust/host_vector.h" "$(@D)/cuda/include/thrust/host_vector.h" && cp "/usr/local/cuda-8.0/include/thrust/version.h" "$(@D)/cuda/include/thrust/version.h" && cp "/usr/local/cuda-8.0/include/thrust/inner_product.h" "$(@D)/cuda/include/thrust/inner_product.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/iterator_traits.h" "$(@D)/cuda/include/thrust/iterator/iterator_traits.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/discard_iterator.h" "$(@D)/cuda/include/thrust/iterator/discard_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/retag.h" "$(@D)/cuda/include/thrust/iterator/retag.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/permutation_iterator.h" "$(@D)/cuda/include/thrust/iterator/permutation_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/transform_iterator.h" "$(@D)/cuda/include/thrust/iterator/transform_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/reverse_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/reverse_iterator.inl" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/zip_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/zip_iterator.inl" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/counting_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/counting_iterator.inl" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/distance_from_result.h" "$(@D)/cuda/include/thrust/iterator/detail/distance_from_result.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/host_system_tag.h" "$(@D)/cuda/include/thrust/iterator/detail/host_system_tag.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_traversal_tags.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_traversal_tags.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/retag.h" "$(@D)/cuda/include/thrust/iterator/detail/retag.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/tagged_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/tagged_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_traits.inl" "$(@D)/cuda/include/thrust/iterator/detail/iterator_traits.inl" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/minimum_category.h" "$(@D)/cuda/include/thrust/iterator/detail/minimum_category.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/discard_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/discard_iterator_base.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_category_to_traversal.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_category_to_traversal.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/zip_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/zip_iterator_base.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/normal_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/normal_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/join_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/join_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/device_system_tag.h" "$(@D)/cuda/include/thrust/iterator/detail/device_system_tag.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/universal_categories.h" "$(@D)/cuda/include/thrust/iterator/detail/universal_categories.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/reverse_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/reverse_iterator_base.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/minimum_system.h" "$(@D)/cuda/include/thrust/iterator/detail/minimum_system.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/tuple_of_iterator_references.h" "$(@D)/cuda/include/thrust/iterator/detail/tuple_of_iterator_references.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/is_iterator_category.h" "$(@D)/cuda/include/thrust/iterator/detail/is_iterator_category.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/permutation_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/permutation_iterator_base.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/any_assign.h" "$(@D)/cuda/include/thrust/iterator/detail/any_assign.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/any_system_tag.h" "$(@D)/cuda/include/thrust/iterator/detail/any_system_tag.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/is_trivial_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/is_trivial_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_category_to_system.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_category_to_system.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_adaptor_base.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_adaptor_base.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/constant_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/constant_iterator_base.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/transform_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/transform_iterator.inl" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_facade_category.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_facade_category.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/detail/iterator_category_with_system_and_traversal.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_category_with_system_and_traversal.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/constant_iterator.h" "$(@D)/cuda/include/thrust/iterator/constant_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/counting_iterator.h" "$(@D)/cuda/include/thrust/iterator/counting_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/iterator_adaptor.h" "$(@D)/cuda/include/thrust/iterator/iterator_adaptor.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/iterator_facade.h" "$(@D)/cuda/include/thrust/iterator/iterator_facade.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/iterator_categories.h" "$(@D)/cuda/include/thrust/iterator/iterator_categories.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/reverse_iterator.h" "$(@D)/cuda/include/thrust/iterator/reverse_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/iterator/zip_iterator.h" "$(@D)/cuda/include/thrust/iterator/zip_iterator.h" && cp "/usr/local/cuda-8.0/include/thrust/logical.h" "$(@D)/cuda/include/thrust/logical.h" && cp "/usr/local/cuda-8.0/include/thrust/tuple.h" "$(@D)/cuda/include/thrust/tuple.h" && cp "/usr/local/cuda-8.0/include/thrust/memory.h" "$(@D)/cuda/include/thrust/memory.h" && cp "/usr/local/cuda-8.0/include/thrust/random.h" "$(@D)/cuda/include/thrust/random.h" && cp "/usr/local/cuda-8.0/include/thrust/fill.h" "$(@D)/cuda/include/thrust/fill.h" && cp "/usr/local/cuda-8.0/include/thrust/transform.h" "$(@D)/cuda/include/thrust/transform.h" && cp "/usr/local/cuda-8.0/include/texture_types.h" "$(@D)/cuda/include/texture_types.h" && cp "/usr/local/cuda-8.0/include/nppversion.h" "$(@D)/cuda/include/nppversion.h" && cp "/usr/local/cuda-8.0/include/cuda_texture_types.h" "$(@D)/cuda/include/cuda_texture_types.h" && cp "/usr/local/cuda-8.0/include/fatbinary.h" "$(@D)/cuda/include/fatbinary.h" && cp "/usr/local/cuda-8.0/include/cublasXt.h" "$(@D)/cuda/include/cublasXt.h" && cp "/usr/local/cuda-8.0/include/cuda_fp16.h" "$(@D)/cuda/include/cuda_fp16.h" && cp "/usr/local/cuda-8.0/include/vector_functions.h" "$(@D)/cuda/include/vector_functions.h" && cp "/usr/local/cuda-8.0/include/cusparse.h" "$(@D)/cuda/include/cusparse.h" && cp "/usr/local/cuda-8.0/include/nppi_filtering_functions.h" "$(@D)/cuda/include/nppi_filtering_functions.h" && cp "/usr/local/cuda-8.0/include/nppi_morphological_operations.h" "$(@D)/cuda/include/nppi_morphological_operations.h" && cp "/usr/local/cuda-8.0/include/sobol_direction_vectors.h" "$(@D)/cuda/include/sobol_direction_vectors.h" && cp "/usr/local/cuda-8.0/include/nvblas.h" "$(@D)/cuda/include/nvblas.h" && cp "/usr/local/cuda-8.0/include/curand_mtgp32dc_p_11213.h" "$(@D)/cuda/include/curand_mtgp32dc_p_11213.h" && cp "/usr/local/cuda-8.0/include/nvcuvid.h" "$(@D)/cuda/include/nvcuvid.h" && cp "/usr/local/cuda-8.0/include/cuda_runtime_api.h" "$(@D)/cuda/include/cuda_runtime_api.h" && cp "/usr/local/cuda-8.0/include/curand_mtgp32_kernel.h" "$(@D)/cuda/include/curand_mtgp32_kernel.h" && cp "/usr/local/cuda-8.0/include/cublas_v2.h" "$(@D)/cuda/include/cublas_v2.h" && cp "/usr/local/cuda-8.0/include/builtin_types.h" "$(@D)/cuda/include/builtin_types.h" && cp "/usr/local/cuda-8.0/include/nppi_geometry_transforms.h" "$(@D)/cuda/include/nppi_geometry_transforms.h" && cp "/usr/local/cuda-8.0/include/npps_support_functions.h" "$(@D)/cuda/include/npps_support_functions.h" && cp "/usr/local/cuda-8.0/include/cufftw.h" "$(@D)/cuda/include/cufftw.h" && cp "/usr/local/cuda-8.0/include/cuda_device_runtime_api.h" "$(@D)/cuda/include/cuda_device_runtime_api.h" && cp "/usr/local/cuda-8.0/include/sm_30_intrinsics.hpp" "$(@D)/cuda/include/sm_30_intrinsics.hpp" && cp "/usr/local/cuda-8.0/include/vector_types.h" "$(@D)/cuda/include/vector_types.h" && cp "/usr/local/cuda-8.0/include/sm_35_atomic_functions.h" "$(@D)/cuda/include/sm_35_atomic_functions.h" && cp "/usr/local/cuda-8.0/include/sm_20_intrinsics.h" "$(@D)/cuda/include/sm_20_intrinsics.h" && cp "/usr/local/cuda-8.0/include/driver_types.h" "$(@D)/cuda/include/driver_types.h" && cp "/usr/local/cuda-8.0/include/nvToolsExtCudaRt.h" "$(@D)/cuda/include/nvToolsExtCudaRt.h" && cp "/usr/local/cuda-8.0/include/curand_globals.h" "$(@D)/cuda/include/curand_globals.h" && cp "/usr/local/cuda-8.0/include/device_atomic_functions.h" "$(@D)/cuda/include/device_atomic_functions.h" && cp "/usr/local/cuda-8.0/include/surface_types.h" "$(@D)/cuda/include/surface_types.h" && cp "/usr/local/cuda-8.0/include/nvrtc.h" "$(@D)/cuda/include/nvrtc.h" && cp "/usr/local/cuda-8.0/include/nppdefs.h" "$(@D)/cuda/include/nppdefs.h" && cp "/usr/local/cuda-8.0/include/sm_60_atomic_functions.h" "$(@D)/cuda/include/sm_60_atomic_functions.h" && cp "/usr/local/cuda-8.0/include/driver_functions.h" "$(@D)/cuda/include/driver_functions.h" && cp "/usr/local/cuda-8.0/include/cusolver_common.h" "$(@D)/cuda/include/cusolver_common.h" && cp "/usr/local/cuda-8.0/include/cublas.h" "$(@D)/cuda/include/cublas.h" && cp "/usr/local/cuda-8.0/include/curand_lognormal.h" "$(@D)/cuda/include/curand_lognormal.h" && cp "/usr/local/cuda-8.0/include/device_atomic_functions.hpp" "$(@D)/cuda/include/device_atomic_functions.hpp" && cp "/usr/local/cuda-8.0/include/crt/device_runtime.h" "$(@D)/cuda/include/crt/device_runtime.h" && cp "/usr/local/cuda-8.0/include/crt/storage_class.h" "$(@D)/cuda/include/crt/storage_class.h" && cp "/usr/local/cuda-8.0/include/crt/func_macro.h" "$(@D)/cuda/include/crt/func_macro.h" && cp "/usr/local/cuda-8.0/include/crt/host_runtime.h" "$(@D)/cuda/include/crt/host_runtime.h" && cp "/usr/local/cuda-8.0/include/nppi_arithmetic_and_logical_operations.h" "$(@D)/cuda/include/nppi_arithmetic_and_logical_operations.h" && cp "/usr/local/cuda-8.0/include/npps_arithmetic_and_logical_operations.h" "$(@D)/cuda/include/npps_arithmetic_and_logical_operations.h" && cp "/usr/local/cuda-8.0/include/nppi_computer_vision.h" "$(@D)/cuda/include/nppi_computer_vision.h" && cp "/usr/local/cuda-8.0/include/surface_functions.hpp" "$(@D)/cuda/include/surface_functions.hpp" && cp "/usr/local/cuda-8.0/include/surface_functions.h" "$(@D)/cuda/include/surface_functions.h" && cp "/usr/local/cuda-8.0/include/curand_normal_static.h" "$(@D)/cuda/include/curand_normal_static.h" && cp "/usr/local/cuda-8.0/include/curand.h" "$(@D)/cuda/include/curand.h" && cp "/usr/local/cuda-8.0/include/math_functions_dbl_ptx3.h" "$(@D)/cuda/include/math_functions_dbl_ptx3.h" && cp "/usr/local/cuda-8.0/include/curand_philox4x32_x.h" "$(@D)/cuda/include/curand_philox4x32_x.h" && cp "/usr/local/cuda-8.0/include/nppi_threshold_and_compare_operations.h" "$(@D)/cuda/include/nppi_threshold_and_compare_operations.h" && cp "/usr/local/cuda-8.0/include/nvml.h" "$(@D)/cuda/include/nvml.h" && cp "/usr/local/cuda-8.0/include/npps.h" "$(@D)/cuda/include/npps.h" && cp "/usr/local/cuda-8.0/include/cuda_vdpau_interop.h" "$(@D)/cuda/include/cuda_vdpau_interop.h" && cp "/usr/local/cuda-8.0/include/sm_61_intrinsics.hpp" "$(@D)/cuda/include/sm_61_intrinsics.hpp" && cp "/usr/local/cuda-8.0/include/cublas_api.h" "$(@D)/cuda/include/cublas_api.h" && cp "/usr/local/cuda-8.0/include/nppi_color_conversion.h" "$(@D)/cuda/include/nppi_color_conversion.h" && cp "/usr/local/cuda-8.0/include/math_functions_dbl_ptx3.hpp" "$(@D)/cuda/include/math_functions_dbl_ptx3.hpp" && cp "/usr/local/cuda-8.0/include/nppcore.h" "$(@D)/cuda/include/nppcore.h" && cp "/usr/local/cuda-8.0/include/cudaGL.h" "$(@D)/cuda/include/cudaGL.h" && cp "/usr/local/cuda-8.0/include/fatBinaryCtl.h" "$(@D)/cuda/include/fatBinaryCtl.h" && cp "/usr/local/cuda-8.0/include/npps_statistics_functions.h" "$(@D)/cuda/include/npps_statistics_functions.h" && cp "/usr/local/cuda-8.0/include/cudaVDPAU.h" "$(@D)/cuda/include/cudaVDPAU.h" && cp "/usr/local/cuda-8.0/include/curand_poisson.h" "$(@D)/cuda/include/curand_poisson.h" && cp "/usr/local/cuda-8.0/include/cusolverDn.h" "$(@D)/cuda/include/cusolverDn.h" && cp "/usr/local/cuda-8.0/include/cuda_profiler_api.h" "$(@D)/cuda/include/cuda_profiler_api.h" && cp "/usr/local/cuda-8.0/include/sm_20_atomic_functions.h" "$(@D)/cuda/include/sm_20_atomic_functions.h" && cp "/usr/local/cuda-8.0/include/nvfunctional" "$(@D)/cuda/include/nvfunctional" +if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-9.0/include/CL/cl.h" "$(@D)/cuda/include/CL/cl.h" && cp "/usr/local/cuda-9.0/include/CL/cl.hpp" "$(@D)/cuda/include/CL/cl.hpp" && cp "/usr/local/cuda-9.0/include/CL/cl_egl.h" "$(@D)/cuda/include/CL/cl_egl.h" && cp "/usr/local/cuda-9.0/include/CL/cl_ext.h" "$(@D)/cuda/include/CL/cl_ext.h" && cp "/usr/local/cuda-9.0/include/CL/cl_gl.h" "$(@D)/cuda/include/CL/cl_gl.h" && cp "/usr/local/cuda-9.0/include/CL/cl_gl_ext.h" "$(@D)/cuda/include/CL/cl_gl_ext.h" && cp "/usr/local/cuda-9.0/include/CL/cl_platform.h" "$(@D)/cuda/include/CL/cl_platform.h" && cp "/usr/local/cuda-9.0/include/CL/opencl.h" "$(@D)/cuda/include/CL/opencl.h" && cp "/usr/local/cuda-9.0/include/builtin_types.h" "$(@D)/cuda/include/builtin_types.h" && cp "/usr/local/cuda-9.0/include/channel_descriptor.h" "$(@D)/cuda/include/channel_descriptor.h" && cp "/usr/local/cuda-9.0/include/common_functions.h" "$(@D)/cuda/include/common_functions.h" && cp "/usr/local/cuda-9.0/include/cooperative_groups.h" "$(@D)/cuda/include/cooperative_groups.h" && cp "/usr/local/cuda-9.0/include/cooperative_groups_helpers.h" "$(@D)/cuda/include/cooperative_groups_helpers.h" && cp "/usr/local/cuda-9.0/include/crt/common_functions.h" "$(@D)/cuda/include/crt/common_functions.h" && cp "/usr/local/cuda-9.0/include/crt/device_double_functions.h" "$(@D)/cuda/include/crt/device_double_functions.h" && cp "/usr/local/cuda-9.0/include/crt/device_double_functions.hpp" "$(@D)/cuda/include/crt/device_double_functions.hpp" && cp "/usr/local/cuda-9.0/include/crt/device_functions.h" "$(@D)/cuda/include/crt/device_functions.h" && cp "/usr/local/cuda-9.0/include/crt/device_functions.hpp" "$(@D)/cuda/include/crt/device_functions.hpp" && cp "/usr/local/cuda-9.0/include/crt/func_macro.h" "$(@D)/cuda/include/crt/func_macro.h" && cp "/usr/local/cuda-9.0/include/crt/host_config.h" "$(@D)/cuda/include/crt/host_config.h" && cp "/usr/local/cuda-9.0/include/crt/host_defines.h" "$(@D)/cuda/include/crt/host_defines.h" && cp "/usr/local/cuda-9.0/include/crt/host_runtime.h" "$(@D)/cuda/include/crt/host_runtime.h" && cp "/usr/local/cuda-9.0/include/crt/math_functions.h" "$(@D)/cuda/include/crt/math_functions.h" && cp "/usr/local/cuda-9.0/include/crt/math_functions.hpp" "$(@D)/cuda/include/crt/math_functions.hpp" && cp "/usr/local/cuda-9.0/include/crt/mma.h" "$(@D)/cuda/include/crt/mma.h" && cp "/usr/local/cuda-9.0/include/crt/mma.hpp" "$(@D)/cuda/include/crt/mma.hpp" && cp "/usr/local/cuda-9.0/include/crt/nvfunctional" "$(@D)/cuda/include/crt/nvfunctional" && cp "/usr/local/cuda-9.0/include/crt/sm_70_rt.h" "$(@D)/cuda/include/crt/sm_70_rt.h" && cp "/usr/local/cuda-9.0/include/crt/sm_70_rt.hpp" "$(@D)/cuda/include/crt/sm_70_rt.hpp" && cp "/usr/local/cuda-9.0/include/crt/storage_class.h" "$(@D)/cuda/include/crt/storage_class.h" && cp "/usr/local/cuda-9.0/include/cuComplex.h" "$(@D)/cuda/include/cuComplex.h" && cp "/usr/local/cuda-9.0/include/cublas.h" "$(@D)/cuda/include/cublas.h" && cp "/usr/local/cuda-9.0/include/cublasXt.h" "$(@D)/cuda/include/cublasXt.h" && cp "/usr/local/cuda-9.0/include/cublas_api.h" "$(@D)/cuda/include/cublas_api.h" && cp "/usr/local/cuda-9.0/include/cublas_v2.h" "$(@D)/cuda/include/cublas_v2.h" && cp "/usr/local/cuda-9.0/include/cuda.h" "$(@D)/cuda/include/cuda.h" && cp "/usr/local/cuda-9.0/include/cudaEGL.h" "$(@D)/cuda/include/cudaEGL.h" && cp "/usr/local/cuda-9.0/include/cudaGL.h" "$(@D)/cuda/include/cudaGL.h" && cp "/usr/local/cuda-9.0/include/cudaProfiler.h" "$(@D)/cuda/include/cudaProfiler.h" && cp "/usr/local/cuda-9.0/include/cudaVDPAU.h" "$(@D)/cuda/include/cudaVDPAU.h" && cp "/usr/local/cuda-9.0/include/cuda_device_runtime_api.h" "$(@D)/cuda/include/cuda_device_runtime_api.h" && cp "/usr/local/cuda-9.0/include/cuda_fp16.h" "$(@D)/cuda/include/cuda_fp16.h" && cp "/usr/local/cuda-9.0/include/cuda_fp16.hpp" "$(@D)/cuda/include/cuda_fp16.hpp" && cp "/usr/local/cuda-9.0/include/cuda_gl_interop.h" "$(@D)/cuda/include/cuda_gl_interop.h" && cp "/usr/local/cuda-9.0/include/cuda_occupancy.h" "$(@D)/cuda/include/cuda_occupancy.h" && cp "/usr/local/cuda-9.0/include/cuda_profiler_api.h" "$(@D)/cuda/include/cuda_profiler_api.h" && cp "/usr/local/cuda-9.0/include/cuda_runtime.h" "$(@D)/cuda/include/cuda_runtime.h" && cp "/usr/local/cuda-9.0/include/cuda_runtime_api.h" "$(@D)/cuda/include/cuda_runtime_api.h" && cp "/usr/local/cuda-9.0/include/cuda_surface_types.h" "$(@D)/cuda/include/cuda_surface_types.h" && cp "/usr/local/cuda-9.0/include/cuda_texture_types.h" "$(@D)/cuda/include/cuda_texture_types.h" && cp "/usr/local/cuda-9.0/include/cuda_vdpau_interop.h" "$(@D)/cuda/include/cuda_vdpau_interop.h" && cp "/usr/local/cuda-9.0/include/cudalibxt.h" "$(@D)/cuda/include/cudalibxt.h" && cp "/usr/local/cuda-9.0/include/cudnn.h" "$(@D)/cuda/include/cudnn.h" && cp "/usr/local/cuda-9.0/include/cufft.h" "$(@D)/cuda/include/cufft.h" && cp "/usr/local/cuda-9.0/include/cufftXt.h" "$(@D)/cuda/include/cufftXt.h" && cp "/usr/local/cuda-9.0/include/cufftw.h" "$(@D)/cuda/include/cufftw.h" && cp "/usr/local/cuda-9.0/include/curand.h" "$(@D)/cuda/include/curand.h" && cp "/usr/local/cuda-9.0/include/curand_discrete.h" "$(@D)/cuda/include/curand_discrete.h" && cp "/usr/local/cuda-9.0/include/curand_discrete2.h" "$(@D)/cuda/include/curand_discrete2.h" && cp "/usr/local/cuda-9.0/include/curand_globals.h" "$(@D)/cuda/include/curand_globals.h" && cp "/usr/local/cuda-9.0/include/curand_kernel.h" "$(@D)/cuda/include/curand_kernel.h" && cp "/usr/local/cuda-9.0/include/curand_lognormal.h" "$(@D)/cuda/include/curand_lognormal.h" && cp "/usr/local/cuda-9.0/include/curand_mrg32k3a.h" "$(@D)/cuda/include/curand_mrg32k3a.h" && cp "/usr/local/cuda-9.0/include/curand_mtgp32.h" "$(@D)/cuda/include/curand_mtgp32.h" && cp "/usr/local/cuda-9.0/include/curand_mtgp32_host.h" "$(@D)/cuda/include/curand_mtgp32_host.h" && cp "/usr/local/cuda-9.0/include/curand_mtgp32_kernel.h" "$(@D)/cuda/include/curand_mtgp32_kernel.h" && cp "/usr/local/cuda-9.0/include/curand_mtgp32dc_p_11213.h" "$(@D)/cuda/include/curand_mtgp32dc_p_11213.h" && cp "/usr/local/cuda-9.0/include/curand_normal.h" "$(@D)/cuda/include/curand_normal.h" && cp "/usr/local/cuda-9.0/include/curand_normal_static.h" "$(@D)/cuda/include/curand_normal_static.h" && cp "/usr/local/cuda-9.0/include/curand_philox4x32_x.h" "$(@D)/cuda/include/curand_philox4x32_x.h" && cp "/usr/local/cuda-9.0/include/curand_poisson.h" "$(@D)/cuda/include/curand_poisson.h" && cp "/usr/local/cuda-9.0/include/curand_precalc.h" "$(@D)/cuda/include/curand_precalc.h" && cp "/usr/local/cuda-9.0/include/curand_uniform.h" "$(@D)/cuda/include/curand_uniform.h" && cp "/usr/local/cuda-9.0/include/cusolverDn.h" "$(@D)/cuda/include/cusolverDn.h" && cp "/usr/local/cuda-9.0/include/cusolverRf.h" "$(@D)/cuda/include/cusolverRf.h" && cp "/usr/local/cuda-9.0/include/cusolverSp.h" "$(@D)/cuda/include/cusolverSp.h" && cp "/usr/local/cuda-9.0/include/cusolverSp_LOWLEVEL_PREVIEW.h" "$(@D)/cuda/include/cusolverSp_LOWLEVEL_PREVIEW.h" && cp "/usr/local/cuda-9.0/include/cusolver_common.h" "$(@D)/cuda/include/cusolver_common.h" && cp "/usr/local/cuda-9.0/include/cusparse.h" "$(@D)/cuda/include/cusparse.h" && cp "/usr/local/cuda-9.0/include/cusparse_v2.h" "$(@D)/cuda/include/cusparse_v2.h" && cp "/usr/local/cuda-9.0/include/device_atomic_functions.h" "$(@D)/cuda/include/device_atomic_functions.h" && cp "/usr/local/cuda-9.0/include/device_atomic_functions.hpp" "$(@D)/cuda/include/device_atomic_functions.hpp" && cp "/usr/local/cuda-9.0/include/device_double_functions.h" "$(@D)/cuda/include/device_double_functions.h" && cp "/usr/local/cuda-9.0/include/device_double_functions.hpp" "$(@D)/cuda/include/device_double_functions.hpp" && cp "/usr/local/cuda-9.0/include/device_functions.h" "$(@D)/cuda/include/device_functions.h" && cp "/usr/local/cuda-9.0/include/device_functions.hpp" "$(@D)/cuda/include/device_functions.hpp" && cp "/usr/local/cuda-9.0/include/device_functions_decls.h" "$(@D)/cuda/include/device_functions_decls.h" && cp "/usr/local/cuda-9.0/include/device_launch_parameters.h" "$(@D)/cuda/include/device_launch_parameters.h" && cp "/usr/local/cuda-9.0/include/device_types.h" "$(@D)/cuda/include/device_types.h" && cp "/usr/local/cuda-9.0/include/driver_functions.h" "$(@D)/cuda/include/driver_functions.h" && cp "/usr/local/cuda-9.0/include/driver_types.h" "$(@D)/cuda/include/driver_types.h" && cp "/usr/local/cuda-9.0/include/dynlink_cuda.h" "$(@D)/cuda/include/dynlink_cuda.h" && cp "/usr/local/cuda-9.0/include/dynlink_cuda_cuda.h" "$(@D)/cuda/include/dynlink_cuda_cuda.h" && cp "/usr/local/cuda-9.0/include/dynlink_cuviddec.h" "$(@D)/cuda/include/dynlink_cuviddec.h" && cp "/usr/local/cuda-9.0/include/dynlink_nvcuvid.h" "$(@D)/cuda/include/dynlink_nvcuvid.h" && cp "/usr/local/cuda-9.0/include/fatBinaryCtl.h" "$(@D)/cuda/include/fatBinaryCtl.h" && cp "/usr/local/cuda-9.0/include/fatbinary.h" "$(@D)/cuda/include/fatbinary.h" && cp "/usr/local/cuda-9.0/include/host_config.h" "$(@D)/cuda/include/host_config.h" && cp "/usr/local/cuda-9.0/include/host_defines.h" "$(@D)/cuda/include/host_defines.h" && cp "/usr/local/cuda-9.0/include/library_types.h" "$(@D)/cuda/include/library_types.h" && cp "/usr/local/cuda-9.0/include/math_constants.h" "$(@D)/cuda/include/math_constants.h" && cp "/usr/local/cuda-9.0/include/math_functions.h" "$(@D)/cuda/include/math_functions.h" && cp "/usr/local/cuda-9.0/include/math_functions.hpp" "$(@D)/cuda/include/math_functions.hpp" && cp "/usr/local/cuda-9.0/include/math_functions_dbl_ptx3.h" "$(@D)/cuda/include/math_functions_dbl_ptx3.h" && cp "/usr/local/cuda-9.0/include/math_functions_dbl_ptx3.hpp" "$(@D)/cuda/include/math_functions_dbl_ptx3.hpp" && cp "/usr/local/cuda-9.0/include/mma.h" "$(@D)/cuda/include/mma.h" && cp "/usr/local/cuda-9.0/include/npp.h" "$(@D)/cuda/include/npp.h" && cp "/usr/local/cuda-9.0/include/nppcore.h" "$(@D)/cuda/include/nppcore.h" && cp "/usr/local/cuda-9.0/include/nppdefs.h" "$(@D)/cuda/include/nppdefs.h" && cp "/usr/local/cuda-9.0/include/nppi.h" "$(@D)/cuda/include/nppi.h" && cp "/usr/local/cuda-9.0/include/nppi_arithmetic_and_logical_operations.h" "$(@D)/cuda/include/nppi_arithmetic_and_logical_operations.h" && cp "/usr/local/cuda-9.0/include/nppi_color_conversion.h" "$(@D)/cuda/include/nppi_color_conversion.h" && cp "/usr/local/cuda-9.0/include/nppi_compression_functions.h" "$(@D)/cuda/include/nppi_compression_functions.h" && cp "/usr/local/cuda-9.0/include/nppi_computer_vision.h" "$(@D)/cuda/include/nppi_computer_vision.h" && cp "/usr/local/cuda-9.0/include/nppi_data_exchange_and_initialization.h" "$(@D)/cuda/include/nppi_data_exchange_and_initialization.h" && cp "/usr/local/cuda-9.0/include/nppi_filtering_functions.h" "$(@D)/cuda/include/nppi_filtering_functions.h" && cp "/usr/local/cuda-9.0/include/nppi_geometry_transforms.h" "$(@D)/cuda/include/nppi_geometry_transforms.h" && cp "/usr/local/cuda-9.0/include/nppi_linear_transforms.h" "$(@D)/cuda/include/nppi_linear_transforms.h" && cp "/usr/local/cuda-9.0/include/nppi_morphological_operations.h" "$(@D)/cuda/include/nppi_morphological_operations.h" && cp "/usr/local/cuda-9.0/include/nppi_statistics_functions.h" "$(@D)/cuda/include/nppi_statistics_functions.h" && cp "/usr/local/cuda-9.0/include/nppi_support_functions.h" "$(@D)/cuda/include/nppi_support_functions.h" && cp "/usr/local/cuda-9.0/include/nppi_threshold_and_compare_operations.h" "$(@D)/cuda/include/nppi_threshold_and_compare_operations.h" && cp "/usr/local/cuda-9.0/include/npps.h" "$(@D)/cuda/include/npps.h" && cp "/usr/local/cuda-9.0/include/npps_arithmetic_and_logical_operations.h" "$(@D)/cuda/include/npps_arithmetic_and_logical_operations.h" && cp "/usr/local/cuda-9.0/include/npps_conversion_functions.h" "$(@D)/cuda/include/npps_conversion_functions.h" && cp "/usr/local/cuda-9.0/include/npps_filtering_functions.h" "$(@D)/cuda/include/npps_filtering_functions.h" && cp "/usr/local/cuda-9.0/include/npps_initialization.h" "$(@D)/cuda/include/npps_initialization.h" && cp "/usr/local/cuda-9.0/include/npps_statistics_functions.h" "$(@D)/cuda/include/npps_statistics_functions.h" && cp "/usr/local/cuda-9.0/include/npps_support_functions.h" "$(@D)/cuda/include/npps_support_functions.h" && cp "/usr/local/cuda-9.0/include/nppversion.h" "$(@D)/cuda/include/nppversion.h" && cp "/usr/local/cuda-9.0/include/nvToolsExt.h" "$(@D)/cuda/include/nvToolsExt.h" && cp "/usr/local/cuda-9.0/include/nvToolsExtCuda.h" "$(@D)/cuda/include/nvToolsExtCuda.h" && cp "/usr/local/cuda-9.0/include/nvToolsExtCudaRt.h" "$(@D)/cuda/include/nvToolsExtCudaRt.h" && cp "/usr/local/cuda-9.0/include/nvToolsExtMeta.h" "$(@D)/cuda/include/nvToolsExtMeta.h" && cp "/usr/local/cuda-9.0/include/nvToolsExtSync.h" "$(@D)/cuda/include/nvToolsExtSync.h" && cp "/usr/local/cuda-9.0/include/nvblas.h" "$(@D)/cuda/include/nvblas.h" && cp "/usr/local/cuda-9.0/include/nvfunctional" "$(@D)/cuda/include/nvfunctional" && cp "/usr/local/cuda-9.0/include/nvgraph.h" "$(@D)/cuda/include/nvgraph.h" && cp "/usr/local/cuda-9.0/include/nvml.h" "$(@D)/cuda/include/nvml.h" && cp "/usr/local/cuda-9.0/include/nvrtc.h" "$(@D)/cuda/include/nvrtc.h" && cp "/usr/local/cuda-9.0/include/sm_20_atomic_functions.h" "$(@D)/cuda/include/sm_20_atomic_functions.h" && cp "/usr/local/cuda-9.0/include/sm_20_atomic_functions.hpp" "$(@D)/cuda/include/sm_20_atomic_functions.hpp" && cp "/usr/local/cuda-9.0/include/sm_20_intrinsics.h" "$(@D)/cuda/include/sm_20_intrinsics.h" && cp "/usr/local/cuda-9.0/include/sm_20_intrinsics.hpp" "$(@D)/cuda/include/sm_20_intrinsics.hpp" && cp "/usr/local/cuda-9.0/include/sm_30_intrinsics.h" "$(@D)/cuda/include/sm_30_intrinsics.h" && cp "/usr/local/cuda-9.0/include/sm_30_intrinsics.hpp" "$(@D)/cuda/include/sm_30_intrinsics.hpp" && cp "/usr/local/cuda-9.0/include/sm_32_atomic_functions.h" "$(@D)/cuda/include/sm_32_atomic_functions.h" && cp "/usr/local/cuda-9.0/include/sm_32_atomic_functions.hpp" "$(@D)/cuda/include/sm_32_atomic_functions.hpp" && cp "/usr/local/cuda-9.0/include/sm_32_intrinsics.h" "$(@D)/cuda/include/sm_32_intrinsics.h" && cp "/usr/local/cuda-9.0/include/sm_32_intrinsics.hpp" "$(@D)/cuda/include/sm_32_intrinsics.hpp" && cp "/usr/local/cuda-9.0/include/sm_35_atomic_functions.h" "$(@D)/cuda/include/sm_35_atomic_functions.h" && cp "/usr/local/cuda-9.0/include/sm_35_intrinsics.h" "$(@D)/cuda/include/sm_35_intrinsics.h" && cp "/usr/local/cuda-9.0/include/sm_60_atomic_functions.h" "$(@D)/cuda/include/sm_60_atomic_functions.h" && cp "/usr/local/cuda-9.0/include/sm_60_atomic_functions.hpp" "$(@D)/cuda/include/sm_60_atomic_functions.hpp" && cp "/usr/local/cuda-9.0/include/sm_61_intrinsics.h" "$(@D)/cuda/include/sm_61_intrinsics.h" && cp "/usr/local/cuda-9.0/include/sm_61_intrinsics.hpp" "$(@D)/cuda/include/sm_61_intrinsics.hpp" && cp "/usr/local/cuda-9.0/include/sobol_direction_vectors.h" "$(@D)/cuda/include/sobol_direction_vectors.h" && cp "/usr/local/cuda-9.0/include/surface_functions.h" "$(@D)/cuda/include/surface_functions.h" && cp "/usr/local/cuda-9.0/include/surface_functions.hpp" "$(@D)/cuda/include/surface_functions.hpp" && cp "/usr/local/cuda-9.0/include/surface_indirect_functions.h" "$(@D)/cuda/include/surface_indirect_functions.h" && cp "/usr/local/cuda-9.0/include/surface_indirect_functions.hpp" "$(@D)/cuda/include/surface_indirect_functions.hpp" && cp "/usr/local/cuda-9.0/include/surface_types.h" "$(@D)/cuda/include/surface_types.h" && cp "/usr/local/cuda-9.0/include/texture_fetch_functions.h" "$(@D)/cuda/include/texture_fetch_functions.h" && cp "/usr/local/cuda-9.0/include/texture_fetch_functions.hpp" "$(@D)/cuda/include/texture_fetch_functions.hpp" && cp "/usr/local/cuda-9.0/include/texture_indirect_functions.h" "$(@D)/cuda/include/texture_indirect_functions.h" && cp "/usr/local/cuda-9.0/include/texture_indirect_functions.hpp" "$(@D)/cuda/include/texture_indirect_functions.hpp" && cp "/usr/local/cuda-9.0/include/texture_types.h" "$(@D)/cuda/include/texture_types.h" && cp "/usr/local/cuda-9.0/include/thrust/adjacent_difference.h" "$(@D)/cuda/include/thrust/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/advance.h" "$(@D)/cuda/include/thrust/advance.h" && cp "/usr/local/cuda-9.0/include/thrust/binary_search.h" "$(@D)/cuda/include/thrust/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/complex.h" "$(@D)/cuda/include/thrust/complex.h" && cp "/usr/local/cuda-9.0/include/thrust/copy.h" "$(@D)/cuda/include/thrust/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/count.h" "$(@D)/cuda/include/thrust/count.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/adjacent_difference.inl" "$(@D)/cuda/include/thrust/detail/adjacent_difference.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/advance.inl" "$(@D)/cuda/include/thrust/detail/advance.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/allocator_traits.h" "$(@D)/cuda/include/thrust/detail/allocator/allocator_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/allocator_traits.inl" "$(@D)/cuda/include/thrust/detail/allocator/allocator_traits.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/copy_construct_range.h" "$(@D)/cuda/include/thrust/detail/allocator/copy_construct_range.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/copy_construct_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/copy_construct_range.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/default_construct_range.h" "$(@D)/cuda/include/thrust/detail/allocator/default_construct_range.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/default_construct_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/default_construct_range.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/destroy_range.h" "$(@D)/cuda/include/thrust/detail/allocator/destroy_range.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/destroy_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/destroy_range.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/fill_construct_range.h" "$(@D)/cuda/include/thrust/detail/allocator/fill_construct_range.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/fill_construct_range.inl" "$(@D)/cuda/include/thrust/detail/allocator/fill_construct_range.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/malloc_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/malloc_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/malloc_allocator.inl" "$(@D)/cuda/include/thrust/detail/allocator/malloc_allocator.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/no_throw_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/no_throw_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/tagged_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/tagged_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/tagged_allocator.inl" "$(@D)/cuda/include/thrust/detail/allocator/tagged_allocator.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/temporary_allocator.h" "$(@D)/cuda/include/thrust/detail/allocator/temporary_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/allocator/temporary_allocator.inl" "$(@D)/cuda/include/thrust/detail/allocator/temporary_allocator.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/binary_search.inl" "$(@D)/cuda/include/thrust/detail/binary_search.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/arithmetic.h" "$(@D)/cuda/include/thrust/detail/complex/arithmetic.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/c99math.h" "$(@D)/cuda/include/thrust/detail/complex/c99math.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/catrig.h" "$(@D)/cuda/include/thrust/detail/complex/catrig.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/catrigf.h" "$(@D)/cuda/include/thrust/detail/complex/catrigf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/ccosh.h" "$(@D)/cuda/include/thrust/detail/complex/ccosh.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/ccoshf.h" "$(@D)/cuda/include/thrust/detail/complex/ccoshf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/cexp.h" "$(@D)/cuda/include/thrust/detail/complex/cexp.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/cexpf.h" "$(@D)/cuda/include/thrust/detail/complex/cexpf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/clog.h" "$(@D)/cuda/include/thrust/detail/complex/clog.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/clogf.h" "$(@D)/cuda/include/thrust/detail/complex/clogf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/complex.inl" "$(@D)/cuda/include/thrust/detail/complex/complex.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/cpow.h" "$(@D)/cuda/include/thrust/detail/complex/cpow.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/cpowf.h" "$(@D)/cuda/include/thrust/detail/complex/cpowf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/cproj.h" "$(@D)/cuda/include/thrust/detail/complex/cproj.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/csinh.h" "$(@D)/cuda/include/thrust/detail/complex/csinh.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/csinhf.h" "$(@D)/cuda/include/thrust/detail/complex/csinhf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/csqrt.h" "$(@D)/cuda/include/thrust/detail/complex/csqrt.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/csqrtf.h" "$(@D)/cuda/include/thrust/detail/complex/csqrtf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/ctanh.h" "$(@D)/cuda/include/thrust/detail/complex/ctanh.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/ctanhf.h" "$(@D)/cuda/include/thrust/detail/complex/ctanhf.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/math_private.h" "$(@D)/cuda/include/thrust/detail/complex/math_private.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/complex/stream.h" "$(@D)/cuda/include/thrust/detail/complex/stream.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config.h" "$(@D)/cuda/include/thrust/detail/config.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/compiler.h" "$(@D)/cuda/include/thrust/detail/config/compiler.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/compiler_fence.h" "$(@D)/cuda/include/thrust/detail/config/compiler_fence.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/config.h" "$(@D)/cuda/include/thrust/detail/config/config.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/debug.h" "$(@D)/cuda/include/thrust/detail/config/debug.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/device_system.h" "$(@D)/cuda/include/thrust/detail/config/device_system.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/exec_check_disable.h" "$(@D)/cuda/include/thrust/detail/config/exec_check_disable.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/forceinline.h" "$(@D)/cuda/include/thrust/detail/config/forceinline.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/global_workarounds.h" "$(@D)/cuda/include/thrust/detail/config/global_workarounds.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/host_device.h" "$(@D)/cuda/include/thrust/detail/config/host_device.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/host_system.h" "$(@D)/cuda/include/thrust/detail/config/host_system.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/config/simple_defines.h" "$(@D)/cuda/include/thrust/detail/config/simple_defines.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/contiguous_storage.h" "$(@D)/cuda/include/thrust/detail/contiguous_storage.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/contiguous_storage.inl" "$(@D)/cuda/include/thrust/detail/contiguous_storage.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/copy.h" "$(@D)/cuda/include/thrust/detail/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/copy.inl" "$(@D)/cuda/include/thrust/detail/copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/copy_if.h" "$(@D)/cuda/include/thrust/detail/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/copy_if.inl" "$(@D)/cuda/include/thrust/detail/copy_if.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/count.inl" "$(@D)/cuda/include/thrust/detail/count.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/cstdint.h" "$(@D)/cuda/include/thrust/detail/cstdint.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_delete.inl" "$(@D)/cuda/include/thrust/detail/device_delete.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_free.inl" "$(@D)/cuda/include/thrust/detail/device_free.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_malloc.inl" "$(@D)/cuda/include/thrust/detail/device_malloc.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_new.inl" "$(@D)/cuda/include/thrust/detail/device_new.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_ptr.inl" "$(@D)/cuda/include/thrust/detail/device_ptr.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_reference.inl" "$(@D)/cuda/include/thrust/detail/device_reference.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/device_vector.inl" "$(@D)/cuda/include/thrust/detail/device_vector.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/dispatch/is_trivial_copy.h" "$(@D)/cuda/include/thrust/detail/dispatch/is_trivial_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/distance.inl" "$(@D)/cuda/include/thrust/detail/distance.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/equal.inl" "$(@D)/cuda/include/thrust/detail/equal.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/execute_with_allocator.h" "$(@D)/cuda/include/thrust/detail/execute_with_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/execution_policy.h" "$(@D)/cuda/include/thrust/detail/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/extrema.inl" "$(@D)/cuda/include/thrust/detail/extrema.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/fill.inl" "$(@D)/cuda/include/thrust/detail/fill.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/find.inl" "$(@D)/cuda/include/thrust/detail/find.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/for_each.inl" "$(@D)/cuda/include/thrust/detail/for_each.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/function.h" "$(@D)/cuda/include/thrust/detail/function.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional.inl" "$(@D)/cuda/include/thrust/detail/functional.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/actor.h" "$(@D)/cuda/include/thrust/detail/functional/actor.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/actor.inl" "$(@D)/cuda/include/thrust/detail/functional/actor.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/argument.h" "$(@D)/cuda/include/thrust/detail/functional/argument.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/composite.h" "$(@D)/cuda/include/thrust/detail/functional/composite.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/arithmetic_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/arithmetic_operators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/assignment_operator.h" "$(@D)/cuda/include/thrust/detail/functional/operators/assignment_operator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/bitwise_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/bitwise_operators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/compound_assignment_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/compound_assignment_operators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/logical_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/logical_operators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/operator_adaptors.h" "$(@D)/cuda/include/thrust/detail/functional/operators/operator_adaptors.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/operators/relational_operators.h" "$(@D)/cuda/include/thrust/detail/functional/operators/relational_operators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/placeholder.h" "$(@D)/cuda/include/thrust/detail/functional/placeholder.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/functional/value.h" "$(@D)/cuda/include/thrust/detail/functional/value.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/gather.inl" "$(@D)/cuda/include/thrust/detail/gather.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/generate.inl" "$(@D)/cuda/include/thrust/detail/generate.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/get_iterator_value.h" "$(@D)/cuda/include/thrust/detail/get_iterator_value.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/host_vector.inl" "$(@D)/cuda/include/thrust/detail/host_vector.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/inner_product.inl" "$(@D)/cuda/include/thrust/detail/inner_product.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/integer_math.h" "$(@D)/cuda/include/thrust/detail/integer_math.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/integer_traits.h" "$(@D)/cuda/include/thrust/detail/integer_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/internal_functional.h" "$(@D)/cuda/include/thrust/detail/internal_functional.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/logical.inl" "$(@D)/cuda/include/thrust/detail/logical.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/detail/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/merge.inl" "$(@D)/cuda/include/thrust/detail/merge.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/minmax.h" "$(@D)/cuda/include/thrust/detail/minmax.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/mismatch.inl" "$(@D)/cuda/include/thrust/detail/mismatch.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/mpl/math.h" "$(@D)/cuda/include/thrust/detail/mpl/math.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/numeric_traits.h" "$(@D)/cuda/include/thrust/detail/numeric_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/overlapped_copy.h" "$(@D)/cuda/include/thrust/detail/overlapped_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/pair.inl" "$(@D)/cuda/include/thrust/detail/pair.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/partition.inl" "$(@D)/cuda/include/thrust/detail/partition.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/pointer.h" "$(@D)/cuda/include/thrust/detail/pointer.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/pointer.inl" "$(@D)/cuda/include/thrust/detail/pointer.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/range/head_flags.h" "$(@D)/cuda/include/thrust/detail/range/head_flags.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/range/tail_flags.h" "$(@D)/cuda/include/thrust/detail/range/tail_flags.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/raw_pointer_cast.h" "$(@D)/cuda/include/thrust/detail/raw_pointer_cast.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/raw_reference_cast.h" "$(@D)/cuda/include/thrust/detail/raw_reference_cast.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/reduce.inl" "$(@D)/cuda/include/thrust/detail/reduce.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/reference.h" "$(@D)/cuda/include/thrust/detail/reference.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/reference.inl" "$(@D)/cuda/include/thrust/detail/reference.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/reference_forward_declaration.h" "$(@D)/cuda/include/thrust/detail/reference_forward_declaration.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/remove.inl" "$(@D)/cuda/include/thrust/detail/remove.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/replace.inl" "$(@D)/cuda/include/thrust/detail/replace.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/reverse.inl" "$(@D)/cuda/include/thrust/detail/reverse.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/scan.inl" "$(@D)/cuda/include/thrust/detail/scan.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/scatter.inl" "$(@D)/cuda/include/thrust/detail/scatter.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/seq.h" "$(@D)/cuda/include/thrust/detail/seq.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/sequence.inl" "$(@D)/cuda/include/thrust/detail/sequence.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/set_operations.inl" "$(@D)/cuda/include/thrust/detail/set_operations.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/sort.inl" "$(@D)/cuda/include/thrust/detail/sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/static_assert.h" "$(@D)/cuda/include/thrust/detail/static_assert.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/static_map.h" "$(@D)/cuda/include/thrust/detail/static_map.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/swap.h" "$(@D)/cuda/include/thrust/detail/swap.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/swap.inl" "$(@D)/cuda/include/thrust/detail/swap.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/swap_ranges.inl" "$(@D)/cuda/include/thrust/detail/swap_ranges.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/tabulate.inl" "$(@D)/cuda/include/thrust/detail/tabulate.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/temporary_array.h" "$(@D)/cuda/include/thrust/detail/temporary_array.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/temporary_array.inl" "$(@D)/cuda/include/thrust/detail/temporary_array.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/detail/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/transform.inl" "$(@D)/cuda/include/thrust/detail/transform.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/transform_reduce.inl" "$(@D)/cuda/include/thrust/detail/transform_reduce.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/transform_scan.inl" "$(@D)/cuda/include/thrust/detail/transform_scan.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/trivial_sequence.h" "$(@D)/cuda/include/thrust/detail/trivial_sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/tuple.inl" "$(@D)/cuda/include/thrust/detail/tuple.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/tuple_meta_transform.h" "$(@D)/cuda/include/thrust/detail/tuple_meta_transform.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/tuple_transform.h" "$(@D)/cuda/include/thrust/detail/tuple_transform.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits.h" "$(@D)/cuda/include/thrust/detail/type_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/algorithm/intermediate_type_from_function_and_iterators.h" "$(@D)/cuda/include/thrust/detail/type_traits/algorithm/intermediate_type_from_function_and_iterators.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/function_traits.h" "$(@D)/cuda/include/thrust/detail/type_traits/function_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/has_member_function.h" "$(@D)/cuda/include/thrust/detail/type_traits/has_member_function.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/has_nested_type.h" "$(@D)/cuda/include/thrust/detail/type_traits/has_nested_type.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/has_trivial_assign.h" "$(@D)/cuda/include/thrust/detail/type_traits/has_trivial_assign.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/is_call_possible.h" "$(@D)/cuda/include/thrust/detail/type_traits/is_call_possible.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/is_metafunction_defined.h" "$(@D)/cuda/include/thrust/detail/type_traits/is_metafunction_defined.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/iterator/is_discard_iterator.h" "$(@D)/cuda/include/thrust/detail/type_traits/iterator/is_discard_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/iterator/is_output_iterator.h" "$(@D)/cuda/include/thrust/detail/type_traits/iterator/is_output_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/minimum_type.h" "$(@D)/cuda/include/thrust/detail/type_traits/minimum_type.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/pointer_traits.h" "$(@D)/cuda/include/thrust/detail/type_traits/pointer_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/type_traits/result_of_adaptable_function.h" "$(@D)/cuda/include/thrust/detail/type_traits/result_of_adaptable_function.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/uninitialized_copy.inl" "$(@D)/cuda/include/thrust/detail/uninitialized_copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/uninitialized_fill.inl" "$(@D)/cuda/include/thrust/detail/uninitialized_fill.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/unique.inl" "$(@D)/cuda/include/thrust/detail/unique.inl" && cp "/usr/local/cuda-9.0/include/thrust/detail/use_default.h" "$(@D)/cuda/include/thrust/detail/use_default.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/util/align.h" "$(@D)/cuda/include/thrust/detail/util/align.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/util/blocking.h" "$(@D)/cuda/include/thrust/detail/util/blocking.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/vector_base.h" "$(@D)/cuda/include/thrust/detail/vector_base.h" && cp "/usr/local/cuda-9.0/include/thrust/detail/vector_base.inl" "$(@D)/cuda/include/thrust/detail/vector_base.inl" && cp "/usr/local/cuda-9.0/include/thrust/device_allocator.h" "$(@D)/cuda/include/thrust/device_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/device_delete.h" "$(@D)/cuda/include/thrust/device_delete.h" && cp "/usr/local/cuda-9.0/include/thrust/device_free.h" "$(@D)/cuda/include/thrust/device_free.h" && cp "/usr/local/cuda-9.0/include/thrust/device_malloc.h" "$(@D)/cuda/include/thrust/device_malloc.h" && cp "/usr/local/cuda-9.0/include/thrust/device_malloc_allocator.h" "$(@D)/cuda/include/thrust/device_malloc_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/device_new.h" "$(@D)/cuda/include/thrust/device_new.h" && cp "/usr/local/cuda-9.0/include/thrust/device_new_allocator.h" "$(@D)/cuda/include/thrust/device_new_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/device_ptr.h" "$(@D)/cuda/include/thrust/device_ptr.h" && cp "/usr/local/cuda-9.0/include/thrust/device_reference.h" "$(@D)/cuda/include/thrust/device_reference.h" && cp "/usr/local/cuda-9.0/include/thrust/device_vector.h" "$(@D)/cuda/include/thrust/device_vector.h" && cp "/usr/local/cuda-9.0/include/thrust/distance.h" "$(@D)/cuda/include/thrust/distance.h" && cp "/usr/local/cuda-9.0/include/thrust/equal.h" "$(@D)/cuda/include/thrust/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/execution_policy.h" "$(@D)/cuda/include/thrust/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/extrema.h" "$(@D)/cuda/include/thrust/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/fill.h" "$(@D)/cuda/include/thrust/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/find.h" "$(@D)/cuda/include/thrust/find.h" && cp "/usr/local/cuda-9.0/include/thrust/for_each.h" "$(@D)/cuda/include/thrust/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/functional.h" "$(@D)/cuda/include/thrust/functional.h" && cp "/usr/local/cuda-9.0/include/thrust/gather.h" "$(@D)/cuda/include/thrust/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/generate.h" "$(@D)/cuda/include/thrust/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/host_vector.h" "$(@D)/cuda/include/thrust/host_vector.h" && cp "/usr/local/cuda-9.0/include/thrust/inner_product.h" "$(@D)/cuda/include/thrust/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/constant_iterator.h" "$(@D)/cuda/include/thrust/iterator/constant_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/counting_iterator.h" "$(@D)/cuda/include/thrust/iterator/counting_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/any_assign.h" "$(@D)/cuda/include/thrust/iterator/detail/any_assign.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/any_system_tag.h" "$(@D)/cuda/include/thrust/iterator/detail/any_system_tag.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/constant_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/constant_iterator_base.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/counting_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/counting_iterator.inl" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/device_system_tag.h" "$(@D)/cuda/include/thrust/iterator/detail/device_system_tag.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/discard_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/discard_iterator_base.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/distance_from_result.h" "$(@D)/cuda/include/thrust/iterator/detail/distance_from_result.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/host_system_tag.h" "$(@D)/cuda/include/thrust/iterator/detail/host_system_tag.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/is_iterator_category.h" "$(@D)/cuda/include/thrust/iterator/detail/is_iterator_category.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/is_trivial_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/is_trivial_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_adaptor_base.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_adaptor_base.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_category_to_system.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_category_to_system.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_category_to_traversal.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_category_to_traversal.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_category_with_system_and_traversal.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_category_with_system_and_traversal.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_facade_category.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_facade_category.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_traits.inl" "$(@D)/cuda/include/thrust/iterator/detail/iterator_traits.inl" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/iterator_traversal_tags.h" "$(@D)/cuda/include/thrust/iterator/detail/iterator_traversal_tags.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/join_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/join_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/minimum_category.h" "$(@D)/cuda/include/thrust/iterator/detail/minimum_category.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/minimum_system.h" "$(@D)/cuda/include/thrust/iterator/detail/minimum_system.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/normal_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/normal_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/permutation_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/permutation_iterator_base.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/retag.h" "$(@D)/cuda/include/thrust/iterator/detail/retag.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/reverse_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/reverse_iterator.inl" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/reverse_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/reverse_iterator_base.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/tagged_iterator.h" "$(@D)/cuda/include/thrust/iterator/detail/tagged_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/transform_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/transform_iterator.inl" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/transform_output_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/transform_output_iterator.inl" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/tuple_of_iterator_references.h" "$(@D)/cuda/include/thrust/iterator/detail/tuple_of_iterator_references.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/universal_categories.h" "$(@D)/cuda/include/thrust/iterator/detail/universal_categories.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/zip_iterator.inl" "$(@D)/cuda/include/thrust/iterator/detail/zip_iterator.inl" && cp "/usr/local/cuda-9.0/include/thrust/iterator/detail/zip_iterator_base.h" "$(@D)/cuda/include/thrust/iterator/detail/zip_iterator_base.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/discard_iterator.h" "$(@D)/cuda/include/thrust/iterator/discard_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/iterator_adaptor.h" "$(@D)/cuda/include/thrust/iterator/iterator_adaptor.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/iterator_categories.h" "$(@D)/cuda/include/thrust/iterator/iterator_categories.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/iterator_facade.h" "$(@D)/cuda/include/thrust/iterator/iterator_facade.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/iterator_traits.h" "$(@D)/cuda/include/thrust/iterator/iterator_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/permutation_iterator.h" "$(@D)/cuda/include/thrust/iterator/permutation_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/retag.h" "$(@D)/cuda/include/thrust/iterator/retag.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/reverse_iterator.h" "$(@D)/cuda/include/thrust/iterator/reverse_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/transform_iterator.h" "$(@D)/cuda/include/thrust/iterator/transform_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/transform_output_iterator.h" "$(@D)/cuda/include/thrust/iterator/transform_output_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/iterator/zip_iterator.h" "$(@D)/cuda/include/thrust/iterator/zip_iterator.h" && cp "/usr/local/cuda-9.0/include/thrust/logical.h" "$(@D)/cuda/include/thrust/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/memory.h" "$(@D)/cuda/include/thrust/memory.h" && cp "/usr/local/cuda-9.0/include/thrust/merge.h" "$(@D)/cuda/include/thrust/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/mismatch.h" "$(@D)/cuda/include/thrust/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/pair.h" "$(@D)/cuda/include/thrust/pair.h" && cp "/usr/local/cuda-9.0/include/thrust/partition.h" "$(@D)/cuda/include/thrust/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/random.h" "$(@D)/cuda/include/thrust/random.h" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/discard_block_engine.inl" "$(@D)/cuda/include/thrust/random/detail/discard_block_engine.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/linear_congruential_engine.inl" "$(@D)/cuda/include/thrust/random/detail/linear_congruential_engine.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/linear_congruential_engine_discard.h" "$(@D)/cuda/include/thrust/random/detail/linear_congruential_engine_discard.h" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/linear_feedback_shift_engine.inl" "$(@D)/cuda/include/thrust/random/detail/linear_feedback_shift_engine.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/linear_feedback_shift_engine_wordmask.h" "$(@D)/cuda/include/thrust/random/detail/linear_feedback_shift_engine_wordmask.h" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/mod.h" "$(@D)/cuda/include/thrust/random/detail/mod.h" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/normal_distribution.inl" "$(@D)/cuda/include/thrust/random/detail/normal_distribution.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/normal_distribution_base.h" "$(@D)/cuda/include/thrust/random/detail/normal_distribution_base.h" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/random_core_access.h" "$(@D)/cuda/include/thrust/random/detail/random_core_access.h" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/subtract_with_carry_engine.inl" "$(@D)/cuda/include/thrust/random/detail/subtract_with_carry_engine.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/uniform_int_distribution.inl" "$(@D)/cuda/include/thrust/random/detail/uniform_int_distribution.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/uniform_real_distribution.inl" "$(@D)/cuda/include/thrust/random/detail/uniform_real_distribution.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/xor_combine_engine.inl" "$(@D)/cuda/include/thrust/random/detail/xor_combine_engine.inl" && cp "/usr/local/cuda-9.0/include/thrust/random/detail/xor_combine_engine_max.h" "$(@D)/cuda/include/thrust/random/detail/xor_combine_engine_max.h" && cp "/usr/local/cuda-9.0/include/thrust/random/discard_block_engine.h" "$(@D)/cuda/include/thrust/random/discard_block_engine.h" && cp "/usr/local/cuda-9.0/include/thrust/random/linear_congruential_engine.h" "$(@D)/cuda/include/thrust/random/linear_congruential_engine.h" && cp "/usr/local/cuda-9.0/include/thrust/random/linear_feedback_shift_engine.h" "$(@D)/cuda/include/thrust/random/linear_feedback_shift_engine.h" && cp "/usr/local/cuda-9.0/include/thrust/random/normal_distribution.h" "$(@D)/cuda/include/thrust/random/normal_distribution.h" && cp "/usr/local/cuda-9.0/include/thrust/random/subtract_with_carry_engine.h" "$(@D)/cuda/include/thrust/random/subtract_with_carry_engine.h" && cp "/usr/local/cuda-9.0/include/thrust/random/uniform_int_distribution.h" "$(@D)/cuda/include/thrust/random/uniform_int_distribution.h" && cp "/usr/local/cuda-9.0/include/thrust/random/uniform_real_distribution.h" "$(@D)/cuda/include/thrust/random/uniform_real_distribution.h" && cp "/usr/local/cuda-9.0/include/thrust/random/xor_combine_engine.h" "$(@D)/cuda/include/thrust/random/xor_combine_engine.h" && cp "/usr/local/cuda-9.0/include/thrust/reduce.h" "$(@D)/cuda/include/thrust/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/remove.h" "$(@D)/cuda/include/thrust/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/replace.h" "$(@D)/cuda/include/thrust/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/reverse.h" "$(@D)/cuda/include/thrust/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/scan.h" "$(@D)/cuda/include/thrust/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/scatter.h" "$(@D)/cuda/include/thrust/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/sequence.h" "$(@D)/cuda/include/thrust/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/set_operations.h" "$(@D)/cuda/include/thrust/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/sort.h" "$(@D)/cuda/include/thrust/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/swap.h" "$(@D)/cuda/include/thrust/swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/cpp/detail/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/cpp/detail/assign_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/cpp/detail/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/copy.h" "$(@D)/cuda/include/thrust/system/cpp/detail/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/cpp/detail/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/count.h" "$(@D)/cuda/include/thrust/system/cpp/detail/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/equal.h" "$(@D)/cuda/include/thrust/system/cpp/detail/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/cpp/detail/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/extrema.h" "$(@D)/cuda/include/thrust/system/cpp/detail/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/fill.h" "$(@D)/cuda/include/thrust/system/cpp/detail/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/find.h" "$(@D)/cuda/include/thrust/system/cpp/detail/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/for_each.h" "$(@D)/cuda/include/thrust/system/cpp/detail/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/gather.h" "$(@D)/cuda/include/thrust/system/cpp/detail/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/generate.h" "$(@D)/cuda/include/thrust/system/cpp/detail/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/get_value.h" "$(@D)/cuda/include/thrust/system/cpp/detail/get_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/cpp/detail/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/cpp/detail/iter_swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/logical.h" "$(@D)/cuda/include/thrust/system/cpp/detail/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/cpp/detail/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/memory.inl" "$(@D)/cuda/include/thrust/system/cpp/detail/memory.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/merge.h" "$(@D)/cuda/include/thrust/system/cpp/detail/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/cpp/detail/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/par.h" "$(@D)/cuda/include/thrust/system/cpp/detail/par.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/partition.h" "$(@D)/cuda/include/thrust/system/cpp/detail/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/reduce.h" "$(@D)/cuda/include/thrust/system/cpp/detail/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/cpp/detail/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/remove.h" "$(@D)/cuda/include/thrust/system/cpp/detail/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/replace.h" "$(@D)/cuda/include/thrust/system/cpp/detail/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/reverse.h" "$(@D)/cuda/include/thrust/system/cpp/detail/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/scan.h" "$(@D)/cuda/include/thrust/system/cpp/detail/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/cpp/detail/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/scatter.h" "$(@D)/cuda/include/thrust/system/cpp/detail/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/sequence.h" "$(@D)/cuda/include/thrust/system/cpp/detail/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/cpp/detail/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/sort.h" "$(@D)/cuda/include/thrust/system/cpp/detail/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/cpp/detail/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/cpp/detail/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/cpp/detail/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/transform.h" "$(@D)/cuda/include/thrust/system/cpp/detail/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/cpp/detail/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/cpp/detail/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/cpp/detail/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/cpp/detail/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/unique.h" "$(@D)/cuda/include/thrust/system/cpp/detail/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/cpp/detail/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/detail/vector.inl" "$(@D)/cuda/include/thrust/system/cpp/detail/vector.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/execution_policy.h" "$(@D)/cuda/include/thrust/system/cpp/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/memory.h" "$(@D)/cuda/include/thrust/system/cpp/memory.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cpp/vector.h" "$(@D)/cuda/include/thrust/system/cpp/vector.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/config.h" "$(@D)/cuda/include/thrust/system/cuda/config.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/cuda/detail/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/cuda/detail/assign_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/cuda/detail/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/copy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/cuda/detail/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/core/agent_launcher.h" "$(@D)/cuda/include/thrust/system/cuda/detail/core/agent_launcher.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/core/alignment.h" "$(@D)/cuda/include/thrust/system/cuda/detail/core/alignment.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/core/triple_chevron_launch.h" "$(@D)/cuda/include/thrust/system/cuda/detail/core/triple_chevron_launch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/core/util.h" "$(@D)/cuda/include/thrust/system/cuda/detail/core/util.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/count.h" "$(@D)/cuda/include/thrust/system/cuda/detail/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cross_system.h" "$(@D)/cuda/include/thrust/system/cuda/detail/cross_system.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_histogram.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_histogram.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_radix_sort_downsweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_radix_sort_downsweep.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_radix_sort_upsweep.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_radix_sort_upsweep.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_reduce_by_key.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_reduce_by_key.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_rle.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_rle.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_scan.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_segment_fixup.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_segment_fixup.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_select_if.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_select_if.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_spmv_csrt.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_spmv_csrt.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_spmv_orig.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_spmv_orig.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/agent_spmv_row_based.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/agent_spmv_row_based.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/agent/single_pass_scan_operators.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/agent/single_pass_scan_operators.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_adjacent_difference.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_adjacent_difference.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_discontinuity.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_discontinuity.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_exchange.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_exchange.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_histogram.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_histogram.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_load.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_load.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_radix_rank.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_radix_rank.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_radix_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_radix_sort.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_raking_layout.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_raking_layout.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_scan.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_shuffle.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_shuffle.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/block_store.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/block_store.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_atomic.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_atomic.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_histogram_sort.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking_commutative_only.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_raking_commutative_only.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_warp_reductions.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_reduce_warp_reductions.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_raking.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_raking.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans2.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans2.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans3.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/block/specializations/block_scan_warp_scans3.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/cub.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/cub.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_histogram.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_histogram.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_partition.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_partition.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_radix_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_radix_sort.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_run_length_encode.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_run_length_encode.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_scan.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_segmented_radix_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_segmented_radix_sort.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_segmented_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_segmented_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_select.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_select.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/device_spmv.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/device_spmv.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_histogram.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_histogram.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_radix_sort.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_radix_sort.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_reduce_by_key.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_reduce_by_key.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_rle.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_rle.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_scan.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_select_if.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_select_if.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_csrt.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_csrt.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_orig.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_orig.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_row_based.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/device/dispatch/dispatch_spmv_row_based.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/grid/grid_barrier.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_barrier.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/grid/grid_even_share.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_even_share.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/grid/grid_mapping.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_mapping.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/grid/grid_queue.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/grid/grid_queue.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/host/mutex.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/host/mutex.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/arg_index_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/arg_index_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/cache_modified_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/cache_modified_output_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/cache_modified_output_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/constant_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/constant_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/discard_output_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/discard_output_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/tex_obj_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/tex_obj_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/tex_ref_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/tex_ref_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/iterator/transform_input_iterator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/iterator/transform_input_iterator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/thread/thread_load.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_load.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/thread/thread_operators.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_operators.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/thread/thread_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/thread/thread_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_scan.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/thread/thread_search.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_search.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/thread/thread_store.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/thread/thread_store.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_allocator.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_allocator.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_arch.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_arch.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_debug.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_debug.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_device.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_device.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_macro.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_macro.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_namespace.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_namespace.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_ptx.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_ptx.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/util_type.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/util_type.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_shfl.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_shfl.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_smem.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_reduce_smem.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_shfl.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_shfl.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_smem.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/specializations/warp_scan_smem.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/warp/warp_reduce.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/warp_reduce.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/cub/warp/warp_scan.cuh" "$(@D)/cuda/include/thrust/system/cuda/detail/cub/warp/warp_scan.cuh" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/equal.h" "$(@D)/cuda/include/thrust/system/cuda/detail/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/error.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/error.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/extrema.h" "$(@D)/cuda/include/thrust/system/cuda/detail/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/fill.h" "$(@D)/cuda/include/thrust/system/cuda/detail/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/find.h" "$(@D)/cuda/include/thrust/system/cuda/detail/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/for_each.h" "$(@D)/cuda/include/thrust/system/cuda/detail/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/gather.h" "$(@D)/cuda/include/thrust/system/cuda/detail/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/generate.h" "$(@D)/cuda/include/thrust/system/cuda/detail/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/get_value.h" "$(@D)/cuda/include/thrust/system/cuda/detail/get_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/guarded_cuda_runtime_api.h" "$(@D)/cuda/include/thrust/system/cuda/detail/guarded_cuda_runtime_api.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/guarded_driver_types.h" "$(@D)/cuda/include/thrust/system/cuda/detail/guarded_driver_types.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/cuda/detail/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/internal/copy_cross_system.h" "$(@D)/cuda/include/thrust/system/cuda/detail/internal/copy_cross_system.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/internal/copy_device_to_device.h" "$(@D)/cuda/include/thrust/system/cuda/detail/internal/copy_device_to_device.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/cuda/detail/iter_swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/logical.h" "$(@D)/cuda/include/thrust/system/cuda/detail/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/cuda/detail/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/memory.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/memory.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/memory_buffer.h" "$(@D)/cuda/include/thrust/system/cuda/detail/memory_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/merge.h" "$(@D)/cuda/include/thrust/system/cuda/detail/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/cuda/detail/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/par.h" "$(@D)/cuda/include/thrust/system/cuda/detail/par.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/par_to_seq.h" "$(@D)/cuda/include/thrust/system/cuda/detail/par_to_seq.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/parallel_for.h" "$(@D)/cuda/include/thrust/system/cuda/detail/parallel_for.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/partition.h" "$(@D)/cuda/include/thrust/system/cuda/detail/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/reduce.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/remove.h" "$(@D)/cuda/include/thrust/system/cuda/detail/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/replace.h" "$(@D)/cuda/include/thrust/system/cuda/detail/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/reverse.h" "$(@D)/cuda/include/thrust/system/cuda/detail/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/scan.h" "$(@D)/cuda/include/thrust/system/cuda/detail/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/cuda/detail/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/scatter.h" "$(@D)/cuda/include/thrust/system/cuda/detail/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/sequence.h" "$(@D)/cuda/include/thrust/system/cuda/detail/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/cuda/detail/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/sort.h" "$(@D)/cuda/include/thrust/system/cuda/detail/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/cuda/detail/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/cuda/detail/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/cuda/detail/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/terminate.h" "$(@D)/cuda/include/thrust/system/cuda/detail/terminate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/transform.h" "$(@D)/cuda/include/thrust/system/cuda/detail/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/cuda/detail/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/cuda/detail/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/cuda/detail/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/cuda/detail/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/unique.h" "$(@D)/cuda/include/thrust/system/cuda/detail/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/cuda/detail/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/util.h" "$(@D)/cuda/include/thrust/system/cuda/detail/util.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/detail/vector.inl" "$(@D)/cuda/include/thrust/system/cuda/detail/vector.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/error.h" "$(@D)/cuda/include/thrust/system/cuda/error.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/execution_policy.h" "$(@D)/cuda/include/thrust/system/cuda/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/experimental/pinned_allocator.h" "$(@D)/cuda/include/thrust/system/cuda/experimental/pinned_allocator.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/memory.h" "$(@D)/cuda/include/thrust/system/cuda/memory.h" && cp "/usr/local/cuda-9.0/include/thrust/system/cuda/vector.h" "$(@D)/cuda/include/thrust/system/cuda/vector.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/detail/adl/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/assign_value.h" "$(@D)/cuda/include/thrust/system/detail/adl/assign_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/adl/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/copy.h" "$(@D)/cuda/include/thrust/system/detail/adl/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/copy_if.h" "$(@D)/cuda/include/thrust/system/detail/adl/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/count.h" "$(@D)/cuda/include/thrust/system/detail/adl/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/equal.h" "$(@D)/cuda/include/thrust/system/detail/adl/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/extrema.h" "$(@D)/cuda/include/thrust/system/detail/adl/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/fill.h" "$(@D)/cuda/include/thrust/system/detail/adl/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/find.h" "$(@D)/cuda/include/thrust/system/detail/adl/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/for_each.h" "$(@D)/cuda/include/thrust/system/detail/adl/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/gather.h" "$(@D)/cuda/include/thrust/system/detail/adl/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/generate.h" "$(@D)/cuda/include/thrust/system/detail/adl/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/get_value.h" "$(@D)/cuda/include/thrust/system/detail/adl/get_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/inner_product.h" "$(@D)/cuda/include/thrust/system/detail/adl/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/iter_swap.h" "$(@D)/cuda/include/thrust/system/detail/adl/iter_swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/logical.h" "$(@D)/cuda/include/thrust/system/detail/adl/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/detail/adl/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/merge.h" "$(@D)/cuda/include/thrust/system/detail/adl/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/mismatch.h" "$(@D)/cuda/include/thrust/system/detail/adl/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/partition.h" "$(@D)/cuda/include/thrust/system/detail/adl/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/reduce.h" "$(@D)/cuda/include/thrust/system/detail/adl/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/detail/adl/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/remove.h" "$(@D)/cuda/include/thrust/system/detail/adl/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/replace.h" "$(@D)/cuda/include/thrust/system/detail/adl/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/reverse.h" "$(@D)/cuda/include/thrust/system/detail/adl/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/scan.h" "$(@D)/cuda/include/thrust/system/detail/adl/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/scan_by_key.h" "$(@D)/cuda/include/thrust/system/detail/adl/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/scatter.h" "$(@D)/cuda/include/thrust/system/detail/adl/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/sequence.h" "$(@D)/cuda/include/thrust/system/detail/adl/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/set_operations.h" "$(@D)/cuda/include/thrust/system/detail/adl/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/sort.h" "$(@D)/cuda/include/thrust/system/detail/adl/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/swap_ranges.h" "$(@D)/cuda/include/thrust/system/detail/adl/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/tabulate.h" "$(@D)/cuda/include/thrust/system/detail/adl/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/detail/adl/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/transform.h" "$(@D)/cuda/include/thrust/system/detail/adl/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/transform_reduce.h" "$(@D)/cuda/include/thrust/system/detail/adl/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/transform_scan.h" "$(@D)/cuda/include/thrust/system/detail/adl/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/detail/adl/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/detail/adl/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/unique.h" "$(@D)/cuda/include/thrust/system/detail/adl/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/adl/unique_by_key.h" "$(@D)/cuda/include/thrust/system/detail/adl/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/bad_alloc.h" "$(@D)/cuda/include/thrust/system/detail/bad_alloc.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/errno.h" "$(@D)/cuda/include/thrust/system/detail/errno.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/error_category.inl" "$(@D)/cuda/include/thrust/system/detail/error_category.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/error_code.inl" "$(@D)/cuda/include/thrust/system/detail/error_code.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/error_condition.inl" "$(@D)/cuda/include/thrust/system/detail/error_condition.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/detail/generic/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/adjacent_difference.inl" "$(@D)/cuda/include/thrust/system/detail/generic/adjacent_difference.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/advance.h" "$(@D)/cuda/include/thrust/system/detail/generic/advance.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/advance.inl" "$(@D)/cuda/include/thrust/system/detail/generic/advance.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/generic/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/binary_search.inl" "$(@D)/cuda/include/thrust/system/detail/generic/binary_search.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/copy.h" "$(@D)/cuda/include/thrust/system/detail/generic/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/copy.inl" "$(@D)/cuda/include/thrust/system/detail/generic/copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/copy_if.h" "$(@D)/cuda/include/thrust/system/detail/generic/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/copy_if.inl" "$(@D)/cuda/include/thrust/system/detail/generic/copy_if.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/count.h" "$(@D)/cuda/include/thrust/system/detail/generic/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/count.inl" "$(@D)/cuda/include/thrust/system/detail/generic/count.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/distance.h" "$(@D)/cuda/include/thrust/system/detail/generic/distance.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/distance.inl" "$(@D)/cuda/include/thrust/system/detail/generic/distance.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/equal.h" "$(@D)/cuda/include/thrust/system/detail/generic/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/equal.inl" "$(@D)/cuda/include/thrust/system/detail/generic/equal.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/extrema.h" "$(@D)/cuda/include/thrust/system/detail/generic/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/extrema.inl" "$(@D)/cuda/include/thrust/system/detail/generic/extrema.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/fill.h" "$(@D)/cuda/include/thrust/system/detail/generic/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/find.h" "$(@D)/cuda/include/thrust/system/detail/generic/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/find.inl" "$(@D)/cuda/include/thrust/system/detail/generic/find.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/for_each.h" "$(@D)/cuda/include/thrust/system/detail/generic/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/gather.h" "$(@D)/cuda/include/thrust/system/detail/generic/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/gather.inl" "$(@D)/cuda/include/thrust/system/detail/generic/gather.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/generate.h" "$(@D)/cuda/include/thrust/system/detail/generic/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/generate.inl" "$(@D)/cuda/include/thrust/system/detail/generic/generate.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/inner_product.h" "$(@D)/cuda/include/thrust/system/detail/generic/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/inner_product.inl" "$(@D)/cuda/include/thrust/system/detail/generic/inner_product.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/logical.h" "$(@D)/cuda/include/thrust/system/detail/generic/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/memory.h" "$(@D)/cuda/include/thrust/system/detail/generic/memory.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/memory.inl" "$(@D)/cuda/include/thrust/system/detail/generic/memory.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/merge.h" "$(@D)/cuda/include/thrust/system/detail/generic/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/merge.inl" "$(@D)/cuda/include/thrust/system/detail/generic/merge.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/mismatch.h" "$(@D)/cuda/include/thrust/system/detail/generic/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/mismatch.inl" "$(@D)/cuda/include/thrust/system/detail/generic/mismatch.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/partition.h" "$(@D)/cuda/include/thrust/system/detail/generic/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/partition.inl" "$(@D)/cuda/include/thrust/system/detail/generic/partition.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/reduce.h" "$(@D)/cuda/include/thrust/system/detail/generic/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/reduce.inl" "$(@D)/cuda/include/thrust/system/detail/generic/reduce.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/detail/generic/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/detail/generic/reduce_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/remove.h" "$(@D)/cuda/include/thrust/system/detail/generic/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/remove.inl" "$(@D)/cuda/include/thrust/system/detail/generic/remove.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/replace.h" "$(@D)/cuda/include/thrust/system/detail/generic/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/replace.inl" "$(@D)/cuda/include/thrust/system/detail/generic/replace.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/reverse.h" "$(@D)/cuda/include/thrust/system/detail/generic/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/reverse.inl" "$(@D)/cuda/include/thrust/system/detail/generic/reverse.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scalar/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/generic/scalar/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scalar/binary_search.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scalar/binary_search.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scan.h" "$(@D)/cuda/include/thrust/system/detail/generic/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scan.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scan.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scan_by_key.h" "$(@D)/cuda/include/thrust/system/detail/generic/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scan_by_key.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scan_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scatter.h" "$(@D)/cuda/include/thrust/system/detail/generic/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/scatter.inl" "$(@D)/cuda/include/thrust/system/detail/generic/scatter.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/select_system.h" "$(@D)/cuda/include/thrust/system/detail/generic/select_system.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/sequence.h" "$(@D)/cuda/include/thrust/system/detail/generic/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/sequence.inl" "$(@D)/cuda/include/thrust/system/detail/generic/sequence.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/set_operations.h" "$(@D)/cuda/include/thrust/system/detail/generic/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/set_operations.inl" "$(@D)/cuda/include/thrust/system/detail/generic/set_operations.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/sort.h" "$(@D)/cuda/include/thrust/system/detail/generic/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/sort.inl" "$(@D)/cuda/include/thrust/system/detail/generic/sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/swap_ranges.h" "$(@D)/cuda/include/thrust/system/detail/generic/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/swap_ranges.inl" "$(@D)/cuda/include/thrust/system/detail/generic/swap_ranges.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/tabulate.h" "$(@D)/cuda/include/thrust/system/detail/generic/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/tabulate.inl" "$(@D)/cuda/include/thrust/system/detail/generic/tabulate.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/tag.h" "$(@D)/cuda/include/thrust/system/detail/generic/tag.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/detail/generic/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/temporary_buffer.inl" "$(@D)/cuda/include/thrust/system/detail/generic/temporary_buffer.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/transform.h" "$(@D)/cuda/include/thrust/system/detail/generic/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/transform.inl" "$(@D)/cuda/include/thrust/system/detail/generic/transform.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/transform_reduce.h" "$(@D)/cuda/include/thrust/system/detail/generic/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/transform_reduce.inl" "$(@D)/cuda/include/thrust/system/detail/generic/transform_reduce.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/transform_scan.h" "$(@D)/cuda/include/thrust/system/detail/generic/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/transform_scan.inl" "$(@D)/cuda/include/thrust/system/detail/generic/transform_scan.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/type_traits.h" "$(@D)/cuda/include/thrust/system/detail/generic/type_traits.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/uninitialized_copy.inl" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/uninitialized_fill.inl" "$(@D)/cuda/include/thrust/system/detail/generic/uninitialized_fill.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/unique.h" "$(@D)/cuda/include/thrust/system/detail/generic/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/unique.inl" "$(@D)/cuda/include/thrust/system/detail/generic/unique.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/unique_by_key.h" "$(@D)/cuda/include/thrust/system/detail/generic/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/generic/unique_by_key.inl" "$(@D)/cuda/include/thrust/system/detail/generic/unique_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/internal/decompose.h" "$(@D)/cuda/include/thrust/system/detail/internal/decompose.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/detail/sequential/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/assign_value.h" "$(@D)/cuda/include/thrust/system/detail/sequential/assign_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/binary_search.h" "$(@D)/cuda/include/thrust/system/detail/sequential/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/copy.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/copy_backward.h" "$(@D)/cuda/include/thrust/system/detail/sequential/copy_backward.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/copy_if.h" "$(@D)/cuda/include/thrust/system/detail/sequential/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/count.h" "$(@D)/cuda/include/thrust/system/detail/sequential/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/equal.h" "$(@D)/cuda/include/thrust/system/detail/sequential/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/execution_policy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/extrema.h" "$(@D)/cuda/include/thrust/system/detail/sequential/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/fill.h" "$(@D)/cuda/include/thrust/system/detail/sequential/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/find.h" "$(@D)/cuda/include/thrust/system/detail/sequential/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/for_each.h" "$(@D)/cuda/include/thrust/system/detail/sequential/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/gather.h" "$(@D)/cuda/include/thrust/system/detail/sequential/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/general_copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/general_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/generate.h" "$(@D)/cuda/include/thrust/system/detail/sequential/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/get_value.h" "$(@D)/cuda/include/thrust/system/detail/sequential/get_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/inner_product.h" "$(@D)/cuda/include/thrust/system/detail/sequential/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/insertion_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/insertion_sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/iter_swap.h" "$(@D)/cuda/include/thrust/system/detail/sequential/iter_swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/logical.h" "$(@D)/cuda/include/thrust/system/detail/sequential/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/detail/sequential/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/merge.h" "$(@D)/cuda/include/thrust/system/detail/sequential/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/merge.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/merge.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/mismatch.h" "$(@D)/cuda/include/thrust/system/detail/sequential/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/partition.h" "$(@D)/cuda/include/thrust/system/detail/sequential/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/reduce.h" "$(@D)/cuda/include/thrust/system/detail/sequential/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/detail/sequential/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/remove.h" "$(@D)/cuda/include/thrust/system/detail/sequential/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/replace.h" "$(@D)/cuda/include/thrust/system/detail/sequential/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/reverse.h" "$(@D)/cuda/include/thrust/system/detail/sequential/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/scan.h" "$(@D)/cuda/include/thrust/system/detail/sequential/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/scan_by_key.h" "$(@D)/cuda/include/thrust/system/detail/sequential/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/scatter.h" "$(@D)/cuda/include/thrust/system/detail/sequential/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/sequence.h" "$(@D)/cuda/include/thrust/system/detail/sequential/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/set_operations.h" "$(@D)/cuda/include/thrust/system/detail/sequential/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/stable_merge_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_merge_sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/stable_merge_sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_merge_sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/stable_primitive_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_primitive_sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/stable_primitive_sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_primitive_sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/stable_radix_sort.h" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_radix_sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/stable_radix_sort.inl" "$(@D)/cuda/include/thrust/system/detail/sequential/stable_radix_sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/swap_ranges.h" "$(@D)/cuda/include/thrust/system/detail/sequential/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/tabulate.h" "$(@D)/cuda/include/thrust/system/detail/sequential/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/detail/sequential/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/transform.h" "$(@D)/cuda/include/thrust/system/detail/sequential/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/transform_reduce.h" "$(@D)/cuda/include/thrust/system/detail/sequential/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/transform_scan.h" "$(@D)/cuda/include/thrust/system/detail/sequential/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/trivial_copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/trivial_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/detail/sequential/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/detail/sequential/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/unique.h" "$(@D)/cuda/include/thrust/system/detail/sequential/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/sequential/unique_by_key.h" "$(@D)/cuda/include/thrust/system/detail/sequential/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/detail/system_error.inl" "$(@D)/cuda/include/thrust/system/detail/system_error.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/error_code.h" "$(@D)/cuda/include/thrust/system/error_code.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/omp/detail/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/omp/detail/assign_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/omp/detail/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/copy.h" "$(@D)/cuda/include/thrust/system/omp/detail/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/copy.inl" "$(@D)/cuda/include/thrust/system/omp/detail/copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/omp/detail/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/copy_if.inl" "$(@D)/cuda/include/thrust/system/omp/detail/copy_if.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/count.h" "$(@D)/cuda/include/thrust/system/omp/detail/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/default_decomposition.h" "$(@D)/cuda/include/thrust/system/omp/detail/default_decomposition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/default_decomposition.inl" "$(@D)/cuda/include/thrust/system/omp/detail/default_decomposition.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/equal.h" "$(@D)/cuda/include/thrust/system/omp/detail/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/omp/detail/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/extrema.h" "$(@D)/cuda/include/thrust/system/omp/detail/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/fill.h" "$(@D)/cuda/include/thrust/system/omp/detail/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/find.h" "$(@D)/cuda/include/thrust/system/omp/detail/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/for_each.h" "$(@D)/cuda/include/thrust/system/omp/detail/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/for_each.inl" "$(@D)/cuda/include/thrust/system/omp/detail/for_each.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/gather.h" "$(@D)/cuda/include/thrust/system/omp/detail/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/generate.h" "$(@D)/cuda/include/thrust/system/omp/detail/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/get_value.h" "$(@D)/cuda/include/thrust/system/omp/detail/get_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/omp/detail/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/omp/detail/iter_swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/logical.h" "$(@D)/cuda/include/thrust/system/omp/detail/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/omp/detail/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/memory.inl" "$(@D)/cuda/include/thrust/system/omp/detail/memory.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/merge.h" "$(@D)/cuda/include/thrust/system/omp/detail/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/omp/detail/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/par.h" "$(@D)/cuda/include/thrust/system/omp/detail/par.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/partition.h" "$(@D)/cuda/include/thrust/system/omp/detail/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/partition.inl" "$(@D)/cuda/include/thrust/system/omp/detail/partition.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reduce.h" "$(@D)/cuda/include/thrust/system/omp/detail/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reduce.inl" "$(@D)/cuda/include/thrust/system/omp/detail/reduce.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reduce_intervals.h" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_intervals.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reduce_intervals.inl" "$(@D)/cuda/include/thrust/system/omp/detail/reduce_intervals.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/remove.h" "$(@D)/cuda/include/thrust/system/omp/detail/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/remove.inl" "$(@D)/cuda/include/thrust/system/omp/detail/remove.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/replace.h" "$(@D)/cuda/include/thrust/system/omp/detail/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/reverse.h" "$(@D)/cuda/include/thrust/system/omp/detail/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/scan.h" "$(@D)/cuda/include/thrust/system/omp/detail/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/omp/detail/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/scatter.h" "$(@D)/cuda/include/thrust/system/omp/detail/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/sequence.h" "$(@D)/cuda/include/thrust/system/omp/detail/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/omp/detail/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/sort.h" "$(@D)/cuda/include/thrust/system/omp/detail/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/sort.inl" "$(@D)/cuda/include/thrust/system/omp/detail/sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/omp/detail/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/omp/detail/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/omp/detail/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/transform.h" "$(@D)/cuda/include/thrust/system/omp/detail/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/omp/detail/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/omp/detail/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/omp/detail/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/omp/detail/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/unique.h" "$(@D)/cuda/include/thrust/system/omp/detail/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/unique.inl" "$(@D)/cuda/include/thrust/system/omp/detail/unique.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/omp/detail/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/unique_by_key.inl" "$(@D)/cuda/include/thrust/system/omp/detail/unique_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/detail/vector.inl" "$(@D)/cuda/include/thrust/system/omp/detail/vector.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/execution_policy.h" "$(@D)/cuda/include/thrust/system/omp/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/memory.h" "$(@D)/cuda/include/thrust/system/omp/memory.h" && cp "/usr/local/cuda-9.0/include/thrust/system/omp/vector.h" "$(@D)/cuda/include/thrust/system/omp/vector.h" && cp "/usr/local/cuda-9.0/include/thrust/system/system_error.h" "$(@D)/cuda/include/thrust/system/system_error.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/adjacent_difference.h" "$(@D)/cuda/include/thrust/system/tbb/detail/adjacent_difference.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/assign_value.h" "$(@D)/cuda/include/thrust/system/tbb/detail/assign_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/binary_search.h" "$(@D)/cuda/include/thrust/system/tbb/detail/binary_search.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/copy.h" "$(@D)/cuda/include/thrust/system/tbb/detail/copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/copy.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/copy.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/copy_if.h" "$(@D)/cuda/include/thrust/system/tbb/detail/copy_if.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/copy_if.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/copy_if.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/count.h" "$(@D)/cuda/include/thrust/system/tbb/detail/count.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/equal.h" "$(@D)/cuda/include/thrust/system/tbb/detail/equal.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/execution_policy.h" "$(@D)/cuda/include/thrust/system/tbb/detail/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/extrema.h" "$(@D)/cuda/include/thrust/system/tbb/detail/extrema.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/fill.h" "$(@D)/cuda/include/thrust/system/tbb/detail/fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/find.h" "$(@D)/cuda/include/thrust/system/tbb/detail/find.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/for_each.h" "$(@D)/cuda/include/thrust/system/tbb/detail/for_each.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/for_each.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/for_each.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/gather.h" "$(@D)/cuda/include/thrust/system/tbb/detail/gather.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/generate.h" "$(@D)/cuda/include/thrust/system/tbb/detail/generate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/get_value.h" "$(@D)/cuda/include/thrust/system/tbb/detail/get_value.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/inner_product.h" "$(@D)/cuda/include/thrust/system/tbb/detail/inner_product.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/iter_swap.h" "$(@D)/cuda/include/thrust/system/tbb/detail/iter_swap.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/logical.h" "$(@D)/cuda/include/thrust/system/tbb/detail/logical.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/malloc_and_free.h" "$(@D)/cuda/include/thrust/system/tbb/detail/malloc_and_free.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/memory.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/memory.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/merge.h" "$(@D)/cuda/include/thrust/system/tbb/detail/merge.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/merge.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/merge.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/mismatch.h" "$(@D)/cuda/include/thrust/system/tbb/detail/mismatch.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/par.h" "$(@D)/cuda/include/thrust/system/tbb/detail/par.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/partition.h" "$(@D)/cuda/include/thrust/system/tbb/detail/partition.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/partition.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/partition.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/reduce.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/reduce.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/reduce_by_key.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/reduce_by_key.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/reduce_intervals.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reduce_intervals.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/remove.h" "$(@D)/cuda/include/thrust/system/tbb/detail/remove.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/remove.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/remove.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/replace.h" "$(@D)/cuda/include/thrust/system/tbb/detail/replace.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/reverse.h" "$(@D)/cuda/include/thrust/system/tbb/detail/reverse.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/scan.h" "$(@D)/cuda/include/thrust/system/tbb/detail/scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/scan.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/scan.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/scan_by_key.h" "$(@D)/cuda/include/thrust/system/tbb/detail/scan_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/scatter.h" "$(@D)/cuda/include/thrust/system/tbb/detail/scatter.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/sequence.h" "$(@D)/cuda/include/thrust/system/tbb/detail/sequence.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/set_operations.h" "$(@D)/cuda/include/thrust/system/tbb/detail/set_operations.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/sort.h" "$(@D)/cuda/include/thrust/system/tbb/detail/sort.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/sort.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/sort.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/swap_ranges.h" "$(@D)/cuda/include/thrust/system/tbb/detail/swap_ranges.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/tabulate.h" "$(@D)/cuda/include/thrust/system/tbb/detail/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/temporary_buffer.h" "$(@D)/cuda/include/thrust/system/tbb/detail/temporary_buffer.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/transform.h" "$(@D)/cuda/include/thrust/system/tbb/detail/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/transform_reduce.h" "$(@D)/cuda/include/thrust/system/tbb/detail/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/transform_scan.h" "$(@D)/cuda/include/thrust/system/tbb/detail/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/uninitialized_copy.h" "$(@D)/cuda/include/thrust/system/tbb/detail/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/uninitialized_fill.h" "$(@D)/cuda/include/thrust/system/tbb/detail/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/unique.h" "$(@D)/cuda/include/thrust/system/tbb/detail/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/unique.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/unique.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/unique_by_key.h" "$(@D)/cuda/include/thrust/system/tbb/detail/unique_by_key.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/unique_by_key.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/unique_by_key.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/detail/vector.inl" "$(@D)/cuda/include/thrust/system/tbb/detail/vector.inl" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/execution_policy.h" "$(@D)/cuda/include/thrust/system/tbb/execution_policy.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/memory.h" "$(@D)/cuda/include/thrust/system/tbb/memory.h" && cp "/usr/local/cuda-9.0/include/thrust/system/tbb/vector.h" "$(@D)/cuda/include/thrust/system/tbb/vector.h" && cp "/usr/local/cuda-9.0/include/thrust/system_error.h" "$(@D)/cuda/include/thrust/system_error.h" && cp "/usr/local/cuda-9.0/include/thrust/tabulate.h" "$(@D)/cuda/include/thrust/tabulate.h" && cp "/usr/local/cuda-9.0/include/thrust/transform.h" "$(@D)/cuda/include/thrust/transform.h" && cp "/usr/local/cuda-9.0/include/thrust/transform_reduce.h" "$(@D)/cuda/include/thrust/transform_reduce.h" && cp "/usr/local/cuda-9.0/include/thrust/transform_scan.h" "$(@D)/cuda/include/thrust/transform_scan.h" && cp "/usr/local/cuda-9.0/include/thrust/tuple.h" "$(@D)/cuda/include/thrust/tuple.h" && cp "/usr/local/cuda-9.0/include/thrust/uninitialized_copy.h" "$(@D)/cuda/include/thrust/uninitialized_copy.h" && cp "/usr/local/cuda-9.0/include/thrust/uninitialized_fill.h" "$(@D)/cuda/include/thrust/uninitialized_fill.h" && cp "/usr/local/cuda-9.0/include/thrust/unique.h" "$(@D)/cuda/include/thrust/unique.h" && cp "/usr/local/cuda-9.0/include/thrust/version.h" "$(@D)/cuda/include/thrust/version.h" && cp "/usr/local/cuda-9.0/include/vector_functions.h" "$(@D)/cuda/include/vector_functions.h" && cp "/usr/local/cuda-9.0/include/vector_functions.hpp" "$(@D)/cuda/include/vector_functions.hpp" && cp "/usr/local/cuda-9.0/include/vector_types.h" "$(@D)/cuda/include/vector_types.h" """, ) @@ -1264,72 +1192,69 @@ genrule( name = "cuda-nvvm", outs = [ "cuda/nvvm/bin/cicc", - "cuda/nvvm/libdevice/libdevice.compute_50.10.bc", - "cuda/nvvm/libdevice/libdevice.compute_30.10.bc", - "cuda/nvvm/libdevice/libdevice.compute_20.10.bc", - "cuda/nvvm/libdevice/libdevice.compute_35.10.bc", - "cuda/nvvm/lib64/libnvvm.so.3", - "cuda/nvvm/lib64/libnvvm.so", - "cuda/nvvm/lib64/libnvvm.so.3.1.0", "cuda/nvvm/include/nvvm.h", - "cuda/nvvm/libnvvm-samples/ptxgen/README.txt", - "cuda/nvvm/libnvvm-samples/ptxgen/ptxgen.c", - "cuda/nvvm/libnvvm-samples/ptxgen/CMakeLists.txt", + "cuda/nvvm/lib64/libnvvm.so", + "cuda/nvvm/lib64/libnvvm.so.3", + "cuda/nvvm/lib64/libnvvm.so.3.2.0", + "cuda/nvvm/libdevice/libdevice.10.bc", + "cuda/nvvm/libnvvm-samples/CMakeLists.txt", + "cuda/nvvm/libnvvm-samples/README.txt", "cuda/nvvm/libnvvm-samples/build.bat", - "cuda/nvvm/libnvvm-samples/cuda-c-linking/README.txt", - "cuda/nvvm/libnvvm-samples/cuda-c-linking/math-funcs.cu", + "cuda/nvvm/libnvvm-samples/build.sh", + "cuda/nvvm/libnvvm-samples/common/include/DDSWriter.h", + "cuda/nvvm/libnvvm-samples/common/include/drvapi_error_string.h", "cuda/nvvm/libnvvm-samples/cuda-c-linking/CMakeLists.txt", + "cuda/nvvm/libnvvm-samples/cuda-c-linking/README.txt", "cuda/nvvm/libnvvm-samples/cuda-c-linking/cuda-c-linking.cpp", - "cuda/nvvm/libnvvm-samples/README.txt", - "cuda/nvvm/libnvvm-samples/simple/simple.c", - "cuda/nvvm/libnvvm-samples/simple/simple-gpu.ll", + "cuda/nvvm/libnvvm-samples/cuda-c-linking/math-funcs.cu", + "cuda/nvvm/libnvvm-samples/ptxgen/CMakeLists.txt", + "cuda/nvvm/libnvvm-samples/ptxgen/README.txt", + "cuda/nvvm/libnvvm-samples/ptxgen/ptxgen.c", + "cuda/nvvm/libnvvm-samples/simple/CMakeLists.txt", "cuda/nvvm/libnvvm-samples/simple/README.txt", + "cuda/nvvm/libnvvm-samples/simple/simple-gpu.ll", "cuda/nvvm/libnvvm-samples/simple/simple-gpu64.ll", - "cuda/nvvm/libnvvm-samples/simple/CMakeLists.txt", - "cuda/nvvm/libnvvm-samples/common/include/DDSWriter.h", - "cuda/nvvm/libnvvm-samples/common/include/drvapi_error_string.h", - "cuda/nvvm/libnvvm-samples/build.sh", - "cuda/nvvm/libnvvm-samples/CMakeLists.txt", + "cuda/nvvm/libnvvm-samples/simple/simple.c", ], cmd = """ -if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-8.0/nvvm/bin/cicc" "$(@D)/cuda/nvvm/bin/cicc" && cp "/usr/local/cuda-8.0/nvvm/libdevice/libdevice.compute_50.10.bc" "$(@D)/cuda/nvvm/libdevice/libdevice.compute_50.10.bc" && cp "/usr/local/cuda-8.0/nvvm/libdevice/libdevice.compute_30.10.bc" "$(@D)/cuda/nvvm/libdevice/libdevice.compute_30.10.bc" && cp "/usr/local/cuda-8.0/nvvm/libdevice/libdevice.compute_20.10.bc" "$(@D)/cuda/nvvm/libdevice/libdevice.compute_20.10.bc" && cp "/usr/local/cuda-8.0/nvvm/libdevice/libdevice.compute_35.10.bc" "$(@D)/cuda/nvvm/libdevice/libdevice.compute_35.10.bc" && cp "/usr/local/cuda-8.0/nvvm/lib64/libnvvm.so.3" "$(@D)/cuda/nvvm/lib64/libnvvm.so.3" && cp "/usr/local/cuda-8.0/nvvm/lib64/libnvvm.so" "$(@D)/cuda/nvvm/lib64/libnvvm.so" && cp "/usr/local/cuda-8.0/nvvm/lib64/libnvvm.so.3.1.0" "$(@D)/cuda/nvvm/lib64/libnvvm.so.3.1.0" && cp "/usr/local/cuda-8.0/nvvm/include/nvvm.h" "$(@D)/cuda/nvvm/include/nvvm.h" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/ptxgen/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/ptxgen/README.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/ptxgen/ptxgen.c" "$(@D)/cuda/nvvm/libnvvm-samples/ptxgen/ptxgen.c" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/ptxgen/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/ptxgen/CMakeLists.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/build.bat" "$(@D)/cuda/nvvm/libnvvm-samples/build.bat" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/cuda-c-linking/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/README.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/cuda-c-linking/math-funcs.cu" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/math-funcs.cu" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/cuda-c-linking/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/CMakeLists.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/cuda-c-linking/cuda-c-linking.cpp" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/cuda-c-linking.cpp" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/README.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/simple/simple.c" "$(@D)/cuda/nvvm/libnvvm-samples/simple/simple.c" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/simple/simple-gpu.ll" "$(@D)/cuda/nvvm/libnvvm-samples/simple/simple-gpu.ll" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/simple/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/simple/README.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/simple/simple-gpu64.ll" "$(@D)/cuda/nvvm/libnvvm-samples/simple/simple-gpu64.ll" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/simple/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/simple/CMakeLists.txt" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/common/include/DDSWriter.h" "$(@D)/cuda/nvvm/libnvvm-samples/common/include/DDSWriter.h" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/common/include/drvapi_error_string.h" "$(@D)/cuda/nvvm/libnvvm-samples/common/include/drvapi_error_string.h" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/build.sh" "$(@D)/cuda/nvvm/libnvvm-samples/build.sh" && cp "/usr/local/cuda-8.0/nvvm/libnvvm-samples/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/CMakeLists.txt" +if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-9.0/nvvm/bin/cicc" "$(@D)/cuda/nvvm/bin/cicc" && cp "/usr/local/cuda-9.0/nvvm/include/nvvm.h" "$(@D)/cuda/nvvm/include/nvvm.h" && cp "/usr/local/cuda-9.0/nvvm/lib64/libnvvm.so" "$(@D)/cuda/nvvm/lib64/libnvvm.so" && cp "/usr/local/cuda-9.0/nvvm/lib64/libnvvm.so.3" "$(@D)/cuda/nvvm/lib64/libnvvm.so.3" && cp "/usr/local/cuda-9.0/nvvm/lib64/libnvvm.so.3.2.0" "$(@D)/cuda/nvvm/lib64/libnvvm.so.3.2.0" && cp "/usr/local/cuda-9.0/nvvm/libdevice/libdevice.10.bc" "$(@D)/cuda/nvvm/libdevice/libdevice.10.bc" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/CMakeLists.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/README.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/build.bat" "$(@D)/cuda/nvvm/libnvvm-samples/build.bat" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/build.sh" "$(@D)/cuda/nvvm/libnvvm-samples/build.sh" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/common/include/DDSWriter.h" "$(@D)/cuda/nvvm/libnvvm-samples/common/include/DDSWriter.h" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/common/include/drvapi_error_string.h" "$(@D)/cuda/nvvm/libnvvm-samples/common/include/drvapi_error_string.h" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/cuda-c-linking/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/CMakeLists.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/cuda-c-linking/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/README.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/cuda-c-linking/cuda-c-linking.cpp" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/cuda-c-linking.cpp" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/cuda-c-linking/math-funcs.cu" "$(@D)/cuda/nvvm/libnvvm-samples/cuda-c-linking/math-funcs.cu" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/ptxgen/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/ptxgen/CMakeLists.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/ptxgen/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/ptxgen/README.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/ptxgen/ptxgen.c" "$(@D)/cuda/nvvm/libnvvm-samples/ptxgen/ptxgen.c" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/simple/CMakeLists.txt" "$(@D)/cuda/nvvm/libnvvm-samples/simple/CMakeLists.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/simple/README.txt" "$(@D)/cuda/nvvm/libnvvm-samples/simple/README.txt" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/simple/simple-gpu.ll" "$(@D)/cuda/nvvm/libnvvm-samples/simple/simple-gpu.ll" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/simple/simple-gpu64.ll" "$(@D)/cuda/nvvm/libnvvm-samples/simple/simple-gpu64.ll" && cp "/usr/local/cuda-9.0/nvvm/libnvvm-samples/simple/simple.c" "$(@D)/cuda/nvvm/libnvvm-samples/simple/simple.c" """, ) genrule( name = "cuda-extras", outs = [ - "cuda/extras/CUPTI/include/cupti_result.h", + "cuda/extras/CUPTI/include/GL/gl.h", + "cuda/extras/CUPTI/include/GL/glew.h", + "cuda/extras/CUPTI/include/GL/glext.h", + "cuda/extras/CUPTI/include/GL/glu.h", + "cuda/extras/CUPTI/include/GL/glut.h", + "cuda/extras/CUPTI/include/GL/glx.h", + "cuda/extras/CUPTI/include/GL/glxext.h", + "cuda/extras/CUPTI/include/GL/wglew.h", + "cuda/extras/CUPTI/include/GL/wglext.h", + "cuda/extras/CUPTI/include/cuda_stdint.h", + "cuda/extras/CUPTI/include/cupti.h", + "cuda/extras/CUPTI/include/cupti_activity.h", + "cuda/extras/CUPTI/include/cupti_callbacks.h", + "cuda/extras/CUPTI/include/cupti_driver_cbid.h", "cuda/extras/CUPTI/include/cupti_events.h", - "cuda/extras/CUPTI/include/openacc/cupti_openacc.h", + "cuda/extras/CUPTI/include/cupti_metrics.h", + "cuda/extras/CUPTI/include/cupti_nvtx_cbid.h", + "cuda/extras/CUPTI/include/cupti_result.h", + "cuda/extras/CUPTI/include/cupti_runtime_cbid.h", "cuda/extras/CUPTI/include/cupti_version.h", - "cuda/extras/CUPTI/include/generated_cuda_gl_interop_meta.h", + "cuda/extras/CUPTI/include/generated_cudaGL_meta.h", "cuda/extras/CUPTI/include/generated_cudaVDPAU_meta.h", - "cuda/extras/CUPTI/include/cupti_activity.h", - "cuda/extras/CUPTI/include/generated_cuda_runtime_api_meta.h", + "cuda/extras/CUPTI/include/generated_cuda_gl_interop_meta.h", "cuda/extras/CUPTI/include/generated_cuda_meta.h", - "cuda/extras/CUPTI/include/cupti_nvtx_cbid.h", - "cuda/extras/CUPTI/include/cuda_stdint.h", - "cuda/extras/CUPTI/include/generated_cudaGL_meta.h", + "cuda/extras/CUPTI/include/generated_cuda_runtime_api_meta.h", "cuda/extras/CUPTI/include/generated_cuda_vdpau_interop_meta.h", - "cuda/extras/CUPTI/include/cupti_metrics.h", - "cuda/extras/CUPTI/include/cupti_callbacks.h", - "cuda/extras/CUPTI/include/cupti_runtime_cbid.h", - "cuda/extras/CUPTI/include/cupti.h", - "cuda/extras/CUPTI/include/GL/glut.h", - "cuda/extras/CUPTI/include/GL/glu.h", - "cuda/extras/CUPTI/include/GL/glxext.h", - "cuda/extras/CUPTI/include/GL/wglext.h", - "cuda/extras/CUPTI/include/GL/glx.h", - "cuda/extras/CUPTI/include/GL/glext.h", - "cuda/extras/CUPTI/include/GL/wglew.h", - "cuda/extras/CUPTI/include/GL/gl.h", - "cuda/extras/CUPTI/include/GL/glew.h", - "cuda/extras/CUPTI/include/cupti_driver_cbid.h", "cuda/extras/CUPTI/include/generated_nvtx_meta.h", + "cuda/extras/CUPTI/include/openacc/cupti_openacc.h", ], cmd = """ -if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_result.h" "$(@D)/cuda/extras/CUPTI/include/cupti_result.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_events.h" "$(@D)/cuda/extras/CUPTI/include/cupti_events.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/openacc/cupti_openacc.h" "$(@D)/cuda/extras/CUPTI/include/openacc/cupti_openacc.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_version.h" "$(@D)/cuda/extras/CUPTI/include/cupti_version.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_cuda_gl_interop_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_gl_interop_meta.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_cudaVDPAU_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cudaVDPAU_meta.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_activity.h" "$(@D)/cuda/extras/CUPTI/include/cupti_activity.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_cuda_runtime_api_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_runtime_api_meta.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_cuda_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_meta.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_nvtx_cbid.h" "$(@D)/cuda/extras/CUPTI/include/cupti_nvtx_cbid.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cuda_stdint.h" "$(@D)/cuda/extras/CUPTI/include/cuda_stdint.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_cudaGL_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cudaGL_meta.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_cuda_vdpau_interop_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_vdpau_interop_meta.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_metrics.h" "$(@D)/cuda/extras/CUPTI/include/cupti_metrics.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_callbacks.h" "$(@D)/cuda/extras/CUPTI/include/cupti_callbacks.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_runtime_cbid.h" "$(@D)/cuda/extras/CUPTI/include/cupti_runtime_cbid.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti.h" "$(@D)/cuda/extras/CUPTI/include/cupti.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/glut.h" "$(@D)/cuda/extras/CUPTI/include/GL/glut.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/glu.h" "$(@D)/cuda/extras/CUPTI/include/GL/glu.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/glxext.h" "$(@D)/cuda/extras/CUPTI/include/GL/glxext.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/wglext.h" "$(@D)/cuda/extras/CUPTI/include/GL/wglext.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/glx.h" "$(@D)/cuda/extras/CUPTI/include/GL/glx.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/glext.h" "$(@D)/cuda/extras/CUPTI/include/GL/glext.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/wglew.h" "$(@D)/cuda/extras/CUPTI/include/GL/wglew.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/gl.h" "$(@D)/cuda/extras/CUPTI/include/GL/gl.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/GL/glew.h" "$(@D)/cuda/extras/CUPTI/include/GL/glew.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/cupti_driver_cbid.h" "$(@D)/cuda/extras/CUPTI/include/cupti_driver_cbid.h" && cp "/usr/local/cuda-8.0/extras/CUPTI/include/generated_nvtx_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_nvtx_meta.h" +if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/gl.h" "$(@D)/cuda/extras/CUPTI/include/GL/gl.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/glew.h" "$(@D)/cuda/extras/CUPTI/include/GL/glew.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/glext.h" "$(@D)/cuda/extras/CUPTI/include/GL/glext.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/glu.h" "$(@D)/cuda/extras/CUPTI/include/GL/glu.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/glut.h" "$(@D)/cuda/extras/CUPTI/include/GL/glut.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/glx.h" "$(@D)/cuda/extras/CUPTI/include/GL/glx.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/glxext.h" "$(@D)/cuda/extras/CUPTI/include/GL/glxext.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/wglew.h" "$(@D)/cuda/extras/CUPTI/include/GL/wglew.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/GL/wglext.h" "$(@D)/cuda/extras/CUPTI/include/GL/wglext.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cuda_stdint.h" "$(@D)/cuda/extras/CUPTI/include/cuda_stdint.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti.h" "$(@D)/cuda/extras/CUPTI/include/cupti.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_activity.h" "$(@D)/cuda/extras/CUPTI/include/cupti_activity.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_callbacks.h" "$(@D)/cuda/extras/CUPTI/include/cupti_callbacks.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_driver_cbid.h" "$(@D)/cuda/extras/CUPTI/include/cupti_driver_cbid.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_events.h" "$(@D)/cuda/extras/CUPTI/include/cupti_events.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_metrics.h" "$(@D)/cuda/extras/CUPTI/include/cupti_metrics.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_nvtx_cbid.h" "$(@D)/cuda/extras/CUPTI/include/cupti_nvtx_cbid.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_result.h" "$(@D)/cuda/extras/CUPTI/include/cupti_result.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_runtime_cbid.h" "$(@D)/cuda/extras/CUPTI/include/cupti_runtime_cbid.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/cupti_version.h" "$(@D)/cuda/extras/CUPTI/include/cupti_version.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_cudaGL_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cudaGL_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_cudaVDPAU_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cudaVDPAU_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_cuda_gl_interop_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_gl_interop_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_cuda_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_cuda_runtime_api_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_runtime_api_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_cuda_vdpau_interop_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_cuda_vdpau_interop_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/generated_nvtx_meta.h" "$(@D)/cuda/extras/CUPTI/include/generated_nvtx_meta.h" && cp "/usr/local/cuda-9.0/extras/CUPTI/include/openacc/cupti_openacc.h" "$(@D)/cuda/extras/CUPTI/include/openacc/cupti_openacc.h" """, ) @@ -1337,26 +1262,21 @@ genrule( name = "cuda-lib", outs = [ "cuda/lib/libcuda.so", - "cuda/lib/libcudart.so.8.0", + "cuda/lib/libcudart.so.9.0", "cuda/lib/libcudart_static.a", - "cuda/lib/libcublas.so.8.0", - "cuda/lib/libcusolver.so.8.0", - "cuda/lib/libcurand.so.8.0", - "cuda/lib/libcufft.so.8.0", - "cuda/lib/libcudnn.so.6", - "cuda/lib/libcupti.so.8.0", + "cuda/lib/libcublas.so.9.0", + "cuda/lib/libcusolver.so.9.0", + "cuda/lib/libcurand.so.9.0", + "cuda/lib/libcufft.so.9.0", + "cuda/lib/libcudnn.so.7", + "cuda/lib/libcupti.so.9.0", ], cmd = """ -if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/stubs/libcuda.so" "$(@D)/cuda/lib/libcuda.so" && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/libcudart.so.8.0.61" "$(@D)/cuda/lib/libcudart.so.8.0" && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/libcudart_static.a" "$(@D)/cuda/lib/libcudart_static.a" && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/libcublas.so.8.0.88" "$(@D)/cuda/lib/libcublas.so.8.0" && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/libcusolver.so.8.0.61" "$(@D)/cuda/lib/libcusolver.so.8.0" && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/libcurand.so.8.0.61" "$(@D)/cuda/lib/libcurand.so.8.0" && cp "/usr/local/cuda-8.0/targets/x86_64-linux/lib/libcufft.so.8.0.61" "$(@D)/cuda/lib/libcufft.so.8.0" && cp "/usr/lib/x86_64-linux-gnu/libcudnn.so.6.0.21" "$(@D)/cuda/lib/libcudnn.so.6" && cp "/usr/local/cuda-8.0/extras/CUPTI/lib64/libcupti.so.8.0.61" "$(@D)/cuda/lib/libcupti.so.8.0" +if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs/libcuda.so" "$(@D)/cuda/lib/libcuda.so" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcudart.so.9.0" "$(@D)/cuda/lib/libcudart.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcudart_static.a" "$(@D)/cuda/lib/libcudart_static.a" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcublas.so.9.0" "$(@D)/cuda/lib/libcublas.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcusolver.so.9.0" "$(@D)/cuda/lib/libcusolver.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcurand.so.9.0" "$(@D)/cuda/lib/libcurand.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcufft.so.9.0" "$(@D)/cuda/lib/libcufft.so.9.0" && cp "/usr/lib/x86_64-linux-gnu/libcudnn.so.7" "$(@D)/cuda/lib/libcudnn.so.7" && cp "/usr/local/cuda-9.0/extras/CUPTI/lib64/libcupti.so.9.0" "$(@D)/cuda/lib/libcupti.so.9.0" """, ) -genrule( +filegroup( name = "cudnn-include", - outs = [ - "cuda/include/cudnn.h", - ], - cmd = """ -if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/include/cudnn.h" "$(@D)/cudnn.h" - """, + srcs = [], ) diff --git a/third_party/toolchains/gpus/py/BUILD b/third_party/toolchains/gpus/py/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..2d5ace93ff5054927cda61b0302af4edd8fe56c1 --- /dev/null +++ b/third_party/toolchains/gpus/py/BUILD @@ -0,0 +1,171 @@ +# A build file to configure python remote repository used with Bazel remote +# execution service +# DO NOT EDIT: automatically generated BUILD file + +licenses(["restricted"]) + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "python_headers", + hdrs = [":python_include"], + data = select({ + ":windows": [":python_import_lib"], + "//conditions:default": [], + }), + includes = ["python_include"], + linkopts = select({ + # TODO(pcloudy): Ideally, this should just go into deps after resolving + # https://github.com/bazelbuild/bazel/issues/3237, + ":windows": ["$(locations :python_import_lib)"], + "//conditions:default": [], + }), +) + +cc_library( + name = "numpy_headers", + hdrs = [":numpy_include"], + includes = ["numpy_include"], +) + +config_setting( + name = "windows", + values = {"cpu": "x64_windows"}, + visibility = ["//visibility:public"], +) + +genrule( + name = "python_include", + outs = [ + "python_include/Python-ast.h", + "python_include/Python.h", + "python_include/abstract.h", + "python_include/asdl.h", + "python_include/ast.h", + "python_include/bitset.h", + "python_include/boolobject.h", + "python_include/bufferobject.h", + "python_include/bytearrayobject.h", + "python_include/bytes_methods.h", + "python_include/bytesobject.h", + "python_include/cStringIO.h", + "python_include/cellobject.h", + "python_include/ceval.h", + "python_include/classobject.h", + "python_include/cobject.h", + "python_include/code.h", + "python_include/codecs.h", + "python_include/compile.h", + "python_include/complexobject.h", + "python_include/datetime.h", + "python_include/descrobject.h", + "python_include/dictobject.h", + "python_include/dtoa.h", + "python_include/enumobject.h", + "python_include/errcode.h", + "python_include/eval.h", + "python_include/fileobject.h", + "python_include/floatobject.h", + "python_include/frameobject.h", + "python_include/funcobject.h", + "python_include/genobject.h", + "python_include/graminit.h", + "python_include/grammar.h", + "python_include/import.h", + "python_include/intobject.h", + "python_include/intrcheck.h", + "python_include/iterobject.h", + "python_include/listobject.h", + "python_include/longintrepr.h", + "python_include/longobject.h", + "python_include/marshal.h", + "python_include/memoryobject.h", + "python_include/metagrammar.h", + "python_include/methodobject.h", + "python_include/modsupport.h", + "python_include/moduleobject.h", + "python_include/node.h", + "python_include/object.h", + "python_include/objimpl.h", + "python_include/opcode.h", + "python_include/osdefs.h", + "python_include/parsetok.h", + "python_include/patchlevel.h", + "python_include/pgen.h", + "python_include/pgenheaders.h", + "python_include/py_curses.h", + "python_include/pyarena.h", + "python_include/pycapsule.h", + "python_include/pyconfig.h", + "python_include/pyctype.h", + "python_include/pydebug.h", + "python_include/pyerrors.h", + "python_include/pyexpat.h", + "python_include/pyfpe.h", + "python_include/pygetopt.h", + "python_include/pymacconfig.h", + "python_include/pymactoolbox.h", + "python_include/pymath.h", + "python_include/pymem.h", + "python_include/pyport.h", + "python_include/pystate.h", + "python_include/pystrcmp.h", + "python_include/pystrtod.h", + "python_include/pythonrun.h", + "python_include/pythread.h", + "python_include/rangeobject.h", + "python_include/setobject.h", + "python_include/sliceobject.h", + "python_include/stringobject.h", + "python_include/structmember.h", + "python_include/structseq.h", + "python_include/symtable.h", + "python_include/sysmodule.h", + "python_include/timefuncs.h", + "python_include/token.h", + "python_include/traceback.h", + "python_include/tupleobject.h", + "python_include/ucnhash.h", + "python_include/unicodeobject.h", + "python_include/warnings.h", + "python_include/weakrefobject.h", + ], + cmd = """ +cp "/usr/include/python2.7/Python-ast.h" "$(@D)/python_include/Python-ast.h" && cp "/usr/include/python2.7/Python.h" "$(@D)/python_include/Python.h" && cp "/usr/include/python2.7/abstract.h" "$(@D)/python_include/abstract.h" && cp "/usr/include/python2.7/asdl.h" "$(@D)/python_include/asdl.h" && cp "/usr/include/python2.7/ast.h" "$(@D)/python_include/ast.h" && cp "/usr/include/python2.7/bitset.h" "$(@D)/python_include/bitset.h" && cp "/usr/include/python2.7/boolobject.h" "$(@D)/python_include/boolobject.h" && cp "/usr/include/python2.7/bufferobject.h" "$(@D)/python_include/bufferobject.h" && cp "/usr/include/python2.7/bytearrayobject.h" "$(@D)/python_include/bytearrayobject.h" && cp "/usr/include/python2.7/bytes_methods.h" "$(@D)/python_include/bytes_methods.h" && cp "/usr/include/python2.7/bytesobject.h" "$(@D)/python_include/bytesobject.h" && cp "/usr/include/python2.7/cStringIO.h" "$(@D)/python_include/cStringIO.h" && cp "/usr/include/python2.7/cellobject.h" "$(@D)/python_include/cellobject.h" && cp "/usr/include/python2.7/ceval.h" "$(@D)/python_include/ceval.h" && cp "/usr/include/python2.7/classobject.h" "$(@D)/python_include/classobject.h" && cp "/usr/include/python2.7/cobject.h" "$(@D)/python_include/cobject.h" && cp "/usr/include/python2.7/code.h" "$(@D)/python_include/code.h" && cp "/usr/include/python2.7/codecs.h" "$(@D)/python_include/codecs.h" && cp "/usr/include/python2.7/compile.h" "$(@D)/python_include/compile.h" && cp "/usr/include/python2.7/complexobject.h" "$(@D)/python_include/complexobject.h" && cp "/usr/include/python2.7/datetime.h" "$(@D)/python_include/datetime.h" && cp "/usr/include/python2.7/descrobject.h" "$(@D)/python_include/descrobject.h" && cp "/usr/include/python2.7/dictobject.h" "$(@D)/python_include/dictobject.h" && cp "/usr/include/python2.7/dtoa.h" "$(@D)/python_include/dtoa.h" && cp "/usr/include/python2.7/enumobject.h" "$(@D)/python_include/enumobject.h" && cp "/usr/include/python2.7/errcode.h" "$(@D)/python_include/errcode.h" && cp "/usr/include/python2.7/eval.h" "$(@D)/python_include/eval.h" && cp "/usr/include/python2.7/fileobject.h" "$(@D)/python_include/fileobject.h" && cp "/usr/include/python2.7/floatobject.h" "$(@D)/python_include/floatobject.h" && cp "/usr/include/python2.7/frameobject.h" "$(@D)/python_include/frameobject.h" && cp "/usr/include/python2.7/funcobject.h" "$(@D)/python_include/funcobject.h" && cp "/usr/include/python2.7/genobject.h" "$(@D)/python_include/genobject.h" && cp "/usr/include/python2.7/graminit.h" "$(@D)/python_include/graminit.h" && cp "/usr/include/python2.7/grammar.h" "$(@D)/python_include/grammar.h" && cp "/usr/include/python2.7/import.h" "$(@D)/python_include/import.h" && cp "/usr/include/python2.7/intobject.h" "$(@D)/python_include/intobject.h" && cp "/usr/include/python2.7/intrcheck.h" "$(@D)/python_include/intrcheck.h" && cp "/usr/include/python2.7/iterobject.h" "$(@D)/python_include/iterobject.h" && cp "/usr/include/python2.7/listobject.h" "$(@D)/python_include/listobject.h" && cp "/usr/include/python2.7/longintrepr.h" "$(@D)/python_include/longintrepr.h" && cp "/usr/include/python2.7/longobject.h" "$(@D)/python_include/longobject.h" && cp "/usr/include/python2.7/marshal.h" "$(@D)/python_include/marshal.h" && cp "/usr/include/python2.7/memoryobject.h" "$(@D)/python_include/memoryobject.h" && cp "/usr/include/python2.7/metagrammar.h" "$(@D)/python_include/metagrammar.h" && cp "/usr/include/python2.7/methodobject.h" "$(@D)/python_include/methodobject.h" && cp "/usr/include/python2.7/modsupport.h" "$(@D)/python_include/modsupport.h" && cp "/usr/include/python2.7/moduleobject.h" "$(@D)/python_include/moduleobject.h" && cp "/usr/include/python2.7/node.h" "$(@D)/python_include/node.h" && cp "/usr/include/python2.7/object.h" "$(@D)/python_include/object.h" && cp "/usr/include/python2.7/objimpl.h" "$(@D)/python_include/objimpl.h" && cp "/usr/include/python2.7/opcode.h" "$(@D)/python_include/opcode.h" && cp "/usr/include/python2.7/osdefs.h" "$(@D)/python_include/osdefs.h" && cp "/usr/include/python2.7/parsetok.h" "$(@D)/python_include/parsetok.h" && cp "/usr/include/python2.7/patchlevel.h" "$(@D)/python_include/patchlevel.h" && cp "/usr/include/python2.7/pgen.h" "$(@D)/python_include/pgen.h" && cp "/usr/include/python2.7/pgenheaders.h" "$(@D)/python_include/pgenheaders.h" && cp "/usr/include/python2.7/py_curses.h" "$(@D)/python_include/py_curses.h" && cp "/usr/include/python2.7/pyarena.h" "$(@D)/python_include/pyarena.h" && cp "/usr/include/python2.7/pycapsule.h" "$(@D)/python_include/pycapsule.h" && cp "/usr/include/python2.7/pyconfig.h" "$(@D)/python_include/pyconfig.h" && cp "/usr/include/python2.7/pyctype.h" "$(@D)/python_include/pyctype.h" && cp "/usr/include/python2.7/pydebug.h" "$(@D)/python_include/pydebug.h" && cp "/usr/include/python2.7/pyerrors.h" "$(@D)/python_include/pyerrors.h" && cp "/usr/include/python2.7/pyexpat.h" "$(@D)/python_include/pyexpat.h" && cp "/usr/include/python2.7/pyfpe.h" "$(@D)/python_include/pyfpe.h" && cp "/usr/include/python2.7/pygetopt.h" "$(@D)/python_include/pygetopt.h" && cp "/usr/include/python2.7/pymacconfig.h" "$(@D)/python_include/pymacconfig.h" && cp "/usr/include/python2.7/pymactoolbox.h" "$(@D)/python_include/pymactoolbox.h" && cp "/usr/include/python2.7/pymath.h" "$(@D)/python_include/pymath.h" && cp "/usr/include/python2.7/pymem.h" "$(@D)/python_include/pymem.h" && cp "/usr/include/python2.7/pyport.h" "$(@D)/python_include/pyport.h" && cp "/usr/include/python2.7/pystate.h" "$(@D)/python_include/pystate.h" && cp "/usr/include/python2.7/pystrcmp.h" "$(@D)/python_include/pystrcmp.h" && cp "/usr/include/python2.7/pystrtod.h" "$(@D)/python_include/pystrtod.h" && cp "/usr/include/python2.7/pythonrun.h" "$(@D)/python_include/pythonrun.h" && cp "/usr/include/python2.7/pythread.h" "$(@D)/python_include/pythread.h" && cp "/usr/include/python2.7/rangeobject.h" "$(@D)/python_include/rangeobject.h" && cp "/usr/include/python2.7/setobject.h" "$(@D)/python_include/setobject.h" && cp "/usr/include/python2.7/sliceobject.h" "$(@D)/python_include/sliceobject.h" && cp "/usr/include/python2.7/stringobject.h" "$(@D)/python_include/stringobject.h" && cp "/usr/include/python2.7/structmember.h" "$(@D)/python_include/structmember.h" && cp "/usr/include/python2.7/structseq.h" "$(@D)/python_include/structseq.h" && cp "/usr/include/python2.7/symtable.h" "$(@D)/python_include/symtable.h" && cp "/usr/include/python2.7/sysmodule.h" "$(@D)/python_include/sysmodule.h" && cp "/usr/include/python2.7/timefuncs.h" "$(@D)/python_include/timefuncs.h" && cp "/usr/include/python2.7/token.h" "$(@D)/python_include/token.h" && cp "/usr/include/python2.7/traceback.h" "$(@D)/python_include/traceback.h" && cp "/usr/include/python2.7/tupleobject.h" "$(@D)/python_include/tupleobject.h" && cp "/usr/include/python2.7/ucnhash.h" "$(@D)/python_include/ucnhash.h" && cp "/usr/include/python2.7/unicodeobject.h" "$(@D)/python_include/unicodeobject.h" && cp "/usr/include/python2.7/warnings.h" "$(@D)/python_include/warnings.h" && cp "/usr/include/python2.7/weakrefobject.h" "$(@D)/python_include/weakrefobject.h" + """, +) + +genrule( + name = "numpy_include", + outs = [ + "numpy_include/numpy/__multiarray_api.h", + "numpy_include/numpy/__ufunc_api.h", + "numpy_include/numpy/_neighborhood_iterator_imp.h", + "numpy_include/numpy/_numpyconfig.h", + "numpy_include/numpy/arrayobject.h", + "numpy_include/numpy/arrayscalars.h", + "numpy_include/numpy/halffloat.h", + "numpy_include/numpy/multiarray_api.txt", + "numpy_include/numpy/ndarrayobject.h", + "numpy_include/numpy/ndarraytypes.h", + "numpy_include/numpy/noprefix.h", + "numpy_include/numpy/npy_1_7_deprecated_api.h", + "numpy_include/numpy/npy_3kcompat.h", + "numpy_include/numpy/npy_common.h", + "numpy_include/numpy/npy_cpu.h", + "numpy_include/numpy/npy_endian.h", + "numpy_include/numpy/npy_interrupt.h", + "numpy_include/numpy/npy_math.h", + "numpy_include/numpy/npy_no_deprecated_api.h", + "numpy_include/numpy/npy_os.h", + "numpy_include/numpy/numpyconfig.h", + "numpy_include/numpy/old_defines.h", + "numpy_include/numpy/oldnumeric.h", + "numpy_include/numpy/ufunc_api.txt", + "numpy_include/numpy/ufuncobject.h", + "numpy_include/numpy/utils.h", + ], + cmd = """ +cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/__multiarray_api.h" "$(@D)/numpy_include/numpy/__multiarray_api.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/__ufunc_api.h" "$(@D)/numpy_include/numpy/__ufunc_api.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/_neighborhood_iterator_imp.h" "$(@D)/numpy_include/numpy/_neighborhood_iterator_imp.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/_numpyconfig.h" "$(@D)/numpy_include/numpy/_numpyconfig.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/arrayobject.h" "$(@D)/numpy_include/numpy/arrayobject.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/arrayscalars.h" "$(@D)/numpy_include/numpy/arrayscalars.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/halffloat.h" "$(@D)/numpy_include/numpy/halffloat.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/multiarray_api.txt" "$(@D)/numpy_include/numpy/multiarray_api.txt" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/ndarrayobject.h" "$(@D)/numpy_include/numpy/ndarrayobject.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/ndarraytypes.h" "$(@D)/numpy_include/numpy/ndarraytypes.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/noprefix.h" "$(@D)/numpy_include/numpy/noprefix.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_1_7_deprecated_api.h" "$(@D)/numpy_include/numpy/npy_1_7_deprecated_api.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_3kcompat.h" "$(@D)/numpy_include/numpy/npy_3kcompat.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_common.h" "$(@D)/numpy_include/numpy/npy_common.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_cpu.h" "$(@D)/numpy_include/numpy/npy_cpu.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_endian.h" "$(@D)/numpy_include/numpy/npy_endian.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_interrupt.h" "$(@D)/numpy_include/numpy/npy_interrupt.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_math.h" "$(@D)/numpy_include/numpy/npy_math.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_no_deprecated_api.h" "$(@D)/numpy_include/numpy/npy_no_deprecated_api.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/npy_os.h" "$(@D)/numpy_include/numpy/npy_os.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/numpyconfig.h" "$(@D)/numpy_include/numpy/numpyconfig.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/old_defines.h" "$(@D)/numpy_include/numpy/old_defines.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/oldnumeric.h" "$(@D)/numpy_include/numpy/oldnumeric.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/ufunc_api.txt" "$(@D)/numpy_include/numpy/ufunc_api.txt" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/ufuncobject.h" "$(@D)/numpy_include/numpy/ufuncobject.h" && cp "/usr/local/lib/python2.7/dist-packages/numpy/core/include/numpy/utils.h" "$(@D)/numpy_include/numpy/utils.h" + """, +)